[Feature] Performance boost for VerlEngine's `update_weights_from_tensor`

[yzlnew]

Checklist

• 1. If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/sglang/discussions/new/choose Otherwise, it will be closed.
• 2. Please use English, otherwise it will be closed.

Motivation

Hi, I'm trying to run DeepSeek-V3 GRPO using verl with sglang backend.

Image and code

whatcanyousee/verl:ngc-cu124-vllm0.8.5-sglang0.4.6-mcore0.12.0-te2.3

volcengine/verl#1694

Launch configs

train_files=$gsm8k_train_path
test_files=$gsm8k_test_path

ALL_OFFLOAD=${ALL_OFFLOAD:-True}
COMMON_PARAM_OFFLOAD=${COMMON_PARAM_OFFLOAD:-$ALL_OFFLOAD}
COMMON_GRAD_OFFLOAD=${COMMON_GRAD_OFFLOAD:-$ALL_OFFLOAD}
COMMON_OPTIMIZER_OFFLOAD=${COMMON_OPTIMIZER_OFFLOAD:-$ALL_OFFLOAD}

ACTOR_PARAM_OFFLOAD=${ACTOR_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}
ACTOR_GRAD_OFFLOAD=${ACTOR_GRAD_OFFLOAD:-$COMMON_GRAD_OFFLOAD}
ACTOR_OPTIMIZER_OFFLOAD=${ACTOR_OPTIMIZER_OFFLOAD:-$COMMON_OPTIMIZER_OFFLOAD}
REF_PARAM_OFFLOAD=${REF_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}
CRITIC_PARAM_OFFLOAD=${CRITIC_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}
CRITIC_GRAD_OFFLOAD=${CRITIC_GRAD_OFFLOAD:-$COMMON_GRAD_OFFLOAD}
CRITIC_OPTIMIZER_OFFLOAD=${CRITIC_OPTIMIZER_OFFLOAD:-$COMMON_OPTIMIZER_OFFLOAD}
RM_PARAM_OFFLOAD=${RM_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}


NODES=32
PP=16
TP=1
EP=16
ETP=1
INFER_TP=16

n_resp_per_prompt=4

# RAY_ADDRESS='auto' ray job submit --working-dir . -- 
python3 -m verl.trainer.main_ppo --config-path=./config --config-name='ppo_megatron_trainer'\
    algorithm.adv_estimator=grpo \
    data.train_files="$train_files" \
    data.val_files="$test_files" \
    data.train_batch_size=512 \
    data.max_prompt_length=1024 \
    data.max_response_length=512 \
    data.filter_overlong_prompts=True \
    data.truncation='error' \
    actor_rollout_ref.model.path=$LLM \
    actor_rollout_ref.actor.optim.lr=1e-6 \
    actor_rollout_ref.actor.ppo_mini_batch_size=128 \
    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \
    actor_rollout_ref.actor.use_kl_loss=False \
    actor_rollout_ref.actor.use_torch_compile=False \
    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \
    actor_rollout_ref.rollout.name=sglang \
    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \
    actor_rollout_ref.rollout.temperature=1.0 \
    actor_rollout_ref.rollout.top_p=1.0 \
    actor_rollout_ref.rollout.top_k=-1 \
    actor_rollout_ref.rollout.tensor_model_parallel_size=$INFER_TP \
    +actor_rollout_ref.rollout.engine_kwargs.sglang.enable_torch_compile=False \
    +actor_rollout_ref.rollout.engine_kwargs.sglang.cuda_graph_max_bs=2 \
    algorithm.use_kl_in_reward=False \
    trainer.logger=['console','tensorboard'] \
    trainer.project_name='verl_megatron_gsm8k_examples' \
    trainer.experiment_name='dsv3-32nodes' \
    trainer.n_gpus_per_node=8 \
    trainer.nnodes=$NODES \
    trainer.save_freq=-1 \
    trainer.test_freq=5 \
    +actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_first_pipeline_stage=4 \
    +actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_last_pipeline_stage=1 \
    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$PP \
    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=$PP \
    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$TP \
    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=$TP \
    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP \
    actor_rollout_ref.ref.megatron.expert_model_parallel_size=$EP \
    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP \
    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=$ETP \
    actor_rollout_ref.actor.megatron.param_offload=${ACTOR_PARAM_OFFLOAD} \
    actor_rollout_ref.actor.megatron.optimizer_offload=${ACTOR_OPTIMIZER_OFFLOAD} \
    actor_rollout_ref.actor.megatron.grad_offload=${ACTOR_GRAD_OFFLOAD} \
    actor_rollout_ref.ref.megatron.param_offload=${REF_PARAM_OFFLOAD} \
    actor_rollout_ref.actor.megatron.use_dist_checkpointing=True \
    actor_rollout_ref.ref.megatron.use_dist_checkpointing=True \
    actor_rollout_ref.actor.megatron.dist_checkpointing_path=$DIST_CKPT_PATH \
    actor_rollout_ref.ref.megatron.dist_checkpointing_path=$DIST_CKPT_PATH \
    trainer.default_hdfs_dir=null \
    trainer.default_local_dir=$CKPT_DIR \
    trainer.val_before_train=False \
    trainer.val_before_train=False \
    trainer.total_epochs=100 $@

