[graph_trainer] Add torch.no_grad() and graph-based SAC to traced execution

[tugsbayasgalan]

Stack from ghstack (oldest at bottom):

• -> [graph_trainer] Add torch.no_grad() and graph-based SAC to traced execution #2766

Execute traced graph under torch.no_grad() since it already contains
explicit backward ops (from torch.autograd.grad traced by make_fx).
Without this, PyTorch builds a redundant autograd graph keeping all
forward intermediates alive via grad_fn references.

Adds SAC option on aot_fx_trace on GraphTrainer.

Adds option to dump the peak memory from CUDA caching allocator.

Adds similar test as test_bitwise_equivalency.py for peak memory

Here is the run comparison between eager and aot_fx_trace with SAC

Llama3B 8b with dp=4 tp=2 on 10 step run before SAC with cuda graph seq_len=1024

Llama3B 8b with dp=4 tp=2 on 10 step run after SAC with cuda graph seq_len=1024

Llama3B 8b with dp=4 tp=2 on 10 step run before SAC without cuda graph with seq_len=8192

Llama3B 8b with dp=4 tp=2 on 10 step run after SAC without cuda graph with seq_len=8192

Deepseek3 16B before SAC

Deepseek3 16B after SAC

[SherlockNoMad]

No, user should apply this if they want.

keep the tracing simple.

[tugsbayasgalan]

Isn't it kinda weird that user annotates nodes with AC and don't see any improvement in the end? It feels like silent incorrectness to me.

[SherlockNoMad]

again, remove this. should be a post processing pass in user's control

[claude]

Claude finished @tugsbayasgalan's task in 2m 1s —— View job

---

CI Failure Analysis

• Check CI status
• Analyze failing jobs
• Report findings

---

There are 4 failed jobs out of 14. Lint and CPU Unit Test both passed. 8 jobs are still queued/in-progress.

---

Failure 1: GraphTrainer 8 GPU Integration Tests (both runs)

Error: AttributeError: module 'torch.compiler' has no attribute '_patch_autograd_grad'

Location: make_fx_tracer.py:374 — the code uses:

with torch.compiler._patch_autograd_grad():

Failed test flavors:

• aot_fx_trace llama3 FSDP+TP+cudagraph
• aot_fx_trace llama3 FSDP+TP+FlexAttn

All 8 ranks crash with the same AttributeError. This API doesn't exist in the PyTorch nightly installed in CI.

This is caused by this PR's changes.

---

Failure 2: GraphTrainer 8 GPU H100 Integration Tests

Same root cause — torch.compiler._patch_autograd_grad missing.

Failed test flavors:

• aot_fx_trace deepseek_v3 FSDP+TP+EP
• aot_fx_trace deepseek_v3 FSDP+TP+EP+FlexAttn
• aot_fx_trace qwen3 FSDP+TP
• aot_fx_trace qwen3 MoE FSDP+TP+EP

All aot_fx_trace configurations fail with the identical AttributeError.

This is caused by this PR's changes.

---

Failure 3: 8 GPU Feature Tests (NOT related to this PR)

3 sub-tests failed, all due to pre-existing infrastructure issues:

Test	Error
FSDP+FLEX_ATTN	InductorError: A compilation subprocess exited unexpectedly — Inductor subprocess crash
FSDP + FLEX + per op SAC	LLVM SLPVectorizer.cpp:2690 assertion failure → SIGABRT — Triton/LLVM compiler bug
SFT ChatDataset	Same LLVM SLPVectorizer assertion → SIGABRT

These are NOT caused by this PR — they're Inductor/LLVM/Triton bugs.

---

Summary

Job	Status	This PR?
Lint	Passed	—
CPU Unit Test	Passed	—
GraphTrainer 8 GPU (run 1)	Failed	Yes — _patch_autograd_grad missing
GraphTrainer 8 GPU (run 2)	Failed	Yes — same
GraphTrainer 8 GPU H100	Failed	Yes — same (deepseek, qwen3)
8 GPU Feature Tests	Failed	No — LLVM/Inductor bugs
8 remaining jobs	Queued/In-progress	—

Single root cause for all PR-related failures: torch.compiler._patch_autograd_grad() at make_fx_tracer.py:374 does not exist in the CI PyTorch nightly. You need to either:

