Fine-Tune Llama 4 70B on AWS SageMaker for Enterprise

Step-by-step guide to fine-tuning Llama 4 70B on AWS SageMaker using LoRA and QLoRA for production enterprise deployment.
·
6 min read
·
Mark
·
Machine Learning Fine-Tuning

Problem: Getting Llama 4 70B to Actually Know Your Business

Generic Llama 4 70B gives generic answers. Your enterprise needs a model that understands your domain — your products, your tone, your internal terminology.

Fine-tuning on SageMaker is the production path, but it's easy to get wrong: wrong instance type, misconfigured LoRA ranks, and you've burned $400 on a failed training job.

You'll learn:

  • How to set up SageMaker Training Jobs for 70B-scale models
  • How to apply QLoRA to fit fine-tuning into a cost-effective GPU footprint
  • How to push a fine-tuned adapter to S3 and serve it behind a SageMaker endpoint

Time: 90 min | Level: Advanced

Why This Happens

Llama 4 70B has 70 billion parameters. Full fine-tuning would require ~140GB of GPU VRAM just to hold the weights in bf16 — that's 4x A100 80GB at minimum.

QLoRA (Quantized LoRA) solves this by loading the base model in 4-bit NF4 precision and training only small low-rank adapter matrices. You get 95%+ of the fine-tuning quality at roughly 20% of the compute cost.

SageMaker wraps this in managed infrastructure: spot instances, S3 checkpointing, and one-click endpoint deployment.

Common symptoms that you need fine-tuning (not just prompting):

  • The model ignores your formatting instructions consistently
  • Domain-specific terms are hallucinated or misunderstood
  • RAG retrieval isn't enough because the reasoning style needs to change

Solution

Step 1: Prepare Your Dataset

SageMaker training scripts read from S3. Your data needs to be in JSONL format with instruction-response pairs.

# prepare_dataset.py
import json

def format_record(instruction: str, response: str) -> dict:
    # Llama 4 uses a specific chat template
    # Fine-tune ON the template, not around it
    return {
        "text": f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n{instruction}<|eot_id|>"
                f"<|start_header_id|>assistant<|end_header_id|>\n{response}<|eot_id|>"
    }

records = [
    format_record("What is our refund policy?", "Orders can be returned within 30 days..."),
    # ... your domain data
]

with open("train.jsonl", "w") as f:
    for r in records:
        f.write(json.dumps(r) + "\n")

Upload to S3:

aws s3 cp train.jsonl s3://your-bucket/llama4-finetune/data/train.jsonl
aws s3 cp val.jsonl s3://your-bucket/llama4-finetune/data/val.jsonl

Expected: At least 500 records for meaningful fine-tuning; 5,000+ for domain adaptation. Under 200 records risks overfitting.

Step 2: Write the SageMaker Training Script

This script runs inside SageMaker's managed container. It reads env vars SageMaker injects automatically.

# train.py
import os
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model, TaskType
from trl import SFTTrainer
from datasets import load_dataset

# SageMaker injects these paths at runtime
DATA_DIR = os.environ.get("SM_CHANNEL_TRAINING", "/opt/ml/input/data/training")
MODEL_DIR = os.environ.get("SM_MODEL_DIR", "/opt/ml/model")
OUTPUT_DIR = os.environ.get("SM_OUTPUT_DATA_DIR", "/opt/ml/output")

BASE_MODEL = "meta-llama/Llama-4-70B-Instruct"  # or your S3 path

def main():
    # 4-bit quantization config — fits 70B into 2x A10G (48GB total)
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,   # saves ~0.4 bits per param
        bnb_4bit_quant_type="nf4",         # NF4 is optimal for normally distributed weights
        bnb_4bit_compute_dtype=torch.bfloat16,
    )

    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
    )
    model.config.use_cache = False  # Required for gradient checkpointing

    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"  # Prevents attention mask issues

    # LoRA config — rank 16 is the sweet spot for 70B domain adaptation
    lora_config = LoraConfig(
        r=16,                    # Rank: higher = more capacity, more VRAM
        lora_alpha=32,           # Scale factor: typically 2x rank
        target_modules=[
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",  # Include MLP layers for domain adaptation
        ],
        lora_dropout=0.05,
        bias="none",
        task_type=TaskType.CAUSAL_LM,
    )

    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()
    # Expect: ~0.5% of parameters trainable (keeps cost low)

    dataset = load_dataset("json", data_files={
        "train": f"{DATA_DIR}/train.jsonl",
        "validation": f"{DATA_DIR}/val.jsonl",
    })

    training_args = TrainingArguments(
        output_dir=MODEL_DIR,
        num_train_epochs=3,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=8,    # Effective batch size = 16
        gradient_checkpointing=True,       # Trade compute for VRAM
        learning_rate=2e-4,
        bf16=True,
        logging_steps=25,
        evaluation_strategy="steps",
        eval_steps=100,
        save_strategy="steps",
        save_steps=100,
        load_best_model_at_end=True,
        warmup_ratio=0.03,
        lr_scheduler_type="cosine",
        report_to="none",                  # Disable wandb inside SageMaker
    )

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["validation"],
        dataset_text_field="text",
        max_seq_length=2048,
        packing=True,  # Packs multiple short examples into one sequence (faster)
    )

    trainer.train()
    trainer.save_model(MODEL_DIR)
    tokenizer.save_pretrained(MODEL_DIR)

if __name__ == "__main__":
    main()

