vLLM weights are not synchronized with vllm_enable_sleep_mode=True

[muupan]

vLLM weights are not synchronized with vllm_enable_sleep_mode=True. As a result, GRPO training with vllm_mode="colocate" and vllm_enable_sleep_mode=True seems not working in the main branch after #4700.

Reproduction

code

# train_grpo_example_sleep.py
import argparse
import subprocess

from datasets import Dataset
from trl import GRPOConfig, GRPOTrainer
from trl.rewards import accuracy_reward


parser = argparse.ArgumentParser()
parser.add_argument("--vllm_enable_sleep_mode", action="store_true", default=False)
args = parser.parse_args()

dataset = Dataset.from_list([{"prompt": [{"role": "user", "content": "Who are you?"}]}])

commit = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode("utf-8").strip()

def length_reward(completions: list[list[dict[str, str]]], **kwargs) -> list[float | None]:
    # Shorter is better.
    return [1.0 / (len(completion[0]["content"]) + 1) for completion in completions]

args = GRPOConfig(
    output_dir=f"outputs_train_grpo_example_sleep/sleep_{args.vllm_enable_sleep_mode}_{commit}",
    use_vllm=True,
    vllm_mode="colocate",
    vllm_enable_sleep_mode=args.vllm_enable_sleep_mode,
    max_steps=10,
    num_train_epochs=1,
    logging_steps=1,
    save_strategy="no",
    report_to="tensorboard",
    learning_rate=1e-5,
)

trainer = GRPOTrainer(
    args=args,
    model="Qwen/Qwen2-0.5B-Instruct",
    reward_funcs=length_reward,
    train_dataset=dataset,
)
trainer.train()

command

uv run --no-sync accelerate launch --num_processes 1 train_grpo_example_sleep.py --vllm_enable_sleep_mode
uv run --no-sync accelerate launch --num_processes 1 train_grpo_example_sleep.py

outputs (sleep mode enabled):

{'loss': 0.3442, 'grad_norm': 11.958919525146484, 'learning_rate': 1e-05, 'num_tokens': 586.0, 'completions/mean_length': 50.25, 'completions/min_length': 18.0, 'completions/max_length': 83.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 50.25, 'completions/min_terminated_length': 18.0, 'completions/max_terminated_length': 83.0, 'rewards/length_reward/mean': 0.005347811616957188, 'rewards/length_reward/std': 0.0037498800083994865, 'reward': 0.005347811616957188, 'reward_std': 0.003749879775568843, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.025552719831466675, 'sampling/sampling_logp_difference/max': 0.2405834197998047, 'sampling/importance_sampling_ratio/min': 0.7255696654
319763, 'sampling/importance_sampling_ratio/mean': 0.9870972633361816, 'sampling/importance_sampling_ratio/max': 1.2755951881408691, 'entropy': 1.353233814239502, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 3.4439458176493645, 'epoch': 1.0}
 10%|█████████████████████████████████▎                                                                                                                                                                                                                                                                                                           | 1/10 [00:03<00:31,  3.45s/it]
Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]
Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  2.40it/s]
                                                                                                                                                                                                                                                                                                                                                                                 {'loss': 0.0997, 'grad_norm': 4.2228617668151855, 'learning_rate': 9e-06, 'num_tokens': 1122.0, 'completions/mean_length': 44.0, 'completions/min_length': 13.0, 'completions/max_length': 86.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 44.0, 'completions/min_terminated_length': 13.0, 'completions/max_terminated_length': 86.0, 'rewards/length_reward/mean': 0.00694679282605648, 'rewards/length_reward/std': 0.005016538314521313, 'reward': 0.00694679282605648, 'reward_std': 0.005016538314521313, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.2853383719921112, 'sampling/sampling_logp_difference/max': 8.178083419799805, 'sampling/importance_sampling_ratio/min': 5.281738708617922e-07, 'sampling/importance_sampling_ratio/mean': 0.20064665377140045, 'sampling/importance_sampling_ratio/max': 1.559876799583435, 'entropy': 1.4431818723678589, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 2.479794653132558, 'epoch': 2.0}
 20%|██████████████████████████████████████████████████████████████████▌                                                                                                                                                                                                                                                                          | 2/10 [00:06<00:23,  2.90s/it]
Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]
Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  2.40it/s]

{'loss': -0.0076, 'grad_norm': 0.5156638026237488, 'learning_rate': 8.000000000000001e-06, 'num_tokens': 1735.0, 'completions/mean_length': 53.625, 'completions/min_length': 11.0, 'completions/max_length': 83.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 53.625, 'completions/min_terminated_length': 11.0, 'completions/max_terminated_length': 83.0, 'rewards/length_reward/mean': 0.006394295021891594, 'rewards/length_reward/std': 0.005767505615949631, 'reward': 0.006394295021891594, 'reward_std': 0.005767505615949631, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.4896039366722107, 'sampling/sampling_logp_difference/max': 15.178083419799805, 'sampling/importance_sampling_ratio/min': 2.6204335688767394e-17, 'sampling/importance_sampling_ratio/mean': 0.018652046099305153, 'sampling/importance_sampling_ratio/max': 0.09242421388626099, 'entropy': 1.193473219871521, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 2.444535758346319, 'epoch': 3.0}
 30%|███████████████████████████████████████████████████████████████████████████████████████████████████▉                                                                                                                                                                                                                                         | 3/10 [00:08<00:18,  2.71s/it]
Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]
Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  2.40it/s]

{'loss': -0.1631, 'grad_norm': 7.1509528160095215, 'learning_rate': 7e-06, 'num_tokens': 2310.0, 'completions/mean_length': 48.875, 'completions/min_length': 32.0, 'completions/max_length': 70.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 48.875, 'completions/min_terminated_length': 32.0, 'completions/max_terminated_length': 70.0, 'reward
s/length_reward/mean': 0.004292851313948631, 'rewards/length_reward/std': 0.0010315432446077466, 'reward': 0.004292851313948631, 'reward_std': 0.0010315432446077466, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.5148501396179199, 'sampling/sampling_logp_difference/max': 11.303083419799805, 'sampling/importance_sampling_ratio/min': 1.6145741
68806726e-10, 'sampling/importance_sampling_ratio/mean': 0.37434086203575134, 'sampling/importance_sampling_ratio/max': 2.9669249057769775, 'entropy': 1.1713554859161377, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 2.439166160300374, 'epoch': 4.0}
 40%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▏                                                                                                                                                                                                       | 4/10 [00:11<00:15,  2.64s/it]
Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]
Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  2.38it/s]

{'loss': 0.0471, 'grad_norm': 1.5320026874542236, 'learning_rate': 6e-06, 'num_tokens': 2702.0, 'completions/mean_length': 26.0, 'completions/min_length': 5.0, 'completions/max_length': 57.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 26.0, 'completions/min_terminated_length': 5.0, 'completions/max_terminated_length': 57.0, 'rewards/lengt
h_reward/mean': 0.024240167811512947, 'rewards/length_reward/std': 0.03596894070506096, 'reward': 0.024240167811512947, 'reward_std': 0.03596893697977066, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.5626580119132996, 'sampling/sampling_logp_difference/max': 7.303083419799805, 'sampling/importance_sampling_ratio/min': 0.0, 'sampling/import
ance_sampling_ratio/mean': 0.234994575381279, 'sampling/importance_sampling_ratio/max': 1.8631082773208618, 'entropy': 1.240384578704834, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 2.4113525077700615, 'epoch': 5.0}
 50%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▌                                                                                                                                                                      | 5/10 [00:13<00:12,  2.57s/it]
Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]
Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  2.43it/s]

{'loss': -0.0086, 'grad_norm': 1.3793572187423706, 'learning_rate': 5e-06, 'num_tokens': 3359.0, 'completions/mean_length': 59.125, 'completions/min_length': 8.0, 'completions/max_length': 114.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 59.125, 'completions/min_terminated_length': 8.0, 'completions/max_terminated_length': 114.0, 'reward
s/length_reward/mean': 0.007446803618222475, 'rewards/length_reward/std': 0.009431127458810806, 'reward': 0.007446803618222475, 'reward_std': 0.009431127458810806, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.6624722480773926, 'sampling/sampling_logp_difference/max': 8.193817138671875, 'sampling/importance_sampling_ratio/min': 4.0311181007
447796e-21, 'sampling/importance_sampling_ratio/mean': 0.02877916395664215, 'sampling/importance_sampling_ratio/max': 0.23017366230487823, 'entropy': 1.674418568611145, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 2.4947521574795246, 'epoch': 6.0}
 60%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▊                                                                                                                                     | 6/10 [00:15<00:10,  2.56s/it]
Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]
Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  2.40it/s]

{'loss': -0.0, 'grad_norm': 0.003574257018044591, 'learning_rate': 4.000000000000001e-06, 'num_tokens': 3856.0, 'completions/mean_length': 39.125, 'completions/min_length': 8.0, 'completions/max_length': 56.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 39.125, 'completions/min_terminated_length': 8.0, 'completions/max_terminated_length':
56.0, 'rewards/length_reward/mean': 0.0103794289752841, 'rewards/length_reward/std': 0.01303927507251501, 'reward': 0.0103794289752841, 'reward_std': 0.01303927507251501, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.8130478858947754, 'sampling/sampling_logp_difference/max': 8.381317138671875, 'sampling/importance_sampling_ratio/min': 2.131
5809459747903e-20, 'sampling/importance_sampling_ratio/mean': 4.2642277549020946e-05, 'sampling/importance_sampling_ratio/max': 0.0003398687986191362, 'entropy': 1.9041533470153809, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 2.4032083693891764, 'epoch': 7.
0}
 70%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████                                                                                                    | 7/10 [00:18<00:07,  2.52s/it]
Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]
Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  2.40it/s]

{'loss': -0.0, 'grad_norm': 5.522433639271185e-07, 'learning_rate': 3e-06, 'num_tokens': 4474.0, 'completions/mean_length': 54.25, 'completions/min_length': 28.0, 'completions/max_length': 95.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 54.25, 'completions/min_terminated_length': 28.0, 'completions/max_terminated_length': 95.0, 'rewards/
length_reward/mean': 0.00464101554825902, 'rewards/length_reward/std': 0.0025500054471194744, 'reward': 0.00464101554825902, 'reward_std': 0.0025500054471194744, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.6805005073547363, 'sampling/sampling_logp_difference/max': 7.5238447189331055, 'sampling/importance_sampling_ratio/min': 5.87165938954
2004e-16, 'sampling/importance_sampling_ratio/mean': 1.5323410096357293e-08, 'sampling/importance_sampling_ratio/max': 8.06930842145448e-08, 'entropy': 1.9354838132858276, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 2.5171711817383766, 'epoch': 8.0}
 80%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▍                                                                  | 8/10 [00:20<00:05,  2.53s/it]
Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]
Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  2.38it/s]

{'loss': 0.0005, 'grad_norm': 0.030856702476739883, 'learning_rate': 2.0000000000000003e-06, 'num_tokens': 4981.0, 'completions/mean_length': 40.375, 'completions/min_length': 6.0, 'completions/max_length': 82.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 40.375, 'completions/min_terminated_length': 6.0, 'completions/max_terminated_length
': 82.0, 'rewards/length_reward/mean': 0.018925247713923454, 'rewards/length_reward/std': 0.03732370585203171, 'reward': 0.018925247713923454, 'reward_std': 0.03732370585203171, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.9693665504455566, 'sampling/sampling_logp_difference/max': 9.467376708984375, 'sampling/importance_sampling_ratio/min'
: 1.8341662912330806e-23, 'sampling/importance_sampling_ratio/mean': 0.003137270687147975, 'sampling/importance_sampling_ratio/max': 0.025024525821208954, 'entropy': 1.7585139274597168, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 2.4585518520325422, 'epoch'
: 9.0}
 90%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▋                                 | 9/10 [00:23<00:02,  2.52s/it]
Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]
Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:00<00:00,  2.40it/s]

