qwen2 72B PPO OOM

[lixsh6]

跑了qwen2 72B的PPO出现的OOM异常

环境：4机32卡，80G显存，2T内存。理论上应该不会oom才对。 运行代码：

ray job submit --address="http://127.0.0.1:8265" \
    --runtime-env-json='{"working_dir": "mycode/OpenRLHF-new","excludes": ["*.log","*.wandb","*.pt"] }' \
    --no-wait \
    -- python3 examples/train_ppo_ray.py \
    --ref_num_nodes 1 \
    --ref_num_gpus_per_node 4 \
    --reward_num_nodes 1 \
    --reward_num_gpus_per_node 4 \
    --critic_num_nodes 1 \
    --critic_num_gpus_per_node 8 \
    --actor_num_nodes 1 \
    --actor_num_gpus_per_node 8 \
    --vllm_num_engines 1 \
    --vllm_tensor_parallel_size 8 \
    --pretrain $sft_model_path \
    --reward_pretrain $reward_model_path \
    --save_path $save_path \
    --micro_train_batch_size 1 \
    --train_batch_size 128 \
    --micro_rollout_batch_size 2 \
    --rollout_batch_size 1024 \
    --max_epochs 1 \
    --prompt_max_len 1024 \
    --generate_max_len 1024 \
    --zero_stage 3 \
    --bf16 \
    --actor_learning_rate 5e-7 \
    --critic_learning_rate 9e-6 \
    --init_kl_coef 0.01 \
    --pretrain_data $ptx_data \
    --prompt_data $prompt_data \
    --max_samples 80000 \
    --normalize_reward \
    --actor_init_on_gpu \
    --adam_offload \
    --gradient_checkpointing \
    --use_wandb $WANDB_TOKENS \
    --wandb_run_name $EXP_NAME

[hijkzzz]

奇怪的问题，可能是没有flash_attn导致的？请问具体是哪个环节OOM了？

补充建议： 可以参考这个脚本合并reward critic 以及 actor ref 速度会更快：https://github.com/OpenLLMAI/OpenRLHF/blob/main/examples/scripts/train_ppo_llama3_ray_colocate.sh vllm engine=4 tp_size = 4 actor gpus=8 critic gpus=8

[catqaq]

贴一下环境信息吧？比如ds版本、flash_attn

[lixsh6]

贴一下环境信息吧？比如ds版本、flash_attn

deepspeed 0.12.3+00df0c19 flash-attn 2.3.3 torch 2.2.1
torchaudio 2.2.1
torchvision 0.17.1 vllm 0.4.1
vllm-nccl-cu12 2.18.1.0.4.0 cuda版本12.2

[lixsh6]

奇怪的问题，可能是没有flash_attn导致的？请问具体是哪个环节OOM了？

补充建议： 可以参考这个脚本合并reward critic 以及 actor ref 速度会更快：https://github.com/OpenLLMAI/OpenRLHF/blob/main/examples/scripts/train_ppo_llama3_ray_colocate.sh vllm engine=4 tp_size = 4 actor gpus=8 critic gpus=8

试了还是一样的问题，你们又遇到类似的问题吗？感觉像是ray的问题。我的ray是2.20版本

[hijkzzz]

奇怪的问题，可能是没有flash_attn导致的？请问具体是哪个环节OOM了？ 补充建议： 可以参考这个脚本合并reward critic 以及 actor ref 速度会更快：https://github.com/OpenLLMAI/OpenRLHF/blob/main/examples/scripts/train_ppo_llama3_ray_colocate.sh vllm engine=4 tp_size = 4 actor gpus=8 critic gpus=8

试了还是一样的问题，你们又遇到类似的问题吗？感觉像是ray的问题。我的ray是2.20版本

CPU 内存不够 不足以放下 adam optimizer? docker 环境推荐试试 https://github.com/OpenLLMAI/OpenRLHF/tree/main/dockerfile