export_hf.py

"""Streaming layer-by-layer HF checkpoint export for quantized models.

Forked from modelopt's unified_export_hf.py to support models that don't fit
in CPU RAM (e.g. GLM-5 at ~1.5TB BF16). Instead of materializing the entire
state dict at once, we iterate over top-level modules, strip accelerate hooks,
move to CPU, extract state dict, and write safetensors shards incrementally.
"""

import collections.abc
import gc
import json
import re
import warnings
from pathlib import Path
from typing import Any

import torch
import torch.nn as nn
from safetensors.torch import save_file

from modelopt.torch.export.convert_hf_config import convert_hf_quant_config_format
from modelopt.torch.export.layer_utils import (
    get_expert_linear_names,
    is_moe,
    set_expert_quantizer_amax,
)
from modelopt.torch.export.quant_utils import (
    get_quant_config,
    postprocess_state_dict,
)
from modelopt.torch.export.unified_export_hf import (
    _process_quantized_modules,
    requantize_resmooth_fused_llm_layers,
)


MAX_SHARD_SIZE = 5 * 1024**3  # 5 GB per shard.
SCALE_SHARD_NAME = "model-inputscales.safetensors"
MTP_SHARD_NAME = "model-mtp.safetensors"

_EXPERT_PROJ_RE = re.compile(
    r"^(.*\.experts)\.(gate_proj|up_proj|down_proj)\.(\d+)\.(.+)$"
)


def _remap_expert_key_to_checkpoint(key: str) -> str:
    """Reverse the _QuantFusedExperts module-tree layout back to checkpoint format.

    Module tree (projection-first): experts.gate_proj.0.weight
    Checkpoint  (expert-first):     experts.0.gate_proj.weight
    """
    m = _EXPERT_PROJ_RE.match(key)
    if m:
        prefix, proj, idx, suffix = m.group(1), m.group(2), m.group(3), m.group(4)
        return f"{prefix}.{idx}.{proj}.{suffix}"
    return key


def _handle_moe_expert_quantizers(model: nn.Module) -> None:
    """Handle input quantizers of MoE experts that were not calibrated.

    Mirrors the logic from _export_transformers_checkpoint lines 646-697.
    """
    for _, sub_module in model.named_modules():
        if is_moe(sub_module) and hasattr(sub_module, "experts"):
            expert_linear_names = get_expert_linear_names(sub_module)
            for linear_name in expert_linear_names:
                if "QuantDbrxExperts" in type(sub_module.experts).__name__:
                    experts_mlp = sub_module.experts.mlp
                    if hasattr(experts_mlp, linear_name):
                        linear_modulelist = getattr(experts_mlp, linear_name)
                        if hasattr(linear_modulelist, "__iter__"):
                            set_expert_quantizer_amax(
                                modules=list(linear_modulelist),
                                quantizer_attrs=["input_quantizer"],
                            )
                elif "QuantGptOssExperts" in type(sub_module.experts).__name__:
                    gpt_oss_linear_names = ["gate_up_proj", "down_proj"]
                    for linear_name in gpt_oss_linear_names:
                        if hasattr(sub_module.experts, linear_name):
                            linear_module = getattr(sub_module.experts, linear_name)
                            if hasattr(linear_module, "input_quantizer"):
                                set_expert_quantizer_amax(
                                    modules=[linear_module],
                                    quantizer_attrs=["input_quantizer"],
                                )
                elif hasattr(sub_module.experts, "gate_proj") and isinstance(
                    getattr(sub_module.experts, "gate_proj", None), nn.ModuleList
                ):
                    # _QuantFusedExperts (GLM5): per-projection ModuleLists
                    # gate_proj[i], up_proj[i], down_proj[i] are quantized nn.Linear
                    # Force device=cpu — storage GPUs are packed with weights and
                    # can't allocate even small tensors for the internal torch.stack.
                    for proj_name in ("gate_proj", "up_proj", "down_proj"):
                        proj_list = getattr(sub_module.experts, proj_name, None)
                        if proj_list is not None and isinstance(proj_list, nn.ModuleList):
                            set_expert_quantizer_amax(
                                modules=list(proj_list),
                                quantizer_attrs=["input_quantizer", "weight_quantizer"],
                                device=torch.device("cpu"),
                            )
                elif isinstance(sub_module.experts, collections.abc.Iterable):
                    try:
                        set_expert_quantizer_amax(
                            modules=[
                                getattr(expert, linear_name)
                                for expert in sub_module.experts
                            ],
                            quantizer_attrs=["input_quantizer", "weight_quantizer"],
                        )
                    except AttributeError as e:
                        expert_types = [
                            type(expert).__name__ for expert in sub_module.experts
                        ]
                        raise AttributeError(
                            f"Failed to access attribute '{linear_name}' on experts. "
                            f"MoE module type: {type(sub_module).__name__}, "
                            f"Expert types: {expert_types}, "
                            f"Expected linear names: {expert_linear_names}. "
                            f"Original error: {e}"
                        ) from e
                else:
                    raise NotImplementedError(
                        f"MoE model with experts type "
                        f"'{type(sub_module.experts).__name__}' is not supported."
                    )


def _strip_hooks(module: nn.Module) -> None:
    """Remove accelerate hooks from a single module and all its children."""
    try:
        from accelerate.hooks import remove_hook_from_module
    except ImportError:
        return
    remove_hook_from_module(module, recurse=True)


def _tensor_size(t: torch.Tensor) -> int:
    return t.numel() * t.element_size()


def _flush_shard(
    shard: dict[str, torch.Tensor],
    shard_idx: int,
    export_dir: Path,
    weight_map: dict[str, str],
) -> int:
    """Write one safetensors shard to disk and update weight_map. Returns next shard_idx."""
    fname = f"model-{shard_idx:05d}.safetensors"
    print(f"  Writing shard {fname} ({len(shard)} tensors, "
          f"{sum(_tensor_size(v) for v in shard.values()) / 1e9:.2f} GB)")
    save_file(shard, str(export_dir / fname))
    for key in shard:
        weight_map[key] = fname
    return shard_idx + 1


def _enumerate_top_level_modules(model: nn.Module):
    """Yield (prefix, module) pairs for every top-level piece of the model.

    This produces the decoder layers one at a time, plus the embedding, norm,
    and lm_head modules. The prefixes are the dotted names as they appear in
    the full model state dict (e.g. "model.layers.0", "model.embed_tokens").

    Handles both standard CausalLM (model.model.layers) and VL models where
    layers are nested under model.model.language_model.layers.
    """
    inner = getattr(model, "model", None)
    if inner is None:
        raise ValueError(
            "Expected model to have a .model attribute (standard HF CausalLM layout)."
        )

    lm = getattr(inner, "language_model", None)
    if lm is not None:
        # VL model: layers are inside model.model.language_model.
        for name, child in inner.named_children():
            if name == "language_model":
                continue
            yield f"model.{name}", child

        layers = getattr(lm, "layers", None)
        if layers is not None:
            for i, layer in enumerate(layers):
                yield f"model.language_model.layers.{i}", layer

        for name, child in lm.named_children():
            if name == "layers":
                continue
            yield f"model.language_model.{name}", child
    else:
        # Standard CausalLM: layers directly on model.model.
        layers = getattr(inner, "layers", None)
        if layers is not None:
            for i, layer in enumerate(layers):
                yield f"model.layers.{i}", layer

        for name, child in inner.named_children():
            if name == "layers":
                continue
            yield f"model.{name}", child

    # Top-level children of the outer model that are not .model (lm_head, ...).
    for name, child in model.named_children():
        if name == "model":
            continue
        yield name, child


