Fix elastification unwrap_model import - #4972
Merged
ko3n1g merged 1 commit intoMay 25, 2026
Merged
Conversation
Contributor
|
/ok to test ea9ee79 |
Contributor
|
/ok to test ea9ee79 |
ko3n1g
approved these changes
May 25, 2026
Contributor
|
@Devil1716 thanks for the contribution! could you remove the import guard with this PR as well? https://github.com/NVIDIA/Megatron-LM/blob/main/tests/unit_tests/elastification/test_loss_func.py#L23 |
Devil1716
force-pushed
the
codex/fix-elastification-unwrap-model-import
branch
from
May 25, 2026 14:54
ea9ee79 to
f2d1c07
Compare
Contributor
Author
|
Thanks, done. I removed the module-level import guard from |
Contributor
|
/ok to test f2d1c07 |
Contributor
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/26416242210 |
santhnm2
pushed a commit
to santhnm2/Megatron-LM
that referenced
this pull request
May 26, 2026
Victarry
pushed a commit
to yanring/Megatron-LM
that referenced
this pull request
May 27, 2026
* origin/main: (50 commits) Drain predecessor reduce-scatter at dispatch time (NVIDIA#4940) ci: Add allow_failure flag to gpt and moe recipes that are failing in nightlies (NVIDIA#4905) fix(tests): initialize num_microbatches calculator in vision cudagraph tests (NVIDIA#4986) test: re-enable test_pp2_create_cudagraphs_first_stage on TE 2.15+ (NVIDIA#4985) ci: Add support for MBridge job gating based on PR labels (NVIDIA#4926) test(ci): re-enable 8experts2parallel_multi_dist_optimizer_instances_1node (NVIDIA#4984) test: re-enable paged stashing MoE tests (NVIDIA#4978) Fix elastification unwrap_model import (NVIDIA#4972) Avoid offsetting functional test master port (NVIDIA#4973) test: enable NVTE_CUTEDSL_FUSED_GROUPED_MLP via pytest fixture (NVIDIA#4931) chore(beep boop 🤖): Bump (main) (2026-05-25) test(release): add release goldens for deepseekv3/nemotron3 and set tp2pp2 exit-interval (NVIDIA#4932) Fix `get_batch` return order to ignore BlendedDataset provenance fields (NVIDIA#4952) ci: restore perf test torchrun logs (NVIDIA#4951) Various training utils (NVIDIA#4872) ci: Update training script paths in BERT and T5 (NVIDIA#4939) [MXFP8/FP4-param-gather] Post processing after forced param AG in eval (NVIDIA#4562) Fix mxfp8 param gather numerical issue when DP overlap is off (NVIDIA#4800) Add TEFusedDenseMLP for Dense+Grouped GEMM fusion on SM100+ (NVIDIA#4318) (NVIDIA#4786) Fix paged stashing test submodules lookup (NVIDIA#4925) ... # Conflicts: # megatron/training/training.py
janEbert
pushed a commit
to janEbert/Megatron-LM
that referenced
this pull request
Jun 2, 2026
mathemakitten
pushed a commit
to mathemakitten/Megatron-LM
that referenced
this pull request
Jun 12, 2026
yhgalaxy
pushed a commit
to yhgalaxy/Megatron-LM
that referenced
this pull request
Jun 17, 2026
Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>
jon-barker
pushed a commit
to jon-barker/Megatron-LM
that referenced
this pull request
Jul 10, 2026
Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>
terminator123
pushed a commit
to 021ai/Megatron-LM
that referenced
this pull request
Aug 3, 2026
svcnvidia-nemo-ci
pushed a commit
to dimapihtar/Megatron-LM
that referenced
this pull request
Aug 4, 2026
Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Fixes the elastification loss function import so
tests/unit_tests/elastification/test_loss_func.pycan be collected again afterunwrap_modelstopped being exported frommegatron.training.utils.Issue tracking
Linked issue: Fixes #4971
Contribution process
Pre-checks
Notes
This is a focused import-path fix plus the related test cleanup requested in review. Existing unit coverage already imports
megatron.elastification.loss_funcintests/unit_tests/elastification/test_loss_func.py, which is the collection path reported in #4971.Root cause:
megatron.elastification.loss_funcimportedunwrap_modelfrommegatron.training.utils, butmegatron/training/utils/__init__.pyno longer exports that symbol. The function is defined inmegatron.core.utils, and current callers throughout the repository already import it from there.The import guard in
tests/unit_tests/elastification/test_loss_func.pyhas also been removed so future import regressions fail during collection instead of being hidden by a module-level skip.Validation performed locally:
python -m py_compile megatron\elastification\loss_func.py tests\unit_tests\elastification\test_loss_func.pyrg "pytest\.skip\(|try:|except ImportError|importorskip" tests\unit_tests\elastification\test_loss_func.py -nreturned no matchesgit diff --checkI could not run
pytest tests/unit_tests/elastification/test_loss_func.pyin this local checkout because the local Python environment is missingpytestandtorch. CI should exercise the exact failing collection path.