feat(distributed): add selective activation checkpointing for FSDP2

docs/guides/gradient-checkpointing.md

@@ -1,6 +1,6 @@
-# Gradient (Activation) Checkpointing
+# Use Gradient (Activation) Checkpointing
 
-Gradient checkpointing, also called _activation checkpointing_, trades a little extra compute for a **large reduction in GPU memory** by recomputing intermediate activations during the backwards pass instead of storing them.  
+Gradient checkpointing, also called _activation checkpointing_, trades a little extra compute for a **large reduction in GPU memory** by recomputing intermediate activations during the backward pass instead of storing them.  
 It is especially powerful when combined with memory-efficient loss functions (e.g., Linear-Cut Cross-Entropy) and parameter sharding using FSDP.
 
 ## Enable Gradient Checkpointing
@@ -23,12 +23,31 @@ distributed:
   ...
 ```
 
+For FSDP2, `activation_checkpointing` also accepts explicit policy strings:
+
+```yaml
+distributed:
+  strategy: fsdp2
+  activation_checkpointing: selective
+```
+
+Use `true` or `full` for full activation checkpointing. Use `selective` for PyTorch selective activation checkpointing on FSDP2 configs. Selective checkpointing saves expensive operations such as attention, collectives, and part of the matrix multiplications while recomputing cheaper operations during backward.
+
+> **Note:** `selective` requires the FSDP2 strategy. Non-FSDP2 strategies (`ddp`, `megatron_fsdp`) raise an error when `selective` is requested. KV-sharing models (e.g., Gemma4) automatically fall back to sub-module checkpointing, because attention cannot be recomputed through the KV cache.
+
+> **Tip:** Selective AC only speeds things up when the model's expensive operations are the ones being saved. To see the per-op save/recompute decisions for your model, set `NEMO_SELECTIVE_AC_TRACE=1`; each unique operation is logged once as `SAVE`, `RECOMPUTE`, or `ALTERNATE`. If an expensive op (e.g., an expert grouped-GEMM) shows up as `RECOMPUTE`, selective AC will not beat full checkpointing for that model.
+
+> **Note (full vs. selective):** Selective AC saves the expensive operations (attention and part of the matmuls) and recomputes only the cheaper ones, so it does less recompute work than full AC while holding more activations in memory. Whether that nets out as faster, and at what memory cost, depends on the model, sequence length, and whether `torch.compile` is enabled, so benchmark full vs. selective for your own setup. When you do, keep the `torch.compile` setting the same on both sides (compare full and selective both compiled, or both uncompiled). `torch.compile` is a large speed lever on its own and helps both modes, so mixing it in makes it hard to tell which gain came from the AC mode.
+
+> **Note (MoE/expert parallelism):** Selective AC is designed for dense transformers and generally does **not** help Mixture-of-Experts models with expert parallelism. In an MoE block the experts dominate the cost (they are cheap to recompute but expensive to store), and the expert-parallel dispatch/communication is opaque to the selective policy, so it is recomputed regardless. As a result, selective AC tends to add activation memory without a corresponding speedup for MoE, matching what reference implementations such as TorchTitan observe. Prefer **full** activation checkpointing (`true`/`full`) for MoE; selective remains supported for MoE and FSDP2 as an opt-in.
+
 ### Configure Programmatically
 ```python
 from nemo_automodel.components.distributed.config import FSDP2Config
 from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager
 
 config = FSDP2Config(activation_checkpointing=True)
+# Use activation_checkpointing="selective" for FSDP2 selective checkpointing.
 # device_mesh is created elsewhere (e.g. by the recipe via setup_distributed)
 manager = FSDP2Manager(config, device_mesh=device_mesh, moe_mesh=moe_mesh)
 model = manager.parallelize(model)
@@ -48,7 +67,7 @@ loss_fn:
   _target_: nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy
 ```
 
-LC-CE and gradient checkpointing target **different memory hot-spots** (output layer vs. transformer blocks) so their benefits stack almost linearly.
+LC-CE and gradient checkpointing target **different memory hot-spots** (output layer vs. transformer blocks), so their benefits stack almost linearly.
 
 ## Example Memory Savings (H100-80GB, Llama-3.2-1B)
 | Technique | Max GPU Mem (GB) | Δ vs Baseline |
@@ -62,7 +81,7 @@ LC-CE and gradient checkpointing target **different memory hot-spots** (output l
 :::{note}
 - Measurements taken with local batch size = 8, sequence len = 2048, AdamW, PyTorch 2.8.
 - Peak memory reported by `torch.cuda.max_memory_allocated()` averaged across DP ranks.
-- Expect ±5 % variance depending on exact model, sequence length and GPU architecture.
+- Expect ±5 % variance depending on exact model, sequence length, and GPU architecture.
 :::
 
 ## Performance Considerations
@@ -79,7 +98,7 @@ automodel --nproc-per-node=8 examples/llm_finetune/llama3_2/llama_3_2_1b_my_fine
 # If running on 1x GPU
 automodel examples/llm_finetune/llama3_2/llama_3_2_1b_my_finetune.yaml
 ```
-If we run with the above settings (activation ckpt = on, lc-ce = on, fsdp = on), look for a log line similar to:
+If you run with the above settings (activation ckpt = on, lc-ce = on, fsdp = on), look for a log line similar to:
 ```
 ... | mem 7.30 GiB | ...
-```
+```

fern/versions/nightly/pages/guides/gradient-checkpointing.mdx

@@ -1,9 +1,9 @@
 ---
-title: "Gradient (Activation) Checkpointing in NeMo AutoModel"
+title: "Use Gradient (Activation) Checkpointing"
 description: ""
 position: 2
 ---
-Gradient checkpointing, also called _activation checkpointing_, trades a little extra compute for a **large reduction in GPU memory** by recomputing intermediate activations during the backwards pass instead of storing them.  
+Gradient checkpointing, also called _activation checkpointing_, trades a little extra compute for a **large reduction in GPU memory** by recomputing intermediate activations during the backward pass instead of storing them.  
 It is especially powerful when combined with memory-efficient loss functions (e.g., Linear-Cut Cross-Entropy) and parameter sharding using FSDP.
 
 ## Enable Gradient Checkpointing
@@ -26,12 +26,39 @@ distributed:
   ...
 ```
 