Issues

1. Using verl's original code, sglang raises

RuntimeError: Cowardly refusing to serialize non-leaf tensor which requires_grad, since autograd does not support crossing process boundaries.  If you just want to transfer the data, call detach() on the tensor before serializing (e.g., putting it on the queue).

2. gpu_memory_utilization can only set to 0.8, otherwise it will fail. Need a more proper configuration list for sglang server.
3. I tried to fix 1 by detaching tensors in https://github.com/volcengine/verl/blob/96903e0e97af6f645610bbe45f255a6c76f46a89/verl/utils/megatron_utils.py#L723-L821 (though I think it should be fixed in the sglang side), but still faces performance issues when verl tries to update_weights_from_tensor.

Not much useful logs, I'm seeing many logs like

Cloning 16 replicas of the shared expert into MoE:   0%|          | 0/58 [00:00<?, ?it/s]^MCloning 16 replicas of the shared expert into MoE:   0%|          | 0/58 [00:00<?, ?it/s]^MCloning 16 replicas of the shared expert into MoE:   0%|          | 0/58 [00:00<?, ?it/s]^MCloning 16 replicas of the shared expert into MoE:   0%|          | 0/58 [00:00<?, ?it/s]^MCloning 16 replicas of the shared expert into MoE: 100%|██████████| 58/58 [00:00<00:00, 701065.22it/s]^MCloning 16 replicas of the shared expert into MoE: 100%|██████████| 58/58 [00:00<00:00, 638502.97it/s]^MCloning 16 replicas of the shared expert into MoE: 100%|██████████| 58/58 [00:00<00:00, 443114.08it/s]

for several hours before

The actor is dead because its worker process has died. Worker exit type: SYSTEM_ERROR Worker exit detail: Worker unexpectedly exits with a connection error code 2. End of file. There are some potential root causes. (1) The process is killed by SIGKILL by OOM killer due to high memory usage. (2) ray stop --force is called. (3) The worker is crashed unexpectedly due to SIGSEGV or other unexpected errors.

Related resources

No response

[zhaochenyang20]

We shall turn off shared experts fusion right now

[Yangruipis]

@zhaochenyang20 hi~ Is this feature ready now? I'm also bothered by the bad performance of update_weights_from_tensor with verl megatron backend.

[zhaochenyang20]

@zhaochenyang20 hi~ Is this feature ready now? I'm also bothered by the bad performance of update_weights_from_tensor with verl megatron backend.

@Yangruipis would you like to add my wechat 18015766633

[Yangruipis]

my env

• 1 X H20
• Moonlight 16B
• sglang 0.4.8post1

$ python -m torch.utils.collect_env

PyTorch version: 2.7.1+cu126                            
Is debug build: False                                                                                                                       
CUDA used to build PyTorch: 12.6                      
ROCM used to build PyTorch: N/A                                       
                                                                      
