From 4f33f202c5ab211937a21001c34d30d45f6ee99d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 21:37:06 +0000 Subject: [PATCH 01/42] Carry the per-rank expert split the no-gather writer calls _split_local_fused_module cuts one rank's shard of a fused-experts module into per-expert keys numbered by the experts' global indices, with no gather. It left the keep-fused fold PR when that PR was split, to land with the writer that calls it, so it comes back here, unchanged, with the two tests that went with it: the global numbering under each expert- parallel layout, and, on two ranks, each rank's own split checked against its experts' part of the per-expert export, with per-expert scales handed in as one plain tensor over every expert too. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 97 ++++++++++++++++++- .../export/test_fused_experts_keep_fused.py | 59 +++++++++-- 2 files changed, 146 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 07e0ddd4129..50b8d909f3b 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -669,9 +669,10 @@ def split_fused_experts_state_dict(state_dict: dict) -> dict: checkpoints. Operates on the (already folded/quantized) tensors only -- pure slicing, no quantization. Non-fused-experts keys pass through unchanged. Every emitted tensor is a copy and the input keeps its fused tensors, so the experts are held twice until it returns: pass one - module's tensors at a time, not a whole model's. Each module's experts are numbered from 0 by - their row in the tensors given, which must be plain: whole ones get their global numbers, and a - caller splitting a rank's local shard (``to_local()``) shifts each by the shard's global offset. + module's tensors at a time, as :func:`_split_local_fused_module` does, not a whole model's. + Each module's experts are numbered from 0 by their row in the tensors given, which must be + plain: whole ones get their global numbers, and a caller splitting a rank's local shard + (``to_local()``) shifts each by the shard's global offset, as that function does. Per fused-experts prefix ``P`` (``...experts``), inputs are (gated shown; ungated drops the gate split): ``P.gate_up_proj`` ``(E, 2*I, H)`` + ``P.gate_up_proj_weight_scale`` + ``P.gate_up_proj_input_scale`` @@ -801,3 +802,93 @@ def _dtensor_dim0_offset(dt) -> int: nshards = mesh.size(mdim) return mesh.get_local_rank(mdim) * ((n + nshards - 1) // nshards) return 0 + + +def _split_local_fused_module( + prefix: str, sharded_sd: dict, ep_rank: int = 0, total_experts: int = 0 +) -> dict: + """Split this rank's local shard of a fused-experts module into per-expert keys. + + Split THIS rank's *local* shard of one fused-experts module into per-expert keys with GLOBAL + expert indices (plain, local tensors -- no all-gather). + + FSDP2/EP shard the fused 3-D expert weight on dim 0 (the expert axis), so every expert this rank + owns is whole and local. We reuse :func:`split_fused_experts_state_dict` on the ``to_local()`` + shard (which numbers experts ``0..E_local-1``) and then shift each expert index by the rank's + global dim-0 offset, so the resulting per-expert keys match the single-process export and stay + on this rank for a direct distributed write (no consolidation). + + The fused expert tensor can arrive in two EP layouts, distinguished by its dim-0 size vs + ``total_experts`` (the config expert count): + * **per-EP-group** (size < total): transformers pre-sliced the param to this rank's EP group + (FSDP x EP -> a dp-sharded DTensor of the group; or a plain ep-shard). ``_dtensor_dim0_offset`` + only gives the WITHIN-group offset, so we add the EP block base ``ep_rank * group_size``. + * **all-experts** (size == total): an EP DTensor sharded across the ep mesh spanning all experts + (classic DP x EP). ``_dtensor_dim0_offset`` already returns the TRUE global offset -> no base. + ``total_experts == 0``, a count the caller could not read off the model's config, takes the + tensor for a per-EP-group slice. + """ + gated = f"{prefix}.gate_up_proj" in sharded_sd + first = "gate_up_proj" if gated else "up_proj" + fw = sharded_sd[f"{prefix}.{first}"] + # dp_offset: this rank's dim-0 start WITHIN its DTensor; n_local: experts it owns. + dp_offset = _dtensor_dim0_offset(fw) if isinstance(fw, DTensor) else 0 + n_local = fw.to_local().shape[0] if isinstance(fw, DTensor) else fw.shape[0] + group_size = int(fw.shape[0]) + # Add the EP block base only when the tensor is a per-EP-group slice (size < total_experts) or + # total is unknown (0). When it spans all experts, dp_offset is already global. + add_ep_base = ( + bool(ep_rank) and group_size > 0 and (total_experts == 0 or group_size < total_experts) + ) + offset = dp_offset + (ep_rank * group_size if add_ep_base else 0) + + def _loc(key): + v = sharded_sd.get(key) + if v is None: + return None + return v.to_local() if isinstance(v, DTensor) else v + + def _loc_scale(key): + # Weight scales are dim-0 (per-expert) tensors. A DTensor scale shards with the weight -> + # to_local() aligns. The keep-fused fold emits DTensor scales for a sharded weight, but a + # scale handed in as a PLAIN tensor over every expert of the group is accepted too: slice + # it by the weight's within-group dp_offset so each scale pairs with its own dp-local expert. + v = sharded_sd.get(key) + if v is None: + return None + if isinstance(v, DTensor): + return v.to_local() + if v.dim() >= 1 and v.shape[0] > n_local and dp_offset + n_local <= v.shape[0]: + return v[dp_offset : dp_offset + n_local] + return v + + # Build a fused sub-dict of this rank's local experts (plain), then split + reindex. + local_fused: dict = { + f"{prefix}.{first}": _loc(f"{prefix}.{first}"), + f"{prefix}.down_proj": _loc(f"{prefix}.down_proj"), + } + for proj in (first, "down_proj"): + ws = _loc_scale(f"{prefix}.{proj}_weight_scale") + if ws is not None: + local_fused[f"{prefix}.{proj}_weight_scale"] = ws + # NVFP4 second-level per-tensor weight scale: also a per-expert (dim-0) tensor, so the same + # dp-offset slicing pairs each scalar with its dp-local expert. None for FP8. + ws2 = _loc_scale(f"{prefix}.{proj}_weight_scale_2") + if ws2 is not None: + local_fused[f"{prefix}.{proj}_weight_scale_2"] = ws2 + ins = sharded_sd.get(f"{prefix}.{proj}_input_scale") # shared scalar (replicated) + if ins is not None: + local_fused[f"{prefix}.{proj}_input_scale"] = ins + + split_local = split_fused_experts_state_dict( + local_fused + ) # keys {prefix}.{i}.* (i in 0..E_local-1) + if offset == 0: + return split_local + + plen = len(prefix) + 1 + out: dict = {} + for key, val in split_local.items(): + i_str, tail = key[plen:].split(".", 1) + out[f"{prefix}.{int(i_str) + offset}.{tail}"] = val + return out 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 0a516db55aa..9621c8f29fd 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -17,9 +17,10 @@ The no-gather writer never splits a fused 3-D expert weight while the model is sharded: ``_export_fused_experts_keep_fused`` folds every expert's quantizer into the fused weight in place, -and ``split_fused_experts_state_dict`` cuts the result into per-expert keys later. That pipeline -must write exactly what the single-process ``_export_fused_experts`` writes, so most tests compare -the two byte for byte on a tiny synthetic fused-experts module quantized with FP8 and NVFP4. +and ``split_fused_experts_state_dict`` / ``_split_local_fused_module`` cut the result into +per-expert keys later. That pipeline must write exactly what the single-process +``_export_fused_experts`` writes, so most tests compare the two byte for byte on a tiny synthetic +fused-experts module quantized with FP8 and NVFP4. Sharded cases run on a gloo process group: one rank in process, and two spawned ranks where a non-zero expert offset matters. """ @@ -50,6 +51,7 @@ _export_fused_experts, _export_fused_experts_keep_fused, _fused_experts_prefixes, + _split_local_fused_module, split_fused_experts_state_dict, ) from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry @@ -566,6 +568,32 @@ def test_fused_experts_prefixes_lists_each_fused_module_once_in_order(): ] +@pytest.mark.parametrize( + ("ep_rank", "total_experts", "first_index"), + [ + (0, 0, 0), + # This EP group's slice of a larger expert set: shift by the group base. + (1, 2 * LOCAL_EXPERTS, LOCAL_EXPERTS), + (3, 4 * LOCAL_EXPERTS, 3 * LOCAL_EXPERTS), + # The tensor already spans every expert, so its indices are global. + (1, LOCAL_EXPERTS, 0), + # Total unknown: the legacy per-group shift. + (1, 0, LOCAL_EXPERTS), + ], +) +def test_split_local_numbers_experts_globally(ep_rank, total_experts, first_index): + sd = _fused_state_dict(gated=True) + unshifted = split_fused_experts_state_dict(dict(sd)) + + out = _split_local_fused_module(PREFIX, sd, ep_rank=ep_rank, total_experts=total_experts) + + want = {} + for key, value in unshifted.items(): + expert, tail = key[len(PREFIX) + 1 :].split(".", 1) + want[f"{PREFIX}.{int(expert) + first_index}.{tail}"] = value + _assert_byte_identical(out, want) + + # ------------------------------------------------------------------ sharded (DTensor) experts @@ -631,6 +659,22 @@ def _two_rank_fold_job(rank, world_size): # Gathered whole, the ranks' folds split into exactly what the per-expert export writes. whole = {key: _whole(value) for key, value in fused.items()} _assert_byte_identical(split_fused_experts_state_dict(whole), want) + mine = { + k: v + for k, v in want.items() + if offset <= int(k[len(PREFIX) + 1 :].split(".")[0]) < offset + n_local + } + assert mine, "this rank owns no experts" + + _assert_byte_identical(_split_local_fused_module(PREFIX, fused), mine) + + # A per-expert scale held as a plain tensor over ALL experts (as keep-fused emits it for an + # unsharded weight) must be cut to this rank's experts before pairing. + plain_scales = { + k: plain[k] if k.endswith(("_weight_scale", "_weight_scale_2")) else v + for k, v in fused.items() + } + _assert_byte_identical(_split_local_fused_module(PREFIX, plain_scales), mine) # One expert over two ranks leaves rank 1 none: both ranks refuse, rather than rank 1 alone # failing while rank 0 goes on to write. @@ -686,10 +730,11 @@ def sharded_keeping_quantizers(kept): def test_two_ranks_fold_their_own_experts_under_global_indices(cpu_dist_workers): """Rank 1 holds experts 2-3: each must be folded with ITS quantizer at its global offset, and the fused tensors stay sharded on the expert axis at their global shape, an uneven shard included; - gathered, they split into the per-expert export. Its quantizers are every expert's (FSDP2) or, - under expert parallelism, its own two, and then the ranks decide on each projection's - quantization together. A layout that leaves a rank no expert, or quantizers that are neither, - is refused on every rank.""" + gathered, they split into the per-expert export, and each rank's own split, numbered globally, + into its experts' part of it. Its quantizers are every expert's (FSDP2) or, under expert + parallelism, its own two, and then the ranks decide on each projection's quantization together. + A layout that leaves a rank no expert, or quantizers that are neither, is refused on every + rank.""" cpu_dist_workers.run(_two_rank_fold_job) From afe23a9bb7c6aa69dad8f06bacfe492df97a21f8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 11 Oct 2026 11:10:42 +0000 Subject: [PATCH 02/42] 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 | 264 ++++++++++- .../torch/export/test_unified_export_hf.py | 21 + 7 files changed, 1074 insertions(+), 12 deletions(-) create mode 100644 tests/gpu/torch/export/test_distributed_hf_export.py diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f6cfb894081..8462007ec14 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -152,6 +152,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 b8aba4de503..08cb43e6d64 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -654,6 +654,8 @@ Qwen checkpoints with `model_type` `qwen3_5` or `qwen3_5_moe` support this path > *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 0f7e63436ea..913ec45e309 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 @@ -47,6 +47,39 @@ # A checkpoint's files: its weights, sharded as save_pretrained names them or in one file, and index. _CHECKPOINT_FILES = re.compile(r"model(-\d{5}-of-\d{5})?\.safetensors(\.index\.json)?") +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. @@ -633,3 +666,211 @@ def commit(names: list) -> bool: f"wrote to {export_dir}: {left[:5]}" ) raise + + +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 7c57c3838f5..97b6f275e11 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 @@ -977,6 +978,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. @@ -992,7 +994,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( @@ -1000,6 +1004,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, ) # For the no-gather export, which is to set keep_fused_experts on this context and let a sharded @@ -1058,13 +1063,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) @@ -1678,6 +1692,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, @@ -1914,6 +1942,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, @@ -2020,6 +2068,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 @@ -2027,6 +2078,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 @@ -2049,6 +2112,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 @@ -2138,8 +2233,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 8dd8ec6413c..8a5b03cfd80 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -13,24 +13,32 @@ # 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.distributed.tensor.placement_types import _StridedShard from torch.multiprocessing.reductions import StorageWeakRef @@ -38,9 +46,10 @@ pytest.importorskip("safetensors") pytest.importorskip("transformers") from safetensors.torch import load_file, save_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.distributed_writer import ( _even_bins, _is_full_world_shard0, @@ -49,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 + ) _WRITER = "modelopt.torch.export.distributed_writer" @@ -81,6 +118,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 = { @@ -111,6 +158,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 @@ -752,3 +821,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 1b8d55373fa..f4217ce380d 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -1095,3 +1095,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 3de204100743da9f16d0e26deda49549695c0e2a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 14 Oct 2026 20:30:58 +0000 Subject: [PATCH 03/42] 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 97b6f275e11..328d6b20e4b 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -96,7 +96,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: @@ -2010,14 +2010,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 {})} @@ -2063,11 +2059,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 8a5b03cfd80..828575a8ebd 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(): @@ -160,11 +161,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 3f2ce928953947ad28d1b7303b8926bf424013a1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 00:49:02 +0000 Subject: [PATCH 04/42] 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 95f0b1a333a9de26842a30e4d6aa9d7dc9b93cb8 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 00:56:54 +0000 Subject: [PATCH 05/42] 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 8462007ec14..ad30f4c0cc3 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -225,6 +225,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 74131830cd4f7944323efa146849346f757b5dc1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 17:31:21 +0000 Subject: [PATCH 06/42] 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 ad30f4c0cc3..c793c369dec 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -24,6 +24,7 @@ Changelog - Speed up NVFP4 QAT and QAD of Hugging Face fused MoE experts: all experts of a projection are now fake-quantized in one kernel launch instead of one per expert, about 2x faster per MoE layer on Qwen3-30B-A3B with FSDP2. - ``hf_ptq.py`` calibrates Qwen3.8 (``qwen4_exp``) without running its QSA indexer's per-query Python loop when every visible key fits the indexer's token budget, with bit-identical results. On Qwen3.8-Flash-Next this saves about two hours of layerwise calibration. - Speed up Hugging Face QAT and QAD with ``QATTrainer``, ``QADTrainer`` or ``--manual_gc``: the ``gc.collect()`` run before each training step now skips the model and optimizer state, saving about 1 s per step (30%) on Qwen3-30B-A3B. +- 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* @@ -59,6 +60,7 @@ Changelog - ``mtq.calibrate`` now raises an error, naming the modules, when a Hugging Face fused-experts forward reaches neither quantizer hook (a hand-written ``bmm`` forward that transformers cannot redirect, say), instead of exporting weight-derived input scales with a warning. - Layerwise-calibration resume (``checkpoint_dir``) now loads checkpoint files with ``weights_only=True``. A checkpoint directory written by an earlier release whose layer inputs hold objects other than tensors, plain containers or Hugging Face dynamic caches can no longer be resumed; recalibrate in a fresh directory, or disable checkpointing. +- ``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** @@ -152,7 +154,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* @@ -225,7 +226,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 4ef871b1ef5ebba37e61ffdb6b36c174743fd82f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 22:31:33 +0000 Subject: [PATCH 07/42] 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 828575a8ebd..0f008a7863f 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 @@ -994,15 +994,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 42836815461e3f22ea3e75ee989cc06a1f5e6db3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 15 Oct 2026 21:04:14 +0000 Subject: [PATCH 08/42] 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 0f008a7863f..1d664bfcbee 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -985,13 +985,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 c97f616edc6ee1abf3c21ec4d93b41577f1a2c9e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:00:20 +0000 Subject: [PATCH 09/42] 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 | 33 ------------------- 3 files changed, 13 insertions(+), 36 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 1d664bfcbee..f041498ae0a 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -46,7 +46,7 @@ pytest.importorskip("safetensors") pytest.importorskip("transformers") from safetensors.torch import load_file, save_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 @@ -923,15 +923,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 9621c8f29fd..6ed2b81ad57 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -25,7 +25,6 @@ non-zero expert offset matters. """ -import contextlib import copy import math import re @@ -44,7 +43,6 @@ from transformers import NemotronHConfig, PretrainedConfig from transformers.models.nemotron_h.modeling_nemotron_h import NemotronHExperts -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, @@ -58,7 +56,6 @@ from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer, TensorQuantizer -from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan NUM_EXPERTS = 4 HIDDEN_DIM = 32 @@ -796,36 +793,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"]) - - class GptOssExperts(nn.Module): """Named as transformers' fused BMM experts, which store their weights transposed.""" From 6ba1e4ac04a692297fd00ef2c5967eb925f86836 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 07:11:50 +0000 Subject: [PATCH 10/42] 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 f041498ae0a..12be9b9d41b 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -953,12 +953,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 @@ -991,7 +1004,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. """ @@ -1004,5 +1018,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 4e0157f61b7d213ff121a5d5939deeab5c930f81 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 18:22:45 +0000 Subject: [PATCH 11/42] 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 12be9b9d41b..83482b6a294 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 @@ -229,29 +227,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( @@ -475,42 +450,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, message): # Rank 0's subset reverts; rank 1's fails, as a merge whose group ownership split does. model = AutoModelForCausalLM.from_config(_mixtral_config()) @@ -828,56 +767,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) @@ -982,25 +904,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 f4217ce380d..ef113dbcc4d 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -40,6 +40,7 @@ sync_tied_input_amax, ) from modelopt.torch.export.unified_export_hf import ( + _copy_non_model_files_from_source, _fuse_shared_input_modules, _resolve_export_dtype, _source_checkpoint, @@ -1091,19 +1092,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 98dda9a687b46fa6f3cf0095572961ed5dcb91bf Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 21:54:22 +0000 Subject: [PATCH 12/42] 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 913ec45e309..a4d8a703f59 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -714,7 +714,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 83482b6a294..931d2d46f69 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 3b6066e9ad7555991ab7f5955f80a3fd22076bd5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 21:55:19 +0000 Subject: [PATCH 13/42] 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 a4d8a703f59..bc0dd1bb1b8 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -871,12 +871,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 931d2d46f69..7c2a0760bde 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -794,15 +794,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"] @@ -892,8 +883,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. @@ -930,6 +919,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 34d62b80d066b9142641532dbed3af8a82cdbe10 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 16 Oct 2026 21:56:15 +0000 Subject: [PATCH 14/42] 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 328d6b20e4b..5b1c6312e38 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1692,20 +1692,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, @@ -1953,9 +1939,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 39a8398f71a2b641f16d4aec435c1003f2480c0d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 17 Oct 2026 01:06:36 +0000 Subject: [PATCH 15/42] 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 bc0dd1bb1b8..1d3dc51caa9 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -871,23 +871,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 + + 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]}" + ) - # 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) + # 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 7c2a0760bde..81a68aa2346 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -874,7 +874,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 @@ -927,7 +927,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 aa4814ee5efd2bfd27c70877cc7f269c5cdb03b7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 22 Oct 2026 20:13:34 +0000 Subject: [PATCH 16/42] 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 87304a936b733094bd9725eaaf8c057d92401240 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 20:57:17 +0000 Subject: [PATCH 17/42] 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 1d3dc51caa9..e887423e31a 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -54,7 +54,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, @@ -672,7 +672,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 5b1c6312e38..f2205736eb7 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -96,7 +96,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: @@ -1935,6 +1940,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: @@ -2045,8 +2052,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 @@ -2108,18 +2117,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 81a68aa2346..9f754299ade 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 ) @@ -159,6 +159,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) @@ -796,7 +817,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 @@ -882,7 +905,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 0e0c4963cef80c833c5ff617cf1ff6d4ef6d41d1 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 24 Oct 2026 22:38:57 +0000 Subject: [PATCH 18/42] 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 e887423e31a..d7f6082d2fd 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -58,7 +58,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 @@ -68,9 +68,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, @@ -676,7 +680,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 @@ -714,13 +718,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 @@ -804,20 +801,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() @@ -861,7 +851,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 @@ -890,3 +880,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 f2205736eb7..45c88369bc1 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -2114,7 +2114,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 @@ -2127,7 +2127,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 9f754299ade..bbb19656894 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -148,6 +148,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.""" @@ -846,6 +863,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"] ) @@ -853,23 +893,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") @@ -985,3 +1009,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 ef9af9fa30fd3c9426382b197331a1e018173cda Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 00:20:22 +0000 Subject: [PATCH 19/42] 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 | 21 +++++++++++-- 2 files changed, 36 insertions(+), 16 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index d7f6082d2fd..ea150f8fefa 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -61,7 +61,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: @@ -699,8 +699,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 @@ -751,15 +754,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 @@ -782,12 +787,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) @@ -815,7 +818,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 @@ -824,7 +827,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 bbb19656894..870baaffc8f 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -48,6 +48,7 @@ from transformers.core_model_loading import Chunk, WeightConverter import modelopt.torch.quantization as mtq +from modelopt.torch.export import distributed_writer, quant_aware_conversion from modelopt.torch.export.distributed_writer import ( _even_bins, _is_full_world_shard0, @@ -805,10 +806,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": @@ -822,7 +835,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)) @@ -830,6 +846,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 59d543d92a11451b501e7e81c1830f674cc54768 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sun, 25 Oct 2026 01:18:06 +0000 Subject: [PATCH 20/42] 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 c793c369dec..fd458896a4f 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -60,7 +60,7 @@ Changelog - ``mtq.calibrate`` now raises an error, naming the modules, when a Hugging Face fused-experts forward reaches neither quantizer hook (a hand-written ``bmm`` forward that transformers cannot redirect, say), instead of exporting weight-derived input scales with a warning. - Layerwise-calibration resume (``checkpoint_dir``) now loads checkpoint files with ``weights_only=True``. A checkpoint directory written by an earlier release whose layer inputs hold objects other than tensors, plain containers or Hugging Face dynamic caches can no longer be resumed; recalibrate in a fresh directory, or disable checkpointing. -- ``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 3ee47659ca133b0f2770436e4bede2e3d438c5ae Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:14:16 +0000 Subject: [PATCH 21/42] 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 45c88369bc1..b820b8a2c45 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -2064,18 +2064,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 @@ -2127,6 +2115,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 870baaffc8f..6215c68cb14 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -872,12 +872,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 32b2ec5406d1b5aa11d3af0d0f01b5f5ba44e554 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:19:29 +0000 Subject: [PATCH 22/42] 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 ea150f8fefa..0607a3ebe91 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -709,6 +709,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 @@ -809,6 +811,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) @@ -864,8 +876,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 6215c68cb14..97935dc6407 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -1041,9 +1041,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 b921ae22cb85d73d6ec33e7fc298e8581d3f6013 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:20:24 +0000 Subject: [PATCH 23/42] 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 0607a3ebe91..e59d99be1e8 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -867,6 +867,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 4b404addd0c24c0eb7c1a9909a8bc2624f426f48 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 01:45:45 +0000 Subject: [PATCH 24/42] 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 e59d99be1e8..f8f3ebd1fe1 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -865,7 +865,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 97935dc6407..c673e42527d 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -1032,6 +1032,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 1f318dc525d4c49ba72856406042ad47a2a2cb8e Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 18:53:11 +0000 Subject: [PATCH 25/42] Refuse a MODELOPT_FSDP2_EXPORT value the export does not know Any value other than nogather quietly took the streaming writer, so a misspelled one (no-gather, NOGATHER, 1) ran the path it was set to avoid. The knob now takes nogather, or streaming, the default, and raises on anything else; rank 0 reads it inside broadcast_from, so the refusal reaches every rank. Signed-off-by: Shengliang Xu --- modelopt/torch/export/unified_export_hf.py | 16 +++++++++++++++- .../unit/torch/export/test_distributed_writer.py | 9 +++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index b820b8a2c45..2dd7eb58438 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1954,6 +1954,20 @@ def _write_base_config(model: nn.Module, export_dir: "Path | str") -> None: warnings.warn(f"Could not save generation_config: {gen_err}") +def _no_gather_export_asked() -> bool: + """Whether ``MODELOPT_FSDP2_EXPORT`` picks the no-gather writer for an FSDP2-sharded export. + + ``nogather`` picks it; unset, empty or ``streaming`` keeps the default streaming path. Any other + value raises: a misspelled one would otherwise quietly take the path it was set to avoid. + """ + value = os.environ.get("MODELOPT_FSDP2_EXPORT") or "streaming" + if value not in ("nogather", "streaming"): + raise ValueError( + f"MODELOPT_FSDP2_EXPORT={value!r}: expected 'nogather', or 'streaming', the default." + ) + return value == "nogather" + + def export_hf_checkpoint( model: Any, dtype: torch.dtype | None = None, @@ -2055,7 +2069,7 @@ def export_hf_checkpoint( # 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" + "MODELOPT_FSDP2_EXPORT", _no_gather_export_asked ) # 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 diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index c673e42527d..c2e66a49005 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -884,6 +884,15 @@ def test_nogather_export_saves_modelopt_state(gloo_world_of_one, tmp_path): assert (tmp_path / "export" / "modelopt_state.pth").is_file() +def test_nogather_export_refuses_a_writer_it_does_not_know(gloo_world_of_one, tmp_path): + """A misspelled ``MODELOPT_FSDP2_EXPORT`` raises, rather than quietly taking the default writer, + the one it was set to avoid.""" + model = _fsdp2_shard(_load(_save(_llama_config(), tmp_path / "src")), gloo_world_of_one) + + with pytest.raises(ValueError, match="MODELOPT_FSDP2_EXPORT='no-gather'"): + _nogather_export(model, tmp_path / "export", knob="no-gather") + + def test_nogather_export_names_modules_in_the_quant_config_as_the_weights_do( gloo_world_of_one, tmp_path ): From 8b2e41251d4f6f32b48fa2fac5fe3b28a8116d4c Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 18:55:38 +0000 Subject: [PATCH 26/42] Test that an uneven inference-mode shard is written with its global shape The two-rank writer job now writes its MoE in inference mode, with a 63-row buffer made there that two ranks shard 32 + 31, so a DTensor reaches _to_normal as an inference tensor. Without the shape and stride the rewrap passes, the two shards re-infer different global shapes and the write fails on every rank. Signed-off-by: Shengliang Xu --- .../torch/export/test_distributed_writer.py | 35 ++++++++++++++----- 1 file changed, 27 insertions(+), 8 deletions(-) diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index c2e66a49005..40c29d4e269 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -942,17 +942,31 @@ def _mtp_tail(): } +def _uneven_rows(): + """63 rows: two ranks shard them unevenly.""" + return torch.arange(63, dtype=torch.float32) + + 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, 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() if rank == 0 else None, - ) + # with the source checkpoint and those tensors. It is written in inference mode, with a buffer + # made there that shards unevenly, as an inference-mode pack makes its scales: the writer copies + # each such shard out of inference mode, and the rewrapped DTensor must keep the global shape, + # which the two uneven shards would otherwise each re-infer differently. + moe = _fsdp2_shard(_load(tmp_dir / "moe"), mesh) + with torch.inference_mode(): + rows = _uneven_rows() + local = rows[:32] if rank == 0 else rows[32:] # Shard(0) of 63 rows: 32 + 31 + moe.model.norm.register_buffer( + "uneven_rows", + DTensor.from_local( + local.clone(), mesh, [Shard(0)], run_check=False, shape=rows.shape, stride=(1,) + ), + ) + _write(moe, tmp_dir / "moe_export", 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 # in separate per-dtype buckets. Every rank quantizes identically, so rank 0's single-process @@ -1024,7 +1038,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, - and the tensors handed to every rank are written once. + and the tensors handed to every rank are written once. A tensor sharded unevenly in inference + mode is written whole. The same job also checks the cross-rank consistency errors. """ @@ -1037,7 +1052,11 @@ 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")), **_mtp_tail()} + want = { + **load_file(str(src / "model.safetensors")), + **_mtp_tail(), + "model.norm.uneven_rows": _uneven_rows(), + } _assert_byte_identical(_load_checkpoint(tmp_path / "moe_export"), want) From 11596190968d8e93690bd1052b6148d0134cb179 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Tue, 27 Oct 2026 19:51:57 +0000 Subject: [PATCH 27/42] Drop the no-gather arm's unused tied-weight map _prepare_model_for_export returns a tied-weight map the no-gather arm never reads: the writer builds its own from the model's names. The arm now discards it, and says so. Signed-off-by: Shengliang Xu --- modelopt/torch/export/unified_export_hf.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 2dd7eb58438..4d028fe17e2 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -2109,7 +2109,8 @@ def export_hf_checkpoint( 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) + # The writer builds its own name-based tied-weight map, so this one is not handed down. + _dtype, _, 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. From 2926d00d0ae6504871e6c353b04467bccbbb8fbc Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 21:45:31 +0000 Subject: [PATCH 28/42] Use the writer's module-scope imports in the no-gather export The writer pieces now import torch.distributed, DTensor, all_reduce and revert_weight_conversion_quant_aware at module scope. The no-gather export's own import of the expert-split helpers joins them, rather than following the checkpoint-file pattern, and its function-local imports of the same three names go. A test that fails the reverse conversion has to patch the name the writer calls, which is now its own module-scope import: patching it on quant_aware_conversion no longer reaches the writer, so the test would see every name reverted and no warning. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 7 +------ tests/unit/torch/export/test_distributed_writer.py | 8 ++++---- 2 files changed, 5 insertions(+), 10 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index f8f3ebd1fe1..5123ecd7874 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -36,6 +36,7 @@ from modelopt.torch.utils.distributed import all_reduce +from .moe_utils import _FUSED_PROJ, _fused_experts_prefixes, _split_local_fused_module from .quant_aware_conversion import ( _QUANT_STATE_SUFFIXES, QuantConversionUnsupportedError, @@ -47,8 +48,6 @@ # A checkpoint's files: its weights, sharded as save_pretrained names them or in one file, and index. _CHECKPOINT_FILES = re.compile(r"model(-\d{5}-of-\d{5})?\.safetensors(\.index\.json)?") -from .moe_utils import _FUSED_PROJ, _fused_experts_prefixes, _split_local_fused_module - def distributed_save_hf_checkpoint( model: nn.Module, @@ -705,11 +704,7 @@ def _distributed_save_hf_checkpoint_impl( 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 - 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 diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 40c29d4e269..9eaa58945e7 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -48,7 +48,7 @@ from transformers.core_model_loading import Chunk, WeightConverter import modelopt.torch.quantization as mtq -from modelopt.torch.export import distributed_writer, quant_aware_conversion +from modelopt.torch.export import distributed_writer from modelopt.torch.export.distributed_writer import ( _even_bins, _is_full_world_shard0, @@ -66,6 +66,8 @@ ) from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan +_WRITER = "modelopt.torch.export.distributed_writer" + def _write(model, export_dir, **kwargs): distributed_save_hf_checkpoint( @@ -79,8 +81,6 @@ def _nogather_export(model, export_dir, knob="nogather", **kwargs): model, dtype=torch.bfloat16, export_dir=export_dir, max_shard_size="16KB", **kwargs ) -_WRITER = "modelopt.torch.export.distributed_writer" - def _raw_bytes(tensor): return tensor.detach().contiguous().reshape(-1).view(torch.uint8) @@ -901,7 +901,7 @@ def test_nogather_export_names_modules_in_the_quant_config_as_the_weights_do( 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, + distributed_writer, "revert_weight_conversion_quant_aware", side_effect=RuntimeError("incomplete merge"), ) From f22b9dc03e504b3b87918268fbdec4731d5bca4a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 21:46:22 +0000 Subject: [PATCH 29/42] Record the plan hf_ptq's FSDP2 load shards the model by export_hf_checkpoint now reads how a model is distributed from its recorded plan and refuses an FSDP2-sharded model that records none, so hf_ptq --use_fsdp2 records DistributionPlan(world_size, dense="fsdp") right after its load. The recording moved here from the calibration PR, whose export did not need it yet, and brings what it switches on there: on more than one GPU the ranks size the calibration batch together and calibration that keeps its statistics on each rank is refused. So the README notes on the refusal and the batch size, the AutoQuantize clause, the backward-breaking changelog entry and the test that the load records its plan come back with it. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + examples/hf_ptq/README.md | 6 +++++- examples/hf_ptq/hf_ptq.py | 5 +++++ tests/examples/hf_ptq/test_hf_ptq_args.py | 24 ++++++++++++++++++++++- 4 files changed, 34 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index fd458896a4f..d10c8da24d8 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -47,6 +47,7 @@ Changelog **Backward Breaking Changes** +- ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` on more than one GPU now refuses calibration that keeps its statistics on each rank, which calibrated the ranks' shards of a weight differently: AWQ-Lite, AWQ-Full, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate (run them without ``--use_fsdp2``). - ``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``. - Drop support for Python 3.10, which has reached end-of-life. ModelOpt now requires Python 3.11 or later. - Bump minimum required PyTorch version to 2.10. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 08cb43e6d64..3f6ea1957ee 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -610,14 +610,18 @@ mtq.calibrate(model, algorithm="max", forward_loop=calibrate_loop) ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for distributed model sharding and calibration, exposed via the `--use_fsdp2` flag on the standard `hf_ptq.py` entry point. -> *Weight and KV-cache AutoQuantize with `--use_fsdp2` are experimental.* +> *Weight and KV-cache AutoQuantize with `--use_fsdp2` are experimental, and refused when a candidate format is one of those refused below.* > *Nemotron-H checkpoints containing MTP weights are not supported with `--use_fsdp2` and are rejected before distributed model loading. Remove `--use_fsdp2` to load and calibrate the MTP modules; the distributed loader does not construct them.* > *AWQ calibration (`awq_lite`, `awq_full`), as the AWQ formats and recipes use, AutoQuantize candidates included, is not supported with `--use_fsdp2` on more than one rank and is rejected before distributed model loading: each rank would fold its own smoothing scale into its weight shard. Drop `--use_fsdp2` (the single-process `device_map` still spreads the model over all the GPUs), run on one rank, or use a non-AWQ format.* +> *With `--use_fsdp2` on more than one GPU, calibration that keeps its statistics on each rank is refused when quantization starts: AWQ and SVDQuant (`--qformat int4_awq`, `w4a8_awq_beta`, `nvfp4_awq_lite`, `nvfp4_awq_full`, `nvfp4_svdquant`), local Hessian (`nvfp4_w4a4_weight_local_hessian`), imatrix (the `iq*` formats), the recipes that use them, and AutoQuantize recipes with such a candidate format. Each rank calibrates on its own share of the data, so the ranks would calibrate their shards of a weight differently. Run these without `--use_fsdp2`.* + > *`nvfp4_awq_clip` still works with `--use_fsdp2` on more than one GPU: it calibrates with AWQ-Clip alone, which folds no smoothing scale into the weights, and its clip search averages each candidate clip's loss over the ranks, so every rank takes the same clip.* +> *With `--use_fsdp2` on more than one GPU and `--batch_size 0`, each rank estimates the calibration batch size from one forward and the ranks use the smallest. That size is not probed for out-of-memory, since a rank that runs out of memory partway through a forward would leave the others waiting in its collectives: set `--batch_size` if calibration runs out of memory.* + ### Usage #### Slurm (recommended) diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 47d6a8b2c19..c79c25ccf98 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -99,6 +99,7 @@ get_max_batch_size, get_supported_datasets, ) +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan from modelopt.torch.utils.memory_monitor import launch_memory_monitor from modelopt.torch.utils.mlflow import add_mlflow_args, resolve_mlflow_args from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -339,6 +340,10 @@ def load_model(args: argparse.Namespace): attn_implementation=args.attn_implementation, hf_config=hf_config, ) + # Calibration sizes its batch, and the export writes, by how the model is distributed. + set_distribution_plan( + full_model, DistributionPlan(args.dist_state.world_size, dense="fsdp") + ) elif args.specdec_offline_dataset is not None or not args.low_memory_mode: full_model = get_model( args.local_checkpoint_path, diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 2f0b425de40..5543e9de902 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -26,7 +26,7 @@ import yaml from _test_utils.mlflow import clean_env # noqa: F401 from _test_utils.torch.transformers_models import get_tiny_qwen3 -from transformers import Qwen3_5MoeConfig, Qwen3MoeConfig +from transformers import Qwen3_5MoeConfig, Qwen3Config, Qwen3MoeConfig from modelopt.recipe import load_recipe from modelopt.recipe.config import ( @@ -38,6 +38,7 @@ from modelopt.torch.quantization import tensor_quant from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.utils import mlflow as mlflow_lib +from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.mlflow import describe_run, run_tags _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" @@ -480,6 +481,27 @@ def distributed_load(*args, **kwargs): hf_ptq.load_model(args) +def test_the_fsdp2_load_records_the_plan_it_loaded_by(monkeypatch, tmp_path): + """Calibration sizes its batch, and the export writes, by the model's distribution plan, and an + FSDP2-sharded model without one is refused rather than taken for a single process's.""" + Qwen3Config(architectures=["Qwen3ForCausalLM"]).save_pretrained(tmp_path) + monkeypatch.setenv("RANK", "0") + hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", str(tmp_path), "--use_fsdp2") + args.dist_state = SimpleNamespace(device="cpu", rank=0, world_size=2) + args.hub_model_id, args.local_checkpoint_path = None, str(tmp_path) # as main() sets them + model = torch.nn.Linear(1, 1) + monkeypatch.setattr(hf_ptq, "parallel_load_and_prepare_fsdp2", lambda *args, **kwargs: model) + + def loaded(model): + raise RuntimeError("loaded") + + monkeypatch.setattr(hf_ptq, "prepare_model_for_calibration", loaded) + + with pytest.raises(RuntimeError, match="loaded"): + hf_ptq.load_model(args) + assert distribution_plan(model) == DistributionPlan(2, dense="fsdp") + + def test_mlflow_flag_defaults_the_experiment_name(monkeypatch): monkeypatch.setattr(getpass, "getuser", lambda: "tester") hf_ptq, args = _parse_hf_ptq_args( From 30ecd64dbf9ea0047b36c63770d1542ba6afccb3 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 21:49:54 +0000 Subject: [PATCH 30/42] Raise a calibration out-of-memory error for a model over several ranks The calibration loop retries a batch that runs out of GPU memory as two halves. For a model whose recorded plan spans several ranks, such as hf_ptq's --use_fsdp2 model now that its load records one, a forward can hold collectives: a rank that retried would run more forwards than its peers, which would wait in their collectives. The loop now raises the error for such a model and keeps the retry for any other. The README and the changelog entry say so, and tell the user to set --batch_size. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- examples/hf_ptq/README.md | 2 +- modelopt/torch/utils/dataset_utils.py | 23 ++++++++++++++--- tests/unit/torch/utils/test_dataset_utils.py | 27 ++++++++++++++++++++ 4 files changed, 49 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index d10c8da24d8..da6c17bc231 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -47,7 +47,7 @@ Changelog **Backward Breaking Changes** -- ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` on more than one GPU now refuses calibration that keeps its statistics on each rank, which calibrated the ranks' shards of a weight differently: AWQ-Lite, AWQ-Full, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate (run them without ``--use_fsdp2``). +- ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` on more than one GPU now refuses calibration that keeps its statistics on each rank, which calibrated the ranks' shards of a weight differently: AWQ-Lite, AWQ-Full, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate (run them without ``--use_fsdp2``). It also raises when a calibration batch runs out of GPU memory, instead of retrying it in smaller pieces: set ``--batch_size`` to one that fits. - ``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``. - Drop support for Python 3.10, which has reached end-of-life. ModelOpt now requires Python 3.11 or later. - Bump minimum required PyTorch version to 2.10. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 3f6ea1957ee..1e726b5f801 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -620,7 +620,7 @@ ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for > *`nvfp4_awq_clip` still works with `--use_fsdp2` on more than one GPU: it calibrates with AWQ-Clip alone, which folds no smoothing scale into the weights, and its clip search averages each candidate clip's loss over the ranks, so every rank takes the same clip.* -> *With `--use_fsdp2` on more than one GPU and `--batch_size 0`, each rank estimates the calibration batch size from one forward and the ranks use the smallest. That size is not probed for out-of-memory, since a rank that runs out of memory partway through a forward would leave the others waiting in its collectives: set `--batch_size` if calibration runs out of memory.* +> *With `--use_fsdp2` on more than one GPU, a calibration batch that runs out of memory is not retried in smaller pieces, since a rank that runs out of memory partway through a forward would leave the others waiting in its collectives. With `--batch_size 0`, each rank estimates the batch size from one forward and the ranks use the smallest, which is not probed for out-of-memory either: set `--batch_size` if calibration runs out of memory.* ### Usage diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 124cded7a38..7f741ad0bad 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -1162,7 +1162,11 @@ def _slice_batch(batch_data, start, end): def _process_batch( - batch_data, infer_method, max_working_batch_size=None, allowed_non_tensor_keys=None + batch_data, + infer_method, + max_working_batch_size=None, + allowed_non_tensor_keys=None, + split_on_oom=True, ): """Process a batch of data through the model's inference method. @@ -1171,6 +1175,8 @@ def _process_batch( infer_method: Model's inference method (either forward or generate) max_working_batch_size: Maximum batch size known to work without OOM allowed_non_tensor_keys: Set of key names whose values may be non-tensor types + split_on_oom: Whether a batch that runs out of GPU memory is retried as two halves; + otherwise the error is raised Returns: The maximum batch size that worked successfully @@ -1193,7 +1199,11 @@ def _process_batch( width = min(max_working_batch_size, batch_size - offset) split_data = _slice_batch(batch_data, offset, offset + width) max_working_batch_size = _process_batch( - split_data, infer_method, max_working_batch_size, allowed_non_tensor_keys + split_data, + infer_method, + max_working_batch_size, + allowed_non_tensor_keys, + split_on_oom, ) offset += width @@ -1208,6 +1218,8 @@ def _process_batch( else max(batch_size, max_working_batch_size) ) # This batch size worked successfully except torch.cuda.OutOfMemoryError: + if not split_on_oom: + raise assert batch_size > 1, ( "CUDA out of memory error occurred while processing a single sample. " "This indicates the model is too large for the available GPU memory. " @@ -1245,6 +1257,11 @@ def _forward_loop( dataloader: DataLoader containing the batched input data allowed_non_tensor_keys: Set of key names whose values may be non-tensor types """ + # A batch that runs out of GPU memory is retried in halves, except for a model whose recorded plan + # spans several ranks: its forward can hold collectives, and a rank that retried would run more + # forwards than its peers, which would wait in them. + plan = recorded_distribution_plan(model) + split_on_oom = plan is None or plan.world_size == 1 with _disable_use_cache(model), torch.no_grad(): is_enc_dec = model_type_is_enc_dec(model) # Call the module (not .forward) so nn.Module.__call__ runs pre/post-forward hooks — this is @@ -1255,7 +1272,7 @@ def _forward_loop( for _, data in enumerate(tqdm(dataloader)): # Process batch and update max working batch size max_working_batch_size = _process_batch( - data, infer_method, max_working_batch_size, allowed_non_tensor_keys + data, infer_method, max_working_batch_size, allowed_non_tensor_keys, split_on_oom ) diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index 138cc35c1da..b77a5d94f1c 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -376,6 +376,33 @@ def _collate(samples): assert model.config.use_cache is True +@pytest.mark.parametrize("world_size", [None, 1, 2], ids=["no-plan", "one-rank", "two-ranks"]) +def test_forward_loop_raises_out_of_memory_for_a_model_over_several_ranks(world_size): + """A model distributed over several ranks can run collectives in its forward, so a rank must not + retry a smaller batch after running out of memory: it would run more forwards than its peers. + Any other model splits the batch, as before.""" + batch_sizes = [] + + class _Model(torch.nn.Module): + def forward(self, input_ids): + batch_sizes.append(input_ids.shape[0]) + if input_ids.shape[0] > 2: + raise torch.cuda.OutOfMemoryError + + model = _Model() + if world_size is not None: + set_distribution_plan(model, DistributionPlan(world_size, dense="fsdp")) + loader = DataLoader([{"input_ids": torch.zeros(8, dtype=torch.long)}] * 4, batch_size=4) + + if world_size == 2: + with pytest.raises(torch.cuda.OutOfMemoryError): + _forward_loop(model, loader) + assert batch_sizes == [4] + else: + _forward_loop(model, loader) + assert batch_sizes == [4, 2, 2] + + def test_disable_use_cache_restores_on_exception(): """Restore must run even if the with-block raises.""" model = torch.nn.Linear(4, 4) From c9ffe8077af4de4ff159c2bdf3f580d5c122f0e7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 21:50:46 +0000 Subject: [PATCH 31/42] Say who records the plan the export reads, now that the loader does not The changelog entry and a GPU test comment said load_pretrained records the plan export_hf_checkpoint reads. Its recording is not below this change, so the entry says the plan is the one recorded on the model, names hf_ptq --use_fsdp2 as a caller that records it, and keeps the advice for a model distributed by hand. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- tests/gpu/torch/export/test_distributed_hf_export.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index da6c17bc231..6bab689bd50 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -61,7 +61,7 @@ Changelog - ``mtq.calibrate`` now raises an error, naming the modules, when a Hugging Face fused-experts forward reaches neither quantizer hook (a hand-written ``bmm`` forward that transformers cannot redirect, say), instead of exporting weight-derived input scales with a warning. - Layerwise-calibration resume (``checkpoint_dir``) now loads checkpoint files with ``weights_only=True``. A checkpoint directory written by an earlier release whose layer inputs hold objects other than tensors, plain containers or Hugging Face dynamic caches can no longer be resumed; recalibrate in a fresh directory, or disable checkpointing. -- ``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. +- ``export_hf_checkpoint`` now reads how a model is distributed from the plan recorded on it, and raises for an FSDP2-sharded model, or one in a multi-rank run, that carries none. ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` records it; 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** diff --git a/tests/gpu/torch/export/test_distributed_hf_export.py b/tests/gpu/torch/export/test_distributed_hf_export.py index 61225f0392d..844a46dd98e 100644 --- a/tests/gpu/torch/export/test_distributed_hf_export.py +++ b/tests/gpu/torch/export/test_distributed_hf_export.py @@ -161,7 +161,7 @@ def _ptq_and_export(rank, size, *, src_dir, export_dir, quant_cfg): fsdp2_wrap(model) set_distribution_plan( model, DistributionPlan(size, dense="fsdp") - ) # wrapped by hand, so recorded here, as load_pretrained would + ) # wrapped by hand, so its plan is recorded by hand # 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. From d43544f9eba38c014a1cffeaeecf4ec2bebf8148 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 21:53:18 +0000 Subject: [PATCH 32/42] Say the no-gather export does not overwrite an earlier checkpoint The writer refuses an export directory that already holds a checkpoint, which the default FSDP2 export overwrites, so the README note that introduces MODELOPT_FSDP2_EXPORT=nogather says so. Signed-off-by: Shengliang Xu --- examples/hf_ptq/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 1e726b5f801..21b69d95bba 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -658,7 +658,7 @@ Qwen checkpoints with `model_type` `qwen3_5` or `qwen3_5_moe` support this path > *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.* +> *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 It does not overwrite a checkpoint already in the export directory.* ## Evaluate Accuracy From 5bfe8112cbb76cfc65f45fe1da4076850b29add5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 21:53:39 +0000 Subject: [PATCH 33/42] Restore the full stop the no-gather README note lost Signed-off-by: Shengliang Xu --- examples/hf_ptq/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 21b69d95bba..44c8e91416f 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -658,7 +658,7 @@ Qwen checkpoints with `model_type` `qwen3_5` or `qwen3_5_moe` support this path > *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 It does not overwrite a checkpoint already in the export directory.* +> *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. It does not overwrite a checkpoint already in the export directory.* ## Evaluate Accuracy From 23c0b55b2158992d0485455c7ecdfe110093b47d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 22:13:09 +0000 Subject: [PATCH 34/42] Record the plan in the GPU tests that export a model they shard by hand export_hf_checkpoint refuses an FSDP2-sharded model with no recorded plan, and neither fsdp2_wrap nor parallel_load_and_prepare_fsdp2 records one. The carry-over export tests wrap their models with fsdp2_wrap, and the loader tests load with parallel_load_and_prepare_fsdp2, so they now record DistributionPlan(size, dense="fsdp") themselves. The lines are those the load-routing PR adds to the same tests: that PR was below this one before the 2026-10-29 re-cut, so whichever of the two lands second merges them unchanged. The changelog entry names parallel_load_and_prepare_fsdp2 next to fsdp2_wrap as a way to distribute a model that records no plan. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- tests/gpu/torch/export/test_export_carry_over.py | 10 +++++++++- tests/gpu/torch/utils/test_model_load_utils.py | 8 +++++++- 3 files changed, 17 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6bab689bd50..a052d081046 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -61,7 +61,7 @@ Changelog - ``mtq.calibrate`` now raises an error, naming the modules, when a Hugging Face fused-experts forward reaches neither quantizer hook (a hand-written ``bmm`` forward that transformers cannot redirect, say), instead of exporting weight-derived input scales with a warning. - Layerwise-calibration resume (``checkpoint_dir``) now loads checkpoint files with ``weights_only=True``. A checkpoint directory written by an earlier release whose layer inputs hold objects other than tensors, plain containers or Hugging Face dynamic caches can no longer be resumed; recalibrate in a fresh directory, or disable checkpointing. -- ``export_hf_checkpoint`` now reads how a model is distributed from the plan recorded on it, and raises for an FSDP2-sharded model, or one in a multi-rank run, that carries none. ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` records it; 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 recorded on it, and raises for an FSDP2-sharded model, or one in a multi-rank run, that carries none. ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` records it; if you distribute a model yourself, for example with ``fsdp2_wrap`` or ``parallel_load_and_prepare_fsdp2``, which record no plan, record one with ``modelopt.torch.utils.distributed.set_distribution_plan`` before calibrating or exporting it. **Deprecations** diff --git a/tests/gpu/torch/export/test_export_carry_over.py b/tests/gpu/torch/export/test_export_carry_over.py index fa528f4cdf9..34a54d23f5e 100644 --- a/tests/gpu/torch/export/test_export_carry_over.py +++ b/tests/gpu/torch/export/test_export_carry_over.py @@ -40,7 +40,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. @@ -108,6 +113,8 @@ def _ptq_and_export(rank, size, *, src_dir, export_dir, quant_cfg, **export_kwar model.eval() fsdp2_wrap(model) + # Wrapped by hand, the model records how it is distributed, as load_pretrained would. + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) # The FSDP2 export path is selected by exactly this predicate, 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. @@ -201,6 +208,7 @@ def _ptq_language_model_and_export(rank, size, *, src_dir, export_dir): ) model.eval() fsdp2_wrap(model) + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) assert is_fsdp2_model(model), "fsdp2_wrap did not shard the VLM" torch.distributed.barrier() diff --git a/tests/gpu/torch/utils/test_model_load_utils.py b/tests/gpu/torch/utils/test_model_load_utils.py index 3296ccbfc86..05b524fb19a 100644 --- a/tests/gpu/torch/utils/test_model_load_utils.py +++ b/tests/gpu/torch/utils/test_model_load_utils.py @@ -26,7 +26,11 @@ from torch.distributed.tensor import DTensor from modelopt.torch.export.unified_export_hf import export_hf_checkpoint -from modelopt.torch.utils.distributed import broadcast_state_dict +from modelopt.torch.utils.distributed import ( + DistributionPlan, + broadcast_state_dict, + set_distribution_plan, +) from modelopt.torch.utils.plugins.model_load_utils import parallel_load_and_prepare_fsdp2 VOCAB_SIZE = 64 @@ -61,6 +65,7 @@ def _test_parallel_load_and_export(rank, size, ckpt_dir, export_dir, cpu_offload size, cpu_offload=cpu_offload, ) + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) # Decoder layers AND root params (embed/lm_head) are sharded DTensors under shard_root=True. decoder_params = list(model.model.layers[0].parameters()) @@ -106,6 +111,7 @@ def _test_carry_over_unplaced_weights(rank, size, ckpt_dir, export_dir, orphan_p """Load, export, and require the weights the model could not place to survive.""" device = torch.device(f"cuda:{rank}") model = parallel_load_and_prepare_fsdp2(ckpt_dir, device, rank, size) + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) unplaced = getattr(model, "_modelopt_unplaced_source_keys", None) assert unplaced, ( From db65c363a9435c7d262acc8a5774d98b140a845f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 22:13:13 +0000 Subject: [PATCH 35/42] Say what sets keep_fused_experts now that the no-gather export does The comment above the held weight windows was written before any caller set keep_fused_experts, and said that until then no window opened there. The no-gather export sets it, so the comment says so. Signed-off-by: Shengliang Xu --- modelopt/torch/export/unified_export_hf.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 4d028fe17e2..1fba9c60777 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1012,9 +1012,9 @@ def _process_quantized_modules( keep_fused_experts=keep_fused_experts, ) - # For the no-gather export, which is to set keep_fused_experts on this context and let a sharded - # model past the assert above: each FSDP2 unit's modules are then packed in one weight window, - # held across them. Until then no window opens here. + # The no-gather export sets keep_fused_experts on this context, which lets a sharded model past + # the assert above: each FSDP2 unit's modules are then packed in one weight window, held across + # them. with _hf_export_handlers.held_fsdp2_weight_windows(ctx): for name, sub_module in model.named_modules(): _dispatch_export_handler(name, sub_module, ctx) From c167a6a4c936bccf9a7df77988730729ce9ba5dd Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:18:21 +0000 Subject: [PATCH 36/42] Drop the expert-parallel scaffolding the FSDP2 writer never sets The writer held ep_rank and dp_idx at 0, so every rank wrote its experts, and the expert count it read off the config could not change an expert's index. Under FSDP2 alone an expert's global index is the DTensor's own dim-0 offset, which _split_local_fused_module reads. The expert-parallel export, which reads both coordinates from the routed experts' ParallelState, adds them back where it sets them. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 57 ++++----------------- 1 file changed, 11 insertions(+), 46 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index 5123ecd7874..bc338f35ea0 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -687,7 +687,7 @@ def _distributed_save_hf_checkpoint_impl( 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. + already disjoint across ranks. 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 @@ -718,37 +718,6 @@ def _distributed_save_hf_checkpoint_impl( export_dir = Path(export_dir) rank, world = dist.get_rank(), dist.get_world_size() - - # 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 inference tensor as a NORMAL one UP FRONT. A caller in inference mode hands @@ -822,21 +791,17 @@ def _is_expert_key(key: str) -> bool: 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). + # FSDP2 shards a fused 3-D expert weight on dim 0, the expert axis, so each per-expert key's global + # index is the DTensor's own dim-0 offset, which _split_local_fused_module reads. local_sd: dict = {} - if write_experts: - for prefix in 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: 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() + for prefix in prefixes: + local_sd.update(_split_local_fused_module(prefix, sharded_sd)) + local_sd = postprocess_state_dict( + local_sd, maxbound, kv_cache_format, is_modelopt_qlora, tied_map=tied_map + ) + # 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 # set so _materialize_owned can tell disjoint experts from replicated dense. expert_key_set = set(local_sd) From c57a8cf9a26b07505a05f20dacbff3db69386da5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:19:53 +0000 Subject: [PATCH 37/42] Import the writer's helpers at module scope, and say why the writer is not The writer imported get_model_state_dict, TiedWeightMap and postprocess_state_dict inside its function, though nothing forces it: it already imports unified_export_hf_streaming at module scope, which imports model_utils and quant_utils itself. They move to the top. The export imports the writer itself inside the branch that runs it, as it does the streaming writers: the writer imports unified_export_hf_streaming, which imports unified_export_hf, so an import at module scope there fails import modelopt.torch.export with a circular import. A comment there now says so. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 8 +++----- modelopt/torch/export/unified_export_hf.py | 1 + 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index bc338f35ea0..0db36d70995 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -32,10 +32,12 @@ import torch.distributed as dist import torch.nn as nn from safetensors.torch import save_file +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 .moe_utils import _FUSED_PROJ, _fused_experts_prefixes, _split_local_fused_module from .quant_aware_conversion import ( _QUANT_STATE_SUFFIXES, @@ -43,6 +45,7 @@ _build_reverse_rules, revert_weight_conversion_quant_aware, ) +from .quant_utils import postprocess_state_dict from .unified_export_hf_streaming import _parse_shard_size # A checkpoint's files: its weights, sharded as save_pretrained names them or in one file, and index. @@ -704,11 +707,6 @@ def _distributed_save_hf_checkpoint_impl( inference tensor in its shard. Placements the point-to-point transport does not cover are rejected rather than routed through an untested slower path. """ - from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict - - 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 diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 1fba9c60777..77946677995 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -2106,6 +2106,7 @@ def export_hf_checkpoint( # 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. + # Imported here: the writer imports unified_export_hf_streaming, which imports this module. from .distributed_writer import distributed_save_hf_checkpoint _qlora = kwargs.get("is_modelopt_qlora", False) From 4ce6a495b458c1c035954fc5738f557974f3c439 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:21:11 +0000 Subject: [PATCH 38/42] Patch the state-dict read where the writer now imports it The two-rank input-scale test hands one rank a state dict without a scale by patching get_model_state_dict. The writer now binds that name at module scope, so the test patches it on distributed_writer instead of on torch's module, which no longer reaches the writer. Signed-off-by: Shengliang Xu --- tests/unit/torch/export/test_distributed_writer.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 9eaa58945e7..265da3a60c7 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -1105,9 +1105,7 @@ def _expert_input_scale_job(rank, world_size, tmp_dir): # 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 + real = distributed_writer.get_model_state_dict def without_a_scale(*args, **kwargs): sd = real(*args, **kwargs) @@ -1118,7 +1116,7 @@ def without_a_scale(*args, **kwargs): try: with ( bounded_collectives(), - mock.patch.object(state_dict, "get_model_state_dict", without_a_scale), + mock.patch.object(distributed_writer, "get_model_state_dict", without_a_scale), ): _nogather_export(_fsdp2_shard(model, mesh), tmp_dir / "diverging") except RuntimeError as error: From f6dfb95d87949c13602a8d0a6fd5c9c07ade0297 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 29 Oct 2026 23:23:13 +0000 Subject: [PATCH 39/42] Write the no-gather export's config files with save_non_weight_artifacts _write_base_config repeated save_non_weight_artifacts without its copy of the source's non-model files, which the end of export_hf_checkpoint then made for this writer alone. The FSDP2 streaming writer calls save_non_weight_artifacts from rank 0 once its shards are written, so the no-gather writer now does the same, and the final copy is skipped for both FSDP2 writers again, as on main. What differed does not change what is written. transformers sets a model's generation_config only when it can generate, so the can_generate() guard decided nothing the generation_config check does not. A config's save_pretrained never reaches revert_weight_conversion, so the patch the helper applies around it changes nothing. The config, generation config, quant config, index and copied files of a quantized Llama and Qwen2-MoE export are byte-identical before and after. Signed-off-by: Shengliang Xu --- modelopt/torch/export/unified_export_hf.py | 28 +++------------------- 1 file changed, 3 insertions(+), 25 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 77946677995..e706b344b83 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1933,27 +1933,6 @@ 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. - """ - # 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: - # 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 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: - warnings.warn(f"Could not save generation_config: {gen_err}") - - def _no_gather_export_asked() -> bool: """Whether ``MODELOPT_FSDP2_EXPORT`` picks the no-gather writer for an FSDP2-sharded export. @@ -2130,7 +2109,7 @@ def export_hf_checkpoint( extra_state_dict=extra_state_dict, ) if rank == 0: - _write_base_config(model, export_dir) + save_non_weight_artifacts(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) @@ -2227,9 +2206,8 @@ 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 and not use_no_gather_writer)): - # The streaming paths already did, from save_non_weight_artifacts; the no-gather - # writer does not. + if not (_offloaded or is_fsdp2_sharded): + # The streaming paths already did, from save_non_weight_artifacts. _copy_non_model_files_from_source(model, export_dir) except Exception as e: From 7c5955dba188815c9269972ab91b650c8ec2b69a Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 00:06:09 +0000 Subject: [PATCH 40/42] Number the local split's experts by their FSDP2 offset alone _split_local_fused_module took an expert-parallel rank and expert count that its one caller, the FSDP2 writer, never passes: an expert's global index is the DTensor's own dim-0 offset. The parameters, the layouts they told apart, and the test of the numbering no caller here can produce go; the expert-parallel export, which sets them, adds them back. The two-rank fold test still checks each rank's own split at its offset. Signed-off-by: Shengliang Xu --- modelopt/torch/export/moe_utils.py | 27 +++---------------- .../export/test_fused_experts_keep_fused.py | 26 ------------------ 2 files changed, 4 insertions(+), 49 deletions(-) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 50b8d909f3b..e12ed147abf 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -804,29 +804,17 @@ def _dtensor_dim0_offset(dt) -> int: return 0 -def _split_local_fused_module( - prefix: str, sharded_sd: dict, ep_rank: int = 0, total_experts: int = 0 -) -> dict: +def _split_local_fused_module(prefix: str, sharded_sd: dict) -> dict: """Split this rank's local shard of a fused-experts module into per-expert keys. Split THIS rank's *local* shard of one fused-experts module into per-expert keys with GLOBAL expert indices (plain, local tensors -- no all-gather). - FSDP2/EP shard the fused 3-D expert weight on dim 0 (the expert axis), so every expert this rank + FSDP2 shards the fused 3-D expert weight on dim 0 (the expert axis), so every expert this rank owns is whole and local. We reuse :func:`split_fused_experts_state_dict` on the ``to_local()`` shard (which numbers experts ``0..E_local-1``) and then shift each expert index by the rank's global dim-0 offset, so the resulting per-expert keys match the single-process export and stay on this rank for a direct distributed write (no consolidation). - - The fused expert tensor can arrive in two EP layouts, distinguished by its dim-0 size vs - ``total_experts`` (the config expert count): - * **per-EP-group** (size < total): transformers pre-sliced the param to this rank's EP group - (FSDP x EP -> a dp-sharded DTensor of the group; or a plain ep-shard). ``_dtensor_dim0_offset`` - only gives the WITHIN-group offset, so we add the EP block base ``ep_rank * group_size``. - * **all-experts** (size == total): an EP DTensor sharded across the ep mesh spanning all experts - (classic DP x EP). ``_dtensor_dim0_offset`` already returns the TRUE global offset -> no base. - ``total_experts == 0``, a count the caller could not read off the model's config, takes the - tensor for a per-EP-group slice. """ gated = f"{prefix}.gate_up_proj" in sharded_sd first = "gate_up_proj" if gated else "up_proj" @@ -834,13 +822,6 @@ def _split_local_fused_module( # dp_offset: this rank's dim-0 start WITHIN its DTensor; n_local: experts it owns. dp_offset = _dtensor_dim0_offset(fw) if isinstance(fw, DTensor) else 0 n_local = fw.to_local().shape[0] if isinstance(fw, DTensor) else fw.shape[0] - group_size = int(fw.shape[0]) - # Add the EP block base only when the tensor is a per-EP-group slice (size < total_experts) or - # total is unknown (0). When it spans all experts, dp_offset is already global. - add_ep_base = ( - bool(ep_rank) and group_size > 0 and (total_experts == 0 or group_size < total_experts) - ) - offset = dp_offset + (ep_rank * group_size if add_ep_base else 0) def _loc(key): v = sharded_sd.get(key) @@ -883,12 +864,12 @@ def _loc_scale(key): split_local = split_fused_experts_state_dict( local_fused ) # keys {prefix}.{i}.* (i in 0..E_local-1) - if offset == 0: + if dp_offset == 0: return split_local plen = len(prefix) + 1 out: dict = {} for key, val in split_local.items(): i_str, tail = key[plen:].split(".", 1) - out[f"{prefix}.{int(i_str) + offset}.{tail}"] = val + out[f"{prefix}.{int(i_str) + dp_offset}.{tail}"] = val return out 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 6ed2b81ad57..a646c3743a3 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -565,32 +565,6 @@ def test_fused_experts_prefixes_lists_each_fused_module_once_in_order(): ] -@pytest.mark.parametrize( - ("ep_rank", "total_experts", "first_index"), - [ - (0, 0, 0), - # This EP group's slice of a larger expert set: shift by the group base. - (1, 2 * LOCAL_EXPERTS, LOCAL_EXPERTS), - (3, 4 * LOCAL_EXPERTS, 3 * LOCAL_EXPERTS), - # The tensor already spans every expert, so its indices are global. - (1, LOCAL_EXPERTS, 0), - # Total unknown: the legacy per-group shift. - (1, 0, LOCAL_EXPERTS), - ], -) -def test_split_local_numbers_experts_globally(ep_rank, total_experts, first_index): - sd = _fused_state_dict(gated=True) - unshifted = split_fused_experts_state_dict(dict(sd)) - - out = _split_local_fused_module(PREFIX, sd, ep_rank=ep_rank, total_experts=total_experts) - - want = {} - for key, value in unshifted.items(): - expert, tail = key[len(PREFIX) + 1 :].split(".", 1) - want[f"{PREFIX}.{int(expert) + first_index}.{tail}"] = value - _assert_byte_identical(out, want) - - # ------------------------------------------------------------------ sharded (DTensor) experts From 55f0bd97ec855d6a75b614bc67964643987b3674 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 30 Oct 2026 22:22:25 +0000 Subject: [PATCH 41/42] Leave AWQ out of the README's list of calibrations refused when quantization starts Main now refuses AWQ with --use_fsdp2 on more than one rank before the model loads (#3128), and its README note says so. The --use_fsdp2 note this PR adds still listed AWQ among the calibrations refused when quantization starts. It now lists the others and points to the AWQ note above it. Signed-off-by: Shengliang Xu --- examples/hf_ptq/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 44c8e91416f..71be2080d53 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -616,7 +616,7 @@ ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for > *AWQ calibration (`awq_lite`, `awq_full`), as the AWQ formats and recipes use, AutoQuantize candidates included, is not supported with `--use_fsdp2` on more than one rank and is rejected before distributed model loading: each rank would fold its own smoothing scale into its weight shard. Drop `--use_fsdp2` (the single-process `device_map` still spreads the model over all the GPUs), run on one rank, or use a non-AWQ format.* -> *With `--use_fsdp2` on more than one GPU, calibration that keeps its statistics on each rank is refused when quantization starts: AWQ and SVDQuant (`--qformat int4_awq`, `w4a8_awq_beta`, `nvfp4_awq_lite`, `nvfp4_awq_full`, `nvfp4_svdquant`), local Hessian (`nvfp4_w4a4_weight_local_hessian`), imatrix (the `iq*` formats), the recipes that use them, and AutoQuantize recipes with such a candidate format. Each rank calibrates on its own share of the data, so the ranks would calibrate their shards of a weight differently. Run these without `--use_fsdp2`.* +> *With `--use_fsdp2` on more than one GPU, calibration other than AWQ that keeps its statistics on each rank is refused when quantization starts: SVDQuant (`--qformat nvfp4_svdquant`), local Hessian (`nvfp4_w4a4_weight_local_hessian`), imatrix (the `iq*` formats), the recipes that use them, and AutoQuantize recipes with such a candidate format. Each rank calibrates on its own share of the data, so the ranks would calibrate their shards of a weight differently. Run these without `--use_fsdp2`. For AWQ, see the note above.* > *`nvfp4_awq_clip` still works with `--use_fsdp2` on more than one GPU: it calibrates with AWQ-Clip alone, which folds no smoothing scale into the weights, and its clip search averages each candidate clip's loss over the ranks, so every rank takes the same clip.* From ac293230a24da64f56c8104bc2f272c90b7c9c7f Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Sat, 31 Oct 2026 00:08:45 +0000 Subject: [PATCH 42/42] Fix stale comments on the carried tensors, consolidation and EP Only rank 0 is handed the tensors the model does not own, yet the two-rank test's comment and docstring, and the writer's comment on them, said every rank holds them. The export's comment on packing named a consolidation step that does not exist and how main once split preparation from packing, and the writer's step (2) comment named an expert-parallel layout this PR does not handle. Comments and a docstring only. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distributed_writer.py | 25 ++++++++++--------- modelopt/torch/export/unified_export_hf.py | 6 ++--- .../torch/export/test_distributed_writer.py | 15 +++++------ 3 files changed, 24 insertions(+), 22 deletions(-) diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index 0db36d70995..21ad6da9f34 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -804,10 +804,11 @@ def _is_expert_key(key: str) -> bool: # 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). + # (2) Dense / non-expert: leave them as they are for now -- FSDP2 leaves DTensors (dim-0 shard), + # and what it does not shard stays a plain tensor, the same on every rank. 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( @@ -827,14 +828,14 @@ def _is_expert_key(key: str) -> bool: 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 - # 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. + # weights read straight off the source checkpoint that the FSDP2 loader drops. They are plain, + # whole CPU tensors, and only rank 0 adds them to its write set (the export entry hands them to + # rank 0 alone); 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. def check_no_overlap() -> None: if rank != 0 or not extra_state_dict: return diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index e706b344b83..9c63130bfd1 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -2091,9 +2091,9 @@ def export_hf_checkpoint( _qlora = kwargs.get("is_modelopt_qlora", False) # The writer builds its own name-based tied-weight map, so this one is not handed down. _dtype, _, 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. + # Preparing the model does not pack its weights, so pack them here. keep_fused_experts + # leaves the MoE experts fused and sharded through the fold, so each rank folds only + # its local experts; the writer then splits each rank's experts into per-expert tensors. pack_quantized_weights(model, _dtype, _qlora, keep_fused_experts=True) if getattr(model, "hf_quantizer", None) is not None: model.hf_quantizer = None diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 265da3a60c7..a35e47086c9 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -951,11 +951,12 @@ 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, 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. It is written in inference mode, with a buffer - # made there that shards unevenly, as an inference-mode pack makes its scales: the writer copies - # each such shard out of inference mode, and the rewrapped DTensor must keep the global shape, - # which the two uneven shards would otherwise each re-infer differently. + # not own, which only rank 0 is handed, as the export entry hands them; the test compares what + # the two ranks wrote between them with the source checkpoint and those tensors. It is written + # in inference mode, with a buffer made there that shards unevenly, as an inference-mode pack + # makes its scales: the writer copies each such shard out of inference mode, and the rewrapped + # DTensor must keep the global shape, which the two uneven shards would otherwise each re-infer + # differently. moe = _fsdp2_shard(_load(tmp_dir / "moe"), mesh) with torch.inference_mode(): rows = _uneven_rows() @@ -1038,8 +1039,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, - and the tensors handed to every rank are written once. A tensor sharded unevenly in inference - mode is written whole. + and the tensors handed to rank 0 alone are written once. A tensor sharded unevenly in + inference mode is written whole. The same job also checks the cross-rank consistency errors. """