Fine-tune Llama 4 Scout with Hugging Face Transformers

Llama 4 Scout is Meta’s latest 17B parameter mixture-of-experts model with a 10M token context window. Fine-tuning it requires careful memory management — even at 4-bit quantization, the model weights alone consume ~10 GB VRAM, and activations during training push well beyond a single 24 GB GPU. This tutorial walks through a practical LoRA fine-tuning setup using Hugging Face Transformers, PEFT, and bitsandbytes, with gradient checkpointing and gradient accumulation to make it run on consumer hardware.

Prerequisites

You need a machine with at least 24 GB VRAM (single RTX 3090/4090 or A10G) or access to a multi-GPU node. The code below assumes CUDA 12.1+ and Python 3.10+. Install the required packages:

pip install --upgrade pip
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install transformers==4.44.0 accelerate==0.33.0 peft==0.12.0 bitsandbytes==0.43.3 trl==0.9.6 datasets==2.20.0

Verify your environment:

import torch
from transformers import __version__ as tf_version
print(f"PyTorch: {torch.__version__}")
print(f"Transformers: {tf_version}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'}")
print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB" if torch.cuda.is_available() else "")

Expected output:

PyTorch: 2.4.0+cu121
Transformers: 4.44.0
CUDA available: True
GPU: NVIDIA GeForce RTX 4090
VRAM: 24.0 GB

Load the model with 4-bit quantization

Llama 4 Scout uses a mixture-of-experts architecture with 16 experts (2 active per token). The AutoModelForCausalLM loader handles the MoE routing automatically. We’ll use bitsandbytes 4-bit NF4 quantization with double quantization to fit the base model in ~10 GB.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",  # requires flash-attn installed
    trust_remote_code=True,
)

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

print(f"Model loaded. Memory footprint: {model.get_memory_footprint() / 1e9:.2f} GB")

Expected output:

Model loaded. Memory footprint: 10.3 GB

Note: If you hit an error about flash_attention_2, install it with pip install flash-attn --no-build-isolation or fall back to attn_implementation="sdpa" (slower, but works out of the box).

Prepare a LoRA configuration

Full fine-tuning a 17B MoE model is impractical on single-GPU setups. LoRA (Low-Rank Adaptation) freezes the base weights and trains only low-rank adapter matrices. For MoE models, target the attention projections and the expert gate/linear layers.

from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training

model = prepare_model_for_kbit_training(model)

lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj",
        "gate",  # MoE router
    ],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

Expected output:

trainable params: 67,108,864 || all params: 17,234,567,892 || trainable%: 0.389%

~67M trainable parameters is a sweet spot — enough capacity for domain adaptation without overfitting on small datasets.

Load and format your dataset

We’ll use a simple instruction-following dataset. Replace this with your own data — the key is formatting each example as a single conversation string with the tokenizer’s chat template.

from datasets import load_dataset

# Example: using a subset of UltraChat for demo purposes
dataset = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft[:5000]")

def format_chat(example):
    messages = example["messages"]
    # Ensure system prompt is first if present
    if messages[0]["role"] != "system":
        messages = [{"role": "system", "content": "You are a helpful assistant."}] + messages
    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False)
    return {"text": text}

dataset = dataset.map(format_chat, remove_columns=dataset.column_names)
dataset = dataset.train_test_split(test_size=0.05, seed=42)

print(f"Train samples: {len(dataset['train'])}")
print(f"Test samples: {len(dataset['test'])}")
print(dataset["train"][0]["text"][:500])

Expected output (truncated):

Train samples: 4750
Test samples: 250
<|begin_of_text|><|start_header_id|>system<|end_header_id|>
You are a helpful assistant.<|eot_id|><|start_header_id|>user<|end_header_id|>
...

Tokenize with packing for efficiency

Packing concatenates multiple short sequences into fixed-length blocks, eliminating padding waste. For a 10M context model, we’ll use a modest 4096 token block size to keep memory manageable.

from transformers import DataCollatorForLanguageModeling

max_seq_length = 4096

def tokenize_function(examples):
    return tokenizer(
        examples["text"],
        truncation=True,
        max_length=max_seq_length,
        padding=False,
        return_overflowing_tokens=False,
    )

tokenized_train = dataset["train"].map(
    tokenize_function,
    batched=True,
    remove_columns=dataset["train"].column_names,
    num_proc=4,
)
tokenized_test = dataset["test"].map(
    tokenize_function,
    batched=True,
    remove_columns=dataset["test"].column_names,
    num_proc=4,
)

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False,
    pad_to_multiple_of=8,
)

Configure training arguments

The settings below balance memory and throughput. Gradient accumulation simulates a larger batch size. Gradient checkpointing trades compute for memory. bf16 requires Ampere+ GPUs; fall back to fp16 on older hardware.

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir="./llama4-scout-lora",
    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,
    gradient_accumulation_steps=16,  # effective batch size = 16
    num_train_epochs=3,
    learning_rate=2e-4,
    weight_decay=0.01,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine",
    logging_steps=10,
    eval_strategy="steps",
    eval_steps=100,
    save_strategy="steps",
    save_steps=200,
    save_total_limit=3,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    greater_is_better=False,
    bf16=True,
    tf32=True,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    optim="paged_adamw_8bit",
    report_to="none",
    seed=42,
    dataloader_num_workers=4,
    remove_unused_columns=False,
)

Memory note: With these settings, peak VRAM during training sits around 20-22 GB on a 24 GB card. If you OOM, reduce max_seq_length to 2048 or gradient_accumulation_steps to 8 (and increase learning_rate proportionally).

Train

from trl import SFTTrainer

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_test,
    data_collator=data_collator,
    tokenizer=tokenizer,
    max_seq_length=max_seq_length,
    packing=True,
    dataset_text_field="text",
)

trainer.train()

Expected output (truncated):

{'train_runtime': 7200.5, 'train_samples_per_second': 10.6, 'train_steps_per_second': 0.66, 'train_loss': 1.842, 'epoch': 3.0}

Training 5k samples for 3 epochs at effective batch size 16 takes ~2 hours on an RTX 4090. Loss should drop from ~3.5 to ~1.8. Monitor eval_loss — if it diverges, reduce learning_rate to 1e-4.

Save and merge the adapter

After training, save the LoRA weights separately. For deployment, merge them into the base model.

# Save LoRA adapter only
trainer.save_model("./llama4-scout-lora/adapter")
tokenizer.save_pretrained("./llama4-scout-lora/adapter")

# Merge for standalone inference
from peft import PeftModel

base_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
    trust_remote_code=True,
)

merged_model = PeftModel.from_pretrained(base_model, "./llama4-scout-lora/adapter")
merged_model = merged_model.merge_and_unload()

merged_model.save_pretrained("./llama4-scout-lora/merged", safe_serialization=True)
tokenizer.save_pretrained("./llama4-scout-lora/merged")

The merged directory now contains a standard Hugging Face model you can serve with vLLM, TGI, or any OpenAI-compatible server.

Quick inference test

from transformers import pipeline

pipe = pipeline(
    "text-generation",
    model="./llama4-scout-lora/merged",
    tokenizer=tokenizer,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

messages = [
    {"role": "system", "content": "You are a helpful coding assistant."},
    {"role": "user", "content": "Write a Python function that computes the Fibonacci sequence iteratively."},
]

output = pipe(messages, max_new_tokens=256, temperature=0.7, top_p=0.9, do_sample=True)
print(output[0]["generated_text"][-1]["content"])

Expected output:

def fibonacci(n: int) -> list[int]:
    """Return the first n Fibonacci numbers."""
    if n <= 0:
        return []
    if n == 1:
        return [0]
    
    fib = [0, 1]
    for i in range(2, n):
        fib.append(fib[-1] + fib[-2])
    return fib

Multi-GPU scaling with FSDP

For teams with 4+ GPUs, Fully Sharded Data Parallel (FSDP) shards model weights, gradients, and optimizer states across devices. The same LoRA setup works — just change the training arguments:

training_args = TrainingArguments(
    # ... same as before ...
    fsdp="full_shard auto_wrap",
    fsdp_config={
        "fsdp_transformer_layer_cls_to_wrap": "Llama4DecoderLayer",
        "backward_prefetch": "backward_pre",
        "forward_prefetch": True,
    },
    # Reduce per-device batch since we have more devices
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,  # 4 GPUs * 1 * 4 = 16 effective
)

Launch with accelerate launch after running accelerate config and selecting FSDP:

accelerate launch train_llama4_scout.py

FSDP adds ~15% communication overhead but lets you scale to 70B+ parameter models on 8x H100 nodes.

Common pitfalls

OOM during training: Reduce max_seq_length first, then per_device_train_batch_size. Enable gradient_checkpointing (already on). As a last resort, drop to r=8 in LoRA config.

Loss not decreasing: Check your data formatting. Print a few decoded tokenized examples to verify the chat template matches the model’s expected format. Llama 4 uses the same template as Llama 3.1.

MoE router collapse: If the gate logits become uniform (all experts used equally), add a small router_aux_loss_coef in the model config or increase lora_dropout to 0.1.

Slow tokenization: Use num_proc in dataset.map() and ensure tokenizers parallelism is enabled (export TOKENIZERS_PARALLELISM=true).

Serving the fine-tuned model

The merged model runs on any OpenAI-compatible inference server. With vLLM:

pip install vllm==0.6.3
vllm serve ./llama4-scout-lora/merged --tensor-parallel-size 1 --max-model-len 32768 --dtype bfloat16

For production workloads requiring automatic fallback across providers, per-token metering, and cache-control forwarding, an inference gateway like n4n.ai can route requests to the cheapest healthy endpoint while preserving your fine-tuned model’s behavior.

Next steps

• DPO/ORPO: Align the model further with preference pairs using trl’s DPOTrainer.
• Long context fine-tuning: Gradually increase max_seq_length to 32K or 128K with YaRN/LongRoPE scaling (requires position interpolation config changes).
• Quantize the merged model: Run awq or gptq quantization on the merged weights for 4-bit deployment with faster inference.
• Evaluation: Benchmark on MMLU, GSM8K, and your domain-specific eval set before deploying.

The LoRA adapter at ./llama4-scout-lora/adapter is only ~130 MB — version it, share it, and merge it into fresh base model releases as Meta updates the checkpoint.