Skip to content

kamend/sagemaker-fine-tunning-example

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

1 Commit
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

Fine-Tuning Tutorial: Qwen 2.5-3B on Restaurant Reviews

A practical, end-to-end walkthrough of fine-tuning a small open-weight LLM on a custom dataset using Amazon SageMaker - from raw data to a deployed endpoint you can call from a web UI.

What you'll build: a specialised model that takes a free-text restaurant review and returns structured JSON (cuisine, rating, price range, highlights, sentiment, recommendation).

Techniques covered:

  • Synthetic dataset creation with Claude (LLM-as-annotator)
  • QLoRA fine-tuning on a quantised base model
  • SageMaker training jobs and HuggingFace endpoints
  • Custom inference handler (model_fn / predict_fn) for chat-style I/O

Prerequisites

  • AWS account with SageMaker access
  • IAM role for SageMaker (see setup below)
  • HuggingFace token with read access to Qwen/Qwen2.5-3B-Instruct
  • Anthropic API key (used once to annotate ~500 reviews)
  • Python 3.12+ and uv
  • GPU quota for ml.g5.2xlarge (training) and ml.g5.xlarge (inference) in your chosen region

AWS IAM role

Create a SageMaker execution role with the AmazonSageMakerFullAccess managed policy attached. Default name in this tutorial is KamendSageMakerExecutionRole

  • change ROLE_NAME in .env to match yours.

Cost estimate

Step Instance Duration Approx. cost
Training ml.g5.2xlarge ~30 min ~$0.65
Endpoint ml.g5.xlarge per hour ~$1.00/hr
Claude annotation API (Haiku) ~500 calls <$0.50

The endpoint keeps billing until you delete it - see Cleanup.


Project layout

fine_tunning_tutorial/
├── config.py               # shared config (reads from .env)
├── .env.example            # copy to .env and fill in
├── pyproject.toml
│
├── data/                   # Stage 1 - dataset prep
│   ├── download_reviews.py   Yelp sampler
│   ├── annotate_reviews.py   Claude annotator + splitter
│   └── schema.json           Output JSON schema
│
├── training/               # Stage 2 - train on SageMaker
│   ├── launch.py             Local launcher (uploads data, starts job)
│   └── sagemaker/            Runs inside the training container
│       ├── train.py
│       ├── inference.py
│       └── requirements.txt
│
├── deployment/             # Stage 3 - deploy endpoint
│   └── deploy.py
│
└── webapp/                 # Stage 4 - Flask UI for testing
    ├── app.py
    └── templates/index.html

Run scripts as modules from the project root, e.g. python -m training.launch. That way from config import X resolves naturally without any sys.path tricks.


Setup

# 1. Install dependencies
uv sync

# 2. Configure secrets and settings
cp .env.example .env
# Edit .env - at minimum set HF_TOKEN and ANTHROPIC_API_KEY

# 3. Verify AWS credentials
aws sts get-caller-identity

All tunable values (region, role name, model ID, instance types, endpoint name) live in .env. The config.py file just loads them with sensible defaults.


Stage 1 - Data preparation

Goal: produce train.jsonl / valid.jsonl / test.jsonl in chat-messages format, where each example is a review paired with the structured JSON the model should learn to emit.

1a. Download raw reviews

python -m data.download_reviews

Pulls the Yelp review dataset from HuggingFace, filters for restaurant-related reviews by keyword matching, and samples ~500 examples balanced across 1–5 star ratings. Writes data/raw_reviews.jsonl.

1b. Annotate with Claude

python -m data.annotate_reviews

Sends each raw review to Claude Haiku with an extraction prompt, validates the returned JSON against data/schema.json, and produces training examples in chat format:

{
  "messages": [
    {"role": "system", "content": "You are a restaurant review analyzer..."},
    {"role": "user", "content": "<review text>"},
    {"role": "assistant", "content": "{\"cuisine\": \"Italian\", ...}"}
  ]
}

This is LLM-as-annotator: we use a strong general model to create training data so a smaller specialised model can learn the same task cheaper and faster.

The script is resumable - if you stop it mid-run, it picks up where it left off via annotated_reviews.jsonl. At the end it shuffles and splits 80/10/10 into train.jsonl, valid.jsonl, test.jsonl.


Stage 2 - Training

Goal: fine-tune Qwen/Qwen2.5-3B-Instruct on your annotated data using QLoRA, then merge the adapter into the base model weights so the result is a drop-in replacement for the base model.

