Skip to content

[graph_trainer][aot_fx_trace][graph_pp] Add dI/dW backward splitting#2727

Merged
sanketpurandare merged 1 commit into
mainfrom
sanketpurandare/stack/6
Jul 2, 2026
Merged

[graph_trainer][aot_fx_trace][graph_pp] Add dI/dW backward splitting#2727
sanketpurandare merged 1 commit into
mainfrom
sanketpurandare/stack/6

Conversation

@sanketpurandare

@sanketpurandare sanketpurandare commented Mar 27, 2026

Copy link
Copy Markdown
Contributor

Stacked PRs:


[graph_trainer][aot_fx_trace][graph_pp] Add dI/dW backward splitting

Add the GraphPP backward split pass for schedules that execute input-gradient and weight-gradient work as separate PP actions. The pass preserves the flat graph calling convention selected by the partitioner and returns explicit activation handoff metadata for the dW graph.

Add unit coverage that reconstructs the full backward result from split dI and dW graphs and validates the no-op path when a graph has no input-gradient outputs.

sanketpurandare pushed a commit that referenced this pull request Mar 27, 2026
stack-info: PR: #2727, branch: sanketpurandare/stack/6
@sanketpurandare sanketpurandare force-pushed the sanketpurandare/stack/6 branch from b54d197 to bed5c29 Compare March 27, 2026 00:49
@sanketpurandare sanketpurandare force-pushed the sanketpurandare/stack/5 branch from 3e47d66 to 19150f9 Compare March 27, 2026 00:49
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Mar 27, 2026
@sanketpurandare sanketpurandare changed the base branch from sanketpurandare/stack/5 to main March 27, 2026 01:13
sanketpurandare pushed a commit that referenced this pull request Mar 27, 2026
stack-info: PR: #2727, branch: sanketpurandare/stack/6
@sanketpurandare sanketpurandare force-pushed the sanketpurandare/stack/6 branch from bed5c29 to e6c34d8 Compare March 27, 2026 01:14
@sanketpurandare sanketpurandare changed the base branch from main to sanketpurandare/stack/5 March 27, 2026 01:14
@sanketpurandare sanketpurandare marked this pull request as draft March 27, 2026 02:24
@sanketpurandare sanketpurandare changed the base branch from sanketpurandare/stack/5 to main March 27, 2026 18:57
sanketpurandare pushed a commit that referenced this pull request Mar 27, 2026
stack-info: PR: #2727, branch: sanketpurandare/stack/6
@sanketpurandare sanketpurandare force-pushed the sanketpurandare/stack/6 branch from e6c34d8 to 89371b2 Compare March 27, 2026 18:57
@sanketpurandare sanketpurandare changed the base branch from main to sanketpurandare/stack/5 March 27, 2026 18:57
@sanketpurandare sanketpurandare changed the base branch from sanketpurandare/stack/5 to main April 1, 2026 16:28
sanketpurandare pushed a commit that referenced this pull request Apr 1, 2026
stack-info: PR: #2727, branch: sanketpurandare/stack/6
@sanketpurandare sanketpurandare force-pushed the sanketpurandare/stack/6 branch from 89371b2 to baa2a42 Compare April 1, 2026 16:28
@sanketpurandare sanketpurandare changed the base branch from main to sanketpurandare/stack/5 April 1, 2026 16:29
@sanketpurandare sanketpurandare changed the base branch from sanketpurandare/stack/5 to main April 23, 2026 20:08
@sanketpurandare sanketpurandare changed the base branch from main to sanketpurandare/stack/5 April 23, 2026 20:08
@sanketpurandare sanketpurandare changed the base branch from sanketpurandare/stack/5 to main April 24, 2026 19:53
@sanketpurandare sanketpurandare force-pushed the sanketpurandare/stack/6 branch from baa2a42 to 0a6ed49 Compare April 24, 2026 19:53
@sanketpurandare sanketpurandare changed the base branch from main to sanketpurandare/stack/5 May 8, 2026 03:13
@sanketpurandare sanketpurandare changed the base branch from sanketpurandare/stack/5 to main May 8, 2026 09:34
@sanketpurandare sanketpurandare force-pushed the sanketpurandare/stack/6 branch from caeee4b to 37de310 Compare May 8, 2026 09:34
@sanketpurandare sanketpurandare changed the base branch from main to sanketpurandare/stack/5 May 8, 2026 09:34
@sanketpurandare sanketpurandare force-pushed the sanketpurandare/stack/5 branch from 74de9a6 to 398c43f Compare May 8, 2026 16:57
sanketpurandare added a commit that referenced this pull request May 8, 2026
stack-info: PR: #2727, branch: sanketpurandare/stack/6
@sanketpurandare sanketpurandare force-pushed the sanketpurandare/stack/6 branch from 37de310 to 47ebb7c Compare May 8, 2026 16:57
@sanketpurandare sanketpurandare changed the base branch from sanketpurandare/stack/5 to main June 9, 2026 03:59
@sanketpurandare sanketpurandare force-pushed the sanketpurandare/stack/6 branch from 47ebb7c to a6db1e3 Compare June 9, 2026 03:59
@sanketpurandare sanketpurandare changed the base branch from main to sanketpurandare/stack/5 June 9, 2026 03:59
@sanketpurandare sanketpurandare marked this pull request as ready for review June 9, 2026 04:16

@aditvenk aditvenk left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM -- some minor comments.

Comment thread torchtitan/experiments/graph_trainer/graph_pp/split_di_dw.py Outdated
Comment thread torchtitan/experiments/graph_trainer/graph_pp/split_di_dw.py
Comment thread torchtitan/experiments/graph_trainer/graph_pp/split_di_dw.py Outdated
Comment thread torchtitan/experiments/graph_trainer/graph_pp/split_di_dw.py Outdated
Add the GraphPP backward split pass for schedules that execute input-gradient and weight-gradient work as separate PP actions. The pass preserves the flat graph calling convention selected by the partitioner and returns explicit activation handoff metadata for the dW graph.

Add unit coverage that reconstructs the full backward result from split dI and dW graphs and validates the no-op path when a graph has no input-gradient outputs.

stack-info: PR: #2727, branch: sanketpurandare/stack/6
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/8gpu CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[graph_trainer] Graph-PP RFC

2 participants