diff --git a/examples/qwen3-0.6b-bdn-weights/training/README.md b/examples/qwen3-0.6b-bdn-weights/training/README.md new file mode 100644 index 0000000..9362a50 --- /dev/null +++ b/examples/qwen3-0.6b-bdn-weights/training/README.md @@ -0,0 +1,26 @@ +# Qwen3 0.6B Fine-Tuning with BDN Weight Loading + +This example fine-tunes the Qwen3 0.6B model using PyTorch on Baseten, with model weights loaded through [BDN](https://docs.baseten.co/training/concepts/storage#load-weights-and-data-with-bdn) instead of downloaded in the training script. BDN also supports loading training data from S3 (`s3://`), GCS (`gs://`), and R2 (`r2://`). + +**Resources:** 1 node, 1x H100 GPU + +## Prerequisites + +1. [Create a Baseten account](https://baseten.co/signup) if you don't already have one. +2. Install the Truss CLI: + ```bash + # pip + pip install -U truss + # or uv + uv add truss + ``` + +## Getting Started + +Initialize the example, navigate into the directory, and push the training job: + +```bash +truss train init --examples qwen3-0.6b-bdn-weights +cd qwen3-0.6b-bdn-weights +truss train push config.py +``` diff --git a/examples/qwen3-0.6b-bdn-weights/training/config.py b/examples/qwen3-0.6b-bdn-weights/training/config.py new file mode 100644 index 0000000..65cb5dc --- /dev/null +++ b/examples/qwen3-0.6b-bdn-weights/training/config.py @@ -0,0 +1,58 @@ +from truss_train import ( + TrainingProject, + TrainingJob, + Image, + Compute, + Runtime, + WeightsSource, + CacheConfig, + CheckpointingConfig, +) +from truss.base.truss_config import AcceleratorSpec + +BASE_IMAGE = "pytorch/pytorch:2.7.0-cuda12.8-cudnn9-runtime" + +training_runtime = Runtime( + start_commands=[ + "pip install transformers datasets accelerate", + "python train.py", + ], + cache_config=CacheConfig(enabled=True), + checkpointing_config=CheckpointingConfig(enabled=True), +) + +training_compute = Compute( + accelerator=AcceleratorSpec(accelerator="H100", count=1), +) + +training_job = TrainingJob( + image=Image(base_image=BASE_IMAGE), + compute=training_compute, + runtime=training_runtime, + weights=[ + # Model weights from Hugging Face, delivered through BDN. + # Available at mount_location before your training script starts. + WeightsSource( + source="hf://Qwen/Qwen3-0.6B", + mount_location="/app/models/Qwen/Qwen3-0.6B", + ), + # Training data from S3, also delivered through BDN. + # Replace with your own bucket and path. + # For private buckets, add auth — see: + # https://docs.baseten.co/development/model/bdn#aws-s3 + # + # WeightsSource( + # source="s3://my-bucket/training-data", + # mount_location="/app/data/training-data", + # auth={ + # "auth_method": "CUSTOM_SECRET", + # "auth_secret_name": "aws_credentials", + # }, + # ), + ], +) + +training_project = TrainingProject( + name="qwen3-0.6b-bdn-weights", + job=training_job, +) diff --git a/examples/qwen3-0.6b-bdn-weights/training/train.py b/examples/qwen3-0.6b-bdn-weights/training/train.py new file mode 100644 index 0000000..eb02ea3 --- /dev/null +++ b/examples/qwen3-0.6b-bdn-weights/training/train.py @@ -0,0 +1,199 @@ +"""Fine-tune Qwen3-0.6B using BDN-mounted weights. + +Model weights are pre-loaded at /app/models/Qwen/Qwen3-0.6B via BDN. +No download code needed — just load from the mount path. +""" + +import os +import math +import torch +import torch.optim as optim +from torch.utils.data import Dataset, DataLoader, random_split +from transformers import ( + AutoTokenizer, + AutoModelForCausalLM, + get_cosine_schedule_with_warmup, +) +from datasets import load_dataset + +# BDN mount path — weights are already here when the container starts. +MODEL_PATH = "/app/models/Qwen/Qwen3-0.6B" +OUTPUT_DIR = os.environ.get("BT_CHECKPOINT_DIR", "/tmp/training_checkpoints") + +BATCH_SIZE = 1 +GRAD_ACCUM_STEPS = 1 +NUM_EPOCHS = 1 +MAX_STEPS = 20 +MAX_LENGTH = 2048 +LEARNING_RATE = 2e-5 +WEIGHT_DECAY = 0.1 +WARMUP_STEPS = 5 +EVAL_STEPS = 10 +VAL_SET_RATIO = 0.05 +SAVE_EVERY_STEPS = 10 + + +def messages_to_text(messages): + chunks = [] + for m in messages: + role = m.get("role", "user") + content = m.get("content", "") + chunks.append(f"{role}: {content}") + return "\n".join(chunks) + + +class ChatDataset(Dataset): + def __init__(self, tokenizer, split="train", max_length=2048): + ds = load_dataset("winglian/pirate-ultrachat-10k", split=split) + ds = ds.filter( + lambda x: x.get("messages") is not None and len(x["messages"]) > 0 + ) + self.ds = ds + self.tokenizer = tokenizer + self.max_length = max_length + + def __len__(self): + return len(self.ds) + + def __getitem__(self, idx): + ex = self.ds[int(idx)] + text = messages_to_text(ex["messages"]) + enc = self.tokenizer( + text, + truncation=True, + padding="max_length", + max_length=self.max_length, + return_tensors="pt", + ) + input_ids = enc["input_ids"].squeeze(0) + attention_mask = enc["attention_mask"].squeeze(0) + labels = input_ids.clone() + labels[attention_mask == 0] = -100 + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels": labels, + } + + +def get_device(model): + return next(model.parameters()).device + + +def validate(model, dataloader): + model.eval() + total_loss = 0.0 + total_count = 0 + device = get_device(model) + with torch.no_grad(): + for batch in dataloader: + input_ids = batch["input_ids"].to(device) + attention_mask = batch["attention_mask"].to(device) + labels = batch["labels"].to(device) + outputs = model( + input_ids=input_ids, attention_mask=attention_mask, labels=labels + ) + total_loss += outputs.loss.item() + total_count += 1 + avg_loss = total_loss / max(total_count, 1) + print(f"Validation loss: {avg_loss:.4f}") + model.train() + + +def main(): + use_cuda = torch.cuda.is_available() + + # Verify BDN mount + print(f"Loading model from BDN mount: {MODEL_PATH}") + assert os.path.exists(MODEL_PATH), f"BDN mount not found at {MODEL_PATH}" + print(f" Files: {os.listdir(MODEL_PATH)}") + + # Load from the BDN mount path — no download happens here. + tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + model = AutoModelForCausalLM.from_pretrained( + MODEL_PATH, + trust_remote_code=True, + torch_dtype=torch.bfloat16, + device_map="auto" if use_cuda else None, + ) + + # Dataset + full_dataset = ChatDataset(tokenizer, max_length=MAX_LENGTH) + val_size = max(int(len(full_dataset) * VAL_SET_RATIO), 1) + train_size = len(full_dataset) - val_size + print(f"Dataset: {len(full_dataset)} total | {train_size} train | {val_size} val") + + train_dataset, val_dataset = random_split(full_dataset, [train_size, val_size]) + train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True) + val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False) + + # Optimizer and scheduler + optimizer = optim.AdamW( + model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY + ) + steps_per_epoch = math.ceil(len(train_loader) / GRAD_ACCUM_STEPS) + total_steps = min(MAX_STEPS, steps_per_epoch * NUM_EPOCHS) + scheduler = get_cosine_schedule_with_warmup( + optimizer, + num_warmup_steps=min(WARMUP_STEPS, total_steps), + num_training_steps=total_steps, + ) + + # Train + model.train() + step = 0 + optimizer.zero_grad(set_to_none=True) + + for epoch in range(NUM_EPOCHS): + for batch_idx, batch in enumerate(train_loader): + if step >= MAX_STEPS: + break + + device = get_device(model) + input_ids = batch["input_ids"].to(device) + attention_mask = batch["attention_mask"].to(device) + labels = batch["labels"].to(device) + + with torch.cuda.amp.autocast(enabled=use_cuda, dtype=torch.bfloat16): + outputs = model( + input_ids=input_ids, + attention_mask=attention_mask, + labels=labels, + ) + loss = outputs.loss / GRAD_ACCUM_STEPS + + loss.backward() + + if (batch_idx + 1) % GRAD_ACCUM_STEPS == 0: + optimizer.step() + optimizer.zero_grad(set_to_none=True) + scheduler.step() + print(f"Step {step} | Loss: {(loss.item() * GRAD_ACCUM_STEPS):.4f}") + + if step > 0 and step % EVAL_STEPS == 0: + validate(model, val_loader) + + if step > 0 and step % SAVE_EVERY_STEPS == 0: + save_path = os.path.join(OUTPUT_DIR, f"checkpoint-{step}") + os.makedirs(save_path, exist_ok=True) + model.save_pretrained(save_path) + tokenizer.save_pretrained(save_path) + print(f"Saved checkpoint to {save_path}") + + step += 1 + + if step >= MAX_STEPS: + break + + # Final save + os.makedirs(OUTPUT_DIR, exist_ok=True) + model.save_pretrained(OUTPUT_DIR) + tokenizer.save_pretrained(OUTPUT_DIR) + print(f"Training complete. Model saved to {OUTPUT_DIR}") + + +if __name__ == "__main__": + main()