[DCP V2 3/4] Integrate dynamic CP with standard train and eval schedules by xiaoyao0115 · Pull Request #7047 · NVIDIA/Megatron-LM

xiaoyao0115 · 2026-09-02T10:03:09Z

I, the PR author, have personally reviewed every line of this PR.

What does this PR do?

This is DCP V2 3/4. Dynamic context parallelism now uses the standard train/evaluation schedules and caller-owned process groups, including packed data on PP, VPP, and MTP stages. The obsolete dedicated Hybrid-CP schedule and data-loader wrapper are removed.

The implementation:

exposes dynamic_context_parallel, retaining a deprecated Hybrid-CP CLI/config alias;
packs text samples into supported runtime CP groups, including singleton and full non-power-of-two groups, and carries runtime CP and TP×CP groups in ProcessGroupCollection;
reroutes samples through CP-lane-local DP all-gather and supplies independent PP/VPP iterators with the fields each stage owns;
returns contiguous/zigzag batch views, packed metadata, and layout-specific padding masks; each Hybrid mask is SP-sharded before reaching its layer;
bounds packing, CP expansion, fallback, and VPP alignment by the actual per-sequence padding cost, and reports unsplittable VPP batches without looping;
supports static Hybrid packing and the default zigzag layout with the cumulative-length shape expected by Transformer Engine; and
ends finite packed evaluation without closing an unopened fault-tolerance step.

For example, a PP first stage no longer tries to CP-slice removed labels; static packing no longer dereferences a missing runtime CP size; and twelve 1024-token rows at CP6/capacity2048/alignment32 are split into batches that fit after padding.

Original change authored by @zhongbozhu in zhongbozhu/Megatron-LM@a708518.

Series and dependencies

Targets main; #7045, #7046, and #7207 have merged. #7048 follows this PR. Source lineage includes #2000, #5154, #5483, and #6378. CUDA graph and router z-loss changes are outside this PR.

October 14 refresh

Rebased onto main at 7a0024e2c, preserving the runtime-group injection, default-zigzag fix, and FT evaluation fix. DCP initialization uses main's current token-rank generator, and activation recomputation retains main's precision-aware checkpoint implementation.

The follow-up fixes the CI fixtures identified in the previous data/models/transformer failures: canonical dynamic-CP config fields and packing capacity, the scheduled ContextParallelBatch interface and flat CUDA batch schema, layout-state mask lookup, and the two added Hybrid MoE configuration fields. Required token, dtype, length, and mask checks remain in place.

Validation

226 CPU regressions passed across two focused runs, including runtime-group injection, PP/CP ownership, static packing, VPP alignment/capacity, layout-specific SP masks, varlen data, and FT evaluation. The local CPU runs exclude GPU-only cases and optional unavailable dataset dependencies.
Black, isort, Pylint, Ruff, whitespace, process-group usage, and kernel-determinism coverage checks passed. Mypy is unavailable locally.
8×H100/NCCL job 20153993 completed with exit code 0: 581 passed and 113 existing conditional skips on every rank, on head dedb587dc. All previously failing data, Hybrid model/config, and residual-recompute cases passed. The run also covers default-zigzag CP2, PP/VPP packing, runtime CP1/2/4/8 router counts/gradient scaling, SP masks, and FT evaluation. The validation container uses Torch 2.12 development and Transformer Engine 2.14 development builds.
GitHub CI for dedb587dc has been triggered. Linting passed; the remaining tests are waiting for the protected test environment queue. The current account cannot approve that environment, and reviewer/codeowner approval is still required.
These targeted tests do not establish full end-to-end training equivalence or complete GitHub CI success.

Contribution process

Pre-checks

I have added relevant unit tests
I have added relevant functional tests
I have added proper typing to my code
I have added relevant documentation
I have run the autoformatter on my PR

copy-pr-bot · 2026-09-02T10:03:13Z

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

xiaoyao0115 · 2026-09-02T10:04:54Z

/ok to test c6c8257

copy-pr-bot · 2026-10-01T08:20:22Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

xiaoyao0115 · 2026-10-01T08:20:57Z

/ok to test 3fc6410

guihong-nv

Please fix the below issues. The rest LGTM.

Signed-off-by: Tailai Ma <58548582+xiaoyao0115@users.noreply.github.com> Signed-off-by: Tailai Ma <tailaim@nvidia.com>

Signed-off-by: Tailai Ma <tailaim@nvidia.com> Signed-off-by: ilml <tolong@nvidia.com> Signed-off-by: Tailai Ma <tailaim@nvidia.com>

