modelscope / modelscope/ms-swift
gemma sft训练到一定steps出现显存剧烈增加直到oom
Open
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以及复现过程，最好有截图)
2卡A100命令：
nproc_per_node=${1:-2}
LR=${3:-4e-6}
DROPOUT=${4:-0.05}
RANK=${5:-32}
EPOCH=${6:-4}
PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True"
CUDA_VISIBLE_DEVICES=${2:-"0,1"}
NPROC_PER_NODE=$nproc_per_node \
swift sft \
--use_hf true \
--model /Models/OSS/google/gemma-3-12b-it \
--train_type lora \
--dataset /data/train/tickets_reply_4class_reviewed_more2_msdata_xml7w_allvali_abnummodified_dataset.jsonl \
--split_dataset_ratio 0.02 \
--torch dtype bfloat16 \
--num_train_epochs $EPOCH \
--per_device_train_batch_size 8 \
--per_device_eval_batch_size 1 \
--learning_rate $LR \
--lr_scheduler_type cosine_with_restarts \
--lr_scheduler_kwargs '{"num_cycles": 3}' \
--lora_rank $RANK \
--lora_alpha $LORA_ALPHA \
--lora_dropout $DROPOUT \
--target_modules all-linear \
--gradient_accumulation_steps 4 \
--eval_steps 50 \
--neptune_noise_alpha 5 \
--save_steps 100 \
--save_total_limit 10 \
--max_length 10240 \
--output_dir /trained_models/$VERSION \
--warmup_ratio 0.03 \
--dataloader_num_workers 32 \
--dataloader_persistent_workers true \
--load_from_cache_file true \
--attn_impl flash_attn \
--logging_steps 1 \
--use_liger_kernel true \
--gradient_checkpointing_kwargs '{"use_reentrant": false}'
每次验证阶段时出现显存剧烈增加，撑了几次爆了：
4卡A100命令类似，更改了per_device_train_batch_size 和gradient_accumulation_steps ，取消了eval（split_dataset_ratio 保持默认0值）：
nproc_per_node=${1:-4}
LR=${3:-4e-6}
DROPOUT=${4:-0.05}
RANK=${5:-32}
EPOCH=${6:-4}
LORA_ALPHA=$((2 * RANK))
VERSION=numcate-lora-r$RANK-d$DROPOUT-lr$LR-bs32-jsonl
VENV_PATH="/venv/ms_swift"
source $VENV_PATH/bin/activate
PYTORCH_CUDA_ALLOC_CONF='expandable_segments:True'
CUDA_VISIBLE_DEVICES=${2:-"0,1,2,3"} \
NPROC_PER_NODE=$nproc_per_node \
swift sft \
--use_hf true \
--model /Models/OSS/google/gemma-3-12b-it \
--train_type lora \
--dataset /data/train/tickets_reply_4class_reviewed_more2_msdata_xml7w_allvali_abnummodified_dataset.jsonl \
--torch_dtype bfloat16 \
--num_train_epochs $EPOCH \
--per_device_train_batch_size 16 \
--learning_rate $LR \
--lr_scheduler_type cosine_with_restarts \
--lr_scheduler_kwargs '{"num_cycles": 3}' \
--lora_rank $RANK \
--lora_alpha $LORA_ALPHA \
--lora_dropout $DROPOUT \
--target_modules all-linear \
--gradient_accumulation_steps 2 \
--eval_steps 50 \
--neftune_noise_alpha 5 \
--save_steps 100 \
--save_total_limit 10 \
--max_length 10240 \
--output_dir /trained_models/$VERSION \
--warmup_ratio 0.03 \
--dataloader_num_workers 32 \
--dataloader_persistent_workers true \
--load_from_cache_file true \
--attn_impl flash_attn \
--logging_steps 1 \
--use_liger_kernel true \
--gradient_checkpointing_kwargs '{"use_reentrant": false}'
总global steps 2180，训练到49 step显存时还是爆了
oom具体信息都类似为：
[rank1]: File "/home/jovyan/workspace-0/ms-swift/swift/cli/sft.py", line 10, in <module>
[rank1]: sft_main()
[rank1]: File "/home/jovyan/workspace-0/ms-swift/swift/llm/train/sft.py", line 331, in sft_main
[rank1]: return SwiftSft(args).main()
[rank1]: ^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/ms-swift/swift/llm/base.py", line 49, in main
[rank1]: result = self.run()
[rank1]: ^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/ms-swift/swift/llm/train/sft.py", line 187, in run
[rank1]: return self.train(trainer)
[rank1]: ^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/ms-swift/swift/llm/train/sft.py", line 235, in train
[rank1]: trainer.train(trainer.args.resume_from_checkpoint)
[rank1]: File "/home/jovyan/workspace-0/ms-swift/swift/trainers/mixin.py", line 676, in train
[rank1]: res = super().train(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/transformers/trainer.py", line 2238, in train
[rank1]: return inner_training_loop(
[rank1]: ^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/transformers/trainer.py", line 2582, in _inner_training_loop
[rank1]: tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/ms-swift/swift/trainers/trainers.py", line 413, in training_step
[rank1]: return super().training_step(model, inputs, *args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/transformers/trainer.py", line 3796, in training_step
[rank1]: loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/ms-swift/swift/trainers/trainers.py", line 327, in compute_loss
[rank1]: outputs = model(**inputs)
[rank1]: ^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/deepspeed/utils/nvtx.py", line 20, in wrapped_fn
[rank1]: ret_val = func(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/deepspeed/runtime/engine.py", line 2105, in forward
[rank1]: loss = self.module(*inputs, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1879, in _call_impl
[rank1]: return inner()
[rank1]: ^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1827, in inner
[rank1]: result = forward_call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/peft/peft_model.py", line 1845, in forward
[rank1]: return self.base_model(
[rank1]: ^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/peft/tuners/tuners_utils.py", line 216, in forward
[rank1]: return self.model.forward(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/transformers/models/gemma3/modeling_gemma3.py", line 1118, in forward
[rank1]: loss = loss_fct(flat_logits, flat_labels)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1773, in _wrapped_call_impl
[rank1]: return self._call_impl(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1784, in _call_impl
[rank1]: return forward_call(*args, **kwargs)
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/torch/nn/modules/loss.py", line 1310, in forward
[rank1]: return F.cross_entropy(
[rank1]: ^^^^^^^^^^^^^^^^
[rank1]: File "/home/jovyan/workspace-0/venv/ms_swift/lib/python3.12/site-packages/torch/nn/functional.py", line 3462, in cross_entropy
[rank1]: return torch._C._nn.cross_entropy_loss(
[rank1]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 8.42 GiB. GPU 1 has a total capacity of 79.10 GiB of which 4.40 GiB is free. Process 3639765 has 74.70 GiB memory in use. Of the allocated memory 66.46 GiB is allocated by PyTorch, and 7.41 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
Your hardware and system info
Write your system info like CUDA version/system/GPU/torch version here(在这里给出硬件信息和系统信息，如CUDA版本，系统，GPU型号和torch版本等)
CUDA12.8 PY3.12，ms-swift环境：
absl-py 2.3.1
accelerate 1.10.0
addict 2.4.0
aiofiles 23.2.1
aiohappyeyeballs 2.5.0
aiohttp 3.11.13
aiosignal 1.3.2
airportsdata 20250224
aliyun-python-sdk-core 2.16.0
aliyun-python-sdk-kms 2.16.5
annotated-types 0.7.0
anthropic 0.49.0
antlr4-python3-runtime 4.13.2
anyio 4.10.0
apex 0.1
apex 0.1
astor 0.8.1
asttokens 3.0.0
attrdict 2.0.1
attrs 25.1.0
autocommand 2.2.2
backports.tarfile 1.2.0
binpacking 1.5.2
blake3 1.0.8
blobfile 3.0.0
build 1.2.2.post1
cachetools 6.2.2
cbor2 5.7.1
certifi 2025.1.31
cffi 1.17.1
charset-normalizer 3.4.1
click 8.2.2
cloudpickle 3.1.1
compressed-tensors 0.11.0
contourpy 1.3.3
cpm-kernels 1.0.11
crcmod 1.7
cryptography 45.0.6
cuda-bindings 12.9.0
cuda-python 12.9.0
cupy-cuda12x 13.6.0
cycler 0.12.1
dacite 1.8.1
datasets 3.6.0
decorator 5.2.1
decord 0.6.0
deepspeed 0.17.2
depyf 0.19.0
dill 0.3.8
diskcache 5.6.3
distro 1.9.0
dnspython 2.8.0
einops 0.8.1
email-validator 2.3.0
executing 2.2.0
fastapi 0.116.1
fastapi-cli 0.0.16
fastapi-cloud-cli 0.3.1
fastrlock 0.8.3
ffmpy 0.5.0
filelock 3.17.0
flash_attn 2.7.4.post1
flash_attn 2.7.4.post1
flashinfer-python 0.2.11.post3
fonttools 4.59.2
frozendict 2.4.6
frozenlist 1.5.0
fsspec 2025.2.0
future 1.0.0
gguf 0.17.1
gradio 5.21.0
gradio_client 1.7.2
groovy 0.1.2
grpcio 1.74.0
h11 0.16.0
hf_transfer 0.1.9
hf-xet 1.1.7
hjson 3.1.0
httpcore 1.0.9
httptools 0.7.1
httpx 0.28.1
huggingface-hub 0.34.4
idna 3.10
importlib_metadata 8.7.0
inflect 7.3.1
interegular 0.3.3
ipython 9.4.0
ipython_pygments_lexers 1.1.1
jaraco.collections 5.1.0
jaraco.context 5.3.0
jaraco.functools 4.0.1
jaraco.text 3.12.1
jedi 0.19.2
jieba 0.42.1
Jinja2 3.1.6
jiter 0.10.0
jmespath 0.10.0
joblib 1.5.2
jsonschema 4.25.1
jsonschema-specifications 2025.4.1
kiwisolver 1.4.9
lark 1.2.2
latex2sympy2_extended 1.0.6
liger_kernel 0.6.4
llguidance 0.7.30
llvmlite 0.44.0
lm-format-enforcer 0.11.3
lxml 6.0.0
Markdown 3.8.2
markdown-it-py 4.0.0
MarkupSafe 2.1.5
math-verify 0.5.2
matplotlib 3.10.5
matplotlib-inline 0.1.7
mdurl 0.1.2
megatron-core 0.13.2
mistral_common 1.8.5
ml_dtypes 0.5.3
modelscope 1.29.1
more-itertools 10.3.0
mpmath 1.3.0
ms_swift 3.8.0.dev0 /ms-swift
msgpack 1.1.1
msgspec 0.19.0
multidict 6.1.0
multiprocess 0.70.16
nest-asyncio 1.6.0
networkx 3.4.2
ninja 1.13.0
nltk 3.9.1
numba 0.61.2
numpy 2.2.6
nvidia-cublas-cu12 12.8.4.1
nvidia-cuda-cupti-cu12 12.8.90
nvidia-cuda-nvcc-cu12 12.9.86
nvidia-cuda-nvrtc-cu12 12.8.93
nvidia-cuda-runtime-cu12 12.8.90
nvidia-cudnn-cu12 9.10.2.21
nvidia-cudnn-frontend 1.13.0
nvidia-cufft-cu12 11.3.3.83
nvidia-cufile-cu12 1.13.1.3
nvidia-curand-cu12 10.3.9.90
nvidia-cusolver-cu12 11.7.3.90
nvidia-cusparse-cu12 12.5.8.93
nvidia-cusparselt-cu12 0.7.1
nvidia-ml-py 12.575.51
nvidia-nccl-cu12 2.27.3
nvidia-nvjitlink-cu12 12.8.93
nvidia-nvtx-cu12 12.8.90
onnx 1.18.0
onnx-ir 0.1.7
onnxscript 0.3.1
openai 1.99.1
openai-harmony 0.0.4
opencv-python-headless 4.12.0.88
opt_einsum 3.4.0
orjson 3.11.2
oss2 2.19.1
outlines 0.1.11
outlines_core 0.2.11
packaging 23.2
pandas 2.2.3
parso 0.8.5
partial-json-parser 0.2.1.1.post4
peft 0.16.0
pexpect 4.9.0
pillow 11.3.0
pip 24.0
pipdeptree 2.9.6
platformdirs 4.2.2
prometheus_client 0.22.1
prometheus-fastapi-instrumentator 7.1.0
prompt_toolkit 3.0.52
propcache 0.3.0
protobuf 6.32.0
psutil 7.0.0
ptyprocess 0.7.0
pure_eval 0.2.3
py-cpuinfo 9.0.0
pyarrow 19.0.1
pybase64 1.4.2
pybind11 2.13.6
pycountry 24.6.1
pycparser 2.22
pycryptodome 3.23.0
pycryptodomex 3.21.0
pydantic 2.11.7
pydantic_core 2.33.2
pydantic-extra-types 2.10.6
pydub 0.25.1
Pygments 2.19.2
pynvml 12.0.0
pyparsing 3.2.3
pyproject_hooks 1.2.0
python-dateutil 2.9.0.post0
python-dotenv 1.2.1
python-json-logger 4.0.0
python-multipart 0.0.20
pytz 2023.4
PyYAML 6.0.2
pyzmq 27.0.1
ray 2.49.1
referencing 0.36.2
regex 2025.7.34
requests 2.32.3
rich 14.1.0
rich-toolkit 0.17.0
rignore 0.6.4
rouge 1.0.1
rpds-py 0.27.0
ruff 0.10.0
safehttpx 0.1.6
safetensors 0.6.2
scipy 1.16.1
semantic-version 2.10.0
sentencepiece 0.2.0
sentry-sdk 2.46.0
setproctitle 1.3.6
setuptools 79.0.1
sgl-kernel 0.3.5
sglang 0.5.1
shellingham 1.5.4
simplejson 3.20.1
six 1.16.0
sniffio 1.3.1
sortedcontainers 2.4.0
soundfile 0.13.1
soxr 0.5.0.post1
stack-data 0.6.3
starlette 0.47.3
sympy 1.14.0
tensorboard 2.20.0
tensorboard-data-server 0.7.2
tiktoken 0.11.0
timm 1.0.16
tokenizers 0.21.4
tomli 2.0.1
tomlkit 0.13.3
torch 2.8.0
torch_memory_saver 0.0.8
torchao 0.9.0
torchaudio 2.8.0
torchvision 0.23.0
tqdm 4.67.1
traitlets 5.14.3
transformer_engine 2.6.0.post1
transformer_engine_cu12 2.6.0.post1
transformer_engine_torch 2.6.0.post1
transformers 4.55.2
transformers-stream-generator 0.0.5
triton 3.4.0
trl 0.19.1
typeguard 4.3.0
typer 0.16.1
typing_extensions 4.12.2
typing-inspection 0.4.1
tzdata 2025.1
urllib3 2.0.7
uvicorn 0.35.0
uvloop 0.21.0
vllm 0.10.2
watchfiles 1.1.1
wcwidth 0.2.13
websockets 15.0.1
Werkzeug 3.1.3
wheel 0.45.1
xformers 0.0.32.post1
xgrammar 0.1.23
xxhash 3.5.0
yarl 1.18.3
zipp 3.23.0
zstandard 0.24.0
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start by reproducing the reported Gemma SFT runs using the two A100 command lines and compare memory during evaluation and training. Trace the failure through swift/llm/train/sft.py, swift/trainers/trainers.py, and the Gemma forward path shown in modeling_gemma3.py; done means the reported step- or evaluation-related memory growth no longer ends in CUDA OOM.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 35/100