modelscope / modelscope/ms-swift

swift=3.10.1 results unexpected OOM during on-policy GKD

Open
#6,639 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
15.7k
Forks
1.7k
Avg merge
1d 16h
Merged PRs (30d)
136

Description

Describe the bug
What the bug is, and how to reproduce, better with screenshots(描述bug以及复现过程,最好有截图)

CUDA_VISIBLE_DEVICES=6 \
swift rollout \
    --model /mnt/qwen/Qwen3-VL-8B-Instruct \
    --vllm_max_model_len 24192


NPROC_PER_NODE=7 \
PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True' \
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6 \
swift rlhf \
    --rlhf_type gkd \
    --model /mnt/qwen/Qwen3-VL-8B-Instruct \
    --teacher_model /mnt/qwen/Qwen3-VL-235B-A22B-Instruct \
    --train_type full \
    --dataset /mnt/images/dataset.jsonl \
    --seq_kd false \
    --lmbda 1 \
    --beta 1 \
    --torch_dtype bfloat16 \
    --num_train_epochs 1 \
    --per_device_train_batch_size 1 \
    --learning_rate 1e-5 \
    --gradient_accumulation_steps 8 \
    --save_steps 500 \
    --save_total_limit 2 \
    --logging_steps 1 \
    --max_length 6000 \
    --max_completion_length 3000 \
    --warmup_ratio 0.05 \
    --save_only_model true \
    --dataloader_num_workers 64 \
    --dataset_num_proc 4 \
    --deepspeed zero2 \
    --teacher_deepspeed zero3 \
    --attn_impl flash_attn \
    --use_vllm true \
    --vllm_mode server \
    --vllm_server_host 127.0.0.1 \
    --vllm_server_port 8000 \
    --output_dir /home/output/ \
Image The GPU memory would be occupied fully once the checkpoint start to load and OOM will occur when hit 31% progress everytime。

Your hardware and system info
Write your system info like CUDA version/system/GPU/torch version here(在这里给出硬件信息和系统信息,如CUDA版本,系统,GPU型号和torch版本等)

Image

Additional context
Add any other context about the problem here(在这里补充其他信息)
The issue could be resolved after downgrading swift version to 3.10.0

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by reproducing the supplied swift rollout and swift rlhf commands with the stated multi-GPU, GKD, vLLM server, and DeepSpeed settings. Compare behavior between ms-swift 3.10.1 and 3.10.0, then trace the checkpoint-loading and training entry points to identify what changes memory usage; done means the regression is explained and the failing configuration no longer OOMs.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.