18: Several devices

E1801

Operation across the batch axis with no lowering for several devices. With several devices each one holds part of the batch, so an operation that combines values across the batch axis needs each device’s part. Sums, means, cross_entropy, and @, matmul, einsum and reduce with sum that contract the batch axis have a distributed lowering; any other reduction over it, an operation along it such as softmax or sort, indexing or slicing it, and causal with it don’t, yet.

devices cuda[0 1 2]

dim D = 4

data Ex:
    x: [D]

fn largest(batch: [b] Ex) -> [D]:
    return max(batch.x, over=b)

Compute it per example, or with one of the operations that have a lowering, or declare one device. See Operations across the batch.

E1802

Batch axis that can’t stay split across the devices. With several devices, a value with the batch axis is split along it, each device holding some of the examples, so the batch axis has to stay an axis of its own, with a generic size that each device’s part gives. Grouping it with another axis, as (b d), a value with the batch axis twice, a batch whose size is fixed, the batch’s size used as a number (on each device, it would be that device’s part’s), and a value with the batch axis passed to a parameter without a shape are errors.

devices cuda[0 1]

dim D = 4

data Ex:
    x: [D]

fn flat(batch: [b] Ex) -> [b*D]:
    return rearrange(batch.x, b d -> (b d))

Keep the batch axis apart, as in b (h w) rather than (b h) w, and use mean rather than a sum divided by b. See The batch axis.

E1803

Device work in a function that runs in the controller. With several devices, a function that loops over a dataset, prints, saves, loads, writes to a sink or uses Python runs in the controller, and so does a function that calls one. value_and_grad, update and changing a model’s field are device work, which runs on the workers, so they can’t be in such a function.

devices cuda[0 1 2]

dim D = 4
dim C = 3

model M:
    param w: [D C] = normal(std=0.02)

data Ex:
    x: [D]
    y: [] Index[C]

fn loss(m: M, batch: [b] Ex) -> []:
    return cross_entropy(batch.x @ m.w, batch.y, over=C)

fn fit(mut m: M, mut opt: AdamW[M], data: Dataset[Ex]):
    for batch in data |> batches(64):
        l, g = value_and_grad(loss, m, batch)
        update(mut m, mut opt, g)

Move the device work into a function of its own, as the README’s train_step is, and call it from the loop. See Where functions run.

E1804

Whole value passed where a function takes a split one. A function on the devices is compiled once. If one call gives it a value with the batch axis, its parameter is split across the devices, and every call must give it one; a value without the batch axis is whole on every device.

devices cuda[0 1]

dim D = 4

model M:
    param w: [D D] = normal(std=0.02)

data Ex:
    x: [D]

fn double(x: [n D]) -> [n D] = x * 2

fn both(m: M, batch: [b] Ex) -> [D D]:
    a = double(batch.x)
    return double(m.w)

Call a function of its own for whole values. See The batch axis.

E1805

Standalone build for several devices. With several devices, the program runs in a controller that starts a worker process for each device, which is the shapescript package’s runtime, so a module built with --standalone, which doesn’t use the package, can’t run on them.

$ cat model.ss
devices cuda[0 1 2]

fn double(x: [n]) -> [n] = x * 2
$ shapescript build --standalone model.ss -o model.py
error[E1805]: a standalone build can't run on 3 devices

Build without --standalone, or for one device. See Building for a device.