def export_hf(
    model: nn.Module,
    export_dir: str | Path,
    dtype: torch.dtype | None = None,
    prepare_fn: Any | None = None,
    extra_mtp_prefixes: list[str] | None = None,
    preserve_remote_code: bool = False,
) -> None:
    """Export a quantized HF model to safetensors, streaming one layer at a time.

    This avoids materializing the full state dict in CPU RAM. For each top-level
    module we: process quantized weights in-place, strip accelerate hooks, move
    to CPU, extract its state dict slice, and flush to a shard file.

    Args:
        prepare_fn: Optional callback called after pre-export processing but
            before the per-layer export loop. Used by StreamingModelLoader to
            remove streaming hooks and install materialization callbacks at the
            right point (after requantize_resmooth's dummy forward pass).
    """
    export_dir = Path(export_dir)
    export_dir.mkdir(parents=True, exist_ok=True)

    if dtype is None:
        dtype = model.config.torch_dtype

    # --- Pre-export processing (all on GPU, hooks still intact) ---
    _handle_moe_expert_quantizers(model)
    resmooth_target = model
    _patched_arch = False
    if hasattr(model, "model") and hasattr(model.model, "language_model"):
        resmooth_target = model.model.language_model
        if getattr(resmooth_target.config, "architectures", None) is None:
            resmooth_target.config.architectures = []
            _patched_arch = True
    # Skipped: requantize_resmooth_fused_llm_layers does a full dummy forward
    # then fuses gate/up amaxes and AWQ pre_quant_scales. For NVFP4 without
    # AWQ/SVDQuant this is all no-ops (resmooth_only returns early for MoE,
    # pre_quant_scale is None, gate/up tying already done in quantize.py).
    # Calling it on streaming models OOMs because the dummy forward + amax
    # stacking keeps too much on GPU 0 simultaneously.
    if _patched_arch:
        del resmooth_target.config.architectures
    quant_config = get_quant_config(model)
    hf_quant_config = convert_hf_quant_config_format(quant_config) if quant_config else None

    kv_cache_max_bound = 448
    kv_cache_format = quant_config["quantization"]["kv_cache_quant_algo"]

    # Transition from streaming execution to export mode. Must happen after
    # pre-processing (which may do dummy forward passes needing streaming hooks)
    # but before the per-layer loop (which needs materialization callbacks).
    if prepare_fn is not None:
        prepare_fn(model)

    # --- Streaming layer-by-layer export ---
    weight_map: dict[str, str] = {}
    shard: dict[str, torch.Tensor] = {}
    scale_shard: dict[str, torch.Tensor] = {}  # input_scale tensors in dedicated shard
    shard_size = 0
    shard_idx = 1
    total_tensors = 0

    print("\nStreaming export — processing layers...")
    for prefix, module in _enumerate_top_level_modules(model):
        print(f"  Processing {prefix}...")

        # Bring everything to CPU before processing.  Meta/disk layers use
        # a materialization callback; GPU-stored layers just .to("cpu").
        materialize_fn = getattr(module, "_streaming_materialize", None)
        if materialize_fn is not None:
            materialize_fn(module)
            delattr(module, "_streaming_materialize")
        else:
            _strip_hooks(module)
            module.to("cpu")
            torch.cuda.empty_cache()

        # Pack quantized weights in-place for this subtree (all on CPU now).
        _process_quantized_modules(module, dtype)

        # Extract this module's state dict with correct full-model key names.
        # input_scale tensors are collected separately into a dedicated shard
        # so they can be updated without touching the large weight shards.
        for local_key, tensor in module.state_dict().items():
            full_key = _remap_expert_key_to_checkpoint(f"{prefix}.{local_key}")
            cpu_tensor = tensor.detach().contiguous().cpu()
            t_size = _tensor_size(cpu_tensor)

            if full_key.endswith(".input_scale"):
                scale_shard[full_key] = cpu_tensor.clone()
            else:
                shard[full_key] = cpu_tensor
                shard_size += t_size

            total_tensors += 1

            if shard_size >= MAX_SHARD_SIZE:
                shard_idx = _flush_shard(shard, shard_idx, export_dir, weight_map)
                shard.clear()
                shard_size = 0

        # Free this module's parameters to reclaim CPU RAM.
        for param in module.parameters():
            param.data = torch.empty(0)
        for buf_name, buf in module.named_buffers():
            buf.data = torch.empty(0)
        gc.collect()

    # Flush remaining weight tensors.
    if shard:
        shard_idx = _flush_shard(shard, shard_idx, export_dir, weight_map)
        shard.clear()

    total_shards = shard_idx - 1

    # Flush input_scale tensors into a dedicated file (not a numbered shard).
    # This file is excluded from postprocess/rename so it keeps its fixed name,
    # making it trivial to update scales without touching weight shards.
    if scale_shard:
        scale_size = sum(_tensor_size(v) for v in scale_shard.values())
        scale_path = export_dir / SCALE_SHARD_NAME
        save_file(scale_shard, str(scale_path))
        for key in scale_shard:
            weight_map[key] = SCALE_SHARD_NAME
        print(f"  Writing {SCALE_SHARD_NAME} ({len(scale_shard)} tensors, "
              f"{scale_size / 1e3:.0f} KB)")
        scale_shard.clear()

    print(f"\nWrote {total_tensors} tensors across {total_shards + 1} files.")

    # Merge MTP (multi-token prediction) weights from source checkpoint if present.
    shard_idx, total_tensors = _merge_mtp_weights(
        model, export_dir, weight_map, shard_idx, total_tensors,
        extra_prefixes=extra_mtp_prefixes or [],
    )

    # Postprocess: filter quantizer metadata keys and rename kv cache scales.
    # We need to do this on the weight_map keys. The actual tensor data has
    # already been written, so we re-read shards that need changes.
    _postprocess_shards(export_dir, weight_map, kv_cache_max_bound, kv_cache_format)

    # Rename shards to final HF naming (model-00001-of-NNNNN.safetensors).
    _rename_shards(export_dir, weight_map, total_shards)

    # --- Save config, tokenizer, and quant metadata ---
    _save_model_metadata(model, export_dir, hf_quant_config, preserve_remote_code=preserve_remote_code)

    print(f"Export complete. Saved to: {export_dir}")


