feat(distributed): add selective activation checkpointing for FSDP2 by yuhezhang-ai · Pull Request #2389 · NVIDIA-NeMo/Automodel

yuhezhang-ai · 2026-06-02T15:44:25Z

Summary

Adds a TorchTitan-style selective activation-checkpointing mode for FSDP2 — dense and Mixture-of-Experts (including expert parallelism), single- and multi-GPU — with torch.compile interoperability.

New selective mode: wraps whole transformer blocks with a per-op policy — matmuls (mm/linear) and MoE expert grouped-GEMMs (_grouped_mm/_scaled_grouped_mm) alternate save/recompute, while attention, communication collectives, and other compute-intensive ops are always saved. Configured via distributed.activation_checkpointing: selective.
Correctness: the matmul counter is keyed on ctx.is_recompute so forward and recompute agree even when a block has an odd number of matmuls (a shared counter silently corrupts gradients).
Op set: seeded from PyTorch's get_default_op_list().compute_intensive_ops plus a curated supplement (attention variants, comm collectives incl. expert-parallel dispatch/combine, the inductor compiled HOP), with a guarded fallback when the private partitioner API is unavailable — rather than a frozen hand-written list.
MoE + expert parallelism: selective AC is now supported for MoE under FSDP2/EP. MoE blocks reuse the shared dense policy, and DeepSeek MLP dispatch was made wrapper-safe so checkpointed MoE blocks work. Note: selective generally does not beat full AC for MoE (experts are cheap to recompute but expensive to store, and EP communication is opaque to the policy), so full AC is recommended for MoE; selective remains an opt-in.
Single-GPU: selective is honored on world_size == 1 instead of silently degrading to full checkpointing.
KV-shared models (e.g. Gemma4 E2B/E4B) cannot checkpoint attention through the DynamicCache, so they fall back to sub-module checkpointing.
torch.compile: selective wrappers are tagged and compiled outer (dense and MoE) so the partitioner honors the SAC recompute tags (compiling the inner layer would collapse selectivity into full recompute); the dynamo LRU cache is disabled (best-effort) for the SAC + compile + PP path.
Diagnostics: set NEMO_SELECTIVE_AC_TRACE=1 to log each op's SAVE/RECOMPUTE/ALTERNATE decision once, to confirm a model's expensive ops are actually being saved.
Code organization: the selective-AC core (op sets, policy, block/sub-module wrappers, KV detection, compile flag) lives in nemo_automodel/components/distributed/activation_checkpointing.py; parallelizer.py keeps only the thin apply_selective_activation_checkpointing entry point, and the opt-in trace helper lives in parallelizer_utils.py.
Validation: selective requires FSDP2 (non-FSDP2 strategies such as ddp/megatron_fsdp are rejected); expert parallelism is allowed.
Docs: the gradient-checkpointing guide (legacy + nightly) documents the mode, single-GPU/MoE behavior, the full-vs-selective trade-off, and the trace flag.

Test plan

Forward/backward numerics parity vs. non-checkpointed baseline (odd and even matmul counts)
Save-op-set composition + partitioner seeding + fallback; grouped-GEMM alternation in the matmul set
Compile-target selection (selective wrapper compiled outer; PP wrapper unwrapped) + LRU-cache guard
Single-GPU selective wrapping + KV-shared fallback
MoE: parallelize_model threads selective; apply_ac wraps blocks with the shared policy and sets the compile-outer flag
Policy trace logs each op once with its verdict
YAML routing/validation for selective/full/true, EP allowed, non-FSDP2 rejected
distributed + moe + recipes unit suites green (197 passed locally)
L1/functional multi-GPU EP parity for selective AC (and under torch.compile) — inductor/cluster-dependent, not covered by CPU unit tests

Real runs:

Dense model qwen3-4b with real data (ac full vs selective, compiled vs uncompiled): https://wandb.ai/Nemo-automodel/yuhez_workspace/groups/qwen3-4b-squad-sdpa/workspace?nw=nwuseryuhez
selective is 5-10% faster.
MOE with EP (where selective is slower in practice): https://wandb.ai/Nemo-automodel/yuhez_workspace/groups/moonlight16b-pretrain-te-ep8-ac-compare/workspace?nw=nwuseryuhez