1. Guard with hasattr(torch.compiler, '_patch_autograd_grad') and make it optional
2. Wait for the PyTorch nightly that includes this API to land in CI
3. Use an alternative if the API was renamed or moved

---

[tugsbayasgalan]

@claude what are the failures?

[claude]

Claude finished @tugsbayasgalan's task in 3m 39s —— View job

---

CI Failure Analysis

• Check CI status
• Analyze failing jobs
• Report findings

---

There are 2 completed failures out of 14 jobs (9 still queued/in-progress). Lint, CPU Unit Test, and 8 GPU Feature Tests all passed.

---

Failure 1 & 2: GraphTrainer 8 GPU Integration Tests (both runs — same errors)

Both runs fail on the same 2 test flavors with different errors:

Test Flavor 1: aot_fx_trace llama3 FSDP+TP+cudagraph

Error: RNG ops in recompute regions

RuntimeError: Activation checkpoint rematerialization in `forward-loss-backward` graph
does not support RNG ops in recompute regions. Please move RNG operations outside of
recompute regions, or use joint graph mode (where partitioner handles RNG).

Stack trace:

trainer.py:132  →  _make_fx_forward_backward_step
passes.py:120   →  apply_graph_passes → pass_fn(gm, example_inputs)
passes.py:491   →  apply_ac_on_fwd_bwd_graph
torch/...remat_using_tags_for_fwd_loss_bwd_graph_pass.py:124  →  raises RuntimeError

Root cause: The remat_using_tags_for_fwd_loss_bwd_graph pass detects RNG operations (e.g., dropout) in regions tagged for recomputation and rejects them. The fwd-loss-bwd graph remat pass doesn't handle RNG state save/restore like the joint graph partitioner does. The SAC tagging in apply_sac_pass is tagging nodes that include RNG ops as PREFER_RECOMPUTE, but the remat pass can't safely recompute those because the random state would differ on replay.

Test Flavor 2: aot_fx_trace llama3 FSDP+TP+FlexAttn

Error: Disjoint backward regions

RuntimeError: Detected 7 disjoint backward regions in the graph but remat only supports
a single backward region. This can happen when non-backward computation appears between
backward sections. Please annotate the real backward with
torch.fx.traceback.annotate({"phase": "backward"}).

Stack trace: Same path through passes.py:491 → remat_using_tags_for_fwd_loss_bwd_graph

Root cause: When FlexAttention is used, the traced graph contains nested backward computations from flex_attention's internal autograd. The _is_backward_node annotation logic (which uses custom["phase"] == "backward") isn't correctly identifying all backward nodes in the FlexAttn case, causing the remat pass to see what looks like multiple disjoint backward regions instead of one contiguous backward phase. This is the exact edge case noted in the _is_backward_node comment at passes.py:48-59 about flex_attention nested backward tracing.

---

Summary

Job	Status	This PR?	Error
Lint	Passed	—	—
CPU Unit Test	Passed	—	—
8 GPU Feature Tests	Passed	—	—
GraphTrainer 8 GPU (run 1)	Failed	Yes	RNG in recompute + disjoint backward
GraphTrainer 8 GPU (run 2)	Failed	Yes	Same
GraphTrainer 8 GPU H100	Queued	—	—
8 other jobs	Queued	—	—

Note: The previous torch.compiler._patch_autograd_grad error is gone — the CI PyTorch nightly now includes that API. The new failures are both in the SAC/remat pass:

1. FSDP+TP+cudagraph: apply_sac_pass tags RNG ops (dropout) for recomputation, but remat_using_tags_for_fwd_loss_bwd_graph doesn't support RNG ops in recompute regions. Fix: either tag RNG ops as MUST_SAVE instead of PREFER_RECOMPUTE, or exclude RNG ops from recomputation tagging in apply_sac_pass.

2. FSDP+TP+FlexAttn: The backward phase annotation doesn't capture all backward nodes when FlexAttention's nested backward is involved, causing the remat pass to see 7 disjoint backward regions. Fix: the backward annotation needs to handle flex_attention's nested autograd.grad calls correctly.

---

[tugsbayasgalan]

@claude I am pretty sure there was no RNG ops to be saved few days ago, is this a recent change?

[claude]

Claude finished @tugsbayasgalan's task in 13m 2s —— View job

---

Analysis: RNG Ops in Recompute Regions

