convert_hf_to_nanotron.py

# ruff: noqa: E402
"""
Converts a HF model to a Nanotron model

Command:
    torchrun --nproc_per_node=1 convert_hf_to_nanotron.py --inp_path state-spaces/mamba-130m-hf --out_path nanotron_weights
"""
import argparse
import json
from dataclasses import asdict
from pathlib import Path
from typing import Dict

import torch
import yaml
from config import MambaConfig, MambaInit, MambaModelConfig
from mamba import MambaForTraining
from nanotron import logging
from nanotron.config import (
    AllForwardAllBackwardPipelineEngine,
    GeneralArgs,
    LoggingArgs,
    ModelArgs,
    ParallelismArgs,
    TensorParallelLinearMode,
    TokenizerArgs,
)
from nanotron.distributed import dist
from nanotron.logging import log_rank, set_ranks_logging_level
from nanotron.models import build_model
from nanotron.parallel import ParallelContext
from nanotron.parallel.parameters import NanotronParameter, sanity_check
from nanotron.serialize import save_meta, save_weights
from nanotron.trainer import mark_tied_parameters
from tqdm import tqdm
from transformers import MambaConfig as HFMambaConfig
from transformers import MambaForCausalLM
from transformers.utils import CONFIG_NAME
from transformers.utils.hub import cached_file

logger = logging.get_logger(__name__)


def load_config_hf(model_name):
    resolved_archive_file = cached_file(model_name, CONFIG_NAME, _raise_exceptions_for_missing_entries=False)
    return json.load(open(resolved_archive_file))


