diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 012c6512b29..37fbdf7248c 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -75,6 +75,7 @@ Changelog - Add FP8 KV-cache cast variants for the partial-NVFP4 and weight-only general PTQ recipes: ``general/ptq/nvfp4_mlp_only-kv_fp8_cast``, ``general/ptq/nvfp4_experts_only-kv_fp8_cast``, ``general/ptq/nvfp4_omlp_only-kv_fp8_cast``, and ``general/ptq/nvfp4_weight_only-kv_fp8_cast``. These compose the same model-quant configs as their ``-kv_fp8`` siblings with the ``kv_fp8_cast`` unit (constant-amax FP8 KV cache, no KV calibration forward pass). - Add Nemotron-3-Super-120B-A12B PTQ recipes ``modelopt_recipes/huggingface/models/nvidia/Nemotron-3-Super-120B-A12B/ptq/nvfp4-mse.yaml`` (MSE-mixed) and ``nvfp4-max-calib.yaml`` (max-calib mixed): NVFP4 W4A4 routed experts + FP8 per-tensor shared experts / Mamba in/out_proj + FP8 KV cache. - Group layerwise calibration options under a nested ``LayerwiseConfig`` and add two knobs: ``get_qdq_activations_from_prev_layer`` (correct GPTQ-Hessian vs max-calib activation semantics — defaults to True for GPTQ, False for max/mse/local_hessian) and ``save_every`` (gate per-window ``next_inputs.pt`` activation-cache writes). Legacy bool ``layerwise`` and flat ``layerwise_checkpoint_dir`` keys still work; the bool form emits a ``DeprecationWarning``. +- Add two layerwise-calibration memory optimizations: ``calib_mutates_weights`` (set False for amax-only algorithms — max/mse/local_hessian — to skip the per-layer weight checkpoint blob and in-memory writeback, persisting only quantizer state), and meta-device skip-layer placeholders (already-calibrated layers emit zero-filled ``meta`` tensors instead of real-device buffers, eliminating their activation memory — models with real-device inter-layer ops on the hidden state are unsupported). - Add ``examples/alpamayo`` showing FP8, NVFP4, and AutoQuantize (mixed-precision) quantization of the Alpamayo (formerly Alpamayo-R1) ~10B vision-language-action model, with a joint VLM + diffusion calibration loop and both fake-quant and ``--real-quant`` packed-checkpoint export. See `examples/alpamayo/README.md `_ for details. - Refactor ``llm_qat`` example with unified YAML-based configuration and flexible dataset blending. ``ModelOptArgParser`` adds ``--config`` YAML support with CLI overrides and auto-generates ``ARGUMENTS.md`` from dataclass definitions. diff --git a/modelopt/torch/quantization/config.py b/modelopt/torch/quantization/config.py index 6c07afab90c..6257623d108 100644 --- a/modelopt/torch/quantization/config.py +++ b/modelopt/torch/quantization/config.py @@ -761,6 +761,17 @@ class LayerwiseConfig(ModeloptBaseConfig): ), ) + calib_mutates_weights: bool = ModeloptField( + default=True, + title="Whether layerwise calibration mutates layer weights.", + description=( + "Set to False only for algorithms that update solely " + "``TensorQuantizer._amax`` (max, mse, local_hessian). Rejected for " + "weight-mutating algorithms (GPTQ, AWQ, SmoothQuant) where it would " + "silently lose updates on resume." + ), + ) + def _coerce_layerwise_input(value): """Normalize a raw ``layerwise`` value to a dict; warn on deprecated bool.""" @@ -784,6 +795,11 @@ def _coerce_layerwise_input(value): class QuantizeAlgorithmConfig(ModeloptBaseConfig): """Calibration algorithm config base.""" + # Whether this algorithm mutates ``layer.weight`` during calibration. Amax-only + # algorithms (max/mse/local_hessian) set this False; it gates whether + # ``layerwise.calib_mutates_weights=False`` is allowed. + _mutates_weights: ClassVar[bool] = True + method: Literal[None] = ModeloptField( None, title="This field specifies the name of the calibration algorithm. If None, no calibration is performed.", @@ -862,6 +878,17 @@ def validate_layerwise_checkpoint_dir(self): ) return self + @model_validator(mode="after") + def _validate_non_mutating_layerwise_supported(self): + """Enforce the ``calib_mutates_weights=False`` whitelist.""" + if not self.layerwise.calib_mutates_weights and self._mutates_weights: + raise ValueError( + f"Algorithm '{self.method}' mutates layer weights in-place; " + "calib_mutates_weights=False would lose those updates on resume. " + "Only max/mse/local_hessian (amax-only) support this flag." + ) + return self + class _SharedStatesConfig(ModeloptBaseConfig): """The ``shared_states`` grouping knob, shared by max / mse / local_hessian calibration.""" @@ -920,6 +947,8 @@ class MaxCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): See `Integer Quantization `_ for the concepts. """ + _mutates_weights: ClassVar[bool] = False + method: Literal["max"] = ModeloptField("max") distributed_sync: bool | None = ModeloptField( @@ -968,6 +997,8 @@ class MseCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): When fp8_scale_sweep is enabled for a supported FP8-scale format, step_size is ignored. """ + _mutates_weights: ClassVar[bool] = False + method: Literal["mse"] = ModeloptField("mse") step_size: float | None = ModeloptField( @@ -1020,6 +1051,8 @@ class LocalHessianCalibConfig(_SharedStatesConfig, QuantizeAlgorithmConfig): """ + _mutates_weights: ClassVar[bool] = False + method: Literal["local_hessian"] = ModeloptField("local_hessian") step_size: float | None = ModeloptField( diff --git a/modelopt/torch/quantization/mode.py b/modelopt/torch/quantization/mode.py index 80be73daa2a..8f58d4c9dbb 100644 --- a/modelopt/torch/quantization/mode.py +++ b/modelopt/torch/quantization/mode.py @@ -230,6 +230,7 @@ def wrapped_calib_func( checkpoint_dir = layerwise_cfg.get("checkpoint_dir") qdq_from_prev = layerwise_cfg.get("get_qdq_activations_from_prev_layer", False) save_every = layerwise_cfg.get("save_every", 1) + calib_mutates_weights = layerwise_cfg.get("calib_mutates_weights", True) if method is not None and "awq" in method: # For backward compatibility kwargs["algorithm"] = method @@ -264,6 +265,7 @@ def wrapped_calib_func( checkpoint_dir=checkpoint_dir, get_qdq_activations_from_prev_layer=qdq_from_prev, save_every=save_every, + calib_mutates_weights=calib_mutates_weights, **kwargs, ) else: diff --git a/modelopt/torch/quantization/model_calib.py b/modelopt/torch/quantization/model_calib.py index 031e766c687..6e3590fc069 100644 --- a/modelopt/torch/quantization/model_calib.py +++ b/modelopt/torch/quantization/model_calib.py @@ -36,7 +36,7 @@ _CheckpointState, ) from modelopt.torch.utils import print_rank_0, warn_rank_0 -from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState +from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState, is_master from modelopt.torch.utils.distributed import is_initialized as dist_is_initialized from modelopt.torch.utils.distributed import size as dist_size from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method @@ -1925,6 +1925,7 @@ def layerwise_calibrate( checkpoint_dir = calib_kwargs.pop("checkpoint_dir", None) qdq_from_prev = calib_kwargs.pop("get_qdq_activations_from_prev_layer", False) save_every = calib_kwargs.pop("save_every", 1) + calib_mutates_weights = calib_kwargs.pop("calib_mutates_weights", True) if forward_loop is None: raise ValueError( @@ -1946,15 +1947,27 @@ def layerwise_calibrate( checkpoint_dir, num_layers, save_every=save_every, + calib_mutates_weights=calib_mutates_weights, ) start_layer = ckpt.start_layer if ckpt else 0 - input_getter = LayerActivationCollector(model) - input_getter._patch_all_layers(decoder_layers=transformer_layers) + layer_pbar = tqdm( + total=num_layers, + initial=start_layer, + desc="Layerwise calibration", + disable=not is_master(), + dynamic_ncols=True, + ) + + def _set_layer_status(status: str): + layer_pbar.set_postfix_str(status, refresh=True) - resumed_inputs = ckpt.setup_resume(transformer_layers) if ckpt and start_layer > 0 else None + input_getter = LayerActivationCollector(model, status_callback=_set_layer_status) try: + input_getter._patch_all_layers(decoder_layers=transformer_layers) + resumed_inputs = ckpt.setup_resume(transformer_layers) if ckpt and start_layer > 0 else None + # Bootstrap: get first layer's inputs (or use resumed inputs). layer_inputs = input_getter.get_first_layer_inputs( start_layer, resumed_inputs, forward_loop @@ -1984,7 +1997,7 @@ def _layer_forward_loop(m, _inputs=layer_inputs): is_last = layer_idx + 1 >= num_layers - with persistent_materialization(layer): + with persistent_materialization(layer, writeback=calib_mutates_weights): # qdq_from_prev=False: capture before calib_func so the forward # replay uses the original FP weights. Disable quantizers too in # case any pre-calibration observer behavior would perturb the @@ -2014,11 +2027,13 @@ def _layer_forward_loop(m, _inputs=layer_inputs): if ckpt: ckpt.save(layer_idx, model, transformer_layers, next_inputs) + layer_pbar.update(1) del layer_inputs torch.cuda.empty_cache() layer_inputs = next_inputs # noqa: F841 (used in next iteration's closure) finally: input_getter._unpatch_all_layers() + layer_pbar.close() if ckpt: ckpt.full_restore(transformer_layers, model) diff --git a/modelopt/torch/quantization/plugins/accelerate.py b/modelopt/torch/quantization/plugins/accelerate.py index f80e2478dc6..156669bfc8f 100644 --- a/modelopt/torch/quantization/plugins/accelerate.py +++ b/modelopt/torch/quantization/plugins/accelerate.py @@ -66,15 +66,22 @@ def _writeback_params_to_weights_map(module, align_hook): # on-disk version. OffloadedWeightsLoader.__getitem__ gives # state_dict priority over index, so this is sufficient. w_map[key] = tensor.detach().cpu() + warnings.warn( + "Accelerate disk-offload writeback is currently kept in CPU state_dict; " + "writing updates back to disk offload files is TODO.", + RuntimeWarning, + stacklevel=2, + ) @contextmanager -def weight_access_and_writeback_context(module): +def weight_access_and_writeback_context(module, writeback: bool = True): """Context manager for weight access and writeback for modules managed by accelerate. Handles CPU-offloaded and disk-offloaded models. Iterates over the module and all - its descendants, materializing weights from any offload hook found and writing them - back on exit. ``pre_forward`` is skipped on modules whose weights are already + its descendants, materializing weights from any offload hook found. If ``writeback`` + is True, writes materialized tensors back on exit. ``pre_forward`` is skipped on + modules whose weights are already materialized (not on meta) to avoid overwriting them with stale CPU copies. """ assert hasattr(module, "_hf_hook") @@ -99,7 +106,8 @@ def weight_access_and_writeback_context(module): finally: for mod, hook, was_materialized in materialized: hook.offload = True - _writeback_params_to_weights_map(mod, hook) + if writeback: + _writeback_params_to_weights_map(mod, hook) if was_materialized: hook.post_forward(mod, None) diff --git a/modelopt/torch/quantization/utils/core_utils.py b/modelopt/torch/quantization/utils/core_utils.py index 80352ced9eb..56eb373a7f6 100644 --- a/modelopt/torch/quantization/utils/core_utils.py +++ b/modelopt/torch/quantization/utils/core_utils.py @@ -29,6 +29,7 @@ from modelopt.torch.quantization.config import QuantizerCfgEntry from modelopt.torch.utils import get_unwrapped_name, print_rank_0 +from modelopt.torch.utils.network import temporarily_remove_accelerate_hook if TYPE_CHECKING: from collections.abc import Generator @@ -486,7 +487,24 @@ def _set_parameter(module: nn.Module, name: str, value: nn.Parameter): @contextmanager -def fsdp2_weight_access_and_writeback_context(module: nn.Module, root_model: nn.Module): +def _fsdp2_unshard_context(fsdp_module: FSDPModule): + """Unshard an FSDP2 module without replacing individual DTensor parameters.""" + fsdp_param_group = fully_shard.state(fsdp_module)._fsdp_param_group + was_sharded = fsdp_param_group.is_sharded + if was_sharded: + fsdp_module.unshard() + try: + with _disable_fsdp_unshard_reshard(fsdp_module): + yield + finally: + if was_sharded: + fsdp_module.reshard() + + +@contextmanager +def fsdp2_weight_access_and_writeback_context( + module: nn.Module, root_model: nn.Module, writeback: bool = True +): """Context manager for FSDP2 weight access and writeback. Gathers sharded DTensor parameters across FSDP/HSDP shards so they can be @@ -501,6 +519,11 @@ def fsdp2_weight_access_and_writeback_context(module: nn.Module, root_model: nn. assert not hasattr(module, "_hf_hook"), "We dont support FSDP2 with HF accelerate hooks" fsdp_module = _get_enclosing_fsdp_module(module, root_model) assert fsdp_module is not None, "Module is not wrapped by FSDP" + if not writeback: + with _fsdp2_unshard_context(fsdp_module): + yield + return + fsdp_device_mesh = _get_fsdp2_mesh(fsdp_module) fsdp_dim = fsdp_device_mesh.ndim @@ -540,7 +563,9 @@ def fsdp2_weight_access_and_writeback_context(module: nn.Module, root_model: nn. @contextmanager -def enable_weight_access_and_writeback(module, root_model, name_to_module: dict | None = None): +def enable_weight_access_and_writeback( + module, root_model, name_to_module: dict | None = None, writeback: bool = True +): """Enable weight access and writeback for a module. Useful for modules with weight not intact such as Linear layer in FSDP wrapped model or @@ -554,16 +579,18 @@ def enable_weight_access_and_writeback(module, root_model, name_to_module: dict total cost when called in a loop. This causes significant CPU overhead on large models, particularly Sparse MoE architectures where each expert is typically implemented as its own module. + writeback: Whether modified weights must be written back to the owning sharded/offload + representation when exiting the context. """ if _get_enclosing_fsdp_module(module, root_model, name_to_module) is not None: - context = fsdp2_weight_access_and_writeback_context(module, root_model) + context = fsdp2_weight_access_and_writeback_context(module, root_model, writeback) elif is_quantized_parallel_linear(module) and hasattr(module, "_hf_tp_plan"): # HF transformers TP sharded linear layer context = module.enable_weight_access_and_writeback() elif hasattr(module, "_hf_hook"): from ..plugins.accelerate import weight_access_and_writeback_context - context = weight_access_and_writeback_context(module) + context = weight_access_and_writeback_context(module, writeback) else: context = nullcontext() @@ -572,18 +599,23 @@ def enable_weight_access_and_writeback(module, root_model, name_to_module: dict @contextmanager -def persistent_materialization(layer): +def persistent_materialization(layer, writeback: bool = True): """Keep all layer weights materialized on GPU for the duration. Suppresses per-forward weight transfers so that N calibration batches pay the cost of one load/unload instead of N. - - **FSDP2**: patches ``FSDPParamGroup.unshard/reshard`` to no-ops, then - gathers weights once via ``enable_weight_access_and_writeback``. - - **Accelerate**: materializes weights and sets ``hook.offload = False`` - so per-forward hooks skip materialization/offloading. + - **FSDP2**: gathers weights once via ``enable_weight_access_and_writeback``, + then patches ``FSDPParamGroup.unshard/reshard`` to no-ops. + - **Accelerate**: materializes weights, sets ``hook.offload = False``, + and bypasses the layer's top-level accelerate hook while the weights are + materialized. """ - with _disable_fsdp_unshard_reshard(layer), enable_weight_access_and_writeback(layer, layer): + with ( + enable_weight_access_and_writeback(layer, layer, writeback=writeback), + _disable_fsdp_unshard_reshard(layer), + temporarily_remove_accelerate_hook(layer), + ): yield diff --git a/modelopt/torch/quantization/utils/layerwise_calib.py b/modelopt/torch/quantization/utils/layerwise_calib.py index 3cc0ff6be8e..070ee521cd5 100644 --- a/modelopt/torch/quantization/utils/layerwise_calib.py +++ b/modelopt/torch/quantization/utils/layerwise_calib.py @@ -42,6 +42,8 @@ ) if TYPE_CHECKING: + from collections.abc import Callable + from modelopt.torch.opt.searcher import ForwardLoop @@ -105,11 +107,15 @@ class LayerActivationCollector: Each decoder layer is patched with a unified forward whose behaviour is governed by a per-layer :class:`_LayerCalibState`: - * **skip** — return a zero-filled dummy whose shape and type match the - layer's real output (reconstructed from lightweight metadata). No - computation is performed. The correctly shaped dummy ensures un-patched - inter-layer operations in the parent forward (e.g. LayerNorm, tuple - unpacking) do not raise shape or type errors. + * **skip** — return a zero-filled meta-device dummy whose shape and type + match the layer's real output (reconstructed from lightweight metadata). + No computation or real-device allocation is performed. Tuple/list + structure is preserved for parent code that unpacks outputs, but + real-device inter-layer tensor operations are intentionally unsupported. + These meta placeholders are used unconditionally (not gated by + ``calib_mutates_weights``): a model whose parent ``forward`` runs + real-device ops on the hidden state *between* decoder blocks is not + supported by layerwise calibration — use non-layerwise calibration for it. * **run** — replay previously captured inputs through the original forward, ignoring whatever the parent passes in. Only the just-calibrated layer uses this mode, so its output reflects updated weights. @@ -124,18 +130,19 @@ class LayerActivationCollector: _decoder_layer_support: list[tuple[Any, Any]] = [] _LAYER_ATTR = "_layerwise_calib" - def __init__(self, model: nn.Module): + def __init__(self, model: nn.Module, status_callback: Callable[[str], None] | None = None): """Initialize the collector for the given model.""" self.model = model self._decoder_layers: nn.ModuleList | None = None self._layer_to_idx: dict[nn.Module, int] = {} self._patched = False + self._status_callback = status_callback def _swap_to_dummy(self, idx: int): """Replace decoder layer *idx* with a parameter-free dummy. ``output_meta`` is intentionally preserved on the original layer: the - ``_SkipLayer`` reads it to produce correctly shaped zero-filled outputs + ``_SkipLayer`` reads it to produce correctly shaped placeholder outputs for the parent forward pass. """ assert self._decoder_layers is not None @@ -173,7 +180,9 @@ def _extract_output_meta(output): Recursively handles tensors, tuples, lists, and non-tensor values (e.g. None). The returned structure can be passed to ``_zeros_from_meta`` to reconstruct a - zero-filled output with identical shape and type. + zero-filled output with identical structure, shape, and dtype. Tensor + placeholders are allocated on the meta device; the recorded device is kept + in metadata for checkpoint compatibility. """ if isinstance(output, torch.Tensor): return ("tensor", output.shape, output.dtype, output.device) @@ -188,11 +197,11 @@ def _extract_output_meta(output): @staticmethod def _zeros_from_meta(meta): - """Reconstruct a zero-filled output from metadata produced by ``_extract_output_meta``.""" + """Reconstruct a zero-filled meta placeholder from ``_extract_output_meta`` metadata.""" tag = meta[0] if tag == "tensor": - _, shape, dtype, device = meta - return torch.zeros(shape, dtype=dtype, device=device) + _, shape, dtype, _device = meta + return torch.zeros(shape, dtype=dtype, device=torch.device("meta")) if tag == "tuple": return tuple(LayerActivationCollector._zeros_from_meta(m) for m in meta[1]) if tag == "list": @@ -258,6 +267,16 @@ def _early_stop_forward(module_self, *args, **kwargs): return module_self._original_forward(*args, **kwargs) except _EarlyStopForwardError: return None + except RuntimeError as e: + if "meta" not in str(e).lower(): + raise + raise RuntimeError( + "Layerwise calibration represents skipped decoder layers with " + "meta-device placeholder outputs, so it does not support models " + "that run real-device operations on the hidden state between " + "decoder blocks. Use non-layerwise calibration for this " + "architecture." + ) from e bind_forward_method(self.model, _early_stop_forward, "_original_forward") except Exception: @@ -322,8 +341,14 @@ def _set_layer_states(self, layer_idx: int): cur.mode = "capture" cur.collected_inputs = [] - def _log_layer_summary(self, layer_idx: int): - """Log a one-line summary of layer modes for the current calibration step.""" + def _emit_status(self, status: str): + if self._status_callback is None: + print_rank_0(status) + else: + self._status_callback(status) + + def _layer_summary(self, layer_idx: int) -> str: + """Return a one-line summary of layer modes for the current calibration step.""" assert self._decoder_layers is not None n = len(self._decoder_layers) groups: dict[str, list[int]] = {} @@ -338,7 +363,10 @@ def _log_layer_summary(self, layer_idx: int): continue ids = groups[mode] parts.append(f"{mode}: {len(ids)}" if mode == "skip" else f"{mode}: {ids}") - print_rank_0(f"Calibrating layer {layer_idx + 1}/{n} | {' | '.join(parts)}") + return f"Calibrating layer {layer_idx + 1}/{n} | {' | '.join(parts)}" + + def _log_layer_summary(self, layer_idx: int): + self._emit_status(self._layer_summary(layer_idx)) @torch.no_grad() def get_input_activations(self, layer: torch.nn.Module, forward_loop: ForwardLoop) -> list: @@ -402,7 +430,8 @@ def get_first_layer_inputs( assert self._decoder_layers is not None if resumed_inputs is not None: - print_rank_0(f"Calibrating layer {start_layer + 1} (resumed)") + n = len(self._decoder_layers) + self._emit_status(f"Calibrating layer {start_layer + 1}/{n} | resumed") for i in range(start_layer): self._swap_to_dummy(i) layer = self._decoder_layers[start_layer] @@ -420,7 +449,8 @@ def cache_outputs_for_next_layer_calib( This puts *layer* into "run" mode (setting its ``output_meta``) and the next layer into "capture" mode, then runs *forward_loop*. Returns the - captured inputs for the next layer. + captured inputs for the next layer. Callers should keep *layer* + materialized for the duration when using offload frameworks. Must be called only when a next layer exists (i.e. *layer* is not the last decoder layer). @@ -429,11 +459,9 @@ def cache_outputs_for_next_layer_calib( layer_idx = self._layer_to_idx[layer] next_idx = layer_idx + 1 assert next_idx < len(self._decoder_layers), "No next layer to capture inputs for." - from .core_utils import persistent_materialization next_layer = self._decoder_layers[next_idx] - with persistent_materialization(layer): - return self.get_input_activations(next_layer, forward_loop) + return self.get_input_activations(next_layer, forward_loop) def _move_to_device(obj: Any, device: torch.device) -> Any: @@ -448,19 +476,6 @@ def _move_to_device(obj: Any, device: torch.device) -> Any: return obj -def _remap_output_metadata_device(meta: tuple, device: torch.device) -> tuple: - """Patch the device field inside output_meta tuples so _zeros_from_meta uses *device*.""" - tag = meta[0] - if tag == "tensor": - _, shape, dtype, _old_device = meta - return ("tensor", shape, dtype, device) - if tag == "tuple": - return ("tuple", tuple(_remap_output_metadata_device(m, device) for m in meta[1])) - if tag == "list": - return ("list", [_remap_output_metadata_device(m, device) for m in meta[1]]) - return meta - - def _read_manifest(checkpoint_dir: str) -> dict | None: """Read manifest.json from *checkpoint_dir*. Returns None if missing or corrupt.""" path = os.path.join(checkpoint_dir, "manifest.json") @@ -478,6 +493,7 @@ def _write_manifest( last_completed_layer: int, num_layers: int, save_every: int, + calib_mutates_weights: bool, ) -> None: """Atomically write manifest.json. Config keys are persisted so resume can detect drift.""" path = os.path.join(checkpoint_dir, "manifest.json") @@ -488,6 +504,7 @@ def _write_manifest( "last_completed_layer": last_completed_layer, "num_layers": num_layers, "save_every": save_every, + "calib_mutates_weights": calib_mutates_weights, }, f, ) @@ -501,21 +518,27 @@ def _layer_dir(checkpoint_dir: str, idx: int) -> str: def _save_layer_files( checkpoint_dir: str, idx: int, - weights: dict, + weights: dict | None, qstate: dict, + quantizer_buffers: dict | None, output_meta: tuple, ) -> None: """Write the per-layer files for layer *idx*. - ``weights.pt``, ``quantizer_state.pt``, and ``output_meta.pt`` are written - every call. ``next_inputs.pt`` and ``manifest.json`` are deferred to window - boundaries in :meth:`_CheckpointState.save`. + Exactly one of ``weights`` (full layer state_dict) or ``quantizer_buffers`` + (just the TensorQuantizer state_dict slice, used when calibration does not mutate weights) + is written; ``full_restore`` falls back to whichever is present. + ``next_inputs.pt`` and ``manifest.json`` are deferred to window boundaries + in :meth:`_CheckpointState.save`. """ d = _layer_dir(checkpoint_dir, idx) if os.path.isdir(d): shutil.rmtree(d) os.makedirs(d) - torch.save(weights, os.path.join(d, "weights.pt")) + if weights is not None: + torch.save(weights, os.path.join(d, "weights.pt")) + elif quantizer_buffers is not None: + torch.save(quantizer_buffers, os.path.join(d, "quantizer_buffers.pt")) torch.save(qstate, os.path.join(d, "quantizer_state.pt")) torch.save(output_meta, os.path.join(d, "output_meta.pt")) @@ -556,6 +579,7 @@ def __init__( num_layers: int, start_layer: int = 0, save_every: int = 1, + calib_mutates_weights: bool = True, ): if dist.is_initialized() and dist.size() > 1: raise RuntimeError( @@ -568,6 +592,7 @@ def __init__( self.num_layers = num_layers self.start_layer = start_layer self.save_every = save_every + self.calib_mutates_weights = calib_mutates_weights # Tracks the most recent saved layer so save() can window-save the layers # since the last save event. Initialized to start_layer - 1 so the first # save event after resume covers the new work only. @@ -579,6 +604,7 @@ def from_folder( checkpoint_dir: str | None, num_layers: int, save_every: int = 1, + calib_mutates_weights: bool = True, ) -> _CheckpointState | None: """Create from folder. Detects resume point. Returns None if no checkpoint_dir.""" if not checkpoint_dir: @@ -587,11 +613,10 @@ def from_folder( info = detect_resume_point(checkpoint_dir) if info is not None: manifest = info[1] - # Pre-0.45 manifests omit save_every; skip the check for keys absent - # from the on-disk manifest. for key, new_value in ( ("num_layers", num_layers), ("save_every", save_every), + ("calib_mutates_weights", calib_mutates_weights), ): ckpt_value = manifest.get(key) if ckpt_value is not None and ckpt_value != new_value: @@ -609,6 +634,7 @@ def from_folder( num_layers, start_layer=start, save_every=save_every, + calib_mutates_weights=calib_mutates_weights, ) def setup_resume(self, layers: nn.ModuleList) -> list | None: @@ -628,8 +654,6 @@ def setup_resume(self, layers: nn.ModuleList) -> list | None: meta = torch.load( os.path.join(d, "output_meta.pt"), map_location="cpu", weights_only=False ) - layer_device = get_module_device(layers[i]) - meta = _remap_output_metadata_device(meta, layer_device) layers[i]._layerwise_calib.output_meta = meta d = _layer_dir(self.checkpoint_dir, last_ckpt) @@ -646,7 +670,10 @@ def full_restore(self, layers: nn.ModuleList, model: nn.Module) -> None: """Restore weights and quantizer state for layers 0..K-1 after the calibration loop.""" from modelopt.torch.quantization.config import QuantizeConfig from modelopt.torch.quantization.conversion import restore_quantizer_state - from modelopt.torch.quantization.utils.core_utils import enable_weight_access_and_writeback + from modelopt.torch.quantization.utils.core_utils import ( + enable_weight_access_and_writeback, + set_quantizer_state_dict, + ) if self.start_layer == 0: return @@ -668,10 +695,31 @@ def full_restore(self, layers: nn.ModuleList, model: nn.Module) -> None: weights_only=False, ) restore_quantizer_state(layer, dummy_config, {"quantizer_state": qstate}) - weights = torch.load( - os.path.join(d, "weights.pt"), map_location=layer_device, weights_only=False - ) - layer.load_state_dict(weights, strict=False, assign=True) + weights_path = os.path.join(d, "weights.pt") + buffers_path = os.path.join(d, "quantizer_buffers.pt") + if os.path.isfile(weights_path): + weights = torch.load( + weights_path, map_location=layer_device, weights_only=False + ) + layer.load_state_dict(weights, strict=False, assign=True) + elif os.path.isfile(buffers_path): + # Non-mutating calibration mode: restore just the TensorQuantizer + # state_dict (carries _amax). The layer's other weights were not + # modified, so the in-memory values already match what would have + # been saved. + quantizer_buffers = torch.load( + buffers_path, map_location=layer_device, weights_only=False + ) + set_quantizer_state_dict(layer, quantizer_buffers) + else: + # restore_quantizer_state freshly registered _amax via torch.empty; + # with neither file to fill it, the layer would silently carry + # uninitialized buffers. Fail loudly instead. + raise FileNotFoundError( + f"Layer {i} checkpoint at {d} is missing both weights.pt and " + "quantizer_buffers.pt; the checkpoint is incomplete. " + "Use a fresh checkpoint directory." + ) print_rank_0(f"Checkpoint: restored {self.start_layer} previously calibrated layers") @@ -692,13 +740,21 @@ def save( previous boundary. """ from modelopt.torch.quantization.conversion import quantizer_state - from modelopt.torch.quantization.utils.core_utils import enable_weight_access_and_writeback + from modelopt.torch.quantization.utils.core_utils import ( + enable_weight_access_and_writeback, + get_quantizer_state_dict, + ) _cpu = torch.device("cpu") layer = layers[layer_idx] - with enable_weight_access_and_writeback(layer, model): + with enable_weight_access_and_writeback(layer, model, writeback=False): qstate = _move_to_device(quantizer_state(layer), _cpu) - weights = _move_to_device(layer.state_dict(), _cpu) + if self.calib_mutates_weights: + weights = _move_to_device(layer.state_dict(), _cpu) + quantizer_buffers = None + else: + weights = None + quantizer_buffers = _move_to_device(get_quantizer_state_dict(layer), _cpu) output_meta = getattr(layer._layerwise_calib, "output_meta", None) if output_meta is None: @@ -710,6 +766,7 @@ def save( layer_idx, weights, qstate, + quantizer_buffers, _move_to_device(output_meta, _cpu), ) @@ -729,6 +786,7 @@ def save( layer_idx, self.num_layers, save_every=self.save_every, + calib_mutates_weights=self.calib_mutates_weights, ) window_start = self._last_saved_layer + 1 self._last_saved_layer = layer_idx diff --git a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py index 03d20d26746..3a2816540a8 100644 --- a/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py +++ b/tests/gpu/torch/quantization/plugins/test_accelerate_gpu.py @@ -474,6 +474,8 @@ def forward(self, x): def test_skip_dummy_has_no_hf_hook(monkeypatch): """Dummies must not carry _hf_hook from the original layer.""" + from contextlib import nullcontext + monkeypatch.setattr( LayerActivationCollector, "_decoder_layer_support", @@ -494,8 +496,12 @@ def forward_loop(m): collector = LayerActivationCollector(model) collector._patch_all_layers() try: - for layer in list(model.layers): - collector.get_input_activations(layer, forward_loop) + for i, layer in enumerate(list(model.layers)): + run_layer_context = ( + persistent_materialization(model.layers[i - 1]) if i > 0 else nullcontext() + ) + with run_layer_context: + collector.get_input_activations(layer, forward_loop) for i in range(2): dummy = model.layers[i] @@ -505,6 +511,26 @@ def forward_loop(m): collector._unpatch_all_layers() +def _assert_persistent_materialization_bypasses_top_hook(layer): + from modelopt.torch.quantization.utils import persistent_materialization + + assert hasattr(layer, "_hf_hook") + original_old_forward = layer._old_forward + + def sentinel_forward(*args, **kwargs): + return "unhooked" + + layer._old_forward = sentinel_forward + try: + with persistent_materialization(layer): + assert layer.forward is sentinel_forward + assert layer("unused") == "unhooked" + + assert layer.forward is not sentinel_forward + finally: + layer._old_forward = original_old_forward + + def test_persistent_materialization_cpu_offloaded(tmp_path): """persistent_materialization keeps CPU-offloaded weights on GPU and writes back modifications.""" model, config, _, inputs = _make_cpu_offloaded_model(tmp_path) @@ -513,6 +539,8 @@ def test_persistent_materialization_cpu_offloaded(tmp_path): # Verify offloaded (meta device) assert all(p.device.type == "meta" for p in offloaded_layer.parameters()) + _assert_persistent_materialization_bypasses_top_hook(offloaded_layer) + # Save reference weight linear = None with enable_weight_access_and_writeback(offloaded_layer, model): @@ -626,6 +654,8 @@ def test_persistent_materialization_disk_offloaded(tmp_path): # Verify offloaded (meta device) assert all(p.device.type == "meta" for p in offloaded_layer.parameters()) + _assert_persistent_materialization_bypasses_top_hook(offloaded_layer) + # Save reference weight linear = None with enable_weight_access_and_writeback(offloaded_layer, model): diff --git a/tests/gpu/torch/quantization/test_fsdp2.py b/tests/gpu/torch/quantization/test_fsdp2.py index 6d47e1620ab..1ad88d087d1 100644 --- a/tests/gpu/torch/quantization/test_fsdp2.py +++ b/tests/gpu/torch/quantization/test_fsdp2.py @@ -16,6 +16,7 @@ """Test of quantization with FSDP2.""" import copy +from contextlib import contextmanager from functools import partial import pytest @@ -213,6 +214,8 @@ def forward(self, x): def _test_layerwise_calibrate_fsdp2(rank, size): """Layerwise calibration on FSDP2-wrapped model matches non-FSDP reference.""" + import modelopt.torch.quantization.model_calib as model_calib + dim = 32 torch.manual_seed(1) model = _SimpleTransformerModel(n_layers=3, dim=dim).cuda() @@ -228,12 +231,28 @@ def _test_layerwise_calibrate_fsdp2(rank, size): ), *old_support, ] + original_persistent_materialization = model_calib.persistent_materialization + materialized_fsdp_layers = 0 + + @contextmanager + def tracked_persistent_materialization(layer, writeback=True): + nonlocal materialized_fsdp_layers + with original_persistent_materialization(layer, writeback=writeback): + if isinstance(layer, torch.distributed.fsdp.FSDPModule) and not writeback: + assert all(not isinstance(param, DTensor) for param in layer.parameters()) + materialized_fsdp_layers += 1 + yield try: + model_calib.persistent_materialization = tracked_persistent_materialization + # Reference: non-FSDP layerwise calibration ref_model = copy.deepcopy(model) seq_cfg = copy.deepcopy(mtq.INT8_DEFAULT_CFG) - seq_cfg["algorithm"] = {"method": "max", "layerwise": True} + seq_cfg["algorithm"] = { + "method": "max", + "layerwise": {"enable": True, "calib_mutates_weights": False}, + } mtq.quantize(ref_model, seq_cfg, lambda m: m(inputs)) output_ref = ref_model(inputs) @@ -245,7 +264,9 @@ def _test_layerwise_calibrate_fsdp2(rank, size): output_test = model(inputs) assert torch.allclose(output_ref, output_test) + assert materialized_fsdp_layers == len(model.layers) finally: + model_calib.persistent_materialization = original_persistent_materialization LayerActivationCollector._decoder_layer_support = old_support @@ -300,6 +321,13 @@ def _test_persistent_materialization(rank, size): with enable_weight_access_and_writeback(layer[0], model): assert torch.allclose(layer[0].weight, ref_weight + 1.0) + with persistent_materialization(layer, writeback=False): + assert not isinstance(layer[0].weight, DTensor) + assert layer[0].weight.device.type == "cuda" + layer(inputs) + + assert isinstance(next(iter(layer.parameters())), DTensor) + def test_persistent_materialization(dist_workers): dist_workers.run(_test_persistent_materialization) diff --git a/tests/gpu/torch/quantization/test_layerwise_calibrate.py b/tests/gpu/torch/quantization/test_layerwise_calibrate.py index d38b82f46fb..0e6751bfb79 100644 --- a/tests/gpu/torch/quantization/test_layerwise_calibrate.py +++ b/tests/gpu/torch/quantization/test_layerwise_calibrate.py @@ -329,3 +329,13 @@ def weight_doubling_calib(layer, layer_forward_loop, **kwargs): # Verify by running model.layers[0] with its updated weights actual = model.layers[0](x) assert torch.allclose(actual, expected) + + +def test_skip_placeholder_uses_meta_device(): + recorded_device = torch.device("cpu") + meta = ("tensor", torch.Size([2, 3]), torch.float32, recorded_device) + + out = LayerActivationCollector._zeros_from_meta(meta) + assert out.device == torch.device("meta") + assert out.shape == torch.Size([2, 3]) + assert out.dtype == torch.float32 diff --git a/tests/unit/torch/quantization/test_config_validation.py b/tests/unit/torch/quantization/test_config_validation.py index 7036f951187..4b969d3259c 100644 --- a/tests/unit/torch/quantization/test_config_validation.py +++ b/tests/unit/torch/quantization/test_config_validation.py @@ -29,11 +29,15 @@ INT4_AWQ_CFG, NVFP4_DEFAULT_CFG, W4A8_AWQ_BETA_CFG, + AWQLiteCalibConfig, GPTQCalibConfig, LayerwiseConfig, + LocalHessianCalibConfig, MaxCalibConfig, + MseCalibConfig, QuantizeConfig, QuantizerAttributeConfig, + SmoothQuantCalibConfig, find_quant_cfg_entry_by_path, need_calibration, normalize_quant_cfg_list, @@ -690,6 +694,7 @@ def test_default_dump_shape(self): "get_qdq_activations_from_prev_layer": False, "checkpoint_dir": None, "save_every": 1, + "calib_mutates_weights": True, } assert "layerwise_checkpoint_dir" not in dumped @@ -697,6 +702,23 @@ def test_save_every_must_be_positive(self): with pytest.raises(ValidationError): MaxCalibConfig(layerwise={"enable": True, "save_every": 0}) + @pytest.mark.parametrize( + "cfg_cls", + [GPTQCalibConfig, AWQLiteCalibConfig, SmoothQuantCalibConfig], + ) + def test_calib_mutates_weights_false_rejected_for_weight_mutating_algorithms(self, cfg_cls): + """Whitelist: only amax-only algorithms (max/mse/local_hessian) may set + calib_mutates_weights=False. Weight-mutating algorithms (GPTQ folds Hessian + updates, AWQ/SmoothQuant fold pre-quant scales) must reject the flag. + """ + with pytest.raises(ValidationError, match="mutates layer weights in-place"): + cfg_cls(layerwise={"enable": True, "calib_mutates_weights": False}) + + @pytest.mark.parametrize("cfg_cls", [MaxCalibConfig, MseCalibConfig, LocalHessianCalibConfig]) + def test_calib_mutates_weights_false_accepted_for_amax_only_algorithms(self, cfg_cls): + cfg = cfg_cls(layerwise={"enable": True, "calib_mutates_weights": False}) + assert cfg.layerwise.calib_mutates_weights is False + class TestFourOverSixBlockSizes: """`four_over_six` is an accepted block_sizes key for NVFP4 4/6 adaptive weight scaling. diff --git a/tests/unit/torch/quantization/test_layerwise_calibrate.py b/tests/unit/torch/quantization/test_layerwise_calibrate.py index 3deae78b70a..1c8234fdb2f 100644 --- a/tests/unit/torch/quantization/test_layerwise_calibrate.py +++ b/tests/unit/torch/quantization/test_layerwise_calibrate.py @@ -313,8 +313,8 @@ def forward_loop(m): collector._unpatch_all_layers() -def test_skip_output_preserves_shape_with_inter_layer_norm(monkeypatch): - """Skip outputs must have correct shape for un-patched LayerNorm between layers.""" +def test_inter_layer_op_raises_descriptive_error(monkeypatch): + """Real-device inter-layer ops on meta skip placeholders raise an actionable error.""" _register_test_discoverer(monkeypatch) model = _InterLayerNormModel(n_layers=5, dim=16) data = [torch.randn(2, 16) for _ in range(3)] @@ -326,9 +326,9 @@ def forward_loop(m): collector = LayerActivationCollector(model) collector._patch_all_layers() try: - for layer in model.layers: - inputs = collector.get_input_activations(layer, forward_loop) - assert len(inputs) == len(data) + with pytest.raises(RuntimeError, match="non-layerwise calibration"): + for layer in model.layers: + collector.get_input_activations(layer, forward_loop) finally: collector._unpatch_all_layers() @@ -612,6 +612,56 @@ def _int8_cfg_with_algorithm(algorithm: dict) -> dict: return cfg +def test_layerwise_calibrate_uses_global_layer_tqdm(monkeypatch): + _register_test_discoverer(monkeypatch) + + class _FakeTqdm: + instances = [] + + def __init__(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + self.postfixes = [] + self.updates = [] + self.closed = False + _FakeTqdm.instances.append(self) + + def set_postfix_str(self, status, refresh=True): + self.postfixes.append((status, refresh)) + + def update(self, n=1): + self.updates.append(n) + + def close(self): + self.closed = True + + monkeypatch.setattr("modelopt.torch.quantization.model_calib.tqdm", _FakeTqdm) + + torch.manual_seed(0) + model = _SimpleTransformerModel(n_layers=3, dim=16) + calib_data = [torch.randint(0, 32, (2, 8)) for _ in range(2)] + + def forward_loop(m): + for batch in calib_data: + m(batch) + + def calib_func(layer, layer_forward_loop): + layer_forward_loop(layer) + + layerwise_calibrate(model, forward_loop, calib_func) + + assert len(_FakeTqdm.instances) == 1 + pbar = _FakeTqdm.instances[0] + assert pbar.kwargs["total"] == 3 + assert pbar.kwargs["initial"] == 0 + assert pbar.kwargs["desc"] == "Layerwise calibration" + assert pbar.kwargs["dynamic_ncols"] is True + assert pbar.updates == [1, 1, 1] + assert pbar.closed + assert any(status.startswith("Calibrating layer 1/3") for status, _ in pbar.postfixes) + assert any(status.startswith("Calibrating layer 3/3") for status, _ in pbar.postfixes) + + def _awq_layerwise_config() -> dict: """INT4 weight-only AWQ config sized for the _DecoderBlock test model.""" cfg = copy.deepcopy(mtq.INT4_AWQ_CFG) @@ -868,20 +918,24 @@ def test_layerwise_save_every_writes_next_inputs_only_at_window_boundaries(monke @pytest.mark.parametrize( - ("n_layers", "save_every", "rewind_to"), + ("scenario", "n_layers", "save_every", "calib_mutates_weights", "rewind_to"), [ + # Pins the quantizer_buffers.pt restore path (no weights.pt on disk). + ("non_mutating", 3, 1, False, 0), # Pins the per-call snapshot fix: each save() captures the # just-calibrated layer's state before the next-layer capture forward # swaps it to _SkipLayer. - (4, 2, 1), + ("save_every", 4, 2, True, 1), ], ) def test_layerwise_checkpoint_resume_matches_one_shot_amax( - monkeypatch, tmp_path, n_layers, save_every, rewind_to + monkeypatch, tmp_path, scenario, n_layers, save_every, calib_mutates_weights, rewind_to ): """Full run → rewind manifest → fresh resume reproduces one-shot ``_amax``. - Covers the per-window save/resume path with always-full-weight saves. + Single test covering both checkpoint optimizations. For the + non-mutating calibration case also asserts the on-disk shape (no + ``weights.pt``, ``quantizer_buffers.pt`` present per layer). """ _register_test_discoverer(monkeypatch) @@ -896,6 +950,7 @@ def build_cfg(ckpt_dir): "enable": True, "checkpoint_dir": str(ckpt_dir), "save_every": save_every, + "calib_mutates_weights": calib_mutates_weights, }, } ) @@ -912,12 +967,19 @@ def build_cfg(ckpt_dir): setup_model = _SimpleTransformerModel(n_layers=n_layers, dim=16) mtq.quantize(setup_model, build_cfg(resume_dir), forward_loop=forward_loop) + if scenario == "non_mutating": + for name in _layer_dir_names(resume_dir): + d = resume_dir / name + assert not (d / "weights.pt").exists() + assert (d / "quantizer_buffers.pt").exists() + (resume_dir / "manifest.json").write_text( json.dumps( { "last_completed_layer": rewind_to, "num_layers": n_layers, "save_every": save_every, + "calib_mutates_weights": calib_mutates_weights, } ) ) @@ -926,7 +988,7 @@ def build_cfg(ckpt_dir): resumed_model = _SimpleTransformerModel(n_layers=n_layers, dim=16) mtq.quantize(resumed_model, build_cfg(resume_dir), forward_loop=forward_loop) - _assert_amax_close(_collect_amax(resumed_model), baseline_amax, "resume") + _assert_amax_close(_collect_amax(resumed_model), baseline_amax, f"{scenario} resume") def test_layerwise_save_every_mid_window_crash_recovers_at_prev_boundary(monkeypatch, tmp_path): @@ -1003,6 +1065,7 @@ def test_layerwise_checkpoint_mismatch_save_every_raises(monkeypatch, tmp_path): "last_completed_layer": 1, "num_layers": 4, "save_every": 2, + "calib_mutates_weights": True, } ) )