Several devices

With several devices, as devices cuda[0 1 2], the program runs once, in one process, and each tensor operation is split across the devices. Nothing in the program changes but the devices line: the same code runs on one device or several, and never asks which device it’s on. On N devices a program computes what it computes on one, up to the order in which floating-point numbers are added; where the compiler can’t keep that guarantee, the program doesn’t compile.

By default this is data parallelism (MGPU-1 and MGPU-2 in the plan): each device has a copy of every model and of the optimizer state, and a part of each batch. The options under devices change how the state is laid out (Options for several devices): shard params gives each device a piece of every param (MGPU-4), and split H splits every value with an axis of size H, such as attention heads, across the devices (MGPU-5).

The controller and the workers

The process that imports the module and runs the program is the controller. Each device has a worker process, which the controller starts when the program first needs the devices, and which runs only the tensor code the compiler made for it, never the program itself. Workers are started fresh, without the main module, as loader workers are (Loading), so the top level of the calling script runs once, and print, save, file writes and Python calls each happen once, in the controller.

The workers form a torch.distributed process group: NCCL with cuda[...], one GPU per worker, and Gloo with cpu[...]. That is the only difference between the two: cpu[...] is for testing, and for running a program meant for several GPUs on a machine without them; it isn’t faster than devices cpu, since the processes share one CPU.

Where functions run

The compiler works out where each fn runs, from what it does:

• In the controller, if it loops over a dataset, prints, saves or loads a checkpoint, writes to a sink, uses Python, takes a dataset or a sink, or calls a function that runs there. Such a function is ordinary Python in the controller; tensors it makes are on the controller’s CPU.
• On the workers, if it uses tensors, models, batches, gradients or optimizer state, and does none of those things. Every worker runs it, each on its part of the data. A call from the controller sends each worker its part of the arguments, and waits for them.
• Anywhere, if it does neither, as a function of ints does. It runs where it’s called.

Device work, value_and_grad, update and changing a model’s field, is on the workers, so it can’t be in a function that runs in the controller (E1803): it goes in a function of its own, as the README’s train_step is, which the training loop calls.

The batch axis

A batch is split across the workers along its batch axis, the first axis of a [b] Digit batch type, and so is every value the types give an axis of that size, such as batch.img, [b D], and x @ m.w1, [b Hd]. Each worker has some of the examples: a batch of n is split as evenly as possible, the first workers taking one more, so 64 over 3 workers is 22, 21 and 21, and 36 is 12 each. Nothing is padded or repeated. Values without the batch axis, such as a model’s params, are whole on every worker.

A function on the workers that is given a value with the batch axis has that parameter split too, so the split follows the data through calls: in the README’s example, forward’s x: [b D] is split because loss passes it batch.img. A function is compiled once, so it can’t take a split value in one call and a whole one in another (E1804). The batch axis stays an axis of its own: grouping it with another, as in (b n), making a value with two batch axes, a batch with a fixed size, or the batch’s size used as a number, which on each worker would be its part’s, is E1802.

Operations across the batch

An operation whose operand is split along the batch axis, but whose result isn’t, combines values across the workers. These have a distributed lowering, which gives what one device would compute, up to the order of the additions:

Each one is an all-reduce: every worker gets the whole result, so a loss or a logged value is the whole batch’s, and an uneven split counts each example once. Anything else that works across the batch axis is E1801: another reduction over it (max, min, std, var, logsumexp, norm, prod, argmax, any, all, reduce with anything but sum), an operation along it (softmax, cumsum, sort, topk, gather, scatter, an extern fn such as layer_norm given it as over=), indexing or slicing it, and causal with it. Operations along other axes, and everything per example, are unchanged.

Gradients and results

value_and_grad on the workers combines the gradients across them: each worker differentiates the value, which is the same on every worker, through its own part of the batch, and the gradients are added up, so they’re the gradient for the whole batch whatever the split. Then every worker applies the same update to its copy of the model, so the copies stay equal. A logs record holds whole values, so writing one to a sink in the controller writes the whole batch’s loss.

A call from the controller gives the first worker’s result, since every worker has the same; a result with the batch axis is every worker’s part, joined in order. Models and optimizer state stay on the workers.

Random draws differ. With Train(seed=…), dropout draws each worker’s mask for its part of the batch, so the masks aren’t one device’s, and results differ from one device’s beyond rounding; with Eval, they don’t. Tying draws to each example’s place in the dataset, so they match, is planned (Q-16).

Placement and communication on hover

With several devices, hovers in the editor show what the compiler planned (MGPU-6). The checker works out the table they show from the placement and the types (Checked::devices, a check::comms::DeviceTable): how each value is laid out, and every collective the compiled program runs, with its operation, its line, how many values it combines and of what dtype.

• A value says how it’s laid out: split across the 3 devices along b, each holding its share of the batch (22/21/21 for a batch of 64), with the example from the file’s batches(n); whole on every device; for a model, a copy on every device: 813,056 bytes (794 KiB) of params and buffers each, or with shard params sharded: each device holds a third of every param of Mlp, 271,024 bytes (265 KiB) of its params; with split H, split along H: each device holds 4 of H = 12; and in a function that runs in the controller, that it’s there. The batch axis’s own name says it’s split.
• A function says where it runs, and on the workers how many bytes it sends per call, its calls included.
• A line that sends anything, hovered anywhere, lists each collective on it, and each call of a function on the workers that sends something:

Sent between the 3 devices on this line, per call:

- all-reduce of the gradients of `Mlp`'s params `w1` and `w2`, which adds up every device's
  part of the batch: 203,264 × f32 = 813,056 bytes (794 KiB)
- inside `loss`: 16 bytes per call

In all: 813,072 bytes (794 KiB).

That is the README’s value_and_grad line on cuda[0 1 2]: the gradients are D·Hd + Hd·C values, and loss’s cross_entropy all-reduces a sum and a count, 8 bytes, forward and backward. The collectives are each lowering’s: the all-reduce of an operation across the batch (mean sends its result’s values and a count, cross_entropy 2 values, batch_norm 2·C + 1 in f64), value_and_grad‘s all-reduce of the whole params’ gradients, one per dtype; with shard params an all-gather of each param where it’s read, padded to a multiple of N, whose backward is a reduce-scatter, and norm and clip_norm‘s all-reduce of one number; with split H each gather (an all-gather) and each sum (an all-reduce) of the plan. A collective in a function that value_and_grad differentiates, of floats, counts twice, since the backward pass runs one as big. Sizes are numbers when the types give them and symbolic otherwise, with the dims’ values put in: 4·D·b·n = 3,072·b·n bytes. They’re the sizes of the values combined; an all-reduce on a ring sends about 2·(N−1)/N times that from each device.

Python calls on several devices

Python calls run once, in the controller, as they would on one device: a function that uses Python runs there (Where functions run). So a Python call can’t be given data that is split across the devices, a batch or a value with its batch axis, which would mean gathering the batch through the controller’s CPU at every step (E1901). That includes a batch in the controller’s own loop over batches(n), which goes to the devices in parts.

There are two ways round it. A function that prepares examples can run in the dataset, as it loads each example. Or it can be declared with extern fn: an extern fn works on each example by itself, along every axis it isn’t given as an Axis[...] argument, so it runs on the devices, each device’s worker calling it with its own examples:

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

The workers import the module an extern fn names (here augment), from the controller’s sys.path, but never the script the controller runs, whose top level runs once: an extern fn naming a function in that script raises ImportError when the workers load the program. A random extern fn draws from each worker’s own random state, so, as with dropout, its draws differ from one device’s.

Options for several devices

Options under the devices line change how the program’s state is laid out over the devices. They change speed and memory use, never results, beyond the order of the additions:

devices cuda[0 1 2]
    shard params
    split H

shard params (MGPU-4) gives each device a piece of every param, about 1/N of it, and so of the optimizer state, as FSDP does: each param is flattened and cut into N equal pieces, the last padded. When the program reads a param, as in x @ m.w1, it’s gathered from the pieces (an all-gather), and its gradient, each device’s part of a sum, is added up and cut into the same pieces (a reduce-scatter), so each device updates only its piece. norm and clip_norm of a gradient add up the pieces’ squares across the devices. save collects the pieces into the file one device would write, and load cuts them again, so checkpoints move between one device, several, and several with shard params. Reading a field from Python, or a model’s state_dict(), gives it whole.

split H (MGPU-5) splits every value with an axis of size H, a dim, across the devices, each holding H / N of it: a model’s fields, the arguments from Python and every value in between. H must be a dim (E1904) that the number of devices divides (E1902), and one split is allowed (E1903). A value is split along its axis of size H, or along a group axis that starts with H, such as (H Dh), whose part on each device is a block of whole heads; a value with H in two places, or inside a group after another axis, is whole. Batches aren’t split with split H: every device has the whole batch.

The compiler carries the split through each operation:

A value that is computed whole but whose type has H is cut to each device’s heads, which needs no communication: in the README’s attention example, split(b n (3 H Dh) -> b H n Dh) cuts q, k and v to each device’s heads, the scores, masked_fill and softmax run on them, and the last matmul(w_o) multiplies each device’s block of (H Dh) by its rows of w_o and adds the parts up. Gradients follow the same plan backwards: a gathered value’s gradient is added up and cut, and a split param’s gradient is its device’s own. value_and_grad adds up the gradients of whole params, which every device computed from its own heads.

