Out of Memory During LoRA Fine-Tuning on LLAMA-4-Scout-17B with H100 (80GB VRAM)

username: bensonbbn
cooked: <p>Dear community,</p>
<p>I am the student currently working on training the LLAMA-4-Scout-17B-16E-Instruct model using LoRA, running on an H100 GPU with 80GB VRAM (on Lambda Labs). However, I have encountered an out of memory error during the training process. I understand that this might fall slightly outside the scope of the course, but despite extensive research and reviewing various community discussions, I have not been able to resolve the issue.</p>
<p>Here is a brief outline of my setup:</p>
<p>Hardware: H100 (80GB VRAM)</p>
<p>Model: LLAMA-4-Scout-17B-16E-Instruct (download on unsloth hugging face)</p>
<p>Training Method: LoRA</p>
<p>Error: CUDA out of memory</p>
<p>Code snippet:<br>
import torch<br>
from transformers import AutoTokenizer,TrainingArguments,Trainer,DataCollatorForLanguageModeling,AutoModelForCausalLM<br>
from peft import LoraConfig, get_peft_model, TaskType<br>
from datasets import load_dataset<br>
from accelerate import dispatch_model<br>
from accelerate import Accelerator<br>
from accelerate.utils import get_balanced_memory, infer_auto_device_map<br>
import os<br>
os.environ[“PYTORCH_CUDA_ALLOC_CONF”] = “expandable_segments:True”</p>
<p>model_path = “/home/ubuntu/llama4”<br>
dataset_path = “llama_nc_instruction_train.jsonl”<br>
output_dir = “./merged_llama4_nccode”</p>
<p>print(“<img src="https://discuss.pytorch.org/images/emoji/apple/brain.png?v=14" title=":brain:" class="emoji" alt=":brain:" loading="lazy" width="20" height="20"> loading tokenizer…”)<br>
tokenizer = AutoTokenizer.from_pretrained(model_path)</p>
<p>print(“<img src="https://discuss.pytorch.org/images/emoji/apple/package.png?v=14" title=":package:" class="emoji" alt=":package:" loading="lazy" width="20" height="20"> loading model…（使用 safetensors）”)<br>
model = AutoModelForCausalLM.from_pretrained(<br>
model_path,<br>
torch_dtype=torch.bfloat16,<br>
low_cpu_mem_usage=True,<br>
trust_remote_code=True<br>
)</p>
<p>print(“<img src="https://discuss.pytorch.org/images/emoji/apple/wrench.png?v=14" title=":wrench:" class="emoji" alt=":wrench:" loading="lazy" width="20" height="20"> applying LoRA setting…”)<br>
lora_config = LoraConfig(<br>
r=8,<br>
lora_alpha=32, <span class="hashtag-raw">#有人用8</span><br>
target_modules=[“q_proj”, “v_proj”],<br>
lora_dropout=0.05,<br>
bias=“none”,<br>
task_type=TaskType.CAUSAL_LM,<br>
)</p>
<p>model = get_peft_model(model, lora_config)</p>
<p>print(“<img src="https://discuss.pytorch.org/images/emoji/apple/page_facing_up.png?v=14" title=":page_facing_up:" class="emoji" alt=":page_facing_up:" loading="lazy" width="20" height="20"> loading data…”)<br>
dataset = load_dataset(“json”, data_files=dataset_path, split=“train”)</p>
<p>def tokenize(example):<br>
tokenized_inputs = tokenizer(<br>
example[“text”],<br>
truncation=True,<br>
padding=“max_length”,<br>
max_length=4196<br>
)<br>
return tokenized_inputs</p>
<p>tokenized_dataset = dataset.map(tokenize, batched=True, remove_columns=[“text”])</p>
<p>print(“<img src="https://discuss.pytorch.org/images/emoji/apple/bullseye.png?v=14" title=":bullseye:" class="emoji" alt=":bullseye:" loading="lazy" width="20" height="20"> establish Trainer…”)<br>
training_args = TrainingArguments(<br>
output_dir=“./lora_tmp”,<br>
num_train_epochs=3,<br>
per_device_train_batch_size=1, <span class="hashtag-raw">#有人用64</span><br>
gradient_accumulation_steps=512,<br>
learning_rate=2e-4,<br>
logging_steps=10,<br>
save_strategy=“no”,<br>
)</p>
<p>trainer = Trainer(<br>
model=model,<br>
args=training_args,<br>
train_dataset=tokenized_dataset,<br>
tokenizer=tokenizer,<br>
data_collator=DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False),<br>
)</p>
<p>print(“<img src="https://discuss.pytorch.org/images/emoji/apple/rocket.png?v=14" title=":rocket:" class="emoji" alt=":rocket:" loading="lazy" width="20" height="20"> training…”)<br>
trainer.train()</p>
<p>print(“<img src="https://discuss.pytorch.org/images/emoji/apple/floppy_disk.png?v=14" title=":floppy_disk:" class="emoji" alt=":floppy_disk:" loading="lazy" width="20" height="20"> merge LoRA weight…”)<br>
model = model.merge_and_unload()</p>
<p>print(“<img src="https://discuss.pytorch.org/images/emoji/apple/package.png?v=14" title=":package:" class="emoji" alt=":package:" loading="lazy" width="20" height="20"> save model to:”, output_dir)<br>
model.save_pretrained(output_dir)<br>
tokenizer.save_pretrained(output_dir)</p>
<p>print(“<img src="https://discuss.pytorch.org/images/emoji/apple/white_check_mark.png?v=14" title=":white_check_mark:" class="emoji" alt=":white_check_mark:" loading="lazy" width="20" height="20"> finish！”)</p>
<p>and this is the error:</p>
<p><img src="https://discuss.pytorch.org/images/emoji/apple/brain.png?v=14" title=":brain:" class="emoji" alt=":brain:" loading="lazy" width="20" height="20"> 載入 tokenizer…<br>
<img src="https://discuss.pytorch.org/images/emoji/apple/package.png?v=14" title=":package:" class="emoji" alt=":package:" loading="lazy" width="20" height="20"> 載入模型…（使用 safetensors）<br>
Loading checkpoint shards: 100%|███████████████████████████████████████████████████████| 50/50 [00:00&lt;00:00, 457.56it/s]<br>
<img src="https://discuss.pytorch.org/images/emoji/apple/wrench.png?v=14" title=":wrench:" class="emoji" alt=":wrench:" loading="lazy" width="20" height="20"> 套用 LoRA 設定…<br>
<img src="https://discuss.pytorch.org/images/emoji/apple/page_facing_up.png?v=14" title=":page_facing_up:" class="emoji" alt=":page_facing_up:" loading="lazy" width="20" height="20"> 載入資料中…<br>
<img src="https://discuss.pytorch.org/images/emoji/apple/bullseye.png?v=14" title=":bullseye:" class="emoji" alt=":bullseye:" loading="lazy" width="20" height="20"> 建立 Trainer…<br>
/home/ubuntu/CNC代碼定義訓練黨TEST.py:68: FutureWarning: tokenizer is deprecated and will be removed in version 5.0.0 for Trainer.<strong>init</strong>. Use processing_class instead.<br>
trainer = Trainer(<br>
Traceback (most recent call last):<br>
File “/home/ubuntu/CNC代碼定義訓練黨TEST.py”, line 68, in<br>
trainer = Trainer(<br>
File “/home/ubuntu/llama_env/lib/python3.10/site-packages/transformers/utils/deprecation.py”, line 172, in wrapped_func<br>
return func(*args, **kwargs)<br>
File “/home/ubuntu/llama_env/lib/python3.10/site-packages/transformers/trainer.py”, line 614, in init<br>
self._move_model_to_device(model, args.device)<br>
File “/home/ubuntu/llama_env/lib/python3.10/site-packages/transformers/trainer.py”, line 901, in _move_model_to_device<br>
model = model.to(device)<br>
File “/home/ubuntu/llama_env/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 1355, in to<br>
return self._apply(convert)<br>
File “/home/ubuntu/llama_env/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 915, in _apply<br>
module._apply(fn)<br>
File “/home/ubuntu/llama_env/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 915, in _apply<br>
module._apply(fn)<br>
File “/home/ubuntu/llama_env/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 915, in _apply<br>
module._apply(fn)<br>
[Previous line repeated 4 more times]<br>
File “/home/ubuntu/llama_env/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 942, in _apply<br>
param_applied = fn(param)<br>
File “/home/ubuntu/llama_env/lib/python3.10/site-packages/torch/nn/modules/module.py”, line 1341, in convert<br>
return t.to(<br>
torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1.25 GiB. GPU 0 has a total capacity of 79.19 GiB of which 359.06 MiB is free. Including non-PyTorch memory, this process has 78.83 GiB memory in use. Of the allocated memory 78.38 GiB is allocated by PyTorch, and 8.21 MiB 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 (<a href="https://pytorch.org/docs/stable/notes/cuda.html#environment-variables" class="inline-onebox" rel="noopener nofollow ugc">CUDA semantics — PyTorch 2.7 documentation</a>)</p>
<p>Would anyone kindly offer any suggestions or best practices to address this issue? Are there specific parameters I should consider adjusting (e.g., batch size, gradient checkpointing, LoRA rank, etc.) to make it fit within the memory constraints? Or is this simply a case of hardware limitation, and even 80GB VRAM is not enough for this model.And I have tried the QLORA method.encountering the same quetion.</p>