Introduce TP in coloc mode

[toslali-ibm]

requirements

• git clone https://github.com/huggingface/open-r1.git

• pip install -e open-r1

• pip install vllm==0.7.3

• pip uninstall -y trl

• git clone -b tpcoloc https://github.com/toslali-ibm/trl.git

• pip install -e trl

• cd open-r1

• run experiment (Tp=4 for 4 GPUs for 32B model Qwen/Qwen2.5-14B-Instruct) for vllm coloc tp

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 accelerate launch \
    --config_file recipes/accelerate_configs/zero3.yaml \
    --num_processes=8 \
    src/open_r1/grpo.py \
    --config config.yaml

See the Config.yaml

# Model arguments
model_name_or_path: Qwen/Qwen2.5-14B-Instruct
model_revision: main
torch_dtype: bfloat16
attn_implementation: flash_attention_2

# Data training arguments
dataset_name: open-r1/OpenR1-Math-220k
system_prompt: "You are a helpful AI Assistant that provides well-reasoned and detailed responses. You first think about the reasoning process as an internal monologue and then provide the user with the answer. Respond in the following format: <think>\n...\n</think>\n<answer>\n...\n</answer>"
dataset_prompt_column: "problem"

# GRPO trainer config
bf16: true
use_vllm: true
vllm_tp: true
vllm_gpu_memory_utilization: 0.35
vllm_enable_prefix_caching: false
vllm_max_model_len: 1536

do_eval: false
eval_strategy: "no"
use_vllm: true
do_eval: false
gradient_accumulation_steps: 4
gradient_checkpointing: true
gradient_checkpointing_kwargs:
  use_reentrant: false

learning_rate: 2.0e-05
log_completions: false
log_level: info
logging_first_step: true
logging_steps: 5
logging_strategy: steps
lr_scheduler_type: cosine
max_grad_norm: 0.2
max_prompt_length: 512
max_completion_length: 1024
max_steps: 10
num_generations: 16
num_train_epochs: 1

overwrite_output_dir: true
per_device_train_batch_size: 16

reward_funcs:
- accuracy
- format
reward_weights:
- 1.0
- 1.0
eval_strategy: "no"
save_strategy: "no"
report_to: none

seed: 42
temperature: 0.7
warmup_ratio: 0.1

[toslali-ibm]

Once you load weights - you have to stay awake -

correct Pattern:

• wake up - no sleep - so generate - then go back to sleep

[liguohao96]

DO NOT sleep with level 2, it will drop the weight.
I tried your tpcol patch and got total random generation results.
After debugging for days, I find that LLM sleep level 2 will drop the weight. This makes training with gradient accumulation go wrong.

vllm's doc on sleep

[toslali-ibm]

DO NOT sleep with level 2, it will drop the weight. I tried your tpcol patch and got total random generation results. After debugging for days, I find that LLM sleep level 2 will drop the weight. This makes training with gradient accumulation go wrong.

vllm's doc on sleep

Hello @liguohao96,
This is great—thank you for testing and debugging the feature and for identifying this issue.

You're absolutely right: during gradient accumulation, update_named_param is not called, so the model isn’t explicitly loaded—we go straight into generate, which wakes up the model. However, the weights were lost after a previous generate call due to level = 2 sleeping.

1. Using Level 1 sleeping stores weights on the CPU, which should help preserve them. That said, we suspect this might slow down training due to the overhead of transferring weights from CPU to GPU (CC @fabianlim).

2. An alternative approach would be to ensure we change the flow slightly for vLLMColocClient, and call update_named_param before each generate, which would preserve the weights. This would slightly change the flow but could potentially be faster than CPU/GPU communication..

3. Third approach, which might be the easiest, is to only calling .sleep(level=2) at the last mini-step within generate.

if (step + 1) % grad_accum_steps == 0:
    self.llm.sleep(level=2)

But this does not release memory for training, during grad accumulations steps...

4. Forth approach can be mixed of sleep level 1 and level 2. We do sleep level 2 before _move_model_to_vllm, because model will be updated, no need to remember weights. We do sleep level 1 during grad accumulation steps to remember weights.

[toslali-ibm]

DO NOT sleep with level 2, it will drop the weight. I tried your tpcol patch and got total random generation results. After debugging for days, I find that LLM sleep level 2 will drop the weight. This makes training with gradient accumulation go wrong.
vllm's doc on sleep

Hello @liguohao96, This is great—thank you for testing and debugging the feature and for identifying this issue.

You're absolutely right: during gradient accumulation, update_named_param is not called, so the model isn’t explicitly loaded—we go straight into generate, which wakes up the model. However, the weights were lost after a previous generate call due to level = 2 sleeping.

1. Using Level 1 sleeping stores weights on the CPU, which should help preserve them. That said, we suspect this might slow down training due to the overhead of transferring weights from CPU to GPU (CC @fabianlim).
2. An alternative approach would be to ensure we change the flow slightly for vLLMColocClient, and call update_named_param before each generate, which would preserve the weights. This would slightly change the flow but could potentially be faster than CPU/GPU communication..
3. Third approach, which might be the easiest, is to only calling .sleep(level=2) at the last mini-step within generate.