OS: Ubuntu 24.04.2 LTS (x86_64)                                 
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04) 13.3.0                                                                                          
Clang version: Could not collect             
CMake version: Could not collect                                                                                                            
Libc version: glibc-2.39                     
                                                                                                                                            
Python version: 3.11.13 (main, Jun  5 2025, 13:12:00) [GCC 11.2.0] (64-bit runtime)                                                         
Python platform: Linux-4.18.0-348.7.1.el8_5.x86_64-x86_64-with-glibc2.39                                                                   
Is CUDA available: True                         
CUDA runtime version: 12.9.86                                                                                                               
CUDA_MODULE_LOADING set to: LAZY                               
GPU models and configuration:                                         
GPU 0: NVIDIA H20-3e                                                                                                                        
GPU 1: NVIDIA H20-3e                             
GPU 2: NVIDIA H20-3e                                                                                                                        
GPU 3: NVIDIA H20-3e                         
GPU 4: NVIDIA H20-3e                                                                                                                        
GPU 5: NVIDIA H20-3e                           
GPU 6: NVIDIA H20-3e                                                                                                                        
GPU 7: NVIDIA H20-3e              
                                                                                                                                            
Nvidia driver version: 570.124.06                
cuDNN version: Probably one of the following:                                                                                               
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.10.2
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.10.2                                                                                            
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.10.2
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.10.2                                                                            
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.10.2
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.10.2                    
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.10.2                                                                                      
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.10.2                     
HIP runtime version: N/A                                                                                                                                                                                                                                                                 
MIOpen runtime version: N/A                                  
Is XNNPACK available: True                                                                                                                                                                                                                                                               
                                                                      
CPU:                                                                                                                                        
Architecture:                    x86_64                                                                                                     
CPU op-mode(s):                  32-bit, 64-bit                                                                                             
Address sizes:                   46 bits physical, 57 bits virtual                                                                          
Byte Order:                      Little Endian                                                                                              
CPU(s):                          208                                                                                                        
On-line CPU(s) list:             0-207                                                                                                      
Vendor ID:                       GenuineIntel                                                                                               
BIOS Vendor ID:                  Intel(R) Corporation                                                                                       
Model name:                      INTEL(R) XEON(R) PLATINUM 8563C                                                                            
BIOS Model name:                 INTEL(R) XEON(R) PLATINUM 8563C  CPU @ 2.6GHz                                                              
BIOS CPU family:                 179                                                                                                        
CPU family:                      6                                                                                                          
Model:                           207                                                                                                        
Thread(s) per core:              2                                                                                                          
Core(s) per socket:              52                                                                                                         
Socket(s):                       2                                                                                                          
Stepping:                        2                                                                                                          
Frequency boost:                 enabled                                                                                                    
CPU(s) scaling MHz:              77%                                                                                                        
CPU max MHz:                     4000.0000                                                                                                  
CPU min MHz:                     800.0000                                                                                                   
BogoMIPS:                        5200.00 
Virtualization:                  VT-x
L1d cache:                       4.9 MiB (104 instances)
L1i cache:                       3.3 MiB (104 instances)
L2 cache:                        208 MiB (104 instances)
L3 cache:                        640 MiB (2 instances)
NUMA node(s):                    2
NUMA node0 CPU(s):               0-51,104-155
NUMA node1 CPU(s):               52-103,156-207
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1:        Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers
Vulnerability Spectre v2:        Vulnerable, IBPB: disabled, STIBP: disabled
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected
Versions of relevant libraries:
[pip3] galore-torch==1.0
[pip3] numpy==2.2.6
[pip3] nvidia-cublas-cu12==12.6.4.1
[pip3] nvidia-cuda-cupti-cu12==12.6.80
[pip3] nvidia-cuda-nvrtc-cu12==12.6.77
[pip3] nvidia-cuda-runtime-cu12==12.6.77
[pip3] nvidia-cudnn-cu12==9.5.1.17
[pip3] nvidia-cufft-cu12==11.3.0.4
[pip3] nvidia-curand-cu12==10.3.7.77
[pip3] nvidia-cusolver-cu12==11.7.1.2
[pip3] nvidia-cusparse-cu12==12.5.4.2
[pip3] nvidia-cusparselt-cu12==0.6.3
[pip3] nvidia-nccl-cu12==2.26.2
[pip3] nvidia-nvjitlink-cu12==12.6.85
[pip3] nvidia-nvtx-cu12==12.6.77
[pip3] torch==2.7.1
[pip3] torch_memory_saver==0.0.8
[pip3] torchao==0.9.0
[pip3] torchaudio==2.7.1
[pip3] torchdata==0.11.0
[pip3] torchmetrics==1.7.3
[pip3] torchvision==0.22.1
[pip3] transformer_engine_torch==2.4.0
[pip3] triton==3.3.1
[conda] galore-torch                                1.0              pypi_0             pypi
[conda] numpy                                       2.2.6            pypi_0             pypi
[conda] nvidia-cublas-cu12                          12.6.4.1         pypi_0             pypi
[conda] nvidia-cuda-cupti-cu12                      12.6.80          pypi_0             pypi
[conda] nvidia-cuda-nvrtc-cu12                      12.6.77          pypi_0             pypi
[conda] nvidia-cuda-runtime-cu12                    12.6.77          pypi_0             pypi
[conda] nvidia-cudnn-cu12                           9.5.1.17         pypi_0             pypi
[conda] nvidia-cufft-cu12                           11.3.0.4         pypi_0             pypi
[conda] nvidia-curand-cu12                          10.3.7.77        pypi_0             pypi
[conda] nvidia-cusolver-cu12                        11.7.1.2         pypi_0             pypi
[conda] nvidia-cusparse-cu12                        12.5.4.2         pypi_0             pypi
[conda] nvidia-cusparselt-cu12                      0.6.3            pypi_0             pypi
[conda] nvidia-nccl-cu12                            2.26.2           pypi_0             pypi
[conda] nvidia-nvjitlink-cu12                       12.6.85          pypi_0             pypi
[conda] nvidia-nvtx-cu12                            12.6.77          pypi_0             pypi
[conda] torch                                       2.7.1            pypi_0             pypi
[conda] torch-memory-saver                          0.0.8            pypi_0             pypi
[conda] torchao                                     0.9.0            pypi_0             pypi
[conda] torchaudio                                  2.7.1            pypi_0             pypi
[conda] torchdata                                   0.11.0           pypi_0             pypi
[conda] torchmetrics                                1.7.3            pypi_0             pypi
[conda] torchvision                                 0.22.1           pypi_0             pypi
[conda] transformer-engine-torch                    2.4.0            pypi_0             pypi
[conda] triton                                      3.3.1            pypi_0             pypi

my launch scripts

#!/bin/bash

set -exo pipefail

NODES=1
PP=2
TP=2
EP=2
ETP=2
CP=1
VLLM_TP=4

DIST_CKPT_PATH=
LLM=
train_files=
test_files=

PROJ_NAME=gsm8k-megatron
EXP_NAME=moonlight_vlm_tp${TP}pp${PP}ep${EP}etp${ETP}

