Skip to content

Comments

nonuniform tensor parallelism#2149

Draft
daiyaanarfeen wants to merge 597 commits intoNVIDIA:mainfrom
daiyaanarfeen:nonuniform-tp
Draft

nonuniform tensor parallelism#2149
daiyaanarfeen wants to merge 597 commits intoNVIDIA:mainfrom
daiyaanarfeen:nonuniform-tp

Conversation

@daiyaanarfeen
Copy link

@daiyaanarfeen daiyaanarfeen commented Nov 5, 2025

What does this PR do ?

Enables training with reduced TP degree on DP ranks with some failed GPUs. Failed ranks can be specified by their coordinate in the parallelism mapping. This can enable training jobs to continue without rescheduling and waiting for available resources.

⚠️ For major changes (either in lines of code or in its impact), please make sure to first share discuss a design-doc with the team.

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]
Loading

Pre-checks

  • I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
  • I have added relevant unit tests
  • I have added relevant functional tests
  • I have added proper typing to my code Typing guidelines
  • I have added relevant documentation
  • I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

For MRs into `main` branch

(Step 1): Add PR label Expert Review

(Step 2): Collect the expert reviewers reviews

  1. Attach the Expert Review label when your PR is ready for review.
  2. GitHub auto-assigns expert reviewers based on your changes. They will get notified and pick up your PR soon.

⚠️ Only proceed to the next step once all reviewers have approved, merge-conflict are resolved and the CI is passing.
Final Review might get declined if these requirements are not fulfilled.

(Step 3): Final Review

  1. Add Final Review label
  2. GitHub auto-assigns final reviewers based on your changes. They will get notified and pick up your PR soon.

(Optional Step 4): Cherry-pick into release branch

If this PR also needs to be merged into core_r* release branches, after this PR has been merged, select Cherry-pick to open a new PR into the release branch.

For MRs into `dev` branch The proposed review process for `dev` branch is under active discussion.

MRs are mergable after one approval by either eharper@nvidia.com or zijiey@nvidia.com.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.

@copy-pr-bot
Copy link

copy-pr-bot bot commented Nov 5, 2025

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.

@ko3n1g
Copy link
Contributor

ko3n1g commented Nov 6, 2025

/ok to test 7a33efd

@ko3n1g ko3n1g added this to the Core 0.16 milestone Nov 6, 2025
@github-actions
Copy link
Contributor

github-actions bot commented Nov 6, 2025

Thank you for your contribution!

NVIDIA Megatron-LM is currently transitioning to development on Github. We will aim to review your PR after we complete our transition and stabilize our Github development process.

Thank you for your understanding.

skyw
skyw reviewed Dec 5, 2025
View reviewed changes
Copy link
Contributor

@skyw skyw left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Recommend to put everything into nununiform_tp.py, inherit from core classes and override member function when needed. That way code can be non-intrusive.

megatron/core/distributed/distributed_data_parallel.py Outdated
param.main_grad.add_(param.grad.data)
param.grad = None

# Nonuniform TP: gather grads from spare GPUs and scatter to core GPUs
Copy link
Contributor

@skyw skyw Dec 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Inherit from DDP, make a new class and override _make_backward_post_hook().

megatron/core/distributed/distributed_data_parallel_config.py Outdated
delay_wgrad_compute: bool = False
"""Delay the weight gradient computation to improve batch-level communication overlapping"""

tp_base: int = 8
Copy link
Contributor

@skyw skyw Dec 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Make a small config class just for NTP.

megatron/core/distributed/nonuniform_tp.py Outdated
ntp_map(layer.mlp, ddp_config, layer.mlp.config.ffn_hidden_size)


def test_ntp():
Copy link
Contributor

@skyw skyw Dec 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

move test to the right place under /tests.

megatron/core/distributed/param_and_grad_buffer.py Outdated
communication_group = self.data_parallel_group

# Coalesce communication kernels across buckets in the bucket group.
# NOTE: only sync on core GPUs (not spares) for nonuniform TP
Copy link
Contributor

