Training LLMs

Hands-on lab with starter code you can copy + run, and guided steps to produce the artifact for this module.

Setup

Open a fresh terminal in a new folder, e.g. mkdir lab-6 && cd lab-6.
Create a Python 3.11 venv: python -m venv .venv && source .venv/bin/activate.
Install the dependencies noted on each starter snippet (the small pip install … line above the code).
Copy each starter into a file with the suggested name, then run it and capture the output.

Starter code

01 · Next-token training loop with gradient accumulation

save as: 06_01_lab.py

python

import torch
from torch.utils.data import DataLoader

opt = torch.optim.AdamW(model.parameters(), lr=3e-4, betas=(0.9, 0.95))
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=1000)
scaler = torch.cuda.amp.GradScaler()
ACCUM = 8

for step, batch in enumerate(loader):
    with torch.cuda.amp.autocast(dtype=torch.bfloat16):
        out = model(**batch, labels=batch["input_ids"])
        loss = out.loss / ACCUM
    scaler.scale(loss).backward()
    if (step + 1) % ACCUM == 0:
        scaler.unscale_(opt)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(opt); scaler.update(); sched.step(); opt.zero_grad()

Run with: python 06_01_lab.py

02 · Wrap a model in FSDP for multi-GPU training

save as: 06_02_lab.py

python

from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
import functools

policy = functools.partial(
    transformer_auto_wrap_policy,
    transformer_layer_cls={Block},   # your transformer block class
)
model = FSDP(model, auto_wrap_policy=policy, mixed_precision=None)

Run with: python 06_02_lab.py

Guided steps

Work through each step in order. Stop when your output matches the expected artifact — then push your work to your lab folder.

1
Train a 10M-param LM on TinyStories
Do — Pre-train a 10M-parameter decoder on the TinyStories dataset for one epoch with AdamW + cosine schedule.
Expected artifact — loss.png going from ~7 → <3.5; a generated story that is grammatical and on-topic.
Run free:
2
Gradient-accumulation correctness
Do — Compare a batch=16 step against accum=4×batch=4. Loss curves should match within noise.
Expected artifact — Two loss curves overlaid; max divergence < 5%.
Run free:
3
FSDP smoke test
Do — Wrap your model in FSDP across 2 GPUs (or 2 CPU processes via gloo) and confirm one training step runs.
Expected artifact — Console prints loss per rank; both ranks report identical loss after all-reduce.
Run free:

Authoritative sources

Independent primary sources used to check and corroborate the guidance on this page.

NIST National Initiative for Cybersecurity Education (NICE)
US national framework for cybersecurity work roles and competencies.
Financial Action Task Force (FATF)
Global AML/CFT standards referenced by financial-crime curricula.
FBI Internet Crime Complaint Center — Annual Reports
Primary loss and typology data used in fraud-investigation training.
US Federal Trade Commission — Consumer Sentinel Data Book
Official complaint statistics for teaching fraud trend analysis.