From aaca100c4df17bb4f1e8154db21a941695e91525 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 11:10:42 +0000 Subject: [PATCH 01/24] Write a sharded model's HF checkpoint without gathering it distributed_save_hf_checkpoint writes an already-processed sharded model as HF safetensors with no rank-0 gather: fused experts are split shard-locally, every other tensor moves whole to one owner rank, the reverse conversion runs there, and each rank writes its own shards. No rank holds the whole model. A model sharded by FSDP2 alone opts in with MODELOPT_FSDP2_EXPORT=nogather; the expert-parallel export routes to it in the next change. On that path export_hf_checkpoint packs with keep_fused_experts; writes config.json, a sanitized generation_config.json and hf_quant_config.json on rank 0; and copies the source's non-model files (tokenizer, processor, remote code, chat templates), which the writer never sees. It rejects save_modelopt_state, which the path cannot honor without save_pretrained, rather than drop it silently. A model built from a config has an empty config._name_or_path, which as a path is the working directory, so the export copied the working directory's tokenizer, processor and code files as if they were the model's. An empty name now means no source checkpoint. The output matches the single-process export tensor for tensor (CPU tests: FP8 and NVFP4, one and two ranks, root-wrapped and layers-only FSDP2). A GPU test PTQs a tiny Qwen3 (dense and MoE) under FSDP2 with FP8 and NVFP4, opting each export into the writer, and checks the checkpoint against the safetensors headers. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + examples/hf_ptq/README.md | 2 + modelopt/torch/export/distributed_writer.py | 245 +++++++++- modelopt/torch/export/unified_export_hf.py | 108 ++++- .../export/test_distributed_hf_export.py | 445 ++++++++++++++++++ .../torch/export/test_distributed_writer.py | 265 ++++++++++- .../torch/export/test_unified_export_hf.py | 21 + 7 files changed, 1075 insertions(+), 12 deletions(-) create mode 100644 tests/gpu/torch/export/test_distributed_hf_export.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 904eac00fa3..51c8523ab8e 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -118,6 +118,7 @@ Changelog - ``KV_QUANT_CFG`` presets in vLLM fake-quant serving now quantize the MLA KV cache; on vLLM 0.16 and later they silently quantized nothing for MLA models. - vLLM fake-quant serving of NVFP4 and cast-mode KV-cache configs no longer needs ``--enforce-eager``. ``examples/vllm_serve/vllm_serve_fakequant.py`` now sets ``VLLM_DISABLE_COMPILE_CACHE=1`` for fake-quant serves by default, because a cached torch.compile graph of the same model without the fake quant would otherwise be reused. - ``examples/deepseek/deepseek_v4/ptq.py`` now supports DeepSeek-V4.1-Flash, whose reference implementation uses 32x32 FP8 blocks and a tokenizer-aware ``Transformer``; this is the setup used to produce ``nvidia/DeepSeek-V4.1-Flash-NVFP4``. DeepSeek-V4-Pro checkpoints work as before. +- Add an opt-in no-gather HuggingFace export for FSDP2-sharded models: with ``MODELOPT_FSDP2_EXPORT=nogather``, ``export_hf_checkpoint`` sends each tensor to a single owner rank and every rank writes its own shards in parallel, instead of gathering the full state dict on one host. MoE experts stay sharded through the export, so each rank only materializes its own experts. *Speculative Decoding* diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 15e3c96ae2b..044337fd22f 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -648,6 +648,8 @@ See [Recipe-based Quantization](#recipe-based-quantization) for the recipe forma > *Performance Note: FSDP2 is designed for training workloads and may result in longer calibration and export times. For faster calibration, maximize the batch size based on available GPU memory and choose the right number of GPUs to avoid unnecessary communication.* +> *Set `MODELOPT_FSDP2_EXPORT=nogather` to have every rank write its own shards of the checkpoint in parallel instead of gathering the full state dict on one host.* + ## Evaluate Accuracy ### TensorRT-LLM Validation diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index 246937ad923..decf61ab99e 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -13,11 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Pieces of a no-gather distributed HF checkpoint writer. +"""No-gather distributed HF checkpoint writer for FSDP2. Every tensor is sent to exactly one owner rank over batched point-to-point sends and written from there, so no rank ever materializes the whole model and no unit is -replicated across ranks. +replicated across ranks. Selected by ``MODELOPT_FSDP2_EXPORT=nogather``. """ import json @@ -28,6 +28,39 @@ import torch import torch.nn as nn +from .moe_utils import _FUSED_PROJ, _fused_experts_prefixes, _split_local_fused_module + + +def distributed_save_hf_checkpoint( + model: nn.Module, + export_dir: "str | Path", + maxbound: float, + kv_cache_format: "str | None", + max_shard_size: "str | int" = "10GB", + is_modelopt_qlora: bool = False, + extra_state_dict: "dict[str, torch.Tensor] | None" = None, +) -> None: + """No-gather distributed HF export. + + Thin wrapper: hf_ptq exports under ``torch.inference_mode()``, so run the entire write with inference mode DISABLED + via a nested ``inference_mode(False)`` context -- the ``.detach()`` in ``get_model_state_dict`` and the version- + counted ops in the write path both reject inference tensors. The context form reliably restores normal mode for this + scope (a decorator did not, under the caller's already-active inference_mode). ``no_grad`` too: + ``inference_mode(False)`` re-ENABLES autograd, so a bare ``clone()``/``detach()`` on an inference tensor would try + to set up grad and read its (missing) version counter -- ``no_grad`` avoids that while the clones still land as + normal tensors because inference mode is off. + """ + with torch.inference_mode(False), torch.no_grad(): + _distributed_save_hf_checkpoint_impl( + model, + export_dir, + maxbound, + kv_cache_format, + max_shard_size, + is_modelopt_qlora, + extra_state_dict, + ) + def _even_bins(keys: list, sizes: dict, n: int, max_bytes: int) -> list: """Partition ``keys`` into ~``n`` EVEN-sized bins (by bytes), each <= ``max_bytes`` where possible. @@ -601,3 +634,211 @@ def write_index() -> bool: # Also the closing barrier: no rank returns before the index is on disk. all_reduce("index written", write_index, all) + + +def _distributed_save_hf_checkpoint_impl( + model: nn.Module, + export_dir: "str | Path", + maxbound: float, + kv_cache_format: "str | None", + max_shard_size: "str | int" = "10GB", + is_modelopt_qlora: bool = False, + extra_state_dict: "dict[str, torch.Tensor] | None" = None, +) -> None: + """Distributed HF safetensors export -- no rank-0 full-model host-RAM gather. + + Writes an already-processed, FSDP2-sharded model to ``export_dir`` as HF safetensors, entirely + distributed. No rank ever holds the whole model: + + 1. Fused expert weights (when present): each rank splits its LOCAL shard into per-expert keys + with global indices and keeps them in place -- no communication. Experts are therefore + already disjoint across ranks. dp-replica ranks under DP x EP skip this. + 2. Dense / non-expert weights: each fqn is assigned to exactly ONE owner rank and moved there + (:func:`_materialize_owned`), so every tensor becomes whole on its owner while the set stays + spread across ranks. Replicated and plain values need no transport; dim-0 shards travel + point-to-point, bucketed per peer. + 3. The reverse weight conversion runs per-rank on those whole tensors -- it cannot run while + they are sharded, because transformers' conversion ops are written against whole tensors + and have no way to express a placement. + 4. Every rank bin-packs its own owned set into even shards, numbered into a global range, and + writes them directly with ``safetensors.save_file``. Rank 0 unions the weight index. + + Peak memory per rank is its owned set, not the model. Placements the point-to-point transport + does not cover are rejected rather than routed through an untested slower path. + """ + import torch.distributed as dist + from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict + from torch.distributed.tensor import DTensor + + from .model_utils import TiedWeightMap + from .quant_utils import postprocess_state_dict + + # Name-based tied-weight dedup, matching the gather path. ``fully_shard`` splits a shared + # nn.Parameter into distinct per-module shards, so a declared tie surfaces here as two + # independent DTensors and both sides would be written. The map is derived from names and is + # identical on every rank, so the alias is dropped consistently; postprocess_state_dict skips a + # group whose canonical is missing from this rank's slice rather than orphaning the alias. + tied_map = TiedWeightMap(model) + + export_dir = Path(export_dir) + rank, world = dist.get_rank(), dist.get_world_size() + device = torch.device(f"cuda:{torch.cuda.current_device()}") + + # This writer targets a FULLY FSDP2-sharded model -- no expert parallelism. FSDP2 + # shards the fused 3-D expert weight on dim 0 (the expert axis), so fused experts are just + # DTensors sharded across the fsdp mesh: each per-expert key's TRUE global index comes from the + # DTensor's own dim-0 offset (``_split_local_fused_module`` reads it), so there is no EP block base + # (ep_rank=0) and no dp-replica dedup (dp_idx=0 -> every rank writes its own experts). The EP-hybrid + # (a2a experts) path -- which resolves ep_rank/dp_idx from the parallel-group provider -- is out of + # scope here. + ep_rank, dp_idx = 0, 0 + # Total config expert count, to tell a per-EP-group expert tensor (size < total) from one that + # spans all experts (classic DP x EP) when computing the per-expert global index offset. + _cfg = getattr(model, "config", None) + + def _experts_of(c): + if c is None: + return 0 + return int( + getattr(c, "num_experts", 0) + or getattr(c, "num_local_experts", 0) + or getattr(c, "n_routed_experts", 0) # nemotron_h + or 0 + ) + + # VLMs nest the text-model fields under ``config.text_config`` (the quantized backbone is + # ``model.language_model``), so fall back to it when the top-level config has no expert count. + total_experts = _experts_of(_cfg) or _experts_of(getattr(_cfg, "text_config", None)) + # DP x EP: experts are dp-REPLICATED across the EDP group (each holds a full EP-sharded copy). Only + # the first replica (dp_idx == 0) writes them; replica ranks (dp_idx > 0) skip the per-expert split. + # dp_idx is the EDP-group rank resolved above -- 0 when experts are not EDP-replicated (dense / + # FSDP2 / EP-without-DP), so every such rank writes. + write_experts = dp_idx == 0 + sharded_sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) + + # Re-materialize every value as a NORMAL tensor UP FRONT. hf_ptq exports under torch.inference_mode(), + # so these are inference tensors; the inference flag is STICKY -- clone/reshape/detach/slice all + # propagate it -- and the expert split, scale reduce and write below all do version-counted ops + # that reject inference tensors. The only way to drop the flag is to copy the data into a freshly + # ALLOCATED tensor (here, under the wrapper's inference_mode(False)). For a DTensor, copy the LOCAL + # shard and rewrap so mesh/placements survive. Bounded to this rank's shard. + def _to_normal(v): + if isinstance(v, DTensor): + _loc = v.to_local() + _new = torch.empty_like(_loc) + _new.copy_(_loc) + # Pass the ORIGINAL global shape/stride: from_local() otherwise re-infers them as + # local x world, which is only right for an EVEN shard. This helper runs over every + # DTensor in the state dict just before classification, so with an uneven shard + # (b_proj: 96 rows over 64 ranks) or a zero-size one (mlp_res_proj) each rank derives a + # DIFFERENT global shape for the same tensor -- and then classifies it differently, + # which breaks the export with no hint of where the shape went wrong. + return DTensor.from_local( + _new, + v.device_mesh, + v.placements, + run_check=False, + shape=v.shape, + stride=v.stride(), + ) + if v.is_inference(): + _new = torch.empty_like(v) + _new.copy_(v) + return _new + return v + + # Re-materialize IN PLACE, not as `{k: _to_normal(v) for ...}`. The comprehension keeps the ENTIRE + # original dict alive while building a full second copy -> a transient 2x of the whole state dict. + # That is fine at 30B (24GB -> 48GB) but OOMs at large MoE: at 235B the write-point state dict is + # ~157GB/rank (fused expert weights are sharded, but the per-expert NVFP4 scale tensors are + # replicated on every rank), and doubling it hits ~314GB > GPU. Overwriting each key frees the old + # value immediately, so peak stays ~1x. (Import of ep-dp's in-place/streaming write handling.) + for _k in list(sharded_sd.keys()): + sharded_sd[_k] = _to_normal(sharded_sd[_k]) + prefixes = _fused_experts_prefixes(sharded_sd) + prefix_set = set(prefixes) + + def _is_expert_key(key: str) -> bool: + for proj in _FUSED_PROJ: + for suffix in ("", "_weight_scale", "_weight_scale_2", "_input_scale"): + if key.endswith(f".experts.{proj}{suffix}"): + if key[: key.rfind(".experts.") + len(".experts")] in prefix_set: + return True + return False + + # Probe the DTensor placement of a sample expert + dense weight: directly shows what state the + # generation forwards left the sharded params in -- FSDP2 (Shard(0)/Replicate/gathered) or EP + # (expert dim-0 Shard, or a plain local shard for the a2a transport). + + # Sync the shared per-module activation (input) scales across ranks (global max), so every + # expert of a module gets the same input_scale regardless of which rank owns it. + in_keys = sorted(k for k in sharded_sd if k.endswith("_input_scale") and _is_expert_key(k)) + if in_keys: + stacked = torch.stack( + [ + sharded_sd[k].detach().to(device=device, dtype=torch.float32).reshape(()) + for k in in_keys + ] + ) + dist.all_reduce(stacked, op=dist.ReduceOp.MAX) + for j, k in enumerate(in_keys): + sharded_sd[k] = stacked[j].clone() + # (1) Experts: split this rank's local shard into per-expert keys (global idx); stays local. + # Skipped on dp-replica ranks under DP x EP (their experts are written by dp group 0). + local_sd: dict = {} + if write_experts: + for pi, prefix in enumerate(prefixes): + local_sd.update( + _split_local_fused_module( + prefix, sharded_sd, ep_rank=ep_rank, total_experts=total_experts + ) + ) + local_sd = postprocess_state_dict( + local_sd, maxbound, kv_cache_format, is_modelopt_qlora, tied_map=tied_map + ) + # In place (see above): avoid a transient 2x copy of the per-expert split. + for _k in list(local_sd.keys()): + local_sd[_k] = local_sd[_k].detach().contiguous() + # At this point local_sd holds ONLY this rank's (already split, per-expert) experts -- capture the + # set so _materialize_owned can tell disjoint experts from replicated dense. + expert_key_set = set(local_sd) + + # (2) Dense / non-expert: leave sharded for now -- FSDP2 leaves DTensors (dim-0 shard); EP leaves + # plain replicated tensors. Do NOT gather to rank 0; _materialize_owned below moves each tensor to + # a single OWNER rank instead, so the set stays spread across ranks. postprocess runs on every rank + # (dict-level key renaming + small scalar scale math -- safe on DTensors). + nonexpert_keys = [k for k in sharded_sd if not _is_expert_key(k)] + dense_sd = {k: sharded_sd[k] for k in nonexpert_keys} + dense_sd = postprocess_state_dict( + dense_sd, maxbound, kv_cache_format, is_modelopt_qlora, tied_map=tied_map + ) + local_sd.update(dense_sd) + # No explicit gc.collect() here: the del drops the last references and CPython frees acyclic + # objects immediately. Collecting a ~900-entry DTensor state dict measured 0.89s of the 0.89s + # this phase used to cost -- and because it ran at different speeds per rank, the next + # collective absorbed the skew and looked like gather cost. + del sharded_sd, dense_sd + # (3) Whole-tensor redistribution -> reverse conversion -> write. Each tensor ends up WHOLE on + # exactly one rank (never gathered to a single rank), which is what lets the reverse run: the + # transformers conversion ops are written against whole tensors and cannot express a placement. + # The reverse must precede the shard layout, since it renames keys and splits tensors. + _materialize_owned(local_sd, expert_key_set, rank, world) + _revert_whole_sd(model, local_sd) + # Tensors the model does not own -- whatever the caller hands in via extra_state_dict, e.g. + # weights read straight off the source checkpoint that the FSDP2 loader drops. They are plain + # CPU tensors, already whole and identical on every rank, so ONE rank adds them to its + # write set; the index union and + # the cross-rank duplicate check below then treat them like any other owned key. Rank 0 takes + # them wholesale: they are small next to the model, and spreading them would cost a collective + # to agree on the split. They go in AFTER the reverse conversion because they carry their + # original hub names and were never converted -- reverting them would rename them wrongly. This + # mirrors the gather path, which merges extra_state_dict into the already-reverted dict. + if extra_state_dict and rank == 0: + collisions = sorted(set(extra_state_dict) & set(local_sd)) + if collisions: + raise RuntimeError( + f"extra_state_dict overlaps the exported model weights on {len(collisions)} key(s), " + f"so the checkpoint would be ambiguous; first few: {collisions[:5]}" + ) + local_sd.update(extra_state_dict) + _write_even_shards(local_sd, export_dir, max_shard_size, rank, world) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index e0ad64e6790..0fdbe31bdf8 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -19,6 +19,7 @@ import copy import importlib import json +import os import re import tempfile import warnings @@ -949,6 +950,7 @@ def _process_quantized_modules( model: nn.Module, dtype: torch.dtype, is_modelopt_qlora: bool = False, + keep_fused_experts: bool = False, ) -> None: """Process all quantized modules in model, export weights in-place. @@ -964,7 +966,9 @@ def _process_quantized_modules( # No per-module dedup cache: tied duplicates are dropped by name in postprocess_state_dict. # The handlers pack whatever weight they are handed, and nothing here materializes a shard -- # FSDP2 models go through collect_export_tensors instead, which packs inside a gather window. - assert not is_fsdp2_model(model), ( + # The no-gather writer is the deliberate exception: it packs each rank's LOCAL shard and never + # materializes a full tensor, which is signalled by keep_fused_experts. + assert not is_fsdp2_model(model) or keep_fused_experts, ( "_process_quantized_modules cannot pack a sharded model; use collect_export_tensors" ) ctx = ExportContext( @@ -972,6 +976,7 @@ def _process_quantized_modules( dtype=dtype, is_modelopt_qlora=is_modelopt_qlora, model_type=hf_model_type(model), + keep_fused_experts=keep_fused_experts, ) # A keep-fused export packs each FSDP2 unit's modules in one weight window, held across them. @@ -1028,13 +1033,22 @@ def _prepare_model_for_export(model, dtype, is_modelopt_qlora): return dtype, tied_map, quant_config -def pack_quantized_weights(model, dtype, is_modelopt_qlora: bool = False) -> None: - """Quantize every module's weight in place, then rebuild the fused MoE linears.""" +def pack_quantized_weights( + model, dtype, is_modelopt_qlora: bool = False, keep_fused_experts: bool = False +) -> None: + """Quantize every module's weight in place, then rebuild the fused MoE linears. + + ``keep_fused_experts`` is for the no-gather distributed FSDP2 export: it leaves MoE experts + fused and sharded through the fold, so each rank materializes only its LOCAL experts instead + of all N on every rank. The per-expert split happens later, in the writer. + """ # Deferred: modelopt.torch.quantization imports modelopt.torch.export, so importing the # per-model PTQ module at module scope here would risk an import cycle. from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear - _process_quantized_modules(model, dtype, is_modelopt_qlora) + _process_quantized_modules( + model, dtype, is_modelopt_qlora, keep_fused_experts=keep_fused_experts + ) _reconstruct_fused_moe_linear(model) @@ -1648,6 +1662,20 @@ def _write_hf_export_config( with open(original_config, "w") as file: json.dump(config_data, file, indent=4) + # generation_config.json: on the plain path save_pretrained emits it next to config.json, but the + # distributed writer writes weight shards only and save_pretrained never runs -- so without this the + # exported checkpoint silently loses the model's generation defaults. Measured on GLM-5.3-BF16: + # byte-identical weight shards on both paths, yet only the non-distributed export carried + # generation_config.json (do_sample / eos_token_id). Silent, and it lands on exactly the knobs that + # decide sampling behaviour at serving time. + gen_cfg = getattr(model, "generation_config", None) + if gen_cfg is not None: + try: + _sanitize_generation_config_for_save(model) + gen_cfg.save_pretrained(str(export_dir)) + except Exception as gen_err: # never fail an otherwise-good export over metadata + warnings.warn(f"Could not save generation_config: {gen_err}") + def _revert_hf_quant_config_names( hf_quant_config: dict, @@ -1884,6 +1912,26 @@ def read_unplaced_weights(model: nn.Module, *, keys_only: bool = False) -> dict[ return out +def _write_base_config(model: nn.Module, export_dir: "Path | str") -> None: + """Write config.json + generation_config for the no-gather distributed export path. + + ``distributed_save_hf_checkpoint`` writes only the weight shards, so rank 0 emits here what + ``save_pretrained`` would have. The weight half of ``save_pretrained`` must NOT run: the + weights are already on disk and the model may still be sharded. + """ + model.config.save_pretrained(export_dir) + if model.can_generate() and getattr(model, "generation_config", None) is not None: + try: + # Same strict-validate trip as every other save site: a config carrying top_p/top_k + # with do_sample unset (e.g. GLM-5.3-BF16) makes save_pretrained raise, and without + # this the only trace is a "Could not save generation_config" warning that reads as + # data loss when the later _write_hf_export_config has in fact written the file. + _sanitize_generation_config_for_save(model) + model.generation_config.save_pretrained(export_dir) + except Exception as gen_err: + warnings.warn(f"Could not save generation_config: {gen_err}") + + def export_hf_checkpoint( model: Any, dtype: torch.dtype | None = None, @@ -1985,6 +2033,9 @@ def export_hf_checkpoint( and torch.distributed.is_initialized() and is_fsdp2_model(model) ) + use_no_gather_writer = ( + is_fsdp2_sharded and os.environ.get("MODELOPT_FSDP2_EXPORT") == "nogather" + ) # Not the global rank: a non-FSDP2 export writes the whole checkpoint from every process. rank = torch.distributed.get_rank() if is_fsdp2_sharded else 0 # Offloaded models take the streaming path: it materializes one layer at a time and @@ -1992,6 +2043,18 @@ def export_hf_checkpoint( # buffer instead of the whole quantized state dict. _offloaded = has_accelerate_offload(model) + # An offloaded model takes the streaming path below; otherwise a distributed FSDP2 model takes + # the no-gather path, which writes weights and config itself instead of going through + # save_pretrained. extra_state_dict IS honoured there (one rank adds those tensors to its write + # set); save_modelopt_state is not, so reject it explicitly rather than dropping it silently + # (the streaming path declines it the same way). + if use_no_gather_writer and not _offloaded: + if save_modelopt_state: + raise NotImplementedError( + "save_modelopt_state=True is not supported by the no-gather FSDP2 distributed " + "export: it writes the checkpoint directly rather than through " + "model.save_pretrained(). Save the ModelOpt state separately with mto.save()." + ) try: if _offloaded: # Imported here rather than at module scope: the streaming exporter imports the @@ -2014,6 +2077,38 @@ def export_hf_checkpoint( if getattr(model, "hf_quantizer", None) is not None: model.hf_quantizer = None hf_quant_config = _revert_quant_config_names_best_effort(model, hf_quant_config) + elif use_no_gather_writer: + # Opt-in alternative to the streaming path below: never replicate a unit across ranks. + # Each tensor is sent to exactly one owner over batched point-to-point sends, instead of + # every rank entering a gather window and all but one discarding the result. Same + # outputs; the win grows with world size, where the replicated gather costs ~world x the + # bytes. Left opt-in so the default path is unchanged. + from .distributed_writer import distributed_save_hf_checkpoint + + _qlora = kwargs.get("is_modelopt_qlora", False) + _dtype, _tied, hf_quant_config = _prepare_model_for_export(model, dtype, _qlora) + # main split preparation from packing, so pack explicitly. keep_fused_experts leaves + # MoE experts fused + sharded through the fold, so each rank folds only its LOCAL + # experts; the writer splits them after consolidation. + pack_quantized_weights(model, _dtype, _qlora, keep_fused_experts=True) + if getattr(model, "hf_quantizer", None) is not None: + model.hf_quantizer = None + distributed_save_hf_checkpoint( + model, + export_dir, + maxbound=448, + kv_cache_format=(hf_quant_config or {}) + .get("quantization", {}) + .get("kv_cache_quant_algo"), + max_shard_size=max_shard_size, + is_modelopt_qlora=_qlora, + extra_state_dict=extra_state_dict, + ) + if rank == 0: + _write_base_config(model, export_dir) + if save_modelopt_state and ModeloptStateManager.is_converted(model): + torch.save(modelopt_state(model), export_dir / _MODELOPT_STATE_SAVE_NAME) + hf_quant_config = _revert_quant_config_names_best_effort(model, hf_quant_config) elif is_fsdp2_sharded: # FSDP2 multi-rank: stream each rank's owned units straight to its own shard files, so # a rank buffers its own share of the model rather than the whole checkpoint, and the @@ -2103,8 +2198,9 @@ def export_hf_checkpoint( if rank == 0: _write_hf_export_config(model, hf_quant_config, export_dir) copy_off_index_safetensors(_source_checkpoint(model), export_dir) - if not (_offloaded or is_fsdp2_sharded): - # The streaming paths already did, from save_non_weight_artifacts. + if not (_offloaded or (is_fsdp2_sharded and not use_no_gather_writer)): + # The streaming paths already did, from save_non_weight_artifacts; the no-gather + # writer does not. _copy_non_model_files_from_source(model, export_dir) except Exception as e: diff --git a/tests/gpu/torch/export/test_distributed_hf_export.py b/tests/gpu/torch/export/test_distributed_hf_export.py new file mode 100644 index 00000000000..9e85e27a0f6 --- /dev/null +++ b/tests/gpu/torch/export/test_distributed_hf_export.py @@ -0,0 +1,445 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""End-to-end checks on the checkpoint written by the no-gather FSDP2 distributed export. + +With ``MODELOPT_FSDP2_EXPORT=nogather``, ``export_hf_checkpoint`` routes an FSDP2-sharded model under +``torch.distributed`` to ``distributed_save_hf_checkpoint`` -- the writer every expert-parallel model +takes: every rank writes its own DCP shards and the files are consolidated in parallel, with no +full-model host-RAM gather on rank 0. That writer reconstructs +the checkpoint from per-rank pieces, so the failure modes it can introduce are structural rather +than numerical -- a key silently dropped, a rank-local shard written where the full tensor belongs, +a sidecar or the weight index never emitted. Those all produce a checkpoint that *looks* fine (no +error, plausible file sizes) but is wrong or unloadable, so they are asserted here explicitly: + +1. every source parameter survives the round trip, and nothing extra is left behind; +2. each quantized weight carries the scales its format needs; +3. every tensor has its full, unsharded shape; +4. the non-weight files are all present -- the ones the exporter writes (config, generation config, + quant config, weight index) and the source's own (tokenizer, vocab, chat template), which the + exporter copies over after the writer has removed its ``sharded/`` staging dir. + +It also covers the two things the writer must not quietly drop on the way: a declared tied +weight still gets deduplicated, and options the path cannot honour are rejected rather than +ignored. + +Run with >=2 GPUs so the expert axis and the FSDP2 shard axis are both actually split. +""" + +import copy +import json +import os +from functools import partial +from pathlib import Path +from unittest import mock + +import pytest +import torch +from _test_utils.torch.transformers_models import create_tiny_qwen3_dir, create_tiny_qwen3_moe_dir +from safetensors.torch import load_file + +import modelopt.torch.quantization as mtq +from modelopt.torch.export.unified_export_hf import export_hf_checkpoint +from modelopt.torch.quantization.utils import patch_fsdp_mp_dtypes +from modelopt.torch.utils.distributed import fsdp2_wrap, is_fsdp2_model + +# Small enough to force the tiny model across several shards, so the multi-file layout -- and the +# weight index that makes it loadable -- are exercised rather than collapsing to one model.safetensors. +# Not smaller: each extra shard is another DCP write + consolidation round trip on a toy model. +MAX_SHARD_SIZE = "512KB" + +# The default tiny Qwen3 is too degenerate to shard: head_dim = hidden_size / num_heads = 32/16 = 2, +# so q_norm/k_norm are 2-element tensors. Split over a 4-rank world most ranks get an empty chunk and +# the DCP planner rejects the coverage outright ("invalid fill tensor-volume"), taking the whole +# worker pool down with it. Size every FSDP2-sharded axis to stay comfortably divisible instead -- +# still a toy model, but one that survives a multi-rank split. Last dims are multiples of 16 so NVFP4 +# block quantization is also well defined. +TINY_KWARGS = { + "hidden_size": 128, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 32, + "intermediate_size": 128, + "max_position_embeddings": 64, + "num_hidden_layers": 2, +} +TINY_MOE_KWARGS = {"moe_intermediate_size": 128, "num_experts": 8, "num_experts_per_tok": 2} + +# Suffixes the exporter ADDS to a checkpoint; every other exported key must correspond to a +# parameter of the source model. +_SCALE_SUFFIXES = ( + "weight_scale", + "weight_scale_2", + "weight_scale_inv", + "input_scale", + "pre_quant_scale", + "k_scale", + "v_scale", +) + +# Exported dtypes that mean "this weight was quantized" (NVFP4 packs two fp4 per uint8). +_QUANTIZED_DTYPES = {"F8_E4M3", "U8", "I8"} +_PACKED_DTYPES = {"U8"} + + +def _safetensors_meta(directory: Path) -> dict[str, tuple[str, tuple[int, ...]]]: + """``{tensor_name: (dtype, shape)}`` across every ``*.safetensors`` in ``directory``. + + Reads the safetensors header directly (8-byte little-endian length, then JSON) so the check + depends only on what is on disk -- no loader, no dequantization, no GPU. + """ + out: dict[str, tuple[str, tuple[int, ...]]] = {} + for path in sorted(directory.glob("*.safetensors")): + with open(path, "rb") as fh: + header = json.loads(fh.read(int.from_bytes(fh.read(8), "little"))) + for name, spec in header.items(): + if name != "__metadata__": + out[name] = (spec["dtype"], tuple(spec["shape"])) + return out + + +def _is_scale(key: str) -> bool: + return key.endswith(_SCALE_SUFFIXES) + + +def _expected_weights(src: dict[str, tuple[str, tuple[int, ...]]]) -> dict[str, tuple[int, ...]]: + """Source ``{name: shape}`` rewritten into the keys the exporter is expected to emit. + + Everything maps 1:1 except a MoE's fused 3-D expert weights, which the exporter splits into + per-expert 2-D weights: ``experts.gate_up_proj`` ``[E, 2I, H]`` becomes ``experts..gate_proj`` + and ``experts..up_proj`` (each ``[I, H]``), and ``experts.down_proj`` ``[E, H, I]`` becomes + ``experts..down_proj`` ``[H, I]``. + """ + expected: dict[str, tuple[int, ...]] = {} + for name, (_dtype, shape) in src.items(): + if name.endswith("mlp.experts.gate_up_proj"): + prefix, experts, two_i, hidden = name[: -len("gate_up_proj")], *shape + assert two_i % 2 == 0, f"{name}: fused gate_up dim {two_i} is not even" + for e in range(experts): + expected[f"{prefix}{e}.gate_proj.weight"] = (two_i // 2, hidden) + expected[f"{prefix}{e}.up_proj.weight"] = (two_i // 2, hidden) + elif name.endswith("mlp.experts.down_proj"): + prefix, experts, hidden, inter = name[: -len("down_proj")], *shape + for e in range(experts): + expected[f"{prefix}{e}.down_proj.weight"] = (hidden, inter) + else: + expected[name] = shape + return expected + + +def _no_gather_export(): + """Opt the FSDP2-only model into the no-gather writer, scoped to this export. + + The workers are a persistent pool, so a bare ``os.environ`` write would leak into later tests. + """ + return mock.patch.dict(os.environ, {"MODELOPT_FSDP2_EXPORT": "nogather"}) + + +def _ptq_and_export(rank, size, *, src_dir, export_dir, quant_cfg, **export_kwargs): + """Load the tiny model on every rank, FSDP2-shard it, PTQ it, and export.""" + from transformers import AutoModelForCausalLM + + with patch_fsdp_mp_dtypes(): + model = AutoModelForCausalLM.from_pretrained(src_dir, dtype=torch.bfloat16).to("cuda") + model.eval() + + fsdp2_wrap(model) + # The no-gather writer is selected by this predicate plus the opt-in below, so assert it here: + # without it the export would silently fall back to the rank-0 gather and the test would pass + # while never touching the code under test. + assert is_fsdp2_model(model), "fsdp2_wrap did not shard the model" + assert torch.distributed.is_initialized() + torch.distributed.barrier() + + input_ids = torch.randint(0, model.config.vocab_size, (2, 8), device="cuda") + mtq.quantize(model, quant_cfg, lambda m: m(input_ids)) + torch.distributed.barrier() + + with _no_gather_export(): + export_hf_checkpoint( + model, export_dir=export_dir, max_shard_size=MAX_SHARD_SIZE, **export_kwargs + ) + torch.distributed.barrier() + + +# Four PTQ + distributed-export round trips; the tests/gpu default of 120s is not enough headroom. +@pytest.mark.timeout(600) +@pytest.mark.parametrize("moe", [False, True], ids=["dense", "moe"]) +@pytest.mark.parametrize( + ("quant_cfg", "algo"), + [(mtq.FP8_DEFAULT_CFG, "FP8"), (mtq.NVFP4_DEFAULT_CFG, "NVFP4")], + ids=["fp8", "nvfp4"], +) +def test_fsdp2_distributed_export_is_complete(dist_workers, tmp_path, moe, quant_cfg, algo): + if torch.cuda.device_count() < 2: + pytest.skip("needs >=2 GPUs to shard the expert and FSDP2 axes") + + make_dir = create_tiny_qwen3_moe_dir if moe else create_tiny_qwen3_dir + kwargs = {**TINY_KWARGS, **(TINY_MOE_KWARGS if moe else {})} + src_dir = Path(make_dir(tmp_path, with_tokenizer=True, **kwargs)) + export_dir = tmp_path / "export" + + # The source's non-weight files (tokenizer, vocab, chat template, ...), which the exporter must + # copy into the export itself: nothing else writes them there. + sidecars = { + p.name + for p in src_dir.iterdir() + if p.is_file() + and not p.name.endswith(".safetensors") + and p.name != "model.safetensors.index.json" + } + + dist_workers.run( + partial(_ptq_and_export, src_dir=src_dir, export_dir=export_dir, quant_cfg=quant_cfg) + ) + + src = _safetensors_meta(src_dir) + exported = _safetensors_meta(export_dir) + assert exported, f"no safetensors written to {export_dir}" + expected = _expected_weights(src) + + # ---- 1. every source parameter survives, and nothing extra is left behind ---- + missing = sorted(set(expected) - set(exported)) + assert not missing, ( + f"{len(missing)} source parameter(s) missing from the export: {missing[:10]}" + ) + + unexpected = sorted(k for k in set(exported) - set(expected) if not _is_scale(k)) + # A fused expert weight left next to the per-expert weights it was split into would show up + # here -- the same tensor exported twice, once unquantized. + assert not unexpected, f"{len(unexpected)} unexpected non-scale key(s): {unexpected[:10]}" + + # ---- 2. + 3. scales present, and every tensor at its full unsharded shape ---- + quantized = 0 + for name, want_shape in expected.items(): + dtype, got_shape = exported[name] + if dtype in _QUANTIZED_DTYPES and name.endswith(".weight"): + quantized += 1 + prefix = name[: -len(".weight")] + assert f"{prefix}.weight_scale" in exported, f"{name}: quantized but no weight_scale" + assert f"{prefix}.input_scale" in exported, f"{name}: quantized but no input_scale" + if algo == "NVFP4": + assert f"{prefix}.weight_scale_2" in exported, f"{name}: NVFP4 needs weight_scale_2" + if dtype in _PACKED_DTYPES: + # Two 4-bit values per byte -> the packed last dim is half the logical one. + want_shape = (*want_shape[:-1], want_shape[-1] // 2) + assert got_shape == want_shape, f"{name}: shape {got_shape}, expected {want_shape}" + assert quantized > 0, f"nothing was quantized -- {algo} config did not take effect" + + # ---- 4. the non-weight files came along ---- + exported_files = {p.name for p in export_dir.iterdir() if p.is_file()} + for required in ("config.json", "generation_config.json", "hf_quant_config.json"): + assert required in exported_files, f"{required} missing from the export" + + # Everything the source shipped alongside its weights (tokenizer, vocab/merges, chat template, + # special-tokens map, ...) is in the export; only the weight files are rewritten. + lost = sorted(sidecars - exported_files) + assert not lost, f"sidecar file(s) missing from the export: {lost}" + assert not (export_dir / "sharded").exists(), "per-rank staging dir left behind" + + # A multi-file checkpoint is unloadable without the index, so require it whenever the writer + # emitted sharded names, and require it to describe every tensor actually on disk. + if any(p.name.startswith("model-") for p in export_dir.glob("*.safetensors")): + index_path = export_dir / "model.safetensors.index.json" + assert index_path.exists(), "sharded export without model.safetensors.index.json" + weight_map = json.loads(index_path.read_text())["weight_map"] + assert set(weight_map) == set(exported), ( + f"index/shard mismatch: {len(set(exported) - set(weight_map))} tensor(s) unindexed, " + f"{len(set(weight_map) - set(exported))} indexed but absent" + ) + for key, fname in weight_map.items(): + assert (export_dir / fname).exists(), f"index points at missing shard {fname} for {key}" + + assert ( + json.loads((export_dir / "hf_quant_config.json").read_text())["quantization"]["quant_algo"] + == algo + ) + + +def _export_with(rank, size, *, src_dir, export_dir, **export_kwargs): + """Export once with ``export_kwargs``, for the rejection checks below.""" + from transformers import AutoModelForCausalLM + + with patch_fsdp_mp_dtypes(): + model = AutoModelForCausalLM.from_pretrained(src_dir, dtype=torch.bfloat16).to("cuda") + model.eval() + fsdp2_wrap(model) + assert is_fsdp2_model(model) + torch.distributed.barrier() + with _no_gather_export(): + export_hf_checkpoint(model, export_dir=export_dir, **export_kwargs) + + +@pytest.mark.timeout(600) +@pytest.mark.parametrize( + ("kwargs", "needle"), + [ + ({"save_modelopt_state": True}, "save_modelopt_state"), + ], + ids=["save_modelopt_state"], +) +def test_fsdp2_distributed_export_rejects_unsupported_options( + dist_workers, tmp_path, kwargs, needle +): + """Options the no-gather path cannot honour must raise, not be silently ignored. + + ``save_modelopt_state`` is forwarded to ``save_pretrained`` on the gather path; the distributed + path writes the checkpoint itself and never calls it, so a caller passing it would otherwise get + a checkpoint quietly missing the ModelOpt state. (``extra_state_dict`` used to be rejected here + too and is now supported -- see the MTP test below.) + """ + if torch.cuda.device_count() < 2: + pytest.skip("needs >=2 GPUs") + + src_dir = Path(create_tiny_qwen3_dir(tmp_path, with_tokenizer=True, **TINY_KWARGS)) + with pytest.raises(Exception, match=needle): + dist_workers.run( + partial(_export_with, src_dir=src_dir, export_dir=tmp_path / "rejected", **kwargs) + ) + + +@pytest.mark.timeout(600) +def test_fsdp2_distributed_export_dedups_tied_weights(dist_workers, tmp_path): + """A declared tie is deduplicated by name in the distributed write, as on the gather path. + + ``fully_shard`` splits the shared parameter into distinct per-module shards, so the tie survives + only as matching names -- both sides reach the writer as independent DTensors and would both be + written unless the writer applies the same ``TiedWeightMap`` the gather path uses. + """ + if torch.cuda.device_count() < 2: + pytest.skip("needs >=2 GPUs") + + src_dir = Path( + create_tiny_qwen3_dir( + tmp_path, with_tokenizer=True, tie_word_embeddings=True, **TINY_KWARGS + ) + ) + export_dir = tmp_path / "export_tied" + dist_workers.run( + partial( + _ptq_and_export, + src_dir=src_dir, + export_dir=export_dir, + quant_cfg=mtq.FP8_DEFAULT_CFG, + ) + ) + + exported = _safetensors_meta(export_dir) + assert exported, "nothing exported" + assert "model.embed_tokens.weight" in exported, "canonical tied weight missing" + assert "lm_head.weight" not in exported, ( + "tied alias 'lm_head.weight' was written alongside its canonical -- " + "name-based dedup did not run in the distributed writer" + ) + + +@pytest.mark.timeout(600) +def test_fsdp2_distributed_export_writes_extra_state_dict(dist_workers, tmp_path): + """Tensors handed in via ``extra_state_dict`` must reach the checkpoint. + + This is how weights the FSDP2 loader drops arrive: the caller reads them straight off the + source checkpoint and passes them through ``extra_state_dict``. They belong to no rank's + shard and no rank's model, so unless the writer deliberately merges them in, such a model + exports silently incomplete -- valid-looking, right size, missing the layer. + + They must also keep their original names: they were never part of the transformers conversion, + so the reverse conversion must not touch them. + """ + if torch.cuda.device_count() < 2: + pytest.skip("needs >=2 GPUs") + + src_dir = Path(create_tiny_qwen3_dir(tmp_path, with_tokenizer=True, **TINY_KWARGS)) + export_dir = tmp_path / "export_extra" + # Shaped and named like an inlined MTP tail (DeepSeek-V3 / GLM-5.1 put it at layers.). + extra = { + "model.layers.99.mtp.weight": torch.arange(64, dtype=torch.bfloat16).reshape(8, 8), + "model.layers.99.mtp.bias": torch.full((8,), 3.0, dtype=torch.bfloat16), + } + dist_workers.run( + partial( + _ptq_and_export, + src_dir=src_dir, + export_dir=export_dir, + quant_cfg=mtq.FP8_DEFAULT_CFG, + extra_state_dict=extra, + ) + ) + + exported = _safetensors_meta(export_dir) + assert exported, "nothing exported" + for name, want in extra.items(): + assert name in exported, f"extra_state_dict tensor '{name}' never reached the checkpoint" + assert exported[name][1] == tuple(want.shape), f"{name}: shape {exported[name][1]}" + + index = json.loads((export_dir / "model.safetensors.index.json").read_text()) + for name in extra: + assert name in index["weight_map"], f"'{name}' written but missing from the weight index" + + # Values must survive intact, not just the names. + merged: dict = {} + for shard in sorted(export_dir.glob("*.safetensors")): + merged.update(load_file(str(shard))) + for name, want in extra.items(): + assert torch.equal(merged[name].cpu(), want), f"{name}: value changed on the way out" + + +@pytest.mark.timeout(600) +def test_fsdp2_distributed_export_keeps_unquantized_submodule(dist_workers, tmp_path): + """A submodule left unquantized must still be exported, in full and unquantized. + + PTQ is routinely applied to part of a model -- the language model of a VLM, or a layer range + excluded by the recipe -- and the checkpoint is expected to be the COMPLETE model, with the + untouched part carried through at its original precision. The writer builds its state dict from + the whole model, so the risk is not that it drops those tensors but that it mangles them: the + expert split, the scale postprocess and the reverse conversion all walk keys that have no + quantizer state here. + """ + if torch.cuda.device_count() < 2: + pytest.skip("needs >=2 GPUs") + + # Disable every quantizer under layer 0; layer 1 stays quantized as the control. + quant_cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) + quant_cfg["quant_cfg"].append({"quantizer_name": "*layers.0*", "enable": False}) + + src_dir = Path(create_tiny_qwen3_dir(tmp_path, with_tokenizer=True, **TINY_KWARGS)) + export_dir = tmp_path / "export_partial" + dist_workers.run( + partial(_ptq_and_export, src_dir=src_dir, export_dir=export_dir, quant_cfg=quant_cfg) + ) + + exported = _safetensors_meta(export_dir) + assert exported, "nothing exported" + + src_meta = _safetensors_meta(src_dir) + layer0_src = {k for k in src_meta if k.startswith("model.layers.0.")} + assert layer0_src, "fixture has no layer 0 weights to check" + + for name in sorted(layer0_src): + assert name in exported, f"unquantized '{name}' missing from the checkpoint" + dtype, shape = exported[name] + assert dtype not in _QUANTIZED_DTYPES, ( + f"'{name}' was excluded from quantization but exported as {dtype}" + ) + assert shape == src_meta[name][1], f"'{name}': shape {shape}, source {src_meta[name][1]}" + + # No scale companions should exist for the excluded layer... + stray = sorted(k for k in exported if k.startswith("model.layers.0.") and _is_scale(k)) + assert not stray, f"excluded layer 0 carries quantization scales: {stray}" + # ...while the control layer must actually be quantized, or the test proves nothing. + quantized_l1 = [ + k + for k, (d, _) in exported.items() + if k.startswith("model.layers.1.") and d in _QUANTIZED_DTYPES + ] + assert quantized_l1, "layer 1 was not quantized -- the exclusion pattern disabled everything" diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index ba0f89f56ba..c8a8f340ded 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -13,32 +13,42 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""The pieces of the no-gather HF checkpoint writer (``modelopt.torch.export.distributed_writer``). +"""The no-gather FSDP2 HF checkpoint writer (``modelopt.torch.export.distributed_writer``) on CPU. The writer moves every tensor to exactly one owner rank and writes it from there, so what it can get wrong is structural: a tensor missing, duplicated, mis-assembled from its shards, renamed, or -left out of the weight index. These tests drive each piece directly, on CPU. +left out of the weight index. The helpers are tested directly; the writer itself runs end to end on +a gloo process group -- one rank in process, two spawned ranks where shards actually travel -- and +the written checkpoint is reloaded and compared with the source checkpoint or with what the +single-process export writes for the same model. """ import contextlib +import copy import json import math +import os import types +from datetime import timedelta +from functools import partial from unittest import mock import pytest import torch -from _test_utils.torch.distributed.utils import bounded_collectives +import torch.distributed as dist +from _test_utils.torch.distributed.utils import bounded_collectives, spawn_multiprocess_job from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor from torch.multiprocessing.reductions import StorageWeakRef pytest.importorskip("safetensors") pytest.importorskip("transformers") from safetensors.torch import load_file -from transformers import AutoModelForCausalLM, LlamaConfig, MixtralConfig +from transformers import AutoModelForCausalLM, LlamaConfig, MixtralConfig, Qwen3MoeConfig from transformers.core_model_loading import Chunk, WeightConverter +import modelopt.torch.quantization as mtq from modelopt.torch.export import quant_aware_conversion from modelopt.torch.export.distributed_writer import ( _even_bins, @@ -48,7 +58,35 @@ _nbytes, _revert_whole_sd, _write_even_shards, + distributed_save_hf_checkpoint, ) +from modelopt.torch.export.quant_aware_conversion import revert_weight_conversion_quant_aware +from modelopt.torch.export.unified_export_hf import ( + _export_transformers_checkpoint, + export_hf_checkpoint, +) + + +def _cpu_device_query(): + """The writer asks CUDA which device to all-reduce MoE expert input scales on. + + Nothing written in this file has expert input scales, so on a CPU-only host any answer works. + """ + return mock.patch.object(torch.cuda, "current_device", return_value=0) + + +def _write(model, export_dir, **kwargs): + with _cpu_device_query(): + distributed_save_hf_checkpoint( + model, export_dir, maxbound=448, kv_cache_format=None, max_shard_size="8KB", **kwargs + ) + + +def _nogather_export(model, export_dir, **kwargs): + with _cpu_device_query(), mock.patch.dict(os.environ, {"MODELOPT_FSDP2_EXPORT": "nogather"}): + export_hf_checkpoint( + model, dtype=torch.bfloat16, export_dir=export_dir, max_shard_size="16KB", **kwargs + ) def _raw_bytes(tensor): @@ -78,6 +116,16 @@ def _load_checkpoint(export_dir): return tensors +def _save(config, path): + torch.manual_seed(0) + AutoModelForCausalLM.from_config(config, dtype=torch.bfloat16).save_pretrained(path) + return path + + +def _load(path): + return AutoModelForCausalLM.from_pretrained(path, dtype=torch.bfloat16).eval() + + def _mixtral_config(**overrides): """A tiny Mixtral: fused experts in memory, per-expert ``block_sparse_moe`` names on the hub.""" kwargs = { @@ -108,6 +156,28 @@ def _llama_config(): ) +def _fsdp2_shard(model, mesh, wrap_root=True): + """Shard each decoder layer, and the root too unless ``wrap_root`` is off (as expert parallelism + leaves it, sharding only the dense layers).""" + for layer in model.model.layers: + fully_shard(layer, mesh=mesh) + if wrap_root: + fully_shard(model, mesh=mesh) + return model + + +def _quantize_with_reference(model, quant_cfg): + """Quantize ``model`` and return what the single-process export writes for a copy of it.""" + torch.manual_seed(1) + batch = torch.randint(0, model.config.vocab_size, (2, 8)) + mtq.quantize(model, quant_cfg, lambda m: m(batch)) + ref = copy.deepcopy(model) + want, _ = _export_transformers_checkpoint(ref, torch.bfloat16) + want = revert_weight_conversion_quant_aware(ref, dict(want)) + assert any(v.dtype in (torch.float8_e4m3fn, torch.uint8) for v in want.values()) + return want + + # ------------------------------------------------------------------ pure helpers @@ -575,3 +645,190 @@ def test_a_name_two_ranks_both_write_fails_the_write_on_every_rank(cpu_dist_work cpu_dist_workers.run(_cross_rank_collision_job, tmp_dir=tmp_path) assert not list(tmp_path.glob("*.safetensors")) + + +# ------------------------------------------------------------------ the writer, one rank + + +@pytest.mark.parametrize("tie", [True, False], ids=["tied", "untied"]) +def test_distributed_save_writes_the_hub_checkpoint_of_a_sharded_moe( + gloo_world_of_one, tmp_path, tie +): + """Fused experts are split, every name is reverse-converted, and a declared tie is dropped.""" + src = _save(_mixtral_config(tie_word_embeddings=tie), tmp_path / "src") + model = _fsdp2_shard(_load(src), gloo_world_of_one) + + _write(model, tmp_path / "export") + + assert len(list((tmp_path / "export").glob("*.safetensors"))) > 1 + want = load_file(str(src / "model.safetensors")) + _assert_byte_identical(_load_checkpoint(tmp_path / "export"), want) + + +def test_distributed_save_runs_from_inside_inference_mode(gloo_world_of_one, tmp_path): + """A caller in inference mode hands the writer inference tensors, which its writes reject.""" + src = _save(_mixtral_config(), tmp_path / "src") + model = _fsdp2_shard(_load(src), gloo_world_of_one) + + with torch.inference_mode(): + # Made in inference mode, like the scales an inference-mode pack registers. + model.model.norm.register_buffer("weight_scale", torch.full((), 0.5)) + _write(model, tmp_path / "export") + + want = { + **load_file(str(src / "model.safetensors")), + "model.norm.weight_scale": torch.tensor(0.5), + } + _assert_byte_identical(_load_checkpoint(tmp_path / "export"), want) + + +def test_distributed_save_writes_extra_state_dict_under_its_own_names(gloo_world_of_one, tmp_path): + """Extra tensors already carry hub names, so the reverse conversion must not rename them. + + The first is named like an in-memory Mixtral expert weight, which the reverse would turn into + ``block_sparse_moe.experts.0.w1``. + """ + src = _save(_mixtral_config(), tmp_path / "src") + model = _fsdp2_shard(_load(src), gloo_world_of_one) + extra = { + "model.layers.9.mlp.experts.0.gate_proj.weight": torch.arange( + 64, dtype=torch.bfloat16 + ).reshape(8, 8), + "model.mtp.norm.weight": torch.full((8,), 3.0, dtype=torch.bfloat16), + } + + _write(model, tmp_path / "export", extra_state_dict=extra) + + want = {**load_file(str(src / "model.safetensors")), **extra} + _assert_byte_identical(_load_checkpoint(tmp_path / "export"), want) + + +def test_distributed_save_rejects_extra_state_dict_that_overlaps_the_model( + gloo_world_of_one, tmp_path +): + model = _fsdp2_shard(_load(_save(_mixtral_config(), tmp_path / "src")), gloo_world_of_one) + + with pytest.raises(RuntimeError, match="overlaps the exported model weights"): + _write(model, tmp_path / "export", extra_state_dict={"model.norm.weight": torch.ones(32)}) + + +@pytest.mark.parametrize("wrap_root", [True, False], ids=["root-wrapped", "layers-only"]) +@pytest.mark.parametrize( + "quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG], ids=["fp8", "nvfp4"] +) +def test_nogather_export_writes_what_the_single_process_export_writes( + gloo_world_of_one, tmp_path, quant_cfg, wrap_root +): + src = _save(_llama_config(), tmp_path / "src") + sidecar = src / "chat_template.jinja" # a source file only the exporter carries over + sidecar.write_text("{{ messages }}") + model = _load(src) + want = _quantize_with_reference(model, quant_cfg) + + _nogather_export(_fsdp2_shard(model, gloo_world_of_one, wrap_root), tmp_path / "export") + + _assert_byte_identical(_load_checkpoint(tmp_path / "export"), want) + for name in ("config.json", "generation_config.json", "hf_quant_config.json"): + assert (tmp_path / "export" / name).is_file(), name + assert (tmp_path / "export" / sidecar.name).read_text() == sidecar.read_text() + + +def test_nogather_export_rejects_save_modelopt_state(gloo_world_of_one, tmp_path): + """The path never runs ``save_pretrained``, which writes that state: refuse, not drop it.""" + model = _fsdp2_shard(_load(_save(_llama_config(), tmp_path / "src")), gloo_world_of_one) + + with pytest.raises(NotImplementedError, match="save_modelopt_state"): + _nogather_export(model, tmp_path / "export", save_modelopt_state=True) + + +@pytest.mark.parametrize( + "quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG], ids=["fp8", "nvfp4"] +) +def test_nogather_export_writes_moe_experts_quantized(gloo_world_of_one, tmp_path, quant_cfg): + config = Qwen3MoeConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=32, + moe_intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + num_experts=4, + num_experts_per_tok=2, + max_position_embeddings=32, + ) + model = _load(_save(config, tmp_path / "src")) + # Weight-only experts: expert input scales would need the writer's CUDA-only all-reduce. + quant_cfg = copy.deepcopy(quant_cfg) + quant_cfg["quant_cfg"].append({"quantizer_name": "*experts*input_quantizer", "enable": False}) + want = _quantize_with_reference(model, quant_cfg) + + _nogather_export(_fsdp2_shard(model, gloo_world_of_one), tmp_path / "export") + + _assert_byte_identical(_load_checkpoint(tmp_path / "export"), want) + + +# ------------------------------------------------------------------ the writer, two ranks + + +def _two_rank_job(rank, world_size, tmp_dir): + mesh = init_device_mesh("cpu", (world_size,)) + + # An unquantized MoE whose vocab rows and experts split unevenly; the test compares what the + # two ranks wrote between them with the source checkpoint. + _write(_fsdp2_shard(_load(tmp_dir / "moe"), mesh), tmp_dir / "moe_export") + + # A quantized dense model through the export entry point: fp8 weights and bf16 norms travel + # in separate per-dtype buckets. Every rank quantizes identically, so rank 0's single-process + # export of its own copy is the reference. + model = _load(tmp_dir / "dense") + want = _quantize_with_reference(model, mtq.FP8_DEFAULT_CFG) + _nogather_export(_fsdp2_shard(model, mesh), tmp_dir / "dense_export") + if rank == 0: + _assert_byte_identical(_load_checkpoint(tmp_dir / "dense_export"), want) + + # Each dense tensor ends up whole on exactly one rank, and the owners are spread out. + full = {f"t{i}": torch.full((4, 2), float(i)) for i in range(8)} + sd = {k: distribute_tensor(v, mesh, [Shard(0)]) for k, v in full.items()} + _materialize_owned(sd, set(), rank, world_size) + owned = [None] * world_size + dist.all_gather_object(owned, sorted(sd)) + assert all(owned), f"a rank owns nothing: {owned}" + assert sorted(k for keys in owned for k in keys) == sorted(full) + assert all(torch.equal(v, full[k]) for k, v in sd.items()) + + # Ranks that disagree must fail with a diagnosis, not hang in a collective or drop a tensor. + sd = {"shared": torch.ones(2), f"rank{rank}": torch.ones(2)} + with pytest.raises(RuntimeError, match="dense key sets differ"): + _materialize_owned(sd, set(), rank, world_size) + w = torch.ones(1, 2) + if rank == 0: + w = DTensor.from_local(w, mesh, [Shard(0)], run_check=False) + with pytest.raises(RuntimeError, match="disagree on which tensors are dim-0 sharded"): + _materialize_owned({"w": w}, set(), rank, world_size) + sd = {"model.norm.weight": torch.full((2,), float(rank))} + if rank == 0: + with pytest.raises(RuntimeError, match="name collision across ranks"): + _write_even_shards(sd, tmp_dir / "collision", "1MB", rank, world_size) + dist.barrier() # releases the other rank from the writer's closing barrier + else: + _write_even_shards(sd, tmp_dir / "collision", "1MB", rank, world_size) + + +def test_two_ranks_write_one_checkpoint_and_reject_inconsistent_state(skip_on_windows, tmp_path): + """Shards cross ranks point to point, each tensor gets one owner, files share one numbering. + + Spawning ranks dominates the cost, so one job also checks the cross-rank consistency errors. + """ + src = _save(_mixtral_config(vocab_size=63, num_local_experts=3), tmp_path / "moe") + _save(_llama_config(), tmp_path / "dense") + + spawn_multiprocess_job(2, partial(_two_rank_job, tmp_dir=tmp_path), backend="gloo") + + shards = sorted(p.name for p in (tmp_path / "moe_export").glob("*.safetensors")) + n = len(shards) + assert n > 2, "both ranks should have written several files" + assert shards == [f"model-{i:05d}-of-{n:05d}.safetensors" for i in range(1, n + 1)] + want = load_file(str(src / "model.safetensors")) + _assert_byte_identical(_load_checkpoint(tmp_path / "moe_export"), want) diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index f5f36627ad0..eba57c4f264 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -1015,3 +1015,24 @@ def test_source_checkpoint_is_none_for_a_model_built_from_config(): """``from_config`` leaves ``_name_or_path`` empty; it must not resolve to the working directory, or every export would copy the cwd as the "source checkpoint".""" assert _source_checkpoint(SimpleNamespace(config=SimpleNamespace(_name_or_path=""))) is None + + +def test_a_model_built_from_a_config_names_no_source_checkpoint(tmp_path, monkeypatch): + """An empty ``_name_or_path`` is the working directory as a path; nothing there is the model's.""" + from modelopt.torch.export.unified_export_hf import ( + _copy_non_model_files_from_source, + _source_checkpoint, + ) + + workdir = tmp_path / "workdir" + workdir.mkdir() + (workdir / "tokenizer.json").write_text("{}") + monkeypatch.chdir(workdir) + model = SimpleNamespace(config=SimpleNamespace(_name_or_path="")) + export_dir = tmp_path / "export" + export_dir.mkdir() + + _copy_non_model_files_from_source(model, export_dir) + + assert _source_checkpoint(model) is None + assert not any(export_dir.iterdir()) From 39b8d8558597e25bdbdc05a2c8e7cc0eb32f049c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 20:30:58 +0000 Subject: [PATCH 02/24] Pick the HF export's route by the model's distribution plan export_hf_checkpoint takes the FSDP2 routes, streaming or no-gather, when the model's recorded plan FSDP2-shards it, and reads the carried-over weights on rank 0 alone for the same reason, rather than looking for an FSDP2 module. A diffusers pipeline is not a module and is never distributed, so it is exported without a plan, as a new test checks. The writer tests record the plan where they shard a model with fully_shard themselves. Signed-off-by: Shengliang Xu --- modelopt/torch/export/unified_export_hf.py | 20 ++++++------------- .../torch/export/test_distributed_writer.py | 4 +++- .../torch/export/test_export_diffusers.py | 17 ++++++++++++++++ 3 files changed, 26 insertions(+), 15 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 0fdbe31bdf8..c8aaddc5c4d 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -92,7 +92,7 @@ from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.dataset_utils import _disable_use_cache -from modelopt.torch.utils.distributed import is_fsdp2_model +from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan, is_fsdp2_model from modelopt.torch.utils.perf import maybe_clear_cuda_cache try: @@ -1980,14 +1980,10 @@ def export_hf_checkpoint( # emits it from rank 0 alone, so on the other ranks a full read is host memory spent and thrown # away -- and the carried set is the large stuff. The KEY list is still resolved everywhere, # because get_quant_config runs per rank and the configs have to agree. - _writes_extra = ( - not ( - torch.distributed.is_available() - and torch.distributed.is_initialized() - and is_fsdp2_model(model) - ) - or torch.distributed.get_rank() == 0 - ) + # How the model is distributed decides how it is exported. A diffusers pipeline is not a module, + # and nothing here distributes one. + plan = distribution_plan(model) if isinstance(model, nn.Module) else DistributionPlan() + _writes_extra = not plan.is_fsdp_dense or torch.distributed.get_rank() == 0 _carried = read_unplaced_weights(model, keys_only=not _writes_extra) if _writes_extra and _carried: extra_state_dict = {**_carried, **(extra_state_dict or {})} @@ -2028,11 +2024,7 @@ def export_hf_checkpoint( ) return - is_fsdp2_sharded = ( - torch.distributed.is_available() - and torch.distributed.is_initialized() - and is_fsdp2_model(model) - ) + is_fsdp2_sharded = plan.is_fsdp_dense use_no_gather_writer = ( is_fsdp2_sharded and os.environ.get("MODELOPT_FSDP2_EXPORT") == "nogather" ) diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index c8a8f340ded..165599251bf 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -65,6 +65,7 @@ _export_transformers_checkpoint, export_hf_checkpoint, ) +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan def _cpu_device_query(): @@ -158,11 +159,12 @@ def _llama_config(): def _fsdp2_shard(model, mesh, wrap_root=True): """Shard each decoder layer, and the root too unless ``wrap_root`` is off (as expert parallelism - leaves it, sharding only the dense layers).""" + leaves it, sharding only the dense layers), and record that the model is FSDP2-sharded.""" for layer in model.model.layers: fully_shard(layer, mesh=mesh) if wrap_root: fully_shard(model, mesh=mesh) + set_distribution_plan(model, DistributionPlan(mesh.size(), dense="fsdp")) return model diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index 001b4245b72..3a39947246a 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -16,6 +16,7 @@ import copy import inspect import json +import types from pathlib import Path from types import SimpleNamespace @@ -239,6 +240,22 @@ def __init__(self, unet, scheduler, setting=None): assert {path.name for path in export_dir.iterdir() if path.is_dir()} == {"unet"} +def test_a_pipeline_is_exported_without_a_distribution_plan(tmp_path, monkeypatch): + """A pipeline is not a module: nothing distributes one, so the export reads no plan of it.""" + exported = [] + monkeypatch.setattr(unified_export_hf, "is_diffusers_object", lambda model: True) + monkeypatch.setattr( + unified_export_hf, + "_export_diffusers_checkpoint", + lambda pipe, *a, **k: exported.append(pipe), + ) + pipe = types.SimpleNamespace(config=None) + + export_hf_checkpoint(pipe, export_dir=tmp_path) + + assert exported == [pipe] + + def test_export_diffusers_unet_quantized_matches_llm_config(tmp_path, monkeypatch): model = get_tiny_unet() export_dir = tmp_path / "export_unet_quant" From 0815e684edb4b4a750fa72be4fc5bd5e2d8d7451 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 00:49:02 +0000 Subject: [PATCH 03/24] Record the plan the no-gather export tests wrap their models by fsdp2_wrap records no plan; load_pretrained does. These tests wrap the model by hand, and the export picks its route from the plan, so each records the FSDP2 plan it wrapped the model by. Signed-off-by: Shengliang Xu --- tests/gpu/torch/export/test_distributed_hf_export.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/gpu/torch/export/test_distributed_hf_export.py b/tests/gpu/torch/export/test_distributed_hf_export.py index 9e85e27a0f6..84b2b6f81c9 100644 --- a/tests/gpu/torch/export/test_distributed_hf_export.py +++ b/tests/gpu/torch/export/test_distributed_hf_export.py @@ -52,7 +52,12 @@ import modelopt.torch.quantization as mtq from modelopt.torch.export.unified_export_hf import export_hf_checkpoint from modelopt.torch.quantization.utils import patch_fsdp_mp_dtypes -from modelopt.torch.utils.distributed import fsdp2_wrap, is_fsdp2_model +from modelopt.torch.utils.distributed import ( + DistributionPlan, + fsdp2_wrap, + is_fsdp2_model, + set_distribution_plan, +) # Small enough to force the tiny model across several shards, so the multi-file layout -- and the # weight index that makes it loadable -- are exercised rather than collapsing to one model.safetensors. @@ -155,6 +160,9 @@ def _ptq_and_export(rank, size, *, src_dir, export_dir, quant_cfg, **export_kwar model.eval() fsdp2_wrap(model) + set_distribution_plan( + model, DistributionPlan(size, dense="fsdp") + ) # wrapped by hand, so recorded here, as load_pretrained would # The no-gather writer is selected by this predicate plus the opt-in below, so assert it here: # without it the export would silently fall back to the rank-0 gather and the test would pass # while never touching the code under test. @@ -275,6 +283,7 @@ def _export_with(rank, size, *, src_dir, export_dir, **export_kwargs): model = AutoModelForCausalLM.from_pretrained(src_dir, dtype=torch.bfloat16).to("cuda") model.eval() fsdp2_wrap(model) + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) assert is_fsdp2_model(model) torch.distributed.barrier() with _no_gather_export(): From 5ffa2d9a619ec18428eb0cdccf3068de33e47190 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 00:56:54 +0000 Subject: [PATCH 04/24] Note in the changelog that a distributed model must carry its plan Batch sizing and the HF export read how a model is distributed from the plan load_pretrained records, and raise for an FSDP2-sharded or multi-rank model without one, where they used to look at the model and the process group. Users who distribute a model themselves, fsdp2_wrap included, now record its plan. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 51c8523ab8e..4868fc93950 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -191,6 +191,7 @@ Changelog - ``get_te_hybrid_stack_spec`` was removed from ``modelopt.torch.nas.plugins.megatron``; it had no use outside tests. Use ``modelopt.torch.utils.plugins.megatron_layer_specs.te_hybrid_stack_spec_sequential_mlp`` for the SequentialMLP layout, or ``megatron.core.models.hybrid.hybrid_layer_specs.hybrid_stack_spec`` for grouped GEMM. - Unified HuggingFace export now fails with ``NotImplementedError`` when it meets an MoE block whose expert projection names it does not know, instead of assuming Mixtral's ``w1``/``w2``/``w3``. If you hit this, register a ``ModelSpec`` for the model under ``modelopt/torch/models/``. Every MoE architecture ModelOpt exported correctly before this change is registered, so no supported model regresses. - ``--recipe`` (and ``modelopt.recipe.load_recipe``) now resolve a recipe path **filesystem-first**: a recipe of the same relative path in the current working directory takes precedence over the shipped built-in of that name, matching how recipe ``$import`` paths already resolve. Previously the built-in won. +- ``export_hf_checkpoint`` and ``modelopt.torch.utils.dataset_utils.get_max_batch_size`` now read how a model is distributed from the plan ``load_pretrained`` records on it, and raise for an FSDP2-sharded model, or one in a multi-rank run, that carries none. If you distribute a model yourself, for example with ``fsdp2_wrap``, which records no plan, record one with ``modelopt.torch.utils.distributed.set_distribution_plan`` before calibrating or exporting it. **Deprecations** From f792fa6c1a44c095574505ce52cbd2d37887cb94 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 17:31:21 +0000 Subject: [PATCH 05/24] File the no-gather export entries under 0.49.0 in the changelog main opened 0.49.0 for what merges after 0.48.0's cut, which this PR does. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4868fc93950..9c475e2d747 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -15,6 +15,7 @@ Changelog - GPTQ no longer runs a max-calibration forward when no activation quantizer needs one, and ``hf_ptq.py`` now sizes the default layerwise calibration batch at 8,192 tokens (16 samples at ``--calib_seq 512``) instead of 1. Layerwise GPTQ on HF fused-MoE models runs over 10x faster; pass ``--batch_size 1`` for the old batching, which pads no sample. - Add a checkpoint-specific NVIDIA Nemotron 3.5 Super VL AutoQuantize recipe with a 4.9-effective-bit NVFP4-MSE/FP8/BF16 search and fixed NVFP4 MTP experts. Select it under ``models/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/auto_quantize/``. - Add ``examples/hf_ptq/convert_ggml_to_int8.py`` to expand a checkpoint exported with a GGML recipe (IQ formats, ``q8_0``) into an INT8 weight-only compressed-tensors checkpoint. Stock vLLM serves it with its integer kernels at about half the BF16 weight memory, reproducing the GGML weights up to rounding their scales to the model dtype. +- Add an opt-in no-gather HuggingFace export for FSDP2-sharded models: with ``MODELOPT_FSDP2_EXPORT=nogather``, ``export_hf_checkpoint`` sends each tensor to a single owner rank and every rank writes its own shards in parallel, instead of gathering the full state dict on one host. MoE experts stay sharded through the export, so each rank only materializes its own experts. *Speculative Decoding* @@ -41,6 +42,7 @@ Changelog - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``. - ``modelopt.torch.utils.distributed.fsdp2_wrap`` moved to ``modelopt.torch.utils.plugins.model_load_utils``, with the other loaders. The old import path still works, but calling it now needs the ``hf`` extra. - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. +- ``export_hf_checkpoint`` and ``modelopt.torch.utils.dataset_utils.get_max_batch_size`` now read how a model is distributed from the plan ``load_pretrained`` records on it, and raise for an FSDP2-sharded model, or one in a multi-rank run, that carries none. If you distribute a model yourself, for example with ``fsdp2_wrap``, which records no plan, record one with ``modelopt.torch.utils.distributed.set_distribution_plan`` before calibrating or exporting it. **Deprecations** @@ -118,7 +120,6 @@ Changelog - ``KV_QUANT_CFG`` presets in vLLM fake-quant serving now quantize the MLA KV cache; on vLLM 0.16 and later they silently quantized nothing for MLA models. - vLLM fake-quant serving of NVFP4 and cast-mode KV-cache configs no longer needs ``--enforce-eager``. ``examples/vllm_serve/vllm_serve_fakequant.py`` now sets ``VLLM_DISABLE_COMPILE_CACHE=1`` for fake-quant serves by default, because a cached torch.compile graph of the same model without the fake quant would otherwise be reused. - ``examples/deepseek/deepseek_v4/ptq.py`` now supports DeepSeek-V4.1-Flash, whose reference implementation uses 32x32 FP8 blocks and a tokenizer-aware ``Transformer``; this is the setup used to produce ``nvidia/DeepSeek-V4.1-Flash-NVFP4``. DeepSeek-V4-Pro checkpoints work as before. -- Add an opt-in no-gather HuggingFace export for FSDP2-sharded models: with ``MODELOPT_FSDP2_EXPORT=nogather``, ``export_hf_checkpoint`` sends each tensor to a single owner rank and every rank writes its own shards in parallel, instead of gathering the full state dict on one host. MoE experts stay sharded through the export, so each rank only materializes its own experts. *Speculative Decoding* @@ -191,7 +192,6 @@ Changelog - ``get_te_hybrid_stack_spec`` was removed from ``modelopt.torch.nas.plugins.megatron``; it had no use outside tests. Use ``modelopt.torch.utils.plugins.megatron_layer_specs.te_hybrid_stack_spec_sequential_mlp`` for the SequentialMLP layout, or ``megatron.core.models.hybrid.hybrid_layer_specs.hybrid_stack_spec`` for grouped GEMM. - Unified HuggingFace export now fails with ``NotImplementedError`` when it meets an MoE block whose expert projection names it does not know, instead of assuming Mixtral's ``w1``/``w2``/``w3``. If you hit this, register a ``ModelSpec`` for the model under ``modelopt/torch/models/``. Every MoE architecture ModelOpt exported correctly before this change is registered, so no supported model regresses. - ``--recipe`` (and ``modelopt.recipe.load_recipe``) now resolve a recipe path **filesystem-first**: a recipe of the same relative path in the current working directory takes precedence over the shipped built-in of that name, matching how recipe ``$import`` paths already resolve. Previously the built-in won. -- ``export_hf_checkpoint`` and ``modelopt.torch.utils.dataset_utils.get_max_batch_size`` now read how a model is distributed from the plan ``load_pretrained`` records on it, and raise for an FSDP2-sharded model, or one in a multi-rank run, that carries none. If you distribute a model yourself, for example with ``fsdp2_wrap``, which records no plan, record one with ``modelopt.torch.utils.distributed.set_distribution_plan`` before calibrating or exporting it. **Deprecations** From f61087048be240b74cfe9fee6797ded0e522e71b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 22:31:33 +0000 Subject: [PATCH 06/24] Run the no-gather writer's two-rank test on the shared gloo workers It spawned two ranks, about 14 s of the linux unit job. Signed-off-by: Shengliang Xu --- tests/unit/torch/export/test_distributed_writer.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 165599251bf..30af5b752ba 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -36,7 +36,7 @@ import pytest import torch import torch.distributed as dist -from _test_utils.torch.distributed.utils import bounded_collectives, spawn_multiprocess_job +from _test_utils.torch.distributed.utils import bounded_collectives from torch.distributed.device_mesh import init_device_mesh from torch.distributed.fsdp import fully_shard from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor @@ -818,15 +818,15 @@ def _two_rank_job(rank, world_size, tmp_dir): _write_even_shards(sd, tmp_dir / "collision", "1MB", rank, world_size) -def test_two_ranks_write_one_checkpoint_and_reject_inconsistent_state(skip_on_windows, tmp_path): +def test_two_ranks_write_one_checkpoint_and_reject_inconsistent_state(cpu_dist_workers, tmp_path): """Shards cross ranks point to point, each tensor gets one owner, files share one numbering. - Spawning ranks dominates the cost, so one job also checks the cross-rank consistency errors. + The same job also checks the cross-rank consistency errors. """ src = _save(_mixtral_config(vocab_size=63, num_local_experts=3), tmp_path / "moe") _save(_llama_config(), tmp_path / "dense") - spawn_multiprocess_job(2, partial(_two_rank_job, tmp_dir=tmp_path), backend="gloo") + cpu_dist_workers.run(_two_rank_job, tmp_dir=tmp_path) shards = sorted(p.name for p in (tmp_path / "moe_export").glob("*.safetensors")) n = len(shards) From 7e2ac38450bd93eeeedac5f6527031aa14cdbbf8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 21:04:14 +0000 Subject: [PATCH 07/24] Expect the cross-rank name collision on every rank in the two-rank writer test The writer now raises it on every rank, so rank 0 no longer has to release rank 1 from the closing barrier; the dedicated test bounds that case, so the two-rank job no longer repeats it. Signed-off-by: Shengliang Xu --- tests/unit/torch/export/test_distributed_writer.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 30af5b752ba..58db13aedd7 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -809,13 +809,6 @@ def _two_rank_job(rank, world_size, tmp_dir): w = DTensor.from_local(w, mesh, [Shard(0)], run_check=False) with pytest.raises(RuntimeError, match="disagree on which tensors are dim-0 sharded"): _materialize_owned({"w": w}, set(), rank, world_size) - sd = {"model.norm.weight": torch.full((2,), float(rank))} - if rank == 0: - with pytest.raises(RuntimeError, match="name collision across ranks"): - _write_even_shards(sd, tmp_dir / "collision", "1MB", rank, world_size) - dist.barrier() # releases the other rank from the writer's closing barrier - else: - _write_even_shards(sd, tmp_dir / "collision", "1MB", rank, world_size) def test_two_ranks_write_one_checkpoint_and_reject_inconsistent_state(cpu_dist_workers, tmp_path): From 1390871c62c52fa917f30d6d0cfd0a936390a0b0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:00:20 +0000 Subject: [PATCH 08/24] Guard the fused-expert fold's FSDP2 window with a real export The window around the fold was guarded only by a test of two stubs' call order. The MoE no-gather export test used Qwen3-MoE, whose layers pack nothing after their experts, so it exported them quantized with or without the window. The window matters when a dense weight in the same FSDP2 unit is packed after the fold: that weight's own window unshards and reshards the unit, which puts FSDP2's original shards back over a fold FSDP2 does not track. Qwen2-MoE's shared expert is such a weight, so the test now exports Qwen2-MoE, which fails without the window, and the stub test goes. The GPU MoE export now requires every expert to be exported quantized, not just some weight. Signed-off-by: Shengliang Xu --- .../export/test_distributed_hf_export.py | 7 +++++ .../torch/export/test_distributed_writer.py | 9 ++++-- .../export/test_fused_experts_keep_fused.py | 31 ------------------- 3 files changed, 13 insertions(+), 34 deletions(-) diff --git a/tests/gpu/torch/export/test_distributed_hf_export.py b/tests/gpu/torch/export/test_distributed_hf_export.py index 84b2b6f81c9..5492e5261df 100644 --- a/tests/gpu/torch/export/test_distributed_hf_export.py +++ b/tests/gpu/torch/export/test_distributed_hf_export.py @@ -244,6 +244,13 @@ def test_fsdp2_distributed_export_is_complete(dist_workers, tmp_path, moe, quant want_shape = (*want_shape[:-1], want_shape[-1] // 2) assert got_shape == want_shape, f"{name}: shape {got_shape}, expected {want_shape}" assert quantized > 0, f"nothing was quantized -- {algo} config did not take effect" + if moe: + # Every expert, not just some weight: the experts are folded on their sharded weights, and a + # fold FSDP2 does not track is undone by its next reshard, leaving bf16 experts beside their + # scales. + experts = [k for k in expected if ".experts." in k and k.endswith(".weight")] + unquantized = [k for k in experts if exported[k][0] not in _QUANTIZED_DTYPES] + assert experts and not unquantized, f"experts exported unquantized: {unquantized[:5]}" # ---- 4. the non-weight files came along ---- exported_files = {p.name for p in export_dir.iterdir() if p.is_file()} diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 58db13aedd7..49d96f52262 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -45,7 +45,7 @@ pytest.importorskip("safetensors") pytest.importorskip("transformers") from safetensors.torch import load_file -from transformers import AutoModelForCausalLM, LlamaConfig, MixtralConfig, Qwen3MoeConfig +from transformers import AutoModelForCausalLM, LlamaConfig, MixtralConfig, Qwen2MoeConfig from transformers.core_model_loading import Chunk, WeightConverter import modelopt.torch.quantization as mtq @@ -747,15 +747,18 @@ def test_nogather_export_rejects_save_modelopt_state(gloo_world_of_one, tmp_path "quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG], ids=["fp8", "nvfp4"] ) def test_nogather_export_writes_moe_experts_quantized(gloo_world_of_one, tmp_path, quant_cfg): - config = Qwen3MoeConfig( + """Qwen2-MoE's shared expert follows its experts in the same FSDP2 unit, and packing it opens a + weight window on that unit. Unless the experts were folded inside a window of their own, which + hands FSDP2 the packed weights, that window's reshard puts FSDP2's unquantized shards back.""" + config = Qwen2MoeConfig( vocab_size=64, hidden_size=32, intermediate_size=32, moe_intermediate_size=32, + shared_expert_intermediate_size=32, num_hidden_layers=1, num_attention_heads=4, num_key_value_heads=2, - head_dim=8, num_experts=4, num_experts_per_tok=2, max_position_embeddings=32, diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index bb308ff1c27..cb0a085e777 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -43,7 +43,6 @@ pytest.importorskip("safetensors") from safetensors.torch import load_file, save_file -import modelopt.torch.export.hf_export_handlers as handlers import modelopt.torch.quantization as mtq from modelopt.torch.export.moe_utils import ( _dtensor_dim0_offset, @@ -743,36 +742,6 @@ def test_keep_fused_refuses_fused_experts_not_named_experts(): torch.testing.assert_close(experts.state_dict(), before, rtol=0, atol=0) -@pytest.mark.parametrize("dense", ["fsdp", "replicate"]) -def test_fused_experts_fold_in_a_weight_window_when_the_plan_shards_the_dense_part( - monkeypatch, dense -): - """FSDP2 owns the fused weights when it shards the dense part around them: folded outside a - weight-update window, they would be overwritten with its original shards at the next reshard. - Outside held windows the handler closes the window itself.""" - events = [] - - @contextlib.contextmanager - def window(root, modules, names=None): - events.append("open") - yield - events.append("close") - - monkeypatch.setattr(handlers, "fsdp2_aware_weight_update", window) - monkeypatch.setattr( - handlers, "_export_fused_experts_keep_fused", lambda module, dtype: events.append("fold") - ) - experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG) - container = nn.Module() - container.experts = experts - set_distribution_plan(container, DistributionPlan(2, ep_size=2, dense=dense)) - ctx = ExportContext(model=container, dtype=torch.bfloat16, keep_fused_experts=True) - - ExportModuleRegistry.match(experts)(PREFIX, experts, ctx) - - assert events == (["open", "fold", "close"] if dense == "fsdp" else ["fold"]) - - @pytest.mark.parametrize("held", [True, False], ids=["held", "per-module"]) @pytest.mark.parametrize("root_sharded", [True, False], ids=["fsdp2", "expert-parallel"]) def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( From 7da69b8bec2b83afc52fd6c26328918464568516 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:11:50 +0000 Subject: [PATCH 09/24] Check on CPU what the GPU-only no-gather export tests checked Every GPU test of the no-gather export needs two GPUs, which the PR runner does not have. What they checked runs on CPU instead, through the same writer and export. The two-rank write is now handed tensors the model does not own on every rank, and must write them once. A linear left unquantized already is checked: every no-gather equivalence case exports the lm_head unquantized, byte for byte as the single-process export does. The GPU completeness test ties the dense model's embeddings, so its expected keys also check that the alias is dropped. The GPU tests these cover go, with the GPU refusal of save_modelopt_state, which the unit test of the same refusal covers. Signed-off-by: Shengliang Xu --- .../export/test_distributed_hf_export.py | 200 +----------------- .../torch/export/test_distributed_writer.py | 24 ++- 2 files changed, 29 insertions(+), 195 deletions(-) diff --git a/tests/gpu/torch/export/test_distributed_hf_export.py b/tests/gpu/torch/export/test_distributed_hf_export.py index 5492e5261df..369385c73d9 100644 --- a/tests/gpu/torch/export/test_distributed_hf_export.py +++ b/tests/gpu/torch/export/test_distributed_hf_export.py @@ -30,14 +30,14 @@ quant config, weight index) and the source's own (tokenizer, vocab, chat template), which the exporter copies over after the writer has removed its ``sharded/`` staging dir. -It also covers the two things the writer must not quietly drop on the way: a declared tied -weight still gets deduplicated, and options the path cannot honour are rejected rather than -ignored. +The dense model ties its embeddings, so the expected keys also check that the writer drops the +alias. What needs no GPU -- tensors handed in through ``extra_state_dict``, a linear left +unquantized (the ``lm_head`` every case there exports), options the path refuses -- is checked on +CPU, in ``tests/unit/torch/export/test_distributed_writer.py``. Run with >=2 GPUs so the expert axis and the FSDP2 shard axis are both actually split. """ -import copy import json import os from functools import partial @@ -47,7 +47,6 @@ import pytest import torch from _test_utils.torch.transformers_models import create_tiny_qwen3_dir, create_tiny_qwen3_moe_dir -from safetensors.torch import load_file import modelopt.torch.quantization as mtq from modelopt.torch.export.unified_export_hf import export_hf_checkpoint @@ -151,7 +150,7 @@ def _no_gather_export(): return mock.patch.dict(os.environ, {"MODELOPT_FSDP2_EXPORT": "nogather"}) -def _ptq_and_export(rank, size, *, src_dir, export_dir, quant_cfg, **export_kwargs): +def _ptq_and_export(rank, size, *, src_dir, export_dir, quant_cfg): """Load the tiny model on every rank, FSDP2-shard it, PTQ it, and export.""" from transformers import AutoModelForCausalLM @@ -175,15 +174,13 @@ def _ptq_and_export(rank, size, *, src_dir, export_dir, quant_cfg, **export_kwar torch.distributed.barrier() with _no_gather_export(): - export_hf_checkpoint( - model, export_dir=export_dir, max_shard_size=MAX_SHARD_SIZE, **export_kwargs - ) + export_hf_checkpoint(model, export_dir=export_dir, max_shard_size=MAX_SHARD_SIZE) torch.distributed.barrier() # Four PTQ + distributed-export round trips; the tests/gpu default of 120s is not enough headroom. @pytest.mark.timeout(600) -@pytest.mark.parametrize("moe", [False, True], ids=["dense", "moe"]) +@pytest.mark.parametrize("moe", [False, True], ids=["dense-tied", "moe"]) @pytest.mark.parametrize( ("quant_cfg", "algo"), [(mtq.FP8_DEFAULT_CFG, "FP8"), (mtq.NVFP4_DEFAULT_CFG, "NVFP4")], @@ -194,7 +191,9 @@ def test_fsdp2_distributed_export_is_complete(dist_workers, tmp_path, moe, quant pytest.skip("needs >=2 GPUs to shard the expert and FSDP2 axes") make_dir = create_tiny_qwen3_moe_dir if moe else create_tiny_qwen3_dir - kwargs = {**TINY_KWARGS, **(TINY_MOE_KWARGS if moe else {})} + # fully_shard splits a tied parameter into per-module shards, so only the names still match: + # the source checkpoint holds no lm_head.weight, and the writer must not write one either. + kwargs = {**TINY_KWARGS, **(TINY_MOE_KWARGS if moe else {"tie_word_embeddings": True})} src_dir = Path(make_dir(tmp_path, with_tokenizer=True, **kwargs)) export_dir = tmp_path / "export" @@ -280,182 +279,3 @@ def test_fsdp2_distributed_export_is_complete(dist_workers, tmp_path, moe, quant json.loads((export_dir / "hf_quant_config.json").read_text())["quantization"]["quant_algo"] == algo ) - - -def _export_with(rank, size, *, src_dir, export_dir, **export_kwargs): - """Export once with ``export_kwargs``, for the rejection checks below.""" - from transformers import AutoModelForCausalLM - - with patch_fsdp_mp_dtypes(): - model = AutoModelForCausalLM.from_pretrained(src_dir, dtype=torch.bfloat16).to("cuda") - model.eval() - fsdp2_wrap(model) - set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) - assert is_fsdp2_model(model) - torch.distributed.barrier() - with _no_gather_export(): - export_hf_checkpoint(model, export_dir=export_dir, **export_kwargs) - - -@pytest.mark.timeout(600) -@pytest.mark.parametrize( - ("kwargs", "needle"), - [ - ({"save_modelopt_state": True}, "save_modelopt_state"), - ], - ids=["save_modelopt_state"], -) -def test_fsdp2_distributed_export_rejects_unsupported_options( - dist_workers, tmp_path, kwargs, needle -): - """Options the no-gather path cannot honour must raise, not be silently ignored. - - ``save_modelopt_state`` is forwarded to ``save_pretrained`` on the gather path; the distributed - path writes the checkpoint itself and never calls it, so a caller passing it would otherwise get - a checkpoint quietly missing the ModelOpt state. (``extra_state_dict`` used to be rejected here - too and is now supported -- see the MTP test below.) - """ - if torch.cuda.device_count() < 2: - pytest.skip("needs >=2 GPUs") - - src_dir = Path(create_tiny_qwen3_dir(tmp_path, with_tokenizer=True, **TINY_KWARGS)) - with pytest.raises(Exception, match=needle): - dist_workers.run( - partial(_export_with, src_dir=src_dir, export_dir=tmp_path / "rejected", **kwargs) - ) - - -@pytest.mark.timeout(600) -def test_fsdp2_distributed_export_dedups_tied_weights(dist_workers, tmp_path): - """A declared tie is deduplicated by name in the distributed write, as on the gather path. - - ``fully_shard`` splits the shared parameter into distinct per-module shards, so the tie survives - only as matching names -- both sides reach the writer as independent DTensors and would both be - written unless the writer applies the same ``TiedWeightMap`` the gather path uses. - """ - if torch.cuda.device_count() < 2: - pytest.skip("needs >=2 GPUs") - - src_dir = Path( - create_tiny_qwen3_dir( - tmp_path, with_tokenizer=True, tie_word_embeddings=True, **TINY_KWARGS - ) - ) - export_dir = tmp_path / "export_tied" - dist_workers.run( - partial( - _ptq_and_export, - src_dir=src_dir, - export_dir=export_dir, - quant_cfg=mtq.FP8_DEFAULT_CFG, - ) - ) - - exported = _safetensors_meta(export_dir) - assert exported, "nothing exported" - assert "model.embed_tokens.weight" in exported, "canonical tied weight missing" - assert "lm_head.weight" not in exported, ( - "tied alias 'lm_head.weight' was written alongside its canonical -- " - "name-based dedup did not run in the distributed writer" - ) - - -@pytest.mark.timeout(600) -def test_fsdp2_distributed_export_writes_extra_state_dict(dist_workers, tmp_path): - """Tensors handed in via ``extra_state_dict`` must reach the checkpoint. - - This is how weights the FSDP2 loader drops arrive: the caller reads them straight off the - source checkpoint and passes them through ``extra_state_dict``. They belong to no rank's - shard and no rank's model, so unless the writer deliberately merges them in, such a model - exports silently incomplete -- valid-looking, right size, missing the layer. - - They must also keep their original names: they were never part of the transformers conversion, - so the reverse conversion must not touch them. - """ - if torch.cuda.device_count() < 2: - pytest.skip("needs >=2 GPUs") - - src_dir = Path(create_tiny_qwen3_dir(tmp_path, with_tokenizer=True, **TINY_KWARGS)) - export_dir = tmp_path / "export_extra" - # Shaped and named like an inlined MTP tail (DeepSeek-V3 / GLM-5.1 put it at layers.). - extra = { - "model.layers.99.mtp.weight": torch.arange(64, dtype=torch.bfloat16).reshape(8, 8), - "model.layers.99.mtp.bias": torch.full((8,), 3.0, dtype=torch.bfloat16), - } - dist_workers.run( - partial( - _ptq_and_export, - src_dir=src_dir, - export_dir=export_dir, - quant_cfg=mtq.FP8_DEFAULT_CFG, - extra_state_dict=extra, - ) - ) - - exported = _safetensors_meta(export_dir) - assert exported, "nothing exported" - for name, want in extra.items(): - assert name in exported, f"extra_state_dict tensor '{name}' never reached the checkpoint" - assert exported[name][1] == tuple(want.shape), f"{name}: shape {exported[name][1]}" - - index = json.loads((export_dir / "model.safetensors.index.json").read_text()) - for name in extra: - assert name in index["weight_map"], f"'{name}' written but missing from the weight index" - - # Values must survive intact, not just the names. - merged: dict = {} - for shard in sorted(export_dir.glob("*.safetensors")): - merged.update(load_file(str(shard))) - for name, want in extra.items(): - assert torch.equal(merged[name].cpu(), want), f"{name}: value changed on the way out" - - -@pytest.mark.timeout(600) -def test_fsdp2_distributed_export_keeps_unquantized_submodule(dist_workers, tmp_path): - """A submodule left unquantized must still be exported, in full and unquantized. - - PTQ is routinely applied to part of a model -- the language model of a VLM, or a layer range - excluded by the recipe -- and the checkpoint is expected to be the COMPLETE model, with the - untouched part carried through at its original precision. The writer builds its state dict from - the whole model, so the risk is not that it drops those tensors but that it mangles them: the - expert split, the scale postprocess and the reverse conversion all walk keys that have no - quantizer state here. - """ - if torch.cuda.device_count() < 2: - pytest.skip("needs >=2 GPUs") - - # Disable every quantizer under layer 0; layer 1 stays quantized as the control. - quant_cfg = copy.deepcopy(mtq.FP8_DEFAULT_CFG) - quant_cfg["quant_cfg"].append({"quantizer_name": "*layers.0*", "enable": False}) - - src_dir = Path(create_tiny_qwen3_dir(tmp_path, with_tokenizer=True, **TINY_KWARGS)) - export_dir = tmp_path / "export_partial" - dist_workers.run( - partial(_ptq_and_export, src_dir=src_dir, export_dir=export_dir, quant_cfg=quant_cfg) - ) - - exported = _safetensors_meta(export_dir) - assert exported, "nothing exported" - - src_meta = _safetensors_meta(src_dir) - layer0_src = {k for k in src_meta if k.startswith("model.layers.0.")} - assert layer0_src, "fixture has no layer 0 weights to check" - - for name in sorted(layer0_src): - assert name in exported, f"unquantized '{name}' missing from the checkpoint" - dtype, shape = exported[name] - assert dtype not in _QUANTIZED_DTYPES, ( - f"'{name}' was excluded from quantization but exported as {dtype}" - ) - assert shape == src_meta[name][1], f"'{name}': shape {shape}, source {src_meta[name][1]}" - - # No scale companions should exist for the excluded layer... - stray = sorted(k for k in exported if k.startswith("model.layers.0.") and _is_scale(k)) - assert not stray, f"excluded layer 0 carries quantization scales: {stray}" - # ...while the control layer must actually be quantized, or the test proves nothing. - quantized_l1 = [ - k - for k, (d, _) in exported.items() - if k.startswith("model.layers.1.") and d in _QUANTIZED_DTYPES - ] - assert quantized_l1, "layer 1 was not quantized -- the exclusion pattern disabled everything" diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 49d96f52262..09c9708b6bd 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -777,12 +777,25 @@ def test_nogather_export_writes_moe_experts_quantized(gloo_world_of_one, tmp_pat # ------------------------------------------------------------------ the writer, two ranks +def _mtp_tail(): + """Tensors the model does not own, named like an inlined MTP tail the loader dropped.""" + return { + "model.layers.99.mtp.weight": torch.arange(64, dtype=torch.bfloat16).reshape(8, 8), + "model.layers.99.mtp.bias": torch.full((8,), 3.0, dtype=torch.bfloat16), + } + + def _two_rank_job(rank, world_size, tmp_dir): mesh = init_device_mesh("cpu", (world_size,)) - # An unquantized MoE whose vocab rows and experts split unevenly; the test compares what the - # two ranks wrote between them with the source checkpoint. - _write(_fsdp2_shard(_load(tmp_dir / "moe"), mesh), tmp_dir / "moe_export") + # An unquantized MoE whose vocab rows and experts split unevenly, written with tensors it does + # not own, which every rank is handed; the test compares what the two ranks wrote between them + # with the source checkpoint and those tensors. + _write( + _fsdp2_shard(_load(tmp_dir / "moe"), mesh), + tmp_dir / "moe_export", + extra_state_dict=_mtp_tail(), + ) # A quantized dense model through the export entry point: fp8 weights and bf16 norms travel # in separate per-dtype buckets. Every rank quantizes identically, so rank 0's single-process @@ -815,7 +828,8 @@ def _two_rank_job(rank, world_size, tmp_dir): def test_two_ranks_write_one_checkpoint_and_reject_inconsistent_state(cpu_dist_workers, tmp_path): - """Shards cross ranks point to point, each tensor gets one owner, files share one numbering. + """Shards cross ranks point to point, each tensor gets one owner, files share one numbering, + and the tensors handed to every rank are written once. The same job also checks the cross-rank consistency errors. """ @@ -828,5 +842,5 @@ def test_two_ranks_write_one_checkpoint_and_reject_inconsistent_state(cpu_dist_w n = len(shards) assert n > 2, "both ranks should have written several files" assert shards == [f"model-{i:05d}-of-{n:05d}.safetensors" for i in range(1, n + 1)] - want = load_file(str(src / "model.safetensors")) + want = {**load_file(str(src / "model.safetensors")), **_mtp_tail()} _assert_byte_identical(_load_checkpoint(tmp_path / "moe_export"), want) From 9f34617cfbdf9c7736b110eb9c898d2a7496be31 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 18:22:45 +0000 Subject: [PATCH 10/24] Trim the no-gather export's tests: fold its end-to-end cases, drop covered ones Merged: - The inference-mode and extra_state_dict writes become cases of test_distributed_save_writes_the_hub_checkpoint_of_a_sharded_moe. - test_a_model_built_from_a_config_names_no_source_checkpoint folds into test_source_checkpoint_is_none_for_a_model_built_from_config, which now also checks that nothing is copied from the working directory. - test_materialize_owned_leaves_every_tensor_whole_and_plain moves into the two-rank job, where ownership is real: its tensors cycle through dim-0 shards, replicated DTensors and plain tensors, plus a zero-row tensor. Dropping, duplicating or not unwrapping any of them, treating experts as dense, or writing the zero-row tensor as float32 fails the job. Dropped, each caught by a remaining test when its code is broken: - test_revert_whole_sd_keeps_a_state_dict_with_nothing_to_revert: wiping a dense state dict fails test_nogather_export_writes_what_the_single_process_export_writes. - test_revert_whole_sd_renames_to_hub_names_in_place: discarding the reverse fails the hub-checkpoint test. - test_a_pipeline_is_exported_without_a_distribution_plan: reading a pipeline's plan fails test_export_diffusers_pipeline_config on a real tiny pipeline. The two-rank job checks the consistency errors with try/except inside bounded_collectives rather than pytest.raises, which takes the shared worker down when nothing is raised. Kept, since no remaining test fails when their code is broken: the split's storage, dense pass-through and prefix tests and the handler test in test_fused_experts_keep_fused.py, the mesh-memo test, and test_write_even_shards_writes_even_shards_and_an_index_of_every_tensor, the only check that written tensors leave the rank's set, of the file count, and of a non-contiguous tensor. Signed-off-by: Shengliang Xu --- .../torch/export/test_distributed_writer.py | 161 ++++++------------ .../torch/export/test_export_diffusers.py | 17 -- .../torch/export/test_unified_export_hf.py | 13 +- 3 files changed, 51 insertions(+), 140 deletions(-) diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 09c9708b6bd..b8d5c09dd75 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -29,8 +29,6 @@ import math import os import types -from datetime import timedelta -from functools import partial from unittest import mock import pytest @@ -227,29 +225,6 @@ def test_is_full_world_shard0_answers_from_the_mesh_memo(gloo_world_of_one): assert _is_full_world_shard0(tensor, 1, {id(gloo_world_of_one): False}) is False -def test_materialize_owned_leaves_every_tensor_whole_and_plain(gloo_world_of_one): - mesh = gloo_world_of_one - full = torch.arange(12.0).reshape(4, 3) - expert_key, expert = "model.layers.0.mlp.experts.0.up_proj.weight", torch.full((2,), 7.0) - sd = { - "plain": torch.ones(3), - "replicated": distribute_tensor(full, mesh, [Replicate()]), - "sharded": distribute_tensor(full, mesh, [Shard(0)]), - "empty": distribute_tensor(torch.zeros(0, 4, dtype=torch.bfloat16), mesh, [Shard(0)]), - expert_key: expert, - } - - _materialize_owned(sd, {expert_key}, 0, 1) - - assert set(sd) == {"plain", "replicated", "sharded", "empty", expert_key} - assert not any(isinstance(v, DTensor) for v in sd.values()) - assert torch.equal(sd["plain"], torch.ones(3)) - assert torch.equal(sd["replicated"], full) - assert torch.equal(sd["sharded"], full) - assert sd["empty"].shape == (0, 4) and sd["empty"].dtype == torch.bfloat16 - assert sd[expert_key] is expert # experts are already disjoint per rank: left alone - - def _padded_shard(mesh): """This rank's 2 rows of a 3-row tensor: a padded local chunk, which would concatenate wrong.""" return DTensor.from_local( @@ -408,42 +383,6 @@ def test_materialize_owned_moves_each_shard_to_its_owner_over_two_ranks(cpu_dist cpu_dist_workers.run(_materialize_two_ranks_job) -def test_revert_whole_sd_keeps_a_state_dict_with_nothing_to_revert(): - """A model without a conversion mapping gets its input dict back from the reverse pass.""" - model = AutoModelForCausalLM.from_config(_llama_config()) - sd = dict(model.state_dict()) - before = dict(sd) - - assert _revert_whole_sd(model, sd) is True - - assert sd == before - - -def test_revert_whole_sd_renames_to_hub_names_in_place(): - model = AutoModelForCausalLM.from_config(_mixtral_config()) - gate, w1, w3, w2 = (torch.full((2, 2), float(i)) for i in range(4)) - sd = { - "model.layers.0.mlp.gate.weight": gate, - "model.layers.0.mlp.experts.0.gate_proj.weight": w1, - "model.layers.0.mlp.experts.0.up_proj.weight": w3, - "model.layers.0.mlp.experts.0.down_proj.weight": w2, - } - same_dict = sd - - assert _revert_whole_sd(model, sd) is True - - assert sd is same_dict - hub = "model.layers.0.block_sparse_moe" - assert set(sd) == { - f"{hub}.gate.weight", - f"{hub}.experts.0.w1.weight", - f"{hub}.experts.0.w3.weight", - f"{hub}.experts.0.w2.weight", - } - assert torch.equal(sd[f"{hub}.experts.0.w1.weight"], w1) - assert torch.equal(sd[f"{hub}.experts.0.w2.weight"], w2) - - def _one_rank_cannot_revert_job(rank, world_size): # Rank 0's subset reverts; rank 1's fails, as a merge whose group ownership split does. model = AutoModelForCausalLM.from_config(_mixtral_config()) @@ -652,56 +591,39 @@ def test_a_name_two_ranks_both_write_fails_the_write_on_every_rank(cpu_dist_work # ------------------------------------------------------------------ the writer, one rank -@pytest.mark.parametrize("tie", [True, False], ids=["tied", "untied"]) +@pytest.mark.parametrize("case", ["tied", "untied", "inference-mode", "extra-state-dict"]) def test_distributed_save_writes_the_hub_checkpoint_of_a_sharded_moe( - gloo_world_of_one, tmp_path, tie + gloo_world_of_one, tmp_path, case ): - """Fused experts are split, every name is reverse-converted, and a declared tie is dropped.""" - src = _save(_mixtral_config(tie_word_embeddings=tie), tmp_path / "src") - model = _fsdp2_shard(_load(src), gloo_world_of_one) + """Fused experts are split, every name is reverse-converted, and a declared tie is dropped. - _write(model, tmp_path / "export") - - assert len(list((tmp_path / "export").glob("*.safetensors"))) > 1 - want = load_file(str(src / "model.safetensors")) - _assert_byte_identical(_load_checkpoint(tmp_path / "export"), want) - - -def test_distributed_save_runs_from_inside_inference_mode(gloo_world_of_one, tmp_path): - """A caller in inference mode hands the writer inference tensors, which its writes reject.""" - src = _save(_mixtral_config(), tmp_path / "src") - model = _fsdp2_shard(_load(src), gloo_world_of_one) - - with torch.inference_mode(): - # Made in inference mode, like the scales an inference-mode pack registers. - model.model.norm.register_buffer("weight_scale", torch.full((), 0.5)) - _write(model, tmp_path / "export") - - want = { - **load_file(str(src / "model.safetensors")), - "model.norm.weight_scale": torch.tensor(0.5), - } - _assert_byte_identical(_load_checkpoint(tmp_path / "export"), want) - - -def test_distributed_save_writes_extra_state_dict_under_its_own_names(gloo_world_of_one, tmp_path): - """Extra tensors already carry hub names, so the reverse conversion must not rename them. - - The first is named like an in-memory Mixtral expert weight, which the reverse would turn into + A caller in inference mode hands the writer inference tensors, which its writes reject. Extra + tensors already carry hub names, so the reverse conversion must not rename them: the first is + named like an in-memory Mixtral expert weight, which the reverse would turn into ``block_sparse_moe.experts.0.w1``. """ - src = _save(_mixtral_config(), tmp_path / "src") + src = _save(_mixtral_config(tie_word_embeddings=case == "tied"), tmp_path / "src") model = _fsdp2_shard(_load(src), gloo_world_of_one) - extra = { - "model.layers.9.mlp.experts.0.gate_proj.weight": torch.arange( - 64, dtype=torch.bfloat16 - ).reshape(8, 8), - "model.mtp.norm.weight": torch.full((8,), 3.0, dtype=torch.bfloat16), - } + want = load_file(str(src / "model.safetensors")) + extra = None + if case == "extra-state-dict": + extra = { + "model.layers.9.mlp.experts.0.gate_proj.weight": torch.arange( + 64, dtype=torch.bfloat16 + ).reshape(8, 8), + "model.mtp.norm.weight": torch.full((8,), 3.0, dtype=torch.bfloat16), + } + want.update(extra) + elif case == "inference-mode": + want["model.norm.weight_scale"] = torch.tensor(0.5) - _write(model, tmp_path / "export", extra_state_dict=extra) + with torch.inference_mode(case == "inference-mode"): + if case == "inference-mode": + # Made in inference mode, like the scales an inference-mode pack registers. + model.model.norm.register_buffer("weight_scale", torch.full((), 0.5)) + _write(model, tmp_path / "export", extra_state_dict=extra) - want = {**load_file(str(src / "model.safetensors")), **extra} + assert len(list((tmp_path / "export").glob("*.safetensors"))) > 1 _assert_byte_identical(_load_checkpoint(tmp_path / "export"), want) @@ -806,25 +728,40 @@ def _two_rank_job(rank, world_size, tmp_dir): if rank == 0: _assert_byte_identical(_load_checkpoint(tmp_dir / "dense_export"), want) - # Each dense tensor ends up whole on exactly one rank, and the owners are spread out. - full = {f"t{i}": torch.full((4, 2), float(i)) for i in range(8)} - sd = {k: distribute_tensor(v, mesh, [Shard(0)]) for k, v in full.items()} + # Each dense tensor ends up whole and plain on exactly one rank, whatever its layout, and the + # owners are spread out. A zero-row tensor has no rows to move: its owner writes it as it is. + full = {f"t{i}": torch.full((4, 2), float(i)) for i in range(9)} + full["t9"] = torch.zeros(0, 4, dtype=torch.bfloat16) + sd = {} + for i, (k, v) in enumerate(full.items()): + layout = ([Shard(0)], [Replicate()], None)[i % 3] + sd[k] = v if layout is None else distribute_tensor(v, mesh, layout) _materialize_owned(sd, set(), rank, world_size) owned = [None] * world_size dist.all_gather_object(owned, sorted(sd)) assert all(owned), f"a rank owns nothing: {owned}" assert sorted(k for keys in owned for k in keys) == sorted(full) - assert all(torch.equal(v, full[k]) for k, v in sd.items()) + for k, v in sd.items(): + assert not isinstance(v, DTensor), k + assert v.dtype == full[k].dtype and torch.equal(v, full[k]), k # Ranks that disagree must fail with a diagnosis, not hang in a collective or drop a tensor. - sd = {"shared": torch.ones(2), f"rank{rank}": torch.ones(2)} - with pytest.raises(RuntimeError, match="dense key sets differ"): - _materialize_owned(sd, set(), rank, world_size) w = torch.ones(1, 2) if rank == 0: w = DTensor.from_local(w, mesh, [Shard(0)], run_check=False) - with pytest.raises(RuntimeError, match="disagree on which tensors are dim-0 sharded"): - _materialize_owned({"w": w}, set(), rank, world_size) + cases = [ + ({"shared": torch.ones(2), f"rank{rank}": torch.ones(2)}, "dense key sets differ"), + ({"w": w}, "disagree on which tensors are dim-0 sharded"), + ] + for sd, diagnosis in cases: + try: + with bounded_collectives(): + _materialize_owned(sd, set(), rank, world_size) + except RuntimeError as error: + raised = str(error) + else: + raised = "no error" + assert diagnosis in raised, raised def test_two_ranks_write_one_checkpoint_and_reject_inconsistent_state(cpu_dist_workers, tmp_path): diff --git a/tests/unit/torch/export/test_export_diffusers.py b/tests/unit/torch/export/test_export_diffusers.py index 3a39947246a..001b4245b72 100644 --- a/tests/unit/torch/export/test_export_diffusers.py +++ b/tests/unit/torch/export/test_export_diffusers.py @@ -16,7 +16,6 @@ import copy import inspect import json -import types from pathlib import Path from types import SimpleNamespace @@ -240,22 +239,6 @@ def __init__(self, unet, scheduler, setting=None): assert {path.name for path in export_dir.iterdir() if path.is_dir()} == {"unet"} -def test_a_pipeline_is_exported_without_a_distribution_plan(tmp_path, monkeypatch): - """A pipeline is not a module: nothing distributes one, so the export reads no plan of it.""" - exported = [] - monkeypatch.setattr(unified_export_hf, "is_diffusers_object", lambda model: True) - monkeypatch.setattr( - unified_export_hf, - "_export_diffusers_checkpoint", - lambda pipe, *a, **k: exported.append(pipe), - ) - pipe = types.SimpleNamespace(config=None) - - export_hf_checkpoint(pipe, export_dir=tmp_path) - - assert exported == [pipe] - - def test_export_diffusers_unet_quantized_matches_llm_config(tmp_path, monkeypatch): model = get_tiny_unet() export_dir = tmp_path / "export_unet_quant" diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index eba57c4f264..c1f89329605 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -39,6 +39,7 @@ sync_tied_input_amax, ) from modelopt.torch.export.unified_export_hf import ( + _copy_non_model_files_from_source, _resolve_export_dtype, _source_checkpoint, read_unplaced_weights, @@ -1011,19 +1012,9 @@ def test_an_unrecorded_model_derives_its_unplaced_keys_once(tmp_path, monkeypatc assert model._modelopt_unplaced_source_keys == [mtp] -def test_source_checkpoint_is_none_for_a_model_built_from_config(): +def test_source_checkpoint_is_none_for_a_model_built_from_config(tmp_path, monkeypatch): """``from_config`` leaves ``_name_or_path`` empty; it must not resolve to the working directory, or every export would copy the cwd as the "source checkpoint".""" - assert _source_checkpoint(SimpleNamespace(config=SimpleNamespace(_name_or_path=""))) is None - - -def test_a_model_built_from_a_config_names_no_source_checkpoint(tmp_path, monkeypatch): - """An empty ``_name_or_path`` is the working directory as a path; nothing there is the model's.""" - from modelopt.torch.export.unified_export_hf import ( - _copy_non_model_files_from_source, - _source_checkpoint, - ) - workdir = tmp_path / "workdir" workdir.mkdir() (workdir / "tokenizer.json").write_text("{}") From 0b59ebb36fb38c2b1f0368a8052f9443f0e7d50c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 21:54:22 +0000 Subject: [PATCH 11/24] Run the writer's input-scale reduce on the CPU under a gloo-only group The no-gather writer took its device from torch.cuda.current_device() on every call, though only the expert input-scale all-reduce needs it. On CPU torch that raised, so every writer unit test patched current_device. Take the device from the backend: the CPU under a gloo-only group, the current CUDA device otherwise (distributed_run sets up cpu:gloo,cuda:nccl). The tests drop the patch, and fail on CPU torch without the change. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 8 +++++++- .../torch/export/test_distributed_writer.py | 17 ++++------------- 2 files changed, 11 insertions(+), 14 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index decf61ab99e..df87045674d 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -682,7 +682,13 @@ def _distributed_save_hf_checkpoint_impl( export_dir = Path(export_dir) rank, world = dist.get_rank(), dist.get_world_size() - device = torch.device(f"cuda:{torch.cuda.current_device()}") + # Where the expert input-scale all-reduce below runs: on the GPU under NCCL (distributed_run + # sets up "cpu:gloo,cuda:nccl"), on the CPU under a gloo-only group. + device = ( + torch.device("cpu") + if dist.get_backend() == "gloo" + else torch.device("cuda", torch.cuda.current_device()) + ) # This writer targets a FULLY FSDP2-sharded model -- no expert parallelism. FSDP2 # shards the fused 3-D expert weight on dim 0 (the expert axis), so fused experts are just diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index b8d5c09dd75..dcf5b266cd1 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -66,23 +66,14 @@ from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan -def _cpu_device_query(): - """The writer asks CUDA which device to all-reduce MoE expert input scales on. - - Nothing written in this file has expert input scales, so on a CPU-only host any answer works. - """ - return mock.patch.object(torch.cuda, "current_device", return_value=0) - - def _write(model, export_dir, **kwargs): - with _cpu_device_query(): - distributed_save_hf_checkpoint( - model, export_dir, maxbound=448, kv_cache_format=None, max_shard_size="8KB", **kwargs - ) + distributed_save_hf_checkpoint( + model, export_dir, maxbound=448, kv_cache_format=None, max_shard_size="8KB", **kwargs + ) def _nogather_export(model, export_dir, **kwargs): - with _cpu_device_query(), mock.patch.dict(os.environ, {"MODELOPT_FSDP2_EXPORT": "nogather"}): + with mock.patch.dict(os.environ, {"MODELOPT_FSDP2_EXPORT": "nogather"}): export_hf_checkpoint( model, dtype=torch.bfloat16, export_dir=export_dir, max_shard_size="16KB", **kwargs ) From 439793a73f2033361e7d0e6661bbe21088ca7f8c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 21:55:19 +0000 Subject: [PATCH 12/24] Check the extra tensors' overlap inside a collective every rank joins distributed_save_hf_checkpoint checked extra_state_dict against the model's keys on rank 0 only, outside any collective. A collision raised on rank 0 while the other ranks went on into the write's first collective and waited there until the timeout, then raised an unrelated timeout. Run the check inside all_reduce, so every rank raises rank 0's error at once. The two-rank writer job covers it, and now checks rank 0's checkpoint bytes after its last collective, so a failed check cannot leave rank 1 waiting either. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 27 +++++++++++----- .../torch/export/test_distributed_writer.py | 31 ++++++++++++------- 2 files changed, 39 insertions(+), 19 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index df87045674d..40da2282597 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -839,12 +839,23 @@ def _is_expert_key(key: str) -> bool: # to agree on the split. They go in AFTER the reverse conversion because they carry their # original hub names and were never converted -- reverting them would rename them wrongly. This # mirrors the gather path, which merges extra_state_dict into the already-reverted dict. - if extra_state_dict and rank == 0: - collisions = sorted(set(extra_state_dict) & set(local_sd)) - if collisions: - raise RuntimeError( - f"extra_state_dict overlaps the exported model weights on {len(collisions)} key(s), " - f"so the checkpoint would be ambiguous; first few: {collisions[:5]}" - ) - local_sd.update(extra_state_dict) + if extra_state_dict: + from modelopt.torch.utils.distributed import all_reduce + + def check_no_overlap() -> None: + if rank != 0: + return + collisions = sorted(set(extra_state_dict) & set(local_sd)) + if collisions: + raise RuntimeError( + f"extra_state_dict overlaps the exported model weights on {len(collisions)} " + f"key(s), so the checkpoint would be ambiguous; first few: {collisions[:5]}" + ) + + # Every rank is handed the same tensors, so every rank checks here: a collision rank 0 + # finds then raises on all of them, instead of leaving the others in the write's first + # collective. + all_reduce("extra_state_dict overlap", check_no_overlap, lambda _: None) + if rank == 0: + local_sd.update(extra_state_dict) _write_even_shards(local_sd, export_dir, max_shard_size, rank, world) diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index dcf5b266cd1..096c5f341a7 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -618,15 +618,6 @@ def test_distributed_save_writes_the_hub_checkpoint_of_a_sharded_moe( _assert_byte_identical(_load_checkpoint(tmp_path / "export"), want) -def test_distributed_save_rejects_extra_state_dict_that_overlaps_the_model( - gloo_world_of_one, tmp_path -): - model = _fsdp2_shard(_load(_save(_mixtral_config(), tmp_path / "src")), gloo_world_of_one) - - with pytest.raises(RuntimeError, match="overlaps the exported model weights"): - _write(model, tmp_path / "export", extra_state_dict={"model.norm.weight": torch.ones(32)}) - - @pytest.mark.parametrize("wrap_root", [True, False], ids=["root-wrapped", "layers-only"]) @pytest.mark.parametrize( "quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG], ids=["fp8", "nvfp4"] @@ -716,8 +707,6 @@ def _two_rank_job(rank, world_size, tmp_dir): model = _load(tmp_dir / "dense") want = _quantize_with_reference(model, mtq.FP8_DEFAULT_CFG) _nogather_export(_fsdp2_shard(model, mesh), tmp_dir / "dense_export") - if rank == 0: - _assert_byte_identical(_load_checkpoint(tmp_dir / "dense_export"), want) # Each dense tensor ends up whole and plain on exactly one rank, whatever its layout, and the # owners are spread out. A zero-row tensor has no rows to move: its owner writes it as it is. @@ -754,6 +743,26 @@ def _two_rank_job(rank, world_size, tmp_dir): raised = "no error" assert diagnosis in raised, raised + # Tensors handed in that collide with the model's: rank 0 alone adds them, and every rank + # raises, not rank 0 alone. + overlap = {k: torch.zeros(1) for k in _load(tmp_dir / "dense").state_dict()} + try: + with bounded_collectives(): + _write( + _fsdp2_shard(_load(tmp_dir / "dense"), mesh), + tmp_dir / "overlap_export", + extra_state_dict=overlap, + ) + except RuntimeError as error: + raised = str(error) + else: + raised = "no error" + assert "overlaps the exported model weights" in raised, raised + + # Last, after every collective: a failing rank 0 then leaves no rank waiting in one. + if rank == 0: + _assert_byte_identical(_load_checkpoint(tmp_dir / "dense_export"), want) + def test_two_ranks_write_one_checkpoint_and_reject_inconsistent_state(cpu_dist_workers, tmp_path): """Shards cross ranks point to point, each tensor gets one owner, files share one numbering, From ff947e2a2e8fa5bb68b1548df6ab4cce20316ba1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 21:56:15 +0000 Subject: [PATCH 13/24] Drop the export config helper's second generation-config write _write_hf_export_config rewrote generation_config.json, which every export route has already written by then: save_pretrained on the plain path, save_non_weight_artifacts on the streaming paths and _write_base_config, behind can_generate(), on the no-gather path. The layerwise path returns before it. Main's helper never wrote it. Signed-off-by: Shengliang Xu --- modelopt/torch/export/unified_export_hf.py | 19 ++----------------- 1 file changed, 2 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index c8aaddc5c4d..4ba4a1fb330 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1662,20 +1662,6 @@ def _write_hf_export_config( with open(original_config, "w") as file: json.dump(config_data, file, indent=4) - # generation_config.json: on the plain path save_pretrained emits it next to config.json, but the - # distributed writer writes weight shards only and save_pretrained never runs -- so without this the - # exported checkpoint silently loses the model's generation defaults. Measured on GLM-5.3-BF16: - # byte-identical weight shards on both paths, yet only the non-distributed export carried - # generation_config.json (do_sample / eos_token_id). Silent, and it lands on exactly the knobs that - # decide sampling behaviour at serving time. - gen_cfg = getattr(model, "generation_config", None) - if gen_cfg is not None: - try: - _sanitize_generation_config_for_save(model) - gen_cfg.save_pretrained(str(export_dir)) - except Exception as gen_err: # never fail an otherwise-good export over metadata - warnings.warn(f"Could not save generation_config: {gen_err}") - def _revert_hf_quant_config_names( hf_quant_config: dict, @@ -1923,9 +1909,8 @@ def _write_base_config(model: nn.Module, export_dir: "Path | str") -> None: if model.can_generate() and getattr(model, "generation_config", None) is not None: try: # Same strict-validate trip as every other save site: a config carrying top_p/top_k - # with do_sample unset (e.g. GLM-5.3-BF16) makes save_pretrained raise, and without - # this the only trace is a "Could not save generation_config" warning that reads as - # data loss when the later _write_hf_export_config has in fact written the file. + # with do_sample unset (e.g. GLM-5.3-BF16) makes save_pretrained raise, and the + # export would then carry the source's generation_config.json instead of this one. _sanitize_generation_config_for_save(model) model.generation_config.save_pretrained(export_dir) except Exception as gen_err: From 8e801408dd84a9199041376f8fd47fc868c2afe6 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 17 Oct 2026 01:06:36 +0000 Subject: [PATCH 14/24] Join the extra tensors' overlap check on every rank, holding extras or not The overlap check ran its all_reduce only where extra_state_dict was non-empty, but a sharded export hands the tensors it carries over (an MTP head, an auxiliary tower) to rank 0 alone. Rank 0 then entered the overlap collective while the others went on to "shard counts", and every sharded export carrying weights failed with "Ranks reduced out of step". Every rank joins now; only rank 0, which holds the extras, checks and adds them. The two-rank writer job hands its extras to rank 0 only, as the export does. Found by the epdp port. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 35 +++++++++---------- .../torch/export/test_distributed_writer.py | 4 +-- 2 files changed, 19 insertions(+), 20 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index 40da2282597..f27ac69483b 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -839,23 +839,22 @@ def _is_expert_key(key: str) -> bool: # to agree on the split. They go in AFTER the reverse conversion because they carry their # original hub names and were never converted -- reverting them would rename them wrongly. This # mirrors the gather path, which merges extra_state_dict into the already-reverted dict. - if extra_state_dict: - from modelopt.torch.utils.distributed import all_reduce - - def check_no_overlap() -> None: - if rank != 0: - return - collisions = sorted(set(extra_state_dict) & set(local_sd)) - if collisions: - raise RuntimeError( - f"extra_state_dict overlaps the exported model weights on {len(collisions)} " - f"key(s), so the checkpoint would be ambiguous; first few: {collisions[:5]}" - ) + from modelopt.torch.utils.distributed import all_reduce - # Every rank is handed the same tensors, so every rank checks here: a collision rank 0 - # finds then raises on all of them, instead of leaving the others in the write's first - # collective. - all_reduce("extra_state_dict overlap", check_no_overlap, lambda _: None) - if rank == 0: - local_sd.update(extra_state_dict) + def check_no_overlap() -> None: + if rank != 0 or not extra_state_dict: + return + collisions = sorted(set(extra_state_dict) & set(local_sd)) + if collisions: + raise RuntimeError( + f"extra_state_dict overlaps the exported model weights on {len(collisions)} " + f"key(s), so the checkpoint would be ambiguous; first few: {collisions[:5]}" + ) + + # The export entry hands the tensors it carries over to rank 0 alone, so whether a rank holds + # any cannot decide whether it joins: every rank does, and a collision rank 0 finds then raises + # on all of them, instead of leaving the others in the write's first collective. + all_reduce("extra_state_dict overlap", check_no_overlap, lambda _: None) + if extra_state_dict and rank == 0: + local_sd.update(extra_state_dict) _write_even_shards(local_sd, export_dir, max_shard_size, rank, world) diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 096c5f341a7..dd0a8465133 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -698,7 +698,7 @@ def _two_rank_job(rank, world_size, tmp_dir): _write( _fsdp2_shard(_load(tmp_dir / "moe"), mesh), tmp_dir / "moe_export", - extra_state_dict=_mtp_tail(), + extra_state_dict=_mtp_tail() if rank == 0 else None, ) # A quantized dense model through the export entry point: fp8 weights and bf16 norms travel @@ -751,7 +751,7 @@ def _two_rank_job(rank, world_size, tmp_dir): _write( _fsdp2_shard(_load(tmp_dir / "dense"), mesh), tmp_dir / "overlap_export", - extra_state_dict=overlap, + extra_state_dict=overlap if rank == 0 else None, ) except RuntimeError as error: raised = str(error) From f9d1e03f276b96eed515e537f7b1d333ac6d5865 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 20:13:34 +0000 Subject: [PATCH 15/24] Say what picks the no-gather writer in the distributed export test The comment credited is_fsdp2_model with choosing the writer; the recorded plan and the MODELOPT_FSDP2_EXPORT opt-in choose it. The assertion stays: it checks the model is really sharded, so the test exercises the shard writer. Signed-off-by: Shengliang Xu --- tests/gpu/torch/export/test_distributed_hf_export.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/gpu/torch/export/test_distributed_hf_export.py b/tests/gpu/torch/export/test_distributed_hf_export.py index 369385c73d9..fe88fe8fd95 100644 --- a/tests/gpu/torch/export/test_distributed_hf_export.py +++ b/tests/gpu/torch/export/test_distributed_hf_export.py @@ -162,9 +162,9 @@ def _ptq_and_export(rank, size, *, src_dir, export_dir, quant_cfg): set_distribution_plan( model, DistributionPlan(size, dense="fsdp") ) # wrapped by hand, so recorded here, as load_pretrained would - # The no-gather writer is selected by this predicate plus the opt-in below, so assert it here: - # without it the export would silently fall back to the rank-0 gather and the test would pass - # while never touching the code under test. + # The no-gather writer is chosen by the plan recorded above plus the opt-in below. It writes + # each rank's FSDP2 shards, so check that fsdp2_wrap sharded the model: unsharded, the test + # would not exercise it. assert is_fsdp2_model(model), "fsdp2_wrap did not shard the model" assert torch.distributed.is_initialized() torch.distributed.barrier() From bf309b607305a5291414d7f43211e879b58c2eba Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 20:57:17 +0000 Subject: [PATCH 16/24] Resolve a mixed-precision KV cache on the no-gather export, and take rank 0's writer choice The no-gather export passed the raw kv_cache_quant_algo to the writer, so a mixed-precision KV cache reached postprocess_state_dict as the "MIXED_PRECISION" sentinel and failed its format assert. It now resolves the per-layer map through _get_kv_cache_postprocess_config, as the gather and streaming exports do; the single-process comparison covers an FP8 and NVFP4 KV cache. Which writer runs is rank 0's MODELOPT_FSDP2_EXPORT, broadcast: the two writers issue different collectives. The two-rank export sets the knob on rank 0 only. The no-gather branch drops a save_modelopt_state write it refused earlier, and _write_base_config says why a config save needs no revert_weight_conversion patch. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 4 +-- modelopt/torch/export/unified_export_hf.py | 25 ++++++++----- .../torch/export/test_distributed_writer.py | 36 ++++++++++++++++--- 3 files changed, 50 insertions(+), 15 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index f27ac69483b..2af31586495 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -35,7 +35,7 @@ def distributed_save_hf_checkpoint( model: nn.Module, export_dir: "str | Path", maxbound: float, - kv_cache_format: "str | None", + kv_cache_format: "str | dict[str, dict[str, str]] | None", max_shard_size: "str | int" = "10GB", is_modelopt_qlora: bool = False, extra_state_dict: "dict[str, torch.Tensor] | None" = None, @@ -640,7 +640,7 @@ def _distributed_save_hf_checkpoint_impl( model: nn.Module, export_dir: "str | Path", maxbound: float, - kv_cache_format: "str | None", + kv_cache_format: "str | dict[str, dict[str, str]] | None", max_shard_size: "str | int" = "10GB", is_modelopt_qlora: bool = False, extra_state_dict: "dict[str, torch.Tensor] | None" = None, diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 4ba4a1fb330..fe4d68fa7b8 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -92,7 +92,12 @@ from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload from modelopt.torch.utils import print_rank_0 from modelopt.torch.utils.dataset_utils import _disable_use_cache -from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan, is_fsdp2_model +from modelopt.torch.utils.distributed import ( + DistributionPlan, + broadcast_from, + distribution_plan, + is_fsdp2_model, +) from modelopt.torch.utils.perf import maybe_clear_cuda_cache try: @@ -1905,6 +1910,8 @@ def _write_base_config(model: nn.Module, export_dir: "Path | str") -> None: ``save_pretrained`` would have. The weight half of ``save_pretrained`` must NOT run: the weights are already on disk and the model may still be sharded. """ + # A config's save_pretrained writes only config.json: transformers reverts weight conversions + # in the model's, so the patch save_non_weight_artifacts takes is not needed here. model.config.save_pretrained(export_dir) if model.can_generate() and getattr(model, "generation_config", None) is not None: try: @@ -2010,8 +2017,10 @@ def export_hf_checkpoint( return is_fsdp2_sharded = plan.is_fsdp_dense - use_no_gather_writer = ( - is_fsdp2_sharded and os.environ.get("MODELOPT_FSDP2_EXPORT") == "nogather" + # Rank 0's choice: the two writers issue different collectives, so ranks that read their + # environments differently would wait in collectives the others never join. + use_no_gather_writer = is_fsdp2_sharded and broadcast_from( + "MODELOPT_FSDP2_EXPORT", lambda: os.environ.get("MODELOPT_FSDP2_EXPORT") == "nogather" ) # Not the global rank: a non-FSDP2 export writes the whole checkpoint from every process. rank = torch.distributed.get_rank() if is_fsdp2_sharded else 0 @@ -2073,18 +2082,16 @@ def export_hf_checkpoint( distributed_save_hf_checkpoint( model, export_dir, - maxbound=448, - kv_cache_format=(hf_quant_config or {}) - .get("quantization", {}) - .get("kv_cache_quant_algo"), + maxbound=448, # the KV cache scale is amax / 448 for FP8 and NVFP4, as above + kv_cache_format=_get_kv_cache_postprocess_config( + (hf_quant_config or {}).get("quantization", {}) + ), max_shard_size=max_shard_size, is_modelopt_qlora=_qlora, extra_state_dict=extra_state_dict, ) if rank == 0: _write_base_config(model, export_dir) - if save_modelopt_state and ModeloptStateManager.is_converted(model): - torch.save(modelopt_state(model), export_dir / _MODELOPT_STATE_SAVE_NAME) hf_quant_config = _revert_quant_config_names_best_effort(model, hf_quant_config) elif is_fsdp2_sharded: # FSDP2 multi-rank: stream each rank's owned units straight to its own shard files, so diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index dd0a8465133..aa1e89694b7 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -72,8 +72,8 @@ def _write(model, export_dir, **kwargs): ) -def _nogather_export(model, export_dir, **kwargs): - with mock.patch.dict(os.environ, {"MODELOPT_FSDP2_EXPORT": "nogather"}): +def _nogather_export(model, export_dir, knob="nogather", **kwargs): + with mock.patch.dict(os.environ, {"MODELOPT_FSDP2_EXPORT": knob}): export_hf_checkpoint( model, dtype=torch.bfloat16, export_dir=export_dir, max_shard_size="16KB", **kwargs ) @@ -157,6 +157,27 @@ def _fsdp2_shard(model, mesh, wrap_root=True): return model +# FP8 weights, with an FP8 KV cache in layer 0 and an NVFP4 one in layer 1: a mixed-precision KV +# cache, whose per-layer formats the export reads from the quantization config. +_FP8_MIXED_KV_CFG = copy.deepcopy(mtq.FP8_DEFAULT_CFG) +_FP8_MIXED_KV_CFG["quant_cfg"] += [ + { + "quantizer_name": "*layers.0.*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "axis": None}, + "enable": True, + }, + { + "quantizer_name": "*layers.1.*[kv]_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + "axis": None, + }, + "enable": True, + }, +] + + def _quantize_with_reference(model, quant_cfg): """Quantize ``model`` and return what the single-process export writes for a copy of it.""" torch.manual_seed(1) @@ -620,7 +641,9 @@ def test_distributed_save_writes_the_hub_checkpoint_of_a_sharded_moe( @pytest.mark.parametrize("wrap_root", [True, False], ids=["root-wrapped", "layers-only"]) @pytest.mark.parametrize( - "quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG], ids=["fp8", "nvfp4"] + "quant_cfg", + [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG, _FP8_MIXED_KV_CFG], + ids=["fp8", "nvfp4", "fp8-mixed-kv-cache"], ) def test_nogather_export_writes_what_the_single_process_export_writes( gloo_world_of_one, tmp_path, quant_cfg, wrap_root @@ -706,7 +729,12 @@ def _two_rank_job(rank, world_size, tmp_dir): # export of its own copy is the reference. model = _load(tmp_dir / "dense") want = _quantize_with_reference(model, mtq.FP8_DEFAULT_CFG) - _nogather_export(_fsdp2_shard(model, mesh), tmp_dir / "dense_export") + # Only rank 0's environment picks the no-gather writer: its choice holds on every rank. + _nogather_export( + _fsdp2_shard(model, mesh), + tmp_dir / "dense_export", + knob="nogather" if rank == 0 else "streaming", + ) # Each dense tensor ends up whole and plain on exactly one rank, whatever its layout, and the # owners are spread out. A zero-row tensor has no rows to move: its owner writes it as it is. From 154841df1620257e33927c01b51ebb34cbe6ac52 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:38:57 +0000 Subject: [PATCH 17/24] Keep the no-gather export's quant config names in step with its weights The no-gather writer keeps the in-memory tensor names on every rank when any rank cannot reverse the conversion mapping (a merge group split across ranks), but export_hf_checkpoint renamed the quant config to hub names regardless, so exclude_modules named modules the checkpoint did not hold and a loader would read an excluded layer as quantized. distributed_save_hf_checkpoint now returns whether the names were reverted, and the config is renamed only when they were. The experts' input-scale all-reduce runs on the scales' own device instead of one picked from the backend name, and the CPU tests now cover it: the one-rank Qwen2-MoE export quantizes expert inputs (its comment called the reduce CUDA-only), and a two-rank export whose ranks calibrated different input amaxes must write the largest on every expert. Also drops a comment that described no code, and the GPU test's description of a staging step this writer never had, with the assertion that nothing left that staging directory behind. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 35 +++----- modelopt/torch/export/unified_export_hf.py | 7 +- .../export/test_distributed_hf_export.py | 19 +++-- .../torch/export/test_distributed_writer.py | 79 +++++++++++++++---- 4 files changed, 89 insertions(+), 51 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index 2af31586495..ae87b419f2b 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -39,7 +39,7 @@ def distributed_save_hf_checkpoint( max_shard_size: "str | int" = "10GB", is_modelopt_qlora: bool = False, extra_state_dict: "dict[str, torch.Tensor] | None" = None, -) -> None: +) -> bool: """No-gather distributed HF export. Thin wrapper: hf_ptq exports under ``torch.inference_mode()``, so run the entire write with inference mode DISABLED @@ -49,9 +49,13 @@ def distributed_save_hf_checkpoint( ``inference_mode(False)`` re-ENABLES autograd, so a bare ``clone()``/``detach()`` on an inference tensor would try to set up grad and read its (missing) version counter -- ``no_grad`` avoids that while the clones still land as normal tensors because inference mode is off. + + Returns: + Whether the tensors carry their original checkpoint names. The reverse conversion is all or + nothing across ranks, so this is the same on every rank. """ with torch.inference_mode(False), torch.no_grad(): - _distributed_save_hf_checkpoint_impl( + return _distributed_save_hf_checkpoint_impl( model, export_dir, maxbound, @@ -644,7 +648,7 @@ def _distributed_save_hf_checkpoint_impl( max_shard_size: "str | int" = "10GB", is_modelopt_qlora: bool = False, extra_state_dict: "dict[str, torch.Tensor] | None" = None, -) -> None: +) -> bool: """Distributed HF safetensors export -- no rank-0 full-model host-RAM gather. Writes an already-processed, FSDP2-sharded model to ``export_dir`` as HF safetensors, entirely @@ -682,13 +686,6 @@ def _distributed_save_hf_checkpoint_impl( export_dir = Path(export_dir) rank, world = dist.get_rank(), dist.get_world_size() - # Where the expert input-scale all-reduce below runs: on the GPU under NCCL (distributed_run - # sets up "cpu:gloo,cuda:nccl"), on the CPU under a gloo-only group. - device = ( - torch.device("cpu") - if dist.get_backend() == "gloo" - else torch.device("cuda", torch.cuda.current_device()) - ) # This writer targets a FULLY FSDP2-sharded model -- no expert parallelism. FSDP2 # shards the fused 3-D expert weight on dim 0 (the expert axis), so fused experts are just @@ -772,20 +769,13 @@ def _is_expert_key(key: str) -> bool: return True return False - # Probe the DTensor placement of a sample expert + dense weight: directly shows what state the - # generation forwards left the sharded params in -- FSDP2 (Shard(0)/Replicate/gathered) or EP - # (expert dim-0 Shard, or a plain local shard for the a2a transport). - # Sync the shared per-module activation (input) scales across ranks (global max), so every - # expert of a module gets the same input_scale regardless of which rank owns it. + # expert of a module gets the same input_scale regardless of which rank owns it. Reduced on the + # scales' own device, as the shard transport moves every tensor: the GPU in a GPU run, the CPU + # under gloo. in_keys = sorted(k for k in sharded_sd if k.endswith("_input_scale") and _is_expert_key(k)) if in_keys: - stacked = torch.stack( - [ - sharded_sd[k].detach().to(device=device, dtype=torch.float32).reshape(()) - for k in in_keys - ] - ) + stacked = torch.stack([sharded_sd[k].detach().reshape(()).float() for k in in_keys]) dist.all_reduce(stacked, op=dist.ReduceOp.MAX) for j, k in enumerate(in_keys): sharded_sd[k] = stacked[j].clone() @@ -829,7 +819,7 @@ def _is_expert_key(key: str) -> bool: # transformers conversion ops are written against whole tensors and cannot express a placement. # The reverse must precede the shard layout, since it renames keys and splits tensors. _materialize_owned(local_sd, expert_key_set, rank, world) - _revert_whole_sd(model, local_sd) + reverted = _revert_whole_sd(model, local_sd) # Tensors the model does not own -- whatever the caller hands in via extra_state_dict, e.g. # weights read straight off the source checkpoint that the FSDP2 loader drops. They are plain # CPU tensors, already whole and identical on every rank, so ONE rank adds them to its @@ -858,3 +848,4 @@ def check_no_overlap() -> None: if extra_state_dict and rank == 0: local_sd.update(extra_state_dict) _write_even_shards(local_sd, export_dir, max_shard_size, rank, world) + return reverted diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index fe4d68fa7b8..ab5d7e2576b 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -2079,7 +2079,7 @@ def export_hf_checkpoint( pack_quantized_weights(model, _dtype, _qlora, keep_fused_experts=True) if getattr(model, "hf_quantizer", None) is not None: model.hf_quantizer = None - distributed_save_hf_checkpoint( + reverted = distributed_save_hf_checkpoint( model, export_dir, maxbound=448, # the KV cache scale is amax / 448 for FP8 and NVFP4, as above @@ -2092,7 +2092,10 @@ def export_hf_checkpoint( ) if rank == 0: _write_base_config(model, export_dir) - hf_quant_config = _revert_quant_config_names_best_effort(model, hf_quant_config) + # The weights keep their in-memory names when any rank could not revert them, and + # the quant config has to name the modules as the weights do. + if reverted: + hf_quant_config = _revert_quant_config_names_best_effort(model, hf_quant_config) elif is_fsdp2_sharded: # FSDP2 multi-rank: stream each rank's owned units straight to its own shard files, so # a rank buffers its own share of the model rather than the whole checkpoint, and the diff --git a/tests/gpu/torch/export/test_distributed_hf_export.py b/tests/gpu/torch/export/test_distributed_hf_export.py index fe88fe8fd95..61225f0392d 100644 --- a/tests/gpu/torch/export/test_distributed_hf_export.py +++ b/tests/gpu/torch/export/test_distributed_hf_export.py @@ -15,20 +15,20 @@ """End-to-end checks on the checkpoint written by the no-gather FSDP2 distributed export. With ``MODELOPT_FSDP2_EXPORT=nogather``, ``export_hf_checkpoint`` routes an FSDP2-sharded model under -``torch.distributed`` to ``distributed_save_hf_checkpoint`` -- the writer every expert-parallel model -takes: every rank writes its own DCP shards and the files are consolidated in parallel, with no -full-model host-RAM gather on rank 0. That writer reconstructs -the checkpoint from per-rank pieces, so the failure modes it can introduce are structural rather -than numerical -- a key silently dropped, a rank-local shard written where the full tensor belongs, -a sidecar or the weight index never emitted. Those all produce a checkpoint that *looks* fine (no -error, plausible file sizes) but is wrong or unloadable, so they are asserted here explicitly: +``torch.distributed`` to ``distributed_save_hf_checkpoint``: every rank writes its own final +safetensors shards in parallel and rank 0 unions the weight index, with no full-model host-RAM +gather on rank 0. That writer reconstructs the checkpoint from per-rank pieces, so the failure modes +it can introduce are structural rather than numerical -- a key silently dropped, a rank-local shard +written where the full tensor belongs, a sidecar or the weight index never emitted. Those all +produce a checkpoint that *looks* fine (no error, plausible file sizes) but is wrong or unloadable, +so they are asserted here explicitly: 1. every source parameter survives the round trip, and nothing extra is left behind; 2. each quantized weight carries the scales its format needs; 3. every tensor has its full, unsharded shape; 4. the non-weight files are all present -- the ones the exporter writes (config, generation config, quant config, weight index) and the source's own (tokenizer, vocab, chat template), which the - exporter copies over after the writer has removed its ``sharded/`` staging dir. + exporter copies over. The dense model ties its embeddings, so the expected keys also check that the writer drops the alias. What needs no GPU -- tensors handed in through ``extra_state_dict``, a linear left @@ -60,7 +60,7 @@ # Small enough to force the tiny model across several shards, so the multi-file layout -- and the # weight index that makes it loadable -- are exercised rather than collapsing to one model.safetensors. -# Not smaller: each extra shard is another DCP write + consolidation round trip on a toy model. +# Not smaller: each extra shard is one more file to write for a toy model. MAX_SHARD_SIZE = "512KB" # The default tiny Qwen3 is too degenerate to shard: head_dim = hidden_size / num_heads = 32/16 = 2, @@ -260,7 +260,6 @@ def test_fsdp2_distributed_export_is_complete(dist_workers, tmp_path, moe, quant # special-tokens map, ...) is in the export; only the weight files are rewritten. lost = sorted(sidecars - exported_files) assert not lost, f"sidecar file(s) missing from the export: {lost}" - assert not (export_dir / "sharded").exists(), "per-rank staging dir left behind" # A multi-file checkpoint is unloadable without the index, so require it whenever the writer # emitted sharded names, and require it to describe every tensor actually on disk. diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index aa1e89694b7..cae0841e9bf 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -146,6 +146,23 @@ def _llama_config(): ) +def _qwen2_moe_config(): + """A tiny Qwen2-MoE: four fused experts and a shared expert in one decoder layer.""" + return Qwen2MoeConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=32, + moe_intermediate_size=32, + shared_expert_intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + num_experts=4, + num_experts_per_tok=2, + max_position_embeddings=32, + ) + + def _fsdp2_shard(model, mesh, wrap_root=True): """Shard each decoder layer, and the root too unless ``wrap_root`` is off (as expert parallelism leaves it, sharding only the dense layers), and record that the model is FSDP2-sharded.""" @@ -670,6 +687,29 @@ def test_nogather_export_rejects_save_modelopt_state(gloo_world_of_one, tmp_path _nogather_export(model, tmp_path / "export", save_modelopt_state=True) +def test_nogather_export_names_modules_in_the_quant_config_as_the_weights_do( + gloo_world_of_one, tmp_path +): + """When the weights keep their in-memory names, as every rank's do once one rank cannot revert + its own, the quant config keeps them too: hub names there would exclude no written module.""" + model = _load(_save(_mixtral_config(), tmp_path / "src")) + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randint(0, 64, (2, 8)))) + failure = mock.patch.object( + quant_aware_conversion, + "revert_weight_conversion_quant_aware", + side_effect=RuntimeError("incomplete merge"), + ) + + with failure, pytest.warns(UserWarning, match="incomplete merge"): + _nogather_export(_fsdp2_shard(model, gloo_world_of_one), tmp_path / "export") + + modules = {key.rsplit(".", 1)[0] for key in _load_checkpoint(tmp_path / "export")} + config = json.loads((tmp_path / "export" / "hf_quant_config.json").read_text()) + excluded = config["quantization"]["exclude_modules"] + assert "model.layers.0.mlp.gate" in excluded + assert set(excluded) <= modules, sorted(set(excluded) - modules) + + @pytest.mark.parametrize( "quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG], ids=["fp8", "nvfp4"] ) @@ -677,23 +717,7 @@ def test_nogather_export_writes_moe_experts_quantized(gloo_world_of_one, tmp_pat """Qwen2-MoE's shared expert follows its experts in the same FSDP2 unit, and packing it opens a weight window on that unit. Unless the experts were folded inside a window of their own, which hands FSDP2 the packed weights, that window's reshard puts FSDP2's unquantized shards back.""" - config = Qwen2MoeConfig( - vocab_size=64, - hidden_size=32, - intermediate_size=32, - moe_intermediate_size=32, - shared_expert_intermediate_size=32, - num_hidden_layers=1, - num_attention_heads=4, - num_key_value_heads=2, - num_experts=4, - num_experts_per_tok=2, - max_position_embeddings=32, - ) - model = _load(_save(config, tmp_path / "src")) - # Weight-only experts: expert input scales would need the writer's CUDA-only all-reduce. - quant_cfg = copy.deepcopy(quant_cfg) - quant_cfg["quant_cfg"].append({"quantizer_name": "*experts*input_quantizer", "enable": False}) + model = _load(_save(_qwen2_moe_config(), tmp_path / "src")) want = _quantize_with_reference(model, quant_cfg) _nogather_export(_fsdp2_shard(model, gloo_world_of_one), tmp_path / "export") @@ -809,3 +833,24 @@ def test_two_ranks_write_one_checkpoint_and_reject_inconsistent_state(cpu_dist_w assert shards == [f"model-{i:05d}-of-{n:05d}.safetensors" for i in range(1, n + 1)] want = {**load_file(str(src / "model.safetensors")), **_mtp_tail()} _assert_byte_identical(_load_checkpoint(tmp_path / "moe_export"), want) + + +def _expert_input_scale_job(rank, world_size, tmp_dir): + mesh = init_device_mesh("cpu", (world_size,)) + model = _load(tmp_dir / "src") + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randint(0, 64, (2, 8)))) + # As if each rank calibrated the experts' shared input quantizer on its own data: each rank folds + # its own experts with its own scale, and only the writer's reduce makes them agree. + model.model.layers[0].mlp.experts.gate_up_proj_input_quantizer.amax = torch.tensor(1.0 + rank) + _nogather_export(_fsdp2_shard(model, mesh), tmp_dir / "export") + + +def test_two_ranks_write_every_expert_with_the_largest_input_scale(cpu_dist_workers, tmp_path): + """The experts' shared input scale is reduced across the ranks that fold them, to its max.""" + _save(_qwen2_moe_config(), tmp_path / "src") + + cpu_dist_workers.run(_expert_input_scale_job, tmp_dir=tmp_path) + + exported = _load_checkpoint(tmp_path / "export") + scales = [exported[f"model.layers.0.mlp.experts.{e}.gate_proj.input_scale"] for e in range(4)] + assert all(torch.allclose(s, torch.tensor(2.0 / 448)) for s in scales), scales From 730d7d2d04d1f80df62fd57079f4d65ef485879b Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 00:20:22 +0000 Subject: [PATCH 18/24] Copy only inference tensors before the no-gather export's transfer The no-gather writer copied every DTensor's local shard up front to drop the inference flag, though only an inference tensor carries it, and its comments called the copy free: the state dict's values are views of the model's own parameters, so each copy adds to the rank's memory instead of replacing a value. A DTensor is now copied only when its local shard is an inference tensor, as a plain tensor already was, so a caller outside inference mode (hf_ptq exports under no_grad) copies nothing. On two gloo ranks exporting a 223 MiB bf16 Llama, the writer's peak above each rank's 112 MiB shard drops from 278 and 246 MiB to 224 and 217 MiB. The docstrings and comments now state the bound the code has. Also drops an unused loop index. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 31 ++++++++++--------- .../torch/export/test_distributed_writer.py | 22 +++++++++++-- 2 files changed, 36 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index ae87b419f2b..63031a23a2d 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -42,7 +42,7 @@ def distributed_save_hf_checkpoint( ) -> bool: """No-gather distributed HF export. - Thin wrapper: hf_ptq exports under ``torch.inference_mode()``, so run the entire write with inference mode DISABLED + Thin wrapper: a caller may be in ``torch.inference_mode()``, so run the entire write with inference mode DISABLED via a nested ``inference_mode(False)`` context -- the ``.detach()`` in ``get_model_state_dict`` and the version- counted ops in the write path both reject inference tensors. The context form reliably restores normal mode for this scope (a decorator did not, under the caller's already-active inference_mode). ``no_grad`` too: @@ -667,8 +667,11 @@ def _distributed_save_hf_checkpoint_impl( 4. Every rank bin-packs its own owned set into even shards, numbered into a global range, and writes them directly with ``safetensors.save_file``. Rank 0 unions the weight index. - Peak memory per rank is its owned set, not the model. Placements the point-to-point transport - does not cover are rejected rather than routed through an untested slower path. + No rank holds the model. On top of its own shard of it, which the model keeps, a rank holds what + it writes -- its experts, split into per-expert copies, and the dense tensors it owns, assembled + as :func:`_materialize_owned` describes -- and, from a caller in inference mode, a copy of each + inference tensor in its shard. Placements the point-to-point transport does not cover are + rejected rather than routed through an untested slower path. """ import torch.distributed as dist from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict @@ -719,15 +722,17 @@ def _experts_of(c): write_experts = dp_idx == 0 sharded_sd = get_model_state_dict(model, options=StateDictOptions(full_state_dict=False)) - # Re-materialize every value as a NORMAL tensor UP FRONT. hf_ptq exports under torch.inference_mode(), - # so these are inference tensors; the inference flag is STICKY -- clone/reshape/detach/slice all + # Re-materialize every inference tensor as a NORMAL one UP FRONT. A caller in inference mode hands + # over inference tensors; the inference flag is STICKY -- clone/reshape/detach/slice all # propagate it -- and the expert split, scale reduce and write below all do version-counted ops # that reject inference tensors. The only way to drop the flag is to copy the data into a freshly # ALLOCATED tensor (here, under the wrapper's inference_mode(False)). For a DTensor, copy the LOCAL - # shard and rewrap so mesh/placements survive. Bounded to this rank's shard. + # shard and rewrap so mesh/placements survive. Anything else is used as it is. def _to_normal(v): if isinstance(v, DTensor): _loc = v.to_local() + if not _loc.is_inference(): + return v _new = torch.empty_like(_loc) _new.copy_(_loc) # Pass the ORIGINAL global shape/stride: from_local() otherwise re-infers them as @@ -750,12 +755,10 @@ def _to_normal(v): return _new return v - # Re-materialize IN PLACE, not as `{k: _to_normal(v) for ...}`. The comprehension keeps the ENTIRE - # original dict alive while building a full second copy -> a transient 2x of the whole state dict. - # That is fine at 30B (24GB -> 48GB) but OOMs at large MoE: at 235B the write-point state dict is - # ~157GB/rank (fused expert weights are sharded, but the per-expert NVFP4 scale tensors are - # replicated on every rank), and doubling it hits ~314GB > GPU. Overwriting each key frees the old - # value immediately, so peak stays ~1x. (Import of ep-dp's in-place/streaming write handling.) + # A copy adds to this rank's memory rather than replacing anything: the values are views of the + # model's own parameters and buffers, which keep their data. Copying every value would hold the + # rank's whole shard twice; copying only inference tensors costs a caller outside inference mode + # (hf_ptq exports under no_grad) nothing. for _k in list(sharded_sd.keys()): sharded_sd[_k] = _to_normal(sharded_sd[_k]) prefixes = _fused_experts_prefixes(sharded_sd) @@ -783,7 +786,7 @@ def _is_expert_key(key: str) -> bool: # Skipped on dp-replica ranks under DP x EP (their experts are written by dp group 0). local_sd: dict = {} if write_experts: - for pi, prefix in enumerate(prefixes): + for prefix in prefixes: local_sd.update( _split_local_fused_module( prefix, sharded_sd, ep_rank=ep_rank, total_experts=total_experts @@ -792,7 +795,7 @@ def _is_expert_key(key: str) -> bool: local_sd = postprocess_state_dict( local_sd, maxbound, kv_cache_format, is_modelopt_qlora, tied_map=tied_map ) - # In place (see above): avoid a transient 2x copy of the per-expert split. + # In place: the split's tensors are its own copies, so one made contiguous replaces the old. for _k in list(local_sd.keys()): local_sd[_k] = local_sd[_k].detach().contiguous() # At this point local_sd holds ONLY this rank's (already split, per-expert) experts -- capture the diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index cae0841e9bf..cdfc5a7c9c6 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -47,7 +47,7 @@ from transformers.core_model_loading import Chunk, WeightConverter import modelopt.torch.quantization as mtq -from modelopt.torch.export import quant_aware_conversion +from modelopt.torch.export import distributed_writer, quant_aware_conversion from modelopt.torch.export.distributed_writer import ( _even_bins, _is_full_world_shard0, @@ -629,10 +629,22 @@ def test_distributed_save_writes_the_hub_checkpoint_of_a_sharded_moe( A caller in inference mode hands the writer inference tensors, which its writes reject. Extra tensors already carry hub names, so the reverse conversion must not rename them: the first is named like an in-memory Mixtral expert weight, which the reverse would turn into - ``block_sparse_moe.experts.0.w1``. + ``block_sparse_moe.experts.0.w1``. Only an inference tensor is copied before the transfer: the + dense shards reach it as the model holds them, since a copy would only add to the rank's memory. """ src = _save(_mixtral_config(tie_word_embeddings=case == "tied"), tmp_path / "src") model = _fsdp2_shard(_load(src), gloo_world_of_one) + shards = {p.to_local().untyped_storage().data_ptr() for p in model.parameters()} + moved = [] # per dense DTensor reaching the transfer: whether it is the model's own shard + + def materialize_owned(local_sd, *args): + moved.extend( + v.to_local().untyped_storage().data_ptr() in shards + for v in local_sd.values() + if isinstance(v, DTensor) + ) + return _materialize_owned(local_sd, *args) + want = load_file(str(src / "model.safetensors")) extra = None if case == "extra-state-dict": @@ -646,7 +658,10 @@ def test_distributed_save_writes_the_hub_checkpoint_of_a_sharded_moe( elif case == "inference-mode": want["model.norm.weight_scale"] = torch.tensor(0.5) - with torch.inference_mode(case == "inference-mode"): + with ( + torch.inference_mode(case == "inference-mode"), + mock.patch.object(distributed_writer, "_materialize_owned", materialize_owned), + ): if case == "inference-mode": # Made in inference mode, like the scales an inference-mode pack registers. model.model.norm.register_buffer("weight_scale", torch.full((), 0.5)) @@ -654,6 +669,7 @@ def test_distributed_save_writes_the_hub_checkpoint_of_a_sharded_moe( assert len(list((tmp_path / "export").glob("*.safetensors"))) > 1 _assert_byte_identical(_load_checkpoint(tmp_path / "export"), want) + assert moved and all(moved), moved @pytest.mark.parametrize("wrap_root", [True, False], ids=["root-wrapped", "layers-only"]) From e26eb9f81a8a56985f4289b8e9c2811ab4139b67 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 01:18:06 +0000 Subject: [PATCH 19/24] Drop get_max_batch_size from the changelog's plan-required entry Since [8/21], get_max_batch_size measures a model with no recorded plan per rank, as it did before, so only export_hf_checkpoint raises for a sharded or multi-rank model that carries no plan. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 9c475e2d747..1e01da0214f 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -42,7 +42,7 @@ Changelog - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``. - ``modelopt.torch.utils.distributed.fsdp2_wrap`` moved to ``modelopt.torch.utils.plugins.model_load_utils``, with the other loaders. The old import path still works, but calling it now needs the ``hf`` extra. - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata. -- ``export_hf_checkpoint`` and ``modelopt.torch.utils.dataset_utils.get_max_batch_size`` now read how a model is distributed from the plan ``load_pretrained`` records on it, and raise for an FSDP2-sharded model, or one in a multi-rank run, that carries none. If you distribute a model yourself, for example with ``fsdp2_wrap``, which records no plan, record one with ``modelopt.torch.utils.distributed.set_distribution_plan`` before calibrating or exporting it. +- ``export_hf_checkpoint`` now reads how a model is distributed from the plan ``load_pretrained`` records on it, and raises for an FSDP2-sharded model, or one in a multi-rank run, that carries none. If you distribute a model yourself, for example with ``fsdp2_wrap``, which records no plan, record one with ``modelopt.torch.utils.distributed.set_distribution_plan`` before calibrating or exporting it. **Deprecations** From db1fe7d1d109796db788373c6c871ce1e76e7ec1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:14:16 +0000 Subject: [PATCH 20/24] Save the ModelOpt state from the no-gather export, as the streaming export does The no-gather FSDP2 writer refused save_modelopt_state=True, so opting into it broke an export the default FSDP2 path completes. That path saves the state from rank 0 with torch.save, not through save_pretrained, and every FSDP2 rank holds the same state: the no-gather path now does the same. Signed-off-by: Shengliang Xu --- modelopt/torch/export/unified_export_hf.py | 15 +++------------ .../unit/torch/export/test_distributed_writer.py | 14 +++++++++----- 2 files changed, 12 insertions(+), 17 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index ab5d7e2576b..9c5c56485c4 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -2029,18 +2029,6 @@ def export_hf_checkpoint( # buffer instead of the whole quantized state dict. _offloaded = has_accelerate_offload(model) - # An offloaded model takes the streaming path below; otherwise a distributed FSDP2 model takes - # the no-gather path, which writes weights and config itself instead of going through - # save_pretrained. extra_state_dict IS honoured there (one rank adds those tensors to its write - # set); save_modelopt_state is not, so reject it explicitly rather than dropping it silently - # (the streaming path declines it the same way). - if use_no_gather_writer and not _offloaded: - if save_modelopt_state: - raise NotImplementedError( - "save_modelopt_state=True is not supported by the no-gather FSDP2 distributed " - "export: it writes the checkpoint directly rather than through " - "model.save_pretrained(). Save the ModelOpt state separately with mto.save()." - ) try: if _offloaded: # Imported here rather than at module scope: the streaming exporter imports the @@ -2092,6 +2080,9 @@ def export_hf_checkpoint( ) if rank == 0: _write_base_config(model, export_dir) + # As the streaming path below: every FSDP2 rank holds the same ModelOpt state. + if save_modelopt_state and ModeloptStateManager.is_converted(model): + torch.save(modelopt_state(model), export_dir / _MODELOPT_STATE_SAVE_NAME) # The weights keep their in-memory names when any rank could not revert them, and # the quant config has to name the modules as the weights do. if reverted: diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index cdfc5a7c9c6..082ae7e0641 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -695,12 +695,16 @@ def test_nogather_export_writes_what_the_single_process_export_writes( assert (tmp_path / "export" / sidecar.name).read_text() == sidecar.read_text() -def test_nogather_export_rejects_save_modelopt_state(gloo_world_of_one, tmp_path): - """The path never runs ``save_pretrained``, which writes that state: refuse, not drop it.""" - model = _fsdp2_shard(_load(_save(_llama_config(), tmp_path / "src")), gloo_world_of_one) +def test_nogather_export_saves_modelopt_state(gloo_world_of_one, tmp_path): + """As the default FSDP2 export does, it writes the ModelOpt state next to the weights.""" + model = _load(_save(_llama_config(), tmp_path / "src")) + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randint(0, 63, (2, 8)))) - with pytest.raises(NotImplementedError, match="save_modelopt_state"): - _nogather_export(model, tmp_path / "export", save_modelopt_state=True) + _nogather_export( + _fsdp2_shard(model, gloo_world_of_one), tmp_path / "export", save_modelopt_state=True + ) + + assert (tmp_path / "export" / "modelopt_state.pth").is_file() def test_nogather_export_names_modules_in_the_quant_config_as_the_weights_do( From 236f95cc20aab80a1eddffa0ac8ed4433a7e9d78 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:19:29 +0000 Subject: [PATCH 21/24] Agree on the expert input-scale keys before the writer reduces them The writer max-reduces the experts' shared input scales over keys each rank derives from its own state dict. A rank holding none skipped the reduce and left the others waiting in it, and lists of different lengths mismatched its size, with nothing naming the cause. The ranks now agree on the keys first, through the labeled all_reduce the writer's other checks use, and a difference fails on every rank, naming the keys. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 14 ++++++++-- .../torch/export/test_distributed_writer.py | 27 ++++++++++++++++++- 2 files changed, 38 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index 63031a23a2d..039f4922584 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -677,6 +677,8 @@ def _distributed_save_hf_checkpoint_impl( from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict from torch.distributed.tensor import DTensor + from modelopt.torch.utils.distributed import all_reduce + from .model_utils import TiedWeightMap from .quant_utils import postprocess_state_dict @@ -777,6 +779,16 @@ def _is_expert_key(key: str) -> bool: # scales' own device, as the shard transport moves every tensor: the GPU in a GPU run, the CPU # under gloo. in_keys = sorted(k for k in sharded_sd if k.endswith("_input_scale") and _is_expert_key(k)) + # Every rank must reduce the same keys: a rank holding none would skip the reduce and leave the + # others waiting in it, and lists of different lengths would mismatch its size, neither saying + # why. So the ranks agree on the keys first, and a difference fails on every rank, naming it. + key_lists = all_reduce("expert input-scale keys", lambda: [in_keys], _in_rank_order) + if any(keys != in_keys for keys in key_lists): + differ = sorted(set().union(*key_lists) - set(key_lists[0]).intersection(*key_lists)) + raise RuntimeError( + "The ranks hold different expert input scales, so they cannot reduce them to their " + f"max; first few held by only some ranks: {differ[:5]}" + ) if in_keys: stacked = torch.stack([sharded_sd[k].detach().reshape(()).float() for k in in_keys]) dist.all_reduce(stacked, op=dist.ReduceOp.MAX) @@ -832,8 +844,6 @@ def _is_expert_key(key: str) -> bool: # to agree on the split. They go in AFTER the reverse conversion because they carry their # original hub names and were never converted -- reverting them would rename them wrongly. This # mirrors the gather path, which merges extra_state_dict into the already-reverted dict. - from modelopt.torch.utils.distributed import all_reduce - def check_no_overlap() -> None: if rank != 0 or not extra_state_dict: return diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 082ae7e0641..a231c67e796 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -864,9 +864,34 @@ def _expert_input_scale_job(rank, world_size, tmp_dir): model.model.layers[0].mlp.experts.gate_up_proj_input_quantizer.amax = torch.tensor(1.0 + rank) _nogather_export(_fsdp2_shard(model, mesh), tmp_dir / "export") + # A scale one rank holds and the other lacks fails on both, named, rather than leave the ranks + # hanging or mismatched in the reduce. + from torch.distributed.checkpoint import state_dict + + real = state_dict.get_model_state_dict + + def without_a_scale(*args, **kwargs): + sd = real(*args, **kwargs) + return {k: v for k, v in sd.items() if rank == 0 or "gate_up_proj_input_scale" not in k} + + model = _load(tmp_dir / "src") + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randint(0, 64, (2, 8)))) + try: + with ( + bounded_collectives(), + mock.patch.object(state_dict, "get_model_state_dict", without_a_scale), + ): + _nogather_export(_fsdp2_shard(model, mesh), tmp_dir / "diverging") + except RuntimeError as error: + raised = str(error) + else: + raised = "no error" + assert "different expert input scales" in raised, raised + def test_two_ranks_write_every_expert_with_the_largest_input_scale(cpu_dist_workers, tmp_path): - """The experts' shared input scale is reduced across the ranks that fold them, to its max.""" + """The experts' shared input scale is reduced across the ranks that fold them, to its max, and + ranks that hold different scales fail alike, naming them.""" _save(_qwen2_moe_config(), tmp_path / "src") cpu_dist_workers.run(_expert_input_scale_job, tmp_dir=tmp_path) From 727894a438ae68f56eabcfdd9faf430a2adf56a0 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:20:24 +0000 Subject: [PATCH 22/24] Keep the blank line ruff format wants before the carried-tensors comment Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index 039f4922584..4f07390d22e 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -835,6 +835,7 @@ def _is_expert_key(key: str) -> bool: # The reverse must precede the shard layout, since it renames keys and splits tensors. _materialize_owned(local_sd, expert_key_set, rank, world) reverted = _revert_whole_sd(model, local_sd) + # Tensors the model does not own -- whatever the caller hands in via extra_state_dict, e.g. # weights read straight off the source checkpoint that the FSDP2 loader drops. They are plain # CPU tensors, already whole and identical on every rank, so ONE rank adds them to its From 1961b455e1e8d3f7e2d13215706afaaa79fd9580 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:45:45 +0000 Subject: [PATCH 23/24] Keep each merge group on one rank in the no-gather export The writer pieces now place a merge group with one owner when they are told the model's groups, but the export did not tell them: a layer's q, k and v could land on different ranks, the merge then failed, and the whole checkpoint kept its in-memory names. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 2 +- .../torch/export/test_distributed_writer.py | 34 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index 4f07390d22e..2289794d23d 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -833,7 +833,7 @@ def _is_expert_key(key: str) -> bool: # exactly one rank (never gathered to a single rank), which is what lets the reverse run: the # transformers conversion ops are written against whole tensors and cannot express a placement. # The reverse must precede the shard layout, since it renames keys and splits tensors. - _materialize_owned(local_sd, expert_key_set, rank, world) + _materialize_owned(local_sd, expert_key_set, rank, world, _merge_groups(model, local_sd)) reverted = _revert_whole_sd(model, local_sd) # Tensors the model does not own -- whatever the caller hands in via extra_state_dict, e.g. diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index a231c67e796..1dff4e333f1 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -855,6 +855,40 @@ def test_two_ranks_write_one_checkpoint_and_reject_inconsistent_state(cpu_dist_w _assert_byte_identical(_load_checkpoint(tmp_path / "moe_export"), want) +def _merge_export_job(rank, world_size, tmp_dir): + model = _load(tmp_dir / "dense") + # A checkpoint that fuses each layer's q, k and v, which the model holds as three projections: + # the reverse merges them back, which needs each layer's three on one rank. + model._weight_conversions = [ + WeightConverter( + source_patterns="self_attn.qkv_proj.weight", + target_patterns=[f"self_attn.{part}_proj.weight" for part in "qkv"], + operations=[Chunk(dim=0)], + ) + ] + with bounded_collectives(): + _write(_fsdp2_shard(model, init_device_mesh("cpu", (world_size,))), tmp_dir / "merged") + + +def test_distributed_save_merges_a_group_its_ranks_hold_apart(cpu_dist_workers, tmp_path): + """Each layer's q, k and v go to one owner, so the written checkpoint fuses them as the gather + export does, instead of keeping every in-memory name because one layer's split over the ranks.""" + _save(_llama_config(), tmp_path / "dense") + cpu_dist_workers.run(_merge_export_job, tmp_dir=tmp_path) + + got = _load_checkpoint(tmp_path / "merged") + src = _load(tmp_path / "dense").state_dict() + layers = sorted({k.rsplit(".self_attn.", 1)[0] for k in src if ".self_attn.q_proj." in k}) + assert layers + for layer in layers: + attn = f"{layer}.self_attn" + assert f"{attn}.q_proj.weight" not in got + want = torch.cat([src[f"{attn}.{part}_proj.weight"] for part in "qkv"]) + assert torch.equal( + got[f"{attn}.qkv_proj.weight"], want.to(got[f"{attn}.qkv_proj.weight"].dtype) + ) + + def _expert_input_scale_job(rank, world_size, tmp_dir): mesh = init_device_mesh("cpu", (world_size,)) model = _load(tmp_dir / "src") From c3114d646ffe2f0d066c0af1162ed4ba4973a5bc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:52:44 +0000 Subject: [PATCH 24/24] Keep the export handlers imported for the keep-fused window test Dropping the stub window test took the module import with it, but the keep-fused fold's window-per-unit test still drives the handlers through it. Signed-off-by: Shengliang Xu --- tests/unit/torch/export/test_fused_experts_keep_fused.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index cb0a085e777..4c09fbb60b4 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -43,6 +43,7 @@ pytest.importorskip("safetensors") from safetensors.torch import load_file, save_file +import modelopt.torch.export.hf_export_handlers as handlers import modelopt.torch.quantization as mtq from modelopt.torch.export.moe_utils import ( _dtensor_dim0_offset,