Skip to content

graph_trainer: support MinimalAsyncEP graph chunking - #3956

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

graph_trainer: support MinimalAsyncEP graph chunking#3956
sanketpurandare wants to merge 4 commits into
gh/sanketpurandare/7/basefrom
gh/sanketpurandare/7/head

Conversation

@sanketpurandare

@sanketpurandare sanketpurandare commented Jul 21, 2026

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Summary

Extend the generic FX chunking path to MinimalAsyncEP. Discover chunk roots
through graph-visible token-exchange launches, then assign MinimalAsyncEP
buffer sets with a separate schema-driven post-chunk pass.

The generic graph cloning algorithm remains backend-independent. Eager and
graph chunking feed the same ownership and launch/wait metadata to the shared
EP scheduler.

Why

Buffer assignment is a backend storage concern, not part of generic region
cloning. Keeping it in a small post-pass preserves the graph chunking contract
and avoids MinimalAsyncEP special cases in the large cloning pass.

Test Plan

  • Validate token-exchange count invariants for the generalized scheduler.
  • Run an exact two-rank DeepSeek-v3 comparison between graph and eager
    MinimalAsyncEP chunking.
  • Compare both loss and gradient norm for 20 steps with forced-balanced routing
    and regional Inductor.

sanketpurandare added a commit that referenced this pull request Jul 21, 2026
ghstack-source-id: 430ba89
Pull Request resolved: #3956
…aph chunking"

## Summary

Extend the generic FX chunking path to MinimalAsyncEP. Discover chunk roots
through graph-visible token-exchange launches, then assign MinimalAsyncEP
buffer sets with a separate schema-driven post-chunk pass.

The generic graph cloning algorithm remains backend-independent. Eager and
graph chunking feed the same ownership and launch/wait metadata to the shared
EP scheduler.

## Why

Buffer assignment is a backend storage concern, not part of generic region
cloning. Keeping it in a small post-pass preserves the graph chunking contract
and avoids MinimalAsyncEP special cases in the large cloning pass.

## Test Plan

- Validate token-exchange count invariants for the generalized scheduler.
- Run an exact two-rank DeepSeek-v3 comparison between graph and eager
  MinimalAsyncEP chunking.
- Compare both loss and gradient norm for 20 steps with forced-balanced routing
  and regional Inductor.

[ghstack-poisoned]
[ghstack-poisoned]
sanketpurandare added a commit that referenced this pull request Aug 1, 2026
## Summary

Extend the generic FX chunking path to MinimalAsyncEP. Discover chunk roots
through graph-visible token-exchange launches, then assign MinimalAsyncEP
buffer sets with a separate schema-driven post-chunk pass.

The generic graph cloning algorithm remains backend-independent. Eager and
graph chunking feed the same ownership and launch/wait metadata to the shared
EP scheduler.

## Why

Buffer assignment is a backend storage concern, not part of generic region
cloning. Keeping it in a small post-pass preserves the graph chunking contract
and avoids MinimalAsyncEP special cases in the large cloning pass.

## Test Plan

- Validate token-exchange count invariants for the generalized scheduler.
- Run an exact two-rank DeepSeek-v3 comparison between graph and eager
  MinimalAsyncEP chunking.
- Compare both loss and gradient norm for 20 steps with forced-balanced routing
  and regional Inductor.

ghstack-source-id: 6ecba47
Pull Request resolved: #3956
sanketpurandare added a commit that referenced this pull request Aug 1, 2026
## Summary

Extend the generic FX chunking path to MinimalAsyncEP. Discover chunk roots
through graph-visible token-exchange launches, then assign MinimalAsyncEP
buffer sets with a separate schema-driven post-chunk pass.

The generic graph cloning algorithm remains backend-independent. Eager and
graph chunking feed the same ownership and launch/wait metadata to the shared
EP scheduler.

## Why

Buffer assignment is a backend storage concern, not part of generic region
cloning. Keeping it in a small post-pass preserves the graph chunking contract
and avoids MinimalAsyncEP special cases in the large cloning pass.

## Test Plan

- Validate token-exchange count invariants for the generalized scheduler.
- Run an exact two-rank DeepSeek-v3 comparison between graph and eager
  MinimalAsyncEP chunking.
- Compare both loss and gradient norm for 20 steps with forced-balanced routing
  and regional Inductor.

ghstack-source-id: 6ecba47
Pull Request resolved: #3956
sanketpurandare added a commit that referenced this pull request Aug 1, 2026
## Summary

Extend the generic FX chunking path to MinimalAsyncEP. Discover chunk roots
through graph-visible token-exchange launches, then assign MinimalAsyncEP
buffer sets with a separate schema-driven post-chunk pass.

The generic graph cloning algorithm remains backend-independent. Eager and
graph chunking feed the same ownership and launch/wait metadata to the shared
EP scheduler.

## Why

Buffer assignment is a backend storage concern, not part of generic region
cloning. Keeping it in a small post-pass preserves the graph chunking contract
and avoids MinimalAsyncEP special cases in the large cloning pass.

## Test Plan

- Validate token-exchange count invariants for the generalized scheduler.
- Run an exact two-rank DeepSeek-v3 comparison between graph and eager
  MinimalAsyncEP chunking.
- Compare both loss and gradient norm for 20 steps with forced-balanced routing
  and regional Inductor.

ghstack-source-id: 6ecba47
Pull Request resolved: #3956
sanketpurandare added a commit that referenced this pull request Aug 1, 2026
## Summary

Extend the generic FX chunking path to MinimalAsyncEP. Discover chunk roots
through graph-visible token-exchange launches, then assign MinimalAsyncEP
buffer sets with a separate schema-driven post-chunk pass.

The generic graph cloning algorithm remains backend-independent. Eager and
graph chunking feed the same ownership and launch/wait metadata to the shared
EP scheduler.

## Why

Buffer assignment is a backend storage concern, not part of generic region
cloning. Keeping it in a small post-pass preserves the graph chunking contract
and avoids MinimalAsyncEP special cases in the large cloning pass.

## Test Plan

- Validate token-exchange count invariants for the generalized scheduler.
- Run an exact two-rank DeepSeek-v3 comparison between graph and eager
  MinimalAsyncEP chunking.
- Compare both loss and gradient norm for 20 steps with forced-balanced routing
  and regional Inductor.

ghstack-source-id: 6ecba47
Pull Request resolved: #3956
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