@skyw skyw Dec 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also subclass then override start_grad_sync

megatron/core/optimizer/optimizer.py Outdated
if hasattr(param, 'main_grad'):
param.grad = param.main_grad
# NOTE: need to make this contiguous for nonuniform TP
param.grad = param.main_grad.contiguous()
Copy link
Contributor

@skyw skyw Dec 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

move this to sync grad, don't touch widely used code.

megatron/core/parallel_state.py Outdated
nccl_comm_cfgs[pg_name][key_value_pair[0]] = key_value_pair[1]


def _get_active_ranks_for_ntp(
Copy link
Contributor

@skyw skyw Dec 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Don't touch parallel_state, move to NTP files.

megatron/core/parallel_state.py Outdated
if rank in ranks:
_MODEL_PARALLEL_GROUP = group
_MODEL_PARALLEL_GLOBAL_RANKS = ranks
_CONTEXT_PARALLEL_GROUP = group
Copy link
Contributor

@skyw skyw Dec 5, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

update to latest, looks like you branched from very old code?

@daiyaanarfeen daiyaanarfeen force-pushed the nonuniform-tp branch 2 times, most recently from 69ea097 to c4a1601 Compare February 12, 2026 20:36
asolergi-nv and others added 16 commits February 23, 2026 10:02
@asolergi-nv @ko3n1g @daiyaanarfeen
Set token_dtype_code init value in GPTDatasetConfig to fix CI (NV…
789a7c1 
…IDIA#2912)

Co-authored-by: oliver könig <okoenig@nvidia.com>
@parthmannan @xiaoyao0115
@kunlunl @daiyaanarfeen
Hybrid Context Parallel Feature (NVIDIA#2282)
87c470c 
Signed-off-by: tailaim <tailaim@nvidia.com>
Signed-off-by: Parth Mannan <pmannan@nvidia.com>
Co-authored-by: Mcore Bot <mcore-bot@nvidia.com>
Co-authored-by: tailaim <tailaim@nvidia.com>
Co-authored-by: kunlunl <kunlunl@nvidia.com>
Co-authored-by: Kunlun Li <94586211+kunlunl@users.noreply.github.com>
@yuzhongw-nvidia @daiyaanarfeen
[main] ci(moe): Add --apply-wd-to-qk-layernorm flag to the gdn test…
5fc44c0 
… case (NVIDIA#2995)
@ko3n1g @daiyaanarfeen
ci: Disable step time on `gpt3_moe_mcore_te_tp2_pp2_ep4_etp1_no_mtp_n… (
80bea7f 
NVIDIA#2991)

Signed-off-by: oliver könig <okoenig@nvidia.com>
@ko3n1g @daiyaanarfeen
ci: Fix workflows on main (NVIDIA#2990)
102f16a 
Signed-off-by: oliver könig <okoenig@nvidia.com>
@shjwudp @cspades @daiyaanarfeen
Make Megatron-FSDP torch.compile compatible (NVIDIA#2425)
30e0260 
Co-authored-by: Cory Ye <44509866+cspades@users.noreply.github.com>
@cspades @daiyaanarfeen
[Megatron-FSDP] Test FP8 activations + parameter sharding with Megatr…
87d5f3f 
…on-FSDP fully-shard. Update README. (NVIDIA#2894)

Signed-off-by: Cory Ye <cye@nvidia.com>
@ko3n1g @daiyaanarfeen
chore: Escape special chars (NVIDIA#3014)
babed52 
Signed-off-by: oliver könig <okoenig@nvidia.com>
@deepakn94 @daiyaanarfeen
Improve memory logging (NVIDIA#2839)
f0586ea 
Signed-off-by: Deepak Narayanan <dnarayanan@nvidia.com>
@github-actions @daiyaanarfeen
chore: rotate oncall schedule
fcdaf0e
@santhnm2 @daiyaanarfeen
Add a wrapper function for FA3 _flash_attn_forward call (NVIDIA#2933)
31bc1e1 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@ko3n1g @daiyaanarfeen
ci(hotfix): Fix unit tests on main
4ef5e0e 
Signed-off-by: oliver könig <okoenig@nvidia.com>
@ko3n1g @daiyaanarfeen
ci(hotfix): Tests on main
fc2b789 
Signed-off-by: oliver könig <okoenig@nvidia.com>
@ko3n1g @daiyaanarfeen
chore: Set umask 0002 (NVIDIA#3027)
cdf232b 
Signed-off-by: oliver könig <okoenig@nvidia.com>
@mathemakitten @daiyaanarfeen
Make attn mask inversion in-place instead of allocating it again (NVI…
be7d033 
…DIA#3019)
@cspades @daiyaanarfeen
[Megatron-FSDP] Fix incorrect gradient scaling target. (NVIDIA#3023)
ec7dc57 
Signed-off-by: Cory Ye <cye@nvidia.com>
janEbert and others added 29 commits February 23, 2026 10:02
@janEbert @ko3n1g @daiyaanarfeen
Switch oncall (NVIDIA#3360)
1b4d434 
Co-authored-by: oliver könig <okoenig@nvidia.com>
@AAnoosheh @daiyaanarfeen
Add missing RMSNorm to llama train script (NVIDIA#3314)
7627a37 
Signed-off-by: Asha Anoosheh <aanoosheh@nvidia.com>
@tdene @daiyaanarfeen
Fix inference for MTP models (NVIDIA#3297)
95e2c9f
@yobibyte @daiyaanarfeen
Add a logprobs test with real gpt model. (NVIDIA#2870)
586a639
@tdene @daiyaanarfeen
Add simple GRPO functional test (NVIDIA#3323)
9416c87
@ko3n1g @daiyaanarfeen
ci: Concurrency control for merge-queue (NVIDIA#3353)
86d1927 
Signed-off-by: oliver könig <okoenig@nvidia.com>
@chtruong814 @daiyaanarfeen
ci: Update golden value download script to work with Github (NVIDIA#3335
8876190 
)

Signed-off-by: Charlie Truong <chtruong@nvidia.com>
@thecaptain789 @Phlip79 @daiyaanarfeen
fix: correct typos 'seperated' and 'recieved' (NVIDIA#3305)
5533713 
Co-authored-by: thecaptain789 <thecaptain789@users.noreply.github.com>
Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>
@shengf-nv @daiyaanarfeen
Improved PyTorch profiler and added PyTorch execution trace (NVIDIA#3273
71dcccb 
)
@megnvidia @daiyaanarfeen
Removing etc from main index page, shifted name of discussions (NVIDI…
5df6e82 
…A#3271)

Signed-off-by: meg miranda <mmiranda@nvidia.com>
@ko3n1g @daiyaanarfeen
build: Bump TE on 2.12 (NVIDIA#3371)
36acfce 
Signed-off-by: oliver könig <okoenig@nvidia.com>
@ko3n1g @daiyaanarfeen
ci(hotfix): job conditions (NVIDIA#3376)
7113c2e 
Signed-off-by: oliver könig <okoenig@nvidia.com>
@github-actions @daiyaanarfeen
Update copy-pr-bot.yaml [skip ci]
9b651a1
@sidsingh-nvidia @ISEEKYAN @daiyaanarfeen
Record moe routing decisions during inference. (NVIDIA#3034)
b3e7f03 
Co-authored-by: litianjian <litianjian@bytedance.com>
Co-authored-by: Yan Bai <baiyan1996@icloud.com>
@Wohox @dimapihtar @daiyaanarfeen
[Main] Fix EP Overlap Bugs for Full-Iter CG (NVIDIA#3164)
39e7e73 
Co-authored-by: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com>
@maanug-nv @daiyaanarfeen
Avoid direct pickle import (NVIDIA#3375)
2cf97a2 
Signed-off-by: Maanu Grover <maanug@nvidia.com>
@Phlip79 @daiyaanarfeen
Delete old pretrain_* files (NVIDIA#3359)
3139c39
@xuwchen @daiyaanarfeen
Add Qwen3-VL support with Megatron-FSDP (NVIDIA#2841)
c7948f7
@santhnm2 @daiyaanarfeen
Refactor Mamba chunked prefill (NVIDIA#3265)
216c12a 
Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>
@jstjohn @ko3n1g @daiyaanarfeen
Improved parallel logging of learning rate (NVIDIA#3319)
e88f2b6 
Co-authored-by: oliver könig <okoenig@nvidia.com>
@lmcafee-nvidia @claude @daiyaanarfeen
Add enhanced event tracking with TTFT measurement and compact seriali…
559834e 
…zation. (NVIDIA#3253)

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
@daiyaanarfeen
feat: non-intrusive nonuniform tensor parallelism implementation
916ac17 
- All NTP logic contained in nonuniform_tp.py as subclasses
- NonuniformTPDistributedDataParallel: inherits from DistributedDataParallel
- NonuniformTPParamAndGradBuffer: handles gradient buffer splitting for NTP
- NonuniformTPOptimizer: wrapper for gradient contiguity
- initialize_nonuniform_tp_process_groups(): reconfigures process groups after init
- Only config changes to core files (distributed_data_parallel_config.py)
- Added comprehensive CLAUDE.md documentation
@daiyaanarfeen
test: add comprehensive unit tests for nonuniform TP
287fd1a 
- Moved test from nonuniform_tp.py to tests/unit_tests/distributed/
- Added TestNonuniformTPUtilities: tests for utility functions
  - compute_uniform_tp_spares_with_parity (3 test cases)
  - get_active_ranks_for_dp (2 test cases)
- Added TestNonuniformTPParameterResharding: tests for parameter resharding
  - ntp_map for no spares, healthy ranks, unhealthy ranks
  - ntp_init for layers with attention and MLP (4 test cases)
- Added TestNonuniformTPOptimizer: tests for optimizer wrapper
  - attribute delegation, prepare_grads, contiguity handling (5 test cases)
- Added TestNonuniformTPIntegration: integration tests
  - DDP initialization and backward hooks (2 test cases)
- Total: 17 test cases covering all major NTP functionality
@daiyaanarfeen
fix: remove read-only shape attribute assignment in test
c7b34d8 
PyTorch tensors have shape as a read-only property, no need to set it
@daiyaanarfeen
fix: remove another shape assignment in test_ntp_backward_hook_core_gpu
c049768
@daiyaanarfeen
test: add end-to-end NTP test for 8 GPUs without mocking
1e87b35 
- Tests 2 DP workers: DP rank 0 with TP=2 (reduced), DP rank 1 with TP=4 (healthy)
- Uses tp_base=4, tp_spares=2 configuration
- Verifies process group reconfiguration
- Tests parameter initialization and gradient computation
- No mocking - actual distributed test with real model
@daiyaanarfeen
fix: skip spare ranks gracefully in end-to-end test
b60f495 
Spare ranks would call sys.exit(0) during NTP initialization,
which pytest treats as a failure. Now spare ranks skip the test
gracefully before that happens.
@daiyaanarfeen
refactor: move NTP config to separate NonuniformTPConfig class
51d98b2 
- Create NonuniformTPConfig dataclass in nonuniform_tp.py
- Remove NTP fields from DistributedDataParallelConfig (non-intrusive)
- Update all NTP functions/classes to use NonuniformTPConfig
- Update all tests to use NonuniformTPConfig
- Update CLAUDE.md documentation

This makes the NTP implementation completely self-contained with
zero modifications to core Megatron files.
@daiyaanarfeen
chore: remove unrelated changes from distributed_data_parallel_config.py
e216358 
Restore file to match upstream main exactly, removing:
- Unused typing imports (Dict, List, Tuple)
- fsdp_db_use_persist_buf_on_alloc_fail field
- fsdp_manual_registration field
- Docstring indentation changes

These were unrelated to NTP and came from other branches.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@skyw skyw skyw left review comments

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

Core 0.16

Development

Successfully merging this pull request may close these issues.