fix(examples): enable activation checkpointing for phi_4_squad

examples/llm_finetune/phi/phi_4_squad.yaml

@@ -49,6 +49,10 @@ checkpoint:
 
 distributed:
   strategy: fsdp2
+  # phi-4 (14B) stores ~60 GiB of activations per training step; without
+  # recomputation a single fwd/bwd peaks at ~74 GiB and OOMs on long batches
+  # (notably the checkpoint-robustness resume phase, which runs extra steps).
+  activation_checkpointing: true
   dp_size: none
   tp_size: 1
   cp_size: 1

nemo_automodel/_transformers/capabilities.py

@@ -29,6 +29,7 @@
 import functools
 import inspect
 import logging
+import weakref
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -158,13 +159,24 @@ class ModelSupports:
         model.supports.pp   # ...
     """
 
-    __slots__ = ("_model", "_model_cls", "_mesh")
+    __slots__ = ("_model_ref", "_model_cls", "_mesh")
 
     def __init__(self, model: "nn.Module", mesh: "MeshContext | None" = None) -> None:
-        self._model = model
+        # Hold the model weakly. ``ModelSupports`` is attached back onto the model
+        # as ``model._supports``; a strong reference here would form a
+        # ``model <-> _supports`` cycle, so the capability descriptor must never be
+        # the reason a (multi-GiB) model stays resident after its owner is dropped.
+        self._model_ref = weakref.ref(model)
         self._model_cls = type(model)
         self._mesh = mesh
 
+    @property
+    def _model(self) -> "nn.Module":
+        model = self._model_ref()
+        if model is None:
+            raise ReferenceError("ModelSupports: underlying model has been garbage-collected")
+        return model
+
     def __repr__(self) -> str:
         names = (
             "tp",

tests/functional_tests/checkpoint_robustness/test_checkpoint_robustness_llm.py

@@ -362,6 +362,36 @@ def _barrier():
         dist.barrier()
 
 
+def _release_recipe_memory(recipe) -> None:
+    """Release a recipe's GPU-resident state between checkpoint-robustness phases.
+
+    Each phase builds a full FSDP2 model + optimizer. A bare ``del`` is not
+    enough: the per-part optimizers are reachable from the model (they are built
+    over ``model.parts``), so the optimizer state (Adam moments are the bulk)
+    lingers. Clear the optimizer state in place and drop the recipe's references,
+    then collect — letting the prior phase's model + optimizer be reclaimed
+    before the next phase allocates its own, keeping the inter-phase baseline low.
+    """
+    if recipe is None:
+        return
+    optimizers = getattr(recipe, "optimizer", None)
+    if not isinstance(optimizers, (list, tuple)):
+        optimizers = [optimizers] if optimizers is not None else []
+    for opt in optimizers:
+        try:
+            opt.state.clear()
+            opt.param_groups.clear()
+        except Exception:
+            pass
+    recipe.model_parts = None
+    recipe.optimizer = None
+    if getattr(recipe, "lr_scheduler", None) is not None:
+        recipe.lr_scheduler = None
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
 def test_checkpoint_robustness():
     """Train -> checkpoint -> reload automodel from consolidated -> reload vanilla HF, compare logits."""
     custom_args, config_argv = _extract_custom_args(sys.argv[1:])
@@ -438,9 +468,8 @@ def test_checkpoint_robustness():
     else:
         original_quantization_config = _raw_qc
 
+    _release_recipe_memory(trainer)
     del trainer
-    gc.collect()
-    torch.cuda.empty_cache()
 
     # Phantom key check: scan consolidated safetensors for leaked quantization keys
     if check_phantom_keys and _rank0():
@@ -493,9 +522,8 @@ def test_checkpoint_robustness():
     )
 
     # Phase 4: Load into vanilla HF (rank 0 only)
+    _release_recipe_memory(restored_trainer)
     del restored_trainer
-    gc.collect()
-    torch.cuda.empty_cache()
     _barrier()  # ensure all ranks free memory before rank 0 loads HF model
 
     if skip_hf_reload:
@@ -652,9 +680,8 @@ def test_checkpoint_robustness():
             f"max per-token KL = {max_kl_cross_tp:.6e} > threshold {cross_tp_kl_threshold:.6e}"
         )
 
+        _release_recipe_memory(cross_tp_trainer)
         del cross_tp_trainer
-        gc.collect()
-        torch.cuda.empty_cache()
         _barrier()
 
     # Phase 6 (optional): Training resumption — verify loss continuity
@@ -693,9 +720,8 @@ def test_checkpoint_robustness():
                     if entry["step"] >= original_max_steps:
                         baseline_losses[entry["step"]] = entry["loss"]
 
+        _release_recipe_memory(baseline_trainer)
         del baseline_trainer
-        gc.collect()
-        torch.cuda.empty_cache()
         shutil.rmtree(baseline_dir, ignore_errors=True)
 
         # Resume: reload from Phase 1 checkpoint and train to resume_max_steps.
@@ -739,9 +765,8 @@ def test_checkpoint_robustness():
             )
             print(f"[Phase 6] Training resumption verified ({matched_steps} steps compared) ✓")
 
+        _release_recipe_memory(resume_trainer)
         del resume_trainer
-        gc.collect()
-        torch.cuda.empty_cache()
         _barrier()
 
     # Skip the atexit-registered destroy_process_group() call. MoE models with expert

tests/unit_tests/_transformers/test_capabilities_magi.py

@@ -59,31 +59,41 @@ def test_uses_magi_attention_no_backend():
 # supports_cp / supports_sequence_packing / supports_cp_with_sequence_packing
 # --------------------------------------------------------------------------- #
 def _supports(attn, cp_size=1):
+    # ModelSupports holds the model weakly (in production the model owns it as
+    # ``model._supports``), so the caller must keep ``model`` alive for the
+    # duration of the capability check -- return it alongside.
+    model = _BackendModel(attn)
     mesh = SimpleNamespace(cp_size=cp_size)
-    return ModelSupports(_BackendModel(attn), mesh)
+    return model, ModelSupports(model, mesh)
 
 
 def test_supports_cp_admits_magi():
-    assert _supports("magi").supports_cp is True
+    model, supports = _supports("magi")
+    assert supports.supports_cp is True
 
 
 def test_supports_cp_rejects_flex_backend():
     """Regression: the gate was not broadened to every custom backend."""
-    assert _supports("flex").supports_cp is False
+    model, supports = _supports("flex")
+    assert supports.supports_cp is False
 
 
 def test_supports_sequence_packing_admits_magi():
-    assert _supports("magi").supports_sequence_packing is True
+    model, supports = _supports("magi")
+    assert supports.supports_sequence_packing is True
 
 
 def test_supports_cp_with_sequence_packing_admits_magi_at_cp2():
-    assert _supports("magi", cp_size=2).supports_cp_with_sequence_packing is True
+    model, supports = _supports("magi", cp_size=2)
+    assert supports.supports_cp_with_sequence_packing is True
 
 
 def test_supports_cp_with_sequence_packing_rejects_flex_at_cp2():
-    assert _supports("flex", cp_size=2).supports_cp_with_sequence_packing is False
+    model, supports = _supports("flex", cp_size=2)
+    assert supports.supports_cp_with_sequence_packing is False
 
 
 def test_supports_cp_with_sequence_packing_cp1_falls_back_to_packing():
     # at cp_size<=1 it reduces to plain sequence-packing support (magi qualifies).
-    assert _supports("magi", cp_size=1).supports_cp_with_sequence_packing is True
+    model, supports = _supports("magi", cp_size=1)
+    assert supports.supports_cp_with_sequence_packing is True