Drop the process-group allowlist entries for HybridCPDataLoaderWrapper and hybrid_cp_schedule.py, which this branch deletes, and point the CUDA-graph validation test at dynamic_context_parallel and its new error message. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: ilml <tolong@nvidia.com> Signed-off-by: Tailai Ma <tailaim@nvidia.com>

megatron/core no longer reads the runtime CP and TP x CP groups from parallel_state inside get_batch_on_this_rank_for_sequence_packing: callers supply dynamic_cp_group_func, dynamic_tp_cp_group_func and tp_cp_group (or pg_collection.tp_cp). Drop the new global cp_group fallback in get_batch_on_this_cp_rank; every caller passes cp_group. This keeps the process-group usage check passing without growing its allowlist. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: ilml <tolong@nvidia.com> Signed-off-by: Tailai Ma <tailaim@nvidia.com>

Signed-off-by: Tailai Ma <tailaim@nvidia.com>

Carry pre-created runtime CP and TPxCP maps in ProcessGroupCollection and use the model collection throughout packed batch scheduling and preparation. Keep MPU bootstrap at the training boundary, and reject scheduling fallback to unavailable groups. Validation: 155 focused CPU tests passed; 8 NCCL batch cases passed on each of two GPU ranks; Black, isort, Ruff, Pylint and process-group usage checks passed. Mypy retains the 21 pre-existing diagnostic sites. Signed-off-by: ilml <tolong@nvidia.com>

With dp_balanced packing, CP > 1 and both CP layouts at the default zigzag, get_batches_on_this_cp_rank shards each document with _get_batch_on_this_cp_rank_per_document_balancing, which indexes row 0 of cu_seqlens_padded. The scheduler batch carried 1-D cu_seqlens, so Transformer Engine received a scalar and failed its cu_seqlens.dim() == 1 check. Pass (1, num_seqs + 1) cu_seqlens and cu_seqlens_padded when building a ContextParallelBatch, matching flatten_batch_for_packed_sequences. Cover the default layout in the static packing CPU test, whose fake partitioner now enforces Transformer Engine's 1-D contract, and in a CP2 test that compares the zigzag batch with the tuple path through Transformer Engine. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: ilml <tolong@nvidia.com> Signed-off-by: Tailai Ma <tailaim@nvidia.com>

The StopIteration branch in evaluate runs before on_eval_step_start, so calling on_eval_step_end there closed a fault-tolerance step section that was never opened once past warmup ("Tried to close step that is not open"), and counted a phantom eval step before it. Only restore config.timers before ending evaluation, and add a CPU test that drives evaluate through the real FT hooks. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: ilml <tolong@nvidia.com> Signed-off-by: Tailai Ma <tailaim@nvidia.com>

Signed-off-by: Tailai Ma <tailaim@nvidia.com>

Use torch.full for Python integer counts to avoid the blocking host-to-device scalar copy. Preserve tensor inputs and the existing loss-scale reductions. Cover integer, implicit, tensor and zero-token counts across runtime CP1/2/4, and replay unmasked per-token sequence auxiliary loss bit-exactly. Signed-off-by: Tailai Ma <tailaim@nvidia.com>

xiaoyao0115 mentioned this pull request Sep 2, 2026

[DCP V2 4/4] Propagate runtime dynamic CP through hybrid model layers #7048

Draft

3 of 6 tasks

xiaoyao0115 added the Run functional tests label Sep 2, 2026 — with ChatGPT Codex Connector

copy-pr-bot Bot temporarily deployed to test September 2, 2026 10:06 Inactive

copy-pr-bot Bot temporarily deployed to public September 2, 2026 10:08 Inactive

lhb8125 mentioned this pull request Sep 7, 2026

Megatron-LM Monthly Activity Report — September 2026 (Final, through September 30, 2026) #7100

Open

xiaoyao0115 force-pushed the dcp-v2-standard-train-eval-schedules branch from c6c8257 to f42fb7f Compare September 7, 2026 14:13

xiaoyao0115 mentioned this pull request Sep 7, 2026

[DCP V2 2/4] Make MoE losses runtime dynamic-CP aware #7046

Merged

3 of 6 tasks

xiaoyao0115 force-pushed the dcp-v2-standard-train-eval-schedules branch 2 times, most recently from 9814b68 to 5a7ad87 Compare September 16, 2026 05:24

xiaoyao0115 mentioned this pull request Sep 16, 2026

[DCP V2 1/4] Define the runtime dynamic-CP group contract #7045

Merged

7 of 10 tasks

FDecaYed mentioned this pull request Sep 18, 2026

[Tracking] DeepSeek-v4 HybridModel merge plan #5795

Draft