if (step + 1) % grad_accum_steps == 0:
    self.llm.sleep(level=2)

But this does not release memory for training, during grad accumulations steps...

4. Forth approach can be mixed of sleep level 1 and level 2. We do sleep level 2 before _move_model_to_vllm, because model will be updated, no need to remember weights. We do sleep level 1 during grad accumulation steps to remember weights.

I tried (1) and got 211.40s/it

I tried (2), but got an error while trying to update model from coloc client as follows:

 [rank6]: Traceback (most recent call last): [rank6]:   File "/workspace/open-r1/src/open_r1/grpo.py", line 187, in <module> [rank6]:     main(script_args, training_args, model_args) [rank6]:   File "/workspace/open-r1/src/open_r1/grpo.py", line 141, in main [rank6]:     train_result = trainer.train(resume_from_checkpoint=checkpoint) [rank6]:   File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 2245, in train [rank6]:     return inner_training_loop( [rank6]:   File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 2556, in _inner_training_loop [rank6]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch) [rank6]:   File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 3712, in training_step [rank6]:     inputs = self._prepare_inputs(inputs) [rank6]:   File "/workspace/trl/trl/extras/profiling.py", line 87, in wrapper [rank6]:     return func(self, *args, **kwargs) [rank6]:   File "/workspace/trl/trl/trainer/grpo_trainer.py", line 666, in _prepare_inputs [rank6]:     inputs = self._generate_and_score_completions(inputs) [rank6]:   File "/workspace/trl/trl/trainer/grpo_trainer.py", line 700, in _generate_and_score_completions [rank6]:     completion_ids = self.vllm_client.generate( [rank6]:   File "/workspace/trl/trl/extras/vllm_client.py", line 475, in generate [rank6]:     self.load_model_during_grad_accumulation() [rank6]:   File "/workspace/trl/trl/extras/vllm_client.py", line 431, in load_model_during_grad_accumulation [rank6]:     with gather_if_zero3([param]): [rank6]:   File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/zero/partition_parameters.py", line 2241, in __exit__ [rank6]:     self.params[0].partition(param_list=self.params, has_been_updated=False) [rank6]:   File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/zero/partition_parameters.py", line 1386, in partition [rank6]:     self._partition(param_list, has_been_updated=has_been_updated) [rank6]:   File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/zero/partition_parameters.py", line 1535, in _partition [rank6]:     self._partition_param(param, has_been_updated=has_been_updated) [rank6]:   File "/usr/local/lib/python3.10/dist-packages/deepspeed/utils/nvtx.py", line 18, in wrapped_fn [rank6]:     ret_val = func(*args, **kwargs) [rank6]:   File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/zero/partition_parameters.py", line 1568, in _partition_param [rank6]:     free_param(param) [rank6]:   File "/usr/local/lib/python3.10/dist-packages/deepspeed/utils/nvtx.py", line 18, in wrapped_fn [rank6]:     ret_val = func(*args, **kwargs) [rank6]:   File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/zero/partition_parameters.py", line 284, in free_param [rank6]:     assert not param.ds_active_sub_modules, param.ds_summary() [rank6]: AssertionError: {'id': 0, 'status': 'AVAILABLE', 'numel': 778567680, 'ds_numel': 778567680, 'shape': (152064, 5120), 'ds_shape': (152064, 5120), 'requires_grad': True, 'grad_shape': None, 'persist': False, 'active_sub_modules': {632}, 'ds_tensor.shape': torch.Size([97320960])}

I tried (3). This seemed easy, however, as we are not sleeping during ACC steps, this is problematic for OOMs. The key idea of using sleep is to free up memory for training process (especially important for larger models)

I tried (4) - and got 181.70s/it. This seems to me as the best option.

What do you think @liguohao96 and @fabianlim ?

[liguohao96]

I use (1) as my current solution. My first thought was sleep level 1 is good enough.
but as your results (4), sleep with level 2 before _move_model_to_vllm is significantly faster that solution (1). but why?

[toslali-ibm]

I use (1) as my current solution. My first thought was sleep level 1 is good enough.
but as your results (4), sleep with level 2 before _move_model_to_vllm is significantly faster that solution (1). but why?

I believe it's faster because sleep level 2 avoids GPU/CPU transfers, which we suspect are a source of slowdown.

Nevertheless, to debug (4), I printed all prompts and generations. Even when I do sleep level 2 only before _move_model_to_vllm, because model will be updated, no need to remember weights—the generations start becoming random.. Because of this, I've also reverted to always using sleep level 1 (approach (1)).

[liguohao96]

Okay.

Another interesing find is that setting max_num_seqs to 1 gives the best speed with in my case.
I am not familiar with vllm, so I really don't know why this happens. maybe you could give a try?

My speculation is: generate with max_num_seqs=1 fits all kv-cache in VRAM, but vllm has to offload some kv-cache to CPU with larger max_num_seqs