Better benchmark with mock data (so seq length steady) with llama3.1-8b:

https://wandb.ai/Nemo-automodel/yuhez_workspace/groups/llama8b-ac-matrix-20260604/workspace?nw=nwuseryuhez
W&B-logged matrix (llama-3.1-8b, train_ft, 50 steps, mock, H100x8, seq4096 bs1)

AC mode	eager tps/gpu	compile tps/gpu	W&B (eager / compile)
none	9,002	9,220	xw4bhe4p / 3xl2ybsv
full	7,961	8,282	8zkn0g8i / snarikge
selective	8,352	8,349	xk0zoths / km144k84

https://wandb.ai/Nemo-automodel/yuhez_workspace/groups/llama8b-batchsweep-20260604/workspace?nw=nwuseryuhez
Batch-size sweep -- llama-3.1-8b, eager, seq4096

lbs (gbs)	full tps/gpu	selective tps/gpu	sel vs full	full mem	sel mem
1 (8)	7,961	8,352	+4.9%	16.0	19.8
2 (16)	9,801	10,326	+5.4%	24.0	31.5
4 (32)	10,665	11,406	+6.9%	40.0	55.0
8 (64)	OOM (~71GB)	OOM (~74GB)	-	-	-

copy-pr-bot · 2026-06-02T15:44:29Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

…ch.compile Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

… GPU Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

yuhezhang-ai · 2026-06-02T16:31:28Z

/ok to test e390fd2

…allelism Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

… trace Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

Extract the TorchTitan-style selective activation checkpointing core out of the central parallelizer.py into a dedicated activation_checkpointing.py: op-set construction, the save/recompute policy, block/sub-module wrappers, KV-sharing detection, and the compile-outer wrapper flag. parallelizer.py keeps only the thin apply_selective_activation_checkpointing entry point, which still needs the heavy, transformers-aware _extract_model_layers, so the dependency stays one-directional (parallelizer -> activation_checkpointing -> parallelizer_utils) with no circular imports. Move the opt-in NEMO_SELECTIVE_AC_TRACE diagnostic out of parallelizer.py into parallelizer_utils.maybe_trace_selective_ac_decision so the hot policy is a single call site instead of trace globals plus a helper. Make the new module's cross-module interface public (drop the leading underscore) and keep internal op-resolution/plumbing private. Update the moe and fsdp2 consumers and the unit tests to import from the new module. Also fix doc wording: clarify that torch.compile must be held fixed when comparing full vs. selective, and refer to TorchTitan as a reference implementation rather than "upstream". Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

yuhezhang-ai · 2026-06-04T15:09:18Z

/ok to test 74d320b

…after AC module split Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

yuhezhang-ai · 2026-06-04T17:10:41Z

/ok to test f04335c

jgerh

Completed tech pubs review of .md and .mdx files and provided a few suggestions

jgerh · 2026-06-08T17:35:44Z

Suggested change

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.

jgerh · 2026-06-08T17:36:53Z

+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.


Suggested change

`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.

`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.

jgerh · 2026-06-08T17:37:17Z

+</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.


Suggested change

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.

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.

jgerh · 2026-06-08T17:39:50Z

+</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.


Suggested change

**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.

**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.

jgerh · 2026-06-08T17:40:57Z

+</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+FSDP2 as an opt-in.


Suggested change

**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+FSDP2 as an opt-in.

**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.

jgerh · 2026-06-08T17:52:22Z

+
+> **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.


Suggested change

