fix: vLLM 0.21 colocate IPC weight sync

[CalvinXKY]

Summary

• Add UpdateVLLMWeightFromTensor for colocated CUDA IPC weight sync (per-engine llm_handle, per-chunk barrier).
• Auto-enable vLLM sleep mode when --colocate; point --worker-extension-cls at merged worker extension.
• Unit tests for colocate lifecycle.
  Stacked on [feat] Support Colocated Weight Sync via CUDA IPC for vime #22 (update_weights_tensor).

Test plan

• pytest tests/unit/backends/megatron_utils/update_weight/test_update_weight_from_tensor_vllm.py
• qwen_0.6b_colocate.sh multi-step training

step6:

[gemini-code-assist]

Code Review

This pull request implements a weight synchronization mechanism for colocated vLLM engines using CUDA IPC via Ray, introducing the UpdateVLLMWeightFromTensor class and a vLLM worker extension. Key feedback includes identifying a logic error in the server launch process that could skip the sleep mode flag for multiple engines and pointing out redundant lifecycle calls that should be restricted to rank 0. Additionally, the reviewer suggested moving a synchronization barrier outside the weight chunk loop to improve performance and recommended against setting global environment variables within the class constructor.

[gemini-code-assist]

There is a logic error here: if multiple engines are launched, the first one will set args.vllm_enable_sleep_mode = True. Subsequent calls to launch_server_process will then skip this block because of the and not getattr(args, "vllm_enable_sleep_mode", False) condition, resulting in subsequent vLLM servers being launched without the --enable-sleep-mode flag. The flag should be added to the command line for every engine if the condition is met, regardless of whether it was already set on the args object.

Suggested change

	if (getattr(args, "offload_rollout", False) or getattr(args, "colocate", False)) and not getattr(
	args, "vllm_enable_sleep_mode", False
	):
	cmd += ["--enable-sleep-mode"]
	args.vllm_enable_sleep_mode = True
	if getattr(args, "offload_rollout", False) or getattr(args, "colocate", False) or getattr(args, "vllm_enable_sleep_mode", False):
	if "--enable-sleep-mode" not in cmd:
	cmd += ["--enable-sleep-mode"]
	args.vllm_enable_sleep_mode = True

[gemini-code-assist]

Setting global environment variables like VLLM_ALLOW_INSECURE_SERIALIZATION inside a class constructor is generally discouraged as it can have unintended side effects on other parts of the application or other tests in the same process. It would be better to set this at the application entry point or ensure it is only set when strictly necessary for the vLLM server process.

[gemini-code-assist]

This lifecycle call is redundant when multiple trainer ranks share the same vLLM engine (e.g., in configurations where rollout_num_gpus_per_engine > 1). Since this triggers an HTTP request to the vLLM server, calling it from every rank is inefficient and could lead to race conditions or unnecessary overhead. It should be called once per engine, ideally by rank 0 for all colocated engines.

Suggested change

	if self._ipc_engine is not None:
	ray.get(self._ipc_engine.start_weight_update.remote(is_checkpoint_format=True))
	dist.barrier(group=get_gloo_group())
	if rank == 0 and self._colocated_engines:
	ray.get([engine.start_weight_update.remote(is_checkpoint_format=True) for engine in self._colocated_engines])

[gemini-code-assist]

This barrier is inside the loop over weight chunks. If the model has many chunks, this will introduce significant synchronization overhead. Since the goal is to ensure all ranks have finished sending weights before signaling the engine to exit update mode, this barrier should be moved outside and after the loop.

[gemini-code-assist]

Similar to start_weight_update, calling finish_weight_update from every rank is redundant and inefficient. It should be called once per engine by rank 0.

Suggested change

	if self._ipc_engine is not None:
	ray.get(self._ipc_engine.finish_weight_update.remote())
	dist.barrier(group=get_gloo_group())
	if rank == 0 and self._colocated_engines:
	ray.get([engine.finish_weight_update.remote() for engine in self._colocated_engines])

[CalvinXKY]

Test script:

export PYTHONPATH=/root/Megatron-LM
SCRIPT_DIR="/data/nfs_87/xky/RL/vime/scripts"
source "${SCRIPT_DIR}/models/qwen3-0.6B.sh"
export PYTHONUNBUFFERED=1
mkdir -p "$TENSORBOARD_DIR" "/data/nfs_87/xky/logs"
export TENSORBOARD_DIR="/data/nfs_87/xky/logs/tb_qwen3_0.6b_$(date +%Y%m%d_%H%M%S)"
LOG_FILE="/data/nfs_87/xky/logs/train_qwen3_0.6b_vllm_$(date +%Y%m%d_%H%M%S).log"

cd /data/nfs_87/xky/RL/vime
 
python train.py \
  --train-backend megatron \
  --actor-num-nodes 1 \
  --actor-num-gpus-per-node 4 \
  --rollout-num-gpus 4 \
  --rollout-num-gpus-per-engine 1 \
  ${MODEL_ARGS[@]} \
  \
  --hf-checkpoint /data/nfs_87/xky/models/Qwen3-0.6B \
  --ref-load /data/nfs_87/xky/models/Qwen3-0.6B_torch_dist \
  \
  --prompt-data /data/nfs_87/xky/datasets/dapo-math-17k/dapo-math-17k.jsonl \
  --input-key prompt \
  --label-key label \
  --apply-chat-template \
  --rollout-shuffle \
  --rm-type deepscaler \
  \
  --colocate \
  \
  --num-rollout 200 \
  --rollout-batch-size 32 \
  --n-samples-per-prompt 8 \
  --rollout-max-response-len 8192 \
  --rollout-temperature 1.0 \
  --global-batch-size 256 \
  --balance-data \
  \
  --advantage-estimator grpo \
  --use-kl-loss \
  --kl-loss-coef 0.0 \
  --kl-loss-type low_var_kl \
  --entropy-coef 0.0 \
  --eps-clip 0.2 \
  --eps-clip-high 0.28 \
  \
  --optimizer adam \
  --lr 1e-6 \
  --lr-decay-style constant \
  --weight-decay 0.1 \
  --adam-beta1 0.9 \
  --adam-beta2 0.98 \
  \
  --tensor-model-parallel-size 1 \
  --pipeline-model-parallel-size 1 \
  --context-parallel-size 1 \
  --expert-model-parallel-size 1 \
  --expert-tensor-parallel-size 1 \
  --recompute-granularity full \
  --recompute-method uniform \
  --recompute-num-layers 1 \
  --use-dynamic-batch-size \
  --max-tokens-per-gpu 8192 \
  \
  --attention-dropout 0.0 \
  --hidden-dropout 0.0 \
  --accumulate-allreduce-grads-in-fp32 \
  --attention-softmax-in-fp32 \
  --attention-backend flash \
  \
  --train-memory-margin-bytes 2147483648 \
  --use-tensorboard \
  2>&1 | tee -a "$LOG_FILE"

[knlnguyen1802]

LGTM, leave some comment about refactor

[knlnguyen1802]

I think we can completely replace this with update_weight_from_tensor.py

[knlnguyen1802]

Can rename it as UpdateWeightFromTensor and remove the old one

[CalvinXKY]

Done