python -m training.launch

What launch.py does (locally)

  1. Uploads data/train.jsonl and data/valid.jsonl to the default SageMaker S3 bucket.
  2. Constructs a HuggingFace estimator pointing at training/sagemaker/ as the source directory.
  3. Calls estimator.fit(...), which provisions a ml.g5.2xlarge instance, pulls the HF training container, installs the extra deps from sagemaker/requirements.txt, and runs sagemaker/train.py.
  4. Streams logs back to your terminal. When it finishes, it prints the S3 URI of model.tar.gz - save this URI, you'll need it for deployment.

What sagemaker/train.py does (remotely)

  1. Loads the base model in 4-bit (BitsAndBytesConfig + NF4 quantisation).
  2. Attaches a LoRA adapter with r=16, targeting all attention and MLP projections.
  3. Uses SFTTrainer with DataCollatorForCompletionOnlyLM so loss is computed only on assistant response tokens, not the prompt.
  4. After training, reloads the base model in fp16, merges the LoRA adapter, and saves the full merged model.
  5. Copies inference.py and a minimal requirements.txt into model/code/ - SageMaker picks these up as the custom inference handler.
  6. SageMaker tars everything in /opt/ml/model/ into model.tar.gz and uploads it to S3.

Why merge the adapter? A merged fp16 model loads faster at inference time and needs no PEFT dependency in the inference container.

Tuning knobs

Hyperparameters live in training/launch.py (not .env - they're experiment-specific):

hyperparameters = {
    "epochs": 3,
    "batch_size": 4,
    "gradient_accumulation_steps": 4,
    "learning_rate": 2e-4,
    "max_seq_length": 1024,
    "lora_r": 16,
    "lora_alpha": 32,
    "lora_dropout": 0.05,
    "warmup_ratio": 0.1,
}

Stage 3 - Deployment

Goal: stand up a real-time HTTPS endpoint backed by a ml.g5.xlarge GPU instance, using the inference.py handler bundled with your model artifact.

python -m deployment.deploy --model-data s3://<bucket>/.../model.tar.gz

The --model-data URI is what launch.py printed at the end of training.

What the inference handler does

training/sagemaker/inference.py implements the four SageMaker hooks:

  • model_fn - loads the merged model + tokenizer once at container startup
  • input_fn - parses the JSON request body
  • predict_fn - applies the chat template, runs model.generate(...), strips the prompt tokens, returns the generated text
  • output_fn - serialises the response to JSON

The endpoint accepts either {"messages": [...]} (chat format) or {"inputs": "<prompt>"} (raw prompt), plus optional parameters (max_new_tokens, temperature).

Deployment takes ~5–10 minutes. When it completes, the endpoint is named ENDPOINT_NAME from your .env (default: qwen-3b-review).


Stage 4 - Try it out

python -m webapp.app

Opens a Flask dev server on http://localhost:5001. Paste a review, hit Analyze, and the UI shows both the parsed structured output and the raw model response.

Under the hood, app.py calls the SageMaker runtime client, parses the response (handling markdown fences / preamble text the model might emit), and renders it.


Cleanup

Endpoints bill by the hour. Delete yours when you're done.

aws sagemaker delete-endpoint --endpoint-name qwen-3b-review
aws sagemaker delete-endpoint-config --endpoint-config-name qwen-3b-review
aws sagemaker delete-model --model-name <model-name>

List what you have:

aws sagemaker list-endpoints

Configuration reference

All settings in .env (see .env.example for the full list):

Variable Purpose
HF_TOKEN Pulls the base model from HuggingFace during training
ANTHROPIC_API_KEY Used by annotate_reviews.py
REGION AWS region for all SageMaker resources
ROLE_NAME SageMaker execution role
MODEL_ID Base model (any causal-LM compatible with HF Transformers works)
TRAINING_INSTANCE Training instance type
INFERENCE_INSTANCE Endpoint instance type
ENDPOINT_NAME Name of the deployed endpoint
S3_PREFIX Prefix under the default SageMaker bucket
TRANSFORMERS_VERSION, *_PYTORCH_VERSION, *_PY_VERSION Pin the HuggingFace container versions

About

A practical, end-to-end walkthrough of fine-tuning a small open-weight LLM on a custom dataset using Amazon SageMaker

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

No releases published

Packages

 
 
 

Contributors