def _merge_mtp_weights(
    model: nn.Module,
    export_dir: Path,
    weight_map: dict[str, str],
    shard_idx: int,
    total_tensors: int,
    extra_prefixes: list[str] | None = None,
) -> tuple[int, int]:
    """Copy MTP weights unquantized from the source checkpoint into our export.

    Keys starting with "mtp." are always included. Additional prefixes can be
    supplied via extra_prefixes for models that store MTP weights elsewhere
    (e.g. GLM-5 uses model.layers.78. for its next-token prediction head).
    """
    from safetensors.torch import load_file

    name_or_path = model.config._name_or_path
    if Path(name_or_path).is_dir():
        source_dir = Path(name_or_path)
    else:
        from huggingface_hub import snapshot_download
        source_dir = Path(snapshot_download(
            name_or_path, allow_patterns=["*.safetensors", "*.json"],
            local_files_only=True,
        ))

    src_index_path = source_dir / "model.safetensors.index.json"
    if not src_index_path.exists():
        print("  No source index found, skipping MTP merge.")
        return shard_idx, total_tensors

    with open(src_index_path) as f:
        src_wm = json.load(f)["weight_map"]

    all_prefixes = ["mtp."] + (extra_prefixes or [])
    mtp_keys = [k for k in src_wm if any(k.startswith(p) for p in all_prefixes)]
    # Don't re-merge keys already exported by the main loop.
    mtp_keys = [k for k in mtp_keys if k not in weight_map]
    if not mtp_keys:
        print("  No MTP keys in source checkpoint, skipping.")
        return shard_idx, total_tensors

    src_shards_needed = sorted(set(src_wm[k] for k in mtp_keys))
    print(f"\nMerging {len(mtp_keys)} MTP weights from {len(src_shards_needed)} source shard(s)...")

    mtp_shard: dict[str, torch.Tensor] = {}
    mtp_size = 0
    for src_shard_name in src_shards_needed:
        src_data = load_file(str(source_dir / src_shard_name))
        for key in mtp_keys:
            if src_wm[key] == src_shard_name and key in src_data:
                t = src_data[key].contiguous()
                mtp_shard[key] = t
                mtp_size += _tensor_size(t)
                total_tensors += 1

    if mtp_shard:
        shard_idx = _flush_shard(mtp_shard, shard_idx, export_dir, weight_map)
        print(f"  Wrote {len(mtp_shard)} MTP tensors ({mtp_size / 1e9:.2f} GB)")
        mtp_shard.clear()

    return shard_idx, total_tensors


def _postprocess_shards(
    export_dir: Path,
    weight_map: dict[str, str],
    kv_cache_max_bound: float,
    kv_cache_format: str | None,
) -> None:
    """Apply postprocess_state_dict logic to already-written shards.

    Re-reads each shard, filters/renames keys, and rewrites if changed.
    """
    from safetensors.torch import load_file

    shard_files = set(weight_map.values())
    new_weight_map: dict[str, str] = {}

    for shard_file in sorted(shard_files):
        path = export_dir / shard_file
        shard_data = load_file(str(path))

        processed = postprocess_state_dict(
            shard_data, kv_cache_max_bound, kv_cache_format
        )

        # Drop dummy KV cache scales/biases from modelopt.
        # These are all 1.0 when only quantizing expert MLPs (not KV cache),
        # and the key naming convention differs between sglang and vLLM,
        # causing spurious warnings in both frameworks.
        processed = {
            k: v for k, v in processed.items()
            if not any(s in k for s in (".k_scale", ".v_scale", ".k_bias", ".v_bias"))
        }

        if set(processed.keys()) != set(shard_data.keys()):
            save_file(processed, str(path))
            print(f"  Postprocessed {shard_file}: "
                  f"{len(shard_data)} -> {len(processed)} keys")

        for key in processed:
            new_weight_map[key] = shard_file

    weight_map.clear()
    weight_map.update(new_weight_map)


