Hungyuehc/omniml 5029 by hychiang-git · Pull Request #2 · ChenhanYu/Megatron-LM

hychiang-git · 2026-06-10T17:09:07Z

This is the patch for OMNIML-5029.
Please merge #1 first where we sync the ChanhanYu to the upstream NVIDIA's branch.

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Co-authored-by: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com>

Signed-off-by: oliver könig <okoenig@nvidia.com>

NVIDIA#4403) Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>

…IA#4227) Signed-off-by: Maanu Grover <maanug@nvidia.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

…4301) Co-authored-by: Siddharth Singh <sidsingh@nvidia.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…nference cuda graph scope for hybrid models (NVIDIA#4440)

…ss curve gaps for latent MoE models (NVIDIA#4433) Signed-off-by: root <jiemingz@nvidia.com>

…VIDIA#4158) Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

…VIDIA#4422) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Signed-off-by: rprenger <rprenger@nvidia.com>

Signed-off-by: qiyuw <qiyuw@nvidia.com> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

Co-authored-by: Siddharth Singh <sidsingh@nvidia.com>

…VIDIA#4427) Signed-off-by: Cory Ye <cye@nvidia.com>

Co-authored-by: root <root@eos0047.eos.clusters.nvidia.com> Co-authored-by: root <root@eos0260.eos.clusters.nvidia.com> Co-authored-by: Dennis(Zhenhuan) Liu <denliu@nvidia.com>

… (NMFW-17) (NVIDIA#4368) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…VIDIA#4330) Co-authored-by: mhh111 <mahonghao1@huawei.com> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

…ing (NVIDIA#4276) Co-authored-by: Hanpeng Hu <haaanpeng@outlook.com> Co-authored-by: Deepak Narayanan <deepakn94@gmail.com> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

…h tests (NVIDIA#4986) Signed-off-by: oliver könig <okoenig@nvidia.com>

… nightlies (NVIDIA#4905)

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Co-authored-by: Robin Zhang <robinz@nvidia.com> Co-authored-by: Dennis Liu <denliu@nvidia.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com> Co-authored-by: Shifang Xu <shifangx@nvidia.com>

Signed-off-by: oliver könig <okoenig@nvidia.com>

…port full-iteration (FWD-BWD) CUDA graphability. (NVIDIA#4663) Signed-off-by: Cory Ye <cye@nvidia.com>

Signed-off-by: oliver könig <okoenig@nvidia.com>

…VIDIA#5022) Signed-off-by: oliver könig <okoenig@nvidia.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…A#5023) Signed-off-by: oliver könig <okoenig@nvidia.com>

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

) Signed-off-by: oliver könig <okoenig@nvidia.com>

Signed-off-by: Maanu Grover <maanug@nvidia.com>

…put_parameter (NVIDIA#5034)

…A#4877)

…VIDIA#4881) Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>

…VIDIA#5040) Signed-off-by: oliver könig <okoenig@nvidia.com>

…DIA#4938) Signed-off-by: Li Ding <liding@nvidia.com>

Signed-off-by: oliver könig <okoenig@nvidia.com>

NVIDIA#5045) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

…sion__ is_nvrx_min_version() unconditionally reads nvrx.__version__, but older nvidia_resiliency_ext builds (notably the one bundled in the NeMo 25.11.nemotron_3_nano container image) don't expose that attribute. The result is an AttributeError at pytest collection time for any test that imports megatron.core.dist_checkpointing, well before the test logic even runs. Switch the read to getattr(nvrx, '__version__', '0.0.0') so the fallback path is reached without crashing. Behavior for builds that DO expose __version__ is unchanged. Encountered in OMNIML-5029 when bind-mounting our packaged Megatron-LM over the container's older one to enable test_moe_sharded_state_dict. Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>

…rting has_nvrx_async_support() asserts is_nvrx_min_version() is True, which crashes the entire dist_checkpointing import chain when the installed nvidia-resiliency-ext is older than the 0.6.0 floor — or, more subtly, when it lacks __version__ entirely (the case on nemo:25.11.nemotron_3_nano, where is_nvrx_min_version() conservatively returns "0.0.0"). Replace the assert with an early `return False` so the function obeys its own contract ("returns whether NVRx async checkpointing is usable") without crashing. The caller already has a non-async fallback path; it just needs to be told that async isn't available. Pairs with the earlier getattr(nvrx, '__version__', '0.0.0') guard so the import succeeds even when the attribute is missing, and the downstream version-floor check becomes recoverable instead of fatal. Signed-off-by: Hung-Yueh Chiang <hungyuehc@nvidia.com>

Phlip79 and others added 30 commits April 23, 2026 04:27

Rename Mamba to Hybrid outside megatron/core (NVIDIA#4159)

60408d5

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Include mtp layers in token per expert logging (NVIDIA#4412)

a52014c

fix: NVRx async compatibility and defer resiliency import (NVIDIA#4420)

32275b2

Co-authored-by: Dmytro Pykhtar <37850217+dimapihtar@users.noreply.github.com>

ci: add base_sha to codecov/codecov-action upload step (NVIDIA#4445)

9bb35a8

Signed-off-by: oliver könig <okoenig@nvidia.com>

Update copy-pr-bot.yaml [skip ci]

3034d86

fix(checkpoint_inspector): allow empty --param-to-param-group-map-json (

f78ed05

NVIDIA#4403) Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

Add the YARN support for hybrid_model (NVIDIA#4244)

4d6cdd5

Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>

[training migration] Add container class for config dataclasses (NVID…

41ffa83

…IA#4227) Signed-off-by: Maanu Grover <maanug@nvidia.com> Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>

Inference: Fix broken functional tests on gitlab (NVIDIA#4454)

a1165fa

SafeUnpickler class for safe pickle usage (NVIDIA#4319)

d4cacef

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

get rid of weights_only=False (NVIDIA#4434)

109feda

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

Inference | Per-block MoE routing storage for prefix caching (NVIDIA#…

64870c1

…4301) Co-authored-by: Siddharth Singh <sidsingh@nvidia.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Add troubleshooting tip for 'access forbidden' (NVIDIA#4449)

017e684

Fix checkpoint loading with rerun state machine (NVIDIA#4448)

3d7bcd3

Add misc CUDA graph sugar to CudaGraphManager (NVIDIA#4425)

9b02206

Inference: Add the embedding and output layer in the full_iteration_i…

35f76df

…nference cuda graph scope for hybrid models (NVIDIA#4440)

Important bugfixes in local CG implementation that were leading to lo…

481efd0

…ss curve gaps for latent MoE models (NVIDIA#4433) Signed-off-by: root <jiemingz@nvidia.com>

fix: Replace polynomial rolling hash with SHA-256 for prefix caching (N…

e9abb6c

…VIDIA#4158) Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

feat(ckpt): expose validate_access_integrity knob on dist-ckpt load (N…

377af02

…VIDIA#4422) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Fix multivalidation (NVIDIA#3388)

241a5ca

Signed-off-by: rprenger <rprenger@nvidia.com>

Add missing knob for reduce_scatter_with_fp32_accumulation (NVIDIA#4410)

f2dcd42

Signed-off-by: qiyuw <qiyuw@nvidia.com> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

Enable CUDA graphs for MTP inference (NVIDIA#4260)

03f4111

Signed-off-by: Keshav Santhanam <ksanthanam@nvidia.com>

chore(beep boop 🤖): Bump (main) (2026-04-27)

1879dc2

checkpoint integrity verification (NVIDIA#4305)

970c254

Signed-off-by: dimapihtar <dpykhtar@nvidia.com>

Fix cache gating (NVIDIA#4455)

ebd70d3

Co-authored-by: Siddharth Singh <sidsingh@nvidia.com>

[Main] Fix FusedAdam.use_decoupled_grad mis-set for Megatron-FSDP. (N…

0447347

…VIDIA#4427) Signed-off-by: Cory Ye <cye@nvidia.com>

add permute fusion into hybrid ep (NVIDIA#4089)

8c5cf05

Co-authored-by: root <root@eos0047.eos.clusters.nvidia.com> Co-authored-by: root <root@eos0260.eos.clusters.nvidia.com> Co-authored-by: Dennis(Zhenhuan) Liu <denliu@nvidia.com>

Add ColocatedBridgeCommunicator for heterogeneous TP/DP MIMO training…

42e396e

… (NMFW-17) (NVIDIA#4368) Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Fix incorrect bias display in extra_repr of Column/RowParallelLinear (N…

6fd6652

…VIDIA#4330) Co-authored-by: mhh111 <mahonghao1@huawei.com> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

Fix assertion logic in combined_1f1b_schedule_for_interleaved_pipelin…

c8a4bfd

…ing (NVIDIA#4276) Co-authored-by: Hanpeng Hu <haaanpeng@outlook.com> Co-authored-by: Deepak Narayanan <deepakn94@gmail.com> Co-authored-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>

ko3n1g and others added 30 commits May 26, 2026 17:23

fix(tests): initialize num_microbatches calculator in vision cudagrap…

6ce6fac

…h tests (NVIDIA#4986) Signed-off-by: oliver könig <okoenig@nvidia.com>

ci: Add allow_failure flag to gpt and moe recipes that are failing in…

859b719

… nightlies (NVIDIA#4905)

Drain predecessor reduce-scatter at dispatch time (NVIDIA#4940)

88e7ab0

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

nightly(ci): Update golden values for functional t5 tests (NVIDIA#4995)

e6b2bd8

chore: rotate oncall schedule

7521ecb

[main] Refactor and Improve MoE Logginginit commit (NVIDIA#3431)

873678a

Co-authored-by: Robin Zhang <robinz@nvidia.com> Co-authored-by: Dennis Liu <denliu@nvidia.com> Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com> Co-authored-by: Shifang Xu <shifangx@nvidia.com>

ci: validate release branch-rules (NVIDIA#4929)

4e52a9e

Signed-off-by: oliver könig <okoenig@nvidia.com>

[Megatron-FSDP] Add conditional param.grad dereferencing logic to sup…

67b2f38

…port full-iteration (FWD-BWD) CUDA graphability. (NVIDIA#4663) Signed-off-by: Cory Ye <cye@nvidia.com>

test: restrict iter-time comparison to steady-state window (NVIDIA#5010)

71223d5

Signed-off-by: oliver könig <okoenig@nvidia.com>

fix(test): pin eval-global-batch-size on 15b gb200 release configs (N…

a6d61fb

…VIDIA#5022) Signed-off-by: oliver könig <okoenig@nvidia.com> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

[fix] Release MTP assertion when EP overlap with PP=1 (NVIDIA#4796)

286445c

fix(test): widen iter-time steady-state window for short tests (NVIDI…

0cb4034

…A#5023) Signed-off-by: oliver könig <okoenig@nvidia.com>

Perf fix (NVIDIA#4996)

805e24e

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Add dev-feature preservation gate and change schedule (NVIDIA#4773)

457e3f7

chore(test): remove orphan nemotron3_super_release_g200 dir (NVIDIA#5024

146d171

) Signed-off-by: oliver könig <okoenig@nvidia.com>

Ignore Claude worktree directory (NVIDIA#5020)

7be1748

Update copy-pr-bot.yaml [skip ci]

9d23a73

ci: update CI workflow conditions for integration tests (NVIDIA#4658)

3e00820

Add NVSkills CI request workflow (NVIDIA#5033)

0010683

DDP wrap pg size fixes (NVIDIA#5006)

2a4f820

Signed-off-by: Maanu Grover <maanug@nvidia.com>

fix(layer_wise): tag MTP-stage word_embeddings as is_embedding_or_out…

e0afe89

…put_parameter (NVIDIA#5034)

Move LTS dependencies from pyproject.toml to Dockerfile.ci.lts (NVIDI…

4f2e7ce

…A#4877)

Use shared ModelOpt calibration loop on 0.45+ with 0.44 fallback fix (N…

88efdf7

…VIDIA#4881) Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>

test(release): skip golden comparison on intermediate resume windows (N…

a5c2d1b

…VIDIA#5040) Signed-off-by: oliver könig <okoenig@nvidia.com>

[mimo] Thread position_ids through MimoModel for multimodal RoPE (NVI…

2ebca1e

…DIA#4938) Signed-off-by: Li Ding <liding@nvidia.com>

build: Switch DSv3 on H100 to HybridEP (NVIDIA#5039)

e6ff4c3

Signed-off-by: oliver könig <okoenig@nvidia.com>

Fix: Import unwrap_model from megatron.core.utils in modelopt examples (

ba9e0ea

NVIDIA#5045) Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Simple and stable Inference APIs (NVIDIA#4697)

3c39d98

Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Navigation Menu

Hungyuehc/omniml 5029 - #2

Hungyuehc/omniml 5029#2
hychiang-git wants to merge 310 commits into
ChenhanYu:mainfrom
hychiang-git:hungyuehc/omniml-5029

hychiang-git commented Jun 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

20 participants

Conversation

hychiang-git commented Jun 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

20 participants