fix: make Qwen3.5 MoE refit work and fail loudly on weight-update errors

[michal2409]

Split out of #3382 per review feedback (smaller, independently reviewable PRs).

Two refit fixes for MoE models, observed on Qwen3.5-397B (GB300, 24 vLLM
engines, collective refit):

1. megatron-bridge exports routed experts as stacked 3D tensors
   (mlp.experts.gate_up_proj [E, 2I, H], mlp.experts.down_proj
   [E, H, I]). vLLM 0.20's fused FusedMoE loader rejects that layout
   ('shard_dim=0 is not a valid data dimension for a 3D tensor'), so
   every engine failed the update while the trainer was already inside
   the NCCL broadcast, deadlocking the first refit. Split the stacked
   exports on the sender into per-expert 2D HF tensors
   (experts.{i}.{gate,up,down}_proj.weight), which vLLM loads through
   its standard per-expert expert_params_mapping. The gate/up split
   follows the HF stacked convention (gate = rows [:I], up = rows [I:]),
   matching vLLM's chunk(2, dim=-2) fused semantics; verified against
   the HF checkpoint safetensors headers. prepare_refit_info and the
   broadcast share the iterator, so metadata and payload always agree.
   No-op for non-MoE models and 2D exports; NRL_REFIT_SPLIT_FUSED_EXPERTS=0
   restores the raw export for a future vLLM that consumes it natively.

2. update_weights_via_ipc_zmq_async / update_weights_from_collective_async
   raise on worker failure instead of print + return False. A silently
   swallowed weight-update failure means subsequent rollouts sample from
   stale weights (off-policy drift) or the collective deadlocks; failures
   now abort the step at the source instead of depending on every caller
   checking the boolean.

Validated as part of the Qwen3.5-397B GB300 GRPO convergence run (eval accuracy 0.789 vs 0.70 target over 20 steps).

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