• Read current passes.py, make_fx_tracer.py, and trainer.py
• Analyze CI failure logs for both test flavors
• Investigate whether RNG ops are from model code or PyTorch internals
• Determine if this is a recent PyTorch nightly change

---

Short answer: This is almost certainly a recent PyTorch nightly change, not something you introduced. The Llama3 debug model has no dropout or explicit RNG ops in its forward pass.

Here's the detailed analysis:

What the CI shows

The SAC pass itself runs successfully — the logs show it tagging nodes with MUST_SAVE/PREFER_RECOMPUTE across all AC regions. The error comes after tagging, when remat_using_tags_for_fwd_loss_bwd_graph (line 124 of the PyTorch pass) validates the tagged graph and finds RNG ops in PREFER_RECOMPUTE regions.

The Llama3 debug model has no RNG ops

I checked the entire model path: TransformerBlock, attention (ScaledDotProductAttention), FeedForward (SwiGLU) — none have dropout, torch.rand, torch.randn, or any other RNG operations. The train_context is also clean (just optional loss_parallel).

What likely changed: _patch_autograd_grad or the remat pass itself

This PR switched from the old _patch_engine_run_backward() to torch.compiler._patch_autograd_grad() (make_fx_tracer.py:355). Two things could be happening:

1. _patch_autograd_grad introduces RNG state-tracking ops into the traced graph. When PyTorch's patched autograd.grad traces backward ops, the nightly may now emit RNG state save/restore operations (like philox_seed_like or aten.bernoulli_.float) as part of its backward decomposition. These ops would appear in the forward portion of the fwd+loss+bwd graph (since the entire torch.autograd.grad call is traced as one flat graph) and get tagged as PREFER_RECOMPUTE by apply_sac_pass.

2. The remat_using_tags_for_fwd_loss_bwd_graph pass recently added or tightened the RNG check. Since you said the RNG error didn't appear a few days ago, and the model code hasn't changed, the most likely explanation is that the PyTorch nightly updated remat_using_tags_for_fwd_loss_bwd_graph to check for RNG ops that it previously didn't check for, or _patch_autograd_grad started emitting new ops that trigger the check.

How to confirm

You can dump the graph right before the remat pass to see exactly which ops are flagged:

# In apply_ac_on_fwd_bwd_graph, before calling remat_using_tags_for_fwd_loss_bwd_graph:
for node in gm.graph.nodes:
    if node.op == "call_function" and "recompute" in node.meta:
        if node.meta["recompute"] == CheckpointPolicy.PREFER_RECOMPUTE:
            # Check what PyTorch considers RNG
            if hasattr(torch._functorch._activation_checkpointing, '_is_rng_node'):
                from torch._functorch._activation_checkpointing import _is_rng_node
                if _is_rng_node(node):
                    logger.info(f"RNG node in recompute region: {node.target}")

Or simpler: print all node targets in the graph and grep for anything RNG-related (bernoulli, uniform, normal, philox, rng).

Two separate failures

1. FSDP+TP+cudagraph → RNG ops error: The remat pass finds RNG ops in recompute regions. Likely caused by new ops from PyTorch nightly in the traced backward.

2. FSDP+TP+FlexAttn → 7 disjoint backward regions: FlexAttention's nested autograd.grad calls create multiple backward sections that _is_backward_node (which checks node.meta.get("autograd_backward", False)) doesn't fully identify. This is a separate issue from the RNG one.

Suggested fix for the RNG issue

If the RNG ops are coming from PyTorch internals (not your model), you likely need to either:

• Tag those specific RNG ops as MUST_SAVE in apply_sac_pass (skip them from PREFER_RECOMPUTE)
• Or coordinate with the PyTorch team on _patch_autograd_grad to not emit RNG ops, or on remat_using_tags_for_fwd_loss_bwd_graph to handle them

---

[tugsbayasgalan]

@claude What is the test command to reproduce these errors?

[claude]

Claude finished @tugsbayasgalan's task in 3m 0s —— View job

---

Test Commands to Reproduce CI Failures

• Check CI workflow files for test commands
• Read run_tests.py for test configurations
• Find the override definitions for GraphTrainer tests
• Provide reproduction commands

---

The two failing test flavors from CI are run via the integration test runner at torchtitan/experiments/graph_trainer/tests/integration_tests.py. Here's how to reproduce each:

Full suite command (what CI runs)

