diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 904eac00fa3..1e01da0214f 100755 --- a/CHANGELOG.rst +++ b/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** diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index 4bdc6b55949..2336120db73 100755 --- a/examples/hf_ptq/README.md +++ b/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 diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index 80b19aafdc6..39af7d9bb0f 100644 --- a/modelopt/torch/export/distributed_writer.py +++ b/modelopt/torch/export/distributed_writer.py @@ -13,11 +13,11 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Pieces of a no-gather distributed HF checkpoint writer. +"""No-gather distributed HF checkpoint writer for FSDP2. Every tensor is sent to exactly one owner rank over batched point-to-point sends and written from there, so no rank ever materializes the whole model and no unit is -replicated across ranks. +replicated across ranks. Selected by ``MODELOPT_FSDP2_EXPORT=nogather``. """ import json @@ -28,6 +28,43 @@ import torch import torch.nn as nn +from .moe_utils import _FUSED_PROJ, _fused_experts_prefixes, _split_local_fused_module + + +def distributed_save_hf_checkpoint( + model: nn.Module, + export_dir: "str | Path", + maxbound: float, + kv_cache_format: "str | dict[str, dict[str, str]] | None", + max_shard_size: "str | int" = "10GB", + is_modelopt_qlora: bool = False, + extra_state_dict: "dict[str, torch.Tensor] | None" = None, +) -> bool: + """No-gather distributed HF export. + + Thin wrapper: a caller may be in ``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 the version- + counted ops in the write path 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. + + Returns: + Whether the tensors carry their original checkpoint names. The reverse conversion is all or + nothing across ranks, so this is the same on every rank. + """ + with torch.inference_mode(False), torch.no_grad(): + return _distributed_save_hf_checkpoint_impl( + model, + export_dir, + maxbound, + kv_cache_format, + max_shard_size, + is_modelopt_qlora, + extra_state_dict, + ) + def _even_bins(keys: list, sizes: dict, n: int, max_bytes: int) -> list: """Partition ``keys`` into ~``n`` EVEN-sized bins (by bytes), each <= ``max_bytes`` where possible. @@ -574,3 +611,228 @@ def write_index() -> bool: # Also the closing barrier: no rank returns before the index is on disk. all_reduce("index written", write_index, all) + + +def _distributed_save_hf_checkpoint_impl( + model: nn.Module, + export_dir: "str | Path", + maxbound: float, + kv_cache_format: "str | dict[str, dict[str, str]] | None", + max_shard_size: "str | int" = "10GB", + is_modelopt_qlora: bool = False, + extra_state_dict: "dict[str, torch.Tensor] | None" = None, +) -> bool: + """Distributed HF safetensors export -- no rank-0 full-model host-RAM gather. + + Writes an already-processed, FSDP2-sharded model to ``export_dir`` as HF safetensors, entirely + distributed. No rank ever holds the whole model: + + 1. Fused expert weights (when present): each rank splits its LOCAL shard into per-expert keys + with global indices and keeps them in place -- no communication. Experts are therefore + already disjoint across ranks. dp-replica ranks under DP x EP skip this. + 2. Dense / non-expert weights: each fqn is assigned to exactly ONE owner rank and moved there + (:func:`_materialize_owned`), so every tensor becomes whole on its owner while the set stays + spread across ranks. Replicated and plain values need no transport; dim-0 shards travel + point-to-point, bucketed per peer. + 3. The reverse weight conversion runs per-rank on those whole tensors -- it cannot run while + they are sharded, because transformers' conversion ops are written against whole tensors + and have no way to express a placement. + 4. Every rank bin-packs its own owned set into even shards, numbered into a global range, and + writes them directly with ``safetensors.save_file``. Rank 0 unions the weight index. + + No rank holds the model. On top of its own shard of it, which the model keeps, a rank holds what + it writes -- its experts, split into per-expert copies, and the dense tensors it owns, assembled + as :func:`_materialize_owned` describes -- and, from a caller in inference mode, a copy of each + inference tensor in its shard. Placements the point-to-point transport does not cover are + rejected rather than routed through an untested slower path. + """ + import torch.distributed as dist + from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict + from torch.distributed.tensor import DTensor + + from modelopt.torch.utils.distributed import all_reduce + + 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() + + # 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, dp_idx = 0, 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 inference tensor as a NORMAL one UP FRONT. A caller in inference mode hands + # over inference tensors; the inference flag is STICKY -- clone/reshape/detach/slice all + # propagate it -- and the expert split, scale reduce and 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. Anything else is used as it is. + def _to_normal(v): + if isinstance(v, DTensor): + _loc = v.to_local() + if not _loc.is_inference(): + return v + _new = torch.empty_like(_loc) + _new.copy_(_loc) + # Pass the ORIGINAL global shape/stride: from_local() otherwise re-infers them as + # local x world, which is only right for an EVEN shard. This helper runs over every + # DTensor in the state dict just before classification, so with an uneven shard + # (b_proj: 96 rows over 64 ranks) or a zero-size one (mlp_res_proj) each rank derives a + # DIFFERENT global shape for the same tensor -- and then classifies it differently, + # which breaks the export with no hint of where the shape went wrong. + return DTensor.from_local( + _new, + v.device_mesh, + v.placements, + run_check=False, + shape=v.shape, + stride=v.stride(), + ) + if v.is_inference(): + _new = torch.empty_like(v) + _new.copy_(v) + return _new + return v + + # A copy adds to this rank's memory rather than replacing anything: the values are views of the + # model's own parameters and buffers, which keep their data. Copying every value would hold the + # rank's whole shard twice; copying only inference tensors costs a caller outside inference mode + # (hf_ptq exports under no_grad) nothing. + 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 + + # 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. Reduced on the + # scales' own device, as the shard transport moves every tensor: the GPU in a GPU run, the CPU + # under gloo. + in_keys = sorted(k for k in sharded_sd if k.endswith("_input_scale") and _is_expert_key(k)) + # Every rank must reduce the same keys: a rank holding none would skip the reduce and leave the + # others waiting in it, and lists of different lengths would mismatch its size, neither saying + # why. So the ranks agree on the keys first, and a difference fails on every rank, naming it. + key_lists = all_reduce("expert input-scale keys", lambda: [in_keys], _in_rank_order) + if any(keys != in_keys for keys in key_lists): + differ = sorted(set().union(*key_lists) - set(key_lists[0]).intersection(*key_lists)) + raise RuntimeError( + "The ranks hold different expert input scales, so they cannot reduce them to their " + f"max; first few held by only some ranks: {differ[:5]}" + ) + if in_keys: + stacked = torch.stack([sharded_sd[k].detach().reshape(()).float() 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 prefix in 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: the split's tensors are its own copies, so one made contiguous replaces the old. + for _k in list(local_sd.keys()): + local_sd[_k] = local_sd[_k].detach().contiguous() + # At this point local_sd holds ONLY this rank's (already split, per-expert) experts -- capture the + # set so _materialize_owned can tell disjoint experts from replicated dense. + expert_key_set = set(local_sd) + + # (2) Dense / non-expert: leave sharded for now -- FSDP2 leaves DTensors (dim-0 shard); EP leaves + # plain replicated tensors. Do NOT gather to rank 0; _materialize_owned below moves each tensor to + # a single OWNER rank instead, so the set stays spread across ranks. 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) + # No explicit gc.collect() here: the del drops the last references and CPython frees acyclic + # objects immediately. Collecting a ~900-entry DTensor state dict measured 0.89s of the 0.89s + # this phase used to cost -- and because it ran at different speeds per rank, the next + # collective absorbed the skew and looked like gather cost. + del sharded_sd, dense_sd + # (3) Whole-tensor redistribution -> reverse conversion -> write. Each tensor ends up WHOLE on + # exactly one rank (never gathered to a single rank), which is what lets the reverse run: the + # transformers conversion ops are written against whole tensors and cannot express a placement. + # The reverse must precede the shard layout, since it renames keys and splits tensors. + _materialize_owned(local_sd, expert_key_set, rank, world, _merge_groups(model, local_sd)) + reverted = _revert_whole_sd(model, local_sd) + + # Tensors the model does not own -- whatever the caller hands in via extra_state_dict, e.g. + # weights read straight off the source checkpoint that the FSDP2 loader drops. They are plain + # CPU tensors, already whole and identical on every rank, so ONE rank adds them to its + # write set; the index union and + # the cross-rank duplicate check below then treat them like any other owned key. Rank 0 takes + # them wholesale: they are small next to the model, and spreading them would cost a collective + # to agree on the split. They go in AFTER the reverse conversion because they carry their + # original hub names and were never converted -- reverting them would rename them wrongly. This + # mirrors the gather path, which merges extra_state_dict into the already-reverted dict. + def check_no_overlap() -> None: + if rank != 0 or not extra_state_dict: + return + collisions = sorted(set(extra_state_dict) & set(local_sd)) + if collisions: + raise RuntimeError( + f"extra_state_dict overlaps the exported model weights on {len(collisions)} " + f"key(s), so the checkpoint would be ambiguous; first few: {collisions[:5]}" + ) + + # The export entry hands the tensors it carries over to rank 0 alone, so whether a rank holds + # any cannot decide whether it joins: every rank does, and a collision rank 0 finds then raises + # on all of them, instead of leaving the others in the write's first collective. + all_reduce("extra_state_dict overlap", check_no_overlap, lambda _: None) + if extra_state_dict and rank == 0: + local_sd.update(extra_state_dict) + _write_even_shards(local_sd, export_dir, max_shard_size, rank, world) + return reverted diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 80448aa2ba0..4e361b9c042 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/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,7 +971,9 @@ 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( @@ -972,6 +981,7 @@ def _process_quantized_modules( 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: 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..61225f0392d --- /dev/null +++ b/tests/gpu/torch/export/test_distributed_hf_export.py @@ -0,0 +1,280 @@ +# 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. + +With ``MODELOPT_FSDP2_EXPORT=nogather``, ``export_hf_checkpoint`` routes an FSDP2-sharded model under +``torch.distributed`` to ``distributed_save_hf_checkpoint``: every rank writes its own final +safetensors shards in parallel and rank 0 unions the weight index, 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 source's own (tokenizer, vocab, chat template), which the + exporter copies over. + +The dense model ties its embeddings, so the expected keys also check that the writer drops the +alias. What needs no GPU -- tensors handed in through ``extra_state_dict``, a linear left +unquantized (the ``lm_head`` every case there exports), options the path refuses -- is checked on +CPU, in ``tests/unit/torch/export/test_distributed_writer.py``. + +Run with >=2 GPUs so the expert axis and the FSDP2 shard axis are both actually split. +""" + +import json +import os +from functools import partial +from pathlib import Path +from unittest import mock + +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 ( + DistributionPlan, + fsdp2_wrap, + is_fsdp2_model, + set_distribution_plan, +) + +# 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 one more file to write for 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 _no_gather_export(): + """Opt the FSDP2-only model into the no-gather writer, scoped to this export. + + The workers are a persistent pool, so a bare ``os.environ`` write would leak into later tests. + """ + return mock.patch.dict(os.environ, {"MODELOPT_FSDP2_EXPORT": "nogather"}) + + +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) + set_distribution_plan( + model, DistributionPlan(size, dense="fsdp") + ) # wrapped by hand, so recorded here, as load_pretrained would + # The no-gather writer is chosen by the plan recorded above plus the opt-in below. It writes + # each rank's FSDP2 shards, so check that fsdp2_wrap sharded the model: unsharded, the test + # would not exercise it. + 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() + + with _no_gather_export(): + 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-tied", "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 + # fully_shard splits a tied parameter into per-module shards, so only the names still match: + # the source checkpoint holds no lm_head.weight, and the writer must not write one either. + kwargs = {**TINY_KWARGS, **(TINY_MOE_KWARGS if moe else {"tie_word_embeddings": True})} + src_dir = Path(make_dir(tmp_path, with_tokenizer=True, **kwargs)) + export_dir = tmp_path / "export" + + # The source's non-weight files (tokenizer, vocab, chat template, ...), which the exporter must + # copy into the export itself: nothing else writes them there. + 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" + } + + 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" + if moe: + # Every expert, not just some weight: the experts are folded on their sharded weights, and a + # fold FSDP2 does not track is undone by its next reshard, leaving bf16 experts beside their + # scales. + experts = [k for k in expected if ".experts." in k and k.endswith(".weight")] + unquantized = [k for k in experts if exported[k][0] not in _QUANTIZED_DTYPES] + assert experts and not unquantized, f"experts exported unquantized: {unquantized[:5]}" + + # ---- 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, ...) is in the export; only the weight files are rewritten. + lost = sorted(sidecars - exported_files) + assert not lost, f"sidecar file(s) missing from the export: {lost}" + + # 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 + ) diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 85fb5fb82c0..e16de11bf57 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -13,33 +13,41 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""The pieces of the no-gather HF checkpoint writer (``modelopt.torch.export.distributed_writer``). +"""The no-gather FSDP2 HF checkpoint writer (``modelopt.torch.export.distributed_writer``) on CPU. The writer moves every tensor to exactly one owner rank and writes it from there, so what it can get wrong is structural: a tensor missing, duplicated, mis-assembled from its shards, renamed, or -left out of the weight index. These tests drive each piece directly, on CPU. +left out of the weight index. The helpers are tested directly; the writer itself runs end to end on +a gloo process group -- one rank in process, two spawned ranks where shards actually travel -- and +the written checkpoint is reloaded and compared with the source checkpoint or with what the +single-process export writes for the same model. """ import contextlib +import copy import json import math +import os import types from unittest import mock import pytest import torch +import torch.distributed as dist from _test_utils.torch.distributed.utils import bounded_collectives from torch.distributed.device_mesh import init_device_mesh +from torch.distributed.fsdp import fully_shard from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor from torch.multiprocessing.reductions import StorageWeakRef pytest.importorskip("safetensors") pytest.importorskip("transformers") from safetensors.torch import load_file -from transformers import AutoModelForCausalLM, LlamaConfig, MixtralConfig +from transformers import AutoModelForCausalLM, LlamaConfig, MixtralConfig, Qwen2MoeConfig from transformers.core_model_loading import Chunk, WeightConverter -from modelopt.torch.export import quant_aware_conversion +import modelopt.torch.quantization as mtq +from modelopt.torch.export import distributed_writer, quant_aware_conversion from modelopt.torch.export.distributed_writer import ( _even_bins, _is_full_world_shard0, @@ -48,7 +56,27 @@ _nbytes, _revert_whole_sd, _write_even_shards, + distributed_save_hf_checkpoint, ) +from modelopt.torch.export.quant_aware_conversion import revert_weight_conversion_quant_aware +from modelopt.torch.export.unified_export_hf import ( + _export_transformers_checkpoint, + export_hf_checkpoint, +) +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan + + +def _write(model, export_dir, **kwargs): + distributed_save_hf_checkpoint( + model, export_dir, maxbound=448, kv_cache_format=None, max_shard_size="8KB", **kwargs + ) + + +def _nogather_export(model, export_dir, knob="nogather", **kwargs): + with mock.patch.dict(os.environ, {"MODELOPT_FSDP2_EXPORT": knob}): + export_hf_checkpoint( + model, dtype=torch.bfloat16, export_dir=export_dir, max_shard_size="16KB", **kwargs + ) def _raw_bytes(tensor): @@ -78,6 +106,16 @@ def _load_checkpoint(export_dir): return tensors +def _save(config, path): + torch.manual_seed(0) + AutoModelForCausalLM.from_config(config, dtype=torch.bfloat16).save_pretrained(path) + return path + + +def _load(path): + return AutoModelForCausalLM.from_pretrained(path, dtype=torch.bfloat16).eval() + + def _mixtral_config(**overrides): """A tiny Mixtral: fused experts in memory, per-expert ``block_sparse_moe`` names on the hub.""" kwargs = { @@ -108,6 +146,67 @@ def _llama_config(): ) +def _qwen2_moe_config(): + """A tiny Qwen2-MoE: four fused experts and a shared expert in one decoder layer.""" + return Qwen2MoeConfig( + vocab_size=64, + hidden_size=32, + intermediate_size=32, + moe_intermediate_size=32, + shared_expert_intermediate_size=32, + num_hidden_layers=1, + num_attention_heads=4, + num_key_value_heads=2, + num_experts=4, + num_experts_per_tok=2, + max_position_embeddings=32, + ) + + +def _fsdp2_shard(model, mesh, wrap_root=True): + """Shard each decoder layer, and the root too unless ``wrap_root`` is off (as expert parallelism + leaves it, sharding only the dense layers), and record that the model is FSDP2-sharded.""" + for layer in model.model.layers: + fully_shard(layer, mesh=mesh) + if wrap_root: + fully_shard(model, mesh=mesh) + set_distribution_plan(model, DistributionPlan(mesh.size(), dense="fsdp")) + return model + + +# FP8 weights, with an FP8 KV cache in layer 0 and an NVFP4 one in layer 1: a mixed-precision KV +# cache, whose per-layer formats the export reads from the quantization config. +_FP8_MIXED_KV_CFG = copy.deepcopy(mtq.FP8_DEFAULT_CFG) +_FP8_MIXED_KV_CFG["quant_cfg"] += [ + { + "quantizer_name": "*layers.0.*[kv]_bmm_quantizer", + "cfg": {"num_bits": (4, 3), "axis": None}, + "enable": True, + }, + { + "quantizer_name": "*layers.1.*[kv]_bmm_quantizer", + "cfg": { + "num_bits": (2, 1), + "block_sizes": {-1: 16, "type": "dynamic", "scale_bits": (4, 3)}, + "axis": None, + }, + "enable": True, + }, +] + + +def _quantize_with_reference(model, quant_cfg): + """Quantize ``model`` and return what the single-process export writes for a copy of it.""" + torch.manual_seed(1) + batch = torch.randint(0, model.config.vocab_size, (2, 8)) + mtq.quantize(model, quant_cfg, lambda m: m(batch)) + ref = copy.deepcopy(model) + want, _ = _export_transformers_checkpoint(ref, torch.bfloat16) + want = revert_weight_conversion_quant_aware(ref, dict(want)) + assert any(v.dtype in (torch.float8_e4m3fn, torch.uint8) for v in want.values()) + return want + + # ------------------------------------------------------------------ pure helpers @@ -155,29 +254,6 @@ def test_is_full_world_shard0_answers_from_the_mesh_memo(gloo_world_of_one): assert _is_full_world_shard0(tensor, 1, {id(gloo_world_of_one): False}) is False -def test_materialize_owned_leaves_every_tensor_whole_and_plain(gloo_world_of_one): - mesh = gloo_world_of_one - full = torch.arange(12.0).reshape(4, 3) - expert_key, expert = "model.layers.0.mlp.experts.0.up_proj.weight", torch.full((2,), 7.0) - sd = { - "plain": torch.ones(3), - "replicated": distribute_tensor(full, mesh, [Replicate()]), - "sharded": distribute_tensor(full, mesh, [Shard(0)]), - "empty": distribute_tensor(torch.zeros(0, 4, dtype=torch.bfloat16), mesh, [Shard(0)]), - expert_key: expert, - } - - _materialize_owned(sd, {expert_key}, 0, 1) - - assert set(sd) == {"plain", "replicated", "sharded", "empty", expert_key} - assert not any(isinstance(v, DTensor) for v in sd.values()) - assert torch.equal(sd["plain"], torch.ones(3)) - assert torch.equal(sd["replicated"], full) - assert torch.equal(sd["sharded"], full) - assert sd["empty"].shape == (0, 4) and sd["empty"].dtype == torch.bfloat16 - assert sd[expert_key] is expert # experts are already disjoint per rank: left alone - - def _padded_shard(mesh): """This rank's 2 rows of a 3-row tensor: a padded local chunk, which would concatenate wrong.""" return DTensor.from_local( @@ -328,42 +404,6 @@ def test_materialize_owned_moves_each_shard_to_its_owner_over_two_ranks(cpu_dist cpu_dist_workers.run(_materialize_two_ranks_job) -def test_revert_whole_sd_keeps_a_state_dict_with_nothing_to_revert(): - """A model without a conversion mapping gets its input dict back from the reverse pass.""" - model = AutoModelForCausalLM.from_config(_llama_config()) - sd = dict(model.state_dict()) - before = dict(sd) - - assert _revert_whole_sd(model, sd) is True - - assert sd == before - - -def test_revert_whole_sd_renames_to_hub_names_in_place(): - model = AutoModelForCausalLM.from_config(_mixtral_config()) - gate, w1, w3, w2 = (torch.full((2, 2), float(i)) for i in range(4)) - sd = { - "model.layers.0.mlp.gate.weight": gate, - "model.layers.0.mlp.experts.0.gate_proj.weight": w1, - "model.layers.0.mlp.experts.0.up_proj.weight": w3, - "model.layers.0.mlp.experts.0.down_proj.weight": w2, - } - same_dict = sd - - assert _revert_whole_sd(model, sd) is True - - assert sd is same_dict - hub = "model.layers.0.block_sparse_moe" - assert set(sd) == { - f"{hub}.gate.weight", - f"{hub}.experts.0.w1.weight", - f"{hub}.experts.0.w3.weight", - f"{hub}.experts.0.w2.weight", - } - assert torch.equal(sd[f"{hub}.experts.0.w1.weight"], w1) - assert torch.equal(sd[f"{hub}.experts.0.w2.weight"], w2) - - def _one_rank_cannot_revert_job(rank, world_size): # Rank 0's subset reverts; rank 1's fails, as a merge whose group ownership split does. model = AutoModelForCausalLM.from_config(_mixtral_config()) @@ -552,3 +592,321 @@ def test_a_name_two_ranks_both_write_fails_the_write_on_every_rank(cpu_dist_work cpu_dist_workers.run(_cross_rank_collision_job, tmp_dir=tmp_path) assert not list(tmp_path.glob("*.safetensors")) + + +# ------------------------------------------------------------------ the writer, one rank + + +@pytest.mark.parametrize("case", ["tied", "untied", "inference-mode", "extra-state-dict"]) +def test_distributed_save_writes_the_hub_checkpoint_of_a_sharded_moe( + gloo_world_of_one, tmp_path, case +): + """Fused experts are split, every name is reverse-converted, and a declared tie is dropped. + + A caller in inference mode hands the writer inference tensors, which its writes reject. Extra + tensors already carry hub names, so the reverse conversion must not rename them: the first is + named like an in-memory Mixtral expert weight, which the reverse would turn into + ``block_sparse_moe.experts.0.w1``. Only an inference tensor is copied before the transfer: the + dense shards reach it as the model holds them, since a copy would only add to the rank's memory. + """ + src = _save(_mixtral_config(tie_word_embeddings=case == "tied"), tmp_path / "src") + model = _fsdp2_shard(_load(src), gloo_world_of_one) + shards = {p.to_local().untyped_storage().data_ptr() for p in model.parameters()} + moved = [] # per dense DTensor reaching the transfer: whether it is the model's own shard + + def materialize_owned(local_sd, *args): + moved.extend( + v.to_local().untyped_storage().data_ptr() in shards + for v in local_sd.values() + if isinstance(v, DTensor) + ) + return _materialize_owned(local_sd, *args) + + want = load_file(str(src / "model.safetensors")) + extra = None + if case == "extra-state-dict": + extra = { + "model.layers.9.mlp.experts.0.gate_proj.weight": torch.arange( + 64, dtype=torch.bfloat16 + ).reshape(8, 8), + "model.mtp.norm.weight": torch.full((8,), 3.0, dtype=torch.bfloat16), + } + want.update(extra) + elif case == "inference-mode": + want["model.norm.weight_scale"] = torch.tensor(0.5) + + with ( + torch.inference_mode(case == "inference-mode"), + mock.patch.object(distributed_writer, "_materialize_owned", materialize_owned), + ): + if case == "inference-mode": + # Made in inference mode, like the scales an inference-mode pack registers. + model.model.norm.register_buffer("weight_scale", torch.full((), 0.5)) + _write(model, tmp_path / "export", extra_state_dict=extra) + + assert len(list((tmp_path / "export").glob("*.safetensors"))) > 1 + _assert_byte_identical(_load_checkpoint(tmp_path / "export"), want) + assert moved and all(moved), moved + + +@pytest.mark.parametrize("wrap_root", [True, False], ids=["root-wrapped", "layers-only"]) +@pytest.mark.parametrize( + "quant_cfg", + [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG, _FP8_MIXED_KV_CFG], + ids=["fp8", "nvfp4", "fp8-mixed-kv-cache"], +) +def test_nogather_export_writes_what_the_single_process_export_writes( + gloo_world_of_one, tmp_path, quant_cfg, wrap_root +): + src = _save(_llama_config(), tmp_path / "src") + sidecar = src / "chat_template.jinja" # a source file only the exporter carries over + sidecar.write_text("{{ messages }}") + model = _load(src) + want = _quantize_with_reference(model, quant_cfg) + + _nogather_export(_fsdp2_shard(model, gloo_world_of_one, wrap_root), tmp_path / "export") + + _assert_byte_identical(_load_checkpoint(tmp_path / "export"), want) + for name in ("config.json", "generation_config.json", "hf_quant_config.json"): + assert (tmp_path / "export" / name).is_file(), name + assert (tmp_path / "export" / sidecar.name).read_text() == sidecar.read_text() + + +def test_nogather_export_saves_modelopt_state(gloo_world_of_one, tmp_path): + """As the default FSDP2 export does, it writes the ModelOpt state next to the weights.""" + model = _load(_save(_llama_config(), tmp_path / "src")) + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randint(0, 63, (2, 8)))) + + _nogather_export( + _fsdp2_shard(model, gloo_world_of_one), tmp_path / "export", save_modelopt_state=True + ) + + assert (tmp_path / "export" / "modelopt_state.pth").is_file() + + +def test_nogather_export_names_modules_in_the_quant_config_as_the_weights_do( + gloo_world_of_one, tmp_path +): + """When the weights keep their in-memory names, as every rank's do once one rank cannot revert + its own, the quant config keeps them too: hub names there would exclude no written module.""" + model = _load(_save(_mixtral_config(), tmp_path / "src")) + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randint(0, 64, (2, 8)))) + failure = mock.patch.object( + quant_aware_conversion, + "revert_weight_conversion_quant_aware", + side_effect=RuntimeError("incomplete merge"), + ) + + with failure, pytest.warns(UserWarning, match="incomplete merge"): + _nogather_export(_fsdp2_shard(model, gloo_world_of_one), tmp_path / "export") + + modules = {key.rsplit(".", 1)[0] for key in _load_checkpoint(tmp_path / "export")} + config = json.loads((tmp_path / "export" / "hf_quant_config.json").read_text()) + excluded = config["quantization"]["exclude_modules"] + assert "model.layers.0.mlp.gate" in excluded + assert set(excluded) <= modules, sorted(set(excluded) - modules) + + +@pytest.mark.parametrize( + "quant_cfg", [mtq.FP8_DEFAULT_CFG, mtq.NVFP4_DEFAULT_CFG], ids=["fp8", "nvfp4"] +) +def test_nogather_export_writes_moe_experts_quantized(gloo_world_of_one, tmp_path, quant_cfg): + """Qwen2-MoE's shared expert follows its experts in the same FSDP2 unit, and packing it opens a + weight window on that unit. Unless the experts were folded inside a window of their own, which + hands FSDP2 the packed weights, that window's reshard puts FSDP2's unquantized shards back.""" + model = _load(_save(_qwen2_moe_config(), tmp_path / "src")) + want = _quantize_with_reference(model, quant_cfg) + + _nogather_export(_fsdp2_shard(model, gloo_world_of_one), tmp_path / "export") + + _assert_byte_identical(_load_checkpoint(tmp_path / "export"), want) + + +# ------------------------------------------------------------------ the writer, two ranks + + +def _mtp_tail(): + """Tensors the model does not own, named like an inlined MTP tail the loader dropped.""" + return { + "model.layers.99.mtp.weight": torch.arange(64, dtype=torch.bfloat16).reshape(8, 8), + "model.layers.99.mtp.bias": torch.full((8,), 3.0, dtype=torch.bfloat16), + } + + +def _two_rank_job(rank, world_size, tmp_dir): + mesh = init_device_mesh("cpu", (world_size,)) + + # An unquantized MoE whose vocab rows and experts split unevenly, written with tensors it does + # not own, which every rank is handed; the test compares what the two ranks wrote between them + # with the source checkpoint and those tensors. + _write( + _fsdp2_shard(_load(tmp_dir / "moe"), mesh), + tmp_dir / "moe_export", + extra_state_dict=_mtp_tail() if rank == 0 else None, + ) + + # A quantized dense model through the export entry point: fp8 weights and bf16 norms travel + # in separate per-dtype buckets. Every rank quantizes identically, so rank 0's single-process + # export of its own copy is the reference. + model = _load(tmp_dir / "dense") + want = _quantize_with_reference(model, mtq.FP8_DEFAULT_CFG) + # Only rank 0's environment picks the no-gather writer: its choice holds on every rank. + _nogather_export( + _fsdp2_shard(model, mesh), + tmp_dir / "dense_export", + knob="nogather" if rank == 0 else "streaming", + ) + + # Each dense tensor ends up whole and plain on exactly one rank, whatever its layout, and the + # owners are spread out. A zero-row tensor has no rows to move: its owner writes it as it is. + full = {f"t{i}": torch.full((4, 2), float(i)) for i in range(9)} + full["t9"] = torch.zeros(0, 4, dtype=torch.bfloat16) + sd = {} + for i, (k, v) in enumerate(full.items()): + layout = ([Shard(0)], [Replicate()], None)[i % 3] + sd[k] = v if layout is None else distribute_tensor(v, mesh, layout) + _materialize_owned(sd, set(), rank, world_size) + owned = [None] * world_size + dist.all_gather_object(owned, sorted(sd)) + assert all(owned), f"a rank owns nothing: {owned}" + assert sorted(k for keys in owned for k in keys) == sorted(full) + for k, v in sd.items(): + assert not isinstance(v, DTensor), k + assert v.dtype == full[k].dtype and torch.equal(v, full[k]), k + + # Ranks that disagree must fail with a diagnosis, not hang in a collective or drop a tensor. + w = torch.ones(1, 2) + if rank == 0: + w = DTensor.from_local(w, mesh, [Shard(0)], run_check=False) + cases = [ + ({"shared": torch.ones(2), f"rank{rank}": torch.ones(2)}, "dense key sets differ"), + ({"w": w}, "disagree on which tensors are dim-0 sharded"), + ] + for sd, diagnosis in cases: + try: + with bounded_collectives(): + _materialize_owned(sd, set(), rank, world_size) + except RuntimeError as error: + raised = str(error) + else: + raised = "no error" + assert diagnosis in raised, raised + + # Tensors handed in that collide with the model's: rank 0 alone adds them, and every rank + # raises, not rank 0 alone. + overlap = {k: torch.zeros(1) for k in _load(tmp_dir / "dense").state_dict()} + try: + with bounded_collectives(): + _write( + _fsdp2_shard(_load(tmp_dir / "dense"), mesh), + tmp_dir / "overlap_export", + extra_state_dict=overlap if rank == 0 else None, + ) + except RuntimeError as error: + raised = str(error) + else: + raised = "no error" + assert "overlaps the exported model weights" in raised, raised + + # Last, after every collective: a failing rank 0 then leaves no rank waiting in one. + if rank == 0: + _assert_byte_identical(_load_checkpoint(tmp_dir / "dense_export"), want) + + +def test_two_ranks_write_one_checkpoint_and_reject_inconsistent_state(cpu_dist_workers, tmp_path): + """Shards cross ranks point to point, each tensor gets one owner, files share one numbering, + and the tensors handed to every rank are written once. + + The same job also checks the cross-rank consistency errors. + """ + src = _save(_mixtral_config(vocab_size=63, num_local_experts=3), tmp_path / "moe") + _save(_llama_config(), tmp_path / "dense") + + cpu_dist_workers.run(_two_rank_job, tmp_dir=tmp_path) + + shards = sorted(p.name for p in (tmp_path / "moe_export").glob("*.safetensors")) + n = len(shards) + assert n > 2, "both ranks should have written several files" + assert shards == [f"model-{i:05d}-of-{n:05d}.safetensors" for i in range(1, n + 1)] + want = {**load_file(str(src / "model.safetensors")), **_mtp_tail()} + _assert_byte_identical(_load_checkpoint(tmp_path / "moe_export"), want) + + +def _merge_export_job(rank, world_size, tmp_dir): + model = _load(tmp_dir / "dense") + # A checkpoint that fuses each layer's q, k and v, which the model holds as three projections: + # the reverse merges them back, which needs each layer's three on one rank. + model._weight_conversions = [ + WeightConverter( + source_patterns="self_attn.qkv_proj.weight", + target_patterns=[f"self_attn.{part}_proj.weight" for part in "qkv"], + operations=[Chunk(dim=0)], + ) + ] + with bounded_collectives(): + _write(_fsdp2_shard(model, init_device_mesh("cpu", (world_size,))), tmp_dir / "merged") + + +def test_distributed_save_merges_a_group_its_ranks_hold_apart(cpu_dist_workers, tmp_path): + """Each layer's q, k and v go to one owner, so the written checkpoint fuses them as the gather + export does, instead of keeping every in-memory name because one layer's split over the ranks.""" + _save(_llama_config(), tmp_path / "dense") + cpu_dist_workers.run(_merge_export_job, tmp_dir=tmp_path) + + got = _load_checkpoint(tmp_path / "merged") + src = _load(tmp_path / "dense").state_dict() + layers = sorted({k.rsplit(".self_attn.", 1)[0] for k in src if ".self_attn.q_proj." in k}) + assert layers + for layer in layers: + attn = f"{layer}.self_attn" + assert f"{attn}.q_proj.weight" not in got + want = torch.cat([src[f"{attn}.{part}_proj.weight"] for part in "qkv"]) + assert torch.equal( + got[f"{attn}.qkv_proj.weight"], want.to(got[f"{attn}.qkv_proj.weight"].dtype) + ) + + +def _expert_input_scale_job(rank, world_size, tmp_dir): + mesh = init_device_mesh("cpu", (world_size,)) + model = _load(tmp_dir / "src") + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randint(0, 64, (2, 8)))) + # As if each rank calibrated the experts' shared input quantizer on its own data: each rank folds + # its own experts with its own scale, and only the writer's reduce makes them agree. + model.model.layers[0].mlp.experts.gate_up_proj_input_quantizer.amax = torch.tensor(1.0 + rank) + _nogather_export(_fsdp2_shard(model, mesh), tmp_dir / "export") + + # A scale one rank holds and the other lacks fails on both, named, rather than leave the ranks + # hanging or mismatched in the reduce. + from torch.distributed.checkpoint import state_dict + + real = state_dict.get_model_state_dict + + def without_a_scale(*args, **kwargs): + sd = real(*args, **kwargs) + return {k: v for k, v in sd.items() if rank == 0 or "gate_up_proj_input_scale" not in k} + + model = _load(tmp_dir / "src") + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, lambda m: m(torch.randint(0, 64, (2, 8)))) + try: + with ( + bounded_collectives(), + mock.patch.object(state_dict, "get_model_state_dict", without_a_scale), + ): + _nogather_export(_fsdp2_shard(model, mesh), tmp_dir / "diverging") + except RuntimeError as error: + raised = str(error) + else: + raised = "no error" + assert "different expert input scales" in raised, raised + + +def test_two_ranks_write_every_expert_with_the_largest_input_scale(cpu_dist_workers, tmp_path): + """The experts' shared input scale is reduced across the ranks that fold them, to its max, and + ranks that hold different scales fail alike, naming them.""" + _save(_qwen2_moe_config(), tmp_path / "src") + + cpu_dist_workers.run(_expert_input_scale_job, tmp_dir=tmp_path) + + exported = _load_checkpoint(tmp_path / "export") + scales = [exported[f"model.layers.0.mlp.experts.{e}.gate_proj.input_scale"] for e in range(4)] + assert all(torch.allclose(s, torch.tensor(2.0 / 448)) for s in scales), scales diff --git a/tests/unit/torch/export/test_fused_experts_keep_fused.py b/tests/unit/torch/export/test_fused_experts_keep_fused.py index febefc06dfe..06ed08e51e3 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -24,7 +24,6 @@ rank in process, and two spawned ranks where a non-zero expert offset matters. """ -import contextlib import copy import math import re @@ -723,36 +722,6 @@ def test_keep_fused_refuses_fused_experts_not_named_experts(): torch.testing.assert_close(experts.state_dict(), before, rtol=0, atol=0) -@pytest.mark.parametrize("dense", ["fsdp", "replicate"]) -def test_fused_experts_fold_in_a_weight_window_when_the_plan_shards_the_dense_part( - monkeypatch, dense -): - """FSDP2 owns the fused weights when it shards the dense part around them: folded outside a - weight-update window, they would be overwritten with its original shards at the next reshard.""" - events = [] - - @contextlib.contextmanager - def window(root, modules, names=None): - events.append("open") - yield - events.append("close") - - monkeypatch.setattr(handlers, "fsdp2_aware_weight_update", window) - monkeypatch.setattr( - handlers, "_export_fused_experts_keep_fused", lambda module, dtype: events.append("fold") - ) - experts = _quantized_experts(_KeepFusedGatedExperts, mtq.FP8_DEFAULT_CFG) - container = nn.Module() - container.experts = experts - set_distribution_plan(container, DistributionPlan(2, ep_size=2, dense=dense)) - ctx = ExportContext(model=container, dtype=torch.bfloat16, keep_fused_experts=True) - - ExportModuleRegistry.match(experts)(PREFIX, experts, ctx) - handlers.close_fsdp2_weight_window(ctx) - - assert events == (["open", "fold", "close"] if dense == "fsdp" else ["fold"]) - - @pytest.mark.parametrize("root_sharded", [True, False], ids=["fsdp2", "expert-parallel"]) def test_quant_linear_handlers_pack_an_fsdp2_unit_in_one_weight_window( monkeypatch, gloo_world_of_one, root_sharded diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index f5f36627ad0..c1f89329605 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -39,6 +39,7 @@ sync_tied_input_amax, ) from modelopt.torch.export.unified_export_hf import ( + _copy_non_model_files_from_source, _resolve_export_dtype, _source_checkpoint, read_unplaced_weights, @@ -1011,7 +1012,18 @@ def test_an_unrecorded_model_derives_its_unplaced_keys_once(tmp_path, monkeypatc assert model._modelopt_unplaced_source_keys == [mtp] -def test_source_checkpoint_is_none_for_a_model_built_from_config(): +def test_source_checkpoint_is_none_for_a_model_built_from_config(tmp_path, monkeypatch): """``from_config`` leaves ``_name_or_path`` empty; it must not resolve to the working directory, or every export would copy the cwd as the "source checkpoint".""" - assert _source_checkpoint(SimpleNamespace(config=SimpleNamespace(_name_or_path=""))) is None + workdir = tmp_path / "workdir" + workdir.mkdir() + (workdir / "tokenizer.json").write_text("{}") + monkeypatch.chdir(workdir) + model = SimpleNamespace(config=SimpleNamespace(_name_or_path="")) + export_dir = tmp_path / "export" + export_dir.mkdir() + + _copy_non_model_files_from_source(model, export_dir) + + assert _source_checkpoint(model) is None + assert not any(export_dir.iterdir())