Step 3: Configure and Launch the Training Job

# launch_training.py
import boto3
import sagemaker
from sagemaker.pytorch import PyTorch

session = sagemaker.Session()
role = sagemaker.get_execution_role()

# ml.p4d.24xlarge = 8x A100 80GB — use for 70B with headroom
# ml.g5.48xlarge = 8x A10G 24GB — cheaper, works with QLoRA
INSTANCE_TYPE = "ml.g5.48xlarge"

estimator = PyTorch(
    entry_point="train.py",
    source_dir="./src",                # Directory containing train.py + requirements.txt
    role=role,
    framework_version="2.2",
    py_version="py310",
    instance_count=1,
    instance_type=INSTANCE_TYPE,
    use_spot_instances=True,           # Up to 90% cost savings
    max_wait=86400,                    # 24hr max wait for spot
    max_run=43200,                     # 12hr max training time
    checkpoint_s3_uri="s3://your-bucket/llama4-finetune/checkpoints/",
    hyperparameters={},
    environment={
        "HUGGING_FACE_HUB_TOKEN": "<your-token>",  # For gated model access
    },
    volume_size=200,                   # GB — 70B weights need space
)

estimator.fit({
    "training": "s3://your-bucket/llama4-finetune/data/"
})

print(f"Model artifacts: {estimator.model_data}")

Launch it:

python launch_training.py

Expected: Job appears in SageMaker console under Training Jobs. On ml.g5.48xlarge with 5k records and 3 epochs, expect 4-6 hours.

If it fails:

  • ResourceLimitExceeded: Request a quota increase for ml.g5.48xlarge in your AWS region — these are high-demand instances
  • CUDA OOM: Reduce per_device_train_batch_size to 1 and increase gradient_accumulation_steps to 16
  • ModuleNotFoundError: Add missing packages to src/requirements.txt (peft, trl, bitsandbytes, datasets)

SageMaker Training Job dashboard showing active job Training job running — watch CloudWatch logs for loss curves

Step 4: Deploy the Fine-Tuned Model

After training, your adapter weights are in S3. Deploy them merged with the base model for inference.

# deploy.py
import sagemaker
from sagemaker.huggingface import HuggingFaceModel

session = sagemaker.Session()
role = sagemaker.get_execution_role()

# Point to the adapter artifacts from your training job
model_data = "s3://your-bucket/llama4-finetune/checkpoints/model.tar.gz"

huggingface_model = HuggingFaceModel(
    model_data=model_data,
    role=role,
    transformers_version="4.37",
    pytorch_version="2.1",
    py_version="py310",
    env={
        "HF_MODEL_ID": "meta-llama/Llama-4-70B-Instruct",
        "HF_TASK": "text-generation",
        "SM_NUM_GPUS": "4",
        "MAX_INPUT_LENGTH": "2048",
        "MAX_TOTAL_TOKENS": "4096",
    },
)

predictor = huggingface_model.deploy(
    initial_instance_count=1,
    instance_type="ml.g5.12xlarge",   # 4x A10G for inference
    endpoint_name="llama4-enterprise-v1",
)

# Test it
response = predictor.predict({
    "inputs": "<|begin_of_text|><|start_header_id|>user<|end_header_id|>\nWhat is our refund policy?<|eot_id|><|start_header_id|>assistant<|end_header_id|>",
    "parameters": {
        "max_new_tokens": 256,
        "temperature": 0.1,    # Low temp for factual enterprise responses
        "do_sample": True,
    }
})

print(response[0]["generated_text"])

Expected: Response using your fine-tuned domain knowledge, not the base model's generic answer.

Verification

Run a quick eval against your validation set to confirm the fine-tuned model outperforms the base:

python evaluate.py \
  --endpoint llama4-enterprise-v1 \
  --val-data s3://your-bucket/llama4-finetune/data/val.jsonl \
  --output eval-results.json

You should see: Improved exact-match or ROUGE scores on domain-specific prompts vs. the base model baseline. A 15-30% improvement on domain accuracy is typical for 5k+ record datasets.

Evaluation results showing fine-tuned vs base model comparison Fine-tuned model (blue) vs. base Llama 4 (gray) on domain accuracy

Estimated Costs

ComponentInstanceApprox. Cost
Training (6hr)ml.g5.48xlarge spot~$35-60
Training (6hr)ml.g5.48xlarge on-demand~$120-180
Inference (per hour)ml.g5.12xlarge~$16/hr

Use spot instances for training — SageMaker's checkpoint support means interrupted jobs resume automatically.

What You Learned

  • QLoRA makes 70B fine-tuning practical: ~0.5% trainable parameters with 95%+ quality retention
  • SageMaker spot instances cut training cost by up to 90% with zero code changes
  • packing=True in SFTTrainer significantly increases throughput for short training examples

Limitations to know:

  • Adapter merging at inference time adds ~2-3 seconds of cold start latency
  • QLoRA adapters are not portable across quantization configs — retrain if you change bnb_config
  • Fine-tuning on fewer than 500 examples often hurts more than it helps; use few-shot prompting instead for small datasets

When NOT to use this approach:

  • You need sub-100ms inference — use a smaller model (Llama 4 8B) or quantized base without adapters
  • Your use case changes weekly — prompting with RAG is faster to iterate than retraining

Tested on Llama 4 70B Instruct, PyTorch 2.2, SageMaker SDK 2.x, Python 3.10 — Ubuntu 22.04 base container