scripts/convert_checkpoint.py

#!/usr/bin/env python3
"""Convert a legacy SALMONN-2 Trainer checkpoint into a Hugging Face checkpoint.

The converter operates one safetensors shard at a time, merges Qwen LoRA weights,
rewrites legacy PEFT parameter names, copies tokenizer and remote-code assets, and
deliberately excludes DeepSpeed, optimizer, scheduler, RNG, and Trainer state.
"""

import argparse
import json
import shutil
from pathlib import Path

import torch
from safetensors import safe_open
from safetensors.torch import save_file


TOKENIZER_FILES = (
    "added_tokens.json",
    "chat_template.jinja",
    "merges.txt",
    "special_tokens_map.json",
    "tokenizer.json",
    "tokenizer_config.json",
    "vocab.json",
)


def read_json(path):
    with path.open(encoding="utf-8") as handle:
        return json.load(handle)


class ShardReader:
    def __init__(self, checkpoint, weight_map):
        self.checkpoint = checkpoint
        self.weight_map = weight_map

    def tensor(self, name):
        shard = self.weight_map[name]
        with safe_open(self.checkpoint / shard, framework="pt", device="cpu") as handle:
            return handle.get_tensor(name)


def merge_and_rename(name, tensor, reader, lora_scale):
    if ".lora_A." in name or ".lora_B." in name:
        return None, None

    if name.endswith(".base_layer.weight"):
        prefix = name[: -len(".base_layer.weight")]
        a_name = prefix + ".lora_A.default.weight"
        b_name = prefix + ".lora_B.default.weight"
        present = (a_name in reader.weight_map, b_name in reader.weight_map)
        if any(present) and not all(present):
            raise ValueError(f"Incomplete LoRA pair for {prefix}: A={present[0]}, B={present[1]}")
        if all(present):
            a = reader.tensor(a_name).float()
            b = reader.tensor(b_name).float()
            tensor = (tensor.float() + (b @ a) * lora_scale).to(tensor.dtype)
        name = prefix + ".weight"

    legacy_prefix = "base_llm.base_model.model."
    if name.startswith(legacy_prefix):
        name = "base_llm." + name[len(legacy_prefix) :]
    return name, tensor.contiguous()


def copy_remote_code(repo_root, output):
    for filename in ("configuration_salmonn.py", "audio.py"):
        shutil.copy2(repo_root / "salmonn" / filename, output / filename)
    modeling = (repo_root / "salmonn" / "modeling_salmonn.py").read_text(encoding="utf-8")
    modeling = modeling.replace("from .zipformer.model import", "from .zipformer_model import")
    modeling = modeling.replace("from .zipformer.scaling import", "from .zipformer_scaling import")
    modeling = modeling.replace("from .zipformer.subsampling import", "from .zipformer_subsampling import")
    modeling = modeling.replace("from .zipformer.zipformer_layerwise import", "from .zipformer_layerwise import")
    (output / "modeling_salmonn.py").write_text(modeling, encoding="utf-8")

    flattened = {
        "model.py": "zipformer_model.py",
        "scaling.py": "zipformer_scaling.py",
        "subsampling.py": "zipformer_subsampling.py",
        "zipformer_layerwise.py": "zipformer_layerwise.py",
    }
    for source_name, output_name in flattened.items():
        content = (repo_root / "salmonn" / "zipformer" / source_name).read_text(encoding="utf-8")
        content = content.replace("from .scaling import", "from .zipformer_scaling import")
        (output / output_name).write_text(content, encoding="utf-8")
    (output / "__init__.py").write_text(
        "from .configuration_salmonn import SalmonnConfig\n"
        "from .modeling_salmonn import SalmonnForConditionalGeneration\n",
        encoding="utf-8",
    )


