Distributed no-gather HF checkpoint export for FSDP2

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 <tracking-uri>`` 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/<checkpoint basename>-<recipe name or --qformat>`` and can be overridden with ``--mlflow_experiment`` / ``--mlflow_run_name``.
 - Add ``--mlflow <tracking-uri>`` 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/<model basename>-<recipe name or quantization config>`` 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**
 
@@ -43,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 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)
 ^^^^^^^^^^^^^^^^^

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))
+
+