diff --git a/CHANGELOG.rst b/CHANGELOG.rst index f6cfb894081..a052d081046 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -24,6 +24,7 @@ Changelog - Speed up NVFP4 QAT and QAD of Hugging Face fused MoE experts: all experts of a projection are now fake-quantized in one kernel launch instead of one per expert, about 2x faster per MoE layer on Qwen3-30B-A3B with FSDP2. - ``hf_ptq.py`` calibrates Qwen3.8 (``qwen4_exp``) without running its QSA indexer's per-query Python loop when every visible key fits the indexer's token budget, with bit-identical results. On Qwen3.8-Flash-Next this saves about two hours of layerwise calibration. - Speed up Hugging Face QAT and QAD with ``QATTrainer``, ``QADTrainer`` or ``--manual_gc``: the ``gc.collect()`` run before each training step now skips the model and optimizer state, saving about 1 s per step (30%) on Qwen3-30B-A3B. +- 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* @@ -46,6 +47,7 @@ Changelog **Backward Breaking Changes** +- ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` on more than one GPU now refuses calibration that keeps its statistics on each rank, which calibrated the ranks' shards of a weight differently: AWQ-Lite, AWQ-Full, SVDQuant, local Hessian and imatrix, also as an AutoQuantize candidate (run them without ``--use_fsdp2``). It also raises when a calibration batch runs out of GPU memory, instead of retrying it in smaller pieces: set ``--batch_size`` to one that fits. - ``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``. - Drop support for Python 3.10, which has reached end-of-life. ModelOpt now requires Python 3.11 or later. - Bump minimum required PyTorch version to 2.10. @@ -59,6 +61,7 @@ Changelog - ``mtq.calibrate`` now raises an error, naming the modules, when a Hugging Face fused-experts forward reaches neither quantizer hook (a hand-written ``bmm`` forward that transformers cannot redirect, say), instead of exporting weight-derived input scales with a warning. - Layerwise-calibration resume (``checkpoint_dir``) now loads checkpoint files with ``weights_only=True``. A checkpoint directory written by an earlier release whose layer inputs hold objects other than tensors, plain containers or Hugging Face dynamic caches can no longer be resumed; recalibrate in a fresh directory, or disable checkpointing. +- ``export_hf_checkpoint`` now reads how a model is distributed from the plan recorded on it, and raises for an FSDP2-sharded model, or one in a multi-rank run, that carries none. ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` records it; if you distribute a model yourself, for example with ``fsdp2_wrap`` or ``parallel_load_and_prepare_fsdp2``, which record 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 b8aba4de503..71be2080d53 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -610,14 +610,18 @@ mtq.calibrate(model, algorithm="max", forward_loop=calibrate_loop) ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for distributed model sharding and calibration, exposed via the `--use_fsdp2` flag on the standard `hf_ptq.py` entry point. -> *Weight and KV-cache AutoQuantize with `--use_fsdp2` are experimental.* +> *Weight and KV-cache AutoQuantize with `--use_fsdp2` are experimental, and refused when a candidate format is one of those refused below.* > *Nemotron-H checkpoints containing MTP weights are not supported with `--use_fsdp2` and are rejected before distributed model loading. Remove `--use_fsdp2` to load and calibrate the MTP modules; the distributed loader does not construct them.* > *AWQ calibration (`awq_lite`, `awq_full`), as the AWQ formats and recipes use, AutoQuantize candidates included, is not supported with `--use_fsdp2` on more than one rank and is rejected before distributed model loading: each rank would fold its own smoothing scale into its weight shard. Drop `--use_fsdp2` (the single-process `device_map` still spreads the model over all the GPUs), run on one rank, or use a non-AWQ format.* +> *With `--use_fsdp2` on more than one GPU, calibration other than AWQ that keeps its statistics on each rank is refused when quantization starts: SVDQuant (`--qformat nvfp4_svdquant`), local Hessian (`nvfp4_w4a4_weight_local_hessian`), imatrix (the `iq*` formats), the recipes that use them, and AutoQuantize recipes with such a candidate format. Each rank calibrates on its own share of the data, so the ranks would calibrate their shards of a weight differently. Run these without `--use_fsdp2`. For AWQ, see the note above.* + > *`nvfp4_awq_clip` still works with `--use_fsdp2` on more than one GPU: it calibrates with AWQ-Clip alone, which folds no smoothing scale into the weights, and its clip search averages each candidate clip's loss over the ranks, so every rank takes the same clip.* +> *With `--use_fsdp2` on more than one GPU, a calibration batch that runs out of memory is not retried in smaller pieces, since a rank that runs out of memory partway through a forward would leave the others waiting in its collectives. With `--batch_size 0`, each rank estimates the batch size from one forward and the ranks use the smallest, which is not probed for out-of-memory either: set `--batch_size` if calibration runs out of memory.* + ### Usage #### Slurm (recommended) @@ -654,6 +658,8 @@ Qwen checkpoints with `model_type` `qwen3_5` or `qwen3_5_moe` support this path > *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. It does not overwrite a checkpoint already in the export directory.* + ## Evaluate Accuracy ### TensorRT-LLM Validation diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 47d6a8b2c19..c79c25ccf98 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -99,6 +99,7 @@ get_max_batch_size, get_supported_datasets, ) +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan from modelopt.torch.utils.memory_monitor import launch_memory_monitor from modelopt.torch.utils.mlflow import add_mlflow_args, resolve_mlflow_args from modelopt.torch.utils.plugins.hf_checkpoint_utils import ( @@ -339,6 +340,10 @@ def load_model(args: argparse.Namespace): attn_implementation=args.attn_implementation, hf_config=hf_config, ) + # Calibration sizes its batch, and the export writes, by how the model is distributed. + set_distribution_plan( + full_model, DistributionPlan(args.dist_state.world_size, dense="fsdp") + ) elif args.specdec_offline_dataset is not None or not args.low_memory_mode: full_model = get_model( args.local_checkpoint_path, diff --git a/modelopt/torch/export/distributed_writer.py b/modelopt/torch/export/distributed_writer.py index 0f7e63436ea..21ad6da9f34 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 @@ -32,22 +32,61 @@ import torch.distributed as dist import torch.nn as nn from safetensors.torch import save_file +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 .moe_utils import _FUSED_PROJ, _fused_experts_prefixes, _split_local_fused_module from .quant_aware_conversion import ( _QUANT_STATE_SUFFIXES, QuantConversionUnsupportedError, _build_reverse_rules, revert_weight_conversion_quant_aware, ) +from .quant_utils import postprocess_state_dict from .unified_export_hf_streaming import _parse_shard_size # A checkpoint's files: its weights, sharded as save_pretrained names them or in one file, and index. _CHECKPOINT_FILES = re.compile(r"model(-\d{5}-of-\d{5})?\.safetensors(\.index\.json)?") +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. @@ -633,3 +672,185 @@ def commit(names: list) -> bool: f"wrote to {export_dir}: {left[:5]}" ) raise + + +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. + 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. + """ + # 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() + 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. + # FSDP2 shards a fused 3-D expert weight on dim 0, the expert axis, so each per-expert key's global + # index is the DTensor's own dim-0 offset, which _split_local_fused_module reads. + local_sd: dict = {} + for prefix in prefixes: + local_sd.update(_split_local_fused_module(prefix, sharded_sd)) + 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 them as they are for now -- FSDP2 leaves DTensors (dim-0 shard), + # and what it does not shard stays a plain tensor, the same on every rank. 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, + # whole CPU tensors, and only rank 0 adds them to its write set (the export entry hands them to + # rank 0 alone); 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/moe_utils.py b/modelopt/torch/export/moe_utils.py index 07e0ddd4129..e12ed147abf 100644 --- a/modelopt/torch/export/moe_utils.py +++ b/modelopt/torch/export/moe_utils.py @@ -669,9 +669,10 @@ def split_fused_experts_state_dict(state_dict: dict) -> dict: checkpoints. Operates on the (already folded/quantized) tensors only -- pure slicing, no quantization. Non-fused-experts keys pass through unchanged. Every emitted tensor is a copy and the input keeps its fused tensors, so the experts are held twice until it returns: pass one - module's tensors at a time, not a whole model's. Each module's experts are numbered from 0 by - their row in the tensors given, which must be plain: whole ones get their global numbers, and a - caller splitting a rank's local shard (``to_local()``) shifts each by the shard's global offset. + module's tensors at a time, as :func:`_split_local_fused_module` does, not a whole model's. + Each module's experts are numbered from 0 by their row in the tensors given, which must be + plain: whole ones get their global numbers, and a caller splitting a rank's local shard + (``to_local()``) shifts each by the shard's global offset, as that function does. 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`` @@ -801,3 +802,74 @@ def _dtensor_dim0_offset(dt) -> int: 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) -> dict: + """Split this rank's local shard of a fused-experts module into per-expert keys. + + 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 shards 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). + """ + 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] + + 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. The keep-fused fold emits DTensor scales for a sharded weight, but a + # scale handed in as a PLAIN tensor over every expert of the group is accepted too: 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 dp_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) + dp_offset}.{tail}"] = val + return out diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 7c57c3838f5..9c63130bfd1 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 @@ -95,7 +96,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: @@ -977,6 +983,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. @@ -992,7 +999,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( @@ -1000,11 +1009,12 @@ def _process_quantized_modules( dtype=dtype, is_modelopt_qlora=is_modelopt_qlora, model_type=hf_model_type(model), + keep_fused_experts=keep_fused_experts, ) - # For the no-gather export, which is to set keep_fused_experts on this context and let a sharded - # model past the assert above: each FSDP2 unit's modules are then packed in one weight window, - # held across them. Until then no window opens here. + # The no-gather export sets keep_fused_experts on this context, which lets a sharded model past + # the assert above: each FSDP2 unit's modules are then packed in one weight window, held across + # them. with _hf_export_handlers.held_fsdp2_weight_windows(ctx): for name, sub_module in model.named_modules(): _dispatch_export_handler(name, sub_module, ctx) @@ -1058,13 +1068,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) @@ -1914,6 +1933,20 @@ def read_unplaced_weights(model: nn.Module, *, keys_only: bool = False) -> dict[ return out +def _no_gather_export_asked() -> bool: + """Whether ``MODELOPT_FSDP2_EXPORT`` picks the no-gather writer for an FSDP2-sharded export. + + ``nogather`` picks it; unset, empty or ``streaming`` keeps the default streaming path. Any other + value raises: a misspelled one would otherwise quietly take the path it was set to avoid. + """ + value = os.environ.get("MODELOPT_FSDP2_EXPORT") or "streaming" + if value not in ("nogather", "streaming"): + raise ValueError( + f"MODELOPT_FSDP2_EXPORT={value!r}: expected 'nogather', or 'streaming', the default." + ) + return value == "nogather" + + def export_hf_checkpoint( model: Any, dtype: torch.dtype | None = None, @@ -1962,14 +1995,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 {})} @@ -2015,10 +2044,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", _no_gather_export_asked ) # 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 @@ -2049,6 +2079,44 @@ 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. + # Imported here: the writer imports unified_export_hf_streaming, which imports this module. + from .distributed_writer import distributed_save_hf_checkpoint + + _qlora = kwargs.get("is_modelopt_qlora", False) + # The writer builds its own name-based tied-weight map, so this one is not handed down. + _dtype, _, hf_quant_config = _prepare_model_for_export(model, dtype, _qlora) + # Preparing the model does not pack its weights, so pack them here. keep_fused_experts + # leaves the MoE experts fused and sharded through the fold, so each rank folds only + # its local experts; the writer then splits each rank's experts into per-expert tensors. + 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: + save_non_weight_artifacts(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 diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 124cded7a38..7f741ad0bad 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -1162,7 +1162,11 @@ def _slice_batch(batch_data, start, end): def _process_batch( - batch_data, infer_method, max_working_batch_size=None, allowed_non_tensor_keys=None + batch_data, + infer_method, + max_working_batch_size=None, + allowed_non_tensor_keys=None, + split_on_oom=True, ): """Process a batch of data through the model's inference method. @@ -1171,6 +1175,8 @@ def _process_batch( infer_method: Model's inference method (either forward or generate) max_working_batch_size: Maximum batch size known to work without OOM allowed_non_tensor_keys: Set of key names whose values may be non-tensor types + split_on_oom: Whether a batch that runs out of GPU memory is retried as two halves; + otherwise the error is raised Returns: The maximum batch size that worked successfully @@ -1193,7 +1199,11 @@ def _process_batch( width = min(max_working_batch_size, batch_size - offset) split_data = _slice_batch(batch_data, offset, offset + width) max_working_batch_size = _process_batch( - split_data, infer_method, max_working_batch_size, allowed_non_tensor_keys + split_data, + infer_method, + max_working_batch_size, + allowed_non_tensor_keys, + split_on_oom, ) offset += width @@ -1208,6 +1218,8 @@ def _process_batch( else max(batch_size, max_working_batch_size) ) # This batch size worked successfully except torch.cuda.OutOfMemoryError: + if not split_on_oom: + raise assert batch_size > 1, ( "CUDA out of memory error occurred while processing a single sample. " "This indicates the model is too large for the available GPU memory. " @@ -1245,6 +1257,11 @@ def _forward_loop( dataloader: DataLoader containing the batched input data allowed_non_tensor_keys: Set of key names whose values may be non-tensor types """ + # A batch that runs out of GPU memory is retried in halves, except for a model whose recorded plan + # spans several ranks: its forward can hold collectives, and a rank that retried would run more + # forwards than its peers, which would wait in them. + plan = recorded_distribution_plan(model) + split_on_oom = plan is None or plan.world_size == 1 with _disable_use_cache(model), torch.no_grad(): is_enc_dec = model_type_is_enc_dec(model) # Call the module (not .forward) so nn.Module.__call__ runs pre/post-forward hooks — this is @@ -1255,7 +1272,7 @@ def _forward_loop( for _, data in enumerate(tqdm(dataloader)): # Process batch and update max working batch size max_working_batch_size = _process_batch( - data, infer_method, max_working_batch_size, allowed_non_tensor_keys + data, infer_method, max_working_batch_size, allowed_non_tensor_keys, split_on_oom ) diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index 2f0b425de40..5543e9de902 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -26,7 +26,7 @@ import yaml from _test_utils.mlflow import clean_env # noqa: F401 from _test_utils.torch.transformers_models import get_tiny_qwen3 -from transformers import Qwen3_5MoeConfig, Qwen3MoeConfig +from transformers import Qwen3_5MoeConfig, Qwen3Config, Qwen3MoeConfig from modelopt.recipe import load_recipe from modelopt.recipe.config import ( @@ -38,6 +38,7 @@ from modelopt.torch.quantization import tensor_quant from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.utils import mlflow as mlflow_lib +from modelopt.torch.utils.distributed import DistributionPlan, distribution_plan from modelopt.torch.utils.mlflow import describe_run, run_tags _EXAMPLES_DIR = Path(__file__).resolve().parents[3] / "examples" / "hf_ptq" @@ -480,6 +481,27 @@ def distributed_load(*args, **kwargs): hf_ptq.load_model(args) +def test_the_fsdp2_load_records_the_plan_it_loaded_by(monkeypatch, tmp_path): + """Calibration sizes its batch, and the export writes, by the model's distribution plan, and an + FSDP2-sharded model without one is refused rather than taken for a single process's.""" + Qwen3Config(architectures=["Qwen3ForCausalLM"]).save_pretrained(tmp_path) + monkeypatch.setenv("RANK", "0") + hf_ptq, args = _parse_hf_ptq_args(monkeypatch, "--pyt_ckpt_path", str(tmp_path), "--use_fsdp2") + args.dist_state = SimpleNamespace(device="cpu", rank=0, world_size=2) + args.hub_model_id, args.local_checkpoint_path = None, str(tmp_path) # as main() sets them + model = torch.nn.Linear(1, 1) + monkeypatch.setattr(hf_ptq, "parallel_load_and_prepare_fsdp2", lambda *args, **kwargs: model) + + def loaded(model): + raise RuntimeError("loaded") + + monkeypatch.setattr(hf_ptq, "prepare_model_for_calibration", loaded) + + with pytest.raises(RuntimeError, match="loaded"): + hf_ptq.load_model(args) + assert distribution_plan(model) == DistributionPlan(2, dense="fsdp") + + def test_mlflow_flag_defaults_the_experiment_name(monkeypatch): monkeypatch.setattr(getpass, "getuser", lambda: "tester") hf_ptq, args = _parse_hf_ptq_args( 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..844a46dd98e --- /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 its plan is recorded by hand + # 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/gpu/torch/export/test_export_carry_over.py b/tests/gpu/torch/export/test_export_carry_over.py index fa528f4cdf9..34a54d23f5e 100644 --- a/tests/gpu/torch/export/test_export_carry_over.py +++ b/tests/gpu/torch/export/test_export_carry_over.py @@ -40,7 +40,12 @@ 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 +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. @@ -108,6 +113,8 @@ def _ptq_and_export(rank, size, *, src_dir, export_dir, quant_cfg, **export_kwar model.eval() fsdp2_wrap(model) + # Wrapped by hand, the model records how it is distributed, as load_pretrained would. + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) # The FSDP2 export path 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. @@ -201,6 +208,7 @@ def _ptq_language_model_and_export(rank, size, *, src_dir, export_dir): ) model.eval() fsdp2_wrap(model) + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) assert is_fsdp2_model(model), "fsdp2_wrap did not shard the VLM" torch.distributed.barrier() diff --git a/tests/gpu/torch/utils/test_model_load_utils.py b/tests/gpu/torch/utils/test_model_load_utils.py index 3296ccbfc86..05b524fb19a 100644 --- a/tests/gpu/torch/utils/test_model_load_utils.py +++ b/tests/gpu/torch/utils/test_model_load_utils.py @@ -26,7 +26,11 @@ from torch.distributed.tensor import DTensor from modelopt.torch.export.unified_export_hf import export_hf_checkpoint -from modelopt.torch.utils.distributed import broadcast_state_dict +from modelopt.torch.utils.distributed import ( + DistributionPlan, + broadcast_state_dict, + set_distribution_plan, +) from modelopt.torch.utils.plugins.model_load_utils import parallel_load_and_prepare_fsdp2 VOCAB_SIZE = 64 @@ -61,6 +65,7 @@ def _test_parallel_load_and_export(rank, size, ckpt_dir, export_dir, cpu_offload size, cpu_offload=cpu_offload, ) + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) # Decoder layers AND root params (embed/lm_head) are sharded DTensors under shard_root=True. decoder_params = list(model.model.layers[0].parameters()) @@ -106,6 +111,7 @@ def _test_carry_over_unplaced_weights(rank, size, ckpt_dir, export_dir, orphan_p """Load, export, and require the weights the model could not place to survive.""" device = torch.device(f"cuda:{rank}") model = parallel_load_and_prepare_fsdp2(ckpt_dir, device, rank, size) + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) unplaced = getattr(model, "_modelopt_unplaced_source_keys", None) assert unplaced, ( diff --git a/tests/unit/torch/export/test_distributed_writer.py b/tests/unit/torch/export/test_distributed_writer.py index 8dd8ec6413c..a35e47086c9 100644 --- a/tests/unit/torch/export/test_distributed_writer.py +++ b/tests/unit/torch/export/test_distributed_writer.py @@ -13,14 +13,18 @@ # 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 @@ -29,8 +33,10 @@ 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.distributed.tensor.placement_types import _StridedShard from torch.multiprocessing.reductions import StorageWeakRef @@ -38,9 +44,11 @@ pytest.importorskip("safetensors") pytest.importorskip("transformers") from safetensors.torch import load_file, save_file -from transformers import AutoModelForCausalLM, LlamaConfig, MixtralConfig +from transformers import AutoModelForCausalLM, LlamaConfig, MixtralConfig, Qwen2MoeConfig from transformers.core_model_loading import Chunk, WeightConverter +import modelopt.torch.quantization as mtq +from modelopt.torch.export import distributed_writer from modelopt.torch.export.distributed_writer import ( _even_bins, _is_full_world_shard0, @@ -49,11 +57,31 @@ _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 _WRITER = "modelopt.torch.export.distributed_writer" +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): return tensor.detach().contiguous().reshape(-1).view(torch.uint8) @@ -81,6 +109,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 = { @@ -111,6 +149,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 @@ -158,29 +257,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( @@ -404,42 +480,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, message): # Rank 0's subset reverts; rank 1's fails, as a merge whose group ownership split does. model = AutoModelForCausalLM.from_config(_mixtral_config()) @@ -752,3 +792,348 @@ 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_refuses_a_writer_it_does_not_know(gloo_world_of_one, tmp_path): + """A misspelled ``MODELOPT_FSDP2_EXPORT`` raises, rather than quietly taking the default writer, + the one it was set to avoid.""" + model = _fsdp2_shard(_load(_save(_llama_config(), tmp_path / "src")), gloo_world_of_one) + + with pytest.raises(ValueError, match="MODELOPT_FSDP2_EXPORT='no-gather'"): + _nogather_export(model, tmp_path / "export", knob="no-gather") + + +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( + distributed_writer, + "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 _uneven_rows(): + """63 rows: two ranks shard them unevenly.""" + return torch.arange(63, dtype=torch.float32) + + +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 only rank 0 is handed, as the export entry hands them; the test compares what + # the two ranks wrote between them with the source checkpoint and those tensors. It is written + # in inference mode, with a buffer made there that shards unevenly, as an inference-mode pack + # makes its scales: the writer copies each such shard out of inference mode, and the rewrapped + # DTensor must keep the global shape, which the two uneven shards would otherwise each re-infer + # differently. + moe = _fsdp2_shard(_load(tmp_dir / "moe"), mesh) + with torch.inference_mode(): + rows = _uneven_rows() + local = rows[:32] if rank == 0 else rows[32:] # Shard(0) of 63 rows: 32 + 31 + moe.model.norm.register_buffer( + "uneven_rows", + DTensor.from_local( + local.clone(), mesh, [Shard(0)], run_check=False, shape=rows.shape, stride=(1,) + ), + ) + _write(moe, 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 rank 0 alone are written once. A tensor sharded unevenly in + inference mode is written whole. + + 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(), + "model.norm.uneven_rows": _uneven_rows(), + } + _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. + real = distributed_writer.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(distributed_writer, "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 0a516db55aa..a646c3743a3 100644 --- a/tests/unit/torch/export/test_fused_experts_keep_fused.py +++ b/tests/unit/torch/export/test_fused_experts_keep_fused.py @@ -17,14 +17,14 @@ The no-gather writer never splits a fused 3-D expert weight while the model is sharded: ``_export_fused_experts_keep_fused`` folds every expert's quantizer into the fused weight in place, -and ``split_fused_experts_state_dict`` cuts the result into per-expert keys later. That pipeline -must write exactly what the single-process ``_export_fused_experts`` writes, so most tests compare -the two byte for byte on a tiny synthetic fused-experts module quantized with FP8 and NVFP4. +and ``split_fused_experts_state_dict`` / ``_split_local_fused_module`` cut the result into +per-expert keys later. That pipeline must write exactly what the single-process +``_export_fused_experts`` writes, so most tests compare the two byte for byte on a tiny synthetic +fused-experts module quantized with FP8 and NVFP4. Sharded cases run on a gloo process group: one rank in process, and two spawned ranks where a non-zero expert offset matters. """ -import contextlib import copy import math import re @@ -43,20 +43,19 @@ from transformers import NemotronHConfig, PretrainedConfig from transformers.models.nemotron_h.modeling_nemotron_h import NemotronHExperts -import modelopt.torch.export.hf_export_handlers as handlers import modelopt.torch.quantization as mtq from modelopt.torch.export.moe_utils import ( _dtensor_dim0_offset, _export_fused_experts, _export_fused_experts_keep_fused, _fused_experts_prefixes, + _split_local_fused_module, split_fused_experts_state_dict, ) from modelopt.torch.export.registry import ExportContext, ExportModuleRegistry from modelopt.torch.quantization.config import QuantizerAttributeConfig from modelopt.torch.quantization.ggml import GGML_FORMAT_REGISTRY from modelopt.torch.quantization.nn import QuantModuleRegistry, SequentialQuantizer, TensorQuantizer -from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan NUM_EXPERTS = 4 HIDDEN_DIM = 32 @@ -631,6 +630,22 @@ def _two_rank_fold_job(rank, world_size): # Gathered whole, the ranks' folds split into exactly what the per-expert export writes. whole = {key: _whole(value) for key, value in fused.items()} _assert_byte_identical(split_fused_experts_state_dict(whole), want) + mine = { + k: v + for k, v in want.items() + if offset <= int(k[len(PREFIX) + 1 :].split(".")[0]) < offset + n_local + } + assert mine, "this rank owns no experts" + + _assert_byte_identical(_split_local_fused_module(PREFIX, fused), mine) + + # A per-expert scale held as a plain tensor over ALL experts (as keep-fused emits it for an + # unsharded weight) must be cut to this rank's experts before pairing. + plain_scales = { + k: plain[k] if k.endswith(("_weight_scale", "_weight_scale_2")) else v + for k, v in fused.items() + } + _assert_byte_identical(_split_local_fused_module(PREFIX, plain_scales), mine) # One expert over two ranks leaves rank 1 none: both ranks refuse, rather than rank 1 alone # failing while rank 0 goes on to write. @@ -686,10 +701,11 @@ def sharded_keeping_quantizers(kept): def test_two_ranks_fold_their_own_experts_under_global_indices(cpu_dist_workers): """Rank 1 holds experts 2-3: each must be folded with ITS quantizer at its global offset, and the fused tensors stay sharded on the expert axis at their global shape, an uneven shard included; - gathered, they split into the per-expert export. Its quantizers are every expert's (FSDP2) or, - under expert parallelism, its own two, and then the ranks decide on each projection's - quantization together. A layout that leaves a rank no expert, or quantizers that are neither, - is refused on every rank.""" + gathered, they split into the per-expert export, and each rank's own split, numbered globally, + into its experts' part of it. Its quantizers are every expert's (FSDP2) or, under expert + parallelism, its own two, and then the ranks decide on each projection's quantization together. + A layout that leaves a rank no expert, or quantizers that are neither, is refused on every + rank.""" cpu_dist_workers.run(_two_rank_fold_job) @@ -751,36 +767,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. - Outside held windows the handler closes the window itself.""" - 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) - - assert events == (["open", "fold", "close"] if dense == "fsdp" else ["fold"]) - - class GptOssExperts(nn.Module): """Named as transformers' fused BMM experts, which store their weights transposed.""" diff --git a/tests/unit/torch/export/test_unified_export_hf.py b/tests/unit/torch/export/test_unified_export_hf.py index 1b8d55373fa..ef113dbcc4d 100644 --- a/tests/unit/torch/export/test_unified_export_hf.py +++ b/tests/unit/torch/export/test_unified_export_hf.py @@ -40,6 +40,7 @@ sync_tied_input_amax, ) from modelopt.torch.export.unified_export_hf import ( + _copy_non_model_files_from_source, _fuse_shared_input_modules, _resolve_export_dtype, _source_checkpoint, @@ -1091,7 +1092,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()) diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index 138cc35c1da..b77a5d94f1c 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -376,6 +376,33 @@ def _collate(samples): assert model.config.use_cache is True +@pytest.mark.parametrize("world_size", [None, 1, 2], ids=["no-plan", "one-rank", "two-ranks"]) +def test_forward_loop_raises_out_of_memory_for_a_model_over_several_ranks(world_size): + """A model distributed over several ranks can run collectives in its forward, so a rank must not + retry a smaller batch after running out of memory: it would run more forwards than its peers. + Any other model splits the batch, as before.""" + batch_sizes = [] + + class _Model(torch.nn.Module): + def forward(self, input_ids): + batch_sizes.append(input_ids.shape[0]) + if input_ids.shape[0] > 2: + raise torch.cuda.OutOfMemoryError + + model = _Model() + if world_size is not None: + set_distribution_plan(model, DistributionPlan(world_size, dense="fsdp")) + loader = DataLoader([{"input_ids": torch.zeros(8, dtype=torch.long)}] * 4, batch_size=4) + + if world_size == 2: + with pytest.raises(torch.cuda.OutOfMemoryError): + _forward_loop(model, loader) + assert batch_sizes == [4] + else: + _forward_loop(model, loader) + assert batch_sizes == [4, 2, 2] + + def test_disable_use_cache_restores_on_exception(): """Restore must run even if the with-block raises.""" model = torch.nn.Linear(4, 4)