def _rename_shards(
    export_dir: Path,
    weight_map: dict[str, str],
    total_shards: int,
) -> None:
    """Rename shards to the standard HF format: model-00001-of-NNNNN.safetensors."""
    del total_shards
    old_names = sorted(
        name
        for name in set(weight_map.values())
        if name not in {SCALE_SHARD_NAME, MTP_SHARD_NAME}
    )
    total_shards = len(old_names)
    rename_map: dict[str, str] = {}

    for i, old_name in enumerate(old_names, 1):
        new_name = f"model-{i:05d}-of-{total_shards:05d}.safetensors"
        if old_name != new_name:
            (export_dir / old_name).rename(export_dir / new_name)
            rename_map[old_name] = new_name

    # Update weight_map with final names.
    for key in list(weight_map.keys()):
        old_shard = weight_map[key]
        if old_shard in rename_map:
            weight_map[key] = rename_map[old_shard]

    # Write the index file.
    index = {
        "metadata": {"total_size": sum(
            (export_dir / f).stat().st_size
            for f in set(weight_map.values())
        )},
        "weight_map": dict(sorted(weight_map.items())),
    }
    with open(export_dir / "model.safetensors.index.json", "w") as f:
        json.dump(index, f, indent=2)

    print(f"  Wrote model.safetensors.index.json ({len(weight_map)} entries)")


def _save_model_metadata(
    model: nn.Module,
    export_dir: Path,
    hf_quant_config: dict | None,
    preserve_remote_code: bool = False,
) -> None:
    """Save config.json (with quantization_config), generation_config, and tokenizer files."""
    # Save config.json from the model's config object.
    config = model.config
    config_dict = config.to_dict()

    # Strip auto_map for models already in upstream transformers. Remote-code
    # models such as MiMo need these entries plus the copied Python files.
    if not preserve_remote_code:
        config_dict.pop("auto_map", None)

    # Ensure architectures is set. Nested text_config may have None; derive
    # from the outer model class (e.g. Qwen3_5MoeForConditionalGeneration).
    if not config_dict.get("architectures"):
        config_dict["architectures"] = [type(model).__name__]

    if hf_quant_config is not None:
        config_dict["quantization_config"] = hf_quant_config

    with open(export_dir / "config.json", "w") as f:
        json.dump(config_dict, f, indent=4)

    # Save hf_quant_config.json for backward compatibility.
    if hf_quant_config is not None:
        with open(export_dir / "hf_quant_config.json", "w") as f:
            json.dump(hf_quant_config, f, indent=4)

    # Save generation_config if present.
    if hasattr(model, "generation_config") and model.generation_config is not None:
        gen_config = model.generation_config.to_dict()
        with open(export_dir / "generation_config.json", "w") as f:
            json.dump(gen_config, f, indent=4)

    # Copy tokenizer and chat template files from the source model.
    import shutil
    name_or_path = model.config._name_or_path
    if Path(name_or_path).is_dir():
        source_dir = Path(name_or_path)
    else:
        from huggingface_hub import snapshot_download
        allow_patterns = ["*.json", "*.jinja", "*.txt", "*.model"]
        if preserve_remote_code:
            allow_patterns.extend(["*.py", "audio_tokenizer/*"])
        source_dir = Path(snapshot_download(
            name_or_path,
            allow_patterns=allow_patterns,
            ignore_patterns=None if preserve_remote_code else ["*.safetensors"],
            local_files_only=True,
        ))
    copy_patterns = [
        "tokenizer.json", "tokenizer_config.json", "chat_template.jinja",
        "special_tokens_map.json", "tokenizer.model", "added_tokens.json",
        "preprocessor_config.json", "video_preprocessor_config.json",
        "merges.txt", "vocab.json",
    ]
    if preserve_remote_code:
        copy_patterns.extend([
            "configuration_mimo_v2.py", "modeling_mimo_v2.py",
        ])
    copied = []
    for name in copy_patterns:
        src = source_dir / name
        if src.exists():
            shutil.copy2(src, export_dir / name)
            copied.append(name)

    if preserve_remote_code:
        for src in sorted(source_dir.glob("*.py")):
            if src.name in copied:
                continue
            shutil.copy2(src, export_dir / src.name)
            copied.append(src.name)

    if preserve_remote_code:
        for dirname in ("audio_tokenizer",):
            src = source_dir / dirname
            if src.exists():
                dst = export_dir / dirname
                if dst.exists():
                    shutil.rmtree(dst)
                shutil.copytree(src, dst)
                copied.append(dirname)

    print(f"  Wrote config.json, generation_config.json, and copied: {', '.join(copied)}")