diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 4f5b89098d4..44bca74e533 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -12,7 +12,7 @@ Changelog - On Linux x86-64 with Python 3.11 and later, the ONNX extra now selects ONNX Runtime 1.30.x and CUDA 13 packages by default. GPU calibration requires a compatible CUDA 13 runtime and driver; Python 3.10 and other platforms retain their existing dependency selections. - Add imatrix calibration (``algorithm: {method: imatrix}``) for the GGML formats with a weighted search: each weight is encoded once, its search weighted by the square root of its inputs' per-column ``E[x^2]`` (llama.cpp's imatrix statistic), and HF fused-MoE experts get one statistic per expert. A weight that already holds final codes is left alone, so ``[{method: gptq}, {method: imatrix}]`` runs GPTQ where it applies and imatrix on the rest; the built-in IQ1 and IQ2 recipes now do this, and new ``_experts_only`` recipes apply imatrix to the routed MoE experts alone. - GPTQ now also quantizes HuggingFace fused MoE experts, each with its own Hessian from the tokens routed to it; they were previously skipped. The built-in IQ1 and IQ2 recipes therefore run GPTQ on the experts too, and large MoEs need layerwise calibration to hold one layer's expert Hessians at a time. -- GPTQ no longer runs a max-calibration forward when no activation quantizer needs one, and ``hf_ptq.py`` now sizes the default layerwise calibration batch at 8,192 tokens (16 samples at ``--calib_seq 512``) instead of 1. Layerwise GPTQ on HF fused-MoE models runs over 10x faster; pass ``--batch_size 1`` for the old batching, which pads no sample. +- GPTQ no longer runs a max-calibration forward when no activation quantizer needs one, and ``hf_ptq.py`` now sizes the default layerwise calibration batch to what one decoder layer's forward fits in GPU memory, up to 32,768 tokens (64 samples at ``--calib_seq 512``), instead of 1. Layerwise GPTQ on HF fused-MoE models runs over 10x faster; pass ``--batch_size 1`` for the old batching. - Add a checkpoint-specific NVIDIA Nemotron 3.5 Super VL AutoQuantize recipe with a 4.9-effective-bit NVFP4-MSE/FP8/BF16 search and fixed NVFP4 MTP experts. Select it under ``models/nvidia/NVIDIA-Nemotron-3.5-Super-VL-120B-A12B-BF16/auto_quantize/``. - Add ``examples/hf_ptq/convert_ggml_to_int8.py`` to expand a checkpoint exported with a GGML recipe (IQ formats, ``q8_0``) into an INT8 weight-only compressed-tensors checkpoint. Stock vLLM serves it with its integer kernels at about half the BF16 weight memory, reproducing the GGML weights up to rounding their scales to the model dtype. @@ -60,6 +60,8 @@ Changelog - Fix KV-cache quantization through ``register_attention_for_kv_quant``, the fallback for Hugging Face attention classes that do not use the attention interface, attaching the ``q``, ``k`` and ``v`` bmm quantizers to the wrong tensors for most attention layouts. Recalibrate and re-export models quantized through this fallback. Attention classes that compute their scores with ``torch.baddbmm`` (e.g. ``BloomAttention``, ``GPTNeoXJapaneseAttention``) are no longer patched, since their operands were always wired wrong (or raised ``IndexError``), and now get the existing warning to register a custom quantized attention class. - Fix GPTQ with ``fused: true`` failing with a CUDA illegal memory access, or a Triton pointer error, on a model that ``device_map="auto"`` spreads over several GPUs once calibration reaches a layer outside the current GPU. - Nested ``only_student_forward()`` and ``only_teacher_forward()`` contexts now preserve the enclosing execution mode, including when an inner context is disabled or raises an exception. +- Fix calibrating a CPU- or disk-offloaded model running out of GPU memory as it progresses: the quantizer state calibration adds to an offloaded layer now leaves the GPU with the layer's weights instead of accumulating there. +- Fix layerwise calibration replaying the wrong cached activations, or failing with ``is in 'run' mode but has no cached inputs to replay``, after a GPU out-of-memory error. It now stops with an error asking for a smaller calibration batch size. - Fix ONNX preprocessing with custom TensorRT plugin libraries on TensorRT 11. - Fix ONNX remote AutoQDQ safety benchmarks returning infinite latency by running the generated engine with ``trtexec_safe`` on the configured target. - Fix ``INT8_DEFAULT_CFG``, ``FP8_DEFAULT_CFG``, ``NVFP4_DEFAULT_CFG`` and the other presets and recipes that import ``configs/ptq/units/default_disabled_quantizers`` and enable ``*input_quantizer``: since 0.44 they also quantized ``nn.LayerNorm`` inputs, which adds quantization error to the residual stream of models built on ``nn.LayerNorm`` such as BERT-style encoders. Re-quantize affected models; a recipe that wants ``nn.LayerNorm`` quantizers must enable them after importing that unit. diff --git a/docs/source/announcements/20260928-single-gpu-ptq.rst b/docs/source/announcements/20260928-single-gpu-ptq.rst index b5e4e7fbac2..58fcf3286e6 100644 --- a/docs/source/announcements/20260928-single-gpu-ptq.rst +++ b/docs/source/announcements/20260928-single-gpu-ptq.rst @@ -133,14 +133,15 @@ at sequence length 512, **batch size 8**. .. note:: - **Pick the batch size on purpose.** For layerwise runs, ``hf_ptq.py`` does not probe memory: - left at its default, ``--batch_size`` becomes 8,192 tokens' worth of samples (16 at sequence - length 512). It was 1 when this run was made, which is why it passed ``--batch_size 8``. An - explicit value passes straight through; set a smaller one if a layer's activations do not - fit your GPU. Going from 1 to 8 cut wall clock roughly in half. It does move the + **Pick the batch size on purpose.** For layerwise runs, ``hf_ptq.py`` left at its default + probes one decoder layer of each kind and takes the most samples whose forward fits in half + the GPU's free memory, up to 32,768 tokens' worth (64 at sequence length 512). It was 1 when + this run was made, which is why it passed ``--batch_size 8``. An explicit value passes + straight through. Going from 1 to 8 cut wall clock roughly in half. It does move the numbers a little. Weights, ``weight_scale`` and ``weight_scale_2`` come out bit-identical, - but activation ``input_scale`` values shift slightly, because batching pads each sequence to - the longest one in the batch and those pad positions feed the activation amax. + but activation ``input_scale`` values shift slightly, because a different batch shape + rounds the activations differently. Padding does not change: ``hf_ptq.py`` pads every + calibration sample to the same length whatever the batch size. Then launch the run: diff --git a/examples/hf_ptq/hf_ptq.py b/examples/hf_ptq/hf_ptq.py index 27e8fb3e8d5..3f41e64b83a 100755 --- a/examples/hf_ptq/hf_ptq.py +++ b/examples/hf_ptq/hf_ptq.py @@ -89,6 +89,7 @@ from modelopt.torch.quantization.config import need_calibration from modelopt.torch.quantization.plugins.accelerate import init_quantized_weights from modelopt.torch.quantization.utils import is_quantized +from modelopt.torch.quantization.utils.layerwise_calib import get_max_layerwise_batch_size from modelopt.torch.speculative.eagle.utils import ( EagleOfflineDataCollator, OfflineSupervisedDataset, @@ -114,12 +115,12 @@ from modelopt.torch.utils.vlm_dataset_utils import get_vlm_dataset_dataloader RAND_SEED = 1234 -# Tokens per calibration batch when layerwise calibration picks the batch size (16 x 512). -LAYERWISE_CALIB_TOKENS = 8192 +# Most tokens per calibration batch when layerwise calibration picks the batch size (64 x 512). +LAYERWISE_CALIB_TOKENS = 32768 def _layerwise_calib_batch_size(calib_seq: int, num_samples: int) -> int: - """Default layerwise batch size: LAYERWISE_CALIB_TOKENS tokens, but 1 to all of the samples.""" + """Largest default layerwise batch: LAYERWISE_CALIB_TOKENS tokens, but 1 to all of the samples.""" return max(1, min(LAYERWISE_CALIB_TOKENS // calib_seq, num_samples)) @@ -998,13 +999,17 @@ def quantize_main( args.batch_size = 1 # Layerwise calibration processes one layer at a time; auto batch-size probing runs a # full-model forward which defeats the point and can OOM on very large models. A batch - # then only holds one layer's activations, so size it by tokens: batch 1 makes an MoE - # layer run its per-expert loop once per sample. + # then only holds one layer's activations, so probe one decoder layer of each kind, up + # to a token budget: batch 1 makes an MoE layer run its per-expert loop once per sample. elif is_layerwise: - args.batch_size = _layerwise_calib_batch_size(args.calib_seq, sum(args.calib_size)) + max_batch_size = _layerwise_calib_batch_size(args.calib_seq, sum(args.calib_size)) + args.batch_size = get_max_layerwise_batch_size( + language_model, args.calib_seq, max_batch_size, sum(args.calib_size), device=device + ) print( - f"Layerwise calibration enabled. Using default batch_size={args.batch_size} " - f"({LAYERWISE_CALIB_TOKENS} tokens per batch) for calibration." + f"Layerwise calibration enabled. Using batch_size={args.batch_size}, the most " + f"that fits a decoder layer's forward, up to {max_batch_size} " + f"({LAYERWISE_CALIB_TOKENS} tokens)." ) else: # Calibration/sparsification will actually take much more memory than regular inference diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py index 47b188ab9fc..fa36b566193 100644 --- a/modelopt/torch/quantization/mode.py +++ b/modelopt/torch/quantization/mode.py @@ -73,6 +73,7 @@ smoothquant, svdquant, ) +from .utils.calib_utils import reset_calib_forward_memory __all__ = ["BaseCalibrateModeDescriptor"] @@ -250,6 +251,7 @@ def wrapped_calib_func( module._moe_calib_experts_ratio = moe_calib_experts_ratio if func is not None: + reset_calib_forward_memory() if layerwise: # All currently implemented PTQ algorithms support layerwise calibration; # future algorithms that need full-model context must add a guard here. diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index f02732a9334..b00c10a8ba5 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -19,6 +19,7 @@ import math import time import warnings +from collections import Counter from collections.abc import Callable, Mapping, Sequence from functools import partial from typing import Any, TypeAlias @@ -69,7 +70,12 @@ persistent_materialization, promote_static_block_weight_quantizers, ) -from .utils.calib_utils import _GPTQ_HELPER_REGISTRY, _IMATRIX_ENCODER_REGISTRY, GPTQHelper +from .utils.calib_utils import ( + _GPTQ_HELPER_REGISTRY, + _IMATRIX_ENCODER_REGISTRY, + GPTQHelper, + track_calib_forward_memory, +) __all__ = [ "CalibratorFactory", @@ -2480,14 +2486,21 @@ def _expert_hook(_input_quantizer, _args, output): input_quantizer.register_forward_hook(_make_expert_hook(module, per_expert)) ) + placement = Counter( + str(h.hessian.device) for h in [*gptq_handles.values(), *expert_handles.values()] + ) print_rank_0( f"Computing Hessians for {len(gptq_handles)} linear layers and " - f"{len(expert_handles)} fused-expert weights ({expert_hessian_bytes / 2**30:.1f} GiB)..." + f"{len(expert_handles)} fused-expert weights ({expert_hessian_bytes / 2**30:.1f} GiB, " + f"on {', '.join(f'{d} x{n}' for d, n in sorted(placement.items()))})..." ) try: - with set_quantizer_by_cfg_context( - model, [{"quantizer_name": "*weight_quantizer", "enable": False}] + with ( + set_quantizer_by_cfg_context( + model, [{"quantizer_name": "*weight_quantizer", "enable": False}] + ), + track_calib_forward_memory(), ): forward_loop(model) finally: diff --git a/modelopt/torch/quantization/plugins/accelerate.py b/modelopt/torch/quantization/plugins/accelerate.py index eb206db2c5c..fc2c37b9d94 100644 --- a/modelopt/torch/quantization/plugins/accelerate.py +++ b/modelopt/torch/quantization/plugins/accelerate.py @@ -28,6 +28,8 @@ import modelopt.torch.quantization as mtq +from ..nn import TensorQuantizer + __all__ = ["init_quantized_weights"] @@ -79,6 +81,85 @@ def _writeback_params_to_weights_map(module, align_hook): w_map[key] = tensor.detach().cpu() +def _own_quantizers(module): + """Quantizers whose state belongs with ``module``'s weights, not with a hooked descendant's.""" + for child in module.children(): + if isinstance(child, TensorQuantizer): + yield child + elif not hasattr(child, "_hf_hook"): + yield from _own_quantizers(child) + + +# A quantizer's calibrators: its amax calibrator, and its bias calibrator under static bias. +_CALIBRATORS = ("_calibrator", "_bias_calibrator") + + +def _quantizer_tensors(module): + """``(key, tensor)`` for the state of ``module``'s own quantizers and their calibrators.""" + for quantizer in _own_quantizers(module): + sources = {"_parameters": quantizer._parameters, "_buffers": quantizer._buffers} + for source in _CALIBRATORS: + calibrator = getattr(quantizer, source, None) + sources[source] = vars(calibrator) if calibrator is not None else {} + for source, tensors in sources.items(): + for name, tensor in tensors.items(): + if isinstance(tensor, torch.Tensor): + yield (id(quantizer), source, name), tensor + + +def _carry_quantizers_with_offload(hook): + """Make an offload ``hook`` move its module's quantizer state on and off the GPU with the weights. + + Accelerate only manages the tensors a module had when it was dispatched, so calibration state + such as NVFP4's per-block ``_amax``, and its calibrator's copy, stayed on the execution device + for every offloaded layer and added up over the model. The CPU tensors stay the master copy: + a forward gets GPU copies, and after it the CPU tensors come back. A quantizer's own tensors + are always copied back, since its setters write through ``.data``, which leaves ``_version`` + unchanged; a calibrator's tensor is copied back only if the forward replaced or modified it. + The packed-weight cache is keyed to the weight just released, so it is dropped. + """ + if getattr(hook, "_modelopt_carries_quantizers", False): + return + pre_forward, post_forward = hook.pre_forward, hook.post_forward + onloaded: dict[tuple, tuple[torch.Tensor, torch.Tensor, int]] = {} + + def carrying_pre_forward(module, *args, **kwargs): + result = pre_forward(module, *args, **kwargs) + if hook.offload: + for key, tensor in _quantizer_tensors(module): + if tensor.device.type == "cpu": + master = tensor.data + tensor.data = master.to(hook.execution_device) + onloaded[key] = (tensor, master, tensor._version) + return result + + def carrying_post_forward(module, output): + output = post_forward(module, output) + if hook.offload: + for key, tensor in _quantizer_tensors(module): + # meta: a hook that places submodules offloaded it already. + if tensor.device.type in ("cpu", "meta"): + continue + entry = onloaded.pop(key, None) + if entry is None or entry[0] is not tensor: + tensor.data = tensor.data.cpu() + continue + master = entry[1] + if key[1] not in _CALIBRATORS or entry[2] != tensor._version: + if master.shape != tensor.shape or master.dtype != tensor.dtype: + tensor.data = tensor.data.cpu() + continue + # Into the existing master: a new CPU tensor costs 8x as much in page faults. + master.copy_(tensor.data) + tensor.data = master + for quantizer in _own_quantizers(module): + quantizer._quantizer_cache = None + return output + + hook.pre_forward, hook.post_forward = carrying_pre_forward, carrying_post_forward + hook._modelopt_carries_quantizers = True + + @contextmanager def weight_access_and_writeback_context(module, writeback: bool = True): """Context manager for weight access and writeback for modules managed by accelerate. @@ -102,6 +183,7 @@ def weight_access_and_writeback_context(module, writeback: bool = True): # weights would be overwritten with stale CPU state_dict values. needs_materialize = any(p.device.type == "meta" for p in mod.parameters()) if needs_materialize: + _carry_quantizers_with_offload(hook) hook.pre_forward(mod) hook.offload = False materialized.append((mod, hook, needs_materialize)) diff --git a/modelopt/torch/quantization/utils/calib_utils.py b/modelopt/torch/quantization/utils/calib_utils.py index 81f91d883e5..4895a31103b 100644 --- a/modelopt/torch/quantization/utils/calib_utils.py +++ b/modelopt/torch/quantization/utils/calib_utils.py @@ -38,36 +38,71 @@ """GPTQ helper and Hessian utilities for calibration.""" +from contextlib import contextmanager + import torch import torch.distributed as dist -from modelopt.torch.utils import print_rank_0 +from modelopt.torch.utils import get_available_gpu_memory, get_process_gpus, print_rank_0 from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method -from modelopt.torch.utils.perf import get_used_gpu_mem_fraction -# Used fraction of a GPU's memory above which GPTQ keeps new Hessians off it. -_HESSIAN_GPU_MEM_LIMIT = 0.65 +# A GPU takes a new Hessian only if enough memory stays available after it for the calibration +# forward that follows. Until GPTQ has measured that forward on the GPU, keep this share free. +_HESSIAN_UNMEASURED_RESERVE = 0.35 +# Once measured, keep this multiple of the most a calibration forward added on the GPU, and at +# least this share of the GPU. +_HESSIAN_RESERVE_MARGIN = 1.5 +_HESSIAN_MIN_RESERVE = 0.1 +# The most memory a GPTQ calibration forward has added on each GPU, by device index. +_calib_forward_peak: dict[int, int] = {} + + +def _gpu_has_room(gpu: int, nbytes: int) -> bool: + """Whether ``nbytes`` more on ``gpu`` leaves the calibration forward the memory it needs.""" + total = torch.cuda.get_device_properties(gpu).total_memory + peak = _calib_forward_peak.get(gpu) + if peak is None: + reserve = _HESSIAN_UNMEASURED_RESERVE * total + else: + reserve = max(_HESSIAN_MIN_RESERVE * total, _HESSIAN_RESERVE_MARGIN * peak) + return get_available_gpu_memory(gpu) - nbytes >= reserve + +def reset_calib_forward_memory(): + """Forget measured calibration forwards; a new run, maybe of another model, measures afresh.""" + _calib_forward_peak.clear() -def _gpu_is_full(device) -> bool: - return get_used_gpu_mem_fraction(device) > _HESSIAN_GPU_MEM_LIMIT + +@contextmanager +def track_calib_forward_memory(): + """Record the memory the enclosed calibration forward adds on each GPU, for Hessian placement. + + PyTorch's peak counters are reset first: otherwise an earlier peak, such as the previous + layer's Hessians and GPTQ updates, would count as forward memory and send Hessians to the CPU + for nothing. + """ + gpus = get_process_gpus() if torch.cuda.is_available() else [] + allocated = {} + for gpu in gpus: + torch.cuda.reset_peak_memory_stats(gpu) + allocated[gpu] = torch.cuda.memory_allocated(gpu) + yield + for gpu in gpus: + added = torch.cuda.max_memory_allocated(gpu) - allocated[gpu] + _calib_forward_peak[gpu] = max(_calib_forward_peak.get(gpu, 0), added) def _hessian_offload_device(nbytes: int) -> torch.device: - """Where a Hessian goes when its weight's GPU is full: the emptiest GPU with room, else CPU. + """Where a Hessian goes when its weight's GPU lacks room: the emptiest GPU with room, else CPU. A model sharded over several GPUs leaves the ones not holding the current layer idle, and a Hessian there costs a device-to-device copy per update instead of a host copy and a CPU rescale of the whole matrix. Only GPUs this process already holds memory on are candidates, - and only outside distributed runs, so no other job's or rank's GPU is touched. A candidate - must also have the Hessian's ``nbytes`` free. + and only outside distributed runs, so no other job's or rank's GPU is touched. """ if not (dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1): - gpus = [i for i in range(torch.cuda.device_count()) if torch.cuda.memory_reserved(i)] - for gpu in sorted(gpus, key=get_used_gpu_mem_fraction): - if _gpu_is_full(gpu): - break - if torch.cuda.mem_get_info(gpu)[0] > nbytes: + for gpu in sorted(get_process_gpus(), key=get_available_gpu_memory, reverse=True): + if _gpu_has_room(gpu, nbytes): return torch.device("cuda", gpu) return torch.device("cpu") @@ -186,8 +221,10 @@ def __init__( device = weight.device if device.type == "meta": device = torch.device("cpu") - elif offload_to_cpu and device.type == "cuda" and _gpu_is_full(device): - device = _hessian_offload_device(4 * in_features**2) + elif offload_to_cpu and device.type == "cuda": + index = torch.cuda.current_device() if device.index is None else device.index + if not _gpu_has_room(index, 4 * in_features**2): + device = _hessian_offload_device(4 * in_features**2) self.hessian = torch.zeros(in_features, in_features, dtype=torch.float32, device=device) self.n_samples = 0 # Set by update_weights(); listed here for documentation. diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py index 7d9609063b6..82cb4669561 100644 --- a/modelopt/torch/quantization/utils/layerwise_calib.py +++ b/modelopt/torch/quantization/utils/layerwise_calib.py @@ -35,7 +35,13 @@ import torch.nn as nn from modelopt.torch.utils import distributed as dist -from modelopt.torch.utils import print_rank_0, warn_rank_0 +from modelopt.torch.utils import ( + get_available_gpu_memory, + get_process_gpus, + print_rank_0, + warn_rank_0, +) +from modelopt.torch.utils.dataset_utils import _disable_use_cache from modelopt.torch.utils.network import ( bind_forward_method, get_module_device, @@ -361,6 +367,15 @@ def _early_stop_forward(module_self, *args, **kwargs): if depth[0] > 1: raise return None + except torch.cuda.OutOfMemoryError as e: + # The forward loop would split the data batch and retry, but the layer in 'run' + # mode replays the batch it cached at full size, one cached batch per call: a + # retry cannot shrink it and would replay the next batch's inputs instead. + raise RuntimeError( + "Layerwise calibration ran out of GPU memory. Splitting the batch now " + "cannot help, since each decoder layer replays its cached inputs at the " + "size they were cached with; rerun with a smaller calibration batch size." + ) from e except RuntimeError as e: if "meta" not in str(e).lower(): raise @@ -577,6 +592,162 @@ def cache_outputs_for_next_layer_calib( return self.get_input_activations(next_layer, forward_loop) +def _cuda_bytes_per_gpu(obj: Any, out: dict[int, int] | None = None) -> dict[int, int]: + """Bytes of the CUDA tensors in a nested args/kwargs structure, per GPU.""" + out = {} if out is None else out + if isinstance(obj, torch.Tensor): + if obj.is_cuda: + out[obj.device.index] = out.get(obj.device.index, 0) + obj.nbytes + elif isinstance(obj, dict): + for value in obj.values(): + _cuda_bytes_per_gpu(value, out) + elif isinstance(obj, (list, tuple)): + for value in obj: + _cuda_bytes_per_gpu(value, out) + return out + + +def _capture_layer_inputs( + model: nn.Module, layer: nn.Module, batch_size: int, seq_len: int, device: torch.device +): + """Run *model* on a dummy batch until *layer* and return the ``(args, kwargs)`` it receives.""" + input_ids = torch.zeros(batch_size, seq_len, dtype=torch.long, device=device) + # One left-pad position, like padded calibration rows, so the model builds the explicit + # attention mask that calibration batches get instead of a cheaper causal-only path. + attention_mask = torch.ones_like(input_ids) + attention_mask[:, 0] = 0 + info = layer._layerwise_calib + info.mode = "capture" + try: + with _disable_use_cache(model): + model(input_ids, attention_mask=attention_mask) + finally: + info.mode = "original" + (inputs,) = info.collected_inputs + info.collected_inputs = [] + return inputs + + +def _forward_peak_per_gpu(layer: nn.Module, args: tuple, kwargs: dict) -> dict[int, int]: + """Peak memory one forward of *layer* adds on each GPU.""" + base = {} + for gpu in get_process_gpus(): + torch.cuda.reset_peak_memory_stats(gpu) + base[gpu] = torch.cuda.memory_allocated(gpu) + layer(*args, **kwargs) + return { + gpu: torch.cuda.max_memory_allocated(gpu) - base.get(gpu, 0) for gpu in get_process_gpus() + } + + +def _gpu_index(device: torch.device) -> int | None: + """CUDA device index of *device*; ``None`` if it is not a GPU.""" + if device.type != "cuda": + return None + return torch.cuda.current_device() if device.index is None else device.index + + +def _layer_batch_fit( + model: nn.Module, + layer: nn.Module, + gpus: set[int], + seq_len: int, + num_samples: int, + memory_fraction: float, + device: torch.device, +) -> int | None: + """Largest batch whose forward through *layer* fits; ``None`` if no GPU memory grows with it. + + Layers of the same kind on other *gpus* cost what *layer* costs on its own GPU, so those GPUs' + memory limits the batch too. + """ + # The first forward only warms up: it also allocates one-off buffers like cuBLAS's workspace. + for _ in range(2): + args, kwargs = _capture_layer_inputs(model, layer, 1, seq_len, device) + cached_per_sample = _cuda_bytes_per_gpu((args, kwargs)) + one = _forward_peak_per_gpu(layer, args, kwargs) + del args, kwargs + try: + args, kwargs = _capture_layer_inputs(model, layer, 2, seq_len, device) + two = _forward_peak_per_gpu(layer, args, kwargs) + except RuntimeError as e: + # Capturing inputs reports running out of memory as a RuntimeError raised from it. + if not any(isinstance(x, torch.cuda.OutOfMemoryError) for x in (e, e.__cause__)): + raise + two = None + finally: + args = kwargs = None + if two is None: + # One sample fits but two do not: release what the failed forward left cached. + torch.cuda.empty_cache() + return 1 + + costs = {} + for gpu, peak in two.items(): + per_sample = peak - one.get(gpu, 0) + if per_sample > 0: + costs[gpu] = (per_sample, one.get(gpu, 0) - per_sample, cached_per_sample.get(gpu, 0)) + home = _gpu_index(get_module_device(layer)) + if home in costs: + for gpu in gpus: + costs.setdefault(gpu, costs[home]) + fits = [] + for gpu, (per_sample, fixed, cached) in costs.items(): + # Layerwise calibration caches every sample's inputs for this layer and the next. + available = get_available_gpu_memory(gpu) - 2 * num_samples * cached + fits.append(int((memory_fraction * available - fixed) // per_sample)) + return min(fits, default=None) + + +@torch.no_grad() +def get_max_layerwise_batch_size( + model: nn.Module, + seq_len: int, + max_batch_size: int, + num_samples: int, + memory_fraction: float = 0.5, + device: torch.device | None = None, +) -> int: + """Largest calibration batch, up to *max_batch_size*, whose decoder-layer forwards fit. + + Layerwise calibration runs one decoder layer at a time, so a full-model batch probe such as + :func:`get_max_batch_size ` both + overstates the memory a batch needs and risks running out of it. Instead, the first layer + of each kind (by module classes and parameter shapes: dense or MoE, full or linear attention, + and per-layer widths such as Nemotron-NAS's) runs on one and two samples of *seq_len* tokens, + and the growth of its peak memory on each GPU is that GPU's cost per sample; every other GPU + that runs a layer of that kind is charged the same. A batch fits + when the layer's forward stays within *memory_fraction* of the GPU's available memory, after + the inputs layerwise calibration caches for all *num_samples*; the rest is left for + quantizers and Hessians. A layer that fits one sample but not two gives a batch of one. The + dummy token ids go to *device*, by default the model's. All ranks return the same, smallest + batch. + """ + batch_size = max_batch_size + layers = LayerActivationCollector.get_decoder_layers(model) + if layers and torch.cuda.is_available() and max_batch_size > 1: + kinds: dict[tuple, list[nn.Module]] = {} + for layer in layers: + classes = tuple(sorted({type(m).__name__ for m in layer.modules()})) + shapes = tuple(tuple(p.shape) for p in layer.parameters()) + kinds.setdefault((classes, shapes), []).append(layer) + device = device or get_module_device(model) + collector = LayerActivationCollector(model) + collector._patch_all_layers(layers) + try: + for same_kind in kinds.values(): + indexes = (_gpu_index(get_module_device(layer)) for layer in same_kind) + gpus = {gpu for gpu in indexes if gpu is not None} + fit = _layer_batch_fit( + model, same_kind[0], gpus, seq_len, num_samples, memory_fraction, device + ) + if fit is not None: + batch_size = min(batch_size, fit) + finally: + collector._unpatch_all_layers() + return max(1, min(dist.allgather(batch_size))) + + def _move_to_device(obj: Any, device: torch.device) -> Any: """Recursively move tensors to *device*. Non-tensors are returned as-is.""" if isinstance(obj, torch.Tensor): diff --git a/modelopt/torch/utils/perf.py b/modelopt/torch/utils/perf.py index 93bd4839c15..196aba04287 100644 --- a/modelopt/torch/utils/perf.py +++ b/modelopt/torch/utils/perf.py @@ -27,7 +27,9 @@ "AccumulatingTimer", "Timer", "clear_cuda_cache", + "get_available_gpu_memory", "get_cuda_memory_stats", + "get_process_gpus", "get_used_gpu_mem_fraction", "maybe_clear_cuda_cache", "report_memory", @@ -91,6 +93,25 @@ def get_used_gpu_mem_fraction(device="cuda:0"): return (total_memory - free_memory) / total_memory +def get_process_gpus() -> list[int]: + """GPUs this process already holds memory on; querying any other would create a context on it.""" + return [i for i in range(torch.cuda.device_count()) if torch.cuda.memory_reserved(i)] + + +def get_available_gpu_memory(gpu: int) -> int: + """Bytes a new allocation on ``gpu`` can get: the driver's free memory and our cached blocks. + + Free pieces of partly used cached blocks are left out: a large allocation cannot use them, + and PyTorch cannot hand them back to the driver. + """ + free, _ = torch.cuda.mem_get_info(gpu) + stats = torch.cuda.memory_stats(gpu) + cached = stats.get("reserved_bytes.all.current", 0) - stats.get( + "allocated_bytes.all.current", 0 + ) + return free + cached - stats.get("inactive_split_bytes.all.current", 0) + + def report_memory(name="", rank=0, device=None): """Simple GPU memory report.""" memory_stats = get_cuda_memory_stats(device) diff --git a/tests/examples/hf_ptq/test_hf_ptq_args.py b/tests/examples/hf_ptq/test_hf_ptq_args.py index b441dcc3ce7..c8f7e1b667e 100644 --- a/tests/examples/hf_ptq/test_hf_ptq_args.py +++ b/tests/examples/hf_ptq/test_hf_ptq_args.py @@ -436,17 +436,20 @@ def test_mlflow_is_off_by_default(monkeypatch): @pytest.mark.parametrize( - ("extra_args", "expected"), + ("extra_args", "memory_fit", "expected"), [ - ((), 16), # 8,192 tokens at --calib_seq 512 - (("--calib_size", "4"), 4), # never more than the samples - (("--calib_seq", "16384"), 1), # at least one sample - (("--batch_size", "1"), 1), # an explicit value is kept + ((), 1000, 64), # 32,768 tokens at --calib_seq 512 + ((), 5, 5), # what fits a decoder layer's forward + (("--calib_size", "4"), 1000, 4), # never more than the samples + (("--calib_seq", "65536"), 1000, 1), # at least one sample + (("--batch_size", "1"), 1000, 1), # an explicit value is kept ], - ids=["token_budget", "sample_cap", "long_sequence", "explicit"], + ids=["token_budget", "memory", "sample_cap", "long_sequence", "explicit"], ) -def test_layerwise_calibration_batch_reaches_the_dataloader(monkeypatch, extra_args, expected): - """A layerwise recipe sizes the default calibration batch by tokens.""" +def test_layerwise_calibration_batch_reaches_the_dataloader( + monkeypatch, extra_args, memory_fit, expected +): + """A layerwise recipe sizes the default calibration batch by tokens and layer memory.""" hf_ptq, args = _parse_hf_ptq_args( monkeypatch, "--pyt_ckpt_path", @@ -468,6 +471,11 @@ def make_calib_dataloader(args, *_args, **_kwargs): raise _DataloaderReachedError monkeypatch.setattr(hf_ptq, "make_calib_dataloader", make_calib_dataloader) + monkeypatch.setattr( + hf_ptq, + "get_max_layerwise_batch_size", + lambda model, seq_len, max_batch_size, num_samples, device: min(max_batch_size, memory_fit), + ) with pytest.raises(_DataloaderReachedError): hf_ptq.quantize_main(args, None, None, None, False, None, None, torch.device("cpu")) assert batch_sizes == [expected] diff --git a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py index 6aab8fe48c8..c92c9b54c1a 100644 --- a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py +++ b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py @@ -39,6 +39,7 @@ import modelopt.torch.quantization as mtq from modelopt.torch.quantization.extensions import get_cuda_ext_mx +from modelopt.torch.quantization.nn import TensorQuantizer from modelopt.torch.quantization.utils import ( enable_weight_access_and_writeback, is_quantized_linear, @@ -140,6 +141,46 @@ def test_cpu_offloaded_tinyllama(tmp_path): assert torch.allclose(output_ref.logits, output_test.logits) +def test_offloaded_quantizer_keeps_what_its_setter_wrote(tmp_path): + """A setter writes through ``.data``, which leaves ``_version`` alone; the offload must keep it.""" + model, _, _, inputs = make_cpu_offloaded_model(tmp_path) + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, lambda model: model(inputs)) + linear = model.model.layers[0].self_attn.q_proj + quantizer = linear.weight_quantizer + with enable_weight_access_and_writeback(linear, model): + pass + assert quantizer._amax.device.type == "cpu" + + amax = quantizer._amax * 2 + with enable_weight_access_and_writeback(linear, model): + quantizer.amax = amax + assert quantizer._amax.device.type == "cpu" + assert torch.equal(quantizer._amax, amax) + + +def test_offloaded_layer_keeps_bias_calibration_off_the_gpu(tmp_path): + """A static-bias quantizer's bias calibrator leaves the GPU with the offloaded layer's weights.""" + model, _, _, inputs = make_cpu_offloaded_model(tmp_path) + input_cfg = {"num_bits": 8, "axis": None, "bias": {-1: None, "type": "static"}} + quant_cfg = { + "quant_cfg": [ + {"quantizer_name": "*", "enable": False}, + {"quantizer_name": "*input_quantizer", "cfg": input_cfg}, + ], + "algorithm": "max", + } + mtq.quantize(model, make_layerwise_cfg(quant_cfg), lambda model: model(inputs)) + bias_state = [ + tensor + for quantizer in model.model.layers[0].modules() + if isinstance(quantizer, TensorQuantizer) and quantizer._bias_calibrator is not None + for tensor in vars(quantizer._bias_calibrator).values() + if isinstance(tensor, torch.Tensor) + ] + assert bias_state + assert all(tensor.device.type == "cpu" for tensor in bias_state) + + @pytest.mark.parametrize("use_checkpoint", [False, True], ids=["no_ckpt", "ckpt"]) def test_layerwise_calibrate_cpu_offloaded(tmp_path, use_checkpoint): """Layerwise calibration on CPU-offloaded model matches GPU-only reference.""" @@ -184,6 +225,20 @@ def test_layerwise_calibrate_cpu_offloaded(tmp_path, use_checkpoint): assert torch.allclose(output_ref.logits, output_test.logits) + # Between uses, the offloaded layer's quantizer state stays off the GPU like its weights. + offloaded_state = [ + tensor + for quantizer in model.model.layers[0].modules() + if isinstance(quantizer, TensorQuantizer) + for tensor in [ + *quantizer.buffers(recurse=False), + *(vars(quantizer._calibrator).values() if quantizer._calibrator is not None else ()), + ] + if isinstance(tensor, torch.Tensor) + ] + assert offloaded_state + assert all(tensor.device.type == "cpu" for tensor in offloaded_state) + if use_checkpoint: manifest_path = os.path.join(ckpt_dir, "manifest.json") assert os.path.isfile(manifest_path) diff --git a/tests/gpu/torch/quantization/test_gptq.py b/tests/gpu/torch/quantization/test_gptq.py index ee6dbec46e1..286883969f9 100644 --- a/tests/gpu/torch/quantization/test_gptq.py +++ b/tests/gpu/torch/quantization/test_gptq.py @@ -33,19 +33,23 @@ gptq_blockwise_update_fused_scalar, update_hessian, ) +from modelopt.torch.utils import get_available_gpu_memory from modelopt.torch.utils.dataset_utils import create_forward_loop, get_dataset_dataloader RAND_SEED = 42 torch.manual_seed(RAND_SEED) -def test_hessian_of_a_full_gpu_goes_to_an_idle_gpu(need_2_gpus, monkeypatch): - """A Hessian whose GPU is full lands on an idle GPU rather than the CPU, with the same values.""" +def _report_available(monkeypatch, available): + """Make each GPU report ``available[gpu]`` bytes for a new allocation, with no forward measured.""" + monkeypatch.setattr(calib_utils, "get_available_gpu_memory", lambda gpu: available[gpu]) + monkeypatch.setattr(calib_utils, "_calib_forward_peak", {}) - def used_fraction(device): - return 0.9 if torch.device(device).index in (0, None) else 0.1 - monkeypatch.setattr(calib_utils, "get_used_gpu_mem_fraction", used_fraction) +def test_hessian_of_a_full_gpu_goes_to_an_idle_gpu(need_2_gpus, monkeypatch): + """A Hessian whose GPU is full lands on an idle GPU rather than the CPU, with the same values.""" + total = torch.cuda.get_device_properties(0).total_memory + _report_available(monkeypatch, {0: total // 10, 1: total * 9 // 10}) linear = torch.nn.Linear(64, 8, bias=False, device="cuda:0") shard = torch.zeros(1, device="cuda:1") # the model also lives on cuda:1 helper = GPTQHelper(linear, "linear", offload_to_cpu=True) @@ -60,16 +64,15 @@ def used_fraction(device): @pytest.mark.parametrize("blocker", ["idle_gpu_full", "idle_gpu_too_small", "multi_rank"]) def test_hessian_of_a_full_gpu_falls_back_to_the_cpu(need_2_gpus, monkeypatch, blocker): """With no other GPU that can take it, or on any rank of a multi-rank run, it goes to the CPU.""" - idle_fraction = 0.9 if blocker == "idle_gpu_full" else 0.1 - - def used_fraction(device): - return 0.9 if torch.device(device).index in (0, None) else idle_fraction - - monkeypatch.setattr(calib_utils, "get_used_gpu_mem_fraction", used_fraction) - if blocker == "idle_gpu_too_small": - # 64 x 64 float32 is 16 KiB; leave 1 KiB free. - monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device=None: (1024, 2**40)) - elif blocker == "multi_rank": + total = torch.cuda.get_device_properties(0).total_memory + reserve = int(calib_utils._HESSIAN_UNMEASURED_RESERVE * total) + idle = { + "idle_gpu_full": reserve // 2, + "idle_gpu_too_small": reserve + 64 * 64 * 4 - 1, # one byte short of reserve + Hessian + "multi_rank": total, + }[blocker] + _report_available(monkeypatch, {0: total // 10, 1: idle}) + if blocker == "multi_rank": monkeypatch.setattr(calib_utils.dist, "is_initialized", lambda: True) monkeypatch.setattr(calib_utils.dist, "get_world_size", lambda group=None: 2) linear = torch.nn.Linear(64, 8, bias=False, device="cuda:0") @@ -78,6 +81,59 @@ def used_fraction(device): assert helper.hessian.device == torch.device("cpu") +@pytest.mark.parametrize( + ("measured_fraction", "stays"), + [(None, False), (0.05, True), (0.25, False)], + ids=["unmeasured", "small_forward", "large_forward"], +) +def test_a_measured_calibration_forward_sets_the_hessian_reserve( + monkeypatch, measured_fraction, stays +): + """At 30% available, a GPU keeps a Hessian only once a measured forward leaves it room.""" + total = torch.cuda.get_device_properties(0).total_memory + _report_available(monkeypatch, {0: total * 3 // 10}) + monkeypatch.setattr(calib_utils, "get_process_gpus", lambda: [0]) + if measured_fraction is not None: + calib_utils._calib_forward_peak[0] = int(measured_fraction * total) + linear = torch.nn.Linear(64, 8, bias=False, device="cuda:0") + helper = GPTQHelper(linear, "linear", offload_to_cpu=True) + assert helper.hessian.device == (torch.device("cuda", 0) if stays else torch.device("cpu")) + + +def test_available_gpu_memory_is_what_one_allocation_can_get(monkeypatch): + """Free pieces of partly used cached blocks do not count: one allocation of all the memory + reported beyond the driver's free memory comes from the cache.""" + torch.cuda.empty_cache() + whole = torch.empty(64 * 2**20, dtype=torch.uint8, device="cuda") + pieces = [torch.empty(3 * 2**20, dtype=torch.uint8, device="cuda") for _ in range(24)] + del whole, pieces[::2] # a whole free block, and free pieces of partly used ones + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda device=None: (0, 2**40)) + available = get_available_gpu_memory(torch.cuda.current_device()) + reserved = torch.cuda.memory_reserved() + block = torch.empty(available, dtype=torch.uint8, device="cuda") # noqa: F841 + assert torch.cuda.memory_reserved() == reserved + + +def test_gptq_records_its_calibration_forward_memory(monkeypatch): + """The next layer's Hessian reserve covers what this layer's calibration forward added, and + not an earlier peak such as the previous layer's Hessians.""" + monkeypatch.setattr(calib_utils, "_calib_forward_peak", {}) + model = torch.nn.Linear(256, 256, device="cuda") + inputs = torch.randn(2, 16, 256, device="cuda") + mtq.quantize(model, mtq.NVFP4_DEFAULT_CFG, forward_loop=lambda m: m(inputs)) + transient = 64 * 2**20 + + def forward_loop(m): + scratch = torch.empty(transient, dtype=torch.uint8, device="cuda") + m(inputs) + del scratch + + earlier = 16 * transient + torch.empty(earlier, dtype=torch.uint8, device="cuda") # freed at once; only the peak stays + gptq(model, forward_loop=forward_loop, perc_damp=0.1, block_size=16) + assert transient <= calib_utils._calib_forward_peak[torch.cuda.current_device()] < earlier // 2 + + @pytest.mark.parametrize( ("block_size", "dim", "model_weight", "expect_weight_change"), [ diff --git a/tests/gpu/torch/quantization/test_layerwise_calibrate.py b/tests/gpu/torch/quantization/test_layerwise_calibrate.py index 0e6751bfb79..0ad88b1fa8a 100644 --- a/tests/gpu/torch/quantization/test_layerwise_calibrate.py +++ b/tests/gpu/torch/quantization/test_layerwise_calibrate.py @@ -15,11 +15,16 @@ """Integration tests for layerwise_calibrate and LayerActivationCollector.""" +import pytest import torch import torch.nn as nn from modelopt.torch.quantization.model_calib import layerwise_calibrate -from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector +from modelopt.torch.quantization.utils import layerwise_calib +from modelopt.torch.quantization.utils.layerwise_calib import ( + LayerActivationCollector, + get_max_layerwise_batch_size, +) class _DecoderBlock(nn.Module): @@ -339,3 +344,98 @@ def test_skip_placeholder_uses_meta_device(): assert out.device == torch.device("meta") assert out.shape == torch.Size([2, 3]) assert out.dtype == torch.float32 + + +class _WideDecoderBlock(_DecoderBlock): + """A heavier kind of decoder layer: its FFN is 8x wider than _DecoderBlock's.""" + + def __init__(self, dim=16): + super().__init__(dim) + self.ffn = nn.Sequential( + nn.Linear(dim, dim * 32, bias=False), + nn.ReLU(), + nn.Linear(dim * 32, dim, bias=False), + ) + + +@pytest.mark.parametrize("same_class", [False, True], ids=["new_class", "same_class"]) +def test_layerwise_batch_probe_fits_the_heaviest_kind_of_layer(monkeypatch, same_class): + """The batch probe runs every kind of decoder layer, not just the first one, including layers + of one class that differ in width, like Nemotron-NAS's.""" + _register_test_discoverer(monkeypatch) + budget = 256 * 2**20 + monkeypatch.setattr(layerwise_calib, "get_available_gpu_memory", lambda gpu: budget) + model = _SimpleTransformerModel(n_layers=3, dim=256).cuda() + + def probe(): + return get_max_layerwise_batch_size(model, seq_len=512, max_batch_size=4096, num_samples=1) + + light = probe() + if same_class: + model.layers[2].ffn = _WideDecoderBlock(256).ffn.cuda() + else: + model.layers[2] = _WideDecoderBlock(256).cuda() + heavy = probe() + assert 1 < heavy < light < 4096 + assert not any(hasattr(layer, "_layerwise_calib") for layer in model.layers) + + # The heavy layer's forward at the probed batch takes at most half of the available memory. + x = torch.randn(heavy, 512, 256, device="cuda") + torch.cuda.reset_peak_memory_stats() + base = torch.cuda.memory_allocated() + with torch.no_grad(): + model.layers[2](x) + assert torch.cuda.max_memory_allocated() - base <= budget / 2 + + +class _OutOfMemoryOnTwoSamples(nn.Module): + """Runs *module*, but runs out of GPU memory on a batch of more than one sample.""" + + def __init__(self, module): + super().__init__() + self.module = module + + def forward(self, x, **kwargs): + if x.shape[0] > 1: + raise torch.cuda.OutOfMemoryError("CUDA out of memory.") + return self.module(x, **kwargs) + + +@pytest.mark.parametrize("where", ["layer_forward", "input_capture"]) +def test_layerwise_batch_probe_falls_back_to_one_sample(monkeypatch, where): + """A decoder layer that fits one sample but not two, or whose inputs for two samples cannot + be captured, gets a batch of one instead of stopping the run.""" + _register_test_discoverer(monkeypatch) + model = _SimpleTransformerModel(n_layers=2, dim=64).cuda() + if where == "layer_forward": + model.layers[1] = _OutOfMemoryOnTwoSamples(model.layers[1]) + else: + model.embed = _OutOfMemoryOnTwoSamples(model.embed) + + assert get_max_layerwise_batch_size(model, seq_len=8, max_batch_size=64, num_samples=1) == 1 + assert not any(hasattr(layer, "_layerwise_calib") for layer in model.layers) + + +class _MultiGPUModel(_SimpleTransformerModel): + """Decoder layers spread over GPUs, as a ``device_map`` places them.""" + + def forward(self, x, **kwargs): + x = self.embed(x) + for layer in self.layers: + x = layer(x.to(layer.norm.weight.device)) + return x + + +def test_layerwise_batch_probe_fits_every_gpu_a_kind_of_layer_runs_on(need_2_gpus, monkeypatch): + """Of two identical layers, only the first is probed, but the second's GPU, with less memory, + sets the batch.""" + _register_test_discoverer(monkeypatch) + model = _MultiGPUModel(n_layers=2, dim=256).cuda() + model.layers[1].to("cuda:1") + + def probe(budgets): + monkeypatch.setattr(layerwise_calib, "get_available_gpu_memory", lambda gpu: budgets[gpu]) + return get_max_layerwise_batch_size(model, seq_len=512, max_batch_size=4096, num_samples=1) + + small, large = 256 * 2**20, 1024 * 2**20 + assert probe({0: large, 1: small}) == probe({0: small, 1: small}) < probe({0: large, 1: large}) diff --git a/tests/unit/torch/quantization/test_gptq.py b/tests/unit/torch/quantization/test_gptq.py index ee1a55a8554..5744348d67e 100644 --- a/tests/unit/torch/quantization/test_gptq.py +++ b/tests/unit/torch/quantization/test_gptq.py @@ -22,6 +22,7 @@ from _test_utils.torch.distributed.utils import spawn_multiprocess_job import modelopt.torch.quantization as mtq +from modelopt.torch.quantization.utils import calib_utils from modelopt.torch.quantization.utils.calib_utils import GPTQHelper, update_hessian from modelopt.torch.utils.distributed import ParallelState @@ -172,3 +173,11 @@ def test_gptq_hessian_dp_sync(skip_on_windows): spawn_multiprocess_job( 2, partial(_gptq_hessian_dp_sync_helper, token_counts=(5, 0)), backend="gloo" ) + + +def test_a_calibration_run_forgets_earlier_forward_measurements(monkeypatch): + """Another run, maybe of another model, must not size its Hessian reserve from this one.""" + monkeypatch.setattr(calib_utils, "_calib_forward_peak", {0: 2**60}) + model = torch.nn.Linear(16, 16) + mtq.quantize(model, mtq.INT8_DEFAULT_CFG, forward_loop=lambda m: m(torch.randn(4, 16))) + assert 2**60 not in calib_utils._calib_forward_peak.values() diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py index 07e14d50f2a..5e7ee03307e 100644 --- a/tests/unit/torch/quantization/test_layerwise_calibrate.py +++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py @@ -31,6 +31,7 @@ _OutsideQuantizerCalibrator, _SkipLayer, ) +from modelopt.torch.utils.dataset_utils import create_forward_loop class _DecoderBlock(nn.Module): @@ -681,6 +682,48 @@ def test_set_layer_states_raises_on_empty_collected_inputs(monkeypatch): collector._unpatch_all_layers() +class _OutOfMemoryOnceBlock(_DecoderBlock): + """Decoder block whose next forward runs out of memory when ``oom`` is set.""" + + oom = False + + def forward(self, x, **kwargs): + if self.oom: + self.oom = False + raise torch.cuda.OutOfMemoryError("out of memory") + return super().forward(x, **kwargs) + + +class _InputIdsModel(_SimpleTransformerModel): + """Takes ``input_ids`` as the calibration forward loop passes it.""" + + def __init__(self): + super().__init__(n_layers=3, dim=16) + self.layers[0] = _OutOfMemoryOnceBlock(16) + + def forward(self, input_ids, **kwargs): + return super().forward(input_ids) + + +def test_out_of_memory_in_a_replayed_layer_is_not_split(monkeypatch): + """A batch split on OOM would replay the next cached batch in its place, so the OOM must stop.""" + _register_test_discoverer(monkeypatch) + model = _InputIdsModel() + batches = [{"input_ids": torch.randint(0, 32, (4, 8))} for _ in range(3)] + forward_loop = create_forward_loop(dataloader=batches) + + collector = LayerActivationCollector(model) + collector._patch_all_layers() + try: + collector.get_input_activations(model.layers[0], forward_loop) + model.layers[0].oom = True # layer 0 replays its cached batches while layer 1 is captured + with pytest.raises(RuntimeError) as error: + collector.get_input_activations(model.layers[1], forward_loop) + assert not isinstance(error.value, torch.cuda.OutOfMemoryError) + finally: + collector._unpatch_all_layers() + + def test_run_asserts_on_empty_cached_inputs(monkeypatch): """A layer in 'run' mode with no cached inputs must raise AssertionError.""" _register_test_discoverer(monkeypatch)