Skip to content

[model] fix: Fix ERNIE VL L1 tests - #4078

Merged
yaoyu-33 merged 2 commits into
mainfrom
chcui/fix-ernie-vl-l1-tests
May 30, 2026
Merged

[model] fix: Fix ERNIE VL L1 tests#4078
yaoyu-33 merged 2 commits into
mainfrom
chcui/fix-ernie-vl-l1-tests

Conversation

@cuichenx

@cuichenx cuichenx commented May 29, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Fix ERNIE VL dual-pool MoE construction under the 26.06 line by accepting TransformerLayer kwargs and forwarding pg_collection, is_mtp_layer, and module names into the inner MoE/shared-expert modules.
  • Fix EP local-expert renumbering for nested ERNIE VL MoE paths such as vision_moe_layer.experts.local_experts.*.
  • Point the ERNIE VL fwd/bwd functional test at the existing example helper and add an H100 L1 launch script.

Related: original ERNIE VL support PR #3263

Implementation Note

ERNIE 4.5 VL still uses SequentialMLP-style experts.local_experts.N paths for its dual-pool MoE; this PR does not remove that support. The separate SequentialMLP fallback block in _megatron_local_name_to_global() was removed because it was unreachable for the nested ERNIE VL paths it was meant to cover: the old is_expert_param check only matched .mlp.experts.local_experts., so ep_group was left as None for paths like .mlp.vision_moe_layer.experts.local_experts.N..., and the fallback's ep_group is not None guard prevented it from running.

The fix broadens the initial EP detection to .experts.local_experts. and lets the main local-expert renumbering path handle both standard SequentialMLP paths and nested ERNIE VL text/vision MoE paths. The new unit coverage checks that decoder.layers.0.mlp.vision_moe_layer.experts.local_experts.1.linear_fc1.weight maps to the correct global expert rank under EP.

Validation

  • git diff --check
  • python3 -m py_compile src/megatron/bridge/models/conversion/model_bridge.py src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/ernie_moe_layer.py tests/unit_tests/models/test_qat_bridge_support.py tests/unit_tests/models/ernie_vl/test_ernie45_vl_bridge.py tests/functional_tests/test_groups/models/ernie_vl/test_ernie45_vl_conversion.py examples/models/vlm/ernie_vl/ernie45_vl_fwd_bwd.py
  • bash -n tests/functional_tests/launch_scripts/h100/active/L1_Launch_models_ernie_vl.sh
  • uv run --no-sync pre-commit run --all-files
  • DFW H100 interactive validation with nvcr.io/nvidian/nemo:26.06.rc1: job 12307622, completed 0:0, log /lustre/fs1/portfolios/coreai/projects/coreai_dlalgo_llm/users/chcui/logs/ernie-vl-l1-fix/ernie_vl_l1_fix_12307622.log
    • tests/unit_tests/models/ernie_vl/test_ernie45_vl_bridge.py tests/unit_tests/models/test_qat_bridge_support.py: 51 passed
    • tests/functional_tests/test_groups/models/ernie_vl/test_ernie45_vl_conversion.py: 8 passed, including TP/PP/EP conversion and single-GPU/EP2 text+vision fwd/bwd cases

uv run pre-commit run --all-files was also attempted, but this host cannot resolve the locked nvidia-resiliency-ext==0.6.0 wheel for manylinux_2_31_x86_64; the package currently publishes manylinux_2_39 wheels.

Blast Radius / Test Recommendation

Changed surface is ERNIE VL MoE construction, shared EP expert-name conversion mapping, and ERNIE VL tests/launch coverage. L0 targeted unit + functional coverage passed. This PR adds an H100 L1 launcher and is labeled needs-more-tests to exercise the L1 CI path. L2 is not required unless CI reports a broader model/runtime regression.

Signed-off-by: Chen Cui <chcui@nvidia.com>
@cuichenx cuichenx added the needs-more-tests Requires additional L0 and L1 test coverage before merge label May 29, 2026
Comment on lines +1 to +2
# CI_TIMEOUT=60
#!/bin/bash

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Nit: #!/bin/bash on line 2 is a no-op — shebangs only work on line 1. Other CI_TIMEOUT scripts (e.g., L0_Launch_megatron_fsdp.sh) don't include a shebang at all since CI invokes them via bash explicitly. Remove it to stay consistent:

Suggested change
# CI_TIMEOUT=60
#!/bin/bash
# CI_TIMEOUT=60

@claude

claude Bot commented May 29, 2026

Copy link
Copy Markdown
Contributor

Light Code Review — The core logic changes look correct. Broadening the expert-param detection from .mlp.experts. to .experts. properly handles nested MoE paths (e.g., vision_moe_layer.experts.local_experts.*) that the old code failed to match. The deleted elif block was dead code -- ep_group was always None when it was reached because the narrower is_expert_param check never fired for nested paths, so _get_ep_group was never called. The ErnieMultiTypeMoE constructor changes correctly forward pg_collection, is_mtp_layer, and name to child modules, matching Megatron-Core TransformerLayer conventions. Finding: The shebang on line 2 of L1_Launch_models_ernie_vl.sh is a no-op (shebangs only work on line 1). See inline comment. Test Coverage: Good -- test_nested_local_expert_name_maps_to_global_expert_rank validates broadened EP renumbering, test_accepts_transformer_layer_kwargs validates kwargs forwarding, and the functional test path is updated. Suggested test cases: No perf tests impacted.

@yaoyu-33
yaoyu-33 merged commit 1ace3e2 into main May 30, 2026
139 checks passed
@yaoyu-33
yaoyu-33 deleted the chcui/fix-ernie-vl-l1-tests branch May 30, 2026 03:51
vasunvidia pushed a commit to vasunvidia/Megatron-Bridge that referenced this pull request Jun 10, 2026
Signed-off-by: Chen Cui <chcui@nvidia.com>
Signed-off-by: Vasudevan Rengasamy <vrengasamy@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:model Model implementations and HF bridge logic bug Something isn't working needs-more-tests Requires additional L0 and L1 test coverage before merge needs-review PR is ready for code review and waiting on a reviewer

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants