Lab Guide

Format: Python script (lab.py) · Time estimate: 60–90 minutes · Prerequisites: PyTorch, Transformers, Datasets, DeepSpeed installed; GPU recommended (Colab T4 sufficient)

Lab Files

⬇ lab.py (student) ⬇ lab_answer.py (solution)

Setup

Install the required packages before starting:

# Core training dependencies
pip install torch transformers datasets accelerate deepspeed

# For plotting training curves
pip install matplotlib tensorboard

If using Google Colab, select a T4 GPU runtime. The exercises use GPT-2 (124M parameters), which fits comfortably on a single T4 with DeepSpeed.

GPU Requirement

Exercise 1 can run on CPU (slowly). Exercises 2 and 3 require a GPU. If you do not have local GPU access, use Google Colab (free T4) or a cloud GPU instance.

Exercise 1: Training with HuggingFace Trainer

Goal

Set up a causal language model training run using HuggingFace Trainer, training GPT-2 on a small text corpus. Understand the core components: model, tokeniser, dataset, data collator, and training arguments.

Steps

Load the gpt2 model and tokeniser from HuggingFace:

from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# GPT-2 has no pad token by default
tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = tokenizer.eos_token_id

Load a small dataset. We use wikitext-2-raw-v1 as a compact, well-known corpus:

from datasets import load_dataset

dataset = load_dataset("wikitext", "wikitext-2-raw-v1")
# dataset has train, validation, and test splits

Tokenise and group texts into fixed-length blocks (context packing):

block_size = 512

def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=False)

def group_texts(examples):
    """Concatenate all texts and split into chunks of block_size."""
    concatenated = {k: sum(examples[k], []) for k in examples.keys()}
    total_length = len(concatenated["input_ids"])
    total_length = (total_length // block_size) * block_size
    result = {
        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
        for k, t in concatenated.items()
    }
    result["labels"] = result["input_ids"].copy()
    return result

tokenized = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
lm_dataset = tokenized.map(group_texts, batched=True)

Configure TrainingArguments and launch training:

from transformers import Trainer, TrainingArguments, DataCollatorForLanguageModeling

training_args = TrainingArguments(
    output_dir="./gpt2-wikitext",
    overwrite_output_dir=True,
    num_train_epochs=1,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    eval_strategy="steps",
    eval_steps=200,
    save_steps=500,
    logging_steps=50,
    learning_rate=5e-5,
    weight_decay=0.01,
    warmup_steps=100,
    bf16=True,  # Use BF16 mixed precision if GPU supports it
    logging_dir="./logs",
    report_to="tensorboard",
)

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=False  # mlm=False for causal LM
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=lm_dataset["train"],
    eval_dataset=lm_dataset["validation"],
    data_collator=data_collator,
)

trainer.train()

After training completes, record the final training loss and evaluation loss. Generate a sample to verify the model works:

import torch

prompt = "The history of artificial intelligence"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
    output = model.generate(**inputs, max_new_tokens=100, temperature=0.7, do_sample=True)
print(tokenizer.decode(output[0], skip_special_tokens=True))

Expected Output

Training Loss (final): ~3.2 - 3.5
Eval Loss: ~3.4 - 3.7
Training time: ~5-10 minutes on T4

Generated text should be coherent English, though quality
will be limited by the small dataset and single epoch.

Key Observations

The group_texts function implements context packing — no padding waste.
bf16=True enables mixed-precision training, halving memory usage.
The model already knows English from its original pre-training; we are doing continued pre-training on WikiText.

Exercise 2: DeepSpeed Configuration

Goal

Configure DeepSpeed ZeRO Stage 2 to reduce memory usage, enabling larger batch sizes or larger models on the same hardware.

Steps

Create a DeepSpeed configuration file ds_config.json:

{
  "bf16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2,
    "offload_optimizer": {
      "device": "none"
    },
    "allgather_partitions": true,
    "allgather_bucket_size": 2e8,
    "overlap_comm": true,
    "reduce_scatter": true,
    "reduce_bucket_size": 2e8,
    "contiguous_gradients": true
  },
  "gradient_accumulation_steps": 4,
  "gradient_clipping": 1.0,
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "wall_clock_breakdown": true
}

Modify the TrainingArguments to use DeepSpeed:

training_args_ds = TrainingArguments(
    output_dir="./gpt2-wikitext-ds",
    overwrite_output_dir=True,
    num_train_epochs=1,
    per_device_train_batch_size=8,    # Doubled thanks to ZeRO
    per_device_eval_batch_size=8,
    gradient_accumulation_steps=4,
    eval_strategy="steps",
    eval_steps=200,
    save_steps=500,
    logging_steps=50,
    learning_rate=5e-5,
    weight_decay=0.01,
    warmup_steps=100,
    bf16=True,
    logging_dir="./logs-ds",
    report_to="tensorboard",
    deepspeed="ds_config.json",       # Enable DeepSpeed
)

Run training with the DeepSpeed config. Compare GPU memory usage and throughput (samples/second) against Exercise 1.
Experiment: change stage from 2 to 3 and observe memory and speed differences. Stage 3 uses less memory but has higher communication overhead.

Expected Output

Without DeepSpeed (Exercise 1):
  GPU memory: ~6-8 GB
  Throughput: ~X samples/sec

With DeepSpeed ZeRO-2:
  GPU memory: ~4-5 GB (reduced)
  Throughput: ~Y samples/sec (similar or slightly lower)
  Effective batch size: 8 * 4 = 32 (larger due to gradient accumulation)

Multi-GPU Note

On a single GPU, ZeRO-2 provides modest memory savings. The real benefits appear with multiple GPUs, where optimiser states and gradients are sharded across devices. If you have access to a multi-GPU machine, try: deepspeed --num_gpus=2 lab.py

Exercise 3: Analysing Training Curves

Goal

Visualise and interpret training metrics to diagnose common training issues: learning rate schedules, loss curves, and throughput.

Steps

Load the training logs from the Trainer's log history:

import matplotlib.pyplot as plt

log_history = trainer.state.log_history

# Extract training loss
train_steps = [log["step"] for log in log_history if "loss" in log]
train_loss = [log["loss"] for log in log_history if "loss" in log]

# Extract eval loss
eval_steps = [log["step"] for log in log_history if "eval_loss" in log]
eval_loss = [log["eval_loss"] for log in log_history if "eval_loss" in log]

# Extract learning rate
lr_steps = [log["step"] for log in log_history if "learning_rate" in log]
lr_values = [log["learning_rate"] for log in log_history if "learning_rate" in log]

Plot the training and evaluation loss curves:

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 5))

ax1.plot(train_steps, train_loss, label="Train Loss", color="#58a6ff")
ax1.plot(eval_steps, eval_loss, label="Eval Loss", color="#f78166", marker="o")
ax1.set_xlabel("Step")
ax1.set_ylabel("Loss")
ax1.set_title("Training & Evaluation Loss")
ax1.legend()
ax1.grid(True, alpha=0.3)

ax2.plot(lr_steps, lr_values, color="#7ee787")
ax2.set_xlabel("Step")
ax2.set_ylabel("Learning Rate")
ax2.set_title("Learning Rate Schedule")
ax2.grid(True, alpha=0.3)

plt.tight_layout()
plt.savefig("training_curves.png", dpi=150, bbox_inches="tight")
plt.show()

Compute and display training throughput:

# Calculate tokens per second
total_tokens = len(lm_dataset["train"]) * block_size
total_time = sum(
    log.get("train_runtime", 0) for log in log_history if "train_runtime" in log
)
if total_time > 0:
    tokens_per_second = total_tokens / total_time
    print(f"Training throughput: {tokens_per_second:,.0f} tokens/sec")
    print(f"Total tokens processed: {total_tokens:,}")
    print(f"Total training time: {total_time:.1f} seconds")

Write a cost estimation function that extrapolates from your run:

def estimate_training_cost(
    target_tokens: int,
    target_params: int,
    measured_tokens_per_sec: float,
    gpu_cost_per_hour: float = 2.00,  # A100 on-demand
    num_gpus: int = 8,
):
    """Estimate cost for a larger training run based on measured throughput."""
    # Scale throughput: roughly linear with GPU count, sublinear with model size
    # FLOPs scale as 6 * N * D; throughput scales inversely with model size
    scale_factor = (124_000_000 / target_params)  # GPT-2 = 124M params
    scaled_tps = measured_tokens_per_sec * num_gpus * scale_factor

    total_seconds = target_tokens / scaled_tps
    total_hours = total_seconds / 3600
    total_cost = total_hours * num_gpus * gpu_cost_per_hour

    return {
        "gpu_hours": total_hours * num_gpus,
        "wall_clock_days": total_hours / 24,
        "estimated_cost_usd": total_cost,
    }

# Example: estimate cost for a 7B model on 140B tokens
estimate = estimate_training_cost(
    target_tokens=140_000_000_000,
    target_params=7_000_000_000,
    measured_tokens_per_sec=tokens_per_second,
)
print(f"Estimated GPU hours: {estimate['gpu_hours']:,.0f}")
print(f"Estimated wall-clock time: {estimate['wall_clock_days']:.1f} days")
print(f"Estimated cost: ${estimate['estimated_cost_usd']:,.0f}")

Expected Output

Training throughput: ~15,000-25,000 tokens/sec (T4)
Total tokens processed: ~2,000,000
Total training time: ~300-600 seconds

7B model estimate (8x A100):
  GPU hours: ~1,200-1,500
  Wall-clock time: ~6-8 days
  Estimated cost: $15,000-$25,000

Note: These are rough estimates. Real-world costs depend on
GPU utilisation, communication overhead, and checkpointing time.

What to Look For in Training Curves

Smooth loss decrease: Healthy training shows a steady decline. Spikes suggest data issues or too-high learning rate.
Eval loss tracking train loss: If eval loss diverges upward while train loss decreases, the model is overfitting.
Warmup phase: The first 100 steps should show the learning rate ramping up. Loss may be unstable during warmup — that is normal.
Diminishing returns: Loss improvement slows as training progresses. This is expected and is the basis of scaling laws.

Finishing Up

Once you have completed all three exercises:

Run the full script end-to-end to verify everything executes without errors.
Save your training curves plot (training_curves.png).
Review your cost estimates — are they in the right ballpark compared to published training costs for similar models?
Compare your implementation with lab_answer.py.

Self-Check

✓ Exercise 1: Training completes without errors and final loss is below 4.0.
✓ Exercise 1: You can generate coherent text from the trained model.
✓ Exercise 2: DeepSpeed config is valid and training runs with reduced memory usage.
✓ Exercise 2: You can explain the difference between ZeRO Stage 2 and Stage 3.
✓ Exercise 3: You have a saved plot showing loss curves and learning rate schedule.
✓ Exercise 3: Your cost estimate for a 7B model is in a realistic range ($10K–$50K on cloud GPUs).