zupengwang mentioned this pull request Sep 19, 2026

Fix context-parallel normalization of masked training losses #7513

Open

4 of 6 tasks

lhb8125 mentioned this pull request Oct 1, 2026

Megatron-LM Monthly Activity Report — September 2026 #7773

Open

xiaoyao0115 force-pushed the dcp-v2-standard-train-eval-schedules branch from 5a7ad87 to 3fc6410 Compare October 1, 2026 08:20

copy-pr-bot Bot deployed to test October 1, 2026 08:23 Active

copy-pr-bot Bot deployed to public October 1, 2026 08:32 Active

lhb8125 mentioned this pull request Oct 5, 2026

Megatron-LM Monthly Activity Report — October 2026 (through October 11, 2026) #7882

Open

ilml marked this pull request as ready for review October 7, 2026 01:48

ilml requested review from a team as code owners October 7, 2026 01:48

copy-pr-bot Bot deployed to public October 9, 2026 09:17 Active

zhongbozhu mentioned this pull request Oct 9, 2026

Add configurable batch and sequence partitioning within GTP_remat #7748

Merged

guihong-nv reviewed Oct 12, 2026

View reviewed changes

Comment thread megatron/core/datasets/data_schedule.py

Comment thread megatron/training/training.py Outdated

guihong-nv approved these changes Oct 12, 2026

View reviewed changes

Phlip79 removed the request for review from a team October 13, 2026 16:46

xiaoyao0115 and others added 10 commits October 14, 2026 11:49

Integrate dynamic CP with standard train and eval schedules

c9689a3

Signed-off-by: Tailai Ma <58548582+xiaoyao0115@users.noreply.github.com> Signed-off-by: Tailai Ma <tailaim@nvidia.com>

Format rebased schedule regression tests

34af996

Signed-off-by: Tailai Ma <tailaim@nvidia.com> Signed-off-by: ilml <tolong@nvidia.com> Signed-off-by: Tailai Ma <tailaim@nvidia.com>

Fix packed CP scheduling and Hybrid SP review regressions

9b79d5a

Signed-off-by: Tailai Ma <tailaim@nvidia.com>

Use bias-free MoE in packed schedule distributed regression

014ace8

Signed-off-by: Tailai Ma <tailaim@nvidia.com>

Align CI fixtures with dynamic CP packing and layout interfaces

27b719c

Signed-off-by: Tailai Ma <tailaim@nvidia.com>

xiaoyao0115 force-pushed the dcp-v2-standard-train-eval-schedules branch from 9fef21a to 27b719c Compare October 14, 2026 03:00

copy-pr-bot Bot had a problem deploying to test October 14, 2026 03:02 Error

copy-pr-bot Bot deployed to public October 14, 2026 03:04 Active

Read optional fields from scheduled context parallel batches

dedb587

Signed-off-by: Tailai Ma <tailaim@nvidia.com>

copy-pr-bot Bot had a problem deploying to test October 14, 2026 03:22 Error

copy-pr-bot Bot deployed to public October 14, 2026 03:24 Active

copy-pr-bot Bot requested a deployment to test October 14, 2026 05:02 Waiting

copy-pr-bot Bot deployed to public October 14, 2026 05:05 Active

Navigation Menu

Repository navigation

[DCP V2 3/4] Integrate dynamic CP with standard train and eval schedules - #7047

[DCP V2 3/4] Integrate dynamic CP with standard train and eval schedules#7047
xiaoyao0115 wants to merge 12 commits into
NVIDIA:mainfrom
xiaoyao0115:dcp-v2-standard-train-eval-schedules

xiaoyao0115 commented Sep 2, 2026 •

edited

Loading

Uh oh!

copy-pr-bot Bot commented Sep 2, 2026

Uh oh!

xiaoyao0115 commented Sep 2, 2026

Uh oh!

copy-pr-bot Bot commented Oct 1, 2026

Uh oh!

xiaoyao0115 commented Oct 1, 2026

Uh oh!

guihong-nv left a comment •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

This branch is waiting to be deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

xiaoyao0115 commented Sep 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Series and dependencies

October 14 refresh

Validation

Contribution process

Pre-checks

Uh oh!

copy-pr-bot Bot commented Sep 2, 2026

Uh oh!

xiaoyao0115 commented Sep 2, 2026

Uh oh!

copy-pr-bot Bot commented Oct 1, 2026

Uh oh!

xiaoyao0115 commented Oct 1, 2026

Uh oh!

guihong-nv left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

This branch is waiting to be deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

xiaoyao0115 commented Sep 2, 2026 •

edited

Loading

guihong-nv left a comment •

edited

Loading