[OMNIML-6083] Write a sharded model's HF checkpoint without gathering it

CHANGELOG.rst

@@ -15,6 +15,7 @@ Changelog
 - GPTQ no longer runs a max-calibration forward when no activation quantizer needs one, and ``hf_ptq.py`` now sizes the default layerwise calibration batch at 8,192 tokens (16 samples at ``--calib_seq 512``) instead of 1. Layerwise GPTQ on HF fused-MoE models runs over 10x faster; pass ``--batch_size 1`` for the old batching, which pads no sample.
 - Add a checkpoint-specific NVIDIA Nemotron 3.5 Super VL AutoQuantize recipe with a 4.9-effective-bit NVFP4-MSE/FP8/BF16 search and fixed NVFP4 MTP experts. Select it under ``models/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/auto_quantize/``.
 - Add ``examples/hf_ptq/convert_ggml_to_int8.py`` to expand a checkpoint exported with a GGML recipe (IQ formats, ``q8_0``) into an INT8 weight-only compressed-tensors checkpoint. Stock vLLM serves it with its integer kernels at about half the BF16 weight memory, reproducing the GGML weights up to rounding their scales to the model dtype.
+- Add an opt-in no-gather HuggingFace export for FSDP2-sharded models: with ``MODELOPT_FSDP2_EXPORT=nogather``, ``export_hf_checkpoint`` sends each tensor to a single owner rank and every rank writes its own shards in parallel, instead of gathering the full state dict on one host. MoE experts stay sharded through the export, so each rank only materializes its own experts.
 
 *Speculative Decoding*
 
@@ -41,6 +42,7 @@ Changelog
 - ``export_mcore_gpt_to_hf`` given a Hugging Face Hub model ID no longer writes the tokenizer, processor, generation config or remote-code files. Pass a local directory instead, e.g. from ``modelopt.torch.export.ensure_local_checkpoint``.
 - ``modelopt.torch.utils.distributed.fsdp2_wrap`` moved to ``modelopt.torch.utils.plugins.model_load_utils``, with the other loaders. The old import path still works, but calling it now needs the ``hf`` extra.
 - Diffusion HF exports now omit attention quantizer buffers by default, except for FLUX.1, which retains enabled Q/K/V ``_amax`` keys for runtime compatibility; other consumers of those keys can opt in with ``export_attention_quantizers=True`` (CLI: ``--extra-param export_attention_quantizers=true``). Disabled attention quantizers remain excluded, and diffusion components, including LTX-2, no longer emit KV-cache quantization metadata.
+- ``export_hf_checkpoint`` now reads how a model is distributed from the plan ``load_pretrained`` records on it, and raises for an FSDP2-sharded model, or one in a multi-rank run, that carries none. If you distribute a model yourself, for example with ``fsdp2_wrap``, which records no plan, record one with ``modelopt.torch.utils.distributed.set_distribution_plan`` before calibrating or exporting it.
 
 **Deprecations**
 

examples/hf_ptq/README.md

@@ -646,6 +646,8 @@ See [Recipe-based Quantization](#recipe-based-quantization) for the recipe forma
 
 > *Performance Note: FSDP2 is designed for training workloads and may result in longer calibration and export times. For faster calibration, maximize the batch size based on available GPU memory and choose the right number of GPUs to avoid unnecessary communication.*
 
+> *Set `MODELOPT_FSDP2_EXPORT=nogather` to have every rank write its own shards of the checkpoint in parallel instead of gathering the full state dict on one host.*
+
 ## Evaluate Accuracy
 
 ### TensorRT-LLM Validation

modelopt/torch/export/unified_export_hf.py

@@ -19,6 +19,7 @@
 import copy
 import importlib
 import json
+import os
 import re
 import tempfile
 import warnings
@@ -91,7 +92,12 @@
 from modelopt.torch.quantization.utils.core_utils import has_accelerate_offload
 from modelopt.torch.utils import print_rank_0
 from modelopt.torch.utils.dataset_utils import _disable_use_cache
-from modelopt.torch.utils.distributed import is_fsdp2_model
+from modelopt.torch.utils.distributed import (
+    DistributionPlan,
+    broadcast_from,
+    distribution_plan,
+    is_fsdp2_model,
+)
 from modelopt.torch.utils.perf import maybe_clear_cuda_cache
 
 try:
@@ -949,6 +955,7 @@ def _process_quantized_modules(
     model: nn.Module,
     dtype: torch.dtype,
     is_modelopt_qlora: bool = False,
+    keep_fused_experts: bool = False,
 ) -> None:
     """Process all quantized modules in model, export weights in-place.
 
@@ -964,14 +971,17 @@ def _process_quantized_modules(
     # No per-module dedup cache: tied duplicates are dropped by name in postprocess_state_dict.
     # The handlers pack whatever weight they are handed, and nothing here materializes a shard --
     # FSDP2 models go through collect_export_tensors instead, which packs inside a gather window.
-    assert not is_fsdp2_model(model), (
+    # The no-gather writer is the deliberate exception: it packs each rank's LOCAL shard and never
+    # materializes a full tensor, which is signalled by keep_fused_experts.
+    assert not is_fsdp2_model(model) or keep_fused_experts, (
         "_process_quantized_modules cannot pack a sharded model; use collect_export_tensors"
     )
     ctx = ExportContext(
         model=model,
         dtype=dtype,
         is_modelopt_qlora=is_modelopt_qlora,
         model_type=hf_model_type(model),
+        keep_fused_experts=keep_fused_experts,
     )
 
     try:
@@ -1031,13 +1041,22 @@ def _prepare_model_for_export(model, dtype, is_modelopt_qlora):
     return dtype, tied_map, quant_config
 
 
-def pack_quantized_weights(model, dtype, is_modelopt_qlora: bool = False) -> None:
-    """Quantize every module's weight in place, then rebuild the fused MoE linears."""
+def pack_quantized_weights(
+    model, dtype, is_modelopt_qlora: bool = False, keep_fused_experts: bool = False
+) -> None:
+    """Quantize every module's weight in place, then rebuild the fused MoE linears.
+
+    ``keep_fused_experts`` is for the no-gather distributed FSDP2 export: it leaves MoE experts
+    fused and sharded through the fold, so each rank materializes only its LOCAL experts instead
+    of all N on every rank. The per-expert split happens later, in the writer.
+    """
     # Deferred: modelopt.torch.quantization imports modelopt.torch.export, so importing the
     # per-model PTQ module at module scope here would risk an import cycle.
     from modelopt.torch.models.step3p5.modeling_ptq import _reconstruct_fused_moe_linear
 
-    _process_quantized_modules(model, dtype, is_modelopt_qlora)
+    _process_quantized_modules(
+        model, dtype, is_modelopt_qlora, keep_fused_experts=keep_fused_experts
+    )
     _reconstruct_fused_moe_linear(model)
 
 
@@ -1887,6 +1906,27 @@ def read_unplaced_weights(model: nn.Module, *, keys_only: bool = False) -> dict[
     return out
 
 
+def _write_base_config(model: nn.Module, export_dir: "Path | str") -> None:
+    """Write config.json + generation_config for the no-gather distributed export path.
+
+    ``distributed_save_hf_checkpoint`` writes only the weight shards, so rank 0 emits here what
+    ``save_pretrained`` would have. The weight half of ``save_pretrained`` must NOT run: the
+    weights are already on disk and the model may still be sharded.
+    """
+    # A config's save_pretrained writes only config.json: transformers reverts weight conversions
+    # in the model's, so the patch save_non_weight_artifacts takes is not needed here.
+    model.config.save_pretrained(export_dir)
+    if model.can_generate() and getattr(model, "generation_config", None) is not None:
+        try:
+            # Same strict-validate trip as every other save site: a config carrying top_p/top_k
+            # with do_sample unset (e.g. GLM-5.3-BF16) makes save_pretrained raise, and the
+            # export would then carry the source's generation_config.json instead of this one.
+            _sanitize_generation_config_for_save(model)
+            model.generation_config.save_pretrained(export_dir)
+        except Exception as gen_err:
+            warnings.warn(f"Could not save generation_config: {gen_err}")
+
+
 def export_hf_checkpoint(
     model: Any,
     dtype: torch.dtype | None = None,
@@ -1935,14 +1975,10 @@ def export_hf_checkpoint(
     # emits it from rank 0 alone, so on the other ranks a full read is host memory spent and thrown
     # away -- and the carried set is the large stuff. The KEY list is still resolved everywhere,
     # because get_quant_config runs per rank and the configs have to agree.
-    _writes_extra = (
-        not (
-            torch.distributed.is_available()
-            and torch.distributed.is_initialized()
-            and is_fsdp2_model(model)
-        )
-        or torch.distributed.get_rank() == 0
-    )
+    # How the model is distributed decides how it is exported. A diffusers pipeline is not a module,
+    # and nothing here distributes one.
+    plan = distribution_plan(model) if isinstance(model, nn.Module) else DistributionPlan()
+    _writes_extra = not plan.is_fsdp_dense or torch.distributed.get_rank() == 0
     _carried = read_unplaced_weights(model, keys_only=not _writes_extra)
     if _writes_extra and _carried:
         extra_state_dict = {**_carried, **(extra_state_dict or {})}
@@ -1983,10 +2019,11 @@ def export_hf_checkpoint(
         )
         return
 
-    is_fsdp2_sharded = (
-        torch.distributed.is_available()
-        and torch.distributed.is_initialized()
-        and is_fsdp2_model(model)
+    is_fsdp2_sharded = plan.is_fsdp_dense
+    # Rank 0's choice: the two writers issue different collectives, so ranks that read their
+    # environments differently would wait in collectives the others never join.
+    use_no_gather_writer = is_fsdp2_sharded and broadcast_from(
+        "MODELOPT_FSDP2_EXPORT", lambda: os.environ.get("MODELOPT_FSDP2_EXPORT") == "nogather"
     )
     # Not the global rank: a non-FSDP2 export writes the whole checkpoint from every process.
     rank = torch.distributed.get_rank() if is_fsdp2_sharded else 0
@@ -2017,6 +2054,42 @@ def export_hf_checkpoint(
             if getattr(model, "hf_quantizer", None) is not None:
                 model.hf_quantizer = None
             hf_quant_config = _revert_quant_config_names_best_effort(model, hf_quant_config)
+        elif use_no_gather_writer:
+            # Opt-in alternative to the streaming path below: never replicate a unit across ranks.
+            # Each tensor is sent to exactly one owner over batched point-to-point sends, instead of
+            # every rank entering a gather window and all but one discarding the result. Same
+            # outputs; the win grows with world size, where the replicated gather costs ~world x the
+            # bytes. Left opt-in so the default path is unchanged.
+            from .distributed_writer import distributed_save_hf_checkpoint
+
+            _qlora = kwargs.get("is_modelopt_qlora", False)
+            _dtype, _tied, hf_quant_config = _prepare_model_for_export(model, dtype, _qlora)
+            # main split preparation from packing, so pack explicitly. keep_fused_experts leaves
+            # MoE experts fused + sharded through the fold, so each rank folds only its LOCAL
+            # experts; the writer splits them after consolidation.
+            pack_quantized_weights(model, _dtype, _qlora, keep_fused_experts=True)
+            if getattr(model, "hf_quantizer", None) is not None:
+                model.hf_quantizer = None
+            reverted = distributed_save_hf_checkpoint(
+                model,
+                export_dir,
+                maxbound=448,  # the KV cache scale is amax / 448 for FP8 and NVFP4, as above
+                kv_cache_format=_get_kv_cache_postprocess_config(
+                    (hf_quant_config or {}).get("quantization", {})
+                ),
+                max_shard_size=max_shard_size,
+                is_modelopt_qlora=_qlora,
+                extra_state_dict=extra_state_dict,
+            )
+            if rank == 0:
+                _write_base_config(model, export_dir)
+                # As the streaming path below: every FSDP2 rank holds the same ModelOpt state.
+                if save_modelopt_state and ModeloptStateManager.is_converted(model):
+                    torch.save(modelopt_state(model), export_dir / _MODELOPT_STATE_SAVE_NAME)
+                # The weights keep their in-memory names when any rank could not revert them, and
+                # the quant config has to name the modules as the weights do.
+                if reverted:
+                    hf_quant_config = _revert_quant_config_names_best_effort(model, hf_quant_config)
         elif is_fsdp2_sharded:
             # FSDP2 multi-rank: stream each rank's owned units straight to its own shard files, so
             # a rank buffers its own share of the model rather than the whole checkpoint, and the
@@ -2106,8 +2179,9 @@ def export_hf_checkpoint(
         if rank == 0:
             _write_hf_export_config(model, hf_quant_config, export_dir)
             copy_off_index_safetensors(_source_checkpoint(model), export_dir)
-            if not (_offloaded or is_fsdp2_sharded):
-                # The streaming paths already did, from save_non_weight_artifacts.
+            if not (_offloaded or (is_fsdp2_sharded and not use_no_gather_writer)):
+                # The streaming paths already did, from save_non_weight_artifacts; the no-gather
+                # writer does not.
                 _copy_non_model_files_from_source(model, export_dir)
 
     except Exception as e: