Several Devices

Every program so far began with devices cpu. This chapter is about the other choices: a GPU, several GPUs, and a GPU rented by the second on Modal, with no change to the code but that line.

One device

devices says where every tensor the program makes lives. There’s no default, and no code below that line mentions a device:

devices cpu          # the CPU
devices mps          # an Apple GPU
devices cuda[0]      # the first NVIDIA GPU

A tensor passed in from Python must already be on that device; it isn’t moved for you, and one on another device raises shapescript.DeviceError, naming the argument and suggesting x.to("cuda:0"). Datasets copy their batches to the device themselves, and logs and print copy values back.

Several devices

Change the line to several devices and the same program trains on all of them, with data parallelism:

devices cuda[0 1 2]

dim D  = 784
dim Hd = 256
dim C  = 10

model Mlp:
    param w1: [D Hd] = normal(std=0.02)
    param w2: [Hd C] = normal(std=0.02)

data Digit:
    img:   [D]
    label: [] Index[C]

fn loss(m: Mlp, batch: [b] Digit) -> [] NonNeg:
    return cross_entropy(batch.img @ m.w1 |> gelu |> matmul(m.w2), batch.label, over=C)

fn step(mut m: Mlp, mut opt: AdamW[Mlp], batch: [b] Digit) -> [] NonNeg:
    l, g = value_and_grad(loss, m, batch)
    update(mut m, mut opt, g)
    return l

fn train(mut m: Mlp, data: Dataset[Digit], epochs: int):
    opt = adamw(m, lr=3e-4)
    for epoch in 0..epochs:
        for batch in data |> shuffle(seed=epoch) |> batches(64):
            print("loss", step(mut m, mut opt, batch))

How it works, and what it promises:

• Your program runs once. One process, the controller, runs it, so print, save and calls into Python happen once. Each device gets a worker process that runs only the compiled tensor operations. You never ask which GPU you’re on.
• Device work and host work are separate functions. A function that loops over a dataset, prints, saves or calls Python runs in the controller; value_and_grad, update and the maths run on the workers. So step above is its own function, which train calls: putting value_and_grad directly in train’s loop is error E1803.
• The results are one device’s, up to the order floating-point numbers are added. A batch of 64 is split 22/21/21, the gradients are combined weighted by each part’s size, and the logged loss is the whole batch’s. batch_norm’s statistics cover the whole batch.
• The cost is visible. In the editor, hovering a line shows how each value is split across the devices and how many bytes each step sends between them: on this example, value_and_grad all-reduces the gradients of w1 and w2, 203,264 f32 values, 813,056 bytes per step.

Operations across the batch that can’t be split exactly, such as a softmax over the batch axis, are compile errors for now.

What works today. All of this has run on three GPUs on Modal, 3×T4 and 3×A100, with NCCL: the program runs once from one process, print prints once, and 100 training steps match a devices cuda[0] run to 2.4e-7 with dropout off. shard params and split H have run there too. Each training step is recorded as a CUDA graph and replayed, on each worker as on one GPU, which the fixed shapes allow (every step has the same shapes, with the last batch padded and masked). On 3×A100 a step of the training example takes 1.37 ms at batch 64, under PyTorch’s DDP on the same machine (2.37 ms), when its data is a TensorDataset that stays on the GPUs; a dataset that loads each batch on the CPU is slower with large batches. Without GPUs, devices cpu[0 1 2] runs a program just as devices cuda[0 1 2] would, with three worker processes on the CPU in place of three GPUs, so you can test a several-GPU program on a laptop. With dropout on, each device draws masks for its own examples, so results differ from one device’s. Coming: training across several machines.

Options

Indented lines under devices change how the work is split. They change speed and memory, never results:

devices cuda[0 1 2]
    shard params    # split parameters and optimizer state across the GPUs, for models too big for one
    split H         # split every parameter axis named H across the GPUs, such as attention heads

dim H = 12

With shard params, each device holds a third of every parameter and of its optimizer state, as FSDP does. With split H, every value with an axis of size H is split along it, so each device computes its own attention heads. H must be a dim that the number of devices divides.

Python on split data

Python code runs once, in the program’s own process, so it can’t run on data that’s split across devices. The checker says so:

devices cuda[0 1 2]

dim D = 784

data Digit:
    img: [D]

import python augment

fn augmented(batch: [b] Digit) -> [b D]:
    return augment.random_crop(batch.img) as [b D]

$ shapescript check split_python.ss
error[E1901]: `batch.img` is split across the 3 devices, but `augment.random_crop` is Python, which runs in the controller
  --> split_python.ss:11:32
   |
11 |     return augment.random_crop(batch.img) as [b D]
   |            ------------------- ^^^^^^^^^ split along its batch axis `b`
   |            |
   |            a Python function
   |
   = note: with 3 devices each one holds part of every batch, and Python calls run once, in the controller; this one would gather the batch through the controller's CPU at every step
   = help: move `random_crop` into the dataset, so it runs on each example as it's loaded, or declare it per example with `extern fn random_crop(...) = "augment.random_crop"`, so each device runs it on its own examples

For more information about E1901, run `shapescript explain E1901`.

An extern fn works on each example independently, so each device’s worker runs it on its own examples:

devices cuda[0 1 2]

dim D = 784

data Digit:
    img: [D]

extern fn random_crop(img: [b D]) -> [b D] = "augment.random_crop"

fn cropped(batch: [b] Digit) -> [b D]:
    return random_crop(batch.img)

Renting a GPU on Modal

No NVIDIA GPU at hand? Modal rents them by the second, and shapescript run --modal runs a whole program there, unchanged:

pip install modal          # into the Python that has shapescript
modal setup                # once: sign in
shapescript run --modal T4 train.py 3

The program’s folder is uploaded, the program runs in a container with that GPU and the same versions of Python, torch and ShapeScript, its output is streamed back as it’s printed, and the files it writes, such as checkpoints and logs, come back when it ends. In the container, devices cuda[0] is the GPU; --modal H100:3 gives three, for devices cuda[0 1 2]. --timeout 10m caps a run, --env KEY passes an environment variable, and --secret NAME a Modal secret. A .ss program is checked here first, so its errors cost nothing. The guide is docs/modal.md.

Exercises

1. You have a laptop and want to test that a program meant for devices cuda[0 1 2] gives the same losses as on one device. What do you change?

2. Why does devices cuda[0 1 2] with --modal T4 fail before anything starts, and what’s the fix?

3. Does this check?

   devices cuda[0 1 2]
       split H
   
   dim H = 8