Skip to content

graph_trainer: unify chunked grad collective normalization - #3953

Draft
sanketpurandare wants to merge 4 commits into
gh/sanketpurandare/4/basefrom
gh/sanketpurandare/4/head
Draft

graph_trainer: unify chunked grad collective normalization#3953
sanketpurandare wants to merge 4 commits into
gh/sanketpurandare/4/basefrom
gh/sanketpurandare/4/head

Conversation

@sanketpurandare

@sanketpurandare sanketpurandare commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Summary

Normalize duplicate chunk-local parameter-gradient collective chains after
either eager or graph chunking. Compatible chunk gradients are accumulated
before a single all-reduce or reduce-scatter, and any wait, view, or dtype-cast
suffix is replayed once on the normalized result.

Why

Dense FSDP and EP schedulers should see the same gradient buckets regardless
of how the two chunks were produced. A shared post-chunking pass keeps that
contract in one place and avoids producer-specific scheduling behavior.

Test Plan

  • Trace real SimpleFSDP eager and graph chunk paths in replicate and
    fully-sharded modes.
  • Verify eager duplicate collectives normalize and graph collectives remain
    canonical.
  • Exercise suffix replay through a naturally produced mixed-precision cast.

…tive normalization"

## Summary

Normalize duplicate chunk-local parameter-gradient collective chains after
either eager or graph chunking. Compatible chunk gradients are accumulated
before a single all-reduce or reduce-scatter, and any wait, view, or dtype-cast
suffix is replayed once on the normalized result.

## Why

Dense FSDP and EP schedulers should see the same gradient buckets regardless
of how the two chunks were produced. A shared post-chunking pass keeps that
contract in one place and avoids producer-specific scheduling behavior.

## Test Plan

- Trace real SimpleFSDP eager and graph chunk paths in replicate and
  fully-sharded modes.
- Verify eager duplicate collectives normalize and graph collectives remain
  canonical.
- Exercise suffix replay through a naturally produced mixed-precision cast.

[ghstack-poisoned]
[ghstack-poisoned]
[ghstack-poisoned]
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.

1 participant