diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 6687ebd31ea..4210b12228a 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** @@ -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) ^^^^^^^^^^^^^^^^^ 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..1a7fbe2f9cd 100644 --- a/modelopt/torch/export/distribute.py +++ b/modelopt/torch/export/distribute.py @@ -303,3 +303,315 @@ 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 .model_utils import TiedWeightMap + from .quant_utils import postprocess_state_dict + + # Name-based tied-weight dedup, matching the gather path. ``fully_shard`` splits a shared + # nn.Parameter into distinct per-module shards, so a declared tie surfaces here as two + # independent DTensors and both sides would be written. The map is derived from names and is + # identical on every rank, so the alias is dropped consistently; postprocess_state_dict skips a + # group whose canonical is missing from this rank's slice rather than orphaning the alias. + tied_map = TiedWeightMap(model) + + export_dir = Path(export_dir) + rank, world = dist.get_rank(), dist.get_world_size() + device = torch.device(f"cuda:{torch.cuda.current_device()}") + + # This writer targets a FULLY FSDP2-sharded model -- no expert parallelism. FSDP2 + # shards the fused 3-D expert weight on dim 0 (the expert axis), so fused experts are just + # DTensors sharded across the fsdp mesh: each per-expert key's TRUE global index comes from the + # DTensor's own dim-0 offset (``_split_local_fused_module`` reads it), so there is no EP block base + # (ep_rank=0) and no dp-replica dedup (dp_idx=0 -> every rank writes its own experts). The EP-hybrid + # (a2a experts) path -- which resolves ep_rank/dp_idx from the parallel-group provider -- is out of + # scope here. + ep_rank, 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 + + # Re-materialize IN PLACE, not as `{k: _to_normal(v) for ...}`. The comprehension keeps the ENTIRE + # original dict alive while building a full second copy -> a transient 2x of the whole state dict. + # That is fine at 30B (24GB -> 48GB) but OOMs at large MoE: at 235B the write-point state dict is + # ~157GB/rank (fused expert weights are sharded, but the per-expert NVFP4 scale tensors are + # replicated on every rank), and doubling it hits ~314GB > GPU. Overwriting each key frees the old + # value immediately, so peak stays ~1x. (Import of ep-dp's in-place/streaming write handling.) + for _k in list(sharded_sd.keys()): + sharded_sd[_k] = _to_normal(sharded_sd[_k]) + + prefixes = _fused_experts_prefixes(sharded_sd) + prefix_set = set(prefixes) + + def _is_expert_key(key: str) -> bool: + for proj in _FUSED_PROJ: + for suffix in ("", "_weight_scale", "_weight_scale_2", "_input_scale"): + if key.endswith(f".experts.{proj}{suffix}"): + if key[: key.rfind(".experts.") + len(".experts")] in prefix_set: + return True + return False + + # Probe the DTensor placement of a sample expert + dense weight: directly shows what state the + # generation forwards left the sharded params in -- FSDP2 (Shard(0)/Replicate/gathered) or EP + # (expert dim-0 Shard, or a plain local shard for the a2a transport). + + + # Sync the shared per-module activation (input) scales across ranks (global max), so every + # expert of a module gets the same input_scale regardless of which rank owns it. + in_keys = sorted(k for k in sharded_sd if k.endswith("_input_scale") and _is_expert_key(k)) + if in_keys: + stacked = torch.stack( + [ + sharded_sd[k].detach().to(device=device, dtype=torch.float32).reshape(()) + for k in in_keys + ] + ) + dist.all_reduce(stacked, op=dist.ReduceOp.MAX) + for j, k in enumerate(in_keys): + sharded_sd[k] = stacked[j].clone() + + # (1) Experts: split this rank's local shard into per-expert keys (global idx); stays local. + # Skipped on dp-replica ranks under DP x EP (their experts are written by dp group 0). + local_sd: dict = {} + if write_experts: + for pi, prefix in enumerate(prefixes): + local_sd.update( + _split_local_fused_module( + prefix, sharded_sd, ep_rank=ep_rank, total_experts=total_experts + ) + ) + local_sd = postprocess_state_dict( + local_sd, maxbound, kv_cache_format, is_modelopt_qlora, tied_map=tied_map + ) + # In place (see above): avoid a transient 2x copy of the per-expert split. + for _k in list(local_sd.keys()): + local_sd[_k] = local_sd[_k].detach().contiguous() + + # (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, tied_map=tied_map + ) + 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) 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() 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 4ce60b192ee..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, @@ -283,3 +418,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 + + 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 77429b1cfaf..acd5fc268b2 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(): @@ -929,8 +935,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. @@ -992,14 +999,39 @@ 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 ( + 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 +1477,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, @@ -1550,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 @@ -1583,12 +1650,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, 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..3b4526ee634 --- /dev/null +++ b/tests/gpu/torch/export/test_distributed_hf_export.py @@ -0,0 +1,341 @@ +# 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. + +It also covers the two things the writer must not quietly drop on the way: a declared tied +weight still gets deduplicated, and options the path cannot honour are rejected rather than +ignored. + +Run with >=2 GPUs so the expert axis and the FSDP2 shard axis are both actually split. +""" + +import 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 + + +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" + )