diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 89fb2a8b8d7..f6cfb894081 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -68,6 +68,7 @@ Changelog **Bug Fixes** - Megatron-Core fake quantization of Gated DeltaNet (Qwen3.5 / Qwen3.6) with a Transformer Engine ``in_proj`` now keeps its beta / alpha gates in BF16 during evaluation, QAD, ``fold_weight`` and the vLLM fake-quant export, matching deployment. Other ``in_proj`` layers stay fully quantized and warn. +- AWQ-Clip's per-tensor clip, which ``nvfp4_awq_clip`` uses, is now searched over every data-parallel rank's data, as its block-wise clip already was, so the replicas of a weight take the same clip (exported amax values of data-parallel runs may change) and ``nvfp4_awq_clip`` works with ``examples/hf_ptq/hf_ptq.py --use_fsdp2`` on several GPUs. - ``hf_quant_config.json`` now declares ``pre_quant_scale`` only for INT4 / W4A8 layers that have the tensor; loading INT4 / W4A8 fused-MoE checkpoints in TensorRT-LLM also needs a release whose MoE loader treats it as optional. AWQ, SmoothQuant and SVDQuant now warn when they skip quantized fused-MoE experts. - Hugging Face fused MoE experts stored transposed, ``(num_experts, in_dim, out_dim)``, without a quantized wrapper of their own (for example Aria's) are now left unquantized with a warning. They were quantized with gate and up split along the wrong axis, so their exported weights were wrong; the Hugging Face export now refuses such a model's MoE block instead. - Calibrating a Megatron-Core MLA model now gives ``linear_q_down_proj`` and ``linear_kv_down_proj`` one per-tensor input and weight range, as vLLM and TensorRT-LLM fuse them into one GEMM with one scale per tensor. Re-calibrate existing MLA checkpoints to get the shared ranges. diff --git a/examples/hf_ptq/README.md b/examples/hf_ptq/README.md index d108dcbf786..b8aba4de503 100755 --- a/examples/hf_ptq/README.md +++ b/examples/hf_ptq/README.md @@ -616,6 +616,8 @@ ModelOpt enables quantization of LLMs across multiple GPU nodes using FSDP2 for > *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.* +> *`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.* + ### Usage #### Slurm (recommended) diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 7397d79eb56..1d9ddf63e34 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -2029,16 +2029,10 @@ def __init__(self, module): clip_ratios = [ round(float(k), 2) for k in torch.arange(min_clip_ratio, 1.0, shrink_step) ] + [1.0] - if self.is_per_tensor_clip(module): - self.loss = {k: torch.tensor(0.0, device=module.weight.device) for k in clip_ratios} - else: - self.loss = { - k: torch.zeros( - (co, math.ceil(ci / self.block_size)), - device=module.weight.device, - ) - for k in clip_ratios - } + shape = () if self.is_per_tensor_clip(module) else (co, math.ceil(ci / self.block_size)) + # One buffer for every candidate's loss, which the data-parallel sync reduces in place. + self.loss_buffer = torch.zeros(len(clip_ratios), *shape, device=module.weight.device) + self.loss = dict(zip(clip_ratios, self.loss_buffer)) self.best_clip_val = None self.best_loss = None @@ -2067,7 +2061,24 @@ def update_best_params(self): indices, self.awq_clip.w_amax * shrink, self.awq_clip.best_clip_val ) + def sync_across_dp(module): + """Average the clip search over the data-parallel ranks, so every replica picks one clip. + + Once, after the forwards: a module's forward count, as an MoE expert's, can differ by rank. + """ + group = module.parallel_state.data_parallel_group + if not group.is_initialized(): + return + losses = module.awq_clip.loss_buffer + dist.all_reduce(losses, op=dist.ReduceOp.SUM, group=group.group) + losses /= group.world_size() + tokens = torch.tensor(module.awq_clip.num_tokens, device=losses.device) + dist.all_reduce(tokens, op=dist.ReduceOp.SUM, group=group.group) + module.awq_clip.num_tokens = int(tokens.item()) + def _clip_search(self, inputs, co_bsz=256, max_tokens=16): + if inputs.numel() == 0: # no tokens to search with; their loss, a mean over none, is NaN + return weight = self.weight self.weight_quantizer.enable() @@ -2126,11 +2137,6 @@ def _clip_search(self, inputs, co_bsz=256, max_tokens=16): # co_bsz, n_block loss = (cur_out - org_out).float().pow(2).mean(dim=1) - parallel_group = self.parallel_state.data_parallel_group - if parallel_group.is_initialized(): - dist.all_reduce(loss, op=dist.ReduceOp.SUM, group=parallel_group.group) - loss /= parallel_group.world_size() - del cur_out, cur_w self.awq_clip.loss[shrink][ co_batch * co_bsz : min((co_batch + 1) * co_bsz, co) @@ -2198,6 +2204,11 @@ def postprocess(module): if module.awq_clip.is_input_quantized: module.input_quantizer.enable() + # Every module's reduction before any weight gather, which an FSDP2 model's postprocess runs. + for module in model.modules(): + if is_quantized_linear(module) and hasattr(module, "awq_clip"): + sync_across_dp(module) + for name, module in model.named_modules(): if is_quantized_linear(module) and hasattr(module, "awq_clip"): if module.awq_clip.num_tokens > 0: diff --git a/modelopt/torch/quantization/model_quant.py b/modelopt/torch/quantization/model_quant.py index 87c8677985c..54b952e2607 100644 --- a/modelopt/torch/quantization/model_quant.py +++ b/modelopt/torch/quantization/model_quant.py @@ -53,7 +53,7 @@ get_kv_cache_auto_quantize_config, ) from .kv_cache_auto_quant import _validate_search_inputs as _validate_kv_cache_search_inputs -from .mode import QuantizeModeRegistry, get_modelike_from_algo_cfg +from .mode import CalibrateModeRegistry, QuantizeModeRegistry, get_modelike_from_algo_cfg from .nn import AnyQuantizer, QuantModule, SequentialQuantizer, TensorQuantizer from .utils import is_quantized from .utils.core_utils import _get_fsdp2_mesh @@ -72,6 +72,116 @@ "temporarily_fold_weights", ] +# Calibration algorithms whose statistics stay on the rank that gathered them. Any other keeps its +# own there too when its config turns distributed_sync off. +_RANK_LOCAL_CALIBRATION = ( + "awq_lite", + "awq_full", + "svdquant", + "local_hessian", # its distributed_sync syncs only the max amax it starts from + "imatrix", +) + + +def _rank_local_calibrations(algorithm: QuantizeAlgoCfgType) -> list[str]: + """The rank-local calibrations ``algorithm`` runs, with those it runs for its scales.""" + found = [] + for mode, options in get_modelike_from_algo_cfg(algorithm): + method = mode.removesuffix("_calibrate") + # The options as the calibration reads them, through its mode's schema: a string such as + # "false" is coerced, and a default counts as if spelled out. + config = CalibrateModeRegistry[mode].config_class.model_validate(options) + if method in _RANK_LOCAL_CALIBRATION: + found.append(method) + elif not getattr(config, "distributed_sync", True): # None skips the sync too + found.append(f"{method} with distributed_sync={config.distributed_sync}") + for key in ("weight_scale_algorithm", "scale_algorithm"): + if getattr(config, key, None) is not None: + found += _rank_local_calibrations(getattr(config, key)) + return found + + +def _group_ranks(group: dist.DistributedProcessGroup) -> list[int]: + """The ranks ``group`` spans, in its rank order; just this rank for a group not set.""" + if not group.is_initialized(): + return [dist.rank()] + return torch.distributed.get_process_group_ranks(group.group or torch.distributed.group.WORLD) + + +def _names_the_experts_groups( + state: "dist.ParallelState | None", plan: dist.DistributionPlan +) -> bool: + """Whether ``state`` names the groups a routed expert of ``plan`` syncs its scales over. + + Those of the ``(edp, ep)`` mesh the expert-parallel load lays the ranks out in: as its + expert-parallel group, this rank's row, whose ranks hold the experts' shards in row order, and + as its data-parallel group, its column, the ranks holding the same shard. + """ + if state is None or not state.expert_model_parallel_group.is_initialized(): + return False + row, column = divmod(dist.rank(), plan.ep_size) + expert_parallel = range(row * plan.ep_size, (row + 1) * plan.ep_size) + data_parallel = range(column, plan.world_size, plan.ep_size) + return _group_ranks(state.expert_model_parallel_group) == list(expert_parallel) and sorted( + _group_ranks(state.data_parallel_group) + ) == list(data_parallel) + + +def _refuse_unsupported_distributed_calibration( + model: nn.Module, algorithm: QuantizeAlgoCfgType +) -> None: + """Refuse a calibration that would come out wrong on a model distributed over several ranks. + + Each rank of such a model calibrates on its own data, so an algorithm that keeps its statistics + on the rank would calibrate the ranks' copies, or shards, of a weight differently. Under expert + parallelism each rank's per-expert quantizers hold its own experts, and calibration syncs them + over the groups their module's ``ParallelState`` names, every rank for a module converted + without one: a module holding routed experts must name the experts' expert-parallel group and, + for data parallelism, only the ranks holding the same experts, or any calibration would merge + different experts' scales. Called before anything changes, so the caller seeds the routed + experts' ``ParallelState`` before quantizing, as the expert-parallel load does; their quantized + modules, whose state calibration reads, are checked again once converted, before they calibrate. + """ + plan = dist.recorded_distribution_plan(model) + if plan is None or plan.world_size == 1: + return + refused = _rank_local_calibrations(algorithm) + if refused: + raise NotImplementedError( + f"Calibrating with {', '.join(refused)} is not supported on a model distributed over " + f"{plan.world_size} ranks. Each rank calibrates on its own data and this calibration " + "keeps what it gathers on the rank, so the ranks would calibrate their copies, or their " + "shards, of a weight differently. Calibrate with max or MSE with distributed_sync on, " + "GPTQ, SmoothQuant or awq_clip, or in a single process." + ) + if algorithm is None or not plan.is_expert_parallel: + return + # Routed experts: the modules carrying an expert-parallel group, and the quantized modules under + # a container named ``experts``, as every expert-parallel plan in transformers names it. + holders = {} + for name, module in model.named_modules(): + state = getattr(module, "_parallel_state", None) + if (state is not None and state.expert_model_parallel_group.is_initialized()) or ( + isinstance(module, QuantModule) and "experts" in name.split(".") + ): + holders[name] = state + if not holders and not any("experts" in name.split(".") for name, _ in model.named_modules()): + return # no routed experts here, as in a VLM's vision tower: no experts' scales to merge + wrong = [name for name, state in holders.items() if not _names_the_experts_groups(state, plan)] + if not holders or wrong: + raise NotImplementedError( + f"Calibrating a model whose experts are sharded over {plan.ep_size} ranks needs each " + "routed expert's ParallelState to name, as its expert-parallel group, this rank's row " + f"of the ({plan.edp}, {plan.ep_size}) expert-parallel mesh and, as its data-parallel " + f"group, its column, the {plan.edp} rank(s) holding the same experts, or calibration " + "would sync their scales over the wrong ranks. " + + ( + f"These modules do not: {wrong[:3]}." + if wrong + else "None of the model's modules carries an expert-parallel group." + ) + ) + def _refuse_awq_on_a_sharded_fsdp2_model( model: nn.Module, algorithm: QuantizeAlgoCfgType, candidate: str | None = None @@ -128,6 +238,7 @@ def calibrate( Returns: The calibrated pytorch model. """ _refuse_awq_on_a_sharded_fsdp2_model(model, algorithm) + _refuse_unsupported_distributed_calibration(model, algorithm) if forward_loop is not None: # get the number of arguments of forward_loop num_args = len(inspect.signature(forward_loop).parameters) @@ -502,6 +613,7 @@ def forward_loop(model) -> None: """ quantize_config = QuantizeConfig(**dict(config)) _refuse_awq_on_a_sharded_fsdp2_model(model, config.get("algorithm")) + _refuse_unsupported_distributed_calibration(model, config.get("algorithm")) if not is_quantized(model): model = apply_mode(model, mode=[("quantize", dict(config))], registry=QuantizeModeRegistry) else: @@ -940,6 +1052,14 @@ def forward_backward_step(model, batch) -> None: if isinstance(quantization_formats, str) or not isinstance(quantization_formats, Sequence): raise TypeError("`quantization_formats` must be a sequence of formats.") quantization_formats = list(quantization_formats) + plan = dist.recorded_distribution_plan(model) + if plan is not None and plan.is_expert_parallel: # before anything converts the model + raise NotImplementedError( + f"auto_quantize does not support a model whose experts are sharded over {plan.ep_size} " + "ranks yet: its forwards would run the experts without the expert-parallel forwards " + "that quantization replaces and distributed_forward_loop restores. Quantize it with " + "mtq.quantize through distributed_forward_loop, or AutoQuantize in a single process." + ) method_name, method_config = _parse_auto_quantize_method(method) is_kv_search = constraints is not None and constraints.get("cost_model") == COST_MODEL_KV_CACHE @@ -1061,6 +1181,7 @@ def forward_backward_step(model, batch) -> None: ) # Every candidate, before the search calibrates any: calibrate() would refuse mid-search. _refuse_awq_on_a_sharded_fsdp2_model(model, algo, candidate=name) + _refuse_unsupported_distributed_calibration(model, algo) # Select the appropriate searcher based on method if method_name not in AUTO_QUANTIZE_SEARCHERS: diff --git a/modelopt/torch/utils/dataset_utils.py b/modelopt/torch/utils/dataset_utils.py index 1b82a045390..124cded7a38 100644 --- a/modelopt/torch/utils/dataset_utils.py +++ b/modelopt/torch/utils/dataset_utils.py @@ -30,6 +30,7 @@ from torch.utils.data.distributed import DistributedSampler from tqdm import tqdm +from .distributed import all_reduce, recorded_distribution_plan from .logging import print_rank_0, warn_rank_0 if TYPE_CHECKING: @@ -1005,17 +1006,60 @@ def get_max_batch_size( """Get the maximum batch size that can be used for the model. ``forward_step(model, inputs)`` optionally probes a custom workload, including during - OOM retries. By default, probe a model forward (or generation for encoder-decoder models). + OOM retries (which a model distributed over several ranks skips, below). By default, probe a + model forward (or generation for encoder-decoder models). + + For a model whose recorded plan distributes it over several ranks + (:func:`~modelopt.torch.utils.distributed.set_distribution_plan`) it is collective: every rank + calls it, and every rank gets the smallest batch size any rank's GPU allows + (:func:`~modelopt.torch.utils.distributed.all_reduce`), so that all run the same number of + calibration steps. Each rank then measures the GPU it computes on, from one forward of the + single-sample input, and does not probe the batch size it arrives at: retrying a smaller batch + after an out-of-memory error would run a different number of forwards, and of the collectives + inside them, on each rank. Probing in lockstep would not make it safe either: a rank that runs out + of memory partway through a forward has left its peers waiting in that forward's collectives. If + the estimate does not fit, pass an explicit batch size instead. A model with no plan recorded, + such as a Megatron model, is measured on every visible GPU, as in a single process. """ + args = ( + model, + max_sample_length, + sample_memory_usage_ratio, + sample_input_single_batch, + enable_grad, + forward_step, + ) + plan = recorded_distribution_plan(model) + if plan is not None and plan.world_size > 1: + # The device this rank computes on, which LOCAL_RANK does not name under a launcher that + # shows each rank one GPU. + device = torch.device("cuda", torch.cuda.current_device()) + return all_reduce("calib_batch_size", lambda: _max_batch_size(*args, device=device), min) + return _max_batch_size(*args) + + +def _max_batch_size( + model: torch.nn.Module, + max_sample_length: int, + sample_memory_usage_ratio: float, + sample_input_single_batch: torch.Tensor | None, + enable_grad: bool, + forward_step: Callable[[torch.nn.Module, torch.Tensor], Any] | None, + device: torch.device | None = None, +) -> int: + """:func:`get_max_batch_size` measured on ``device`` with one probe, or on every visible GPU.""" + devices = range(torch.cuda.device_count()) if device is None else [device] def _get_free_gpu_mem(): - min_gpu_free_mem = torch.cuda.get_device_properties(0).total_memory + min_gpu_free_mem = torch.cuda.get_device_properties( + 0 if device is None else device + ).total_memory max_allocated_mem = 0 - for device in range(torch.cuda.device_count()): - free_mem = torch.cuda.mem_get_info(device)[0] + for gpu in devices: + free_mem = torch.cuda.mem_get_info(gpu)[0] if free_mem < min_gpu_free_mem: min_gpu_free_mem = free_mem - max_allocated_mem = torch.cuda.max_memory_allocated(device) + max_allocated_mem = torch.cuda.max_memory_allocated(gpu) return min_gpu_free_mem, max_allocated_mem torch.cuda.empty_cache() @@ -1058,27 +1102,27 @@ def infer_method(inputs): else: target_data_batch = max(int(free_mem_before / mem_diff_per_data_batch), 1) - def _expand_to(batch: int) -> torch.Tensor: - return sample_input_single_batch.expand( - [ - batch if index == 0 else dim - for index, dim in enumerate(sample_input_single_batch.shape) - ] - ) - - target_input = _expand_to(target_data_batch) - # For some models on multi GPU, we observe the memory per batch is not a constant. - # So we just test the target batch size and make sure we do not go OOM. - while target_data_batch > 1: - with torch.set_grad_enabled(enable_grad): - try: - infer_method(target_input) - break - except torch.cuda.OutOfMemoryError: # pragma: no cover - GPU OOM retry path - target_data_batch = target_data_batch // 2 # pragma: no cover - target_input = _expand_to(target_data_batch) # pragma: no cover - torch.cuda.empty_cache() # pragma: no cover + # So we just test the target batch size and make sure we do not go OOM. A model distributed + # over several ranks is measured once, not probed (see get_max_batch_size). + if device is None: + + def _expand_to(batch: int) -> torch.Tensor: + return sample_input_single_batch.expand( + [ + batch if index == 0 else dim + for index, dim in enumerate(sample_input_single_batch.shape) + ] + ) + + while target_data_batch > 1: + with torch.set_grad_enabled(enable_grad): + try: + infer_method(_expand_to(target_data_batch)) + break + except torch.cuda.OutOfMemoryError: # pragma: no cover - GPU OOM retry path + target_data_batch = target_data_batch // 2 # pragma: no cover + torch.cuda.empty_cache() # pragma: no cover # Regulate the data batch target to be 1, 2, 4, 8, 12, ..., capped at 64 if target_data_batch < 2: diff --git a/modelopt/torch/utils/distributed.py b/modelopt/torch/utils/distributed.py index 26892ef112d..7528af9e420 100644 --- a/modelopt/torch/utils/distributed.py +++ b/modelopt/torch/utils/distributed.py @@ -16,12 +16,15 @@ """Utility functions for using torch.distributed.""" import functools +import gc import io +import itertools import os import sys import time import traceback -from collections.abc import Callable, Iterator +import types +from collections.abc import Callable, Iterator, Mapping from contextlib import contextmanager, suppress from dataclasses import dataclass from datetime import timedelta @@ -32,6 +35,15 @@ import torch.distributed from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, fully_shard from torch.distributed.tensor import DTensor +from torch.utils.data import ( + BatchSampler, + DataLoader, + IterableDataset, + RandomSampler, + SequentialSampler, +) +from torch.utils.data.distributed import DistributedSampler +from tqdm import tqdm __all__ = [ "DistributedProcessGroup", @@ -42,6 +54,7 @@ "backend", "barrier", "broadcast_from", + "distributed_forward_loop", "distributed_run", "distribution_plan", "fsdp2_wrap", @@ -54,7 +67,9 @@ "is_initialized", "is_master", "rank", + "recorded_distribution_plan", "set_distribution_plan", + "shard_dataloader", "size", ] @@ -396,7 +411,7 @@ def distribution_plan(model: torch.nn.Module) -> DistributionPlan: a launch sets its group up, each of its processes runs alone with a model of its own, so a model with no plan reads as a single process's there. """ - plan = getattr(model, _PLAN_ATTR, None) + plan = recorded_distribution_plan(model) if plan is not None: return plan sharded = is_fsdp2_model(model) @@ -410,6 +425,15 @@ def distribution_plan(model: torch.nn.Module) -> DistributionPlan: return DistributionPlan() +def recorded_distribution_plan(model: torch.nn.Module) -> DistributionPlan | None: + """The plan :func:`set_distribution_plan` recorded on ``model``, or ``None`` if none was. + + For code that also serves models distributed without a plan, such as Megatron's, which + :func:`distribution_plan` refuses in a run of several ranks. + """ + return getattr(model, _PLAN_ATTR, None) + + _V = TypeVar("_V") @@ -666,6 +690,269 @@ def distributed_run( _destroy_group(exc) +class _StridedIterableShard(IterableDataset): + """Rank-strided shard of a streaming ``IterableDataset`` (no ``len()`` required). + + ``DistributedSampler`` only works with map-style datasets. Streaming calibration datasets (e.g. the + Nemotron image-text ``NemotronTarPlusJsonlIterable``) are iterable, so each rank instead builds its + own copy of the stream, the same on every rank as :func:`shard_dataloader` requires, + and takes slot ``rank`` from every complete group of ``world_size`` consecutive samples. Any + trailing incomplete group is dropped so every rank yields exactly the same count -- preserving + the equal-per-rank-batch invariant that keeps per-step EP collectives from hanging when one rank + runs short. A stream shorter than ``world_size`` has no complete group, so it is refused rather + than leaving every rank with nothing to calibrate on. It is iterated in the loader's own process + (:func:`shard_dataloader` refuses workers for a stream). Every rank still reads and decodes every + sample, keeping one in ``world_size``: for a stream costly to decode, a map-style dataset, or a + source that shards itself, saves that work. + ``vlm_dataset_utils._ShardedIterable`` cannot stand in: it pads every rank to a count its own loader + knows, while this wraps any loader's stream, which may know none. + """ + + def __init__(self, base: IterableDataset, rank: int, world_size: int): + self.base = base + self.rank = rank + self.world_size = world_size + + def __iter__(self): + it = iter(self.base) + for index in itertools.count(): + group = list(itertools.islice(it, self.world_size)) + if len(group) < self.world_size: + if index == 0: + # shard_dataloader requires the same stream on every rank and worker, so all of + # them refuse it together. + raise ValueError( + f"The calibration stream has {len(group)} sample(s), fewer than the " + f"{self.world_size} ranks: no rank would get a sample to calibrate on." + ) + return # drop incomplete trailing group -> equal counts across ranks + yield group[self.rank] + + +def shard_dataloader(loader: DataLoader, rank: int, world_size: int) -> DataLoader: + """Wrap a DataLoader so each rank sees a unique shard. + + For map-style datasets a ``DistributedSampler`` is used. ``drop_last=False`` keeps per-rank batch + counts equal (else a rank exits calibration early and hangs the others on the per-step collectives), + at the cost of the sampler repeating up to ``world_size - 1`` samples to pad the even split. It + replaces the loader's sampler, so a shuffling one becomes the rank's unshuffled share, and one that + picks samples itself (a subset, a sampler over another source than the loader's dataset, a + random sampler drawing its own number of samples or with replacement, or a batch sampler other + than a plain ``BatchSampler``) is refused rather than silently widened. For streaming + ``IterableDataset`` datasets (which have no ``len()`` + and so can't use a sampler) each rank strides its own copy of the stream via + :class:`_StridedIterableShard`. That stream must yield the same samples in the same order on + every rank, or the shards overlap and the ranks' step counts can differ, which hangs the + collectives: one that shuffles needs a fixed seed of its own, as the ranks' RNGs need not agree. + A stream loaded with workers is refused: a dataset that splits itself across them by + ``get_worker_info()``, as a Hugging Face streaming dataset does, would be split twice and + silently lose samples. + + Forwards all non-sampler DataLoader settings from ``loader`` (batching, workers, pinning, + prefetch, init fn, generator, ...). With one rank there is nothing to shard: ``loader`` itself + comes back, its sampler and shuffling included. + """ + if world_size == 1: + return loader + batching = loader.batch_sampler + if not isinstance(loader.dataset, IterableDataset): + # A plain BatchSampler, built from batch_size or passed in, batches what its sampler picks; + # with auto-batching off there is none, and the loader's sampler picks each item. + if type(batching) is BatchSampler: + picker = batching.sampler + else: + picker = loader.sampler if batching is None else batching + known = type(picker) in (SequentialSampler, RandomSampler, DistributedSampler) + # A sampler over another source than the loader's dataset, or a RandomSampler drawing its + # own number of samples or with replacement, picks which and how many it reads, as a subset + # does; one that orders or shuffles the whole dataset does not. + source = getattr(picker, "data_source", getattr(picker, "dataset", ())) + narrowing = known and ( + len(source) != len(loader.dataset) + or ( + isinstance(picker, RandomSampler) + and (picker.replacement or picker.num_samples != len(loader.dataset)) + ) + ) + if not known or narrowing: + raise ValueError( + f"shard_dataloader: the loader's {type(picker).__name__} picks which samples it " + "reads, and the rank's DistributedSampler would replace it; pass a loader over " + "just the samples to calibrate on (a torch.utils.data.Subset) instead." + ) + if isinstance(loader.dataset, IterableDataset): + if loader.num_workers: + raise ValueError( + f"shard_dataloader: a streaming loader with num_workers={loader.num_workers} is " + "refused: a dataset that splits itself across the workers by get_worker_info() " + "would be split twice, by worker and by rank, and lose samples. Load it with " + "num_workers=0." + ) + return DataLoader( + _StridedIterableShard(loader.dataset, rank, world_size), + batch_size=loader.batch_size, + drop_last=loader.drop_last, + collate_fn=loader.collate_fn, + num_workers=loader.num_workers, + pin_memory=loader.pin_memory, + timeout=loader.timeout, + worker_init_fn=loader.worker_init_fn, + multiprocessing_context=loader.multiprocessing_context, + generator=loader.generator, + prefetch_factor=loader.prefetch_factor, + persistent_workers=loader.persistent_workers, + pin_memory_device=getattr(loader, "pin_memory_device", ""), + ) + sampler = DistributedSampler( + loader.dataset, + num_replicas=world_size, + rank=rank, + shuffle=False, + drop_last=False, + ) + return DataLoader( + loader.dataset, + batch_size=loader.batch_size if batching is None else batching.batch_size, + drop_last=loader.drop_last if batching is None else batching.drop_last, + sampler=sampler, + collate_fn=loader.collate_fn, + num_workers=loader.num_workers, + pin_memory=loader.pin_memory, + timeout=loader.timeout, + worker_init_fn=loader.worker_init_fn, + multiprocessing_context=loader.multiprocessing_context, + generator=loader.generator, + prefetch_factor=loader.prefetch_factor, + persistent_workers=loader.persistent_workers, + pin_memory_device=getattr(loader, "pin_memory_device", ""), + ) + + +def _reinstall_a2a_forwards(model, ep_mesh) -> int: + """Re-plant the torch_a2a TP forward wrappers that quantizer insertion clobbered. + + transformers installs its TP input/output transforms by REPLACING ``module.forward`` with an + instance attribute (TensorParallelLayer.install_forward). For torch_a2a_experts that wrapper is + what passes the EP ``process_group`` into the experts forward. ModelOpt's quantizer insertion + rebinds ``module.forward`` on the same instance, dropping the wrapper -- after which + ``torch_a2a_experts_forward`` receives ``process_group=None`` and raises. + + Called at the top of the calibration loop: after conversion (so it wraps the QUANTIZED forward) + and before the first forward. Only a forward quantizer insertion rebound is planted again: a + wrapper still in place -- transformers' from the load, on a module quantization left alone, or + one planted here on an earlier pass -- is kept, since wrapping it again would run its transforms + twice. ``ep_mesh`` is the mesh the routed experts are sharded over; without one there is nothing + to re-plant. + """ + if ep_mesh is None: + return 0 + try: + from transformers.distributed.tensor_parallel import ALL_PARALLEL_STYLES + except ImportError: # older transformers + from transformers.integrations.tensor_parallel import ALL_PARALLEL_STYLES + try: + from transformers.distributed.tensor_parallel import _get_parameter_tp_plan + except ImportError: + from transformers.integrations.tensor_parallel import _get_parameter_tp_plan + + plan = getattr(model, "tp_plan", None) or {} + if not plan: + return 0 + + n = 0 + for name, mod in model.named_modules(): + style_name = _get_parameter_tp_plan(parameter_name=name, tp_plan=plan, is_weight=False) + if style_name not in ("torch_a2a_experts", "torch_a2a_router"): + continue + # install_forward() -- and the clobber it causes -- only exist on transformers >=5.16. Older + # releases install their TP input/output transforms as forward PRE-HOOKS, which survive + # ModelOpt rebinding ``module.forward``: nothing to re-plant. Checked per style, the one + # about to be called. + style = ALL_PARALLEL_STYLES.get(style_name) + if not hasattr(style, "install_forward"): + continue + # install_forward stores a plain function as the instance's forward, and quantizer + # insertion (bind_forward_method) a bound method over it: only the latter lost its wrapper. + if not isinstance(vars(mod).get("forward"), types.MethodType): + continue + style.install_forward(mod, ep_mesh) + n += 1 + if n and is_master(): + print( + f"[ep] re-installed {n} torch_a2a forward wrapper(s) clobbered by quantizer insertion", + flush=True, + ) + return n + + +def _release_cached_cuda_memory(): + """Hand the caching allocator's reserved-but-unused pool back to the driver. + + Only a process that started CUDA has a pool to hand back; elsewhere the collection alone costs + about a second. + """ + if torch.cuda.is_initialized(): + gc.collect() + torch.cuda.empty_cache() + + +def distributed_forward_loop( + model: torch.nn.Module, + dataloader: DataLoader, + device: torch.device | str | None = None, + ep_mesh: "torch.distributed.device_mesh.DeviceMesh | None" = None, +) -> Callable[[torch.nn.Module], None]: + """Build an ``mtq.quantize`` ``forward_loop`` for a distributed (e.g. expert-parallel) model. + + Moves each batch onto ``device``, when given, and runs the captured ``model`` forward (the EP + all-reduce / all-to-all collectives fire inside that forward). A model whose plan + (:func:`distribution_plan`) is expert-parallel passes the ``ep_mesh`` its experts are sharded + over, for the all-to-all forwards that quantizer insertion drops to be re-installed before the + first forward; any other model passes none. ``mtq.quantize`` hands the + ``forward_loop`` callback an inner module that this closure ignores, calling the captured + top-level ``model`` instead so the full distributed forward (and its collectives) runs. + + The forwards run without gradients and with KV/state caching disabled (``_disable_use_cache``), as + the single-node ``_forward_loop`` runs them: a hybrid Mamba/attention model mutates its cache in + place. Each batch is a mapping of the forward's keyword arguments, as ``model(**batch)`` takes; + only its tensor values move to ``device``, so tensors nested in a value (a list of images, say) + must already be there. + """ + # Local import to avoid any import cycle within modelopt.torch.utils. + from modelopt.torch.utils.dataset_utils import _disable_use_cache + + plan = distribution_plan(model) + if plan.is_expert_parallel != (ep_mesh is not None): + raise ValueError( + f"{plan!r}: an expert-parallel model calibrates with the ep_mesh its routed experts " + "are sharded over, and any other model without one." + ) + if ep_mesh is not None and ep_mesh.size() != plan.ep_size: + raise ValueError(f"{plan!r}: the ep_mesh spans {ep_mesh.size()} ranks, not {plan.ep_size}.") + + def calibrate(_inner_model): + _reinstall_a2a_forwards(model, ep_mesh) + with _disable_use_cache(model), torch.no_grad(): + # Hand the caching allocator's unused pool back: NCCL allocates its buffers outside it. + _release_cached_cuda_memory() + for batch in tqdm(dataloader, desc="Calibrating", disable=not is_master()): + if not isinstance(batch, Mapping): + raise TypeError( + f"distributed_forward_loop: the loader yields {type(batch).__name__}, not a " + "mapping of the forward's keyword arguments." + ) + if device is not None: + batch = { + k: (v.to(device) if isinstance(v, torch.Tensor) else v) + for k, v in batch.items() + } + model(**batch) + # Again for what runs next (amax sync, export collectives), past the forwards' pool. + _release_cached_cuda_memory() + + return calibrate + + def cleanup(): """Cleans up the distributed environment. diff --git a/tests/gpu/torch/quantization/test_calib_cuda.py b/tests/gpu/torch/quantization/test_calib_cuda.py index 677a03dc0e2..5ee12738cb9 100644 --- a/tests/gpu/torch/quantization/test_calib_cuda.py +++ b/tests/gpu/torch/quantization/test_calib_cuda.py @@ -16,9 +16,11 @@ """Calibration tests.""" import torch +import torch.distributed as dist import torch.nn as nn import modelopt.torch.quantization as mtq +from modelopt.torch.utils.distributed import DistributionPlan, ParallelState, set_distribution_plan class _TwoBranchModel(nn.Module): @@ -71,3 +73,56 @@ def _forward_loop(m): if uncal_q.amax is not None: assert uncal_q.axis is None assert uncal_q.amax.numel() == 1 + + +def _nvfp4_linears(device, count): + """``count`` copies of one linear: tiny weights, whose FP8 block scales a tighter per-tensor + amax resolves better, and a large one per row on input channel 0, which a tighter amax clips.""" + model = nn.Sequential(*(nn.Linear(64, 16, bias=False) for _ in range(count))).to(device) + for linear in model: + with torch.no_grad(): + linear.weight.copy_(torch.linspace(-1e-3, 1e-3, 16 * 64, device=device).view(16, 64)) + linear.weight[:, 0] = 100.0 + return model + + +def _nvfp4_rank_input(rank, device): + x = torch.randn(8, 64, device=device, generator=torch.Generator(device).manual_seed(rank)) + if rank % 2 == 0: + x[:, 0] = 0 # never reaches the large weights: the tightest clip suits it + else: + x[:, 1:] = 0 # reaches only them: no clip suits it + return x + + +def _nvfp4_awq_clip_on_every_rank_job(rank, size): + device = torch.device(f"cuda:{rank}") + model = _nvfp4_linears(device, 2) + set_distribution_plan(model, DistributionPlan(size)) + x = _nvfp4_rank_input(rank, device) + + def forward_loop(m): + m[0](x) + # The second runs on no tokens on the odd ranks, as an MoE expert no token reaches can. + m[1](x if rank % 2 == 0 else x[:0]) + + mtq.quantize(model, mtq.NVFP4_AWQ_CLIP_CFG, forward_loop) + amax = torch.stack([linear.weight_quantizer.amax.float().reshape(()) for linear in model]) + clips = [torch.empty_like(amax) for _ in range(size)] + dist.all_gather(clips, amax) + assert all(torch.equal(clip, amax) for clip in clips), clips + # The second's clip is the one the even ranks' tokens pick, calibrated in one process. + alone = _nvfp4_linears(device, 1) + alone[0]._parallel_state = ParallelState(data_parallel_group=-1) + inputs = [_nvfp4_rank_input(even, device) for even in range(0, size, 2)] + mtq.quantize(alone, mtq.NVFP4_AWQ_CLIP_CFG, lambda m: [m(even) for even in inputs]) + assert amax[1] > 0, amax + assert torch.equal(amax[1], alone[0].weight_quantizer.amax.float().reshape(())), amax + + +def test_nvfp4_awq_clip_picks_the_same_clip_on_every_rank(need_2_gpus, dist_workers): + """NVFP4's awq_clip clips the per-tensor amax. Each rank of a model distributed over several + ranks calibrates on its own data, and averaging each candidate clip's loss over the ranks gives + every rank's copy of a weight the same clip, though each rank's own data would pick another. A + rank that runs a weight on no tokens leaves it the clip the other ranks' tokens pick.""" + dist_workers.run(_nvfp4_awq_clip_on_every_rank_job) diff --git a/tests/gpu/torch/quantization/test_fsdp2.py b/tests/gpu/torch/quantization/test_fsdp2.py index aec8e37f813..632a4ccddac 100644 --- a/tests/gpu/torch/quantization/test_fsdp2.py +++ b/tests/gpu/torch/quantization/test_fsdp2.py @@ -21,10 +21,12 @@ import pytest import torch +import torch.distributed as dist import torch.nn as nn -from _test_utils.torch.distributed.utils import synchronize_state_dict +from _test_utils.torch.distributed.utils import bounded_collectives, synchronize_state_dict from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard from torch.distributed.tensor import DTensor +from torch.utils.data import DataLoader import modelopt.torch.quantization as mtq import modelopt.torch.quantization.model_calib as model_calib @@ -35,7 +37,13 @@ persistent_materialization, ) from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector -from modelopt.torch.utils.dataset_utils import _forward_loop +from modelopt.torch.utils.dataset_utils import _forward_loop, get_max_batch_size +from modelopt.torch.utils.distributed import ( + DistributionPlan, + distributed_forward_loop, + set_distribution_plan, + shard_dataloader, +) def _test_fsdp2_simple_linear(rank, size): @@ -461,3 +469,65 @@ def _test_sharded_root_calibration(rank, size): def test_sharded_root_calibration(dist_workers): dist_workers.run(_test_sharded_root_calibration) + + +def _test_planned_batch_size_then_max_calibration(rank, size): + torch.manual_seed(0) + model = nn.Sequential(nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 64)).cuda(rank) + synchronize_state_dict(model) + fully_shard(model[0]) + fully_shard(model[2]) + model = fully_shard(model) + set_distribution_plan(model, DistributionPlan(size, dense="fsdp")) + forwards = [] + model.register_forward_pre_hook(lambda module, args: forwards.append(rank)) + # What a sample holds in the measured forward: a fortieth of rank 0's free GPU memory, so its + # own estimate is about 40 samples, and a sixth of every other rank's, about 6 (4 regulated). + held = torch.cuda.mem_get_info()[0] // (40 if rank == 0 else 6) + + def forward_step(model, inputs): + model(inputs) + return torch.empty(held, dtype=torch.uint8, device=inputs.device) + + # A rank left waiting in the batch size's reduction, or in a gather below, gives up in seconds. + with bounded_collectives(seconds=60): + batch_size = get_max_batch_size( + model, sample_input_single_batch=torch.ones(1, 64).cuda(rank), forward_step=forward_step + ) + sizes = [None] * size + dist.all_gather_object(sizes, batch_size) + # Every rank takes the smallest any rank's GPU allows, not its own (rank 0's is about 40). + assert sizes == [batch_size] * size and batch_size <= 4, sizes + + # An uneven split, which the shard pads so that every rank runs the same steps. + data = torch.randn( + 2 * size * batch_size + 1, 64, generator=torch.Generator().manual_seed(1) + ) + loader = DataLoader( + data, batch_size=batch_size, collate_fn=lambda rows: {"input": torch.stack(rows)} + ) + forwards.clear() + forward_loop = distributed_forward_loop( + model, shard_dataloader(loader, rank, size), device=torch.device("cuda", rank) + ) + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, forward_loop) + amaxes = { + name: module.amax.float().cpu() + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) and module.is_enabled and module.amax is not None + } + gathered = [None] * size + dist.all_gather_object(gathered, (len(forwards), amaxes)) + + assert amaxes, "no quantizer was calibrated" + for steps, theirs in gathered: + assert steps == len(forwards) == 3, gathered # two whole batches and one of one sample + assert theirs.keys() == amaxes.keys() + assert all(torch.equal(theirs[name], amax) for name, amax in amaxes.items()), theirs + + +def test_planned_batch_size_then_max_calibration(need_2_gpus, dist_workers): + """A model whose plan spreads it over several ranks sizes its calibration batch once for all of + them, the smallest any rank's GPU allows, then calibrates each rank's shard with it: every rank + runs the same steps, and max calibration gives every rank the same scales.""" + dist_workers.run(_test_planned_batch_size_then_max_calibration) diff --git a/tests/gpu/torch/utils/test_moe_a2a.py b/tests/gpu/torch/utils/test_moe_a2a.py index 9b8a3be53a4..dfd8c704d27 100644 --- a/tests/gpu/torch/utils/test_moe_a2a.py +++ b/tests/gpu/torch/utils/test_moe_a2a.py @@ -16,13 +16,29 @@ """GPU tests for the ``torch_a2a`` expert-parallel MoE transport (``a2a_moe``).""" import copy +import dataclasses +from datetime import timedelta import pytest import torch import torch.distributed as dist +from _test_utils.torch.distributed.utils import bounded_collectives from packaging.version import Version +from torch.distributed.device_mesh import init_device_mesh +import modelopt.torch.quantization as mtq import modelopt.torch.utils.moe as moe +from modelopt.torch.quantization.nn import TensorQuantizer +from modelopt.torch.utils.distributed import ( + DistributionPlan, + ParallelState, + distributed_forward_loop, + set_distribution_plan, +) + +transformers = pytest.importorskip("transformers") +from transformers import AutoConfig, AutoModelForCausalLM, Qwen3MoeConfig +from transformers.distributed import DistributedConfig # One MoE layer's expert stage, big enough that the order of each token's top-k adds shows in bf16. TOKENS, HIDDEN, EXPERTS, TOP_K = 4096, 1024, 16, 8 @@ -71,3 +87,132 @@ def test_torch_a2a_gives_grouped_mms_result_bit_for_bit(dist_workers): if Version(transformers.__version__) < Version("5.6"): pytest.skip("expert parallelism needs transformers >= 5.6") dist_workers.run(_grouped_mm_bit_for_bit_job) + + +NUM_EXPERTS = 8 + + +def _save_tiny_qwen3_moe(ckpt_dir): + config = Qwen3MoeConfig( + hidden_size=64, + intermediate_size=128, + moe_intermediate_size=32, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + num_experts=NUM_EXPERTS, + num_experts_per_tok=2, + vocab_size=128, + max_position_embeddings=64, + ) + torch.manual_seed(0) + AutoModelForCausalLM.from_config(config).to(torch.bfloat16).save_pretrained(ckpt_dir) + + +def _load_expert_parallel(ckpt_dir, size): + """The model with its experts sharded over every rank by transformers' own expert-parallel load, + on the ``torch_a2a`` styles and experts forward, and the mesh they are sharded over.""" + moe.register_torch_a2a() + config = AutoConfig.from_pretrained(ckpt_dir) + swap = {"ep_router": "torch_a2a_router", "moe_tp_experts": "torch_a2a_experts"} + config.base_model_ep_plan = {k: swap.get(v, v) for k, v in config.base_model_ep_plan.items()} + mesh = init_device_mesh("cuda", (size,), mesh_dim_names=("ep",)) + # Each option where this release takes it: on DistributedConfig when it has a field for it. + fields = {field.name for field in dataclasses.fields(DistributedConfig)} + options = {"tp_plan": "auto", "tp_size": size} + model = AutoModelForCausalLM.from_pretrained( + ckpt_dir, + config=config, + dtype=torch.bfloat16, + device_mesh=mesh, + distributed_config=DistributedConfig( + enable_expert_parallel=True, **{k: v for k, v in options.items() if k in fields} + ), + experts_implementation="torch_a2a", + **({} if "tp_plan" in fields else {"tp_plan": "auto"}), + ) + for module in model.modules(): # the experts forward reads it at forward time + if hasattr(getattr(module, "config", None), "_experts_implementation"): + module.config._experts_implementation = "torch_a2a" + return model, mesh + + +def _calibrated(model): + return { + name: module.amax.float().cpu() + for name, module in model.named_modules() + if isinstance(module, TensorQuantizer) and module.is_enabled and module.amax is not None + } + + +def _expert_parallel_calibration_job(rank, size, ckpt_dir): + device = torch.device("cuda", rank) + batches = [ + torch.randint(0, 128, (2, 16), generator=torch.Generator().manual_seed(r)) + for r in range(size) + ] + # The scales' syncs give up in seconds on this group of every rank: the experts' expert-parallel + # group (they span every rank), and the dense modules' data-parallel one. + group = dist.new_group(timeout=timedelta(seconds=120)) + dist.barrier(group=group) + try: + with bounded_collectives(seconds=120): + model, mesh = _load_expert_parallel(ckpt_dir, size) + for name, module in model.named_modules(): + module._parallel_state = ( + ParallelState(data_parallel_group=-1, expert_model_parallel_group=group) + if ".mlp.experts" in name + else ParallelState(data_parallel_group=group) + ) + set_distribution_plan(model, DistributionPlan(size, ep_size=size)) + # Each rank its own batch, through the quantized model's re-planted torch_a2a forwards. + forward_loop = distributed_forward_loop( + model, [{"input_ids": batches[rank]}], device=device, ep_mesh=mesh + ) + mtq.quantize(model, mtq.FP8_DEFAULT_CFG, forward_loop) + + # The whole model, calibrated in one process on every rank's batch. + whole = AutoModelForCausalLM.from_pretrained(ckpt_dir, dtype=torch.bfloat16).to(device) + for module in whole.modules(): + module._parallel_state = ParallelState(data_parallel_group=-1) + mtq.quantize(whole, mtq.FP8_DEFAULT_CFG, lambda m: [m(b.to(device)) for b in batches]) + + sharded, reference = _calibrated(model), _calibrated(whole) + problems, compared = [] if sharded else ["no quantizer was calibrated"], [] + per_rank = NUM_EXPERTS // size + for name, amax in sharded.items(): + if "_weight_quantizers." in name: # this rank's local expert -> the global one + prefix, index = name.rsplit(".", 1) + name = f"{prefix}.{rank * per_rank + int(index)}" + compared.append(name) + want = reference.get(name) + if want is None: + problems.append(f"{name} is calibrated here but not in the whole model") + continue + # bf16 activations, computed in another order: compare by relative norm. + error = ((amax - want).norm() / want.norm().clamp_min(1e-12)).item() + if not error < 2e-2: + problems.append(f"{name}: amax off by {error:.3g} (relative)") + gathered = [None] * size + dist.all_gather_object(gathered, (problems, compared)) + finally: + dist.destroy_process_group(group) + + found = [problem for problems, _ in gathered for problem in problems] + assert not found, "\n".join(found[:20]) + experts = {name for name in reference if "_weight_quantizers." in name} + assert experts and experts <= {name for _, names in gathered for name in names}, experts + + +def test_expert_parallel_calibration_matches_the_whole_model(need_2_gpus, dist_workers, tmp_path): + """transformers' own expert-parallel load on the torch_a2a transport, quantized and calibrated + through distributed_forward_loop, each rank on its own batch: every expert's weight scales and + every activation scale match the whole model calibrated on every rank's batch. Quantizer + insertion drops the transport's forward wrappers on transformers >= 5.16, which the loop + re-plants.""" + if Version(transformers.__version__) < Version("5.16"): + pytest.skip("the torch_a2a forward wrappers are pre-hooks below transformers 5.16") + if NUM_EXPERTS % torch.cuda.device_count(): + pytest.skip(f"{NUM_EXPERTS} experts do not shard over {torch.cuda.device_count()} GPUs") + _save_tiny_qwen3_moe(tmp_path) + dist_workers.run(_expert_parallel_calibration_job, ckpt_dir=str(tmp_path)) diff --git a/tests/unit/torch/quantization/test_calib.py b/tests/unit/torch/quantization/test_calib.py index 27a3dae1776..69c56ba8770 100644 --- a/tests/unit/torch/quantization/test_calib.py +++ b/tests/unit/torch/quantization/test_calib.py @@ -15,10 +15,12 @@ """Unittests for AWQ and SVDQuant""" +from datetime import timedelta from functools import partial import pytest import torch +import torch.distributed as dist import torch.nn as nn from _test_utils.torch.quantization.quantize_common import get_awq_config @@ -33,6 +35,7 @@ ) from modelopt.torch.quantization.nn import QuantLinear, TensorQuantizer from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector +from modelopt.torch.utils.distributed import DistributionPlan, ParallelState, set_distribution_plan class _SimpleMLP(nn.Module): @@ -266,6 +269,70 @@ def test_awq_clip(): _awq_clip_tester(model, model_ref, ref_data, config) +class _Experts(nn.Module): + """Two experts, each run only on the inputs routed to it, as an MoE runs them.""" + + def __init__(self, data_parallel_group): + super().__init__() + self.experts = nn.ModuleList(nn.Linear(16, 4, bias=False) for _ in range(2)) + for expert in self.experts: + with torch.no_grad(): + # Every rank's copy: small weights, and a large one per row on input channel 0. + expert.weight.copy_(torch.linspace(-0.1, 0.1, 64).view(4, 16)) + expert.weight[:, 0] = 1.0 + expert._parallel_state = ParallelState(data_parallel_group=data_parallel_group) + + def forward(self, calls): + for index, x in calls: + self.experts[index](x) + + +def _rank_input(rank): + x = torch.randn(8, 16, generator=torch.Generator().manual_seed(rank)) + if rank % 2 == 0: + x[:, 0] = 0 # never reaches the large weights: a tighter clip only sharpens the rest + else: + x[:, 1:] = 0 # reaches only them: clipping them costs the most + return x + + +def _awq_clip_on_every_rank_job(rank, size): + x = _rank_input(rank) + # The same number of steps on every rank, but expert 1 runs on both of rank 0's and on one of + # rank 1's, there on no tokens: an expert's forwards, unlike the steps, are each rank's own. + steps = [[(0, x), (1, x)]] * 2 if rank == 0 else [[(0, x), (1, x[:0])], [(0, x)]] + # A rank left waiting in the clip search's sync gives up in seconds on this group. + group = dist.new_group(timeout=timedelta(seconds=10)) + dist.barrier(group=group) + model = _Experts(group) + set_distribution_plan(model, DistributionPlan(size)) + try: + mtq.quantize(model, get_awq_config("awq_clip"), lambda m: [m(calls) for calls in steps]) + clips = [expert.weight_quantizer.amax for expert in model.experts] + gathered = [None] * size + dist.all_gather_object(gathered, clips, group=group) + finally: + dist.destroy_process_group(group) + assert all( + torch.equal(mine, theirs) for other in gathered for mine, theirs in zip(clips, other) + ), gathered + # Expert 1's clips are those rank 0's tokens pick alone, which rank 1's empty call leaves be. + alone = _Experts(data_parallel_group=-1) + x0 = _rank_input(0) + mtq.quantize(alone, get_awq_config("awq_clip"), lambda m: [m([(1, x0)]) for _ in range(2)]) + assert (clips[1] > 0).all(), clips[1] + assert torch.equal(clips[1], alone.experts[1].weight_quantizer.amax) + + +def test_awq_clip_picks_the_same_clips_on_every_rank(cpu_dist_workers): + """Each rank of a model distributed over several ranks calibrates on its own data. awq_clip + averages each candidate clip's loss over the data-parallel ranks, once all the forwards have + run, so the ranks' copies of a weight take the same clips, though each rank's own data would + pick others, and though an MoE's expert runs a different number of forwards on each rank, on + one of them on no tokens, which leaves the clips the other ranks' tokens pick.""" + cpu_dist_workers.run(_awq_clip_on_every_rank_job) + + @pytest.mark.parametrize("use_weight_scale", [None, False]) def test_awq_full(use_weight_scale): """Test awq.""" diff --git a/tests/unit/torch/utils/test_dataset_utils.py b/tests/unit/torch/utils/test_dataset_utils.py index 1666308e1d4..138cc35c1da 100644 --- a/tests/unit/torch/utils/test_dataset_utils.py +++ b/tests/unit/torch/utils/test_dataset_utils.py @@ -14,10 +14,12 @@ # limitations under the License. import json +from types import SimpleNamespace from unittest.mock import Mock, patch import pytest import torch +from _test_utils.torch.distributed.utils import bounded_collectives from _test_utils.torch.transformers_models import get_tiny_tokenizer from torch.utils.data import DataLoader @@ -34,6 +36,7 @@ get_max_batch_size, prepare_messages_for_chat_template, ) +from modelopt.torch.utils.distributed import DistributionPlan, set_distribution_plan def test_prepare_messages_for_chat_template(): @@ -436,6 +439,68 @@ def fake_call(x): assert result == 4 +class _BatchSizeProbe(torch.nn.Module): + """Records the batch size of every forward.""" + + def __init__(self): + super().__init__() + self.batch_sizes = [] + + def forward(self, x): + self.batch_sizes.append(x.shape[0]) + + +def _batch_size_job(rank, world_size, planned): + model = _BatchSizeProbe() + if planned: + set_distribution_plan(model, DistributionPlan(world_size)) + # A sample takes 100 bytes: rank 0's GPU has room for 4 of them, rank 1's for 10 (regulated to 8). + free = [400, 1000][rank] + # What one measurement reads: the rank's own GPU, or every visible one. + gpus = 1 if planned else 4 + measured = [] + + def mem_get_info(device): + measured.append(device) + return (free if len(measured) <= gpus else free - 100, 0) + + with ( + # A rank left waiting in the reduction gives up in seconds. + bounded_collectives(), + patch("torch.cuda.empty_cache"), + patch( + "torch.cuda.get_device_properties", + return_value=SimpleNamespace(total_memory=10**12), + ), + patch("torch.cuda.device_count", return_value=4), + # A launcher that shows each rank one GPU: every rank computes on cuda:0. + patch("torch.cuda.current_device", return_value=0), + patch("torch.cuda.mem_get_info", side_effect=mem_get_info), + patch("torch.cuda.max_memory_allocated", return_value=0), + ): + batch_size = get_max_batch_size( + model, max_sample_length=8, sample_input_single_batch=torch.ones((1, 8)) + ) + + if planned: + assert measured == [torch.device("cuda", 0)] * 2 # the GPU it computes on, not LOCAL_RANK + # One probe: an OOM retry would run more forwards on some ranks. + assert model.batch_sizes == [1] + assert batch_size == 4 # rank 0's + else: + assert measured == list(range(4)) * 2 # every visible GPU, as in a single process + assert batch_size == [4, 8][rank] # its own + + +@pytest.mark.parametrize("planned", [True, False], ids=["planned", "no-plan"]) +def test_get_max_batch_size_in_a_distributed_run(cpu_dist_workers, planned): + """Every rank must run the same number of calibration steps, so ranks agree on one batch size + for a model whose plan distributes it over them: the smallest any rank's GPU allows. A model + with no plan recorded, such as a Megatron model, is measured by each rank as a single process + measures it, and is not refused.""" + cpu_dist_workers.run(_batch_size_job, planned=planned) + + def test_get_dataset_samples_with_unsupported_dataset(make_toy_hf_dataset): """A dataset not in ``SUPPORTED_DATASET_CONFIG`` loads via the auto-detect path.""" dataset_name = make_toy_hf_dataset() # basename never matches a registered key diff --git a/tests/unit/torch/utils/test_distributed_calibration.py b/tests/unit/torch/utils/test_distributed_calibration.py new file mode 100644 index 00000000000..2efef9fe86d --- /dev/null +++ b/tests/unit/torch/utils/test_distributed_calibration.py @@ -0,0 +1,816 @@ +# 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. + +"""Unit tests for distributed calibration: per-rank data shards and the distributed forward loop. + +Each rank of a distributed (e.g. expert-parallel) PTQ run calibrates on its own shard of the data, +and the EP collectives inside every forward hang unless all ranks run the same number of batches. +The torch_a2a TP forward wrappers come from transformers >= 5.16 and need a process group, so they +are faked here; everything runs in one CPU process. +""" + +import re +import sys +import types +from collections import Counter, OrderedDict + +import pytest +import torch +from _test_utils.torch.distributed.utils import bounded_collectives +from torch import nn +from torch.utils.data import ( + BatchSampler, + DataLoader, + DistributedSampler, + IterableDataset, + RandomSampler, + SequentialSampler, + SubsetRandomSampler, +) + +import modelopt.torch.quantization as mtq +import modelopt.torch.utils.distributed as mto_dist +from modelopt.torch.quantization.model_quant import _refuse_unsupported_distributed_calibration +from modelopt.torch.utils.distributed import ( + DistributionPlan, + ParallelState, + _reinstall_a2a_forwards, + distributed_forward_loop, + set_distribution_plan, + shard_dataloader, +) + +_EXPERT_PARALLEL = DistributionPlan(world_size=2, ep_size=2) + +# ------------------------------------------------------------------------------------- data helpers + + +def _collate(ids): + """Calibration-style collate: a dict of batched tensors, as ``model(**batch)`` expects.""" + return {"input_ids": torch.tensor(ids).unsqueeze(-1)} + + +class _Stream(IterableDataset): + """A deterministic stream with no ``len()``, like the streaming image-text calibration sets.""" + + def __init__(self, n): + self.n = n + + def __iter__(self): + return iter(range(self.n)) + + +def _loader(dataset, batch_size=2): + return DataLoader(dataset, batch_size=batch_size, collate_fn=_collate) + + +def _batches(loader): + """The sample ids in each batch the loader yields.""" + return [batch["input_ids"].flatten().tolist() for batch in loader] + + +def _ids(loader): + return [i for batch in _batches(loader) for i in batch] + + +def _seed_worker(worker_id): + torch.manual_seed(worker_id) + + +def _dataset(kind, n): + """Samples ``0 .. n - 1`` as a list (map-style) or as a stream with no ``len()``.""" + return list(range(n)) if kind == "map-style" else _Stream(n) + + +# --------------------------------------------------------------------------------- shard_dataloader + + +@pytest.mark.parametrize("kind", ["map-style", "iterable"]) +@pytest.mark.parametrize("world_size", [1, 2, 3]) +def test_shards_are_strided_cover_every_sample_once_and_replay(kind, world_size): + """Rank ``r`` gets samples ``r, r + W, ...``, batched with the original size and collate_fn, and + a second pass replays the same shard: multi-pass calibration algorithms (e.g. AWQ) iterate the + loader again. A lone rank's shard is the whole (unshuffled) loader, batch for batch. + """ + n = 12 + shards = [ + shard_dataloader(_loader(_dataset(kind, n)), r, world_size) for r in range(world_size) + ] + + assert sorted(i for shard in shards for i in _ids(shard)) == list(range(n)) + for r, shard in enumerate(shards): + ids = list(range(r, n, world_size)) + assert _batches(shard) == _batches(shard) == [ids[i : i + 2] for i in range(0, len(ids), 2)] + + +@pytest.mark.parametrize( + ("kind", "padded", "dropped"), + [("map-style", 2, set()), ("iterable", 0, {8, 9})], + ids=["map-style", "iterable"], +) +def test_an_uneven_split_gives_every_rank_the_same_batch_count(kind, padded, dropped): + """No rank runs short: it would hang the others' EP collectives. A map-style loader repeats a + few samples to even the ranks out (``drop_last=False``); a stream, with no ``len()`` to pad to, + drops its trailing partial group instead. + """ + n, world_size = 10, 4 + shards = [ + _batches(shard_dataloader(_loader(_dataset(kind, n)), r, world_size)) + for r in range(world_size) + ] + + assert len({len(batches) for batches in shards}) == 1 + seen = Counter(i for batches in shards for batch in batches for i in batch) + assert set(seen) == set(range(n)) - dropped + assert seen.total() - (n - len(dropped)) == padded + + +def test_a_stream_shorter_than_the_world_is_refused_on_every_rank(): + """With fewer samples than ranks there is no complete group, so every rank would calibrate on + nothing; each rank refuses the stream instead, before any forward.""" + for rank in range(4): + with pytest.raises(ValueError, match="fewer than the 4 ranks"): + _ids(shard_dataloader(_loader(_Stream(3)), rank, 4)) + + +def test_a_stream_loaded_with_workers_is_refused(): + """A stream that splits itself across the loader's workers, as a Hugging Face streaming dataset + does, would be split again by rank and silently lose samples. A lone rank keeps its loader.""" + loader = DataLoader(_Stream(16), batch_size=2, collate_fn=_collate, num_workers=2) + + with pytest.raises(ValueError, match="num_workers=2 is refused"): + shard_dataloader(loader, 0, 2) + assert shard_dataloader(loader, 0, 1) is loader + + +@pytest.mark.parametrize( + "picker", + [ + {"sampler": SubsetRandomSampler(range(4)), "batch_size": 2}, + {"sampler": SubsetRandomSampler(range(4)), "batch_size": None}, + {"batch_sampler": [[0, 1], [2, 3]]}, + {"batch_sampler": BatchSampler(SubsetRandomSampler(range(4)), 2, drop_last=False)}, + {"sampler": RandomSampler(range(8), num_samples=4), "batch_size": 2}, + {"sampler": RandomSampler(range(8), replacement=True), "batch_size": 2}, + {"sampler": SequentialSampler(range(4)), "batch_size": 2}, + {"sampler": RandomSampler(range(4)), "batch_size": 2}, + {"sampler": DistributedSampler(range(4), num_replicas=2, rank=0), "batch_size": 2}, + ], + ids=[ + "subset-sampler", + "subset-sampler-unbatched", + "batch-sampler", + "subset-in-a-plain-batch-sampler", + "random-sampler-drawing-fewer", + "random-sampler-with-replacement", + "sequential-sampler-over-fewer", + "random-sampler-over-fewer", + "distributed-sampler-over-fewer", + ], +) +def test_a_loader_that_picks_its_own_samples_is_refused(picker): + """The rank's DistributedSampler would replace what it picks, silently calibrating on the whole + dataset instead, however the loader batches the picks. A lone rank has nothing to shard and + keeps the loader as it is.""" + loader = DataLoader(list(range(8)), collate_fn=_collate, **picker) + + with pytest.raises(ValueError, match="picks which samples it reads"): + shard_dataloader(loader, 0, 2) + assert shard_dataloader(loader, 0, 1) is loader + + +@pytest.mark.parametrize( + ("kind", "batching", "batches"), + [ + ( + "map-style", + {"batch_sampler": BatchSampler(SequentialSampler(range(10)), 2, drop_last=True)}, + [[1, 3], [5, 7]], + ), + ("map-style", {"batch_size": None}, [[1], [3], [5], [7], [9]]), + ("iterable", {"batch_size": None}, [[1], [3], [5], [7], [9]]), + ("map-style", {"batch_size": 2, "shuffle": True}, [[1, 3], [5, 7], [9]]), + ], + ids=["plain-batch-sampler", "auto-batching-off", "auto-batching-off-stream", "shuffled"], +) +def test_the_shard_batches_as_the_loader_does(kind, batching, batches): + """A loader given a plain BatchSampler reports ``batch_size=None``; its shard still batches by the + sampler's size and drop_last. A loader with auto-batching off reads one item per step, and so + does its shard. A loader shuffling the whole dataset gives the rank its unshuffled share.""" + loader = DataLoader(_dataset(kind, 10), collate_fn=_collate, **batching) + + assert _batches(shard_dataloader(loader, 1, 2)) == batches + + +@pytest.mark.parametrize("dataset", [list(range(8)), _Stream(8)], ids=["map-style", "iterable"]) +def test_sharded_loader_keeps_the_original_loader_settings(dataset): + """Batching, pinning, timeout, worker init and RNG carry over to the shard, and for a map-style + loader its workers and prefetching too (a stream is refused workers).""" + workers = ( + {} + if isinstance(dataset, _Stream) + else { + "num_workers": 2, + "multiprocessing_context": "spawn", + "prefetch_factor": 3, + "persistent_workers": True, + } + ) + loader = DataLoader( + dataset, + batch_size=4, + drop_last=True, + collate_fn=_collate, + pin_memory=True, + timeout=7.0, + worker_init_fn=_seed_worker, + generator=torch.Generator(), + **workers, + ) + shard = shard_dataloader(loader, 1, 2) + + for setting in ( + "batch_size", + "drop_last", + "collate_fn", + "num_workers", + "pin_memory", + "timeout", + "worker_init_fn", + "multiprocessing_context", + "generator", + "prefetch_factor", + "persistent_workers", + ): + assert getattr(shard, setting) == getattr(loader, setting), setting + + +# -------------------------------------------------------------------------- _reinstall_a2a_forwards + +_TP_MODULE = "transformers.distributed.tensor_parallel" +_LEGACY_TP_MODULE = "transformers.integrations.tensor_parallel" + + +class _FakeMesh: + """The EP device mesh; the TP wrappers only need its process group, the loop its size.""" + + def __init__(self, ranks=2): + self.group = object() + self.ranks = ranks + + def get_group(self): + return self.group + + def size(self): + return self.ranks + + +class _WrapperStyle: + """A transformers >= 5.16 TP style: ``install_forward`` wraps the module's current forward.""" + + def __init__(self, inject_group=False): + self.inject_group = inject_group + self.installed = [] + + def install_forward(self, module, mesh): + self.installed.append(module) + inner = module.forward + extra = (mesh.get_group(),) if self.inject_group else () + + def forward(*args, **kwargs): + return inner(*args, *extra, **kwargs) + + module.forward = forward + + +class _PreHookStyle: + """A transformers < 5.16 TP style: its transforms are forward pre-hooks, no ``install_forward``.""" + + +def _get_parameter_tp_plan(parameter_name, tp_plan, is_weight=True): + """transformers' plan lookup: layer indices match ``*``; a weight also takes its module's entry.""" + name = re.sub(r"\.\d+(\.|$)", lambda m: ".*" + m.group(1), parameter_name) + if name in tp_plan: + return tp_plan[name] + if is_weight and "." in name and (module_name := name.rsplit(".", 1)[0]) in tp_plan: + return tp_plan[module_name] + return None + + +def _install_fake_tp(monkeypatch, styles=None, where=_TP_MODULE): + """Serve ``styles`` as transformers' ``ALL_PARALLEL_STYLES`` from the module ``where``.""" + if styles is None: + styles = { + "torch_a2a_experts": _WrapperStyle(inject_group=True), + "torch_a2a_router": _WrapperStyle(), + "colwise": _WrapperStyle(), + } + module = types.ModuleType(where) + module.ALL_PARALLEL_STYLES = styles + module._get_parameter_tp_plan = _get_parameter_tp_plan + if where != _TP_MODULE: + monkeypatch.setitem(sys.modules, _TP_MODULE, None) # importing it now raises ImportError + monkeypatch.setitem(sys.modules, where, module) + return styles + + +class _Experts(nn.Module): + """A torch_a2a experts module: its forward needs the EP group its TP wrapper injects.""" + + def __init__(self): + super().__init__() + self.weight_quantizer = nn.Identity() # quantizer insertion adds children like this + self.groups = [] + + def forward(self, hidden_states, top_k_index, top_k_weights, process_group=None): + if process_group is None: + raise ValueError("torch_a2a experts forward requires the EP `process_group`") + self.groups.append(process_group) + return hidden_states + + +class _Block(nn.Module): + def __init__(self): + super().__init__() + self.gate = nn.Identity() # the router + self.proj = nn.Identity() # a dense layer under a non-a2a TP style + self.experts = _Experts() + + def forward(self, hidden_states): + return self.experts(self.proj(self.gate(hidden_states)), None, None) + + +class _MoEModel(nn.Module): + """Two MoE layers whose plan sends the experts and routers through torch_a2a.""" + + def __init__(self): + super().__init__() + self.layers = nn.ModuleList([_Block(), _Block()]) + self.config = types.SimpleNamespace(use_cache=True) + self.tp_plan = { + "layers.*.gate": "torch_a2a_router", + "layers.*.proj": "colwise", + "layers.*.experts": "torch_a2a_experts", + } + + def forward(self, input_ids): + hidden_states = input_ids.float() + for layer in self.layers: + hidden_states = layer(hidden_states) + return hidden_states + + +def _load_tp_wrappers(model, styles, mesh): + """Do what transformers' expert-parallel load does: wrap every planned module's forward.""" + for name, module in model.named_modules(): + if (style := _get_parameter_tp_plan(name, model.tp_plan, is_weight=False)) is not None: + styles[style].install_forward(module, mesh) + for style in styles.values(): + style.installed.clear() + + +def _drop_tp_wrapper(module): + """Do what quantizer insertion does to the instance: rebind the class forward over the wrapper. + + (``DynamicModule.convert`` -> ``bind_forward_method``.) + """ + module.forward = types.MethodType(type(module).forward, module) + + +def test_reinstall_is_a_noop_without_an_ep_mesh(monkeypatch): + """A single-process or non-EP model is left untouched, even with the torch_a2a styles around.""" + styles = _install_fake_tp(monkeypatch) + + assert _reinstall_a2a_forwards(_MoEModel(), None) == 0 + assert not any(style.installed for style in styles.values()) + + +@pytest.mark.parametrize( + ("experts_style", "planted"), + [(_PreHookStyle(), 0), (_WrapperStyle(inject_group=True), 2)], + ids=["all-pre-hooks", "only-the-experts-wrap"], +) +def test_reinstall_skips_styles_whose_transforms_are_pre_hooks(monkeypatch, experts_style, planted): + """transformers < 5.16 has no ``install_forward``; its pre-hooks survive quantizer insertion. + + Each style is checked for it, the one about to be called: a style without one is skipped even + where the other has one. Those releases keep the styles in ``transformers.integrations``: the + fallback import. + """ + _install_fake_tp( + monkeypatch, + styles={"torch_a2a_experts": experts_style, "torch_a2a_router": _PreHookStyle()}, + where=_LEGACY_TP_MODULE, + ) + model = _MoEModel() + for layer in model.layers: # quantizer insertion rebinds the planned modules' forwards + _drop_tp_wrapper(layer.gate) + _drop_tp_wrapper(layer.experts) + assert _reinstall_a2a_forwards(model, _FakeMesh()) == planted + + +def test_reinstall_leaves_a_module_the_load_did_not_wrap(monkeypatch): + """Only a forward quantizer insertion rebound is planted again: a planned module whose forward + the load did not wrap keeps running as it did, without the style's transforms.""" + styles = _install_fake_tp(monkeypatch) + + assert _reinstall_a2a_forwards(_MoEModel(), _FakeMesh()) == 0 + assert not any(style.installed for style in styles.values()) + + +# ------------------------------------------------------------------------- distributed_forward_loop + + +class _Recorder(nn.Module): + """Records each forward's kwargs and the ``use_cache`` flag it ran under.""" + + def __init__(self, on_forward=None): + super().__init__() + self.config = types.SimpleNamespace(use_cache=True) + self.calls = [] + self.use_cache_seen = [] + self.on_forward = on_forward + + def forward(self, **batch): + self.calls.append(batch) + self.use_cache_seen.append(self.config.use_cache) + if self.on_forward is not None: + self.on_forward() + + +def test_forward_loop_runs_this_ranks_shard_through_the_captured_model(): + """Every batch of the rank's shard reaches the top-level model, not the module ``mtq`` passes. + + Only the full model's forward runs the EP collectives; without a device the batches stay put. + """ + model, inner = _Recorder(), _Recorder() + shard = shard_dataloader(_loader(list(range(8))), 0, world_size=2) + + distributed_forward_loop(model, shard)(inner) + + calls = [call["input_ids"].flatten().tolist() for call in model.calls] + assert calls == [[0, 2], [4, 6]] + assert inner.calls == [] + + +def test_forward_loop_moves_tensors_to_the_device_it_is_given(): + """Batch tensors go to ``device`` when one is given; other values pass through. + + The meta device makes the move observable on a CPU-only host. (Without a device the batches stay + put: the shard-routing test reads their values back.) + """ + batches = [{"input_ids": torch.zeros(2, 3), "logits_to_keep": 1}] * 2 + model = _Recorder() + distributed_forward_loop(model, batches, device="meta")(model) + + assert [call["input_ids"].device.type for call in model.calls] == ["meta"] * 2 + assert [call["logits_to_keep"] for call in model.calls] == [1, 1] + + +def test_forward_loop_disables_use_cache_only_while_calibrating(): + """KV/state caching is off for the calibration forwards (it breaks hybrid Mamba models).""" + model = _Recorder() + distributed_forward_loop(model, [{"input_ids": torch.zeros(1, 1)}] * 2)(model) + + assert model.use_cache_seen == [False, False] + assert model.config.use_cache is True + + +@pytest.mark.parametrize("cuda_started", [True, False], ids=["cuda", "cpu-only"]) +def test_forward_loop_frees_the_allocator_cache_around_the_forwards(monkeypatch, cuda_started): + """Re-plant the a2a wrappers and free the cached pool before the first forward, and again after. + + Building the loop does neither: quantizers are inserted between building and calling it. NCCL + allocates its buffers outside torch's caching allocator, so a pool hoarded by the model load or + by the calibration activations starves the next collective. A process that never started CUDA + has no pool to free and skips the garbage collection, which alone takes about a second. + """ + events = [] + monkeypatch.setattr( + mto_dist, "_reinstall_a2a_forwards", lambda *args: events.append(("reinstall", *args)) + ) + monkeypatch.setattr(torch.cuda, "is_initialized", lambda: cuda_started) + monkeypatch.setattr(mto_dist.gc, "collect", lambda: events.append("collect")) + monkeypatch.setattr(torch.cuda, "empty_cache", lambda: events.append("empty_cache")) + model, mesh = _Recorder(on_forward=lambda: events.append("forward")), _FakeMesh() + set_distribution_plan(model, _EXPERT_PARALLEL) + forward_loop = distributed_forward_loop( + model, [{"input_ids": torch.zeros(1, 1)}] * 2, ep_mesh=mesh + ) + assert events == [] + + forward_loop(model) + + release = ["collect", "empty_cache"] if cuda_started else [] + assert events == [("reinstall", model, mesh), *release, "forward", "forward", *release] + + +def test_forward_loop_replants_a2a_wrappers_dropped_after_it_was_built(monkeypatch): + """hf_ptq's order -- load, build the loop, insert quantizers, calibrate -- hands the experts the + EP group. + + It holds on every pass (multi-pass algorithms call the loop again), with one wrapper per module, + and again once another conversion drops the wrappers anew. Only the wrappers conversion dropped + are planted again: a router quantization left alone keeps the one from the load, since wrapping + it again would run its transforms twice. Other TP styles keep their forward, and so do the + quantizers under an experts module, since the plan is looked up per module (``is_weight=False``) + and children do not inherit their parent's style. + """ + styles = _install_fake_tp(monkeypatch) + model, mesh = _MoEModel(), _FakeMesh() + _load_tp_wrappers(model, styles, mesh) + set_distribution_plan(model, _EXPERT_PARALLEL) + forward_loop = distributed_forward_loop( + model, [{"input_ids": torch.ones(2, 4)}] * 3, ep_mesh=mesh + ) + for layer in model.layers: + _drop_tp_wrapper(layer.experts) + + forward_loop(model) + forward_loop(model) + for layer in model.layers: + _drop_tp_wrapper(layer.experts) + forward_loop(model) + + assert [layer.experts.groups for layer in model.layers] == [[mesh.group] * 9] * 2 + assert styles["torch_a2a_experts"].installed == [layer.experts for layer in model.layers] * 2 + assert styles["torch_a2a_router"].installed == [] + assert styles["colwise"].installed == [] + + +@pytest.mark.parametrize( + ("plan", "mesh", "match"), + [ + (_EXPERT_PARALLEL, None, "calibrates with the ep_mesh"), + (DistributionPlan(world_size=2), _FakeMesh(), "calibrates with the ep_mesh"), + (DistributionPlan(world_size=4, ep_size=2), _FakeMesh(ranks=4), "spans 4 ranks"), + ], + ids=["expert-parallel-without-mesh", "mesh-for-a-replicated-model", "mesh-of-the-whole-world"], +) +def test_the_ep_mesh_is_the_plans_expert_parallel_one(plan, mesh, match): + """Without its mesh, the experts' forward would run with no EP group at the first step; a mesh + for a model whose plan shards no experts would be ignored; and a mesh over other ranks than an + expert-parallel group would hand the all-to-all the wrong group.""" + model = _Recorder() + set_distribution_plan(model, plan) + + with pytest.raises(ValueError, match=match): + distributed_forward_loop(model, [{"input_ids": torch.zeros(1, 1)}], ep_mesh=mesh) + + +def test_forward_loop_refuses_a_batch_that_is_not_a_mapping(): + """The loop runs ``model(**batch)``; a loader of bare tensors is refused as such at its first + batch, rather than failing on ``batch.items()`` inside the loop.""" + model = _Recorder() + + with pytest.raises(TypeError, match="yields Tensor"): + distributed_forward_loop(model, [torch.zeros(1, 1)], device="meta")(model) + + +# ------------------------------------------------------------------------ rank-local calibration + + +@pytest.mark.parametrize( + ("entry", "algorithm", "refused"), + [ + ("quantize", {"method": "awq_lite"}, "awq_lite"), + ("calibrate", "awq_lite", "awq_lite"), + ( + "calibrate", + {"method": "nvfp4_act_headroom", "weight_scale_algorithm": {"method": "local_hessian"}}, + "local_hessian", + ), + ( + "calibrate", + {"method": "max", "distributed_sync": "false"}, + "max with distributed_sync=False", + ), + ( + "calibrate", + { + "method": "nvfp4_act_headroom", + "weight_scale_algorithm": {"method": "mse", "distributed_sync": None}, + }, + "mse with distributed_sync=None", + ), + ("auto_quantize", {"method": "awq_lite"}, "awq_lite"), + ], + ids=[ + "quantize", + "calibrate", + "for-the-weight-scales", + "max-without-its-sync", + "for-the-weight-scales-without-their-sync", + "auto_quantize", + ], +) +def test_rank_local_calibration_is_refused_before_it_changes_a_distributed_model( + entry, algorithm, refused +): + """Each rank calibrates on its own data, and AWQ keeps what its search gathers on the rank: the + ranks' copies of a weight would come out different. So would local_hessian, run by another + algorithm for its weight scales, and max or MSE told not to sync, by False in any spelling its + schema reads as False, or by None, which skips the sync too. The refusal comes before any + change.""" + model = nn.Sequential(nn.Linear(16, 16)) + set_distribution_plan(model, DistributionPlan(world_size=2)) + if entry == "calibrate": + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(torch.randn(2, 16))) + modules = [type(module) for module in model.modules()] + state = {name: tensor.clone() for name, tensor in model.state_dict().items()} + + with pytest.raises(NotImplementedError, match=rf"{refused} is not supported .* over 2 ranks"): + if entry == "quantize": + config = {**mtq.INT4_AWQ_CFG, "algorithm": algorithm} + mtq.quantize(model, config, lambda m: m(torch.randn(2, 16))) + elif entry == "calibrate": + mtq.calibrate(model, algorithm, lambda m: m(torch.randn(2, 16))) + else: + mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[({**mtq.INT4_AWQ_CFG, "algorithm": algorithm}, "awq")], + data_loader=[torch.randn(2, 16)], + forward_step=lambda m, batch: m(batch), + loss_func=lambda output, batch: output.sum(), + ) + + assert [type(module) for module in model.modules()] == modules + after = model.state_dict() + assert after.keys() == state.keys() + assert all(torch.equal(after[name], tensor) for name, tensor in state.items()) + + +# {module: (data_parallel_group, expert_model_parallel_group)} each layout seeds. With the experts +# sharded over every rank, their expert-parallel group is every rank (None), and no other rank holds +# the same experts (-1: no data-parallel group). +_EXPERT_LAYOUTS = { + "seeded": {"experts": (-1, None)}, + "without-the-group": {}, + "on-the-router-only": {"gate": (-1, None)}, + "data-parallel-over-every-rank": {"experts": (None, None)}, +} + + +def _expert_parallel_calibration_job(rank, size, layout): + # A router and a routed expert, named as expert-parallel plans name the experts' container. + model = nn.Sequential(OrderedDict(gate=nn.Linear(16, 16), experts=nn.Linear(16, 16))) + for name, (dp, ep) in _EXPERT_LAYOUTS[layout].items(): + model.get_submodule(name)._parallel_state = ParallelState( + data_parallel_group=dp, expert_model_parallel_group=ep + ) + set_distribution_plan(model, DistributionPlan(world_size=size, ep_size=size)) + raised = None + with bounded_collectives(): + # Not pytest.raises: when nothing is raised it fails with a BaseException, which would take + # the shared worker down instead of failing this test alone. + try: + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda m: m(torch.randn(2, 16))) + except NotImplementedError as e: + raised = str(e) + if layout == "seeded": + assert raised is None, raised + return + assert raised is not None and "expert-parallel group" in raised, raised + # Refused before converting anything, but for unseeded experts beside a seeded router: those + # are told apart once quantization has converted them, and refused before they calibrate. + assert (type(model.experts) is nn.Linear) == (layout != "on-the-router-only") + + +@pytest.mark.parametrize("layout", list(_EXPERT_LAYOUTS)) +def test_expert_parallel_calibration_needs_the_experts_expert_parallel_group( + cpu_dist_workers, layout +): + """Each rank's per-expert quantizers hold its own experts, and max syncs them over the groups + their module's ParallelState names, every rank unless set, which would merge different experts' + scales. Calibration is refused until the experts carry their groups: with no expert-parallel + group at all, with one on the router alone, and with a data-parallel group spanning ranks that + hold other experts.""" + cpu_dist_workers.run(_expert_parallel_calibration_job, layout=layout) + + +def _expert_parallel_auto_quantize_job(rank, size): + model = nn.Sequential(OrderedDict(gate=nn.Linear(16, 16), experts=nn.Linear(16, 16))) + dp, ep = _EXPERT_LAYOUTS["seeded"]["experts"] + model.experts._parallel_state = ParallelState( + data_parallel_group=dp, expert_model_parallel_group=ep + ) + set_distribution_plan(model, DistributionPlan(world_size=size, ep_size=size)) + modules = [type(module) for module in model.modules()] + state = {name: tensor.clone() for name, tensor in model.state_dict().items()} + raised = None + with bounded_collectives(): + try: # not pytest.raises, as in _expert_parallel_calibration_job + mtq.auto_quantize( + model, + constraints={"effective_bits": 8.0}, + quantization_formats=[(mtq.INT8_DEFAULT_CFG, "int8")], + data_loader=[torch.randn(2, 16)], + forward_step=lambda m, batch: m(batch), + loss_func=lambda output, batch: output.sum(), + ) + except NotImplementedError as e: + raised = str(e) + assert raised is not None and f"experts are sharded over {size} ranks" in raised, raised + assert [type(module) for module in model.modules()] == modules + after = model.state_dict() + assert after.keys() == state.keys() + assert all(torch.equal(after[name], tensor) for name, tensor in state.items()) + + +def test_auto_quantize_refuses_an_expert_parallel_model_before_changing_it(cpu_dist_workers): + """auto_quantize's forwards would run the experts without the expert-parallel forwards that + quantization replaces, which distributed_forward_loop restores. So it refuses such a model, with + its experts seeded and a candidate every rank syncs, before it converts anything.""" + cpu_dist_workers.run(_expert_parallel_auto_quantize_job) + + +# Rank 0 of four at EP 2, in the (edp, ep) = (2, 2) mesh: ranks 0 and 1 shard the experts between +# them, ranks 2 and 3 do the same, and rank 2 holds the same shard as rank 0. The two-rank pool cannot +# build these groups, so the process group answers as a four-rank launch's would. +_FOUR_RANK_GROUPS = { + "0,1": [0, 1], + "1,0": [1, 0], + "0,2": [0, 2], + "0,3": [0, 3], + "all": [0, 1, 2, 3], +} + + +@pytest.mark.parametrize( + ("ep", "dp", "seeded"), + [ + ("0,1", "0,2", True), + ("0,2", "0,1", False), + ("0,1", "0,1", False), + ("0,1", "0,3", False), + ("1,0", "0,2", False), + ("all", "0,2", False), + ], + ids=[ + "seeded", + "groups-transposed", + "data-parallel-over-the-expert-parallel-ranks", + "data-parallel-over-another-shard", + "expert-parallel-in-another-order", + "expert-parallel-over-every-rank", + ], +) +def test_the_experts_groups_are_checked_by_their_ranks(monkeypatch, ep, dp, seeded): + """A data-parallel group of the plan's edp ranks is not enough: over the experts' + expert-parallel ranks, or ranks that meet them only at this one but hold another shard (rank 3 + holds rank 1's), max's and GPTQ's syncs would merge different experts' scales. Nor is an + expert-parallel group over more ranks than the plan shards over, or in another order, which + gives this rank another shard, nor the two groups transposed, the mesh's column as the + expert-parallel group and its row as the data-parallel one, which match in size when edp is + ep_size.""" + monkeypatch.setattr(mto_dist, "backend", lambda: "torch") + monkeypatch.setattr(mto_dist, "rank", lambda group=None: 0) + monkeypatch.setattr( + torch.distributed, + "get_world_size", + lambda group=None: 4 if group is None else len(_FOUR_RANK_GROUPS[group]), + ) + monkeypatch.setattr( + torch.distributed, "get_process_group_ranks", lambda group: _FOUR_RANK_GROUPS[group] + ) + model = nn.Sequential(OrderedDict(experts=nn.Linear(16, 16))) + model.experts._parallel_state = ParallelState( + data_parallel_group=dp, expert_model_parallel_group=ep + ) + set_distribution_plan(model, DistributionPlan(world_size=4, ep_size=2)) + + if seeded: + _refuse_unsupported_distributed_calibration(model, "max") + else: + with pytest.raises(NotImplementedError, match=r"row of the \(2, 2\) expert-parallel mesh"): + _refuse_unsupported_distributed_calibration(model, "max") + + +@pytest.mark.parametrize( + "config", + [ + {"quant_cfg": [{"quantizer_name": "*", "enable": False}], "algorithm": "max"}, + mtq.INT8_DEFAULT_CFG, + ], + ids=["all-disabled", "quantized"], +) +def test_a_module_without_routed_experts_calibrates_under_an_expert_parallel_plan(config): + """hf_ptq quantizes a VLM's vision tower apart from its language model, every quantizer + disabled. It holds no routed experts, so it has no experts' scales to merge, quantized or not.""" + model = nn.Sequential(nn.Linear(16, 16)) + set_distribution_plan(model, _EXPERT_PARALLEL) + + mtq.quantize(model, config, lambda m: m(torch.randn(2, 16)))