[main] [3/5] Qwen3.5 support: SharedExpertMLP meta init

[wplf]

Qwen3.5 support series

This is part of a 5-PR series adding Qwen3.5-VL support, split for review clarity.

Main PRs (this series):

• [1/5] MTP packed-seq CP+THD fix — fix(mtp): use padded cu_seqlens in MTP roll for THD with CP #4495
• [2/5] FSDP DTensor Bridge checkpoint compatibility — [main] [2/5] Qwen3.5 support: FSDP DTensor Bridge checkpoint compatibility #4753
• [3/5] SharedExpertMLP meta init — [main] [3/5] Qwen3.5 support: SharedExpertMLP meta init #4754 ← this PR
• [4/5] Interleaved MRoPE layout — [main] [4/5] Qwen3.5 support: Interleaved MRoPE layout #4755
• [5/5] Qwen3.5-VL training example — [main] [5/5] Qwen3.5 support: Qwen3.5-VL training example #4756

Dev PRs (corresponding mirrors):

• [1/5] [Dev] fix(mtp): use padded cu_seqlens in MTP roll for THD with CP #4494
• [2/5] [dev] [2/5] Qwen3.5 support: FSDP DTensor Bridge checkpoint compatibility #4748
• [3/5] [dev] [3/5] Qwen3.5 support: SharedExpertMLP meta init #4749
• [4/5] [dev] [4/5] Qwen3.5 support: Interleaved MRoPE layout #4750
• [5/5] [dev] [5/5] Qwen3.5 support: Qwen3.5-VL training example #4751

---

Summary

Add _reset_parameters to SharedExpertMLP so the directly-owned gate_weight is materialized off the meta device when use_shared_expert_gate=True.

Why

Without this, meta-init leaves gate_weight on the meta device and the first forward fails with a meta-tensor error. Submodules already have their own _reset_parameters, so only the directly-owned gate_weight needs handling.

The implementation mirrors the standard per-parameter init pattern:

• run init_method when config.perform_initialization
• cast to config.params_dtype
• set sequence_parallel attribute

Risk

Only fires when use_shared_expert_gate=True and gate_weight is not None — no effect on existing paths.

Notes

Mirror of #4749 (same patch, targeting main instead of dev).

🤖 Generated with Claude Code

[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.

[wplf]

/ok to test eab5149