+For FSDP2, `activation_checkpointing` also accepts explicit policy strings:
+
+```yaml
+distributed:
+  strategy: fsdp2
+  activation_checkpointing: selective
+```
+
+Use `true` or `full` for full activation checkpointing. Use `selective` for PyTorch selective activation checkpointing on FSDP2 configs. Selective checkpointing saves expensive operations such as attention, collectives, and part of the matrix multiplications while recomputing cheaper operations during backward.
+
+<Note>
+`selective` requires the FSDP2 strategy. Non-FSDP2 strategies (`ddp`, `megatron_fsdp`) raise an error when `selective` is requested. KV-sharing models (e.g., Gemma4) automatically fall back to sub-module checkpointing, because attention cannot be recomputed through the KV cache.
+</Note>
+
+<Tip>
+Selective AC only speeds things up when the model's expensive operations are the ones being saved. To see the per-op save/recompute decisions for your model, set `NEMO_SELECTIVE_AC_TRACE=1`; each unique operation is logged once as `SAVE`, `RECOMPUTE`, or `ALTERNATE`. If an expensive op (e.g., an expert grouped-GEMM) shows up as `RECOMPUTE`, selective AC will not beat full checkpointing for that model.
+</Tip>
+
+<Note>
+**Full vs. selective:** Selective AC saves the expensive operations (attention and part of the matmuls) and recomputes only the cheaper ones, so it does less recompute work than full AC while holding more activations in memory. Whether that nets out as faster, and at what memory cost, depends on the model, sequence length, and whether `torch.compile` is enabled, so benchmark full vs. selective for your own setup. When you do, keep the `torch.compile` setting the same on both sides (compare full and selective both compiled, or both uncompiled). `torch.compile` is a large speed lever on its own and helps both modes, so mixing it in makes it hard to tell which gain came from the AC mode.
+</Note>
+
+<Note>
+**MoE/expert parallelism:** Selective AC is designed for dense transformers and generally does **not** help Mixture-of-Experts models with expert parallelism. In an MoE block the experts dominate the cost (they are cheap to recompute but expensive to store), and the expert-parallel dispatch/communication is opaque to the selective policy, so it is recomputed regardless. As a result, selective AC tends to add activation memory without a corresponding speedup for MoE, matching what reference implementations such as TorchTitan observe. Prefer **full** activation checkpointing (`true`/`full`) for MoE; selective remains supported for MoE and FSDP2 as an opt-in.
+</Note>
+
 ### Configure Programmatically
 ```python
 from nemo_automodel.components.distributed.config import FSDP2Config
 from nemo_automodel.components.distributed.fsdp2 import FSDP2Manager
 
 config = FSDP2Config(activation_checkpointing=True)
+# Use activation_checkpointing="selective" for FSDP2 selective checkpointing.
 # device_mesh is created elsewhere (e.g. by the recipe via setup_distributed)
 manager = FSDP2Manager(config, device_mesh=device_mesh, moe_mesh=moe_mesh)
 model = manager.parallelize(model)
@@ -51,7 +78,7 @@ loss_fn:
   _target_: nemo_automodel.components.loss.linear_ce.FusedLinearCrossEntropy
 ```
 
-LC-CE and gradient checkpointing target **different memory hot-spots** (output layer vs. transformer blocks) so their benefits stack almost linearly.
+LC-CE and gradient checkpointing target **different memory hot-spots** (output layer vs. transformer blocks), so their benefits stack almost linearly.
 
 ## Example Memory Savings (H100-80GB, Llama-3.2-1B)
 | Technique | Max GPU Mem (GB) | Δ vs Baseline |
@@ -65,7 +92,7 @@ LC-CE and gradient checkpointing target **different memory hot-spots** (output l
 <Note>
 - Measurements taken with local batch size = 8, sequence len = 2048, AdamW, PyTorch 2.8.
 - Peak memory reported by `torch.cuda.max_memory_allocated()` averaged across DP ranks.
-- Expect ±5 % variance depending on exact model, sequence length and GPU architecture.
+- Expect ±5 % variance depending on exact model, sequence length, and GPU architecture.
 
 </Note>
 
@@ -83,7 +110,7 @@ automodel --nproc-per-node=8 examples/llm_finetune/llama3_2/llama_3_2_1b_my_fine
 # If running on 1x GPU
 automodel examples/llm_finetune/llama3_2/llama_3_2_1b_my_finetune.yaml
 ```
-If we run with the above settings (activation ckpt = on, lc-ce = on, fsdp = on), look for a log line similar to:
+If you run with the above settings (activation ckpt = on, lc-ce = on, fsdp = on), look for a log line similar to:
 ```
 ... | mem 7.30 GiB | ...
 ```

nemo_automodel/_transformers/infrastructure.py

@@ -261,7 +261,7 @@ def instantiate_infrastructure(
     pipeline_config: Optional[PipelineConfig] = None,
     qat_config: Optional[QATConfig] = None,
     moe_config: Optional[MoEParallelizerConfig] = None,
-    activation_checkpointing: bool = False,
+    activation_checkpointing: bool | str = False,
     device: Optional[torch.device] = None,
     mesh: Optional[MeshContext] = None,
     # Deprecated -- prefer passing ``mesh`` directly