Qwen3-vl-embedding-2B训练爆显存

[ma-biao]

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，然后一直增加直到爆显存。

[z0o0ey]

我按照你的描述测试了如下命令：
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]

我按照你的描述测试了如下命令： 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'，并没有出现一直增加的情况 可以提供给我更多的信息吗

deepspeed是zero1，我把其他所有版本包对齐之后，显存会比较稳定，但是占用差距较大。在swift4.1.0版本下per_device_train_batch_size 可以开到16，显存占用50G左右，swift4.4.2版本下只能开到8，开到16会爆显存。

[z0o0ey]

我按照你的描述测试了如下命令： 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'，并没有出现一直增加的情况 可以提供给我更多的信息吗

deepspeed是zero1，我把其他所有版本包对齐之后，显存会比较稳定，但是占用差距较大。在swift4.1.0版本下per_device_train_batch_size 可以开到16，显存占用50G左右，swift4.4.2版本下只能开到8，开到16会爆显存。

老师你在4.1.0版本里面把--mrl_dims {"1024": 1.0}参数去掉了对吗
我去掉这个参数之后--deepspeed zero1、--per_device_train_batch_size 8的情况下显存占用稳定在'memory(GiB)': '50.21'，和4.4.2是一致的

[z0o0ey]

{'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'}
持续观察也无增涨