[megatron] enable bucketed weight sync for non-colocated nccl weight sync in megatron

[erictang000]

Enable packed broadcast for non-colocated Megatron weight sync

Bucketing was previously only enabled for the CUDA IPC strategy (colocated mode). This extends it to the broadcast strategy (non-colocated mode), packing all tensors in each bucket into a single contiguous buffer and broadcasting it in one NCCL operation — matching how CUDA IPC already works. This reduces both per-tensor NCCL overhead and HTTP round-trips, which matters most for MoE models with many small expert parameters.

Changes

• broadcast_strategy.py: Add sizes field to BroadcastWeightUpdateRequest; sender packs tensors into one buffer and broadcasts once per bucket; receiver unpacks using sizes
• megatron_worker.py: Always enable bucketing, removing the CudaIpcTransferStrategy-only guard
• test_megatron_worker.py: Add non_colocated_moe test entry for Moonlight-16B-A3B-Instruct

Results

Partial weight sync on 4xL40S (4 layers only)

Before:

After:

Full weight sync on 8xH100 (4 inf, 4 train) with Moonlight-16B-A3B

Before:

After:

Test

uv run --isolated --extra dev --extra megatron -- pytest -s tests/backends/skyrl_train/gpu/gpu_ci/test_megatron_worker.py::test_megatron_policy_weight_sync -k non_colocated

---

[SumanthRH]

Nice