Repository navigation
Qwen3-vl-embedding-2B训练爆显存 #9885
Description
Checklist / 检查清单
- I have searched existing issues, and this is a new question or discussion topic. / 我已经搜索过现有的 issues，确认这是一个新的问题与讨论。
Question Description / 问题描述
8*A100
swift4.1+deepspeed0.18.9可以正常训练
swift4.4.2+deepspeed0.18.9/0.18.0/0.16.0都会爆显存
训练脚本 examples/train/embedding/qwen3/qwen3_vl_emb.sh 中都改成全参--tuner_type full、长度max_length 1024、gradient accumulation 4以及在swift4.4.2上加上了--mrl_dims '{"1024": 1.0}'
swift4.4.2显存第一个step单卡显存占用直接到50G，然后一直增加直到爆显存。
Activity
z0o0ey commentedon Aug 13, 2026
我按照你的描述测试了如下命令：
CUDA_VISIBLE_DEVICES=0,1
INFONCE_TEMPERATURE=0.1
NPROC_PER_NODE=2
swift sft
--model Qwen/Qwen3-Embedding-4B
--task_type embedding
--tuner_type full
--lora_rank 8
--lora_alpha 32
--learning_rate 5e-5
--target_modules all-linear
--dataset sentence-transformers/stsb:positive
--attn_impl flash_attn
--padding_free true
--torch_dtype bfloat16
--load_from_cache_file true
--split_dataset_ratio 0.02
--eval_strategy steps
--output_dir output
--save_steps 50
--eval_steps 50
--save_total_limit 2
--logging_steps 5
--num_train_epochs 5
--max_length 1024
--per_device_train_batch_size 8
--per_device_eval_batch_size 8
--gradient_accumulation_steps 4
--dataloader_num_workers 4
--dataset_num_proc 4
--warmup_ratio 0.05
--loss_type infonce
--dataloader_drop_last true
--mrl_dims '{"1024": 1.0}'
--deepspeed zero2
显存占用一直稳定在'memory(GiB)': '48.75'，并没有出现一直增加的情况
可以提供给我更多的信息吗
ma-biao commentedon Aug 13, 2026
deepspeed是zero1，我把其他所有版本包对齐之后，显存会比较稳定，但是占用差距较大。在swift4.1.0版本下per_device_train_batch_size 可以开到16，显存占用50G左右，swift4.4.2版本下只能开到8，开到16会爆显存。
z0o0ey commentedon Aug 14, 2026
老师你在4.1.0版本里面把--mrl_dims {"1024": 1.0}参数去掉了对吗
我去掉这个参数之后--deepspeed zero1、--per_device_train_batch_size 8的情况下显存占用稳定在'memory(GiB)': '50.21'，和4.4.2是一致的
z0o0ey commentedon Aug 14, 2026
{'loss': '0.01419', 'grad_norm': '0.5742', 'learning_rate': '3.457e-05', 'epoch': '2.037', 'global_step/max_steps': '55/135', 'elapsed_time': '1m 58s', 'remaining_time': '2m 51s', 'memory(GiB)': '50.21','train_speed(s/it)': '2.143'}
{'loss': '0.005914', 'grad_norm': '0.01436', 'learning_rate': '3.167e-05', 'epoch': '2.224', 'global_step/max_steps': '60/135', 'elapsed_time': '2m 4s', 'remaining_time': '2m 35s', 'memory(GiB)': '50.21', 'train_speed(s/it)': '2.067'}
{'loss': '0.004235', 'grad_norm': '0.102', 'learning_rate': '2.867e-05', 'epoch': '2.411', 'global_step/max_steps': '65/135', 'elapsed_time': '2m 11s', 'remaining_time': '2m 21s', 'memory(GiB)': '50.21','train_speed(s/it)': '2.01'}
{'loss': '0.005672', 'grad_norm': '0.03346', 'learning_rate': '2.561e-05', 'epoch': '2.598', 'global_step/max_steps': '70/135', 'elapsed_time': '2m 16s', 'remaining_time': '2m 6s', 'memory(GiB)': '50.21', 'train_speed(s/it)': '1.939'}
{'loss': '0.004322', 'grad_norm': '0.06222', 'learning_rate': '2.255e-05', 'epoch': '2.785', 'global_step/max_steps': '75/135', 'elapsed_time': '2m 21s', 'remaining_time': '1m 53s', 'memory(GiB)': '50.21', 'train_speed(s/it)': '1.877'}
{'loss': '0.003611', 'grad_norm': '0.01675', 'learning_rate': '1.952e-05', 'epoch': '2.972', 'global_step/max_steps': '80/135', 'elapsed_time': '2m 26s', 'remaining_time': '1m 40s', 'memory(GiB)': '50.21', 'train_speed(s/it)': '1.822'}
{'loss': '0.00629', 'grad_norm': '0.01475', 'learning_rate': '1.658e-05', 'epoch': '3.15', 'global_step/max_steps': '85/135', 'elapsed_time': '2m 31s', 'remaining_time': '1m 29s', 'memory(GiB)': '50.21','train_speed(s/it)': '1.773'}
{'loss': '0.005863', 'grad_norm': '0.02261', 'learning_rate': '1.376e-05', 'epoch': '3.336', 'global_step/max_steps': '90/135', 'elapsed_time': '2m 36s', 'remaining_time': '1m 18s', 'memory(GiB)': '50.21', 'train_speed(s/it)': '1.731'}
{'loss': '0.002068', 'grad_norm': '0.01391', 'learning_rate': '1.111e-05', 'epoch': '3.523', 'global_step/max_steps': '95/135', 'elapsed_time': '2m 41s', 'remaining_time': '1m 8s', 'memory(GiB)': '50.21', 'train_speed(s/it)': '1.691'}
{'loss': '0.006392', 'grad_norm': '0.01193', 'learning_rate': '8.67e-06', 'epoch': '3.71', 'global_step/max_steps': '100/135', 'elapsed_time': '2m 46s', 'remaining_time': '58s', 'memory(GiB)': '50.21', 'train_speed(s/it)': '1.657'}
持续观察也无增涨
ma-biao commentedon Aug 17, 2026
我重新对齐了环境，把gcc, cuda, numpy, torch等环境包版本对齐后，在swift4.1.0版本下显存稳定无增长。相同训练指令和配置下，swift4.4.0（mrl应该是这个版本才开始支持的）环境下启动会爆显存：
z0o0ey commentedon Aug 19, 2026
将4.4.0版本命令中的--mrl_dims去掉试一下呢老师
ma-biao commentedon Aug 19, 2026
显存一样的，会直接升到80G然后退出。而且升级到4.4.0就是为了mrl训练。
z0o0ey commentedon Aug 19, 2026
在4.4.0的训练命令里显式定义一下--vit_gradient_checkpointing true
ma-biao commentedon Aug 20, 2026
添加这个参数之后显存占用稳定了
z0o0ey commentedon Aug 20, 2026
可以尝试一下main分支
https://github.com/modelscope/ms-swift/pull/9945这个pr修复了这个bug
使用main分支的代码时不需要添加--vit_gradient_checkpointing true