DeepSpeed
26a49fd3 - AutoEP: share one DeepEP buffer across MoE layers instead of one per layer (#8557)

Commit
1 day ago
AutoEP: share one DeepEP buffer across MoE layers instead of one per layer (#8557) ## Summary Share a DeepEP `ElasticBuffer` across compatible MoE layers within each AutoEP-converted model, instead of constructing one buffer per layer. This removes redundant buffer construction and unblocks the 48-layer EP32 configuration tested here. ## Motivation In the tested four-node EP32 setup, buffer construction failed at the twenty-eighth buffer inside `ncclDevCommCreate`. An isolated buffer-construction probe reproduced the failure without a model, pointing to a resource limit associated with creating many buffers rather than model depth itself. This is an observed limit in this environment, not a general DeepEP limit. With sharing enabled, the 48-layer EP32 benchmark completes successfully. ## Design - Sharing is scoped to one converted model. Separately converted models retain independent buffers, even when their EP group and configuration match. - Layers share only when their EP group, expert count, top-k, hidden size, token capacity, SM budget, and QP margin match. - References are released during engine teardown. The final release destroys the buffer and removes its registry entry. - Dispatch outputs are separate tensors, and each autograd node retains its own handle for backward. Sharing relies on the current serialized DeepEP execution path; this PR does not introduce communication overlap. ## Testing Done - 43 unit tests pass, covering sharing, configuration separation, model isolation, and teardown. - Four-GPU parity testing: 4 passed, 2 failed. The same two optimizer-delta comparisons also fail on the pre-change control; the detailed comparison is documented in the testing comment. - The 48-layer EP32 paired benchmark completes with the shared-buffer implementation. - EP16 with an 8192-token sequence also completes. ## Scope This change addresses redundant communication-buffer construction and the observed EP32 initialization failure. It does not claim a steady-state throughput improvement or a reduction in PyTorch peak allocated memory: these DeepEP buffers are allocated outside the PyTorch caching allocator. --------- Signed-off-by: Han Yu <hanyu@linkedin.com> Signed-off-by: yh0903 <helloyu0903@gmail.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Masahiro Tanaka <81312776+tohtana@users.noreply.github.com>
Author
Parents
Loading