cmd=(
    +actor_rollout_ref.megatron.extra.num_layers_in_last_pipeline_stage=13
    +critic.megatron.extra.num_layers_in_last_pipeline_stage=13

    actor_rollout_ref.actor.megatron.dist_checkpointing_path=$DIST_CKPT_PATH
    actor_rollout_ref.actor.megatron.grad_offload=True
    actor_rollout_ref.actor.megatron.optimizer_offload=True
    actor_rollout_ref.actor.megatron.param_offload=True
    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$PP
    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP
    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP
    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$TP
    actor_rollout_ref.actor.megatron.use_dist_checkpointing=True
    actor_rollout_ref.actor.optim.lr=1e-6

    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2
    actor_rollout_ref.actor.ppo_mini_batch_size=32
    actor_rollout_ref.actor.use_kl_loss=True

    actor_rollout_ref.model.enable_gradient_checkpointing=True
    actor_rollout_ref.model.path=$LLM
    actor_rollout_ref.model.trust_remote_code=True

    actor_rollout_ref.ref.log_prob_micro_batch_size=16
    actor_rollout_ref.ref.megatron.dist_checkpointing_path=$DIST_CKPT_PATH
    actor_rollout_ref.ref.megatron.param_offload=True
    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=$PP
    actor_rollout_ref.ref.megatron.expert_model_parallel_size=$EP
    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=$ETP
    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=$TP
    actor_rollout_ref.ref.megatron.use_dist_checkpointing=True

    actor_rollout_ref.rollout.mode=sync # async
    actor_rollout_ref.rollout.gpu_memory_utilization=0.65
    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2
    actor_rollout_ref.rollout.name=sglang
    actor_rollout_ref.rollout.tensor_model_parallel_size=$VLLM_TP

    algorithm.adv_estimator=grpo
    algorithm.use_kl_in_reward=False

    data.filter_overlong_prompts=False
    data.max_prompt_length=1024
    data.max_response_length=512
    data.train_batch_size=128
    data.train_files="$train_files"
    data.val_files="$test_files"
    data.truncation='error'
    data.trust_remote_code=True

    trainer.critic_warmup=0
    trainer.experiment_name=${EXP_NAME}
    trainer.logger=['console','tensorboard']
    trainer.n_gpus_per_node=8
    trainer.nnodes=1
    trainer.project_name=${PROJ_NAME}
    trainer.save_freq=-1
    trainer.test_freq=-1
    trainer.total_epochs=7
    trainer.val_before_train=False
)

dt=`date '+%Y-%m-%d_%H-%M-%S'`
log_dir=./saves/${PROJ_NAME}/${EXP_NAME}/log/${dt}
mkdir -p ${log_dir}

python3 -m verl.trainer.main_ppo --config-path=config \
        --config-name='ppo_megatron_trainer.yaml' "${cmd[@]}" 2>&1 | tee ${log_dir}/main.log

torch profile results for per tensor

	megatron per_tensor_generator	sglang update_weights_from_tensor
trace json	qs-26125-730588-raycluster-tlmpx-main-worker-z97k6_102636.1752135752961638207.pt.trace.json	qs-26125-730588-raycluster-tlmpx-main-worker-z97k6_129525.1752136836706965956.pt.trace.json
profiler step position	megatron_sglang.py#L92	sglang/python/sglang/srt/model_executor/model_runner.py Line 838 in 4a0d191 def update_weights_from_tensor(

cc @zhaochenyang20

[fzyzcjy]

If the torch profile mimics the primary slowness we care about, then besides directly optimizing cpu operations like reducing calling mesh[...] and load_weights and per_tensor_generator and do batching on tensors, we may also pre-create tensors shared across processes and then do copy to it instead of creating and disposing them over and over again as a workaround (that will introduce extra same-gpu copy overhead, but if we use 128 gpu for 671B DeepSeek, then we will copy ~6GB per gpu, thus ~2ms overall, which seems small compared to the total 40s update time; would be great to UT whether mem handle is really this slow or not as well). Also doing overlapping between the two processes may help (just remove the await from update_weights_from_tensor and only await it in the next iteration and ensure tensors are not prematurely disposed).

Btw, I am curious why other systems still need 40s to update weights. Even if we let all ranks receive all 671B through 50GB rdma, it is only 14s (and iirc we do not need such broadcast).

[hebiao064]

A potential optimization would be:

• Learn from Slime's bucket update params: https://github.com/THUDM/slime/blob/8ebfaf683f149c7948dcbb19ee2f28e468cd4481/slime/ray/ppo_actor.py#L112

As we can see from the screenshot, batch update tensors is much faster then per tensor update

[Yangruipis]

my apologizes, the previous profile is not the latest verl branch and has a bug which is solved by volcengine/verl#2336 two days ago.

here is the latest profile for per tensor update_weights_from_tensor

qs-26125-731861-raycluster-hcmgb-main-worker-sxv8c_116150.1752200881675989264.pt.trace.json