With shard params as well, the params with an axis of size H are split along it, and the others sharded.

Several devices from Python

From Python, a model on several devices is a handle, a shapescript.parallel.RemoteModel: its tensors stay on the workers, and reading a field, as model.w1, copies the first worker’s to the CPU. model.state_dict() copies them all. A handle can’t be changed from Python: a ShapeScript function that takes the model as mut changes it. Mlp.init(seed=…) makes the model on every worker, the same on each; adamw(m, …) and sgd(m, …) of a handle make optimizer state on every worker (RemoteOptimizer). save writes the first worker’s copies, in the same file format, and load makes the models and state on every worker; a checkpoint saved on several devices loads on one, and the other way round.

A function on the workers takes tensors and batches from Python as on one device, on any device: they’re copied to the workers, split as its parameters are. Errors raise in the controller as they would on one device, as a ShapeError with the .ss location, say, with the worker’s traceback as a note.

Speed on several devices

What a step costs on several devices, beyond the devices’ own work:

• Batches are loaded by the workers. A loop over a dataset’s batches in the controller, as the README’s run, doesn’t load them: the dataset is sent to every worker once, when the loop first needs the workers (again only if its length changes, so a dataset is taken not to change while it’s used), and each call sends a batch’s number. Each worker makes the loop’s order from its shuffles’ seeds, loads its own part, padded and split as the controller would split it (Padding on several devices), and checks its items as the loader does. It loads the next batches in threads of its own while the step runs, and on a GPU copies them there from pinned memory on a stream of its own. Before the step the workers agree that each has its part, so an item that doesn’t fit raises the loader’s ShapeError or DomainError in the controller and none of them starts the step; the models are as they were. A dataset defined in the main script can’t be sent, since workers never import it: the controller loads it, with a warning, and sends each worker its part, which costs a copy of every batch through the pipes. Reading a field of such a batch in the controller (it’s a shapescript.parallel.RemoteBatch) loads it there.
• A TensorDataset is kept on each GPU (Datasets kept on the device): each worker gathers its part of each batch from its own copy, the next batch while the controller makes the next call, and, since the whole dataset was checked when it was kept, doesn’t wait for the others to agree that each has its part. This is what DDP’s usual measurement does, with its data on the GPUs, and what makes a large batch as fast as DDP’s.
• A dataset can load a batch at once. If it has __getitems__(indices), PyTorch’s protocol for that, and returns a dict of stacked fields (each [n ...] for n indices), the batch is checked as a whole, which is much faster than item by item; a list of items is checked one by one. This holds on one device too.
• Tensors from Python are sent. A function on the workers called from Python with tensors gets them through the pipes, as their bytes: fine for small values, and costly for big batches (a 6144-example batch of 784 floats, 19 MB, takes about 20 ms a step). Loop over a dataset in ShapeScript instead.
• Steps are recorded on each worker (Recorded steps), with the gradients’ NCCL all-reduce in the graph, so a step is one launch per GPU.
• Each call is one small message per worker, the same for every worker when it can be, pickled once, and an answer from each; the first worker’s carries the step’s logs, which are read on the host every step when a sink writes them.

Measured on 3×A100-40GB with the README example (the plan’s MGPU-7 has the table), its data a TensorDataset: 1.37 ms a step at batch 64 and 1.68 ms at batch 6144, against 2.37 and 2.42 ms for DDP on the same machine, whose data is on the GPUs. With a dataset each worker loads from the CPU, 1.98 ms at 64 and 3.45 ms at 6144: copying 2048 examples per worker from the CPU every step is what costs. DDP’s step time differs between machines of the same kind (from 1.5 to 2.4 ms), more than ShapeScript’s.

When a worker fails

If a worker stops, as when it’s killed or runs out of memory, the controller raises shapescript.WorkerError straight away rather than waiting, naming the worker, how it stopped and what it was running. The other workers are stopped too, and what was on the devices is gone: new models, or a checkpoint, are needed to go on.

A worker waits for the others where they combine values. If they don’t all get there within the collective timeout, the program stops with a WorkerError that says where each worker was, by its .ss location:

WorkerError: a worker failed while running `train_step` (train.ss:32:4), and worker 1 didn't finish:
  worker 0: was waiting for the other workers at train.ss:30:12
  worker 1: didn't answer; it was idle
  worker 2: was waiting for the other workers at train.ss:30:12

The timeout is 600 seconds, or $SHAPESCRIPT_COLLECTIVE_TIMEOUT; shapescript.collective_timeout(seconds) changes it for workers started afterwards. Ctrl-C stops the workers along with the program.