{'loss': 0.0, 'grad_norm': 4.947049092152156e-06, 'learning_rate': 1.0000000000000002e-06, 'num_tokens': 5587.0, 'completions/mean_length': 52.75, 'completions/min_length': 15.0, 'completions/max_length': 85.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 52.75, 'completions/min_terminated_length': 15.0, 'completions/max_terminated_length':
 85.0, 'rewards/length_reward/mean': 0.006554092280566692, 'rewards/length_reward/std': 0.006052262615412474, 'reward': 0.006554092280566692, 'reward_std': 0.006052262149751186, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.8090547323226929, 'sampling/sampling_logp_difference/max': 7.428083419799805, 'sampling/importance_sampling_ratio/min'
: 3.197867121484542e-21, 'sampling/importance_sampling_ratio/mean': 3.2516194892195927e-07, 'sampling/importance_sampling_ratio/max': 2.595229943835875e-06, 'entropy': 2.1421802043914795, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 3.329218603670597, 'epoch
': 10.0}
{'train_runtime': 26.87, 'train_samples_per_second': 2.977, 'train_steps_per_second': 0.372, 'train_loss': 0.031205586939371255, 'epoch': 10.0}
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [00:26<00:00,  2.69s/it]

outputs (sleep mode disabled):

{'loss': 0.3442, 'grad_norm': 11.958919525146484, 'learning_rate': 1e-05, 'num_tokens': 586.0, 'completions/mean_length': 50.25, 'completions/min_length': 18.0, 'completions/max_length': 83.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 50.25, 'completions/min_terminated_length': 18.0, 'completions/max_terminated_length': 83.0, 'rewards/length_reward/mean': 0.005347811616957188, 'rewards/length_reward/std': 0.0037498800083994865, 'reward': 0.005347811616957188, 'reward_std': 0.003749879775568843, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.025552719831466675, 'sampling/sampling_logp_difference/max': 0.2405834197998047, 'sampling/importance_sampling_ratio/min': 0.7255696654319763, 'sampling/importance_sampling_ratio/mean': 0.9870972633361816, 'sampling/importance_sampling_ratio/max': 1.2755951881408691, 'entropy': 1.353233814239502, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 2.475095007568598, 'epoch': 1.0}                  {'loss': 0.8298, 'grad_norm': 36.34674072265625, 'learning_rate': 9e-06, 'num_tokens': 984.0, 'completions/mean_length': 26.75, 'completions/min_length': 12.0, 'completions/max_length': 46.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 26.75, 'completions/min_terminated_length': 12.0, 'completions/max_terminated_length': 46.0, 'rewards/len
gth_reward/mean': 0.011167747899889946, 'rewards/length_reward/std': 0.0069429934956133366, 'reward': 0.011167747899889946, 'reward_std': 0.0069429934956133366, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.022812088951468468, 'sampling/sampling_logp_difference/max': 0.27854061126708984, 'sampling/importance_sampling_ratio/min': 0.733009517
1928406, 'sampling/importance_sampling_ratio/mean': 1.1807054281234741, 'sampling/importance_sampling_ratio/max': 1.9025675058364868, 'entropy': 0.836448609828949, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.500464078038931, 'epoch': 2.0}
{'loss': -0.2303, 'grad_norm': 180.8460693359375, 'learning_rate': 8.000000000000001e-06, 'num_tokens': 1257.0, 'completions/mean_length': 11.125, 'completions/min_length': 8.0, 'completions/max_length': 12.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 11.125, 'completions/min_terminated_length': 8.0, 'completions/max_terminated_length':
12.0, 'rewards/length_reward/mean': 0.022984981536865234, 'rewards/length_reward/std': 0.005988886579871178, 'reward': 0.022984981536865234, 'reward_std': 0.005988887045532465, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.04025040194392204, 'sampling/sampling_logp_difference/max': 0.6134061813354492, 'sampling/importance_sampling_ratio/min
': 0.6521506905555725, 'sampling/importance_sampling_ratio/mean': 0.9877357482910156, 'sampling/importance_sampling_ratio/max': 1.9371200799942017, 'entropy': 0.4185393154621124, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.4148744363337755, 'epoch': 3.0}
{'loss': 0.0762, 'grad_norm': 201.97996520996094, 'learning_rate': 7e-06, 'num_tokens': 1499.0, 'completions/mean_length': 7.25, 'completions/min_length': 6.0, 'completions/max_length': 8.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 7.25, 'completions/min_terminated_length': 6.0, 'completions/max_terminated_length': 8.0, 'rewards/length_
reward/mean': 0.0457589291036129, 'rewards/length_reward/std': 0.013862923718988895, 'reward': 0.0457589291036129, 'reward_std': 0.01386292278766632, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.009500036016106606, 'sampling/sampling_logp_difference/max': 0.07426857948303223, 'sampling/importance_sampling_ratio/min': 0.9473329186439514, 's
ampling/importance_sampling_ratio/mean': 0.988956093788147, 'sampling/importance_sampling_ratio/max': 1.0583280324935913, 'entropy': 0.2068965584039688, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.3868972603231668, 'epoch': 4.0}
{'loss': -0.0992, 'grad_norm': 97.04837799072266, 'learning_rate': 6e-06, 'num_tokens': 1731.0, 'completions/mean_length': 6.0, 'completions/min_length': 6.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 6.0, 'completions/min_terminated_length': 6.0, 'completions/max_terminated_length': 6.0, 'rewards/length_re
ward/mean': 0.0616319440305233, 'rewards/length_reward/std': 0.0024552317336201668, 'reward': 0.0616319440305233, 'reward_std': 0.0024552324321120977, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.016970515251159668, 'sampling/sampling_logp_difference/max': 0.34142231941223145, 'sampling/importance_sampling_ratio/min': 0.7267212271690369, '
sampling/importance_sampling_ratio/mean': 1.0187859535217285, 'sampling/importance_sampling_ratio/max': 1.0605095624923706, 'entropy': 0.19921875, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.3645392376929522, 'epoch': 5.0}
{'loss': 0.0, 'grad_norm': 0.0, 'learning_rate': 5e-06, 'num_tokens': 1963.0, 'completions/mean_length': 6.0, 'completions/min_length': 6.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 6.0, 'completions/min_terminated_length': 6.0, 'completions/max_terminated_length': 6.0, 'rewards/length_reward/mean': 0.0625
, 'rewards/length_reward/std': 0.0, 'reward': 0.0625, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'sampling/sampling_logp_difference/mean': 0.0010602406691759825, 'sampling/sampling_logp_difference/max': 0.003823213279247284, 'sampling/importance_sampling_ratio/min': 0.9938066601753235, 'sampling/importance_sampling_ratio/mean': 0.9938066601753235, 'sampling/impo
rtance_sampling_ratio/max': 0.9938066601753235, 'entropy': 0.0478515625, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.325540291145444, 'epoch': 6.0}
{'loss': 0.0, 'grad_norm': 0.0, 'learning_rate': 4.000000000000001e-06, 'num_tokens': 2195.0, 'completions/mean_length': 6.0, 'completions/min_length': 6.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 6.0, 'completions/min_terminated_length': 6.0, 'completions/max_terminated_length': 6.0, 'rewards/length_rewa
rd/mean': 0.0625, 'rewards/length_reward/std': 0.0, 'reward': 0.0625, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'sampling/sampling_logp_difference/mean': 0.004386706743389368, 'sampling/sampling_logp_difference/max': 0.021789629012346268, 'sampling/importance_sampling_ratio/min': 0.9740239977836609, 'sampling/importance_sampling_ratio/mean': 0.9740239977836609,
 'sampling/importance_sampling_ratio/max': 0.9740239977836609, 'entropy': 0.0582682304084301, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.324770538136363, 'epoch': 7.0}
{'loss': 0.0, 'grad_norm': 0.0, 'learning_rate': 3e-06, 'num_tokens': 2427.0, 'completions/mean_length': 6.0, 'completions/min_length': 6.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 6.0, 'completions/min_terminated_length': 6.0, 'completions/max_terminated_length': 6.0, 'rewards/length_reward/mean': 0.0625
, 'rewards/length_reward/std': 0.0, 'reward': 0.0625, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'sampling/sampling_logp_difference/mean': 0.0014839638024568558, 'sampling/sampling_logp_difference/max': 0.005859885364770889, 'sampling/importance_sampling_ratio/min': 0.9913414716720581, 'sampling/importance_sampling_ratio/mean': 0.9913414716720581, 'sampling/impo
rtance_sampling_ratio/max': 0.9913414716720581, 'entropy': 0.0576171875, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.332946429029107, 'epoch': 8.0}
{'loss': 0.0, 'grad_norm': 0.0, 'learning_rate': 2.0000000000000003e-06, 'num_tokens': 2659.0, 'completions/mean_length': 6.0, 'completions/min_length': 6.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 6.0, 'completions/min_terminated_length': 6.0, 'completions/max_terminated_length': 6.0, 'rewards/length_rew
ard/mean': 0.0625, 'rewards/length_reward/std': 0.0, 'reward': 0.0625, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'sampling/sampling_logp_difference/mean': 0.0009608844993636012, 'sampling/sampling_logp_difference/max': 0.002919524908065796, 'sampling/importance_sampling_ratio/min': 1.000125765800476, 'sampling/importance_sampling_ratio/mean': 1.000125765800476,
 'sampling/importance_sampling_ratio/max': 1.000125765800476, 'entropy': 0.0768229141831398, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.339396871626377, 'epoch': 9.0}
{'loss': 0.091, 'grad_norm': 133.67459106445312, 'learning_rate': 1.0000000000000002e-06, 'num_tokens': 2888.0, 'completions/mean_length': 5.625, 'completions/min_length': 3.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 5.625, 'completions/min_terminated_length': 3.0, 'completions/max_terminated_length': 6.0
, 'rewards/length_reward/mean': 0.0725446417927742, 'rewards/length_reward/std': 0.028410542756319046, 'reward': 0.0725446417927742, 'reward_std': 0.028410542756319046, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.017341550439596176, 'sampling/sampling_logp_difference/max': 0.2619295120239258, 'sampling/importance_sampling_ratio/min': 0.92
93384552001953, 'sampling/importance_sampling_ratio/mean': 0.9763498306274414, 'sampling/importance_sampling_ratio/max': 1.305429220199585, 'entropy': 0.10833333432674408, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.3455794714391232, 'epoch': 10.0}
{'train_runtime': 15.2597, 'train_samples_per_second': 5.243, 'train_steps_per_second': 0.655, 'train_loss': 0.10115763023495675, 'epoch': 10.0}

When vllm_enable_sleep_mode=False, the model completion length gets soon shorter, leading increased rewards. When vllm_enable_sleep_mode=True, completion lengths and rewards are not changing.

I will write the cause and how it can be fixed below.

System Info

trl env did not work: ImportError: cannot import name '__version__' from 'trl' (unknown location)

trl b86b760
torch 2.8.0
vllm 0.11.0
gpu: NVIDIA H200 x 1

Checklist

• I have checked that my issue isn't already filed (see open issues)
• I have included my system information
• Any code provided is minimal, complete, and reproducible (more on MREs)
• Any code provided is properly formatted in code blocks, (no screenshot, more on code blocks)
• Any traceback provided is complete

[muupan]

What causes this problem

It seems like the PR #4700 caused it.

Before the PR, self._move_model_to_vllm() is called after self.llm.collective_rpc("reload_weights"), which ensure the weights are synchronized.

trl/trl/trainer/grpo_trainer.py

Lines 1306 to 1316 in 226ef57

	if self.vllm_mode == "colocate" and self.args.vllm_enable_sleep_mode:
	# wake up colocated vLLM instances if needed
	torch.cuda.empty_cache() # required to avoid OOM in some cases
	self.llm.wake_up(tags=["weights"])
	# Work around for https://github.com/vllm-project/vllm/issues/29341
	self.llm.collective_rpc("reload_weights")
	# First, update the vLLM weights if needed
	if self.state.global_step != self._last_loaded_step:
	self._move_model_to_vllm()
	self._last_loaded_step = self.state.global_step

After the PR, the order is reversed. VLLMGeneration.sync_weights is called first, and then self.llm.collective_rpc("reload_weights"), which seems to revert the synchronized weights. It seems like the comment "idempotent if already awake from sync_weights" does not apply to self.llm.collective_rpc("reload_weights").

trl/trl/trainer/grpo_trainer.py

Line 1567 in b86b760

self.vllm_generation.sync_weights()

trl/trl/generation/vllm_generation.py

Lines 568 to 577 in b86b760

	# Wake up colocated vLLM weights if needed (idempotent if already awake from sync_weights)
	if self.mode == "colocate" and self.enable_sleep_mode:
	empty_cache() # required to avoid OOM in some cases
	self.llm.wake_up(tags=["weights"])
	# Work around for https://github.com/vllm-project/vllm/issues/29341
	try:
	self.llm.collective_rpc("reload_weights")
	except NotImplementedError:
	# Non-CUDA vLLM backends (e.g., vllm-ascend's NPUWorkerV1), don't implement reload_weights
	pass

How to fix

I verified that calling sync_weights instead of self.llm.collective_rpc("reload_weights") resolves the issue. I opened a PR: #5313.

diff --git a/trl/generation/vllm_generation.py b/trl/generation/vllm_generation.py
index f93d3c32..d908366c 100644
--- a/trl/generation/vllm_generation.py
+++ b/trl/generation/vllm_generation.py
@@ -565,16 +565,9 @@ class VLLMGeneration:
         repetition_penalty = self.repetition_penalty
         max_completion_length = self.max_completion_length

-        # Wake up colocated vLLM weights if needed (idempotent if already awake from sync_weights)
-        if self.mode == "colocate" and self.enable_sleep_mode:
-            empty_cache()  # required to avoid OOM in some cases
-            self.llm.wake_up(tags=["weights"])
-            # Work around for https://github.com/vllm-project/vllm/issues/29341
-            try:
-                self.llm.collective_rpc("reload_weights")
-            except NotImplementedError:
-                # Non-CUDA vLLM backends (e.g., vllm-ascend's NPUWorkerV1), don't implement reload_weights
-                pass
+        # Since vLLM sleep level 2 discards weights, we need to sync weights to wake up the engine
+        if self.mode == "colocate" and self.enable_sleep_mode and self.llm.llm_engine.is_sleeping():
+            self.sync_weights()

         # Generate completions using vLLM: gather all prompts and use them in a single call in the main process
         if self.mode == "server":

outputs (sleep mode enabled):

{'loss': 0.3442, 'grad_norm': 11.958919525146484, 'learning_rate': 1e-05, 'num_tokens': 586.0, 'completions/mean_length': 50.25, 'completions/min_length': 18.0, 'completions/max_length': 83.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 50.25, 'completions/
min_terminated_length': 18.0, 'completions/max_terminated_length': 83.0, 'rewards/length_reward/mean': 0.005347811616957188, 'rewards/length_reward/std': 0.0037498800083994865, 'reward': 0.005347811616957188, 'reward_std': 0.003749879775568843, 'frac_reward_zero_std': 0.0, 'sampling/s
ampling_logp_difference/mean': 0.025552719831466675, 'sampling/sampling_logp_difference/max': 0.2405834197998047, 'sampling/importance_sampling_ratio/min': 0.7255696654319763, 'sampling/importance_sampling_ratio/mean': 0.9870972633361816, 'sampling/importance_sampling_ratio/max': 1.27
55951881408691, 'entropy': 1.353233814239502, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 3.2483372762799263, 'epoch': 1.0}
{'loss': 0.8298, 'grad_norm': 36.34674072265625, 'learning_rate': 9e-06, 'num_tokens': 984.0, 'completions/mean_length': 26.75, 'completions/min_length': 12.0, 'completions/max_length': 46.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 26.75, 'completions/m
in_terminated_length': 12.0, 'completions/max_terminated_length': 46.0, 'rewards/length_reward/mean': 0.011167747899889946, 'rewards/length_reward/std': 0.0069429934956133366, 'reward': 0.011167747899889946, 'reward_std': 0.0069429934956133366, 'frac_reward_zero_std': 0.0, 'sampling/s
ampling_logp_difference/mean': 0.022812088951468468, 'sampling/sampling_logp_difference/max': 0.27854061126708984, 'sampling/importance_sampling_ratio/min': 0.7330095171928406, 'sampling/importance_sampling_ratio/mean': 1.1807054281234741, 'sampling/importance_sampling_ratio/max': 1.9
025675058364868, 'entropy': 0.836448609828949, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.9199919868260622, 'epoch': 2.0}
{'loss': -0.2303, 'grad_norm': 180.8460693359375, 'learning_rate': 8.000000000000001e-06, 'num_tokens': 1257.0, 'completions/mean_length': 11.125, 'completions/min_length': 8.0, 'completions/max_length': 12.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 11.
125, 'completions/min_terminated_length': 8.0, 'completions/max_terminated_length': 12.0, 'rewards/length_reward/mean': 0.022984981536865234, 'rewards/length_reward/std': 0.005988886579871178, 'reward': 0.022984981536865234, 'reward_std': 0.005988887045532465, 'frac_reward_zero_std':
0.0, 'sampling/sampling_logp_difference/mean': 0.04025040194392204, 'sampling/sampling_logp_difference/max': 0.6134061813354492, 'sampling/importance_sampling_ratio/min': 0.6521506905555725, 'sampling/importance_sampling_ratio/mean': 0.9877357482910156, 'sampling/importance_sampling_r
atio/max': 1.9371200799942017, 'entropy': 0.4185393154621124, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.8387768473476171, 'epoch': 3.0}
{'loss': 0.0762, 'grad_norm': 201.97996520996094, 'learning_rate': 7e-06, 'num_tokens': 1499.0, 'completions/mean_length': 7.25, 'completions/min_length': 6.0, 'completions/max_length': 8.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 7.25, 'completions/min
_terminated_length': 6.0, 'completions/max_terminated_length': 8.0, 'rewards/length_reward/mean': 0.0457589291036129, 'rewards/length_reward/std': 0.013862923718988895, 'reward': 0.0457589291036129, 'reward_std': 0.01386292278766632, 'frac_reward_zero_std': 0.0, 'sampling/sampling_log
p_difference/mean': 0.009500036016106606, 'sampling/sampling_logp_difference/max': 0.07426857948303223, 'sampling/importance_sampling_ratio/min': 0.9473329186439514, 'sampling/importance_sampling_ratio/mean': 0.988956093788147, 'sampling/importance_sampling_ratio/max': 1.0583280324935
913, 'entropy': 0.2068965584039688, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 3.378942394629121, 'epoch': 4.0}
{'loss': -0.0992, 'grad_norm': 97.04837799072266, 'learning_rate': 6e-06, 'num_tokens': 1731.0, 'completions/mean_length': 6.0, 'completions/min_length': 6.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 6.0, 'completions/min_t
erminated_length': 6.0, 'completions/max_terminated_length': 6.0, 'rewards/length_reward/mean': 0.0616319440305233, 'rewards/length_reward/std': 0.0024552317336201668, 'reward': 0.0616319440305233, 'reward_std': 0.0024552324321120977, 'frac_reward_zero_std': 0.0, 'sampling/sampling_lo
gp_difference/mean': 0.016970515251159668, 'sampling/sampling_logp_difference/max': 0.34142231941223145, 'sampling/importance_sampling_ratio/min': 0.7267212271690369, 'sampling/importance_sampling_ratio/mean': 1.0187859535217285, 'sampling/importance_sampling_ratio/max': 1.06050956249
23706, 'entropy': 0.19921875, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.8896218035370111, 'epoch': 5.0}
{'loss': 0.0, 'grad_norm': 0.0, 'learning_rate': 5e-06, 'num_tokens': 1963.0, 'completions/mean_length': 6.0, 'completions/min_length': 6.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 6.0, 'completions/min_terminated_length':
 6.0, 'completions/max_terminated_length': 6.0, 'rewards/length_reward/mean': 0.0625, 'rewards/length_reward/std': 0.0, 'reward': 0.0625, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'sampling/sampling_logp_difference/mean': 0.0010602406691759825, 'sampling/sampling_logp_difference
/max': 0.003823213279247284, 'sampling/importance_sampling_ratio/min': 0.9938066601753235, 'sampling/importance_sampling_ratio/mean': 0.9938066601753235, 'sampling/importance_sampling_ratio/max': 0.9938066601753235, 'entropy': 0.0478515625, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_
min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.7937180362641811, 'epoch': 6.0}
{'loss': 0.0, 'grad_norm': 0.0, 'learning_rate': 4.000000000000001e-06, 'num_tokens': 2195.0, 'completions/mean_length': 6.0, 'completions/min_length': 6.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 6.0, 'completions/min_ter
minated_length': 6.0, 'completions/max_terminated_length': 6.0, 'rewards/length_reward/mean': 0.0625, 'rewards/length_reward/std': 0.0, 'reward': 0.0625, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'sampling/sampling_logp_difference/mean': 0.004386706743389368, 'sampling/sampling_
logp_difference/max': 0.021789629012346268, 'sampling/importance_sampling_ratio/min': 0.9740239977836609, 'sampling/importance_sampling_ratio/mean': 0.9740239977836609, 'sampling/importance_sampling_ratio/max': 0.9740239977836609, 'entropy': 0.0582682304084301, 'clip_ratio/low_mean':
0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.7724495846778154, 'epoch': 7.0}
{'loss': 0.0, 'grad_norm': 0.0, 'learning_rate': 3e-06, 'num_tokens': 2427.0, 'completions/mean_length': 6.0, 'completions/min_length': 6.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 6.0, 'completions/min_terminated_length':
 6.0, 'completions/max_terminated_length': 6.0, 'rewards/length_reward/mean': 0.0625, 'rewards/length_reward/std': 0.0, 'reward': 0.0625, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'sampling/sampling_logp_difference/mean': 0.0014839638024568558, 'sampling/sampling_logp_difference
/max': 0.005859885364770889, 'sampling/importance_sampling_ratio/min': 0.9913414716720581, 'sampling/importance_sampling_ratio/mean': 0.9913414716720581, 'sampling/importance_sampling_ratio/max': 0.9913414716720581, 'entropy': 0.0576171875, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_
min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.7585801929235458, 'epoch': 8.0}
{'loss': 0.0, 'grad_norm': 0.0, 'learning_rate': 2.0000000000000003e-06, 'num_tokens': 2659.0, 'completions/mean_length': 6.0, 'completions/min_length': 6.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 6.0, 'completions/min_te
rminated_length': 6.0, 'completions/max_terminated_length': 6.0, 'rewards/length_reward/mean': 0.0625, 'rewards/length_reward/std': 0.0, 'reward': 0.0625, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'sampling/sampling_logp_difference/mean': 0.0009608844993636012, 'sampling/samplin
g_logp_difference/max': 0.002919524908065796, 'sampling/importance_sampling_ratio/min': 1.000125765800476, 'sampling/importance_sampling_ratio/mean': 1.000125765800476, 'sampling/importance_sampling_ratio/max': 1.000125765800476, 'entropy': 0.0768229141831398, 'clip_ratio/low_mean': 0
.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 2.0748476274311543, 'epoch': 9.0}
{'loss': 0.091, 'grad_norm': 133.67459106445312, 'learning_rate': 1.0000000000000002e-06, 'num_tokens': 2888.0, 'completions/mean_length': 5.625, 'completions/min_length': 3.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 5.625, 'completions/min_terminated_length': 3.0, 'completions/max_terminated_length': 6.0, 'rewards/length_reward/mean': 0.0725446417927742, 'rewards/length_reward/std': 0.028410542756319046, 'reward': 0.0725446417927742, 'reward_std': 0.028410542756319046, 'frac_reward_zero_std': 0.0, 'sampling/sampling_logp_difference/mean': 0.017341550439596176, 'sampling/sampling_logp_difference/max': 0.2619295120239258, 'sampling/importance_sampling_ratio/min': 0.9293384552001953, 'sampling/importance_sampling_ratio/mean': 0.9763498306274414, 'sampling/importance_sampling_ratio/ma
x': 1.305429220199585, 'entropy': 0.10833333432674408, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 2.111778823658824, 'epoch': 10.0}
{'train_runtime': 22.357, 'train_samples_per_second': 3.578, 'train_steps_per_second': 0.447, 'train_loss': 0.10115763023495675, 'epoch': 10.0}

outputs (sleep mode disabled):

{'loss': 0.3442, 'grad_norm': 11.958919525146484, 'learning_rate': 1e-05, 'num_tokens': 586.0, 'completions/mean_length': 50.25, 'completions/min_length': 18.0, 'completions/max_length': 83.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 50.25, 'completions/
min_terminated_length': 18.0, 'completions/max_terminated_length': 83.0, 'rewards/length_reward/mean': 0.005347811616957188, 'rewards/length_reward/std': 0.0037498800083994865, 'reward': 0.005347811616957188, 'reward_std': 0.003749879775568843, 'frac_reward_zero_std': 0.0, 'sampling/s
ampling_logp_difference/mean': 0.025552719831466675, 'sampling/sampling_logp_difference/max': 0.2405834197998047, 'sampling/importance_sampling_ratio/min': 0.7255696654319763, 'sampling/importance_sampling_ratio/mean': 0.9870972633361816, 'sampling/importance_sampling_ratio/max': 1.27
55951881408691, 'entropy': 1.353233814239502, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 2.3536373246461153, 'epoch': 1.0}
{'loss': 0.8298, 'grad_norm': 36.34674072265625, 'learning_rate': 9e-06, 'num_tokens': 984.0, 'completions/mean_length': 26.75, 'completions/min_length': 12.0, 'completions/max_length': 46.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 26.75, 'completions/m
in_terminated_length': 12.0, 'completions/max_terminated_length': 46.0, 'rewards/length_reward/mean': 0.011167747899889946, 'rewards/length_reward/std': 0.0069429934956133366, 'reward': 0.011167747899889946, 'reward_std': 0.0069429934956133366, 'frac_reward_zero_std': 0.0, 'sampling/s
ampling_logp_difference/mean': 0.022812088951468468, 'sampling/sampling_logp_difference/max': 0.27854061126708984, 'sampling/importance_sampling_ratio/min': 0.7330095171928406, 'sampling/importance_sampling_ratio/mean': 1.1807054281234741, 'sampling/importance_sampling_ratio/max': 1.9
025675058364868, 'entropy': 0.836448609828949, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.3384274262934923, 'epoch': 2.0}
{'loss': -0.2303, 'grad_norm': 180.8460693359375, 'learning_rate': 8.000000000000001e-06, 'num_tokens': 1257.0, 'completions/mean_length': 11.125, 'completions/min_length': 8.0, 'completions/max_length': 12.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 11.
125, 'completions/min_terminated_length': 8.0, 'completions/max_terminated_length': 12.0, 'rewards/length_reward/mean': 0.022984981536865234, 'rewards/length_reward/std': 0.005988886579871178, 'reward': 0.022984981536865234, 'reward_std': 0.005988887045532465, 'frac_reward_zero_std':
0.0, 'sampling/sampling_logp_difference/mean': 0.04025040194392204, 'sampling/sampling_logp_difference/max': 0.6134061813354492, 'sampling/importance_sampling_ratio/min': 0.6521506905555725, 'sampling/importance_sampling_ratio/mean': 0.9877357482910156, 'sampling/importance_sampling_r
atio/max': 1.9371200799942017, 'entropy': 0.4185393154621124, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.253534099087119, 'epoch': 3.0}
{'loss': 0.0762, 'grad_norm': 201.97996520996094, 'learning_rate': 7e-06, 'num_tokens': 1499.0, 'completions/mean_length': 7.25, 'completions/min_length': 6.0, 'completions/max_length': 8.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 7.25, 'completions/min
_terminated_length': 6.0, 'completions/max_terminated_length': 8.0, 'rewards/length_reward/mean': 0.0457589291036129, 'rewards/length_reward/std': 0.013862923718988895, 'reward': 0.0457589291036129, 'reward_std': 0.01386292278766632, 'frac_reward_zero_std': 0.0, 'sampling/sampling_log
p_difference/mean': 0.009500036016106606, 'sampling/sampling_logp_difference/max': 0.07426857948303223, 'sampling/importance_sampling_ratio/min': 0.9473329186439514, 'sampling/importance_sampling_ratio/mean': 0.988956093788147, 'sampling/importance_sampling_ratio/max': 1.0583280324935
913, 'entropy': 0.2068965584039688, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.2504651956260204, 'epoch': 4.0}
{'loss': -0.0992, 'grad_norm': 97.04837799072266, 'learning_rate': 6e-06, 'num_tokens': 1731.0, 'completions/mean_length': 6.0, 'completions/min_length': 6.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 6.0, 'completions/min_t
erminated_length': 6.0, 'completions/max_terminated_length': 6.0, 'rewards/length_reward/mean': 0.0616319440305233, 'rewards/length_reward/std': 0.0024552317336201668, 'reward': 0.0616319440305233, 'reward_std': 0.0024552324321120977, 'frac_reward_zero_std': 0.0, 'sampling/sampling_lo
gp_difference/mean': 0.016970515251159668, 'sampling/sampling_logp_difference/max': 0.34142231941223145, 'sampling/importance_sampling_ratio/min': 0.7267212271690369, 'sampling/importance_sampling_ratio/mean': 1.0187859535217285, 'sampling/importance_sampling_ratio/max': 1.06050956249
23706, 'entropy': 0.19921875, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.2208976987749338, 'epoch': 5.0}
{'loss': 0.0, 'grad_norm': 0.0, 'learning_rate': 5e-06, 'num_tokens': 1963.0, 'completions/mean_length': 6.0, 'completions/min_length': 6.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 6.0, 'completions/min_terminated_length':
 6.0, 'completions/max_terminated_length': 6.0, 'rewards/length_reward/mean': 0.0625, 'rewards/length_reward/std': 0.0, 'reward': 0.0625, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'sampling/sampling_logp_difference/mean': 0.0010602406691759825, 'sampling/sampling_logp_difference
/max': 0.003823213279247284, 'sampling/importance_sampling_ratio/min': 0.9938066601753235, 'sampling/importance_sampling_ratio/mean': 0.9938066601753235, 'sampling/importance_sampling_ratio/max': 0.9938066601753235, 'entropy': 0.0478515625, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_
min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.1684445310384035, 'epoch': 6.0}
{'loss': 0.0, 'grad_norm': 0.0, 'learning_rate': 4.000000000000001e-06, 'num_tokens': 2195.0, 'completions/mean_length': 6.0, 'completions/min_length': 6.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 6.0, 'completions/min_ter
minated_length': 6.0, 'completions/max_terminated_length': 6.0, 'rewards/length_reward/mean': 0.0625, 'rewards/length_reward/std': 0.0, 'reward': 0.0625, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'sampling/sampling_logp_difference/mean': 0.004386706743389368, 'sampling/sampling_
logp_difference/max': 0.021789629012346268, 'sampling/importance_sampling_ratio/min': 0.9740239977836609, 'sampling/importance_sampling_ratio/mean': 0.9740239977836609, 'sampling/importance_sampling_ratio/max': 0.9740239977836609, 'entropy': 0.0582682304084301, 'clip_ratio/low_mean':
0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.182079752907157, 'epoch': 7.0}
{'loss': 0.0, 'grad_norm': 0.0, 'learning_rate': 3e-06, 'num_tokens': 2427.0, 'completions/mean_length': 6.0, 'completions/min_length': 6.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 6.0, 'completions/min_terminated_length':
 6.0, 'completions/max_terminated_length': 6.0, 'rewards/length_reward/mean': 0.0625, 'rewards/length_reward/std': 0.0, 'reward': 0.0625, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'sampling/sampling_logp_difference/mean': 0.0014839638024568558, 'sampling/sampling_logp_difference
/max': 0.005859885364770889, 'sampling/importance_sampling_ratio/min': 0.9913414716720581, 'sampling/importance_sampling_ratio/mean': 0.9913414716720581, 'sampling/importance_sampling_ratio/max': 0.9913414716720581, 'entropy': 0.0576171875, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_
min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.187234789133072, 'epoch': 8.0}
{'loss': 0.0, 'grad_norm': 0.0, 'learning_rate': 2.0000000000000003e-06, 'num_tokens': 2659.0, 'completions/mean_length': 6.0, 'completions/min_length': 6.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 6.0, 'completions/min_te
rminated_length': 6.0, 'completions/max_terminated_length': 6.0, 'rewards/length_reward/mean': 0.0625, 'rewards/length_reward/std': 0.0, 'reward': 0.0625, 'reward_std': 0.0, 'frac_reward_zero_std': 1.0, 'sampling/sampling_logp_difference/mean': 0.0009608844993636012, 'sampling/samplin
g_logp_difference/max': 0.002919524908065796, 'sampling/importance_sampling_ratio/min': 1.000125765800476, 'sampling/importance_sampling_ratio/mean': 1.000125765800476, 'sampling/importance_sampling_ratio/max': 1.000125765800476, 'entropy': 0.0768229141831398, 'clip_ratio/low_mean': 0
.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.1681837420910597, 'epoch': 9.0}
{'loss': 0.091, 'grad_norm': 133.67459106445312, 'learning_rate': 1.0000000000000002e-06, 'num_tokens': 2888.0, 'completions/mean_length': 5.625, 'completions/min_length': 3.0, 'completions/max_length': 6.0, 'completions/clipped_ratio': 0.0, 'completions/mean_terminated_length': 5.625
, 'completions/min_terminated_length': 3.0, 'completions/max_terminated_length': 6.0, 'rewards/length_reward/mean': 0.0725446417927742, 'rewards/length_reward/std': 0.028410542756319046, 'reward': 0.0725446417927742, 'reward_std': 0.028410542756319046, 'frac_reward_zero_std': 0.0, 'sa
mpling/sampling_logp_difference/mean': 0.017341550439596176, 'sampling/sampling_logp_difference/max': 0.2619295120239258, 'sampling/importance_sampling_ratio/min': 0.9293384552001953, 'sampling/importance_sampling_ratio/mean': 0.9763498306274414, 'sampling/importance_sampling_ratio/ma
x': 1.305429220199585, 'entropy': 0.10833333432674408, 'clip_ratio/low_mean': 0.0, 'clip_ratio/low_min': 0.0, 'clip_ratio/high_mean': 0.0, 'clip_ratio/high_max': 0.0, 'clip_ratio/region_mean': 0.0, 'step_time': 1.1852372493594885, 'epoch': 10.0}
{'train_runtime': 13.7076, 'train_samples_per_second': 5.836, 'train_steps_per_second': 0.73, 'train_loss': 0.10115763023495675, 'epoch': 10.0}

With this patch, both vllm_enable_sleep_mode=False and vllm_enable_sleep_mode=True result in the same completion length and reward curves.

[qgallouedec]

Nice! Thanks!