alibaba / alibaba/ROLL

Torch 2.8.0 + vLLM 0.10.2 开启 use_kl_loss 后训练qwen2.5-vl-3B-agentic崩溃,Torch 2.6.0 + vLLM 0.8.4 正常

Open
#398 17 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
3.4k
Forks
312
Avg merge
1h 2m
Merged PRs (30d)
2

Description

相同配置,在设置use_kl_loss为true后,使用torch2.8.0 + vLLM0.10.2镜像训练到一定步数后会崩溃,但是使用torch2.6.0 + vLLM0.8.4镜像则可以正常训练,训练曲线对比如下:

Image

yaml配置如下:
```
defaults:
- ../config/vl_traj_envs@_here_
- ../config/deepspeed_zero@_here_
- ../config/deepspeed_zero2@_here_
- ../config/deepspeed_zero3@_here_
- ../config/deepspeed_zero3_cpuoffload@_here_

hydra:
run:
dir: .
output_subdir: null

exp_name: "agentic_pipeline_3b"
seed: 42
logging_dir: ./output/logs
output_dir: ./output
render_save_dir: ./output/render
system_envs:
USE_MODELSCOPE: '1'

track_with: wandb
tracker_kwargs:
api_key: ""
project: ""
name: ${exp_name}_sokoban
notes: "agentic_pipeline"
tags:
- agentic
- roll
- baseline

# track_with: tensorboard
# tracker_kwargs:
# log_dir: /data/oss_bucket_0/yali/llm/tensorboard/roll_exp/agentic_sokoban

checkpoint_config:
type: file_system
output_dir: /data/cpfs_0/rl_examples/models/${exp_name}

num_gpus_per_node: 8

max_steps: 1024
save_steps: 10000
logging_steps: 1
eval_steps: 10
resume_from_checkpoint: false

rollout_batch_size: 1024
val_batch_size: 1024
sequence_length: 8192

advantage_clip: 0.2
ppo_epochs: 1
adv_estimator: "grpo"
#pg_clip: 0.1
#dual_clip_loss: True
init_kl_coef: 0.0
whiten_advantages: true
entropy_loss_coef: 0
max_grad_norm: 1.0
enable_old_logprobs_recompute: true
use_kl_loss: true
kl_loss_coef: 0.01

pretrain: qwen/Qwen2.5-VL-3B-Instruct
reward_pretrain: qwen/Qwen2.5-VL-3B-Instruct
actor_train:
model_args:
attn_implementation: fa2
disable_gradient_checkpointing: false
dtype: bf16
model_type: ~
training_args:
learning_rate: 1.0e-6
weight_decay: 0
per_device_train_batch_size: 2
gradient_accumulation_steps: 8
warmup_steps: 10
lr_scheduler_type: cosine
data_args:
template: qwen2_5
strategy_args:
strategy_name: deepspeed_train
strategy_config: ${deepspeed_zero3}
# strategy_name: megatron_train
# strategy_config:
# tensor_model_parallel_size: 1
# pipeline_model_parallel_size: 1
# expert_model_parallel_size: 1
# use_distributed_optimizer: true
# recompute_granularity: full
device_mapping: list(range(0,64))
infer_batch_size: 2

actor_infer:
model_args:
disable_gradient_checkpointing: true
dtype: bf16
generating_args:
max_new_tokens: 128 # single-turn response length
top_p: 0.99
top_k: 100
num_beams: 1
temperature: 0.99
num_return_sequences: 1
data_args:
template: qwen2_5
strategy_args:
strategy_name: vllm
strategy_config:
gpu_memory_utilization: 0.8
block_size: 16
load_format: auto
enable_prefix_caching: false
limit_mm_per_prompt:
image: ${max_actions_per_traj}
device_mapping: list(range(0,64))

reference:
model_args:
attn_implementation: fa2
disable_gradient_checkpointing: true
dtype: bf16
model_type: ~
data_args:
template: qwen2_5
strategy_args:
strategy_name: deepspeed_infer
strategy_config: ${deepspeed_zero3}
device_mapping: list(range(0,64))
infer_batch_size: 2

reward_normalization:
grouping: traj_group_id # 可以tags(env_type)/traj_group_id(group)/batch(rollout_batch)... group_by计算reward/adv
method: mean_std # asym_clip / identity / mean_std

train_env_manager:
format_penalty: -0.15 # sokoban env penalty_for_step=-0.1
max_env_num_per_worker: 16
num_env_groups: 128
# under the same group, the env config and env seed are ensured to be equal
group_size: 8
tags: [SimpleSokoban]
num_groups_partition: [128] # If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation

val_env_manager:
max_env_num_per_worker: 32
num_env_groups: 1024
group_size: 1 # should be set to 1 because val temperature is set to 0 and same prompt leads to same output
tags: [SimpleSokoban, LargerSokoban, SokobanDifferentGridVocab, FrozenLake]
num_groups_partition: [256, 256, 256, 256] # TODO: If not set, all env names divide nums equally. Under the same group, the env config and env seed (prompt) are equal in each generation

# Here, you can override variables defined in the imported envs. max_tokens_per_step: 128 in custom_env.SimpleSokoban, here replaced by 64
max_tokens_per_step: 64

custom_envs:
SimpleSokoban:
${custom_env.SimpleSokoban}
LargerSokoban:
${custom_env.LargerSokoban}
SokobanDifferentGridVocab:
${custom_env.SokobanDifferentGridVocab}
FrozenLake:
${custom_env.FrozenLake}
FrozenLakeThink:
${custom_env.FrozenLakeThink}
```

dockerfile如下:
```
FROM roll-registry-vpc.cn-hangzhou.cr.aliyuncs.com/roll/pytorch:nvcr-25.06-py3-torch280-vllm0102

RUN apt-get update && apt-get install --fix-missing -y -o Acquire::Retries=3 iproute2 tmux zip
COPY requirements_torch280_vllm.txt ./requirements_torch280_vllm.txt
COPY requirements_common.txt ./requirements_common.txt
COPY requirements_vision.txt ./requirements_vision.txt
COPY mcore_adapter ./mcore_adapter
RUN pip install -r requirements_torch280_vllm.txt
```

辛苦解答,感谢🙏

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reproducing the supplied YAML configuration with use_kl_loss enabled in the Torch 2.8.0/vLLM 0.10.2 Docker image, then compare it with the working Torch 2.6.0/vLLM 0.8.4 setup. Collect the crash traceback and inspect the training path reached by the use_kl_loss setting; done means identifying the version-related failure and documenting or fixing it with a regression check.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.