# This runs ALL graph_trainer_default tests (including the two that fail)
python -m torchtitan.experiments.graph_trainer.tests.integration_tests \
    --test_suite graph_trainer_default \
    --gpu_arch_type cuda \
    /tmp/test_output \
    --ngpu 8

Individual failing tests

Failure 1: aot_fx_trace llama3 FSDP+TP+cudagraph (RNG ops in recompute regions)

python -m torchtitan.experiments.graph_trainer.tests.integration_tests \
    --test_suite graph_trainer_default \
    --test_name aot_fx_trace_llama3_fsdp_tp \
    --gpu_arch_type cuda \
    /tmp/test_output \
    --ngpu 8

Which under the hood runs ./run_train.sh with these overrides:

--module graph_trainer.llama3
--config graph_trainer_llama3_debugmodel
--compile.mode aot_fx_trace
--parallelism.data_parallel_shard_degree 4
--parallelism.tensor_parallel_degree 2

Failure 2: aot_fx_trace llama3 FSDP+TP+FlexAttn (Disjoint backward regions)

python -m torchtitan.experiments.graph_trainer.tests.integration_tests \
    --test_suite graph_trainer_default \
    --test_name aot_fx_trace_llama3_fsdp_tp_flexattn \
    --gpu_arch_type cuda \
    /tmp/test_output \
    --ngpu 8

Which runs ./run_train.sh with:

--module graph_trainer.llama3
--config graph_trainer_llama3_debugmodel_flex_attn
--compile.mode aot_fx_trace
--parallelism.data_parallel_shard_degree 4
--parallelism.tensor_parallel_degree 2

Equivalent direct run_train.sh commands

If you want to run them directly without the test harness:

# Failure 1: FSDP+TP+cudagraph (RNG error)
NGPU=8 LOG_RANK=0,1,2,3,4,5,6,7 MODULE=graph_trainer.llama3 CONFIG=graph_trainer_llama3_debugmodel \
  ./run_train.sh --dump_folder /tmp/test_output/aot_fx_trace_llama3_fsdp_tp \
    --compile.mode aot_fx_trace \
    --parallelism.data_parallel_shard_degree 4 \
    --parallelism.tensor_parallel_degree 2

# Failure 2: FSDP+TP+FlexAttn (disjoint backward regions)
NGPU=8 LOG_RANK=0,1,2,3,4,5,6,7 MODULE=graph_trainer.llama3 CONFIG=graph_trainer_llama3_debugmodel_flex_attn \
  ./run_train.sh --dump_folder /tmp/test_output/aot_fx_trace_llama3_fsdp_tp_flexattn \
    --compile.mode aot_fx_trace \
    --parallelism.data_parallel_shard_degree 4 \
    --parallelism.tensor_parallel_degree 2

CI also runs these unit tests (which passed)

# Bitwise deterministic test
pytest torchtitan/experiments/graph_trainer/tests/test_bitwise_deterministic.py -v

# SAC peak memory test
pytest torchtitan/experiments/graph_trainer/tests/test_sac_peak_memory.py -v

# SAC pass unit tests
pytest torchtitan/experiments/graph_trainer/tests/test_passes.py::TestApplySACPass -v

The test definitions are in torchtitan/experiments/graph_trainer/tests/integration_tests.py — look for the build_graph_trainer_default_test_list() function to see all OverrideDefinitions.

---

[SherlockNoMad]

not "outer dynamo trace", it's the "current make_fx trace"?

[SherlockNoMad]

apply after graph clean up pass

[SherlockNoMad]

please ask claude to come up with a more tasteful name, also follow naming convention, ends with *_pass

[SherlockNoMad]

follow up in future diff, I also think it' better to have remat_using_tags_for_fwd_loss_bwd_graph_pass in titan, since no one it's using it in core, and it's applicable senario is niche.

[tugsbayasgalan]

we do use it in torch.compile(fullgraph=True) today.

[SherlockNoMad]

I am not a fan of needing _non_strict_tracing_context coupling with _patch_autograd_grad.
What breaks if we don't have this context?
Is there other ways to fix this?

[tugsbayasgalan]

We only want to enable patching autograd in the cases we know it is carefully coupled with rest of the stack. So we added check in patch_autograd that it should only succeed if this specific context is on. The context is _non_strict_tracing_context. For example, we don't want to do this non-strict export.

[SherlockNoMad]

I don't see meaningful change in this file

[fegin]

nit: why delete this line?