lora don't work! OOM

[zhangguoxin1]

Reproduction

python examples/scripts/nash_md.py \
    --model_name_or_path ../../autodl-tmp/models/qwen2_7B  \
    --reward_model ../../autodl-tmp/models/qwen2_reward_0_5B \
    --dataset_name ../datas/hh-rlhf-helpful-base \
    --max_length 1024 \
    --learning_rate 1.0e-5 \
    --output_dir ../../autodl-tmp/saves/qwen2_7B_nash_KLfree \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 8 \
    --num_train_epochs 1 \
    --max_new_tokens 128 \
    --warmup_ratio 0.1 \
    --use_peft True \
    --lora_task_type "CAUSAL_LM" \
    --lora_dropout 0.1 \
    --lora_r 8 \
    --lora_alpha 16 \
    --lora_target_modules "q_proj,k_proj,v_proj" \
    --mixture_coef 0 \
    --run_name nash \
    --temperature 0.7 \
    --max_steps 375 \
    --report_to wandb

System Info

• Platform: Linux-5.15.0-78-generic-x86_64-with-glibc2.35
• Python version: 3.10.8
• PyTorch version: 2.1.2+cu118
• CUDA device(s): NVIDIA GeForce RTX 4090
• Transformers version: 4.46.1
• Accelerate version: 1.0.1
• Accelerate config: not found
• Datasets version: 3.1.0
• HF Hub version: 0.28.1
• TRL version: 0.15.0.dev0
• bitsandbytes version: not installed
• DeepSpeed version: not installed
• Diffusers version: not installed
• Liger-Kernel version: not installed
• LLM-Blender version: 0.0.2
• OpenAI version: not installed
• PEFT version: 0.12.0

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

[JohnConnor123]

try using deepspeed zero stage 3 and trim the length of the input text or the length of the generated text.

[MilchstraB]

I also encountered this problem. I solved it by downgrading the trl version to 0.13.0.