From 8d3936f0015e0800a382d4f7a68b9fa39f9e2139 Mon Sep 17 00:00:00 2001
From: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Date: Tue, 18 Aug 2026 05:58:42 +0000
Subject: [PATCH 1/2] rebased

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
---
 modelopt/torch/export/hf_export_handlers.py   |  49 +++-
 modelopt/torch/export/moe_utils.py            | 181 +++++++++++++++
 modelopt/torch/export/unified_export_hf.py    | 219 +++++++++++++++---
 .../torch/quantization/utils/core_utils.py    | 210 ++++++++++++++++-
 4 files changed, 620 insertions(+), 39 deletions(-)

diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py
index 21a8a3fe246..7a0bfe54408 100644
--- a/modelopt/torch/export/hf_export_handlers.py
+++ b/modelopt/torch/export/hf_export_handlers.py
@@ -17,14 +17,16 @@
 
 import collections.abc
 import warnings
+from contextlib import nullcontext
 
 import torch.nn as nn
 
-from modelopt.torch.quantization.utils import fsdp2_aware_weight_update
+from modelopt.torch.quantization.utils import fsdp2_shard_local_pack
+from modelopt.torch.utils.distributed import is_fsdp2_model
 
 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
 
@@ -36,6 +38,11 @@ def _has_fused_experts_quantizers(module: nn.Module) -> bool:
     return hasattr(module, f"{first_proj_attr}_weight_quantizers")
 
 
+def _use_shard_local(model: nn.Module) -> bool:
+    """Whether to use shard-local packing (FSDP2 only)."""
+    return is_fsdp2_model(model)
+
+
 def _export_weight(
     module: nn.Module,
     ctx: ExportContext,
@@ -130,20 +137,34 @@ def _export_moe_linear(name: str, module: nn.Module, ctx: ExportContext) -> None
 def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContext) -> None:
     """Split and quantize a fused-experts module with plural weight quantizers.
 
+    Under FSDP2 the fused weight is ``Shard(0)`` on the expert dim, so the destructive per-rank split
+    (``_export_fused_experts``) would drop non-owner experts from the gather. Instead pack this rank's
+    experts IN PLACE and keep them fused (``_export_fused_experts_keep_fused`` inside
+    ``fsdp2_shard_local_pack``); the gate/up split is deferred to write time
+    (``_split_fused_experts_state_dict`` in the gather). Non-FSDP is unchanged.
+
     Tied experts are packed independently and their duplicate keys are dropped by name
     in postprocess_state_dict; no per-module dedup cache is used.
     """
-    with fsdp2_aware_weight_update(ctx.model, module, reshard=False):
+    if _use_shard_local(ctx.model):
+        with fsdp2_shard_local_pack(ctx.model, module):
+            _export_fused_experts_keep_fused(module, ctx.dtype)
+    else:
         _export_fused_experts(module, ctx.dtype)
 
 
 @ExportModuleRegistry.register(predicate=is_quantlinear)
 def _export_quant_linear(name: str, module: nn.Module, ctx: ExportContext) -> None:
-    """Export a standard quantized linear layer."""
+    """Export a standard quantized linear layer.
+
+    ``fsdp2_shard_local_pack`` packs this rank's ``Shard(0)`` slice in place (no unshard) under FSDP2,
+    and is a no-op for non-FSDP models -- so the single-process path is unchanged.
+    """
     if get_quantization_format(module) == QUANTIZATION_NONE:
         return
+    cm = fsdp2_shard_local_pack(ctx.model, module) if _use_shard_local(ctx.model) else nullcontext()
     try:
-        with fsdp2_aware_weight_update(ctx.model, module, reshard=False):
+        with cm:
             _export_weight(module, ctx)
     except AssertionError as e:
         raise AssertionError(
@@ -173,8 +194,12 @@ def _export_quant_embedding(name: str, module: nn.Module, ctx: ExportContext) ->
             "The embedding will be exported as its fake-quantized float weight."
         )
         return
+    # The embedding lives in the root FSDP unit (reshard_after_forward=False -> its params may be
+    # unsharded at export); fsdp2_shard_local_pack reshards them to Shard(0) first so this rank packs
+    # only its vocab slice in place, no full unshard. Non-FSDP is unchanged.
+    cm = fsdp2_shard_local_pack(ctx.model, module) if _use_shard_local(ctx.model) else nullcontext()
     try:
-        with fsdp2_aware_weight_update(ctx.model, module, reshard=False):
+        with cm:
             _export_weight(module, ctx)
     except AssertionError as e:
         raise AssertionError(
@@ -184,7 +209,14 @@ def _export_quant_embedding(name: str, module: nn.Module, ctx: ExportContext) ->
 
 @ExportModuleRegistry.register("Llama4TextExperts", "GptOssExperts")
 def _export_bmm_experts(name: str, module: nn.Module, ctx: ExportContext) -> None:
-    """Export fused BMM-style expert weights and quantization metadata."""
+    """Export fused BMM-style expert weights and quantization metadata.
+
+    The fused ``gate_up_proj``/``down_proj`` are ``[E, ...]`` (``Shard(0)`` on E under FSDP2) and each
+    is packed WHOLE by ``_export_quantized_weight`` with its singular weight quantizer -- whose amax is
+    calibrated over all experts (a replicated buffer, i.e. global), so packing a single rank's
+    ``[E/world, ...]`` slice yields the same bytes. So under FSDP2 this uses ``fsdp2_shard_local_pack``
+    (pack this rank's slice in place, no unshard); non-FSDP is unchanged.
+    """
     if get_quantization_format(module) == QUANTIZATION_NONE:
         return
     # TODO: consolidate uncalibrated experts handling logic
@@ -196,6 +228,7 @@ def _export_bmm_experts(name: str, module: nn.Module, ctx: ExportContext) -> Non
         modules=module,
         quantizer_attrs=["gate_up_proj_input_quantizer", "down_proj_input_quantizer"],
     )
-    with fsdp2_aware_weight_update(ctx.model, module, reshard=False):
+    cm = fsdp2_shard_local_pack(ctx.model, module) if _use_shard_local(ctx.model) else nullcontext()
+    with cm:
         for weight_name in ["gate_up_proj", "down_proj"]:
             _export_weight(module, ctx, weight_name)
diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py
index 4ce60b192ee..bee0a130f76 100644
--- a/modelopt/torch/export/moe_utils.py
+++ b/modelopt/torch/export/moe_utils.py
@@ -228,6 +228,187 @@ def _export_fused_experts(
     _delete_fused_moe_source_attrs(module)
 
 
+def _export_fused_experts_keep_fused(module, dtype):
+    """Shard-local keep-fused NVFP4/FP8 pack for one rank's experts.
+
+    Pack THIS rank's experts IN PLACE and keep the weight FUSED (``Shard(0)`` on E),
+    instead of splitting into per-expert child modules.
+
+    Self-adapts to FSDP2 vs single-process via ``module._shard_local_start`` (set by
+    :func:`fsdp2_shard_local_pack`):
+
+    * FSDP2 shard-local: the context has already ``to_local``'d ``gate_up_proj``/``down_proj`` to this
+      rank's ``[E/world, ...]`` block and recorded ``start`` (this rank's global expert offset), so the
+      per-expert weight quantizers are indexed ``[start + i]`` (the quantizer ModuleList is global
+      length-E and replicated -- verified in ``spike_moe_indexing.py``).
+    * single-process: ``start == 0`` and the tensors are the full ``[E, ...]`` block.
+
+    Each local expert's fused first-projection ``[2I, H]`` is packed WHOLE by reusing the standard
+    ``_export_quantized_weight`` (so the per-tensor ``weight_scale_2`` is naturally shared by the gate
+    and up halves -- what vLLM's load-time fusion requires -- and per-block scales are row-local). The
+    packed per-expert results are stacked back into fused ``[local_n, ...]`` tensors + per-expert
+    scale buffers; the gate/up split is deferred to write time (``_split_fused_experts_state_dict``).
+    Non-destructive: the fused params survive, so :func:`fsdp2_shard_local_pack` re-registers them.
+    """
+    from modelopt.torch.export.unified_export_hf import _export_quantized_weight
+
+    first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj")
+    start = getattr(module, "_shard_local_start", {}).get(first_proj_attr, 0)
+    first_proj_wq = getattr(module, f"{first_proj_attr}_weight_quantizers")
+    down_wq = module.down_proj_weight_quantizers
+    first_proj_iq = getattr(module, f"{first_proj_attr}_input_quantizer")
+    down_iq = module.down_proj_input_quantizer
+
+    first_proj = getattr(
+        module, first_proj_attr
+    ).data  # local [local_n, 2I, H] (or full [E, 2I, H])
+    down = module.down_proj.data
+    local_n = first_proj.shape[0]
+
+    def _pack_one(weight_2d, w_quant_src, i_quant):
+        wrapper = nn.Module()
+        wrapper.weight = nn.Parameter(weight_2d.contiguous(), requires_grad=False)
+        # deepcopy so packing does not mutate the shared calibrated quantizer state.
+        wrapper.weight_quantizer = copy.deepcopy(w_quant_src)
+        wrapper.input_quantizer = i_quant
+        wq = wrapper.weight_quantizer
+        if getattr(wq, "is_enabled", False) and (
+            not hasattr(wq, "_amax") or wq._amax is None or torch.all(wq._amax == 0)
+        ):
+            # Uncalibrated expert (received no tokens): fall back to the weight's own amax.
+            wq.amax = weight_2d.abs().amax().to(torch.float32)
+        _export_quantized_weight(wrapper, dtype)
+        return (
+            wrapper.weight.data,
+            getattr(wrapper, "weight_scale", None),
+            getattr(wrapper, "weight_scale_2", None),
+            getattr(wrapper, "input_scale", None),
+        )
+
+    fp_w, fp_s, fp_s2 = [], [], []
+    dp_w, dp_s, dp_s2 = [], [], []
+    fp_input_scale = dp_input_scale = None
+    for i in range(local_n):
+        g = start + i
+        w, s, s2, isc = _pack_one(first_proj[i], first_proj_wq[g], first_proj_iq)
+        fp_w.append(w)
+        fp_s.append(s)
+        fp_s2.append(s2)
+        fp_input_scale = isc if isc is not None else fp_input_scale
+        w, s, s2, isc = _pack_one(down[i], down_wq[g], down_iq)
+        dp_w.append(w)
+        dp_s.append(s)
+        dp_s2.append(s2)
+        dp_input_scale = isc if isc is not None else dp_input_scale
+
+    def _register(attr, weights, scales, scale2s):
+        setattr(module, attr, nn.Parameter(torch.stack(weights), requires_grad=False))
+        if scales[0] is not None:
+            module.register_buffer(f"{attr}_weight_scale", torch.stack(scales))
+        if scale2s[0] is not None:
+            module.register_buffer(
+                f"{attr}_weight_scale_2", torch.stack([x.reshape(()) for x in scale2s])
+            )
+
+    _register(first_proj_attr, fp_w, fp_s, fp_s2)
+    _register("down_proj", dp_w, dp_s, dp_s2)
+    if fp_input_scale is not None:
+        module.register_buffer(f"{first_proj_attr}_input_scale", fp_input_scale)
+    if dp_input_scale is not None:
+        module.register_buffer("down_proj_input_scale", dp_input_scale)
+
+
+def _split_fused_experts_state_dict(state_dict, model):
+    """Split keep-fused expert tensors in a gathered ``state_dict`` into per-expert keys.
+
+    Converts the ``_export_fused_experts_keep_fused`` output (fused ``{name}.gate_up_proj [E,2I,H/2]``
+    + per-expert scales) into the same per-expert deployment keys ``_export_fused_experts`` emits
+    (``{name}.{e}.gate_proj.weight`` / ``up_proj`` / ``down_proj`` + scales). Byte-identical: packing
+    is row-independent, so slicing the packed fused tensor equals packing each half; gate/up share the
+    per-tensor ``weight_scale_2`` (packed together with one amax). Runs on the gathered (full) dict.
+    """
+    from modelopt.torch.quantization.plugins.huggingface import _get_fused_expert_intermediate_dim
+
+    for name, module in model.named_modules():
+        first_proj_attr = getattr(module, "_first_proj_attr", "gate_up_proj")
+        if not hasattr(module, f"{first_proj_attr}_weight_quantizers"):
+            continue
+        prefix = f"{name}." if name else ""
+        fp_w = state_dict.pop(prefix + first_proj_attr, None)
+        if fp_w is None:
+            continue  # not keep-fused here (e.g. already split) -> nothing to do
+        is_gated = getattr(module, "_is_gated", True)
+        n = module.num_experts
+        fp_s = state_dict.pop(prefix + f"{first_proj_attr}_weight_scale", None)
+        fp_s2 = state_dict.pop(prefix + f"{first_proj_attr}_weight_scale_2", None)
+        fp_in = state_dict.pop(prefix + f"{first_proj_attr}_input_scale", None)
+        dp_w = state_dict.pop(prefix + "down_proj", None)
+        dp_s = state_dict.pop(prefix + "down_proj_weight_scale", None)
+        dp_s2 = state_dict.pop(prefix + "down_proj_weight_scale_2", None)
+        dp_in = state_dict.pop(prefix + "down_proj_input_scale", None)
+
+        # Drop leftover fused quantizer buffers (keep-fused does not delete them, unlike the
+        # per-expert path which calls _delete_fused_moe_source_attrs).
+        for k in [
+            k for k in state_dict if k.startswith(prefix) and "_quantizer" in k[len(prefix) :]
+        ]:
+            state_dict.pop(k)
+
+        edim = _get_fused_expert_intermediate_dim(module) if is_gated else None
+
+        def _emit(e, proj, w, s, s2, insc):
+            # clone/contiguous so each per-expert/projection key is a DISTINCT tensor object. The
+            # shared scales (one input_scale across all experts; one weight_scale_2 across gate|up)
+            # would otherwise share a data_ptr and get collapsed by postprocess_state_dict's tied-
+            # weight dedup -- producing fewer keys than the per-expert path, which builds separate
+            # equal-valued objects. Cloning matches that format exactly (values are identical).
+            p = f"{prefix}{e}.{proj}."
+            state_dict[p + "weight"] = w.contiguous()
+            if s is not None:
+                state_dict[p + "weight_scale"] = s.contiguous()
+            if s2 is not None:
+                state_dict[p + "weight_scale_2"] = s2.clone()
+            if insc is not None:
+                state_dict[p + "input_scale"] = insc.clone()
+
+        for e in range(n):
+            if is_gated:
+                _emit(
+                    e,
+                    "gate_proj",
+                    fp_w[e, :edim],
+                    fp_s[e, :edim] if fp_s is not None else None,
+                    fp_s2[e] if fp_s2 is not None else None,
+                    fp_in,
+                )
+                _emit(
+                    e,
+                    "up_proj",
+                    fp_w[e, edim:],
+                    fp_s[e, edim:] if fp_s is not None else None,
+                    fp_s2[e] if fp_s2 is not None else None,
+                    fp_in,
+                )
+            else:
+                _emit(
+                    e,
+                    "up_proj",
+                    fp_w[e],
+                    fp_s[e] if fp_s is not None else None,
+                    fp_s2[e] if fp_s2 is not None else None,
+                    fp_in,
+                )
+            _emit(
+                e,
+                "down_proj",
+                dp_w[e],
+                dp_s[e] if dp_s is not None else None,
+                dp_s2[e] if dp_s2 is not None else None,
+                dp_in,
+            )
+    return state_dict
+
+
 def save_expert_token_count_table(model: nn.Module, output_dir: str | Path | None = None):
     """Collect expert_token_count from all quantized MoE layers and save as an HTML table.
 
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index 77429b1cfaf..b485dc55d74 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -51,7 +51,6 @@
 except ImportError:
     HAS_DIFFUSERS = False
 
-from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict
 from torch.distributed.fsdp import FSDPModule
 
 from modelopt.torch.quantization import set_quantizer_by_cfg_context
@@ -61,6 +60,7 @@
 from modelopt.torch.quantization.qtensor.nvfp4_tensor import _cast_per_block_scale_to_fp8
 from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, quantizer_attr_names
 from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload
+from modelopt.torch.utils import distributed as _dist
 from modelopt.torch.utils.dataset_utils import _disable_use_cache
 from modelopt.torch.utils.distributed import is_fsdp2_model
 
@@ -999,11 +999,10 @@ def _export_transformers_checkpoint(
     _reconstruct_fused_moe_linear(model)
 
     if is_fsdp2_model(model):
-        # FSDP2: gather the full (unsharded) state_dict to CPU on rank 0.
-        quantized_state_dict = get_model_state_dict(
-            model,
-            options=StateDictOptions(full_state_dict=True, cpu_offload=True),
-        )
+        # FSDP2: each rank packed its own Shard(0) slice. Gather only THIS rank's owned decoder-layer
+        # units to full CPU tensors (bounded host memory; the write is parallel per-rank). At world=1
+        # rank 0 owns every unit, so this one path covers any world size.
+        quantized_state_dict = _gather_owned_units(model)
     else:
         # Non-FSDP2: assumes a replicated model (rank 0 has the full state dict).
         quantized_state_dict = model.state_dict()
@@ -1488,6 +1487,152 @@ def _write_hf_export_config(
         json.dump(config_data, file, indent=4)
 
 
+# ---------------------------------------------------------------------------
+# FSDP2 shard-local parallel checkpoint save (see docs/fsdp2_export_changes.md).
+# Each rank writes its owned decoder-layer units concurrently; host memory is
+# bounded to its owned subset and the disk write parallelizes across ranks.
+# ---------------------------------------------------------------------------
+
+
+def _materialize_cpu(v: torch.Tensor) -> torch.Tensor:
+    """Convert a gathered state-dict value to a plain, contiguous CPU tensor.
+
+    A gathered DTensor here is fully replicated (``full_tensor``/replicated buffer), so
+    ``to_local()`` yields the complete tensor. A ``QTensorWrapper`` (compressed weight) unwraps
+    to its packed data via ``.data``.
+    """
+    from torch.distributed.tensor import DTensor
+
+    if isinstance(v, DTensor):
+        v = v.to_local()
+    v = getattr(v, "data", v)  # QTensorWrapper -> packed data; plain tensor -> itself
+    return v.detach().to("cpu").contiguous()
+
+
+def _enumerate_export_units(model, id_to_name):
+    """Ordered export units, identical on every rank (this order drives ownership).
+
+    Each decoder layer is one unit; a trailing "root-leaves" unit holds every module that owns
+    parameters directly and is not under a decoder-layer prefix (embed / lm_head / final norm).
+    Returns ``[(modules, is_root), ...]``. If decoder layers cannot be discovered, the whole model
+    becomes a single root unit (correct, but no parallelism).
+    """
+    from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
+
+    decoder_layers = LayerActivationCollector.get_decoder_layers(model)
+    units: list[tuple[list[nn.Module], bool]] = []
+    layer_prefixes: tuple[str, ...] = ()
+    if decoder_layers is not None:
+        layer_prefixes = tuple(id_to_name[id(layer)] + "." for layer in decoder_layers)
+        units = [([layer], False) for layer in decoder_layers]
+    root_leaves = [
+        m
+        for n, m in model.named_modules()
+        if next(m.parameters(recurse=False), None) is not None
+        and not (layer_prefixes and n.startswith(layer_prefixes))
+    ]
+    units.append((root_leaves, True))
+    return units
+
+
+def _gather_unit(modules, id_to_name, is_owner):
+    """Gather one unit's params to the owner via ``full_tensor()`` (all-gather on every rank).
+
+    Only the owner keeps the materialized CPU copy. Non-owners join the collective and drop the
+    result, so the gate is on the *keep*, never the collective (skipping it on non-owners deadlocks).
+    """
+    from torch.distributed.tensor import DTensor
+
+    unit_sd: dict[str, torch.Tensor] = {}
+    for m in modules:
+        base = id_to_name.get(id(m), "")
+        prefix = (base + ".") if base else ""
+        for name, v in m.state_dict().items():
+            full = v.full_tensor() if isinstance(v, DTensor) else v  # COLLECTIVE (all ranks)
+            if is_owner:
+                unit_sd[prefix + name] = _materialize_cpu(full)
+    return unit_sd if is_owner else None
+
+
+def _write_owned_shards(owned, export_dir, rank_id, max_shard_size):
+    """Write this rank's tensors as safetensors shards via the HF splitter.
+
+    Returns ``(weight_map {key: filename}, total_bytes)`` for the index merge. Filenames embed the rank
+    so ranks never collide; the HF splitter handles ``max_shard_size`` and never splits a single tensor
+    (a >cap tensor gets its own over-cap shard).
+    """
+    from huggingface_hub import split_torch_state_dict_into_shards
+
+    pattern = f"model-r{rank_id:02d}{{suffix}}.safetensors"
+    split = split_torch_state_dict_into_shards(
+        owned, filename_pattern=pattern, max_shard_size=max_shard_size
+    )
+    for fname, keys in split.filename_to_tensors.items():
+        save_file(
+            {k: owned[k] for k in keys}, str(Path(export_dir) / fname), metadata={"format": "pt"}
+        )
+    return split.tensor_to_filename, split.metadata["total_size"]
+
+
+def _gather_owned_units(model):
+    """Each rank gathers only the decoder-layer units it owns (``i % world``) to full CPU tensors.
+
+    All ranks walk every unit in lockstep (``_gather_unit``'s ``full_tensor`` is collective); each rank
+    keeps just its owned units, so host memory is bounded to the owned subset (+ one transient unit on
+    GPU). Keep-fused MoE experts are split into per-expert keys once the owned units are assembled.
+    """
+    from .moe_utils import _split_fused_experts_state_dict
+
+    rank, world = _dist.rank(), _dist.size()
+    id_to_name = {id(m): n for n, m in model.named_modules()}
+    units = _enumerate_export_units(model, id_to_name)
+    my_sd: dict[str, torch.Tensor] = {}
+    for i, (modules, _is_root) in enumerate(units):
+        owned = _gather_unit(
+            modules, id_to_name, is_owner=(i % world == rank)
+        )  # COLLECTIVE all ranks
+        if owned is not None:
+            my_sd.update(owned)
+    _split_fused_experts_state_dict(
+        my_sd, model
+    )  # only this rank's owned keep-fused keys are present
+    return my_sd
+
+
+def _finalize_index(local_maps, export_dir, rank, world):
+    """gather_object the per-rank HF weight_maps to rank 0; write the standard index. No sidecars."""
+    from huggingface_hub.constants import SAFETENSORS_INDEX_FILE
+
+    gathered: list[Any] | None
+    if world > 1:
+        gathered = [None] * world if rank == 0 else None
+        torch.distributed.gather_object(local_maps, gathered, dst=0)
+    else:
+        gathered = [local_maps]
+    if rank != 0 or gathered is None:
+        return
+    weight_map: dict[str, str] = {}
+    total = 0
+    for rank_maps in gathered:
+        for t2f, nbytes in rank_maps:
+            weight_map.update(t2f)
+            total += nbytes
+    index = {"metadata": {"total_size": total}, "weight_map": weight_map}
+    (Path(export_dir) / SAFETENSORS_INDEX_FILE).write_text(json.dumps(index, indent=2))
+
+
+def _parallel_write(my_sd, export_dir, max_shard_size):
+    """Distributed write: each rank writes its own safetensors shards concurrently.
+
+    Rank 0 then merges the single ``model.safetensors.index.json`` from the gathered per-rank
+    weight_maps.
+    """
+    rank, world = _dist.rank(), _dist.size()
+    local_map = _write_owned_shards(my_sd, export_dir, rank, max_shard_size)  # (weight_map, nbytes)
+    _dist.barrier()
+    _finalize_index([local_map], export_dir, rank, world)
+
+
 def export_hf_checkpoint(
     model: Any,
     dtype: torch.dtype | None = None,
@@ -1589,7 +1734,12 @@ def export_hf_checkpoint(
         if getattr(model, "hf_quantizer", None) is not None:
             model.hf_quantizer = None
 
-        export_state_dict = {**post_state_dict, **(extra_state_dict or {})}
+        # extra_state_dict (e.g. MTP) is added on rank 0 only under distributed (it owns that slot);
+        # non-owner ranks hold only their own gathered units.
+        is_rank0 = (not is_distributed) or torch.distributed.get_rank() == 0
+        export_state_dict = dict(post_state_dict)
+        if extra_state_dict and is_rank0:
+            export_state_dict.update(extra_state_dict)
 
         # transformers may have applied a load-time conversion_mapping (fused gate_up_proj,
         # renamed MoE leaves, reordered model/language_model prefix), so the in-memory names
@@ -1613,30 +1763,43 @@ def export_hf_checkpoint(
                 "names may not match the original HF hub checkpoint."
             )
 
-        # Under torch.distributed only rank 0 writes; others sync at the finally barrier.
-        if is_distributed and torch.distributed.get_rank() != 0:
-            return
-
-        # Keep transformers' own revert_weight_conversion disabled (the quant-aware reverse
-        # above replaces it): it can't handle quantized state dicts (RuntimeError on 0-d scalar
-        # scale tensors). Patch both the source and importing module since modeling_utils does
-        # `from core_model_loading import revert_weight_conversion`.
-        _patches = _patch_revert_weight_conversion()
-
         _sanitize_generation_config_for_save(model)
 
-        # TODO: parallelize the disk write across ranks (avoid single-process speed + rank-0 OOM).
-        try:
-            model.save_pretrained(
-                export_dir,
-                state_dict=export_state_dict,
-                save_modelopt_state=save_modelopt_state,
-                max_shard_size=max_shard_size,
-            )
-        finally:
-            _unpatch_revert_weight_conversion(_patches)
+        if is_distributed:
+            # FSDP2 multi-rank: every rank writes its owned shards concurrently (the gather/index
+            # collectives are inside _parallel_write, so all ranks must call it); rank 0 then writes
+            # config / generation / modelopt_state and embeds the quant config. Replaces the old
+            # rank-0 full-gather + single-process save_pretrained (fixes the OOM + serialized write).
+            _parallel_write(export_state_dict, export_dir, max_shard_size)
+            if is_rank0:
+                model.config.save_pretrained(str(export_dir))
+                if getattr(model, "generation_config", None) is not None:
+                    model.generation_config.save_pretrained(str(export_dir))
+                if save_modelopt_state:
+                    # Model-level (independent of the weight write), so rank 0 writes it directly —
+                    # matching what save_pretrained(save_modelopt_state=True) does internally.
+                    from modelopt.torch.opt.conversion import ModeloptStateManager, modelopt_state
+
+                    if ModeloptStateManager.is_converted(model):
+                        torch.save(modelopt_state(model), export_dir / "modelopt_state.pth")
+                _write_hf_export_config(model, hf_quant_config, export_dir)
+        else:
+            # Keep transformers' own revert_weight_conversion disabled (the quant-aware reverse
+            # above replaces it): it can't handle quantized state dicts (RuntimeError on 0-d scalar
+            # scale tensors). Patch both the source and importing module since modeling_utils does
+            # `from core_model_loading import revert_weight_conversion`.
+            _patches = _patch_revert_weight_conversion()
+            try:
+                model.save_pretrained(
+                    export_dir,
+                    state_dict=export_state_dict,
+                    save_modelopt_state=save_modelopt_state,
+                    max_shard_size=max_shard_size,
+                )
+            finally:
+                _unpatch_revert_weight_conversion(_patches)
 
-        _write_hf_export_config(model, hf_quant_config, export_dir)
+            _write_hf_export_config(model, hf_quant_config, export_dir)
 
     except Exception as e:
         warnings.warn(
diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py
index 3d1e017b3a5..8bd52a53f3b 100644
--- a/modelopt/torch/quantization/utils/core_utils.py
+++ b/modelopt/torch/quantization/utils/core_utils.py
@@ -137,7 +137,7 @@ def convert_quantization_axis_to_reduce_axis(input, axis):
     """
     if axis is None:
         return None
-    axis = axis if isinstance(axis, (list, tuple)) else [axis]
+    axis = axis if isinstance(axis, list | tuple) else [axis]
     # Handle positive and negative axis.
     reduce_axis = [i for i in range(input.dim()) if i not in axis and (i - input.dim()) not in axis]
     return reduce_axis
@@ -230,7 +230,7 @@ def representative_weight_quantizer(module: nn.Module, weight_name: str = "weigh
 
     singular = quantizer_attr_names(weight_name).weight_quantizer
     q = getattr(module, singular, None)
-    if isinstance(q, (TensorQuantizer, SequentialQuantizer)):
+    if isinstance(q, TensorQuantizer | SequentialQuantizer):
         return q
     if isinstance(q, GroupedQuantizer) and len(q) > 0:
         return q[0]
@@ -238,7 +238,7 @@ def representative_weight_quantizer(module: nn.Module, weight_name: str = "weigh
     plural = getattr(module, singular + "s", None)
     if isinstance(plural, nn.ModuleList) and len(plural) > 0:
         first = plural[0]
-        if isinstance(first, (TensorQuantizer, SequentialQuantizer)):
+        if isinstance(first, TensorQuantizer | SequentialQuantizer):
             return first
     return None
 
@@ -1046,6 +1046,210 @@ def fsdp2_aware_weight_update(root_model, modules_to_update, reshard=True):
                     root_module.reshard()
 
 
+_ShardInfo = namedtuple("_ShardInfo", ["name", "old", "mesh", "placements"])
+
+
+def _shard_start(param):
+    """Global offset along shard dim 0 for a ``Shard(0)`` DTensor on a 1-D FSDP mesh.
+
+    Returns 0 for non-DTensor params. Uses torch's even-split convention: the first ``dim0 % world``
+    ranks get an extra row, so rank ``r``'s offset is ``r * (dim0 // world) + min(r, dim0 % world)``.
+    Correct for both even and uneven sharding.
+    """
+    if not isinstance(param, DTensor):
+        return 0
+    dim0 = param.shape[0]  # global size
+    world = param.device_mesh.size()
+    r = param.device_mesh.get_local_rank()
+    base, rem = divmod(dim0, world)
+    return r * base + min(r, rem)
+
+
+def _rebuild_fsdp_param_from_shard(old_fp, packed_local):
+    """Re-register a PACKED weight as an FSDPParam from full shape + this rank's packed shard.
+
+    Does so WITHOUT materializing the full weight.
+
+    The obvious route -- ``DTensor.from_local(shard)`` then ``FSDPParam(dtensor, ...)`` -- crashes:
+    ``_init_sharded_param`` takes a DTensor down its tensor-parallel branch
+    (``DeviceMesh._concatenate([dp_mesh, tp_mesh])``) which has no ``tp_mesh`` on a plain 1-D FSDP
+    mesh. Instead build the FSDPParam from a ``meta``-device full-shape tensor (a plain tensor ->
+    non-DTensor path; ``meta`` is explicitly allowed by ``_init_sharded_param``), which computes every
+    size/stride/spec field with no real full data and no crash, then overwrite only the two
+    data-holding fields (``_sharded_param_data``, ``sharded_param``) with this rank's packed shard.
+    Only dim-0 (the ``Shard(0)`` axis) differs between shard and full; inner dims are identical.
+    Validated by ``spike_inplace_pack.py`` (world=2, NVFP4).
+    """
+    from modelopt.torch.quantization.qtensor.base_qtensor import QFSDPParam, QTensorWrapper
+
+    full_shape = (old_fp._orig_size[0], *packed_local.shape[1:])
+    is_qtw = isinstance(packed_local, QTensorWrapper)
+    param_class = QFSDPParam if is_qtw else FSDPParam
+    mp = MixedPrecisionPolicy(
+        param_dtype=packed_local.dtype,
+        reduce_dtype=None,
+        output_dtype=None,
+        cast_forward_inputs=False,
+    )
+
+    # (a) meta full-shape skeleton -> computes sharded_size / padded_sharded_param_size / specs
+    skeleton = torch.empty(full_shape, dtype=packed_local.dtype, device="meta")
+    if is_qtw:
+        skeleton = QTensorWrapper(skeleton, metadata={**packed_local.metadata, "shape": full_shape})
+    new_fp = param_class(
+        nn.Parameter(skeleton, requires_grad=False),
+        old_fp._module_info,
+        old_fp.mesh_info,
+        old_fp.post_forward_mesh_info,
+        old_fp.device,
+        None,
+        mp,
+        None,
+    )
+    if param_class is FSDPParam:
+        new_fp.init_dtype_attrs(mp)
+
+    # (b) overwrite the two data fields with the real packed shard (pad the 0th shard as FSDP does)
+    local = (packed_local.data if is_qtw else packed_local).contiguous().to(old_fp.device)
+    padded = local.new_zeros(new_fp.padded_sharded_param_size)
+    padded.narrow(0, 0, new_fp.sharded_size[0]).copy_(local)
+    new_fp._sharded_param_data = padded.view(-1)
+    new_fp.sharded_param = nn.Parameter(
+        new_fp.to_sharded_dtensor(padded.narrow(0, 0, new_fp.sharded_size[0])), requires_grad=False
+    )
+    new_fp._setattr_on_modules(new_fp.sharded_param)
+    return new_fp
+
+
+def _rewrap_scale_buffers_shard0(module, captured, local_dim0):
+    """Re-wrap packed per-shard scale buffers as ``Shard(0)`` DTensors.
+
+    A later ``full_tensor()`` gather then reconstructs the full scale.
+
+    Scale buffers are ordinary registered buffers (not FSDP params), so they can use
+    ``DTensor.from_local`` directly (the FSDPParam-constructor crash does not apply). A buffer is
+    treated as sharded iff its dim-0 matches the local shard row/expert count (``local_dim0``):
+    ``weight_scale [rows/world, ...]`` and per-expert ``weight_scale_2 [E/world]`` -> ``Shard(0)``;
+    per-tensor scalars (``weight_scale_2`` for a plain linear, ``input_scale``) stay replicated.
+    Validated by ``spike_scale_gather.py`` (world=2, NVFP4).
+    """
+    info = next(iter(captured.values()))
+    mesh, placements = info.mesh, info.placements
+    # The scale's dim-0 matches the weight's global dim-0 (per-row / per-expert). Pass the true global
+    # shape+stride to from_local: without it, from_local infers global = local * world (even-shard
+    # assumption), so under UNEVEN sharding ranks disagree on the global shape and full_tensor()
+    # deadlocks. sharded_param DTensors (weights) already carry the right shape via their spec.
+    global_dim0 = info.old._orig_size[0]
+    for bname, buf in list(module._buffers.items()):
+        if buf is None or isinstance(buf, DTensor) or buf.dim() == 0:
+            continue
+        if buf.shape[0] == local_dim0:
+            local = buf.contiguous()
+            global_shape = torch.Size((global_dim0, *local.shape[1:]))
+            global_stride = torch.empty(global_shape, device="meta").stride()
+            module._buffers[bname] = DTensor.from_local(
+                local, mesh, placements, shape=global_shape, stride=global_stride
+            )
+
+
+@contextmanager
+def fsdp2_shard_local_pack(root_model, module):
+    """Pack a module's ``Shard(0)`` weights on the LOCAL shard, in place, keeping them sharded.
+
+    Works for plain quant-linears (``Shard(0)`` on ``out``) and keep-fused experts (``Shard(0)`` on
+    ``E``). No unshard, no reshard, no collective -> leaves a valid packed FSDP module. No-ops for
+    non-FSDP models (so the same handler code covers the single-process/standard path).
+
+    Enter: ``to_local`` each weight param (recording its old FSDPParam + mesh/placements and this
+    rank's global shard offset in ``module._shard_local_start``); the wrapped handler packs the plain
+    local block in place. Exit: re-register each packed weight via :func:`_rebuild_fsdp_param_from_shard`
+    and re-wrap the packed scale buffers via :func:`_rewrap_scale_buffers_shard0`; no reshard.
+    """
+    if not isinstance(root_model, FSDPModule):
+        yield
+        return
+
+    root_module = _get_enclosing_fsdp_module(module, root_model)
+    group = fully_shard.state(root_module)._fsdp_param_group
+    # The root FSDP module keeps reshard_after_forward=False, so a prior forward (e.g. the export
+    # resmooth) can leave its own params (embed/lm_head/norm) UNSHARDED (full, non-DTensor). Reshard so
+    # every shardable param is Shard(0) and gets captured + FSDPParam-rebuilt below -- otherwise an
+    # in-place pack of an unsharded param is silently discarded when state_dict re-materializes from the
+    # (stale) FSDPParam. Idempotent: a no-op for already-sharded decoder layers.
+    if group is not None and not group.is_sharded:
+        root_module.reshard()
+    mapping = create_fsdp_param_mapping(group.fsdp_params, root_model)
+
+    captured = {}
+    module._shard_local_start = {}
+    for pname, param in list(module.named_parameters(recurse=False)):
+        name = f"{_get_module_name(module, root_model)}.{pname}"
+        if name not in mapping:
+            continue
+        # A non-DTensor param that survives the reshard above is genuinely replicated (not row-sharded,
+        # e.g. shard_root=False root params): shard-local packing does not apply -> leave it for the
+        # handler to pack in place. Only Shard(0) DTensors get the shard-local path.
+        if not isinstance(param, DTensor):
+            continue
+        # Fail fast + clear on uneven sharding rather than deadlocking the gather later. Symmetric
+        # (global shape + world are identical on every rank) so this raise can't itself hang.
+        world = param.device_mesh.size()
+        if param.shape[0] % world != 0:
+            raise NotImplementedError(
+                f"fsdp2_shard_local_pack does not support uneven sharding: '{name}' has dim0="
+                f"{param.shape[0]}, not divisible by world size {world}. Use a world size that "
+                f"divides every sharded dim-0."
+            )
+        captured[pname] = _ShardInfo(name, mapping[name], param.device_mesh, param.placements)
+        module._shard_local_start[pname] = _shard_start(param)
+        module._parameters[pname] = nn.Parameter(param.to_local(), requires_grad=False)
+    try:
+        yield
+    finally:
+        local_dim0 = None
+        with no_requires_grad(), enable_fake_quant(module):
+            for pname, info in captured.items():
+                packed_local = getattr(module, pname)
+                if local_dim0 is None:
+                    local_dim0 = packed_local.shape[0]
+                mapping[info.name] = _rebuild_fsdp_param_from_shard(info.old, packed_local)
+                info.old._post_load_hook_handle.remove()
+        if captured:
+            _rewrap_scale_buffers_shard0(module, captured, local_dim0)
+            group.fsdp_params = list(mapping.values())
+        module.__dict__.pop("_shard_local_start", None)
+
+
+@contextmanager
+def materialize_fsdp2_root(model: nn.Module):
+    """Unshard a sharded FSDP2 root's own params (embed/lm_head/norm) for a calibration forward.
+
+    The calibration loop calls ``model.forward(**batch)`` directly (``dataset_utils._forward_loop``)
+    rather than ``model(**batch)``, so ``nn.Module.__call__`` is bypassed and the root's FSDP2
+    forward pre-hook never fires. Its own params (embed/lm_head/norm) stay sharded DTensors and the
+    forward hits ``aten.embedding: mixed Tensor and DTensor``. Decoder layers are unaffected: they are
+    called as ``layer(...)`` inside ``forward``, so their pre-hooks fire and they unshard normally.
+
+    Unshard the root up front so its params are full tensors, then reshard on exit. The bypass also
+    skips the root's post-forward hook, so nothing reshards it mid-calibration and a single unshard
+    holds across all batches. Cheap: only the root's own param group is gathered, not the decoder
+    layers. No-op for non-FSDP2 or already-replicated roots.
+    """
+    root_sharded = False
+    if isinstance(model, FSDPModule):
+        pg = fully_shard.state(model)._fsdp_param_group
+        root_sharded = pg is not None and pg.is_sharded
+    if root_sharded:
+        with enable_fake_quant(model):
+            model.unshard()
+    try:
+        yield
+    finally:
+        if root_sharded:
+            with enable_fake_quant(model):
+                model.reshard()
+
+
 def update_quant_cfg_with_kv_cache_quant(
     quant_cfg: dict[str, Any], kv_cache_quant_cfg: list[QuantizerCfgEntry]
 ) -> dict[str, Any]:

From db91e0d63c8646077aac55a34c098d36a8f32b1d Mon Sep 17 00:00:00 2001
From: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
Date: Wed, 19 Aug 2026 16:59:35 +0000
Subject: [PATCH 2/2] cleanup

Signed-off-by: Suguna Velury <178320438+sugunav14@users.noreply.github.com>
---
 modelopt/torch/export/hf_export_handlers.py   | 37 +++----
 modelopt/torch/export/moe_utils.py            |  8 +-
 modelopt/torch/export/unified_export_hf.py    | 99 ++++++-------------
 .../torch/quantization/utils/core_utils.py    | 55 ++---------
 4 files changed, 55 insertions(+), 144 deletions(-)

diff --git a/modelopt/torch/export/hf_export_handlers.py b/modelopt/torch/export/hf_export_handlers.py
index 7a0bfe54408..be9dd92ec55 100644
--- a/modelopt/torch/export/hf_export_handlers.py
+++ b/modelopt/torch/export/hf_export_handlers.py
@@ -26,7 +26,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, _export_fused_experts_keep_fused
+from .moe_utils import _export_fused_experts, _pack_fused_experts_shard_local
 from .quant_utils import get_quantization_format
 from .registry import ExportContext, ExportModuleRegistry, PrepareMoEInputsRegistry
 
@@ -38,11 +38,6 @@ def _has_fused_experts_quantizers(module: nn.Module) -> bool:
     return hasattr(module, f"{first_proj_attr}_weight_quantizers")
 
 
-def _use_shard_local(model: nn.Module) -> bool:
-    """Whether to use shard-local packing (FSDP2 only)."""
-    return is_fsdp2_model(model)
-
-
 def _export_weight(
     module: nn.Module,
     ctx: ExportContext,
@@ -137,18 +132,15 @@ def _export_moe_linear(name: str, module: nn.Module, ctx: ExportContext) -> None
 def _export_fused_experts_module(name: str, module: nn.Module, ctx: ExportContext) -> None:
     """Split and quantize a fused-experts module with plural weight quantizers.
 
-    Under FSDP2 the fused weight is ``Shard(0)`` on the expert dim, so the destructive per-rank split
-    (``_export_fused_experts``) would drop non-owner experts from the gather. Instead pack this rank's
-    experts IN PLACE and keep them fused (``_export_fused_experts_keep_fused`` inside
-    ``fsdp2_shard_local_pack``); the gate/up split is deferred to write time
-    (``_split_fused_experts_state_dict`` in the gather). Non-FSDP is unchanged.
+    Under FSDP2 each rank holds only some experts, so it packs just those (kept fused) and the
+    split into per-expert keys is deferred until the gather brings all experts together.
 
     Tied experts are packed independently and their duplicate keys are dropped by name
     in postprocess_state_dict; no per-module dedup cache is used.
     """
-    if _use_shard_local(ctx.model):
+    if is_fsdp2_model(ctx.model):
         with fsdp2_shard_local_pack(ctx.model, module):
-            _export_fused_experts_keep_fused(module, ctx.dtype)
+            _pack_fused_experts_shard_local(module, ctx.dtype)
     else:
         _export_fused_experts(module, ctx.dtype)
 
@@ -162,7 +154,7 @@ def _export_quant_linear(name: str, module: nn.Module, ctx: ExportContext) -> No
     """
     if get_quantization_format(module) == QUANTIZATION_NONE:
         return
-    cm = fsdp2_shard_local_pack(ctx.model, module) if _use_shard_local(ctx.model) else nullcontext()
+    cm = fsdp2_shard_local_pack(ctx.model, module) if is_fsdp2_model(ctx.model) else nullcontext()
     try:
         with cm:
             _export_weight(module, ctx)
@@ -194,10 +186,8 @@ def _export_quant_embedding(name: str, module: nn.Module, ctx: ExportContext) ->
             "The embedding will be exported as its fake-quantized float weight."
         )
         return
-    # The embedding lives in the root FSDP unit (reshard_after_forward=False -> its params may be
-    # unsharded at export); fsdp2_shard_local_pack reshards them to Shard(0) first so this rank packs
-    # only its vocab slice in place, no full unshard. Non-FSDP is unchanged.
-    cm = fsdp2_shard_local_pack(ctx.model, module) if _use_shard_local(ctx.model) else nullcontext()
+    # fsdp2_shard_local_pack reshards the unsharded root embedding to Shard(0); no-op for non-FSDP.
+    cm = fsdp2_shard_local_pack(ctx.model, module) if is_fsdp2_model(ctx.model) else nullcontext()
     try:
         with cm:
             _export_weight(module, ctx)
@@ -209,13 +199,10 @@ def _export_quant_embedding(name: str, module: nn.Module, ctx: ExportContext) ->
 
 @ExportModuleRegistry.register("Llama4TextExperts", "GptOssExperts")
 def _export_bmm_experts(name: str, module: nn.Module, ctx: ExportContext) -> None:
-    """Export fused BMM-style expert weights and quantization metadata.
+    """Export fused BMM-style expert weights (Llama4 / GPT-OSS).
 
-    The fused ``gate_up_proj``/``down_proj`` are ``[E, ...]`` (``Shard(0)`` on E under FSDP2) and each
-    is packed WHOLE by ``_export_quantized_weight`` with its singular weight quantizer -- whose amax is
-    calibrated over all experts (a replicated buffer, i.e. global), so packing a single rank's
-    ``[E/world, ...]`` slice yields the same bytes. So under FSDP2 this uses ``fsdp2_shard_local_pack``
-    (pack this rank's slice in place, no unshard); non-FSDP is unchanged.
+    Its weight quantizer has one amax covering all experts, so under FSDP2 each rank can pack
+    just the experts it owns and produce identical bytes -- no gather or unshard needed.
     """
     if get_quantization_format(module) == QUANTIZATION_NONE:
         return
@@ -228,7 +215,7 @@ def _export_bmm_experts(name: str, module: nn.Module, ctx: ExportContext) -> Non
         modules=module,
         quantizer_attrs=["gate_up_proj_input_quantizer", "down_proj_input_quantizer"],
     )
-    cm = fsdp2_shard_local_pack(ctx.model, module) if _use_shard_local(ctx.model) else nullcontext()
+    cm = fsdp2_shard_local_pack(ctx.model, module) if is_fsdp2_model(ctx.model) else nullcontext()
     with cm:
         for weight_name in ["gate_up_proj", "down_proj"]:
             _export_weight(module, ctx, weight_name)
diff --git a/modelopt/torch/export/moe_utils.py b/modelopt/torch/export/moe_utils.py
index bee0a130f76..ef275f68078 100644
--- a/modelopt/torch/export/moe_utils.py
+++ b/modelopt/torch/export/moe_utils.py
@@ -228,7 +228,7 @@ def _export_fused_experts(
     _delete_fused_moe_source_attrs(module)
 
 
-def _export_fused_experts_keep_fused(module, dtype):
+def _pack_fused_experts_shard_local(module, dtype):
     """Shard-local keep-fused NVFP4/FP8 pack for one rank's experts.
 
     Pack THIS rank's experts IN PLACE and keep the weight FUSED (``Shard(0)`` on E),
@@ -247,7 +247,7 @@ def _export_fused_experts_keep_fused(module, dtype):
     ``_export_quantized_weight`` (so the per-tensor ``weight_scale_2`` is naturally shared by the gate
     and up halves -- what vLLM's load-time fusion requires -- and per-block scales are row-local). The
     packed per-expert results are stacked back into fused ``[local_n, ...]`` tensors + per-expert
-    scale buffers; the gate/up split is deferred to write time (``_split_fused_experts_state_dict``).
+    scale buffers; the gate/up split is deferred to write time (``_split_packed_fused_experts``).
     Non-destructive: the fused params survive, so :func:`fsdp2_shard_local_pack` re-registers them.
     """
     from modelopt.torch.export.unified_export_hf import _export_quantized_weight
@@ -318,10 +318,10 @@ def _register(attr, weights, scales, scale2s):
         module.register_buffer("down_proj_input_scale", dp_input_scale)
 
 
-def _split_fused_experts_state_dict(state_dict, model):
+def _split_packed_fused_experts(state_dict, model):
     """Split keep-fused expert tensors in a gathered ``state_dict`` into per-expert keys.
 
-    Converts the ``_export_fused_experts_keep_fused`` output (fused ``{name}.gate_up_proj [E,2I,H/2]``
+    Converts the ``_pack_fused_experts_shard_local`` output (fused ``{name}.gate_up_proj [E,2I,H/2]``
     + per-expert scales) into the same per-expert deployment keys ``_export_fused_experts`` emits
     (``{name}.{e}.gate_proj.weight`` / ``up_proj`` / ``down_proj`` + scales). Byte-identical: packing
     is row-independent, so slicing the packed fused tensor equals packing each half; gate/up share the
diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py
index b485dc55d74..6f62666f6be 100644
--- a/modelopt/torch/export/unified_export_hf.py
+++ b/modelopt/torch/export/unified_export_hf.py
@@ -27,8 +27,11 @@
 
 import torch
 import torch.nn as nn
+from huggingface_hub import split_torch_state_dict_into_shards
+from huggingface_hub.constants import SAFETENSORS_INDEX_FILE
 from safetensors import safe_open
 from safetensors.torch import save_file
+from torch.distributed.tensor import DTensor
 
 from .diffusers_utils import build_layerwise_quant_metadata, pad_nvfp4_weights, swizzle_nvfp4_scales
 
@@ -60,10 +63,13 @@
 from modelopt.torch.quantization.qtensor.nvfp4_tensor import _cast_per_block_scale_to_fp8
 from modelopt.torch.quantization.utils import fsdp2_aware_weight_update, quantizer_attr_names
 from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload
+from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
 from modelopt.torch.utils import distributed as _dist
 from modelopt.torch.utils.dataset_utils import _disable_use_cache
 from modelopt.torch.utils.distributed import is_fsdp2_model
 
+from .moe_utils import _split_packed_fused_experts
+
 try:
     from modelopt.torch.sparsity.attention_sparsity.conversion import export_sparse_attention_config
 except ImportError:
@@ -999,9 +1005,7 @@ def _export_transformers_checkpoint(
     _reconstruct_fused_moe_linear(model)
 
     if is_fsdp2_model(model):
-        # FSDP2: each rank packed its own Shard(0) slice. Gather only THIS rank's owned decoder-layer
-        # units to full CPU tensors (bounded host memory; the write is parallel per-rank). At world=1
-        # rank 0 owns every unit, so this one path covers any world size.
+        # Each rank gathers only its owned layers to CPU (bounded memory; parallel per-rank write).
         quantized_state_dict = _gather_owned_units(model)
     else:
         # Non-FSDP2: assumes a replicated model (rank 0 has the full state dict).
@@ -1487,62 +1491,36 @@ def _write_hf_export_config(
         json.dump(config_data, file, indent=4)
 
 
-# ---------------------------------------------------------------------------
-# FSDP2 shard-local parallel checkpoint save (see docs/fsdp2_export_changes.md).
-# Each rank writes its owned decoder-layer units concurrently; host memory is
-# bounded to its owned subset and the disk write parallelizes across ranks.
-# ---------------------------------------------------------------------------
-
-
 def _materialize_cpu(v: torch.Tensor) -> torch.Tensor:
-    """Convert a gathered state-dict value to a plain, contiguous CPU tensor.
-
-    A gathered DTensor here is fully replicated (``full_tensor``/replicated buffer), so
-    ``to_local()`` yields the complete tensor. A ``QTensorWrapper`` (compressed weight) unwraps
-    to its packed data via ``.data``.
-    """
-    from torch.distributed.tensor import DTensor
-
+    """Convert a gathered weight into a plain, contiguous CPU tensor."""
     if isinstance(v, DTensor):
         v = v.to_local()
     v = getattr(v, "data", v)  # QTensorWrapper -> packed data; plain tensor -> itself
     return v.detach().to("cpu").contiguous()
 
 
-def _enumerate_export_units(model, id_to_name):
-    """Ordered export units, identical on every rank (this order drives ownership).
+def _enumerate_export_units(model):
+    """List the module groups to export, one per decoder layer plus a final "leftovers" group.
 
-    Each decoder layer is one unit; a trailing "root-leaves" unit holds every module that owns
-    parameters directly and is not under a decoder-layer prefix (embed / lm_head / final norm).
-    Returns ``[(modules, is_root), ...]``. If decoder layers cannot be discovered, the whole model
-    becomes a single root unit (correct, but no parallelism).
+    The leftovers group is the weight-owning modules outside any layer (embeddings, lm_head, norm).
+    Every rank builds the same list, so they agree on who exports what without communicating.
     """
-    from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
-
-    decoder_layers = LayerActivationCollector.get_decoder_layers(model)
-    units: list[tuple[list[nn.Module], bool]] = []
-    layer_prefixes: tuple[str, ...] = ()
-    if decoder_layers is not None:
-        layer_prefixes = tuple(id_to_name[id(layer)] + "." for layer in decoder_layers)
-        units = [([layer], False) for layer in decoder_layers]
+    decoder_layers = LayerActivationCollector.get_decoder_layers(model) or []
+    in_layer = {id(sm) for layer in decoder_layers for sm in layer.modules()}
     root_leaves = [
         m
-        for n, m in model.named_modules()
-        if next(m.parameters(recurse=False), None) is not None
-        and not (layer_prefixes and n.startswith(layer_prefixes))
+        for m in model.modules()
+        if id(m) not in in_layer and next(m.parameters(recurse=False), None) is not None
     ]
-    units.append((root_leaves, True))
-    return units
+    return [[layer] for layer in decoder_layers] + [root_leaves]
 
 
 def _gather_unit(modules, id_to_name, is_owner):
-    """Gather one unit's params to the owner via ``full_tensor()`` (all-gather on every rank).
+    """Collect one group's weights onto the owner rank.
 
-    Only the owner keeps the materialized CPU copy. Non-owners join the collective and drop the
-    result, so the gate is on the *keep*, never the collective (skipping it on non-owners deadlocks).
+    Every rank takes part in the all-gather (it is a collective, so skipping it on the non-owners
+    would hang); only the owner keeps the result.
     """
-    from torch.distributed.tensor import DTensor
-
     unit_sd: dict[str, torch.Tensor] = {}
     for m in modules:
         base = id_to_name.get(id(m), "")
@@ -1555,14 +1533,11 @@ def _gather_unit(modules, id_to_name, is_owner):
 
 
 def _write_owned_shards(owned, export_dir, rank_id, max_shard_size):
-    """Write this rank's tensors as safetensors shards via the HF splitter.
+    """Write this rank's tensors to its own safetensors shard files.
 
-    Returns ``(weight_map {key: filename}, total_bytes)`` for the index merge. Filenames embed the rank
-    so ranks never collide; the HF splitter handles ``max_shard_size`` and never splits a single tensor
-    (a >cap tensor gets its own over-cap shard).
+    Returns the ``{tensor_name: filename}`` map and total byte size, which rank 0 uses to build the
+    index. Filenames include the rank number so ranks never overwrite each other.
     """
-    from huggingface_hub import split_torch_state_dict_into_shards
-
     pattern = f"model-r{rank_id:02d}{{suffix}}.safetensors"
     split = split_torch_state_dict_into_shards(
         owned, filename_pattern=pattern, max_shard_size=max_shard_size
@@ -1575,34 +1550,28 @@ def _write_owned_shards(owned, export_dir, rank_id, max_shard_size):
 
 
 def _gather_owned_units(model):
-    """Each rank gathers only the decoder-layer units it owns (``i % world``) to full CPU tensors.
+    """Gather the layers this rank owns onto it, as CPU tensors.
 
-    All ranks walk every unit in lockstep (``_gather_unit``'s ``full_tensor`` is collective); each rank
-    keeps just its owned units, so host memory is bounded to the owned subset (+ one transient unit on
-    GPU). Keep-fused MoE experts are split into per-expert keys once the owned units are assembled.
+    Every rank walks all layers together (the gather is a collective) but keeps only the ones it owns,
+    so host memory holds just this rank's share. Fused MoE experts are split into per-expert keys at
+    the end.
     """
-    from .moe_utils import _split_fused_experts_state_dict
-
     rank, world = _dist.rank(), _dist.size()
     id_to_name = {id(m): n for n, m in model.named_modules()}
-    units = _enumerate_export_units(model, id_to_name)
+    units = _enumerate_export_units(model)
     my_sd: dict[str, torch.Tensor] = {}
-    for i, (modules, _is_root) in enumerate(units):
+    for i, modules in enumerate(units):
         owned = _gather_unit(
             modules, id_to_name, is_owner=(i % world == rank)
         )  # COLLECTIVE all ranks
         if owned is not None:
             my_sd.update(owned)
-    _split_fused_experts_state_dict(
-        my_sd, model
-    )  # only this rank's owned keep-fused keys are present
+    _split_packed_fused_experts(my_sd, model)  # only this rank's owned keep-fused keys are present
     return my_sd
 
 
 def _finalize_index(local_maps, export_dir, rank, world):
-    """gather_object the per-rank HF weight_maps to rank 0; write the standard index. No sidecars."""
-    from huggingface_hub.constants import SAFETENSORS_INDEX_FILE
-
+    """Collect every rank's tensor-to-file map onto rank 0 and write model.safetensors.index.json."""
     gathered: list[Any] | None
     if world > 1:
         gathered = [None] * world if rank == 0 else None
@@ -1622,11 +1591,7 @@ def _finalize_index(local_maps, export_dir, rank, world):
 
 
 def _parallel_write(my_sd, export_dir, max_shard_size):
-    """Distributed write: each rank writes its own safetensors shards concurrently.
-
-    Rank 0 then merges the single ``model.safetensors.index.json`` from the gathered per-rank
-    weight_maps.
-    """
+    """Each rank writes its own shard files at once; rank 0 then writes the combined index."""
     rank, world = _dist.rank(), _dist.size()
     local_map = _write_owned_shards(my_sd, export_dir, rank, max_shard_size)  # (weight_map, nbytes)
     _dist.barrier()
diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py
index 8bd52a53f3b..7dc55e23646 100644
--- a/modelopt/torch/quantization/utils/core_utils.py
+++ b/modelopt/torch/quantization/utils/core_utils.py
@@ -1050,19 +1050,16 @@ def fsdp2_aware_weight_update(root_model, modules_to_update, reshard=True):
 
 
 def _shard_start(param):
-    """Global offset along shard dim 0 for a ``Shard(0)`` DTensor on a 1-D FSDP mesh.
+    """Global row offset of this rank's ``Shard(0)`` slice (0 for a non-DTensor param).
 
-    Returns 0 for non-DTensor params. Uses torch's even-split convention: the first ``dim0 % world``
-    ranks get an extra row, so rank ``r``'s offset is ``r * (dim0 // world) + min(r, dim0 % world)``.
-    Correct for both even and uneven sharding.
+    Sharding is always even here (``fsdp2_shard_local_pack`` rejects uneven), so it is just
+    ``rank * rows_per_rank``.
     """
     if not isinstance(param, DTensor):
         return 0
-    dim0 = param.shape[0]  # global size
     world = param.device_mesh.size()
     r = param.device_mesh.get_local_rank()
-    base, rem = divmod(dim0, world)
-    return r * base + min(r, rem)
+    return r * (param.shape[0] // world)
 
 
 def _rebuild_fsdp_param_from_shard(old_fp, packed_local):
@@ -1135,21 +1132,13 @@ def _rewrap_scale_buffers_shard0(module, captured, local_dim0):
     """
     info = next(iter(captured.values()))
     mesh, placements = info.mesh, info.placements
-    # The scale's dim-0 matches the weight's global dim-0 (per-row / per-expert). Pass the true global
-    # shape+stride to from_local: without it, from_local infers global = local * world (even-shard
-    # assumption), so under UNEVEN sharding ranks disagree on the global shape and full_tensor()
-    # deadlocks. sharded_param DTensors (weights) already carry the right shape via their spec.
-    global_dim0 = info.old._orig_size[0]
     for bname, buf in list(module._buffers.items()):
         if buf is None or isinstance(buf, DTensor) or buf.dim() == 0:
             continue
+        # A buffer whose dim-0 matches the shard row/expert count is sharded; wrap it Shard(0) so a
+        # later full_tensor() rebuilds it. Even sharding, so from_local infers the global shape.
         if buf.shape[0] == local_dim0:
-            local = buf.contiguous()
-            global_shape = torch.Size((global_dim0, *local.shape[1:]))
-            global_stride = torch.empty(global_shape, device="meta").stride()
-            module._buffers[bname] = DTensor.from_local(
-                local, mesh, placements, shape=global_shape, stride=global_stride
-            )
+            module._buffers[bname] = DTensor.from_local(buf.contiguous(), mesh, placements)
 
 
 @contextmanager
@@ -1220,36 +1209,6 @@ def fsdp2_shard_local_pack(root_model, module):
         module.__dict__.pop("_shard_local_start", None)
 
 
-@contextmanager
-def materialize_fsdp2_root(model: nn.Module):
-    """Unshard a sharded FSDP2 root's own params (embed/lm_head/norm) for a calibration forward.
-
-    The calibration loop calls ``model.forward(**batch)`` directly (``dataset_utils._forward_loop``)
-    rather than ``model(**batch)``, so ``nn.Module.__call__`` is bypassed and the root's FSDP2
-    forward pre-hook never fires. Its own params (embed/lm_head/norm) stay sharded DTensors and the
-    forward hits ``aten.embedding: mixed Tensor and DTensor``. Decoder layers are unaffected: they are
-    called as ``layer(...)`` inside ``forward``, so their pre-hooks fire and they unshard normally.
-
-    Unshard the root up front so its params are full tensors, then reshard on exit. The bypass also
-    skips the root's post-forward hook, so nothing reshards it mid-calibration and a single unshard
-    holds across all batches. Cheap: only the root's own param group is gathered, not the decoder
-    layers. No-op for non-FSDP2 or already-replicated roots.
-    """
-    root_sharded = False
-    if isinstance(model, FSDPModule):
-        pg = fully_shard.state(model)._fsdp_param_group
-        root_sharded = pg is not None and pg.is_sharded
-    if root_sharded:
-        with enable_fake_quant(model):
-            model.unshard()
-    try:
-        yield
-    finally:
-        if root_sharded:
-            with enable_fake_quant(model):
-                model.reshard()
-
-
 def update_quant_cfg_with_kv_cache_quant(
     quant_cfg: dict[str, Any], kv_cache_quant_cfg: list[QuantizerCfgEntry]
 ) -> dict[str, Any]: