Setting Up a Training Environment on Each Platform

By the end of this lesson you will have a Python environment on your own machine that can train, not merely run, a model: the right PyTorch build for your track, the Hugging Face training libraries beside it, and a script that has actually taken fifty optimiser steps on your accelerator and told you what they cost. You will also know what the honest support status of your track is, and on what date that was checked, so that a failure later is a fact about a dated statement rather than a mystery.

Part 1 built an environment that could run a small model. This lesson extends it in three ways: the training libraries, a per-track recipe with its awkward edges named, and a verification step that fails loudly rather than quietly falling back to the CPU.

What has to be true before a training run

Four things, in this order. Each can fail on its own, and each fails differently.

What has to be working, from the bottom up

1. Driver and runtimeThe NVIDIA driver, the ROCm stack, or macOS itself. Checked with nvidia-smi, rocm-smi, or by macOS version.
2. A PyTorch build that matches itCUDA, ROCm and CPU builds are different wheels. Installing the wrong one gives a working import and no accelerator.
3. The training librariestransformers, datasets, trl, peft and accelerate, in the same virtual environment as that PyTorch.
4. A real training stepForward, backward, optimiser step, on the accelerator, in the precision you intend to use.

A training run needs all four. The verification script at the end of this lesson checks the top two, and the per-track sections say how to check the bottom two.

The gap that catches people is the second layer. A CPU-only wheel imports cleanly, trains correctly and takes a hundred times longer, and nothing in the output says so unless you look. That is why every script in this course prints its device on the first line.

The libraries, once, on every track

Whatever your track, the training libraries themselves are the same four packages installed into the same virtual environment that already holds PyTorch, using uv from Part 1, pinned here at uv 0.12.11 · verified 2026-09-08.

RunnableAll tracks

add the training libraries to the Part 1 environment

source ~/llm-course/.venv/bin/activate
uv pip install transformers datasets trl peft accelerate

The versions this part was written against are transformers transformers 5.16.1 · verified 2026-09-08, TRL TRL 1.12.0 · verified 2026-09-08 and PEFT PEFT 0.20.0 · verified 2026-09-08. Record what you actually got:

RunnableAll tracks

record the versions

python -c "import torch, transformers, trl, peft, datasets; print(torch.__version__, transformers.__version__, trl.__version__, peft.__version__, datasets.__version__)"

⌚Version note

TRL’s own release page notes that 1.12.0 is an accidental duplicate of 1.11.0. That does not change anything you type, but if a version string surprises you when you compare two machines, that is why. Checked 2026-09-09.

Per-track recipes

Track S — NVIDIA DGX Spark

The container is the supported path. A DGX Spark runs DGX OS with Docker and the NVIDIA container toolkit already configured, and NVIDIA’s own route to PyTorch on it is the NGC PyTorch container, whose catalog page gives the tag form xx.xx-py3, the run command docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:xx.xx-py3, and states that the image is multi-arch, which is what makes it work on the Spark’s aarch64 CPU. The page also states that the software stack in the container “has been validated for compatibility, and does not require any additional installation”.

Part 1’s setup-env-spark.sh already pulls that image and mounts your course directory at /workspace/course. For this part, start it the same way and add the training libraries inside the container, where they will sit beside the container’s own PyTorch:

RunnableTrack S · DGX Spark

training libraries inside the NGC container

pip install transformers datasets trl peft accelerate

Do not install PyTorch inside the container. It is already there, built for this hardware, and a pip install that replaces it is the most common way to turn a working Spark into a broken one.

Anything you want to keep must live under /workspace/course, which is your host directory. The rest of the container is discarded when the shell exits, because Part 1’s script passes --rm.

Track X — AMD Ryzen AI Max+ 395Partial

Training on this machine needs the ROCm build of PyTorch. The ROCm 10.0.0 compatibility matrix names gfx1151 without a support-tier qualifier, but AMD's own PyTorch install page does not mention the chip; the two pages were read on the same day and did not agree.

Read the two pages, then decide. The ROCm 10.0.0 compatibility matrix, which carries the date 2026-08-14 and was read for this lesson on 2026-09-09, lists the AMD Ryzen AI Max+ 300 series with the Radeon 8060S on the gfx1151 target in its APU table, giving Ubuntu 26.04 and Ubuntu 24.04.4 with the HWE kernel, and Windows 11 25H2. There is no support-tier qualifier and no footnote on that row. AMD’s “Install PyTorch for ROCm” page, read on the same day, documents ROCm 7.2.4, is dated 2026-07-15, and does not mention gfx1151, the Radeon 8060S or the Ryzen AI Max+ anywhere.

That is the whole honest status. Part 5 and Part 8 record the same disagreement for llama.cpp; it matters more here, because Vulkan is a perfectly good inference path on this machine and is not a training path at all. Training needs ROCm.

The install page gives three routes: a wheels package, currently from a nightly index (pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm7.2), a Docker image (docker pull rocm/pytorch:latest), and a source build. Part 1’s setup-env.sh takes the wheel route with TRACK=strix, and passes --pre for exactly this reason. Check the index on the page before running it, because a nightly index moves.

The page’s own verification is two commands:

RunnableTrack X · Ryzen AI Max+

AMD's own check, from the install page

python3 -c 'import torch' 2> /dev/null && echo 'Success' || echo 'Failure'
python3 -c 'import torch; print(torch.cuda.is_available())'

The second one printing True is what you need. The cuda in the name is not a mistake: the ROCm build exposes AMD GPUs through the same device name.

If ROCm will not cooperate, the CPU build finishes this part’s lab in minutes rather than seconds, and every later Level 3 lab states its CPU fallback. Record in the notebook that you are on the CPU and why; that note is worth more later than a run that silently used it.

Track M — Apple silicon

Two frameworks, and they are for different jobs. MLX is the native path for training on Apple silicon and the one Track M uses in this part’s lab, because mlx-lm ships a LoRA fine-tuning command. Its install page requires Apple silicon, macOS 14.0 or later and a native Python 3.10 or newer, and the install is pip install mlx. Part 1’s setup-env.sh with TRACK=mac installs it alongside PyTorch. Add mlx-lm, pinned at mlx-lm 0.31.3 · verified 2026-09-08:

RunnableTrack M · Apple silicon

the training path on a Mac

source ~/llm-course/.venv/bin/activate
uv pip install mlx-lm

PyTorch is still worth having on a Mac, and its MPS backend, documented as enabling “high-performance training on GPU for macOS devices with Metal programming framework”, runs the verification script below and most small PyTorch code. What it is not is the path this course takes for fine-tuning language models on a Mac: the transformers entry in the course’s version file records that there is no MLX target, so the Hugging Face training stack on a Mac means PyTorch MPS, and the mlx-lm route is better supported for this job.

Use PyTorch MPS for reading and adapting PyTorch code, and mlx-lm for the fine-tuning itself. Part 13 revisits the choice with more at stake.

Track N — NVIDIA desktop or laptop

The straightforward case. Install your distribution’s NVIDIA driver and nothing else; the PyTorch wheel brings its own CUDA runtime libraries. Part 1’s setup-env.sh with TRACK=nvidia does the install, and nvidia-smi is the first thing to check when the accelerator does not appear.

On Windows the primary path is WSL2 with Ubuntu, and the one firm rule from Part 1 still applies: install the NVIDIA Windows driver only, and no Linux GPU driver inside WSL. Everything in this part then happens inside the WSL2 terminal exactly as on Linux.

Two things differ from the other tracks and both are about capacity rather than support. Memory is VRAM and only VRAM, so the arithmetic in the next lesson is against the card’s own figure rather than the machine’s. And a card that is also driving your displays has less of that figure available than its specification suggests, which is worth measuring once rather than discovering during a long run.

The container escape hatch

Every track has a container route, and it is the right answer more often than pride suggests. On Track S it is the primary path. On Track X, AMD’s rocm/pytorch:latest image is described on the install page as the latest ROCm-tested PyTorch, which sidesteps the question of whether your system’s ROCm version matches your wheel. On Track N, the same NGC PyTorch image runs on x86.

The trade is the usual one. A container fixes the whole stack below your code at versions someone else validated together, at the cost of a large download and a layer between you and the machine. When you have spent two hours on a wheel index, take the container, and write in the notebook that you did: a run inside a container and a run outside it are two different environments and their results are not interchangeable.

✱Hardware note

Track M does not have this escape hatch. Docker’s own page on GPU support states that “GPU support in Docker Desktop is only available on Windows with the WSL2 backend”, and offers no macOS instructions at all, so a containerised PyTorch on a Mac is a CPU-only PyTorch. It is another reason the Mac training path in this course is MLX.

Editors, notebooks and long sessions

Three habits make the rest of Level 3 less painful.

Notebooks for exploring, scripts for running. JupyterLab installs with pip install jupyterlab and starts with jupyter lab. It is the right tool for looking at a tokeniser’s output or plotting a loss curve. It is the wrong tool for a run you intend to record, because a notebook’s execution order is not recoverable from the file and a run you cannot reconstruct is not a measurement. Every training run in this course is a script with arguments.

Edit remotely, run locally to the machine. VS Code’s Remote - SSH extension opens a folder on any machine with an SSH server, with the editor’s features intact; you connect with Remote-SSH: Connect to Host… from the command palette and give it user@hostname. Its documented requirements on the remote host are modest, so a Spark or a mini PC in another room is a perfectly comfortable place to work.

Detach before you start anything long. A training run tied to an SSH session dies with the session. Start long runs inside tmux or screen so that a closed laptop lid is not a lost afternoon. This part’s lab is short enough not to need it; Part 12’s is not.

Verify a real training step

Now prove it. The script below builds a small stack of linear layers, runs fifty optimiser steps on whichever device your machine offers, prints the loss falling and reports peak device memory where the backend can report it. It downloads nothing and writes nothing.

RunnableAll tracks

verify-training-step.py

"""Prove that a real training step runs on this machine's accelerator, and say what it cost.

Purpose: the per-track smoke test for a training environment. Picks the device the way
         every later script does, runs a short optimisation of a small stack of linear
         layers in the chosen precision, prints the loss falling, and reports peak
         device memory where the backend can report it.
Platform: all (CUDA on Tracks S and N, ROCm on Track X, MPS on Track M, or the CPU)
Minimum memory: 8 GB
Assumes: torch installed in the active environment. Nothing is downloaded and nothing
         is written to disk.

Usage: python verify-training-step.py [--steps 50] [--precision auto|bf16|fp32] [--size 2048]
"""
from __future__ import annotations

import argparse
import contextlib
import time

import torch
from torch import nn


def pick_device() -> torch.device:
    if torch.cuda.is_available():
        return torch.device("cuda")
    mps = getattr(torch.backends, "mps", None)
    if mps is not None and mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def describe(device: torch.device) -> str:
    if device.type == "cuda":
        # The ROCm build of PyTorch also reports AMD GPUs through the cuda device name.
        return f"cuda / {torch.cuda.get_device_name(0)} (torch {torch.__version__})"
    if device.type == "mps":
        return f"mps / Apple silicon GPU (torch {torch.__version__})"
    return f"cpu (torch {torch.__version__})"


def bf16_supported(device: torch.device) -> bool:
    if device.type == "cuda":
        return torch.cuda.is_bf16_supported()
    return False


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--steps", type=int, default=50)
    parser.add_argument("--size", type=int, default=2048, help="width of each linear layer")
    parser.add_argument("--layers", type=int, default=8)
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument("--precision", choices=["auto", "bf16", "fp32"], default="auto")
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    torch.manual_seed(args.seed)
    device = pick_device()
    print(f"device: {describe(device)}")

    want_bf16 = args.precision == "bf16" or (args.precision == "auto" and bf16_supported(device))
    if want_bf16 and not bf16_supported(device) and device.type != "cpu":
        print("note: bfloat16 requested but this device does not report support for it; continuing anyway")
    print(f"precision: {'bfloat16 autocast' if want_bf16 else 'float32'}")

    layers: list[nn.Module] = []
    for _ in range(args.layers):
        layers += [nn.Linear(args.size, args.size), nn.GELU()]
    model = nn.Sequential(*layers, nn.Linear(args.size, 1)).to(device)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"parameters: {n_params:,}")
    print(f"weights at float32: {n_params * 4 / 1e9:.3f} GB   "
          f"gradients: {n_params * 4 / 1e9:.3f} GB   Adam states: {n_params * 8 / 1e9:.3f} GB")

    optimiser = torch.optim.AdamW(model.parameters(), lr=1e-3)
    inputs = torch.randn(args.batch_size, args.size, device=device)
    targets = torch.randn(args.batch_size, 1, device=device)
    loss_fn = nn.MSELoss()

    if device.type == "cuda":
        torch.cuda.reset_peak_memory_stats()

    # torch.amp documents autocast for device types including cuda and cpu; rather than assume
    # anything about the others, the script only enters autocast when it actually wants bfloat16.
    def precision_context():
        if not want_bf16:
            return contextlib.nullcontext()
        return torch.autocast(device_type=device.type, dtype=torch.bfloat16)

    first = last = None
    started = time.time()
    for step in range(1, args.steps + 1):
        with precision_context():
            loss = loss_fn(model(inputs), targets)   # forward
        optimiser.zero_grad(set_to_none=True)
        loss.backward()                              # backward
        optimiser.step()                             # step
        value = loss.item()
        first = value if first is None else first
        last = value
        if step == 1 or step % 10 == 0:
            print(f"step {step:3d}  loss {value:.6f}")
    if device.type == "cuda":
        torch.cuda.synchronize()
    elapsed = time.time() - started

    print(f"loss went from {first:.6f} to {last:.6f} in {args.steps} steps ({elapsed:.1f} s)")
    if device.type == "cuda":
        print(f"peak device memory: {torch.cuda.max_memory_allocated() / 1e9:.3f} GB allocated, "
              f"{torch.cuda.max_memory_reserved() / 1e9:.3f} GB reserved")
    else:
        print("peak device memory: not reported by this backend; watch the system memory monitor instead")

    if last is not None and first is not None and last >= first:
        raise SystemExit("the loss did not fall: something is wrong with this environment, not with the model")
    print("training step verified on this machine")


if __name__ == "__main__":
    main()

Download verify-training-step.py121 lines

RunnableAll tracks

one real training step, on your accelerator

python verify-training-step.py --steps 50

Output — what you should see

device: cuda / NVIDIA GeForce RTX xxxx (torch 2.x.x)
precision: bfloat16 autocast
parameters: 33,562,625
weights at float32: 0.134 GB   gradients: 0.134 GB   Adam states: 0.268 GB
step   1  loss 1.xxxxxx
step  10  loss 0.xxxxxx
...
loss went from 1.xxxxxx to 0.xxxxxx in 50 steps (x.x s)
peak device memory: x.xxx GB allocated, x.xxx GB reserved
training step verified on this machine

The digits marked x vary; the shape is what matters. Four things in that output are worth reading rather than skimming.

The device line. If it says cpu on a machine with an accelerator, stop here and fix that first. Everything after this lesson assumes it does not.

The precision line. The script asks for bfloat16 autocast only where the device reports support for it, and says so either way. On Track X and Track M you may see float32, and that is the honest default rather than a failure.

The three memory figures. They are the arithmetic from the next lesson, on a model small enough to check by hand: weights at four bytes per parameter, gradients the same, and Adam’s two states at eight bytes between them. Three times the model, before a single activation.

The peak memory line. CUDA reports it; the other backends do not, and the script says so rather than printing a zero. When you get to a run that does not fit, this is the number you will wish you had recorded from a run that did.

=Expected result

The loss must fall. The script exits with an error if it does not, because a loss that stays flat on a fifty-step optimisation of random data is a broken environment rather than a hard problem. If it errors, the fault is below your code: an accelerator build that half-works, or a precision the device does not really support. Rerun with --precision fp32 to find out which.

What to record

Add a section to the lab notebook from Part 1 now, because the next lesson’s run-log format assumes it exists:

• the track, the accelerator the verification script named, and the precision it chose;
• the versions of torch, transformers, trl, peft and datasets that the version command printed;
• whether you are working inside a container, and which image tag if so;
• for Track X, the ROCm version you installed and the date you read the compatibility matrix;
• the peak memory and wall-clock of the verification run, as a baseline you can compare against when something later is unexpectedly slow.

Make the smoke test representative of training

A useful environment smoke test includes allocation, forward computation, backward computation and an optimiser update on the intended device. Inference alone cannot exercise every operation needed for training. Test the dtype you plan to use, rather than checking float32 and assuming a lower-precision path behaves identically.

Before the test, stop competing inference services and record available memory. Confirm Python’s executable and package paths inside the actual execution environment. A notebook kernel can retain an older environment after a package change; restart it and print versions again before trusting its results.

Keep the smoke output with the environment lock or container digest. It should identify the device, dtype, loss behaviour and memory observation. The next checkpoint is a small real model using the dataset format and adapter method required by your lab. If that succeeds, scale sequence length and batch separately. This sequence prevents a large model download or a long training run from becoming your first diagnostic of whether the stack works at all.

Recap

A training environment is four layers: driver, a matching PyTorch build, the training libraries, and a proven training step. The libraries are the same everywhere; the PyTorch build is not. Track S uses NVIDIA’s multi-arch NGC container and installs the training libraries inside it, never PyTorch. Track X needs ROCm rather than Vulkan, and the honest status is that the ROCm 10.0.0 compatibility matrix dated 2026-08-14 lists gfx1151 without a support-tier qualifier while the PyTorch install page read the same day does not mention the chip at all. Track M trains with mlx-lm and keeps PyTorch MPS for reading PyTorch code, because containers on macOS cannot reach Metal. Track N is the straightforward case, with VRAM as the ceiling. Whichever track you are on, the verification script’s four output lines, device, precision, memory arithmetic and peak usage, are what turn “it seems to be installed” into “it took fifty steps and here is what they cost”.

Check your understanding

Question 1. On a DGX Spark, why should you not pip install torch inside the NGC PyTorch container?

• Because pip is not available in the container
• Because the container already ships a PyTorch built for this hardware, and replacing it is the usual way to break a working Spark
• Because the container has no network access
• Because torch cannot be installed on aarch64

Show the answer and why

Answer: Because the container already ships a PyTorch built for this hardware, and replacing it is the usual way to break a working Spark

The catalog page states that the container's software stack has been validated for compatibility and does not require additional installation. Add the training libraries beside it; leave PyTorch alone.

Question 2. What is the accurate statement about ROCm support for the Ryzen AI Max+ 395 as of 2026-09-09?

• AMD lists it as a preview tier
• The ROCm 10.0.0 compatibility matrix lists gfx1151 with no support-tier qualifier, while the PyTorch install page read the same day does not mention the chip
• AMD states that it is unsupported
• There is no published information

Show the answer and why

Answer: The ROCm 10.0.0 compatibility matrix lists gfx1151 with no support-tier qualifier, while the PyTorch install page read the same day does not mention the chip

Two AMD pages read on the same day disagree in what they cover. The course reports both, with their dates, rather than inventing a tier label that neither page uses.

Question 3. The verification script prints "device: cpu" on a machine with a GPU. What has almost certainly happened?

• The script is broken
• A CPU-only PyTorch wheel is installed, or the driver is not visible; either way the accelerator layer is the problem, not the code
• The model is too small for the GPU
• bfloat16 is unsupported

Show the answer and why

Answer: A CPU-only PyTorch wheel is installed, or the driver is not visible; either way the accelerator layer is the problem, not the code

A CPU-only wheel imports and trains perfectly, just slowly, and says nothing about it. That is exactly why the script prints its device first and why you check it before starting anything long.

Question 4. Which of these are reasons the course prefers scripts over notebooks for recorded runs? Select all that apply.

• A notebook's execution order cannot be recovered from the file
• Notebooks cannot use a GPU
• A run that cannot be reconstructed is not a measurement
• Arguments on a command line are a record of what was run

Show the answer and why

Answer: A notebook's execution order cannot be recovered from the file, A run that cannot be reconstructed is not a measurement, Arguments on a command line are a record of what was run

Notebooks use accelerators perfectly well and are the right tool for exploring. The objection is to recording a result from a session whose order and state are not recoverable.

Question 5. Why is a container not an escape hatch for training on Track M?

• Docker will not install on macOS
• Containers on macOS run in a virtual machine that does not expose Metal, so PyTorch inside one is CPU-only
• MLX cannot be installed in a container
• Apple silicon cannot run containers

Show the answer and why

Answer: Containers on macOS run in a virtual machine that does not expose Metal, so PyTorch inside one is CPU-only

The GPU is unreachable from the container's virtual machine, so the container gives you a reproducible CPU environment and nothing else. It is one of the reasons the Mac training path in this course is MLX.

Check my answersClear