Add checkpoint save/resume for sequential calibration of large models

Per-layer checkpoints allow sequential calibration to resume from the last
completed layer after a crash or preemption. Also extends sequential
calibration to work with FSDP2 and accelerate CPU-offloaded models.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Signed-off-by: realAsma <akuriparambi@nvidia.com>

Author: realAsma

modelopt/torch/quantization/config.py

@@ -1227,6 +1227,16 @@ class QuantizeAlgorithmConfig(ModeloptBaseConfig):
         ),
     )
 
+    checkpoint_dir: str | None = ModeloptField(
+        default=None,
+        title="Checkpoint directory for sequential calibration.",
+        description=(
+            "If set together with use_sequential=True, per-layer checkpoints are saved to this "
+            "directory during calibration. On restart, calibration resumes from the last "
+            "completed layer."
+        ),
+    )
+
 
 class MaxCalibConfig(QuantizeAlgorithmConfig):
     """The config for max calibration algorithm.

modelopt/torch/quantization/mode.py

@@ -223,6 +223,7 @@ def wrapped_calib_func(
     kwargs = config.model_dump()
     method = kwargs.pop("method")
     sequential = kwargs.pop("use_sequential", False)
+    checkpoint_dir = kwargs.pop("checkpoint_dir", None)
     if method is not None and "awq" in method:
         # For backward compatibility
         kwargs["algorithm"] = method
@@ -240,14 +241,12 @@ def wrapped_calib_func(
         if sequential:
             if forward_loop is None:
                 raise ValueError("forward_loop is required for calibration but got None.")
-            assert method in ["max", "gptq"], (
-                f"Sequential calibration currently only supports max and gptq calibration, got {method}"
-            )
             # Wrap with sequential processing
             sequential_calibrate(
                 model,
                 forward_loop=forward_loop,
                 calib_func=func,
+                checkpoint_dir=checkpoint_dir,
                 **kwargs,
             )
         else:

modelopt/torch/quantization/model_calib.py

@@ -28,7 +28,7 @@
 from tqdm import tqdm
 
 from modelopt.torch.opt.searcher import ForwardLoop
-from modelopt.torch.quantization.utils.activation_collector import LayerActivationCollector
+from modelopt.torch.quantization.utils.layerwise_calib import LayerActivationCollector
 from modelopt.torch.utils import print_rank_0
 from modelopt.torch.utils.distributed import DistributedProcessGroup, ParallelState
 from modelopt.torch.utils.network import bind_forward_method, unpatch_forward_method
@@ -1563,7 +1563,15 @@ def sequential_calibrate(
     Runs the full model forward per layer but patches decoder layers with a
     skip / run / capture strategy so that inter-layer logic in parent modules
     (e.g. mask construction) executes naturally without model-specific hooks.
+
+    If ``checkpoint_dir`` is passed (via ``calib_kwargs``), per-layer checkpoints
+    are saved after each layer completes. On restart, calibration resumes from
+    the last completed layer.
     """
