[dev] lite: Sync Megatron Lite experimental updates for deepseek v4 - #6397
Conversation
Signed-off-by: Yan Bai <bayan@nvidia.com>
Signed-off-by: Yan Bai <bayan@nvidia.com>
Signed-off-by: Yan Bai <bayan@nvidia.com>
737381c to
58bcc78
Compare
pytest requires `pytestmark` to be a Mark or a list of them; a tuple raises `TypeError: got (MarkDecorator(...), MarkDecorator(...)) instead of Mark` during collection, which fails the whole run before any test executes. Observed on one 8xH100 node with `tests/run_tests.sh` against this branch: 878 tests collected, 2 errors -- both this. Every other file in the tree that carries multiple marks already uses a list. Signed-off-by: Yan Bai <bayan@nvidia.com>
f610632 to
5ee3345
Compare
`experimental/lite/examples/` and Megatron-LM's top-level `examples/` are both
regular packages, so whichever appears first on `sys.path` wins. The harness set
`PYTHONPATH=REPO_ROOT:LITE_ROOT`, which was harmless while lite lived in a
lite-only fork whose root holds nothing but README.md and experimental/. Inside
the full tree the root package wins and lite's own examples become unreachable:
ModuleNotFoundError: No module named 'examples.bench'
(tests/unit/runtime/backends/bridge/test_packed_batch_bridge.py)
The worker is launched as `[sys.executable, worker, *args]`, i.e. a script, so
`sys.path[0]` is the harness directory and the working directory never
participates in resolution. Changing the subprocess cwd therefore does not help;
the PYTHONPATH order is what decides.
Verified with `importlib.util.find_spec` on one 8xH100 node with the lite root
first:
megatron.core -> <root>/megatron/core/__init__.py
megatron.lite -> <lite>/megatron/lite/__init__.py
tests -> <root>/tests/__init__.py
`megatron` is unaffected because neither directory carries an `__init__.py`, so
both are implicit namespace packages and merge. `tests` is unaffected because
lite's `tests/` has no `__init__.py`, so the root's regular package still wins.
Only `examples` changes hands, which is the intent.
Signed-off-by: Yan Bai <bayan@nvidia.com>
5ee3345 to
db83120
Compare
Signed-off-by: Yan Bai <bayan@nvidia.com>
The merge kept dev's `add_cross_entropy_fusion`, which reads the flag with a plain `getattr(model, "cross_entropy_fusion", False)`, and dropped the fork's `_resolve_cross_entropy_fusion`, which walks the `.module` chain. Megatron's `_BaseDataParallel` keeps the wrapped module in `self.module` and does not proxy unknown attributes, so once the runtime wraps each chunk the plain getattr misses the flag and silently takes the unfused branch -- materialising the full `[tokens, vocab/tp]` logits in fp32 while every config dump still reports `cross_entropy_fusion=True`. Also updates two contract tests that pin surfaces dev has since moved: - `test_lite_csa_imports_core_csa_kernel_namespace` matched the literal `...experimental_attention_variant.csa_kernels import`, which no longer exists after dev's CSA restructure (1c44a57); it now pins the `csa_utils` package. - `test_attention_public_api_is_narrow` gained `MagiAttentionConfig` and `MagiDotProductAttention`, the two names MagiAttention (NVIDIA#6478) exports. Both are genuine public surface, unlike the internal helpers the same test asserts stay unexported. Signed-off-by: Yan Bai <bayan@nvidia.com>
`test_gqa_magi_branch_uses_dispatched_positions_and_runtime` builds a `GQAttention` through `__new__` and assigns the attributes the forward pass needs. It came from upstream NVIDIA#6478, where `GQAttention` has no `_replicate_kv`; lite's copy does, because it replicates KV heads when `num_key_value_heads < tp_size`, and every forward path reads the flag. After the merge the stub is missing it: AttributeError: 'GQAttention' object has no attribute '_replicate_kv' Set it to False, matching the test's 4 query / 2 KV head configuration, which is what the real `__init__` would compute. No assertion is relaxed. Signed-off-by: Yan Bai <bayan@nvidia.com>
`_test_harness/markers.py` treats a missing `gpus` marker as "this is a CPU
test", so all five CUDA tests in this file were scheduled into the `gpus=0` cpu
suite, where `torch.cuda.is_available()` is false and they skipped.
`pytest_worker.py` counts any skip as a failure and `runner.py` stops at the
first non-PASS suite, so the cpu suite reported FAIL with `736 passed, 0 failed,
5 skipped` and every GPU suite after it never ran at all.
Two of the five call `_require_two_ranks()` and assert
`dist.get_world_size() == 2`, so they need `gpus(2)`; `_gpus_marker` reads
`next(item.iter_markers("gpus"))`, closest-first, so a per-test marker overrides
the module-level one. The sibling NCCL test in the fsdp2 file already uses the
same pattern.
Signed-off-by: Yan Bai <bayan@nvidia.com>
5fe9601 to
a850af5
Compare
|
Brings Only Why now
This PR fixes that as part of the sync. How the merge was doneThe two trees last agreed at c8239f9 (#5862, 2026-07-20). Merging from the
What arrives from the fork
Conflict resolution worth reviewing closelySix source files were changed on both sides. Two of them had a defect in the
The other four ( Test harnessTwo harness changes were needed for the suite to run inside the full repository
ValidationOne 8xH100 node, Nothing fails. The one non-PASS suite is an optional-dependency skip: A second run with the FLA stack present (FLA 0.5.0) confirms the file still does Unmodified |
|
/ok to test a850af5 |
`tools/check_copyright.py` requires every file to open with the current-year NVIDIA header, and CI applies it to the files a PR touches. Four Python files in this change did not: - `verl_mlite/config/ref/__init__.py` carried the header on `dev` and lost it in the merge -- the fork's copy of this file is empty. Restored. - `tests/unit/model/test_deepseek_v4_hf_export_dtype.py` is new and had none. - `primitive/kernels/dsa_kernels.py` still had the older `Copyright (c) 2025, NVIDIA CORPORATION.` form; bumped to the current wording. - `examples/verl/scripts/validate_deepseek_v4_dapo.py` starts with a shebang. The header goes above it, matching `docs/add_copyright_header.py`, because the check tests `content.startswith(...)`. Shell and YAML files touched here open with a shebang or a plain comment on `dev` as well and are left as they are. Signed-off-by: Yan Bai <bayan@nvidia.com>
Keeps the branch current with the base. dev's only new commit (d59d4b9, paged stash with TE whole-MoE) does not touch experimental/lite, so this is a fast-forward of the non-lite tree. Signed-off-by: Yan Bai <bayan@nvidia.com>
|
/ok to test 3888693 |
2 participants
Summary
This PR syncs the current
experimental/**Megatron Lite tree fromISEEKYAN/Megatron-LM:mainonto NVIDIAdevin one PR with two reviewable commits.43124b60c26bb233bc9ba0cf75b2f637d0387c144145ea889e7035e2e6d9e34c709a87eb04e8ebb13833103b4ec7cb64f7f1109d317247f340ca6c6e737381c95c4982f67423b351781869b2b436a764experimental/lite/**Full diff: dev...ISEEKYAN:Megatron-LM:sync/fork-main-to-nvidia-dev
Consolidated diff: ISEEKYAN@3833103
DS4 QAT HF-load diff: ISEEKYAN@737381c
The source and target histories have diverged substantially, so retaining the original commits would also retain thousands of unrelated fork commits. The sync is therefore intentionally squashed; the original implementation notes remain linked below.
Diff contents
QAT and router replay
Dynamic context parallelism
Checkpoint and weight streaming
Model and training compatibility
veRL integration and validation
Original source notes
Credits
0968605a4f,476fda588d,710b0a454e).0cf711bcaa).Validation
git diff --check: passed.experimental/**: passed.114 passed, 14 skipped.23 passed.Known source-test gaps
No code changes were added while preparing this sync. Two existing fork tests are stale relative to the synced implementation:
test_all_model_qat_r3_contracts.pystub the Core CSA namespace without the newercsa_kernelsmodule.test_glm52_checkpoint_mapping_skips_shared_indexer_without_testill calls the removed private_load_attentionhelper instead of the currentGlm5WeightSpecmapping contract.These are reported as source-tree test gaps rather than being silently fixed in this sync PR.