[Victarry]

Summary

• Enable latent MoE to use shared expert overlap with A2A/Flex dispatchers by running shared-expert pre-comm on full hidden states in MoELayer, while keeping shared-expert FC launches inside dispatcher overlap windows.
• Keep dispatcher-owned shared-expert pre/post handling for the non-latent path only, and add the latent shared-expert output after fc2_latent_proj.
• Add forward/backward parity coverage comparing shared-expert overlap vs no-overlap for latent A2A and HybridEP.

Test plan

• PASS: uv run unit-test run tests/unit_tests/transformer/moe/test_shared_experts.py --cluster cw --platform h100 --environment dev --time-limit 0:45:00 (SLURM 12467070, 11 passed; includes latent A2A parity and test_latent_shared_expert_hybridep_forward_backward).
• PASS: CHECK_ONLY=true BASE_REF=main UV_CACHE_DIR=/project/agentic-mcore-dev/.uv_cache uv run --project /project/agentic-mcore-dev/.claude/worktrees/latent-moe-shared-expert-overlap --extra mcore-lint -- bash tools/autoformat.sh (Black/isort/pylint/ruff clean; wrapper exits 0; mypy still prints existing missing-dependency/type issues from this local lint env).
• PASS: tools/check_copyright.py on changed Python files.

Notes

• A separate main-base replay in a temporary CW worktree reached import before test collection but was blocked by the local mcore-ci-dev image having nvidia-resiliency-ext < 0.6.0, while current origin/main requires nvidia-resiliency-ext==0.6.0 (SLURM 12467214).

[copy-pr-bot]

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.

[github-actions]

This PR has been automatically converted to draft because all PRs must start as drafts.

When you are ready for review, click Ready for Review to begin the review process. This will:

1. Add the oncall reviewer (optional reviewer)
2. Add required review teams based on your changes

See the contribution guide for more details.

[kvareddy]

@sidsingh-nvidia can you take a look at this MR. You added this support on inference side.

[Victarry]

Folded the FSDP shared-expert stream synchronization fix into this latent MoE shared expert overlap PR. I also adjusted the implementation to use an opt-in parameter marker from shared expert overlap, so Megatron FSDP handles a generic side-stream grad processing requirement instead of importing MoE shared expert classes or matching parameter names.

[Phlip79]

/claude fix

[svcnvidia-nemo-ci]

🛠️ Claude fix commit 6b1d53cdb18d (attempt 1)

⚠️ This explanation is AI-generated and may be inaccurate; the exact commit is authoritative.

What changed
Resolved the only merge conflict (3 hunks) in tests／unit＿tests／transformer／moe／test＿shared＿experts.py as a union: kept the PR's is＿hybrid＿ep＿available and its ＊＊config＿kwargs helper signature, plus main's TE fakes and new tests, while adopting main's ＿set＿random＿seed seeding call. All four test methods retained under one teardown＿method; removed the unused model＿parallel＿cuda＿manual＿seed import. Markers gone, ast.parse passes, 14 tests, no over-length lines. Not committed or staged.

Files changed by Claude

• tests/unit_tests/transformer/moe/test_shared_experts.py

Why
Steering was empty and the CI log dir was empty, so the conflict was the only actionable failure. The PR hoisted the inlined test body into a ＊＊config＿kwargs helper; main independently migrated seeding to ＿set＿random＿seed and added fixtures／tests. Either raw side of the middle hunk breaks the file: main's references an out-of-scope dispatcher＿type inside the helper, and the PR's calls a seed helper whose import main deleted.

DCO
@Victarry, please fix any DCO failures on your commits before merge. DCO does not block this workflow.

View exact commit

Sanitized and posted by svcnvidia-nemo-ci.

[svcnvidia-nemo-ci]

/ok to test 6b1d53c

[svcnvidia-nemo-ci]

🛠️ Claude fix commit 3edbd80c1bf3 (attempt 2)

⚠️ This explanation is AI-generated and may be inaccurate; the exact commit is authoritative.

What changed
Edited tests／unit＿tests／transformer／moe／test＿shared＿experts.py line 540: replaced pre＿forward＿comm(hidden＿states) with register＿input(hidden＿states) plus pre＿forward＿comm(). One file, 2 added, 1 removed, uncommitted. Base merge was already clean, so no conflicts. Left the determinism-coverage lint alone; it needs a non-PR file or a maintainer label. Verified statically only, as torch and pytest are absent.

Files changed by Claude

• tests/unit_tests/transformer/moe/test_shared_experts.py

Why
PR commit 93fcd2f dropped the input parameter from SharedExpertMLP.pre＿forward＿comm, moving capture into a new register＿input(). It fixed both token＿dispatcher.py callers but missed this test, which the PR also edits. The call passed hidden＿states positionally, binding it to wait＿current＿stream and tripping the new assert. CI had not shown it, as unit tests were skipped behind the lint failure.

DCO
@Victarry, please fix any DCO failures on your commits before merge. DCO does not block this workflow.

View exact commit

Sanitized and posted by svcnvidia-nemo-ci.

[svcnvidia-nemo-ci]

/ok to test 3edbd80

[svcnvidia-nemo-ci]

🛠️ Claude fix commit d96f0b9723c4 (attempt 3)

⚠️ This explanation is AI-generated and may be inaccurate; the exact commit is authoritative.

What changed
Bumped the copyright header in pr-head／megatron／core／transformer／moe／shared＿experts.py from 2025 to 2026 (one line, line 1). This clears the failing tools／check＿copyright.py check, which now passes on all four files the PR changes. Left the kernel determinism coverage gate failing and made no other edits. Changes are uncommitted in pr-head／.

Files changed by Claude

• megatron/core/transformer/moe/shared_experts.py

Why
The base merge was already committed and clean (no conflicts／markers; call sites consistent). Unit tests were "skipped", so the CICD failure was only downstream of linting. The reported determinism-coverage failure is out of policy to fix: its remedies are editing a determinism test the PR never touched (broadening scope) or adding a PR label (needs network). The fixable failure was check＿copyright.py: the PR's own commit b4076ab bumped 3 of its 4 changed files but missed shared＿experts.py.

DCO
@Victarry, please fix any DCO failures on your commits before merge. DCO does not block this workflow.

View exact commit

Sanitized and posted by svcnvidia-nemo-ci.

[svcnvidia-nemo-ci]

/ok to test d96f0b9

[svcnvidia-nemo-ci]

❌ Claude fix stopped after attempt 3; supported lint or unit tests still fail. View exact-SHA CI.