[Bug] Qwen3.5 MoE online weight update sends ETP-sharded TEGroupedLinear weights without all-gather #2280

Bug Description

Online Megatron-to-SGLang weight updates fail for Qwen3.5-35B-A3B when expert tensor parallelism is enabled. At the first weight update, SGLang reports that a 256-element shard is being sliced as if it were the full 512-element tensor:

Failed to update parameter online: start (256) + length (256) exceeds dimension size (256)..
The full weights of the ModelRunner are partially updated. Please discard the whole weights.

The failure appears to be caused by missing tensor-parallel metadata on
TEGroupedLinear expert parameters. slime only enters its ETP all-gather path when param.tensor_model_parallel is true, but the Megatron-LM revision in the image predates the upstream fix that stamps this metadata for
TEGroupedLinear.

Steps to Reproduce

1. Use slime commit 06ffdbe22be068b52f9ed0fc318c473f7030197e with:

   slime                  0.3.1
   SGLang                 0.5.15.post1
   Transformer Engine     2.16.1
   Megatron-LM            1dcf0dafa884ad52ffb243625717a3471643e087
   

2. Launch Qwen3.5-35B-A3B using
   scripts/models/qwen3.5-35B-A3B.sh and enable online weight updates with
   the following relevant arguments:

   source scripts/models/qwen3.5-35B-A3B.sh
   
   python3 /root/slime/train_async.py \
     "${MODEL_ARGS[@]}" \
     --actor-num-nodes 4 \
     --actor-num-gpus-per-node 8 \
     --rollout-num-gpus 16 \
     --update-weights-interval 1 \
     --tensor-model-parallel-size 2 \
     --pipeline-model-parallel-size 1 \
     --context-parallel-size 8 \
     --expert-model-parallel-size 8 \
     --expert-tensor-parallel-size 2 \
     --rollout-num-gpus-per-engine 2 \
     <the usual checkpoint, data, rollout, and optimizer arguments>

3. Wait for the first Megatron-to-SGLang online weight update.

4. Observe the error on the SGLang TP1 ranks. The same error is subsequently
   repeated across the rollout engines.

The reproduction does not depend on the prompt data: it happens during weight
synchronization before useful training begins.

Expected Behavior

slime should all-gather ETP-sharded expert weights across the expert tensor parallel group before converting/sending them to SGLang. SGLang should receive the full 512-element tensor and all ModelRunner ranks should complete the update successfully.

Actual Behavior

For affected TEGroupedLinear expert parameters, slime sends a local 256-element ETP shard without gathering it into the full 512-element tensor. SGLang TP1 then attempts to take the [256:512] slice from a dimension of size 256, raises the error above, and leaves the ModelRunner partially updated.

This occurred in 4/4 launches using:

TP=2, CP=8, EP=8, ETP=2, PP=1
SGLang TP / rollout GPUs per engine=2

Environment

- slime version: 0.3.1, git 06ffdbe22be068b52f9ed0fc318c473f7030197e
- Python version: 3.12.3
- PyTorch version: image default (not separately recorded)
- CUDA/ROCm version: CUDA; NCCL 2.28.9+cuda12.9
- GPU type and count: 48 x H200 (32 actor + 16 rollout)
- OS: slime container image
- SGLang version: 0.5.15.post1
- Megatron-LM version: 1dcf0dafa884ad52ffb243625717a3471643e087
- Transformer Engine version: 2.16.1

Logs

Additional Context

At slime commit 06ffdbe, all_gather_param returns param.data without an
all-gather when param.tensor_model_parallel is false:

The image uses Megatron-LM 1dcf0dafa, which does not contain the later
TEGroupedLinear expert TP-metadata fix. NVIDIA/Megatron-LM PR #5916 explicitly
states that TEGroupedLinear did not mark expert parameters with
param.tensor_model_parallel when it should have, including the ETP != TP
case:

NVIDIA/Megatron-LM#5916

The corresponding upstream merge commit is:

NVIDIA/Megatron-LM@cd4afff

This also explains why the issue becomes visible with newer Transformer Engine
versions: the old Megatron revision does not restore all metadata required by
slime's weight-gather decision.

A possible fix is to update/backport the relevant Megatron-LM
TEGroupedLinear metadata change, or make slime's expert-weight gather logic
robust to this missing attribute. A regression test could assert that an
ETP=2 TEGroupedLinear expert weight is gathered from size 256 to size 512
before it is sent to SGLang.

This is different from #1679: that issue has the same SGLang "partially
updated" suffix, but its cause is an FP8/BF16 dtype mismatch rather than a
missing ETP all-gather.

Pre-submission Checklist

• I have read the CONTRIBUTING.md and understand the collaboration scope.
• I have read the documentation and my issue is not addressed there.
• I have searched for existing issues and this is not a duplicate.
• I have provided a minimal, reproducible example.