From 3a1ebda548b542dc056f045217ca9804b5635aa5 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 29 Jul 2026 21:18:13 -0700 Subject: [PATCH 01/10] feat(export): no-gather distributed HF writer for FSDP2 Add the shard-local, no-rank-0-gather HF safetensors export, decoupled from any EP/session substrate -- for a FULLY FSDP2-sharded model (no expert parallelism): - _export_common.py (new): _size_to_bytes leaf helper. - moe_utils.py: fused-experts HF-format helpers (_fused_experts_prefixes, split_fused_experts_state_dict, _dtensor_dim0_offset, _split_local_fused_module, _FUSED_PROJ) -- split a rank's LOCAL FSDP2 shard of a fused 3-D expert weight into per-expert HF keys, using the DTensor's own dim-0 offset as the global expert index (ep_rank=0, so every rank writes its own experts). - distribute.py: distributed_save_hf_checkpoint (+ _bin_pack_fqn_to_index, _finfo_accepts_int_dtypes) -- DCP HuggingFaceStorageWriter(save_distributed=True) per-rank write + parallel consolidation, no full-model host-RAM gather. The caller (hf_ptq) exports under torch.inference_mode(), so the writer runs under `inference_mode(False) + no_grad()` and re-materializes each state-dict value into a fresh normal tensor up front (empty_like+copy_; DTensor via to_local->from_local to keep mesh/placements) -- inference tensors have no version counter, which the fused-experts split's slicing and DCP's SavePlanner both require. Validated on hecate, 2 nodes x 4 GPU: Qwen3-0.6B FP8 (1 shard) and Qwen3-30B-A3B FP8 (3 shards, 128 experts sharded 16/rank, correct per-expert global indices 0..127) -- both written entirely no-gather. Signed-off-by: Shengliang Xu (cherry picked from commit 6737e0d172f67d81ba1499aafa4e6782dc5b2d46) --- modelopt/torch/export/_export_common.py | 43 ++++ modelopt/torch/export/distribute.py | 273 ++++++++++++++++++++++++ modelopt/torch/export/moe_utils.py | 241 +++++++++++++++++++++ 3 files changed, 557 insertions(+) create mode 100644 modelopt/torch/export/_export_common.py diff --git a/modelopt/torch/export/_export_common.py b/modelopt/torch/export/_export_common.py new file mode 100644 index 00000000000..7cbba0dad4d --- /dev/null +++ b/modelopt/torch/export/_export_common.py @@ -0,0 +1,43 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 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. + +"""Small general helpers shared across the HF export modules. + +Home for leaf utilities used by more than one export module (e.g. moe_utils and distribute) +that belong to neither -- keeping them here avoids a cross-module dependency between siblings. This +module must not import other export modules, so it stays a safe common dependency for all of them. +""" + +import torch + + +def _size_to_bytes(size: "str | int") -> int: + """Parse an HF-style shard-size string (``"5GB"``, ``"500MB"``, ``"1GiB"``) to bytes. + + Matches transformers' decimal convention (GB == 10**9). Bare ints pass through. + """ + if isinstance(size, int): + return size + s = str(size).strip().upper() + units = { + "KIB": 2**10, "MIB": 2**20, "GIB": 2**30, "TIB": 2**40, + "KB": 10**3, "MB": 10**6, "GB": 10**9, "TB": 10**12, + } # fmt: skip + for unit in ("KIB", "MIB", "GIB", "TIB", "KB", "MB", "GB", "TB"): + if s.endswith(unit): + return int(float(s[: -len(unit)]) * units[unit]) + return int(float(s)) + + diff --git a/modelopt/torch/export/distribute.py b/modelopt/torch/export/distribute.py index 8e457c1abfa..4350c62b325 100644 --- a/modelopt/torch/export/distribute.py +++ b/modelopt/torch/export/distribute.py @@ -303,3 +303,276 @@ def _get_weights_nbytes(weights_dict: dict[str, torch.Tensor]): if shm_writer is not None: shm_writer.close() shm_writer.unlink() + +# ----- No-gather distributed HF export (FSDP2): DCP per-rank shard write ----- +import torch.nn as nn # noqa: E402 (distributed_save_hf_checkpoint's signature uses nn.Module) + +from ._export_common import _size_to_bytes +from .moe_utils import _FUSED_PROJ, _fused_experts_prefixes, _split_local_fused_module + +@contextmanager +def _finfo_accepts_int_dtypes(): + """Work around a torch DCP consolidation bug on integer dtypes. + + ``torch.distributed.checkpoint._consolidate_hf_safetensors._parse_input_metadata`` computes each + tensor's byte size as ``torch.finfo(dtype).bits // 8`` unconditionally, which raises ``TypeError`` + for integer dtypes -- e.g. the ``uint8``-packed NVFP4 weights. ``torch.iinfo`` also exposes + ``.bits`` and consolidation only reads ``.bits``, so we temporarily shim ``torch.finfo`` to fall + back to ``torch.iinfo`` for non-floating dtypes. Present in torch 2.9-2.13 (the supported range); + remove once fixed upstream. Scoped to the ``dcp.save`` call. + """ + _orig_finfo = torch.finfo + + def _finfo_or_iinfo(dtype): + try: + return _orig_finfo(dtype) + except TypeError: + return torch.iinfo(dtype) + + torch.finfo = _finfo_or_iinfo + try: + yield + finally: + torch.finfo = _orig_finfo + + +def _bin_pack_fqn_to_index(sizes: dict, max_shard_size: "str | int") -> dict: + """Deterministically bin-pack tensor FQNs into ~``max_shard_size`` files -> ``{fqn: file_index}``. + + file_index is 1..N; HuggingFaceStorageWriter turns it into ``model--of-.safetensors``. Keys + are sorted by name so every rank computes the SAME mapping from the same global size map. Each file + gets at least one tensor even if a single tensor exceeds ``max_shard_size``. + """ + max_bytes = _size_to_bytes(max_shard_size) + mapping: dict = {} + file_idx, cur = 1, 0 + for k in sorted(sizes): + sz = sizes[k] + if cur > 0 and cur + sz > max_bytes: + file_idx += 1 + cur = 0 + mapping[k] = file_idx + cur += sz + return mapping + + +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, +) -> 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 DCP's version-counter reads 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 + ) + + +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, +) -> None: + """Distributed HF safetensors export via torch DCP -- no rank-0 full-model host-RAM gather. + + Writes an already-processed, FSDP2-sharded model to ``export_dir`` as + consolidated HF safetensors, entirely distributed: + - Fused expert weights (when present): each rank splits its LOCAL shard into per-expert keys with global indices + and keeps them in place -- no gather. dp-replica ranks under DP x EP skip this. + - Dense / non-expert weights: kept sharded (FSDP2 DTensor dim-0 shards, or EP replicated plain + tensors), NOT gathered to rank 0. + + All keys are bin-packed into a global ``fqn_to_index_mapping`` (identical on every rank); then + ``dcp.save`` with ``HuggingFaceStorageWriter(save_distributed=True)`` has each rank write only its + own keys into ``export_dir/sharded/``. ``consolidate_safetensors_files_on_every_rank`` merges + those into the final ``model-XXXXX-of-NNNNN.safetensors`` with the output files partitioned across + ranks (parallel); rank 0 writes only the index. No rank ever holds the full model in host RAM. + """ + import gc + import shutil + + import torch.distributed as dist + import torch.distributed.checkpoint as dcp + from torch.distributed.checkpoint import HuggingFaceStorageWriter + from torch.distributed.checkpoint._consolidate_hf_safetensors import ( + consolidate_safetensors_files_on_every_rank, + ) + from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict + from torch.distributed.tensor import DTensor + + from .quant_utils import postprocess_state_dict + + 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, ep_size, dp_idx = 0, 1, 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 + DCP 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) + return DTensor.from_local(_new, v.device_mesh, v.placements, run_check=False) + if v.is_inference(): + _new = torch.empty_like(v) + _new.copy_(v) + return _new + return v + + sharded_sd = {k: _to_normal(v) for k, v in sharded_sd.items()} + + 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) + local_sd = {k: v.detach().contiguous() for k, v in local_sd.items()} + + # (2) Dense / non-expert: keep sharded -- FSDP2 leaves DTensors (dim-0 shard); EP leaves plain + # replicated tensors. Do NOT gather to rank 0: DCP writes each DTensor's shards per-rank (parallel, + # no rank-0 full-model host-RAM gather) and dedups replicated plain tensors. 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) + local_sd.update(dense_sd) + del sharded_sd, dense_sd + gc.collect() + + # torch DCP HuggingFaceStorageWriter consolidation writes 0-dim (scalar) tensors as ZERO (shape () + # -> value dropped); shape (1,) and larger survive. modelopt stores per-tensor scales (input_scale/ + # weight_scale/weight_scale_2) as 0-dim scalars, so promote any 0-dim tensor to (1,) to preserve its + # value. Per-tensor () vs (1,) scales are equivalent for deployment; applies to plain + DTensor. + # (Values are already normal tensors -- re-materialized right after get_model_state_dict.) + local_sd = {k: (v.reshape(1) if v.dim() == 0 else v) for k, v in local_sd.items()} + + # (3) Global shard layout: bin-pack ALL keys into ~max_shard_size files so every rank passes the + # SAME fqn_to_index_mapping to the writer. A DTensor's byte size is its GLOBAL size (numel() is + # global) -- the consolidated file holds the whole tensor; plain keys use their own size. The + # gather dedups replicated/DTensor keys (present on every rank) and unions the disjoint expert keys. + local_sizes = {k: int(v.numel() * v.element_size()) for k, v in local_sd.items()} + gathered: list = [None] * world + dist.all_gather_object(gathered, local_sizes) + all_sizes: dict = {} + for g in gathered: + all_sizes.update(g) + fqn_to_index_mapping = _bin_pack_fqn_to_index(all_sizes, max_shard_size) + n_files = max(fqn_to_index_mapping.values()) if fqn_to_index_mapping else 1 + + # (4) Distributed write, then DISTRIBUTED consolidation. save_distributed has each rank write only + # its own keys (dense DTensor shards + this rank's plain experts) into export_dir/sharded/. We set + # enable_consolidation=False because the writer's built-in consolidation runs on RANK 0 ONLY -- a + # serial re-read+rewrite of the whole model that dominated export time for large models (~700s for + # Kimi-K2 vs ~50s writing). Instead consolidate_safetensors_files_on_every_rank partitions the + # output files across ranks (idx % world_size) so every rank merges its own subset in parallel; + # rank 0 then writes only the small index. The finfo shim covers the integer-dtype (NVFP4 uint8) + # consolidation bug; 0-dim scales were already promoted to (1,) above. + sharded_dir = export_dir / "sharded" + writer = HuggingFaceStorageWriter( + str(sharded_dir), + fqn_to_index_mapping=fqn_to_index_mapping, + save_distributed=True, + enable_consolidation=False, + thread_count=8, + ) + with _finfo_accepts_int_dtypes(): + dcp.save(local_sd, storage_writer=writer) + dist.barrier() + consolidate_safetensors_files_on_every_rank( + input_dir=str(sharded_dir), + output_dir=str(export_dir), + fqn_to_index_mapping=fqn_to_index_mapping, + num_threads=8, + ) + + # (5) Drop the intermediate per-rank sharded/ dir. + if rank == 0: + shutil.rmtree(sharded_dir, ignore_errors=True) + dist.barrier() diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index 4ce60b192ee..bf844ee8eee 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -283,3 +283,244 @@ def save_expert_token_count_table(model: nn.Module, output_dir: str | Path | Non output_path = Path(output_dir) / ".moe.html" output_path.write_text(html_content, encoding="utf-8") print(f"\033[1mExpert token count table saved to {output_path}\033[0m") + +# ----- Distributed-export helpers (fused-experts split for the no-gather FSDP2 write) ----- +_FUSED_FIRST_PROJ = ("gate_up_proj", "up_proj") # gated first / ungated first +_FUSED_PROJ = ("gate_up_proj", "up_proj", "down_proj") + + +def _fused_experts_prefixes(state_dict: dict) -> list[str]: + """Return the sorted set of ``...experts`` prefixes that carry a fused 3-D expert weight + (``.gate_up_proj`` / ``up_proj`` / ``down_proj``) in ``state_dict``. + """ + prefixes: set[str] = set() + for key, tensor in state_dict.items(): + for proj in _FUSED_PROJ: + if key.endswith(f".experts.{proj}") and hasattr(tensor, "dim") and tensor.dim() == 3: + prefixes.add(key[: -(len(proj) + 1)]) + break + return sorted(prefixes) + + +def split_fused_experts_state_dict(state_dict: dict) -> dict: + """Rewrite fused-experts entries in a (consolidated) HF state dict into per-expert entries. + + The distributed export writes experts *fused* (3-D ``experts.`` weights + fused scale + buffers, kept sharded for DCP). This post-pass splits them into the same per-expert layout the + in-model :func:`_export_fused_experts` produces, so the two export paths yield identical + checkpoints. Operates on the (already folded/quantized) tensors only -- pure slicing, no + quantization. Non-fused-experts keys pass through unchanged. + + 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`` + ``P.down_proj`` ``(E, O, I)`` + ``P.down_proj_weight_scale`` + ``P.down_proj_input_scale`` + Outputs, for each expert ``i``: + ``P.{i}.gate_proj.weight`` ``(I, H)`` (+ ``.weight_scale``/``.input_scale``) [gated only] + ``P.{i}.up_proj.weight`` ``(I, H)`` (+ ``.weight_scale``/``.input_scale``) + ``P.{i}.down_proj.weight`` ``(O, I)`` (+ ``.weight_scale``/``.input_scale``) + """ + prefixes = _fused_experts_prefixes(state_dict) + if not prefixes: + return state_dict + + # Keys consumed (and thus removed) while emitting per-expert keys. + consumed: set[str] = set() + out: dict = {} + + for prefix in prefixes: + gated = f"{prefix}.gate_up_proj" in state_dict + first = "gate_up_proj" if gated else "up_proj" + first_w = state_dict[f"{prefix}.{first}"] + down_w = state_dict[f"{prefix}.down_proj"] + n_experts = first_w.shape[0] + + def _scales(proj_name): + ws = state_dict.get(f"{prefix}.{proj_name}_weight_scale") + ins = state_dict.get(f"{prefix}.{proj_name}_input_scale") + # NVFP4 second-level (per-tensor) weight scale; None for FP8. (E,) one scalar per expert. + ws2 = state_dict.get(f"{prefix}.{proj_name}_weight_scale_2") + return ws, ins, ws2 + + first_ws, first_in, first_ws2 = _scales(first) + down_ws, down_in, down_ws2 = _scales("down_proj") + + def _slice_weight_scale(ws, row_slice, fused_rows): + # per-tensor-per-expert scale -> scalar; per-output-channel -> slice the rows. + if ws is None: + return None + if ws.dim() <= 1: # (E,) one scalar per expert + return ws[expert_idx] + sub = ws[expert_idx] # (fused_rows,) or (fused_rows, ...) + return sub if row_slice is None else sub[row_slice] + + for expert_idx in range(n_experts): + if gated: + inter = first_w.shape[1] // 2 + # gate and up are halves of the fused gate_up: they share its single per-tensor + # weight_scale_2 (first_ws2), matching the in-model export + vLLM's W1/W3 fusion. + projections = [ + ( + "gate_proj", + first_w[expert_idx, :inter, :], + first_ws, + first_in, + first_ws2, + slice(0, inter), + ), + ( + "up_proj", + first_w[expert_idx, inter:, :], + first_ws, + first_in, + first_ws2, + slice(inter, None), + ), + ("down_proj", down_w[expert_idx], down_ws, down_in, down_ws2, None), + ] + else: + projections = [ + ("up_proj", first_w[expert_idx], first_ws, first_in, first_ws2, None), + ("down_proj", down_w[expert_idx], down_ws, down_in, down_ws2, None), + ] + for proj_name, weight, ws, ins, ws2, row_slice in projections: + base = f"{prefix}.{expert_idx}.{proj_name}" + # .clone() every emitted tensor so no two keys share storage. gate/up share the + # per-tensor scale object (and slices alias the fused parent); without cloning, + # safetensors/save_pretrained shared-tensor dedup silently drops the duplicate + # (e.g. up_proj.weight_scale/input_scale would go missing). + out[f"{base}.weight"] = weight.detach().clone().contiguous() + sliced_ws = _slice_weight_scale(ws, row_slice, weight.shape[0]) + if sliced_ws is not None: + out[f"{base}.weight_scale"] = sliced_ws.detach().clone().contiguous() + # weight_scale_2 is a per-expert scalar (shared by gate/up): _slice_weight_scale's + # dim<=1 branch returns ws2[expert_idx] regardless of row_slice. NVFP4 only; None for FP8. + sliced_ws2 = _slice_weight_scale(ws2, row_slice, weight.shape[0]) + if sliced_ws2 is not None: + out[f"{base}.weight_scale_2"] = sliced_ws2.detach().clone().contiguous() + if ins is not None: + out[f"{base}.input_scale"] = ins.detach().clone() + + # Mark the fused tensors + their scale buffers consumed. + for proj in (first, "down_proj"): + for suffix in ("", "_weight_scale", "_weight_scale_2", "_input_scale"): + consumed.add(f"{prefix}.{proj}{suffix}") + + for key, tensor in state_dict.items(): + if key not in consumed: + out[key] = tensor + return out + + + +def _dtensor_dim0_offset(dt) -> int: + """Global dim-0 start index of this rank's local shard of a ``Shard(0)`` DTensor.""" + from torch.distributed.tensor import Shard + + mesh, placements = dt.device_mesh, dt.placements + try: + from torch.distributed.tensor._utils import compute_local_shape_and_global_offset + + _, global_offset = compute_local_shape_and_global_offset(tuple(dt.shape), mesh, placements) + return int(global_offset[0]) + except Exception: + n = dt.shape[0] + for mdim, p in enumerate(placements): + if isinstance(p, Shard) and p.dim == 0: + 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 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`` keeps the legacy per-EP-group behavior. + """ + from torch.distributed.tensor import DTensor + + 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] if fw is not None else 0) + ) + group_size = int(fw.shape[0]) if fw is not None else 0 + # Add the EP block base only when the tensor is a per-EP-group slice (size < total_experts) or + # total is unknown (legacy). 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. But under EP the keep-fused fold emits the scale as a PLAIN buffer + # spanning the whole EP group (not dp-sharded like the weight param); 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 + + From 43a26fcf4bac255a35431409d22c433e7401e001 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Wed, 29 Jul 2026 21:18:13 -0700 Subject: [PATCH 02/10] feat(export): route FSDP2 export through the no-gather writer Replace the FSDP2 gather-to-rank-0 export with the per-rank DCP write for a distributed run. _export_transformers_checkpoint gains defer_distributed_fsdp2_write (default False; only export_hf_checkpoint sets it): when the model is FSDP2 under a live process group it does the quant processing but returns a None state dict instead of gathering the full model to CPU on rank 0. export_hf_checkpoint then has all ranks call distributed_save_hf_checkpoint (per-rank shard write, no host-RAM gather); rank 0 writes config.json + generation_config (_write_base_config) and folds the deployment quant/sparse config into config.json -- mirroring the gather-path tail. hf_spec_export's callers keep the gather behavior. Known open items (flagged in-code): the distributed path does not run revert_weight_conversion_quant_aware (non-expert transformers>=5 key renames), and extra_state_dict is not merged into the distributed write. Signed-off-by: Shengliang Xu (cherry picked from commit 29383f8b7613cca66e9677ef52e512d8d952af38) --- modelopt/torch/export/unified_export_hf.py | 88 +++++++++++++++++++++- 1 file changed, 85 insertions(+), 3 deletions(-) diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 77429b1cfaf..19380b19f27 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -929,8 +929,9 @@ def _export_transformers_checkpoint( model: nn.Module, dtype: torch.dtype | None = None, is_modelopt_qlora: bool = False, + defer_distributed_fsdp2_write: bool = False, **kwargs, -) -> tuple[dict[str, Any], dict[str, Any]]: +) -> tuple[dict[str, Any] | None, dict[str, Any]]: """Exports the torch model to the packed checkpoint with original HF naming. The packed checkpoint will be consumed by the TensorRT-LLM unified converter. @@ -998,8 +999,21 @@ def _export_transformers_checkpoint( _process_quantized_modules(model, dtype, is_modelopt_qlora) _reconstruct_fused_moe_linear(model) + if ( + defer_distributed_fsdp2_write + and is_fsdp2_model(model) + and torch.distributed.is_available() + and torch.distributed.is_initialized() + ): + # No-gather distributed FSDP2 export: leave the model SHARDED and signal the caller + # (export_hf_checkpoint) to write it per-rank via distributed_save_hf_checkpoint -- the + # kv-cache/postprocess fold runs inside that writer. Returning a None state dict avoids the + # full-model rank-0 host-RAM gather below (which does not scale to 100s of GB). Gated on the + # flag so other callers (e.g. hf_spec_export) keep the gather behavior. + return None, quant_config + if is_fsdp2_model(model): - # FSDP2: gather the full (unsharded) state_dict to CPU on rank 0. + # FSDP2 without the deferred-write opt-in (or no live process group): gather to CPU on rank 0. quantized_state_dict = get_model_state_dict( model, options=StateDictOptions(full_state_dict=True, cpu_offload=True), @@ -1445,6 +1459,22 @@ def _sanitize_generation_config_for_save(model: torch.nn.Module) -> None: gc.do_sample = True +def _write_base_config(model: nn.Module, export_dir: "Path | str") -> None: + """Write the base config.json + generation_config for the distributed (no-gather) export path. + + ``distributed_save_hf_checkpoint`` writes only the weight shards (no config), so on that path rank 0 + calls this to emit what ``save_pretrained`` would have written on the gather path. The weight half of + ``save_pretrained`` must NOT run here -- the weights are already on disk and the model may still be + sharded -- so only the config files are written. + """ + model.config.save_pretrained(export_dir) + if model.can_generate() and getattr(model, "generation_config", None) is not None: + try: + model.generation_config.save_pretrained(export_dir) + except Exception as gen_err: + warnings.warn(f"Could not save generation_config: {gen_err}") + + def export_speculative_decoding( model: torch.nn.Module, dtype: torch.dtype | None = None, @@ -1583,12 +1613,64 @@ def export_hf_checkpoint( _write_hf_export_config(model, hf_quant_config, export_dir) return - post_state_dict, hf_quant_config = _export_transformers_checkpoint(model, dtype, **kwargs) + post_state_dict, hf_quant_config = _export_transformers_checkpoint( + model, dtype, defer_distributed_fsdp2_write=True, **kwargs + ) # Remove hf_quantizer from model so post_state_dict can be exported. if getattr(model, "hf_quantizer", None) is not None: model.hf_quantizer = None + if post_state_dict is None: + # FSDP2 no-gather distributed export (signalled by _export_transformers_checkpoint returning + # a None state dict). Every rank writes its own weight shards via torch DCP -- collective, + # so this MUST run on all ranks, before the rank-0-only config work -- with no full-model + # rank-0 host-RAM gather. Rank 0 then writes config.json + the deployment quant config, + # mirroring the gather path's tail below. + # + # NOTE: unlike the gather path, this does NOT run revert_weight_conversion_quant_aware on + # the weights. The fused->per-expert un-fusing is handled inside the writer's expert + # split; any *other* transformers>=5 key renames are not reverted here. Validate on a model + # that triggers those renames before relying on it. extra_state_dict is likewise not merged + # into the distributed write yet. + from .distribute import distributed_save_hf_checkpoint + + kv_cache_format = (hf_quant_config or {}).get("quantization", {}).get( + "kv_cache_quant_algo" + ) + distributed_save_hf_checkpoint( + model, + export_dir, + maxbound=448, + kv_cache_format=kv_cache_format, + max_shard_size=max_shard_size, + ) + if not (is_distributed and torch.distributed.get_rank() != 0): + _write_base_config(model, export_dir) + quantization_details = (hf_quant_config or {}).get("quantization", {}) + is_quantized_export = ( + quantization_details.get("quant_algo") is not None + or quantization_details.get("kv_cache_quant_algo") is not None + ) + folded_quant_config = None + if is_quantized_export: + with open(f"{export_dir}/hf_quant_config.json", "w") as file: + json.dump(hf_quant_config, file, indent=4) + folded_quant_config = convert_hf_quant_config_format(hf_quant_config) + original_config = f"{export_dir}/config.json" + with open(original_config) as file: + config_data = json.load(file) + sanitize_hf_config_for_deployment(config_data, model) + if folded_quant_config is not None: + config_data["quantization_config"] = folded_quant_config + if export_sparse_attention_config is not None: + sparse_attn_config = export_sparse_attention_config(model) + if sparse_attn_config is not None: + config_data["sparse_attention_config"] = sparse_attn_config + with open(original_config, "w") as file: + json.dump(config_data, file, indent=4) + return + export_state_dict = {**post_state_dict, **(extra_state_dict or {})} # transformers may have applied a load-time conversion_mapping (fused gate_up_proj, From 4134b9f0052cb2f197a8d6f2e5628951de37ea48 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 30 Jul 2026 12:21:19 -0700 Subject: [PATCH 03/10] fix(export): shard-local fused-experts fold for FSDP2 no-gather export The no-gather distributed export ran the in-model _export_fused_experts, which under FSDP2 all-gathers the Shard(0) fused expert weight (via the export handler's unshard) and then materializes full per-expert weights for ALL experts on EVERY rank -- de-sharding the whole model. Fine at 30B (fits) but OOMs large MoE: at 235B the write-point state dict was ~157GB/rank (~113GB of replicated per-expert weights), and the footprint did not shrink with more nodes (32-GPU OOM'd at the same point as 16-GPU). Add _export_fused_experts_keep_fused: fold the per-expert quantizers into the FUSED 3-D weight in place. On a Shard(0) DTensor it quantizes ONLY this rank's local experts (to_local() + the DTensor dim-0 global offset for the matching per-expert weight quantizer) and rewraps the result Shard(0), so the model stays sharded and DCP writes per-rank shards; the per-expert split then happens on the sharded fused weight (split_fused_experts_state_dict). On a plain tensor (single process) it quantizes all experts -- byte-identical to _export_fused_experts. Dispatched via a keep_fused_experts flag (ExportContext -> the fused-experts export handler), set only when deferring to the no-gather writer; the single-process / gather paths keep the split-in-place behavior unchanged. Validated on hecate (NVFP4, calib 1): Qwen3-30B-A3B export drops 262s -> 62s with a byte-identical checkpoint (74401 keys, 128 experts, 17GB); Qwen3-235B-A22B now exports in 343s (write-point base 157GB -> 22GB) where it previously OOM'd on both 16 and 32 GPUs. Signed-off-by: Shengliang Xu (cherry picked from commit 234024a20d3eba09be4e49d5c68eb878d12a83ea) --- modelopt/torch/export/hf_export_handlers.py | 12 +- modelopt/torch/export/moe_utils.py | 135 ++++++++++++++++++++ modelopt/torch/export/registry.py | 4 + modelopt/torch/export/unified_export_hf.py | 22 +++- 4 files changed, 169 insertions(+), 4 deletions(-) diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py index 21a8a3fe246..5468108a167 100644 --- a/modelopt/torch/export/hf_export_handlers.py +++ b/modelopt/torch/export/hf_export_handlers.py @@ -24,7 +24,7 @@ from .layer_utils import get_expert_linear_names, is_quantlinear, set_expert_quantizer_amax from .model_config import QUANTIZATION_NONE -from .moe_utils import _export_fused_experts +from .moe_utils import _export_fused_experts, _export_fused_experts_keep_fused from .quant_utils import get_quantization_format from .registry import ExportContext, ExportModuleRegistry, PrepareMoEInputsRegistry @@ -132,9 +132,17 @@ def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContex Tied experts are packed independently and their duplicate keys are dropped by name in postprocess_state_dict; no per-module dedup cache is used. + + Distributed no-gather export (``ctx.keep_fused_experts``) folds the quantizers into the fused + weight IN PLACE, quantizing only this rank's local experts and keeping the result ``Shard(0)`` -- + the per-expert split happens later on the sharded weight (``split_fused_experts_state_dict``). + Otherwise (single process, or gather export) it splits into full per-expert submodules here. """ with fsdp2_aware_weight_update(ctx.model, module, reshard=False): - _export_fused_experts(module, ctx.dtype) + if ctx.keep_fused_experts: + _export_fused_experts_keep_fused(module, ctx.dtype) + else: + _export_fused_experts(module, ctx.dtype) @ExportModuleRegistry.register(predicate=is_quantlinear) diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py index bf844ee8eee..d3282468763 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -45,6 +45,141 @@ def _delete_fused_moe_source_attrs(module: nn.Module) -> None: delattr(module, attr) +def _export_fused_experts_keep_fused(module: nn.Module, dtype: torch.dtype) -> None: + """Fold per-expert quantizers into the FUSED 3-D expert weights *in place*, staying sharded. + + Distributed (DCP) counterpart to :func:`_export_fused_experts`. That function slices the fused + weight into full per-expert tensors -- under FSDP2 it must first all-gather the ``Shard(0)`` + fused weight (via the export handler's ``unshard``) and then materializes every expert on every + rank, putting the whole model on each GPU (OOMs at 235B/480B). This keeps each projection a + single fused param: on a ``Shard(0)`` DTensor it quantizes ONLY this rank's local experts + (``to_local()`` + the DTensor's dim-0 global offset for the matching per-expert weight quantizer) + and rewraps the result ``Shard(0)``, so the model stays sharded and DCP writes per-rank shards. + On a plain tensor (single process) it quantizes all ``num_experts`` -- identical to the in-model + path. Quantizing the fused ``[i]`` slice with its per-expert quantizer is byte-identical to + the split-then-quantize, so ``_export_fused_experts_keep_fused`` + :func:`split_fused_experts_state_dict` + reproduce :func:`_export_fused_experts`'s output exactly. Registers, per projection ``P``: + ``module.