def get_weight_from_hf(
    name: str,
    ref_module_state_dict: Dict[str, torch.Tensor],
    ref_module: MambaForCausalLM,
    nanotron_to_hf: Dict[str, str],
    get_grad: bool = False,
    param_is_tp_sharded: bool = False,
) -> torch.Tensor:
    """From our brrr implementation, we get the equivalent tensor in transformers implementation"""

    def _interleave_pattern(N):
        """
        interleave_pattern(4) -> [0, 2, 1, 3]
        interleave_pattern(8) -> [0, 4, 1, 5, 2, 6, 3, 7]
        """
        assert N % 2 == 0, "N must be even"
        pattern = []
        for i in range(N // 2):
            pattern.append(i)
            pattern.append(i + N // 2)
        return pattern

    hf_name = nanotron_to_hf[name]

    if get_grad is False:

        def _get_tensor(path: str):
            return ref_module_state_dict[path]

    else:

        def _get_tensor(path: str):
            param = ref_module.get_parameter(path)
            return param.grad

    param = _get_tensor(hf_name)

    if "in_proj" in hf_name:
        # In Nanotron, we do tensor parallel column so weight need to be split in the column dimension (i.e: xz.view(...))
        # However, the HF weights was trained such that it expected xz.chunk(...) to split the tensor in the row dimension
        # Thus, we need to interleaved the HF weights to make it compatible with Nanotron
        log_rank(
            f"Interleaving {hf_name} to make it compatible with Nanotron", logger=logger, level=logging.INFO, rank=0
        )
        return param[_interleave_pattern(param.shape[0]), :]

    return param


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Convert HF weights from states-space repo to brrr weights")
    parser.add_argument("--inp_path", type=str, default="state-spaces/mamba-130m-hf")
    parser.add_argument("--out_path", type=str, default="nanotron_weight")
    parser.add_argument("--dp", type=int, default=1)
    parser.add_argument("--pp", type=int, default=1)
    parser.add_argument("--tp", type=int, default=1)
    args = parser.parse_args()

    out_path = Path(args.out_path)

    parallel_config = ParallelismArgs(
        dp=args.dp,
        pp=args.pp,
        tp=args.tp,
        pp_engine=AllForwardAllBackwardPipelineEngine(),
        tp_mode=TensorParallelLinearMode.ALL_REDUCE,
        tp_linear_async_communication=False,
    )
    assert (
        parallel_config.tp_mode == TensorParallelLinearMode.ALL_REDUCE
        and parallel_config.tp_linear_async_communication is False
    )

    parallel_context = ParallelContext(
        data_parallel_size=parallel_config.dp,
        pipeline_parallel_size=parallel_config.pp,
        tensor_parallel_size=parallel_config.tp,
    )

    # Set log log levels
    logging_config = LoggingArgs(
        log_level="info",
        log_level_replica="info",
    )

    # Set log levels
    set_ranks_logging_level(parallel_context=parallel_context, logging_config=logging_config)

    hf_config = HFMambaConfig.from_pretrained(args.inp_path)

    dtype_str = "float32"

    # TODO(fmom): Add support for ssm_cfg
    yaml_content = f"""
    is_mamba_config: true
    d_model: {hf_config.hidden_size}
    dtype: {dtype_str}
    fused_add_norm: true
    is_mamba_config: true
    num_hidden_layers: {hf_config.num_hidden_layers}
    pad_token_id: null
    pad_vocab_size_multiple: 8
    residual_in_fp32: true
    rms_norm: true
    rms_norm_eps: 1.0e-05
    ssm_cfg: null
    vocab_size: {hf_config.vocab_size}
    """

    dtype = getattr(torch, dtype_str)
    device = torch.device("cuda")

    attrs = yaml.safe_load(yaml_content)
    model_config = MambaModelConfig(**attrs)

    model_ref = MambaForCausalLM.from_pretrained(args.inp_path)
    model_ref.to(device, dtype=dtype)
    model_ref.eval()

    nanotron_model = build_model(
        model_builder=lambda: MambaForTraining(
            config=model_config,
            parallel_context=parallel_context,
            parallel_config=parallel_config,
            random_states=None,
        ),
        parallel_context=parallel_context,
        dtype=dtype,
        device=device,
    )

    device_map = {}
    current_pp_rank = dist.get_rank(parallel_context.pp_pg)

    tied_embs_ranks = [nanotron_model.model.token_position_embeddings.rank, nanotron_model.model.lm_head.rank]

    device_map["backbone.embedding"] = (
        nanotron_model.model.token_position_embeddings.rank if current_pp_rank in tied_embs_ranks else "meta"
    )

    for i in range(model_config.num_hidden_layers):
        device_map[f"backbone.layers[{i}]"] = (
            nanotron_model.model.decoder[i].rank if current_pp_rank == nanotron_model.model.decoder[i].rank else "meta"
        )

    device_map["lm_head"] = nanotron_model.model.lm_head.rank if current_pp_rank in tied_embs_ranks else "meta"

    # Get mapping of Nanotron layer to HF layer
    nanotron_to_hf = {}

    # Static mappings
    nanotron_to_hf["token_position_embeddings.pp_block.token_embedding.weight"] = "backbone.embeddings.weight"
    nanotron_to_hf["final_layer_norm.pp_block.weight"] = "backbone.norm_f.weight"
    nanotron_to_hf["lm_head.pp_block.weight"] = "lm_head.weight"

    # Dynamic mappings within a loop
    for i in range(model_config.num_hidden_layers):
        nanotron_to_hf[f"decoder.{i}.pp_block.mixer.A_log"] = f"backbone.layers.{i}.mixer.A_log"
        nanotron_to_hf[f"decoder.{i}.pp_block.mixer.D"] = f"backbone.layers.{i}.mixer.D"
        nanotron_to_hf[f"decoder.{i}.pp_block.mixer.in_proj.weight"] = f"backbone.layers.{i}.mixer.in_proj.weight"
        nanotron_to_hf[f"decoder.{i}.pp_block.mixer.conv1d.weight"] = f"backbone.layers.{i}.mixer.conv1d.weight"
        nanotron_to_hf[f"decoder.{i}.pp_block.mixer.conv1d.bias"] = f"backbone.layers.{i}.mixer.conv1d.bias"
        nanotron_to_hf[f"decoder.{i}.pp_block.mixer.x_proj.weight"] = f"backbone.layers.{i}.mixer.x_proj.weight"
        nanotron_to_hf[f"decoder.{i}.pp_block.mixer.x_proj.bias"] = f"backbone.layers.{i}.mixer.x_proj.bias"
        nanotron_to_hf[f"decoder.{i}.pp_block.mixer.dt_proj.weight"] = f"backbone.layers.{i}.mixer.dt_proj.weight"
        nanotron_to_hf[f"decoder.{i}.pp_block.mixer.dt_proj.bias"] = f"backbone.layers.{i}.mixer.dt_proj.bias"
        nanotron_to_hf[f"decoder.{i}.pp_block.mixer.out_proj.weight"] = f"backbone.layers.{i}.mixer.out_proj.weight"
        nanotron_to_hf[f"decoder.{i}.pp_block.mixer.out_proj.bias"] = f"backbone.layers.{i}.mixer.out_proj.bias"
        nanotron_to_hf[f"decoder.{i}.pp_block.norm.weight"] = f"backbone.layers.{i}.norm.weight"

    # Sync weights
    ref_state_dict = model_ref.state_dict()
    for name, param in tqdm(
        nanotron_model.model.named_parameters(),
        total=len(list(nanotron_model.model.named_parameters())),
        desc="Converting",
    ):
        param_is_tp_sharded = (
            isinstance(param, NanotronParameter)
            and param.is_sharded
            and parallel_context.world_ranks_to_pg[param.get_sharded_info().global_ranks] == parallel_context.tp_pg
        )

        ref_param = get_weight_from_hf(
            name=name,
            ref_module_state_dict=ref_state_dict,
            ref_module=model_ref,
            nanotron_to_hf=nanotron_to_hf,
            param_is_tp_sharded=param_is_tp_sharded,
        )

        if param_is_tp_sharded:
            sharded_info = param.get_sharded_info()
            # copy param data (not just the reference)
            with torch.no_grad():
                for local_global_slices_pair in sharded_info.local_global_slices_pairs:
                    local_slices = local_global_slices_pair.local_slices
                    global_slices = local_global_slices_pair.global_slices
                    param[local_slices].copy_(ref_param[global_slices])
        else:
            assert (
                ref_param.shape == param.shape
            ), f"Parameter shape don't match for {name}\n{ref_param.shape} != {param.shape}"
            # copy param data (not just the reference)
            with torch.no_grad():
                param.copy_(ref_param)
                ref_param = None
                torch.cuda.empty_cache()

    # Marks parameters as NanotronParameters
    mark_tied_parameters(model=nanotron_model, parallel_context=parallel_context)

    sanity_check(root_module=nanotron_model)

    save_weights(model=nanotron_model, parallel_context=parallel_context, root_folder=out_path)
    checkpoint_metadata = {
        "last_train_step": 0,
        "consumed_train_samples": 0,
    }
    save_meta(root_folder=out_path, parallel_context=parallel_context, checkpoint_metadata=checkpoint_metadata)

    if dist.get_rank() == 0:
        with open(out_path / "config.yaml", "w") as f:
            config = MambaConfig(
                general=GeneralArgs(project="test", run="mamba"),
                parallelism=parallel_config,
                model=ModelArgs(
                    init_method=MambaInit(),
                    model_config=model_config,
                ),
                tokenizer=TokenizerArgs(args.inp_path),
            )
            log_rank("Saving config ...", logger=logger, level=logging.INFO, rank=0)
            yaml.dump(config.as_dict(), f)

        with open(out_path / "model_config.json", "w") as f:
            log_rank("Saving model config ...", logger=logger, level=logging.INFO, rank=0)
            json.dump(asdict(model_config), f)