+    from modelopt.torch.quantization.utils.layerwise_calib import _CheckpointState
+
+    checkpoint_dir = calib_kwargs.pop("checkpoint_dir", None)
+
     if forward_loop is None:
         raise ValueError(
             "forward_loop must not be None for sequential calibration. "
@@ -1577,30 +1585,77 @@ def sequential_calibrate(
             "Sequential calibration requires a model with identifiable transformer layers."
         )
 
-    print_rank_0(f"Sequential calibration: Found {len(transformer_layers)} transformer layers")
+    num_layers = len(transformer_layers)
+    print_rank_0(f"Sequential calibration: Found {num_layers} transformer layers")
+
+    ckpt = _CheckpointState.from_folder(checkpoint_dir, num_layers)
+    start_layer = ckpt.start_layer if ckpt else 0
 
     input_getter = LayerActivationCollector(model)
     input_getter._patch_all_layers(decoder_layers=transformer_layers)
 
+    resumed_inputs = ckpt.setup_resume(transformer_layers) if ckpt and start_layer > 0 else None
+
     try:
-        for layer_idx, layer in enumerate(transformer_layers):
-            print_rank_0(f"Calibrating layer {layer_idx + 1}/{len(transformer_layers)}")
-            layer_inputs = input_getter.get_input_activations(layer, forward_loop)
+        for layer_idx, layer in enumerate(list(transformer_layers)):
+            if layer_idx < start_layer:
+                continue
+
+            layer_inputs = _get_layer_inputs(
+                layer_idx, start_layer, resumed_inputs, layer, input_getter, forward_loop
+            )
+            if ckpt:
+                ckpt.save_prev(transformer_layers, layer_inputs)
 
             def _layer_forward_loop(m, _inputs=layer_inputs):
                 for args, kwargs_input in _inputs:
                     m(*args, **kwargs_input)
 
             calib_func(layer, _layer_forward_loop, **calib_kwargs)
 
+            if ckpt:
+                ckpt.stash(layer_idx, layer, model)
+
             del layer_inputs
             torch.cuda.empty_cache()
+
+        if ckpt:
+            ckpt.save_last(transformer_layers)
     finally:
         input_getter._unpatch_all_layers()
 
+    if ckpt:
+        ckpt.full_restore(transformer_layers, model)
+
     print_rank_0("Sequential calibration completed")
 
 
+def _get_layer_inputs(
+    layer_idx: int,
+    start_layer: int,
+    resumed_inputs: list | None,
+    layer: nn.Module,
+    input_getter: LayerActivationCollector,
+    forward_loop: ForwardLoop,
+) -> list:
+    """Get inputs for a layer, using resumed_inputs for the first resumed layer."""
+    if layer_idx == start_layer and resumed_inputs is not None:
+        print_rank_0(f"Calibrating layer {layer_idx + 1} (resumed)")
+        # Manually set skip mode on all already-calibrated layers (output_meta
+        # was loaded by setup_resume). Don't call _set_layer_states which
+        # assumes the normal sequential progression with collected_inputs.
+        assert input_getter._decoder_layers is not None
+        for i in range(start_layer):
+            input_getter._swap_to_dummy(i)
+        # Seed collected_inputs so the next _set_layer_states call can
+        # transition this layer to "run" mode.
+        layer._seq_calib.collected_inputs = resumed_inputs
+        layer._seq_calib.mode = "original"
+        return resumed_inputs
+
+    return input_getter.get_input_activations(layer, forward_loop)
+
+
 @torch.no_grad()
 def gptq(
     model: nn.Module,
@@ -1663,8 +1718,10 @@ def gptq(
         handle.cleanup()
 
     print_rank_0("Updating weights using GPTQ algorithm...")
+    name_to_module = dict(model.named_modules())
     for handle in gptq_handles.values():
-        handle.update_weights(block_size, perc_damp)
+        with enable_weight_access_and_writeback(handle.module, model, name_to_module):
+            handle.update_weights(block_size, perc_damp)
         handle.free()
     del gptq_handles
 

modelopt/torch/quantization/plugins/accelerate.py

@@ -33,7 +33,7 @@
 
 def _get_cpu_offload_hook(hook):
     if isinstance(hook, AlignDevicesHook) and hook.offload and hook.weights_map is not None:
-        assert "weight" in hook.weights_map
+        assert len(hook.weights_map) > 0
         if (
             isinstance(hook.weights_map, PrefixedDataset)
             and hook.weights_map.prefix + "weight" not in hook.weights_map.dataset.state_dict
@@ -50,32 +50,68 @@ def _get_cpu_offload_hook(hook):
     return None
 
 
+def _writeback_params_to_weights_map(module, align_hook):
+    """Write all non-meta parameters back to the hook's CPU weights_map."""
+    for name, param in module.named_parameters():
+        if param.device.type == "meta":
+            continue
+        if isinstance(align_hook.weights_map, PrefixedDataset):
+            key = align_hook.weights_map.prefix + name
+            w_map = align_hook.weights_map.dataset.state_dict
+        else:
+            w_map = align_hook.weights_map
+            key = name
+        if key in w_map:
+            w_map[key] = param.data.to(w_map[key].device, dtype=w_map[key].dtype)
+
+
 @contextmanager
 def weight_access_and_writeback_context(module):
-    """Context manager for weight access and writeback for modules managed by accelerate."""
+    """Context manager for weight access and writeback for modules managed by accelerate.
+
+    Handles two cases:
+    1. **Single-module**: the module's own ``_hf_hook`` is an offload hook.
+    2. **Sub-module**: the module's hook is non-offloading, but its children have
+       offload hooks (common with ``SequentialHook`` on sub-modules placed by
+       ``load_checkpoint_and_dispatch``).
+
+    For the sub-module case, ``pre_forward`` is skipped on sub-modules whose weights
+    are already materialized (not on meta).  This allows the context manager to be
+    used as a pure writeback after weight-modifying algorithms.
+    """
     assert hasattr(module, "_hf_hook")
     align_hook = _get_cpu_offload_hook(module._hf_hook)
 
     if align_hook:
-        # Accelerate uses AlignDevicesHook to offload weights to CPU/Disk and then reload them in the forward pass
-        # The CPU/Disk offloaded weights are managed by PrefixDataset and OffloadedWeightsLoader
-        # See https://github.com/huggingface/accelerate/blame/f48d95c4939b281505a45b3d6e0bf554b65cc1ea/src/accelerate/utils/offload.py#L104-L141
-        # TODO: Add support for disk-offloaded models if needed (they will be really slow, hence low priority)
-
-        # This will load the weights from CPU state_dict and move it to the GPU from meta device
         align_hook.pre_forward(module)
+        try:
+            yield
+        finally:
+            _writeback_params_to_weights_map(module, align_hook)
+            align_hook.post_forward(module, None)
+        return
+
+    materialized: list[tuple[torch.nn.Module, AlignDevicesHook, bool]] = []
+    for mod in module.modules():
+        if mod is module or not hasattr(mod, "_hf_hook"):
+            continue
+        hook = _get_cpu_offload_hook(mod._hf_hook)
+        if hook is None:
+            continue
+        # Only call pre_forward if weights need materializing; already-materialized
+        # 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:
+            hook.pre_forward(mod)
+        materialized.append((mod, hook, needs_materialize))
+
     try:
         yield
     finally:
-        if align_hook:
-            # Update the weight in the CPU state_dict
-            if isinstance(align_hook.weights_map, PrefixedDataset):
-                key = align_hook.weights_map.prefix + "weight"
-                w_map = align_hook.weights_map.dataset.state_dict
-            else:
-                key, w_map = "weight", align_hook.weights_map
-            w_map[key] = module.weight.data.to(w_map[key].device, dtype=w_map[key].dtype)
-            align_hook.post_forward(module, None)
+        for mod, hook, was_materialized in materialized:
+            _writeback_params_to_weights_map(mod, hook)
+            if was_materialized:
+                hook.post_forward(mod, None)
 
 
 @contextmanager

modelopt/torch/quantization/plugins/huggingface.py

@@ -39,7 +39,7 @@
 from ..nn.modules.quant_linear import _QuantLinear
 from ..triton import IS_AVAILABLE as IS_TRITON_AVAILABLE
 from ..utils import replace_function, sync_moe_expert_amax
-from ..utils.activation_collector import LayerActivationCollector
+from ..utils.layerwise_calib import LayerActivationCollector
 from .attention import register_attention_for_kv_quant
 from .custom import CUSTOM_MODEL_PLUGINS, _ParallelLinear, _QuantFunctionalMixin
 

modelopt/torch/quantization/utils/__init__.py

@@ -16,8 +16,8 @@
 # ruff: noqa: F405
 """Quantization utilities."""
 
-from .activation_collector import LayerActivationCollector
 from .core_utils import *
+from .layerwise_calib import LayerActivationCollector
 
 __all__ = [
     "EXPORT_MODE",

modelopt/torch/quantization/utils/calib_utils.py

@@ -96,7 +96,7 @@ def __init__(self, module, name, offload_to_cpu=False):
         self.name = name
         in_features = module.weight.shape[-1]
         device = module.weight.device
-        if offload_to_cpu and get_used_gpu_mem_fraction(device) > 0.65:
+        if device.type == "meta" or (offload_to_cpu and get_used_gpu_mem_fraction(device) > 0.65):
             device = "cpu"
         self.hessian = torch.zeros(in_features, in_features, dtype=torch.float32, device=device)
         self.n_samples = 0

modelopt/torch/quantization/utils/core_utils.py

@@ -423,47 +423,70 @@ def _get_enclosing_fsdp_module(
         return root_model
 
 
+def _set_parameter(module: nn.Module, name: str, value: nn.Parameter):
+    """Set a parameter on a module by dotted name (e.g. ``self_attn.q_proj.weight``)."""
+    parts = name.rsplit(".", 1)
+    if len(parts) == 2:
+        parent = module.get_submodule(parts[0])
+        attr = parts[1]
+    else:
+        parent = module
+        attr = name
+    parent._parameters[attr] = value
+
+
 @contextmanager
 def fsdp2_weight_access_and_writeback_context(module: nn.Module, root_model: nn.Module):
     """Context manager for FSDP2 weight access and writeback.
 
-    Note this context will gather the weight across FSDP/HSDP shards. If TP is implemented with DTensor,
-    the weight will be a local tensor of the TP DTensor under this context.
+    Gathers sharded DTensor parameters across FSDP/HSDP shards so they can be
+    read or modified. Works for both leaf modules (single ``weight``) and
+    composite modules like decoder layers (all ``named_parameters``).
+
+    If TP is implemented with DTensor, the weight will be a local tensor of the
+    TP DTensor under this context.
     """
     assert isinstance(root_model, torch.distributed.fsdp.FSDPModule), "We only support FSDP2"
 
     assert not hasattr(module, "_hf_hook"), "We dont support FSDP2 with HF accelerate hooks"
-    assert isinstance(module.weight, torch.distributed.tensor.DTensor)
     fsdp_module = _get_enclosing_fsdp_module(module, root_model)
     assert fsdp_module is not None, "Module is not wrapped by FSDP"
     fsdp_device_mesh = _get_fsdp2_mesh(fsdp_module)
     fsdp_dim = fsdp_device_mesh.ndim
 
-    original_placements = module.weight.placements
-    original_device_mesh = module.weight.device_mesh
-    original_weight = module.weight
-    # Assuming the first fsdp_dim dimensions are for FSDP/HSDP, we only collect the tensor over FSDP/HSDP dimension,
-    # the TP will be handled by the TP reduction.
-    if fsdp_dim != original_device_mesh.ndim:
-        assert fsdp_device_mesh.mesh_dim_names == original_device_mesh.mesh_dim_names[:fsdp_dim], (
-            "FSDP2 mesh should be a slice of DTesnor's device mesh."
+    # Collect all DTensor parameters, replacing them with local replicated copies.
+    originals: dict[str, tuple] = {}
+    for name, param in module.named_parameters():
+        if not isinstance(param, torch.distributed.tensor.DTensor):
+            continue
+        original_placements = param.placements
+        original_device_mesh = param.device_mesh
+        if fsdp_dim != original_device_mesh.ndim:
+            assert (
+                fsdp_device_mesh.mesh_dim_names == original_device_mesh.mesh_dim_names[:fsdp_dim]
+            ), "FSDP2 mesh should be a slice of DTensor's device mesh."
+        collected = param.redistribute(
+            placements=[Replicate()] * fsdp_dim + list(original_placements[fsdp_dim:]),
+            device_mesh=original_device_mesh,
         )
-
-    weight_collected = original_weight.redistribute(
-        placements=[Replicate()] * fsdp_dim + list(original_placements[fsdp_dim:]),
-        device_mesh=original_device_mesh,
-    )
-    new_weight = nn.Parameter(weight_collected.to_local())
-    module._parameters["weight"] = new_weight
+        originals[name] = (param, collected, original_placements, original_device_mesh)
+        _set_parameter(module, name, nn.Parameter(collected.to_local()))
 
     yield
 
-    original_weight.to_local().data.copy_(
-        weight_collected.redistribute(
-            placements=original_placements, device_mesh=original_device_mesh
-        ).to_local()
-    )
-    module._parameters["weight"] = original_weight
+    # Write back and restore original DTensor parameters.
+    for name, (
+        original_param,
+        collected,
+        original_placements,
+        original_device_mesh,
+    ) in originals.items():
+        original_param.to_local().data.copy_(
+            collected.redistribute(
+                placements=original_placements, device_mesh=original_device_mesh
+            ).to_local()
+        )
+        _set_parameter(module, name, original_param)
 
 
 @contextmanager