`` (quantized fused weight), ``module.

_weight_scale`` (stacked per-expert scales), + ``module.

_weight_scale_2`` (NVFP4 per-expert scalar), ``module.

_input_scale`` (shared). + """ + from modelopt.torch.export.unified_export_hf import _export_quantized_weight + + try: + from torch.distributed.tensor import DTensor, Replicate, Shard + except Exception: # pragma: no cover - non-distributed install + DTensor = None + try: + from torch.distributed.tensor._utils import compute_local_shape_and_global_offset + except Exception: # pragma: no cover - older torch: fall back to even-split offset + compute_local_shape_and_global_offset = None + + first = getattr(module, "_first_proj_attr", "gate_up_proj") + n = module.num_experts + for proj in (first, "down_proj"): + weight = getattr(module, proj) + weight_quantizers = getattr(module, f"{proj}_weight_quantizers") + input_quantizer = getattr(module, f"{proj}_input_quantizer", None) + + # Under FSDP2 the fused 3-D weight is a DTensor sharded on dim 0 (the expert axis). Quantize + # ONLY this rank's local experts and rewrap the result sharded identically -- indexing the + # DTensor (weight[i]) would all-gather the whole fused weight onto every rank, de-sharding + # the model and OOMing the GPU at scale. Plain tensor (single process): all experts. + wdata = weight.data + is_dt = DTensor is not None and isinstance(wdata, DTensor) + if is_dt: + mesh, placements = wdata.device_mesh, wdata.placements + local_w = wdata.to_local() + offset = None + if compute_local_shape_and_global_offset is not None: + try: + _, global_offset = compute_local_shape_and_global_offset( + tuple(wdata.shape), mesh, placements + ) + offset = int(global_offset[0]) + except Exception: + offset = None + if offset is None: # even-split fallback (dim-0 Shard on a 1-D mesh) + offset, nshards = 0, 1 + for mdim, p in enumerate(placements): + if isinstance(p, Shard) and p.dim == 0: + nshards = mesh.size(mdim) + offset = mesh.get_local_rank(mdim) * ((n + nshards - 1) // nshards) + break + scale_placements = [ + p if (isinstance(p, Shard) and p.dim == 0) else Replicate() for p in placements + ] + local_indices = list(range(offset, offset + local_w.shape[0])) + else: + local_w = wdata + local_indices = list(range(n)) + + fp8_weights: list[torch.Tensor] = [] + weight_scales: list[torch.Tensor] = [] + weight_scale_2s: list[torch.Tensor] = [] + input_scale: torch.Tensor | None = None + for local_idx, global_idx in enumerate(local_indices): + w_quantizer = weight_quantizers[global_idx] + w_slice = local_w[local_idx] + # Uncalibrated-expert fallback: derive amax from this expert's local weight slice + # (matches _export_fused_experts), so a never-routed expert still exports sane scales. + if ( + hasattr(w_quantizer, "is_enabled") + and w_quantizer.is_enabled + and ( + not hasattr(w_quantizer, "_amax") + or w_quantizer._amax is None + or torch.all(w_quantizer._amax == 0) + ) + ): + w_quantizer.amax = w_slice.abs().amax().to(torch.float32) + warnings.warn( + f"Expert {global_idx} {proj} weight quantizer was not calibrated (amax missing " + f"or zero); using weight-derived amax. Increase calibration size to activate all " + f"experts.", + stacklevel=2, + ) + wrapper = nn.Module() + wrapper.weight = nn.Parameter(w_slice.contiguous(), requires_grad=False) + wrapper.weight_quantizer = w_quantizer + if input_quantizer is not None: + wrapper.input_quantizer = input_quantizer + _export_quantized_weight(wrapper, dtype) + fp8_weights.append(wrapper.weight.data) + weight_scales.append(wrapper.weight_scale) + # NVFP4 carries a SECOND per-tensor weight scale (weight_scale_2) that dequantizes the + # per-block weight_scale; FP8 has none. Keep it so the fused buffer (and the per-expert + # split) stay dequantizable. + if hasattr(wrapper, "weight_scale_2"): + weight_scale_2s.append(wrapper.weight_scale_2) + if hasattr(wrapper, "input_scale"): + input_scale = wrapper.input_scale + local_fp8 = torch.stack(fp8_weights, dim=0) + local_ws = torch.stack(weight_scales, dim=0) + if is_dt: + local_fp8 = DTensor.from_local(local_fp8, mesh, placements, run_check=False) + local_ws = DTensor.from_local(local_ws, mesh, scale_placements, run_check=False) + setattr(module, proj, nn.Parameter(local_fp8, requires_grad=False)) + module.register_buffer(f"{proj}_weight_scale", local_ws) + if weight_scale_2s: + # Per-expert scalar (E,) -- shards on the expert axis exactly like the weight scale. + local_ws2 = torch.stack(weight_scale_2s, dim=0) + if is_dt: + local_ws2 = DTensor.from_local(local_ws2, mesh, scale_placements, run_check=False) + module.register_buffer(f"{proj}_weight_scale_2", local_ws2) + if input_scale is not None: + module.register_buffer(f"{proj}_input_scale", input_scale) + + # Drop the quantizer modules -- their info now lives in the fused weight + scale buffers. + for attr in ( + f"{first}_weight_quantizers", + f"{first}_input_quantizer", + "down_proj_weight_quantizers", + "down_proj_input_quantizer", + ): + if hasattr(module, attr): + delattr(module, attr) + + def _export_fused_experts( module: nn.Module, dtype: torch.dtype, diff --git a/modelopt/torch/export/registry.py b/modelopt/torch/export/registry.py index 5af2a8c2c0e..ed1c9fda06f 100644 --- a/modelopt/torch/export/registry.py +++ b/modelopt/torch/export/registry.py @@ -54,6 +54,10 @@ class ExportContext: model: nn.Module dtype: torch.dtype is_modelopt_qlora: bool = False + # Distributed (no-gather) export: keep MoE experts FUSED + sharded instead of splitting them + # into full per-expert tensors, so each rank materializes only its local experts (the split + # happens post-consolidation on the sharded fused weight). See _export_fused_experts_keep_fused. + keep_fused_experts: bool = False ExportHandler = Callable[[str, nn.Module, ExportContext], None] diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 19380b19f27..71eaefc1ef1 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -895,6 +895,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. @@ -908,7 +909,12 @@ def _process_quantized_modules( If True, modules with base_layer attribute are skipped. """ # No per-module dedup cache: tied duplicates are dropped by name in postprocess_state_dict. - ctx = ExportContext(model=model, dtype=dtype, is_modelopt_qlora=is_modelopt_qlora) + ctx = ExportContext( + model=model, + dtype=dtype, + is_modelopt_qlora=is_modelopt_qlora, + keep_fused_experts=keep_fused_experts, + ) fsdp_module_to_reshard = None for name, sub_module in model.named_modules(): @@ -993,10 +999,22 @@ def _export_transformers_checkpoint( f"{synced_input} tied module group(s)" ) + # No-gather distributed FSDP2 export keeps MoE experts FUSED + sharded during the quantize fold + # (each rank materializes only its LOCAL experts instead of all N on every rank -- the split + # happens later on the sharded weight in the writer). Same condition as the deferred write below. + keep_fused_experts = ( + defer_distributed_fsdp2_write + and is_fsdp2_model(model) + and torch.distributed.is_available() + and torch.distributed.is_initialized() + ) + # Process all quantized modules and export weights from modelopt.torch.quantization.plugins.huggingface 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) if ( From 8da50e483867fa01f02903a0036a7e77436346a4 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Thu, 30 Jul 2026 12:21:19 -0700 Subject: [PATCH 04/10] perf(export): re-materialize the no-gather state dict in place distributed_save_hf_checkpoint rebuilt the sharded state dict as a dict comprehension `{k: _to_normal(v) for ...}`, which keeps the ENTIRE original dict alive while building a full second copy -- a transient 2x of the write-point state dict. Overwrite each key in place instead, so each old value is freed as it is replaced and peak stays ~1x. Same for the per-expert local_sd detach pass. Signed-off-by: Shengliang Xu (cherry picked from commit 3147fe86523cbdd98b107c37bdb54225e7bfa0af) --- modelopt/torch/export/distribute.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/export/distribute.py b/modelopt/torch/export/distribute.py index 4350c62b325..bab829d555a 100644 --- a/modelopt/torch/export/distribute.py +++ b/modelopt/torch/export/distribute.py @@ -470,7 +470,14 @@ def _to_normal(v): return _new return v - sharded_sd = {k: _to_normal(v) for k, v in sharded_sd.items()} + # 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) @@ -513,7 +520,9 @@ def _is_expert_key(key: str) -> bool: ) ) local_sd = postprocess_state_dict(local_sd, maxbound, kv_cache_format, is_modelopt_qlora) - local_sd = {k: v.detach().contiguous() for k, v in local_sd.items()} + # 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() # (2) Dense / non-expert: keep sharded -- FSDP2 leaves DTensors (dim-0 shard); EP leaves plain # replicated tensors. Do NOT gather to rank 0: DCP writes each DTensor's shards per-rank (parallel, From 4770881bb2ceeb7733809f8ab991ec34fae988ec Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 21 Aug 2026 11:22:10 -0700 Subject: [PATCH 05/10] fix(export): write model.safetensors.index.json in the no-gather FSDP2 writer The no-gather distributed export left the checkpoint without a weight index. distribute.py delegates the index to torch's consolidate_safetensors_files_on_every_rank ("rank 0 writes only the index"), but that helper does not emit model.safetensors.index.json -- on torch 2.11.0a0+eb65b36914.nv26.02 it writes only the merged shards. Without the index transformers / vLLM cannot map keys to shards, so every multi-shard export was silently unloadable: a Qwen3-30B-A3B FP8 run produced 4 correct safetensors, no index, and no error anywhere in the log. Build the index on rank 0 after consolidation. No extra collective is needed -- fqn_to_index_mapping and all_sizes are already unioned across ranks by the all_gather_object above and are identical on every rank -- so rank 0 maps each FQN to the file the writer placed it in and sums the global byte sizes. This also gives n_files (computed and unused until now) its purpose. The shard names must match what HuggingFaceStorageWriter derived from the same mapping: model--of-.safetensors, 5-digit zero-padded. The ep-dp branch's _write_even_shards_no_consolidation does this with two all_gather_object calls because there each rank knows only its own shards; here the mapping is already global, so the gather is unnecessary. Validated on lyris (Qwen3-30B-A3B, FP8 weights + FP8 KV, 2 nodes x 4 GPU, pure FSDP2, job 2755390): the index covers all 56211 keys with 0 missing, 0 extra and 0 file mismatches when checked against the actual safetensors headers; total_size 31167395712 sits one header-size below the on-disk total, matching the gather path's convention. The weight shards are byte-identical in size to the pre-fix run, so the change adds the index without perturbing the write. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distribute.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/modelopt/torch/export/distribute.py b/modelopt/torch/export/distribute.py index bab829d555a..ef37a044f58 100644 --- a/modelopt/torch/export/distribute.py +++ b/modelopt/torch/export/distribute.py @@ -581,7 +581,25 @@ def _is_expert_key(key: str) -> bool: num_threads=8, ) - # (5) Drop the intermediate per-rank sharded/ dir. + # (5) Write the HF weight index. torch's consolidation helper does NOT emit + # model.safetensors.index.json, and without it transformers / vLLM cannot map keys to shards, so a + # multi-shard checkpoint is unloadable. No extra collective is needed: fqn_to_index_mapping and + # all_sizes were already unioned across ranks above (all_gather_object) and are identical on every + # rank, so rank 0 builds the index locally. The file names must match what HuggingFaceStorageWriter + # emitted from the same mapping -- model--of-.safetensors, 5-digit zero-padded. + if rank == 0: + weight_map = { + fqn: f"model-{idx:05d}-of-{n_files:05d}.safetensors" + for fqn, idx in fqn_to_index_mapping.items() + } + index = { + "metadata": {"total_size": sum(all_sizes.values())}, + "weight_map": weight_map, + } + with open(export_dir / "model.safetensors.index.json", "w") as f: + json.dump(index, f, indent=2) + + # (6) Drop the intermediate per-rank sharded/ dir. if rank == 0: shutil.rmtree(sharded_dir, ignore_errors=True) dist.barrier() From 79a99c05d561d304f0a84ac898d00c6fee3ab95d Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 21 Aug 2026 13:19:37 -0700 Subject: [PATCH 06/10] test(export): structural checks on the no-gather FSDP2 distributed checkpoint The distributed writer rebuilds a checkpoint from per-rank DCP shards, so the ways it can go wrong are structural, not numerical: a key silently dropped, a rank-local shard written where the full tensor belongs, a sidecar or the weight index never emitted. All of those produce a checkpoint that looks fine -- no error, plausible file sizes -- but is wrong or unloadable. Nothing covered that. Add a GPU test that PTQs a tiny Qwen3 (dense and MoE) under FSDP2 across the whole world with the default FP8 and NVFP4 configs, then asserts against the safetensors headers on disk: 1. every source parameter survives the round trip, and nothing extra is left behind -- checked both directions, so a fused expert weight left next to the per-expert weights it was split into is a failure, not a silent 2.9x checkpoint; 2. each quantized weight carries the scales its format needs (weight_scale + input_scale, plus weight_scale_2 for NVFP4); 3. every tensor is at its full unsharded shape, with NVFP4's uint8 packing accounted for; 4. the files the exporter writes are present (config, generation config, quant config, and the weight index whenever the layout is sharded, cross-checked against the tensors actually on disk), and the sidecars it must leave alone -- tokenizer/vocab/chat template -- survived. The test asserts is_fsdp2_model() inside the workers: the no-gather path is selected by exactly that predicate, and without the check the export would quietly fall back to the rank-0 gather and the test would pass while never touching the code under test. Model dimensions are set explicitly rather than taking the tiny-Qwen3 defaults. Those give head_dim = hidden_size / num_heads = 32/16 = 2, and a 2-element q_norm split over a 4-rank world leaves ranks holding empty chunks, which the DCP planner rejects outright ("invalid fill tensor-volume") and which kills the whole worker pool. Every sharded axis is now comfortably divisible and last dims are multiples of 16 so NVFP4 blocks are well defined. Verified on lyris (4x GB200): 4 passed in ~32s on this branch. Against main (a2fbac7ba) the two MoE cases fail on check 1 with the leftover fused experts -- the toy-model form of the 96 stale BF16 tensors that made a real Qwen3-30B-A3B FP8 export 89.2 GB instead of 31.2 GB -- while both dense cases pass, so the test discriminates rather than failing blanket. Signed-off-by: Shengliang Xu --- .../export/test_distributed_hf_export.py | 258 ++++++++++++++++++ 1 file changed, 258 insertions(+) create mode 100644 tests/gpu/torch/export/test_distributed_hf_export.py 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..adb12bbe305 --- /dev/null +++ b/tests/gpu/torch/export/test_distributed_hf_export.py @@ -0,0 +1,258 @@ +# 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. + +``export_hf_checkpoint`` routes an FSDP2-sharded model under ``torch.distributed`` to +``distributed_save_hf_checkpoint``: 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 ones it must leave alone. The tokenizer/vocab/chat-template + files are written into the export dir by ``hf_ptq.py`` (``tokenizer.save_pretrained`` + + ``copy_custom_model_files``), not by ``export_hf_checkpoint``, so what is asserted of them here + is survival: the writer creates and removes a ``sharded/`` staging dir inside ``export_dir``, and + must not take the sidecars with it. + +Run with >=2 GPUs so the expert axis and the FSDP2 shard axis are both actually split. +""" + +import json +import shutil +from functools import partial +from pathlib import Path + +import pytest +import torch +from _test_utils.torch.transformers_models import ( + create_tiny_qwen3_dir, + create_tiny_qwen3_moe_dir, +) + +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 _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 + + 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 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. + 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() + + 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( + ("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" + + # Seed the export dir with the source's non-weight files, standing in for the + # tokenizer/sidecar save that hf_ptq.py performs around the export call. They must still be + # there afterwards: the writer stages per-rank shards in export_dir/sharded/ and deletes that + # tree when it consolidates, which is exactly the operation that could take them out. + 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" + } + export_dir.mkdir(parents=True, exist_ok=True) + for name in sidecars: + shutil.copy2(src_dir / name, export_dir / name) + + 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, ...) survived the export; only the weight files are rewritten. + lost = sorted(sidecars - exported_files) + assert not lost, f"sidecar file(s) lost: {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 From 78dcfad85dcd430fb7ba0ddb3747125302fb7973 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 21 Aug 2026 13:29:46 -0700 Subject: [PATCH 07/10] docs(changelog): note the no-gather FSDP2 distributed HF export New feature entry for the 0.47 section: export_hf_checkpoint now routes an FSDP2 model under torch.distributed through the per-rank DCP writer instead of the rank-0 host-RAM gather. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6687ebd31ea..cafa167db73 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -22,6 +22,7 @@ Changelog - Add ``modelopt.torch.utils.mlflow.MlflowRunLogger`` for recording a script run on an MLflow tracking server: the invocation, the ModelOpt version, the run log (captured by teeing ``stdout``/``stderr``) and any caller-supplied artifacts, with configuration as searchable params. ``mlflow`` is an optional dependency, imported only when tracking is enabled. - Add ``--mlflow `` to ``examples/hf_ptq/hf_ptq.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too). A tracked run records the invocation, the resolved recipe (``$import``\ s expanded), the run log and the quantization summaries, with every command-line argument as a searchable param; failed runs are recorded with their traceback. The experiment defaults to ``$USER/hf_ptq/-`` and can be overridden with ``--mlflow_experiment`` / ``--mlflow_run_name``. - Add ``--mlflow `` to ``examples/vllm_serve/vllm_serve_fakequant.py`` (MLflow's own ``MLFLOW_TRACKING_URI`` is honoured too), so a fake-quant serve records what it quantized and an evaluation of that endpoint can be traced back to a recipe. A tracked run uploads the launcher command, the resolved ``RECIPE_PATH`` (or the merged ``QUANT_CFG``/``KV_QUANT_CFG`` when presets are used), the worker log and the quantizer summary; the experiment defaults to ``$USER/vllm_serve_fakequant/-`` and can be overridden with ``--mlflow-experiment`` / ``--mlflow-run-name``. +- Export an FSDP2-sharded model to a HuggingFace checkpoint without gathering it: ``export_hf_checkpoint`` now detects an FSDP2 model under ``torch.distributed`` and has every rank write its own weight shards through torch DCP (``distributed_save_hf_checkpoint``), consolidating them in parallel, instead of gathering the full state dict to host RAM on rank 0. MoE experts stay fused and sharded through the quantizer fold, so each rank materializes only its local experts rather than all of them. This removes the rank-0 host-RAM ceiling that made large-MoE export OOM, and cuts export time on a 2-node Qwen3-30B-A3B FP8 run from 301s to 81s (94% of the old cost was the gather). **Backward Breaking Changes** From e58f8c8458302189d589a5e1d42dcbd78353c431 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 21 Aug 2026 13:39:58 -0700 Subject: [PATCH 08/10] docs(changelog): record the stale fused-expert weights in FSDP2 MoE export Bug Fixes entry for 0.47: the FSDP2 HF export split each fused expert weight into per-expert quantized weights but left the original unquantized fused tensors in the checkpoint, so every expert was written twice. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index cafa167db73..3aeab99561b 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -44,6 +44,7 @@ Changelog - Update HuggingFace checkpoint export to use name-based tied-weight deduplication instead of the previous address-based approach. The address-based deduplication could incorrectly drop an untied weight that happened to share memory with a tied one, producing an incomplete checkpoint (observed as a false positive on MiniMax-M2.7). - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``. - Polygraphy minimum dependency upgraded to ``0.53.4`` to solve ONNX AutoCast failures when marking optional graph outputs. +- Fix HuggingFace export of an FSDP2-sharded MoE leaving the fused expert weights in the checkpoint. The exporter splits each fused 3-D expert weight into per-expert quantized weights, but the original unquantized ``mlp.experts.gate_up_proj`` / ``mlp.experts.down_proj`` tensors were still written alongside them, so every expert appeared twice and a consumer keying off the fused names loaded unquantized weights. On Qwen3-30B-A3B FP8 this was 96 stale BF16 tensors and 58 GB of a 89 GB checkpoint, which is 31 GB once they are gone. 0.46 (2026-08-17) ^^^^^^^^^^^^^^^^^ From 0c1d301bb29fc6ab1b90bd3503e1bbc847abec04 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 21 Aug 2026 14:21:18 -0700 Subject: [PATCH 09/10] docs(changelog): the stale fused experts are unquantized weights, not duplicates Reword the 0.47 bug-fix entry: the leftover mlp.experts.gate_up_proj / down_proj tensors are the model's original BF16 pre-quantization parameters, carrying no scales, so a quantized export shipped a full unquantized copy of every expert rather than a redundant copy of the quantized data. Signed-off-by: Shengliang Xu --- CHANGELOG.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 3aeab99561b..4210b12228a 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -44,7 +44,7 @@ Changelog - Update HuggingFace checkpoint export to use name-based tied-weight deduplication instead of the previous address-based approach. The address-based deduplication could incorrectly drop an untied weight that happened to share memory with a tied one, producing an incomplete checkpoint (observed as a false positive on MiniMax-M2.7). - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``. - Polygraphy minimum dependency upgraded to ``0.53.4`` to solve ONNX AutoCast failures when marking optional graph outputs. -- Fix HuggingFace export of an FSDP2-sharded MoE leaving the fused expert weights in the checkpoint. The exporter splits each fused 3-D expert weight into per-expert quantized weights, but the original unquantized ``mlp.experts.gate_up_proj`` / ``mlp.experts.down_proj`` tensors were still written alongside them, so every expert appeared twice and a consumer keying off the fused names loaded unquantized weights. On Qwen3-30B-A3B FP8 this was 96 stale BF16 tensors and 58 GB of a 89 GB checkpoint, which is 31 GB once they are gone. +- Fix HuggingFace export of an FSDP2-sharded MoE writing the model's **unquantized** weights into the quantized checkpoint. The exporter splits each fused 3-D expert weight into per-expert quantized weights, but the original BF16 ``mlp.experts.gate_up_proj`` / ``mlp.experts.down_proj`` parameters were still written alongside them -- with no scales, and under the names ``transformers`` binds directly -- so a quantized export also carried a full pre-quantization copy of every expert. On Qwen3-30B-A3B FP8 the experts are ~29B of the 30.5B parameters, so this was 58 GB of an 89 GB checkpoint that should be 31 GB. 0.46 (2026-08-17) ^^^^^^^^^^^^^^^^^ From 51a8fde8d9258827ce3956aca33a2370369bc3e7 Mon Sep 17 00:00:00 2001 From: Shengliang Xu Date: Fri, 21 Aug 2026 15:21:02 -0700 Subject: [PATCH 10/10] fix(export): close three silent gaps in the no-gather FSDP2 export The no-gather path auto-engages for every FSDP2 export under torch.distributed, so anything the gather path does that it does not is a silent behaviour change for callers who never opted in. Review of #2228 found three; this closes them. extra_state_dict and save_modelopt_state were accepted and ignored. The gather path merges the former into the exported state dict and forwards the latter to save_pretrained; the distributed path writes the checkpoint itself and returns before both, so a caller passing either got a checkpoint quietly missing what they asked for. Reject them with NotImplementedError instead, mirroring how the streaming-offload path already declines save_modelopt_state. The check sits outside the export try/except so the failure surfaces as itself rather than behind the generic "Cannot export model to the model_config" warning, and is gated on `not _offloaded` so an offloaded model still reaches the streaming path. Tied-weight dedup was not applied. The gather path passes a TiedWeightMap into postprocess_state_dict; the writer called it positionally, leaving tied_map=None, so the name-based drop added in #2194 silently did not run for FSDP2 models. fully_shard splits a shared nn.Parameter into distinct per-module shards, so a declared tie reaches the writer as two independent DTensors and both were written. Build the map from the model inside the writer and pass it to both postprocess calls. The map is name-derived and identical on every rank, so the alias is dropped consistently, and postprocess already skips a group whose canonical is absent from this rank's slice rather than orphaning it. Tests: parametrized rejection cases for both options, and a tied-embedding export asserting the alias is gone and the canonical kept. Verified the tied test discriminates -- with tied_map removed from the two call sites it fails with "tied alias 'lm_head.weight' was written alongside its canonical", and passes with it. 7 passed on 4x GB200. Signed-off-by: Shengliang Xu --- modelopt/torch/export/distribute.py | 16 +++- modelopt/torch/export/unified_export_hf.py | 19 +++++ .../export/test_distributed_hf_export.py | 83 +++++++++++++++++++ 3 files changed, 116 insertions(+), 2 deletions(-) diff --git a/modelopt/torch/export/distribute.py b/modelopt/torch/export/distribute.py index ef37a044f58..1a7fbe2f9cd 100644 --- a/modelopt/torch/export/distribute.py +++ b/modelopt/torch/export/distribute.py @@ -413,8 +413,16 @@ 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 .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()}") @@ -519,7 +527,9 @@ def _is_expert_key(key: str) -> bool: 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) + 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() @@ -530,7 +540,9 @@ def _is_expert_key(key: str) -> bool: # 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) + dense_sd = postprocess_state_dict( + dense_sd, maxbound, kv_cache_format, is_modelopt_qlora, tied_map=tied_map + ) local_sd.update(dense_sd) del sharded_sd, dense_sd gc.collect() diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 71eaefc1ef1..acd5fc268b2 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1598,6 +1598,25 @@ 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 and returns before the gather path's extra_state_dict merge. Neither option + # can be honoured there, so reject explicitly rather than dropping them silently (the streaming + # path declines save_modelopt_state the same way). + if is_distributed and not _offloaded: + if extra_state_dict: + raise NotImplementedError( + "extra_state_dict is not supported by the no-gather FSDP2 distributed export: " + "each rank writes its own shards, so the extra tensors are never merged in. " + "Export without it, or outside torch.distributed to take the gather path." + ) + 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 diff --git a/tests/gpu/torch/export/test_distributed_hf_export.py b/tests/gpu/torch/export/test_distributed_hf_export.py index adb12bbe305..3b4526ee634 100644 --- a/tests/gpu/torch/export/test_distributed_hf_export.py +++ b/tests/gpu/torch/export/test_distributed_hf_export.py @@ -32,6 +32,10 @@ is survival: the writer creates and removes a ``sharded/`` staging dir inside ``export_dir``, and must not take the sidecars with it. +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. """ @@ -256,3 +260,82 @@ def test_fsdp2_distributed_export_is_complete(dist_workers, tmp_path, moe, quant 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() + export_hf_checkpoint(model, export_dir=export_dir, **export_kwargs) + + +@pytest.mark.timeout(600) +@pytest.mark.parametrize( + ("kwargs", "needle"), + [ + ({"extra_state_dict": {"extra.tensor": torch.zeros(1)}}, "extra_state_dict"), + ({"save_modelopt_state": True}, "save_modelopt_state"), + ], + ids=["extra_state_dict", "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. + + Both are handled by the gather path (``extra_state_dict`` is merged into the exported state + dict, ``save_modelopt_state`` is forwarded to ``save_pretrained``). The distributed path writes + the checkpoint itself and returns before either, so a caller passing them would otherwise get a + checkpoint quietly missing what they asked for. + """ + 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" + )