> **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 (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.

jgerh · 2026-06-08T17:52:39Z

+
+> **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+FSDP2 as an opt-in.


Suggested change

> **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+FSDP2 as an opt-in.

> **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.

jgerh · 2026-06-08T17:52:58Z

Suggested change

LC-CE and gradient checkpointing target **different memory hot-spots** (output layer vs. transformer blocks), so their benefits stack almost linearly.

jgerh · 2026-06-08T17:53:30Z

Suggested change

- Expect ±5 % variance depending on exact model, sequence length, and GPU architecture.

jgerh · 2026-06-08T17:53:43Z

@@ -82,4 +101,4 @@ 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:


Suggested change

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:

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

…activation-checkpointing Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

yuhezhang-ai · 2026-06-08T19:50:11Z

/ok to test ded10ca

akoumpa · 2026-06-08T22:10:06Z

ci passed https://github.com/NVIDIA-NeMo/Automodel/actions/runs/27162875997/job/80200006070?pr=2389

yuhezhang-ai added 2 commits June 2, 2026 11:48

feat(distributed): add selective activation checkpointing for FSDP2

cb4d002

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

fix(distributed): support selective activation checkpointing with tor…

18d6870

…ch.compile Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

yuhezhang-ai force-pushed the yuhez/feat/selective-activation-checkpointing branch from eb01d92 to 18d6870 Compare June 2, 2026 15:49

yuhezhang-ai added 2 commits June 2, 2026 11:59

docs(fern): drop selective AC from frozen v0.4 snapshot

73fcda3

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

feat(distributed): honor selective activation checkpointing on single…

e390fd2

… GPU Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

copy-pr-bot Bot temporarily deployed to test June 2, 2026 16:31 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci June 2, 2026 16:32 Inactive

copy-pr-bot Bot temporarily deployed to public June 2, 2026 16:32 Inactive

copy-pr-bot Bot temporarily deployed to public June 2, 2026 16:38 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci June 2, 2026 16:40 Inactive

copy-pr-bot Bot temporarily deployed to public June 2, 2026 16:44 Inactive

yuhezhang-ai added 5 commits June 3, 2026 14:54

feat(moe): support selective activation checkpointing with expert par…

5bd754c

…allelism Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

fix(model): make DeepSeek MLP dispatch wrapper-safe

286be84

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

fix(distributed): save expert grouped-GEMM in selective AC and add op…

0e90101

… trace Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

feat(moe): compile selective activation checkpointing wrappers outer

8bb8595

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

copy-pr-bot Bot had a problem deploying to test June 4, 2026 15:09 Error

copy-pr-bot Bot temporarily deployed to nemo-ci June 4, 2026 15:09 Inactive

copy-pr-bot Bot temporarily deployed to public June 4, 2026 16:16 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci June 4, 2026 16:18 Inactive

copy-pr-bot Bot temporarily deployed to public June 4, 2026 16:23 Inactive

test(distributed): patch activation_checkpointing.checkpoint_wrapper …

f04335c

…after AC module split Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

copy-pr-bot Bot temporarily deployed to nemo-ci June 4, 2026 17:11 Inactive

copy-pr-bot Bot temporarily deployed to public June 4, 2026 17:11 Inactive

copy-pr-bot Bot temporarily deployed to test June 4, 2026 17:11 Inactive

copy-pr-bot Bot temporarily deployed to public June 4, 2026 17:14 Inactive

copy-pr-bot Bot temporarily deployed to nemo-ci June 4, 2026 17:16 Inactive

jgerh reviewed Jun 8, 2026

View reviewed changes

yuhezhang-ai added 3 commits June 8, 2026 15:01

docs: apply tech-writer edits to gradient-checkpointing guide

669484f

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

Merge remote-tracking branch 'origin/main' into yuhez/feat/selective-…

e31b35a

…activation-checkpointing Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

Merge remote-tracking branch 'origin/main' into yuhez/feat/selective-…

ded10ca

…activation-checkpointing Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

akoumpa approved these changes Jun 8, 2026

View reviewed changes


	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.

	`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.
	`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 (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+FSDP2 as an opt-in.


	LC-CE and gradient checkpointing target different memory hot-spots (output layer vs. transformer blocks), so their benefits stack almost linearly.


	- Expect ±5 % variance depending on exact model, sequence length, and GPU architecture.

		@@ -82,4 +101,4 @@ 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 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:

Uh oh!

Conversation

yuhezhang-ai commented Jun 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Uh oh!

copy-pr-bot Bot commented Jun 2, 2026

Uh oh!

yuhezhang-ai commented Jun 2, 2026

Uh oh!

yuhezhang-ai commented Jun 4, 2026

Uh oh!

yuhezhang-ai commented Jun 4, 2026

Uh oh!

jgerh left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yuhezhang-ai commented Jun 8, 2026

Uh oh!

akoumpa commented Jun 8, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

yuhezhang-ai commented Jun 2, 2026 •

edited

Loading