[bug] invalid tensor size at index 0 when scatter_to_tp_ranks

Problem

Running verl + TransferQueue + Megatron-Bridge, Megatron/SGLang: invalid tensor size at index 0 when scatter_to_tp_ranks

Minimal repro

1. Use latest verl and docker verlai/verl:sgl0512.dev2
2. Use


pip install git+https://github.com/Ascend/TransferQueue.git@v0.1.7 --no-deps --no-build-isolation && \
pip3 install git+https://github.com/NVIDIA-NeMo/Megatron-Bridge.git@0a21da4 --no-deps --no-build-isolation  --ignore-requires-python


(Megatron-Bridge main also failed)

3. Run verl's script: `examples/grpo_trainer/run_qwen2-7b_math_megatron_fsdp.sh`, change  inference backend to sglang. Where:

TP=${TP:-4}
PP=${PP:-1}
GEN_TP=${GEN_TP:-4}

Expected behavior

No error and smoothly run, as vanilla bridge.

Affected area

area:training

Regression?

Yes

Environment

main/0a21da4
GPU: A100
pip list:

megatron-bridge==0.5.0+0a21da47
megatron-core==0.19.0+bf32f4415
sglang==0.5.12
sglang-kernel==0.4.2.post2+cu130
torch==2.11.0+cu130
torch_c_dlpack_ext==0.1.5
torch_memory_saver==0.0.9.post1
torchao==0.17.0+cu130
torchaudio==2.11.0+cu130
torchcodec==0.11.1+cu130
torchdata==0.11.0
torchvision==0.26.0+cu130
tos==2.8.7
tqdm==4.67.3
traitlets==5.15.0
TransferQueue==0.1.7
transformer_engine==2.15.0+42b84005
transformers==5.6.0

Logs

File "/opt/tiger/verl/verl/utils/transferqueue_utils.py", line 329, in inner
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])     return func(*args, **kwargs)
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])            ^^^^^^^^^^^^^^^^^^^^^
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])   File "/opt/tiger/verl/verl/workers/engine_workers.py", line 590, in init_model
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])     self.actor.reset()
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])   File "/opt/tiger/verl/verl/single_controller/base/decorator.py", line 431, in inner
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])     return func(*args, **kwargs)
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])            ^^^^^^^^^^^^^^^^^^^^^
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])   File "/opt/tiger/verl/verl/utils/transferqueue_utils.py", line 329, in inner
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])     return func(*args, **kwargs)
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])            ^^^^^^^^^^^^^^^^^^^^^
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])   File "/opt/tiger/verl/verl/workers/engine_workers.py", line 170, in reset
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])     self.engine.initialize()
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])   File "/opt/tiger/verl/verl/workers/engine/megatron/transformer_impl.py", line 399, in initialize
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])     self.module = self._build_megatron_module()
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])   File "/opt/tiger/verl/verl/workers/engine/megatron/transformer_impl.py", line 332, in _build_megatron_module
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])     self.bridge.load_hf_weights(
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])   File "/usr/local/lib/python3.12/dist-packages/megatron/bridge/models/conversion/auto_bridge.py", line 439, in load_hf_weights
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])     bridge.load_weights_hf_to_megatron(pre_trained, model, allowed_mismatched_params=allowed_mismatched_params)
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])   File "/usr/local/lib/python3.12/dist-packages/megatron/bridge/models/conversion/model_bridge.py", line 1037, in load_weights_hf_to_megatron
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])     converted_weights = task.mapping.hf_to_megatron(hf_weights, task.megatron_module)
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])   File "/usr/local/lib/python3.12/dist-packages/megatron/bridge/models/conversion/param_mapping.py", line 1356, in hf_to_megatron
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])     return self._mapping.hf_to_megatron(hf_weights, megatron_module)
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])   File "/usr/local/lib/python3.12/dist-packages/megatron/bridge/models/conversion/param_mapping.py", line 930, in hf_to_megatron
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])     return self.scatter_to_tp_ranks(
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])            ^^^^^^^^^^^^^^^^^^^^^^^^^
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])   File "/usr/local/lib/python3.12/dist-packages/megatron/bridge/models/conversion/param_mapping.py", line 588, in scatter_to_tp_ranks
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])     torch.distributed.scatter(
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/c10d_logger.py", line 83, in wrapper
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])     return func(*args, **kwargs)
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])            ^^^^^^^^^^^^^^^^^^^^^
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/distributed_c10d.py", line 4506, in scatter
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])     work = group.scatter(output_tensors, input_tensors, opts)
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a])            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(TaskRunnerV1 pid=588296, ip=[2605:340:cdb1:6400::a7c:2d0a]) ValueError: ProcessGroupNCCL::scatter: invalid tensor size at index 0 (expected (38016, 3584), got (152064, 3584))

Start with examples/grpo_trainer/run_qwen2-7b_math_megatron_fsdp.sh using the stated TransferQueue, Megatron-Bridge, SGLang, TP=4 and GEN_TP=4 setup. Trace initialization through verl/workers/engine/megatron/transformer_impl.py and the reported scatter_to_tp_ranks call in Megatron-Bridge param_mapping.py. Done means the reproduction initializes and runs without the invalid tensor-size error.