Fix sender device alignment in vLLM weight synchronization

[0KEAHA]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Submission base: modelscope/ms-swift@174f15e2a812c2669171c1006d8441348ecf35e5.
The real GPU baseline/patch A/B was run on the earlier main snapshot
2895611fd77ded82c042724cb0bfc633ca748079; intervening upstream commits did
not modify the source or test paths in this PR.

What does this PR do?

Ensures trainer-side tensors sent to external vLLM servers satisfy the device contract of PyNcclCommunicator.

Problem

In Megatron GRPO, the sender communicator is initialized on the last trainer rank's current device, while bridge-exported model tensors may reside on another local CUDA device. vLLM requires the input of PyNcclCommunicator.broadcast() to be on communicator.device and otherwise raises:

AssertionError:
this nccl communicator is created to work on cuda:1,
but the input tensor is on cuda:0

Weight updates run in ThreadPoolExecutor workers, so the worker's current CUDA device/stream also cannot be assumed to match the communicator.

Fix

This change adds a sender-only broadcast wrapper in vllm_client.py that:

• synchronizes each non-CPU source device once;
• keeps same-device tensors unchanged (zero-copy path);
• copies only mismatched tensors to communicator.device;
• broadcasts under the communicator device context;
• synchronizes the communicator device before temporary tensors are released;
• is shared by named, flattened full-model, flattened LoRA, and non-flattened LoRA sender paths.

The generic broadcast helper remains unchanged because it is also used by rollout receivers that require in-place updates to their original receive buffers.

Experiment results

Tests

Added CPU-CI-compatible mock tests for:

• same-device tensors do not call .to();
• mismatched tensors are aligned before broadcast;
• a worker starting on cuda:0 broadcasts a cuda:1 tensor through the actual helper with a cuda:1 current stream;
• CPU/offloaded tensors do not call accelerator synchronize on CPU;
• multiple tensors synchronize each unique source device once.

Local checks completed:

PASS  dependency-free before/after invariant reproducer
PASS  sender contract harness and four-path coverage
PASS  compileall
PASS  flake8
PASS  YAPF check
PASS  repository pre-commit hooks
PASS  git diff --check
PASS  pytest regression (4 tests, GPU environment)

Real CUDA/NCCL validation on 3 x RTX 4090:

PASS  unmodified main snapshot 2895611f VLLMClient.update_flattened_params reproduces
      the exact cuda:1 communicator / cuda:0 tensor assertion
PASS  identical patched VLLMClient/HTTP/rollout request completes
PASS  real two-rank PyNccl receiver validates 3 broadcasts after alignment
PASS  Qwen2.5-0.5B Megatron GRPO full-model sync, 3 training/sync cycles

For the 64 MiB mismatch tensor test, active target allocation returned to zero
after every cycle, while allocator-reserved memory plateaued after the first
cycle.

Scope

This PR intentionally does not change the Megatron exporter, bucket sizing, rollout scheduling, GRPO logic, or receiver-side weight loading.