def build_config(qwen_config, model_args):
    qwen_config = dict(qwen_config)
    qwen_config.pop("architectures", None)
    qwen_config.pop("_name_or_path", None)
    return {
        "model_type": "salmonn_2",
        "architectures": ["SalmonnForConditionalGeneration"],
        "auto_map": {
            "AutoConfig": "configuration_salmonn.SalmonnConfig",
            "AutoModel": "modeling_salmonn.SalmonnForConditionalGeneration",
            "AutoModelForCausalLM": "modeling_salmonn.SalmonnForConditionalGeneration",
        },
        "qwen_config": qwen_config,
        "zipformer_checkpoint": None,
        "freeze_audio_encoder": True,
        "connector_hidden_size": model_args["connector_hid_size"],
        "connector_segment_size": model_args["connector_seg_size"],
        "concatenate_encoder_layers": model_args["concat_encoder_features"],
        "inject_temporal_embedding_nl": model_args.get("inject_temporal_embedding_nl", False),
        "temporal_granularity": model_args.get("temporal_granularity", 2.0),
        "encoder_frame_rate": model_args.get("encoder_frame_rate", 50),
        "dtype": qwen_config.get("dtype", qwen_config.get("torch_dtype", "bfloat16")),
        "transformers_version": qwen_config.get("transformers_version"),
    }


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", required=True, type=Path, help="Legacy checkpoint-N directory")
    parser.add_argument("--output", required=True, type=Path, help="New, empty output directory")
    parser.add_argument(
        "--training-config",
        type=Path,
        help="Experiment config.json containing model_args; defaults to INPUT/../config.json",
    )
    parser.add_argument("--lora-r", type=int, help="Override the saved LoRA rank")
    parser.add_argument("--lora-alpha", type=float, help="Override the saved LoRA alpha")
    args = parser.parse_args()

    checkpoint = args.input.resolve()
    output = args.output.resolve()
    training_config_path = (args.training_config or checkpoint.parent / "config.json").resolve()
    for required in (checkpoint / "config.json", checkpoint / "model.safetensors.index.json", training_config_path):
        if not required.is_file():
            raise FileNotFoundError(required)
    if output.exists() and any(output.iterdir()):
        raise FileExistsError(f"Output directory must be empty: {output}")
    output.mkdir(parents=True, exist_ok=True)

    training_config = read_json(training_config_path)
    model_args = training_config["model_args"]
    if model_args.get("encoder_type") != "zipformer2" or model_args.get("llm_type") != "Qwen":
        raise ValueError("This converter only supports the released Zipformer2 + Qwen model")
    if model_args.get("use_reasoning_network") or model_args.get("num_pause_steps", 0):
        raise ValueError("Refusing to discard an enabled reasoning network or pause embeddings")
    if model_args.get("encoder_lora"):
        raise ValueError("Encoder LoRA conversion is not implemented")

    lora_r = args.lora_r or model_args.get("lora_rank")
    lora_alpha = args.lora_alpha if args.lora_alpha is not None else model_args.get("lora_alpha")
    if not model_args.get("lora") or not lora_r or lora_alpha is None:
        raise ValueError("The saved model_args do not describe a LoRA checkpoint")
    lora_scale = float(lora_alpha) / int(lora_r)

    old_index = read_json(checkpoint / "model.safetensors.index.json")
    weight_map = old_index["weight_map"]
    reader = ShardReader(checkpoint, weight_map)
    new_weight_map = {}
    total_size = 0
    merged_pairs = 0

    for shard_name in sorted(set(weight_map.values())):
        destination_tensors = {}
        with safe_open(checkpoint / shard_name, framework="pt", device="cpu") as source:
            metadata = source.metadata()
            for old_name in source.keys():
                tensor = source.get_tensor(old_name)
                new_name, tensor = merge_and_rename(old_name, tensor, reader, lora_scale)
                if new_name is None:
                    continue
                if old_name.endswith(".base_layer.weight"):
                    prefix = old_name[: -len(".base_layer.weight")]
                    if prefix + ".lora_A.default.weight" in weight_map:
                        merged_pairs += 1
                if new_name in new_weight_map or new_name in destination_tensors:
                    raise ValueError(f"Parameter-name collision after conversion: {new_name}")
                destination_tensors[new_name] = tensor
                new_weight_map[new_name] = shard_name
                total_size += tensor.numel() * tensor.element_size()
        save_file(destination_tensors, output / shard_name, metadata=metadata or {"format": "pt"})
        print(f"Converted {shard_name}: {len(destination_tensors)} tensors")

    if merged_pairs == 0:
        raise RuntimeError("No LoRA modules were merged; refusing to produce a misleading export")
    with (output / "model.safetensors.index.json").open("w", encoding="utf-8") as handle:
        json.dump({"metadata": {"total_size": total_size}, "weight_map": new_weight_map}, handle, indent=2)
        handle.write("\n")

    qwen_config = read_json(checkpoint / "config.json")
    with (output / "config.json").open("w", encoding="utf-8") as handle:
        json.dump(build_config(qwen_config, model_args), handle, indent=2)
        handle.write("\n")
    generation_config = {
        "_from_model_config": True,
        "bos_token_id": qwen_config.get("bos_token_id"),
        "eos_token_id": qwen_config.get("eos_token_id"),
        "pad_token_id": qwen_config.get("pad_token_id"),
        "transformers_version": qwen_config.get("transformers_version"),
    }
    generation_config = {key: value for key, value in generation_config.items() if value is not None}
    with (output / "generation_config.json").open("w", encoding="utf-8") as handle:
        json.dump(generation_config, handle, indent=2)
        handle.write("\n")

    for filename in TOKENIZER_FILES:
        source = checkpoint / filename
        if source.is_file():
            shutil.copy2(source, output / filename)
    repo_root = Path(__file__).resolve().parents[1]
    copy_remote_code(repo_root, output)
    shutil.copy2(repo_root / "LICENSE", output / "LICENSE")
    (output / "README.md").write_text(
        "# SALMONN-2 inference checkpoint\n\n"
        "This checkpoint contains merged Qwen LoRA weights, Zipformer2, the audio connector, "
        "tokenizer assets, and pinned custom model code. Load it with `trust_remote_code=True`.\n",
        encoding="utf-8",
    )
    print(f"Merged {merged_pairs} LoRA modules with scale alpha/r={lora_scale:g}")
    print(f"Exported {len(new_weight_map)} tensors ({total_size / 1024**3:.2f} GiB) to {output}")


if __name__ == "__main__":
    main()