Skip to content

Add native transport for dynamic context parallelism - #3489

Draft
xiaoyao0115 wants to merge 1 commit into
NVIDIA:release_v2.14from
xiaoyao0115:dcp-native-transport-minimal
Draft

Add native transport for dynamic context parallelism#3489
xiaoyao0115 wants to merge 1 commit into
NVIDIA:release_v2.14from
xiaoyao0115:dcp-native-transport-minimal

Conversation

@xiaoyao0115

@xiaoyao0115 xiaoyao0115 commented Sep 6, 2026

Copy link
Copy Markdown

What

  • Add an opt-in NCCL Device API transport for context-parallel P2P rings.
  • Borrow one already-materialized parent NCCL communicator instead of creating a ProcessGroup for every dynamic CP size.
  • Reuse one symmetric arena for forward KV, backward dKV, and auxiliary-loss communication.
  • Accept lightweight logical CP descriptors while leaving the existing ProcessGroup path unchanged.

Scope

  • The native path currently supports P2P context parallelism with FP16/BF16 and vanilla softmax.
  • FP8, non-P2P CP algorithms, and max-logit output are rejected explicitly. CUDA Graph capture is not validated by this PR.
  • The extension is built only when NCCL Device API headers and libraries are available.

Validation

  • Exact-source H100 build and import checks passed with NCCL 2.29.7.
  • FusedAttention forward/backward matched the no-CP reference for SBHD CP1-16 and THD CP1-14.
  • Deterministic THD CP15/16 cases without inter-sequence padding passed for both NCCL P2P and native transport; CP16 also passed with random padding.
  • A deterministic random-padding THD CP15 case exceeds the existing BF16 tolerance on both the original NCCL P2P path and the native path. This is not reported as a native-only pass; the raw comparison is retained as a known test-path limitation.
  • Qwen3-30B-A3B on LongAlign-10k completed 100 fixed-CP16 and 100 native arbitrary-DCP training steps from identical sampled parameters and identical real batches. LM loss max/mean absolute differences were 0.000648/0.000131 (Pearson 0.999999589); sequence auxiliary loss differences were 0.005567/0.001644 (Pearson 0.999958020). Different CP layouts change BF16 reduction order, so bitwise equality is not claimed.
  • Fresh-process 4x4 GB200 NVL72 accounting (32 P2P channels) shows that activating every bounded-ring CP1-16 path adds 0 MiB after native transport creation. The measured native fixed cost was 42/46/54 MiB per rank for parent sizes 4/8/16 respectively; parent-16 used 54 MiB with the minimum 2 MiB window, or 182 MiB with a 128 MiB payload arena. The corresponding ProcessGroupNCCL paths retained 768 MiB on rank 0 (two peers) and 1536 MiB on the median/max rank (four peers); directly activating all 15 peers retained 5812 MiB on rank 0.
  • An isolated H100 CP8 FusedAttention forward+backward A/B measured 10.178 ms for ProcessGroupNCCL and 10.327 ms for native transport (+1.47%). The transport is memory-oriented; end-to-end speedup comes from allowing the scheduler to use better non-power-of-two layouts, not from claiming a faster individual ring.
  • The LongAlign run measured 7181.77 token/s for native arbitrary DCP versus 5811.30 token/s for fixed CP16 (+23.58%). A separate eight-GB200 CP5+CP3 case reduced median step time by 18.75% versus the same native transport constrained to CP8.

The dependent NVIDIA/Megatron-LM#7097 PR enables this path with --use-native-cp-transport.

@github-actions github-actions Bot added the community-contribution PRs from external contributor outside the core maintainers, representing community-driven work. label Sep 6, 2026
@xiaoyao0115
xiaoyao0115 force-pushed the dcp-native-transport-minimal branch from e861f1f to e2dd2bb Compare September 6, 2026 14:21
Signed-off-by: Tailai Ma <tailaim@ferries-work.nvidia.com>
@xiaoyao0115
xiaoyao0115 force-pushed the dcp-native-transport-minimal branch from e2dd2bb to f839ade Compare September 6, 2026 14:29
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

community-contribution PRs from external contributor outside the core maintainers, representing community-driven work.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant