Fixed shapes and padding

A training step whose shapes never change can be recorded once, as a CUDA graph, and replayed: one launch per step instead of one per operation (Recorded steps). Shapes change where a size varies: the last batch of an epoch is smaller, sequences have different lengths, and a boolean selection’s size depends on the data. ShapeScript fixes them by padding: the varying axis is padded up to a fixed size, the padded positions are masked out, and every operation across the axis is rewritten so that its result is exactly what it would be without padding, up to the order in which floating-point numbers are added. Where the checker can’t prove an operation exact, the axis isn’t padded and the program runs as it would otherwise. Padding never changes what a program computes.

Padding makes every step the same, which recording needs: with it, an epoch of the README’s training example replays one recorded graph.

What is padded

An axis is padded when its size is one symbol from one of three places:

• The last batch of a loop over batches (FAST-2). In for batch in train |> batches(64):, the loader pads a last batch of 36 to 64 with copies of its own examples, and every batch carries a mask of its real examples, all of them for a full batch. So every step of an epoch has the same shapes and does the same operations, the full batches included. This is on by default, wherever the checker finds it exact; where it doesn’t, the loop has hint H2001, which says why.
• The generic axes of a function that Python calls with varying sizes (FAST-3). shapescript.padded(f) wraps a ShapeScript function so that each call pads each generic axis it can, such as n in x: [n D], up to the next of a few fixed sizes: 1, 2, 3, 4, 5, 7, 9, 12, 15, 19, 24, 30, 38, 48, 60, 80, 104, 136, 176, 224, 280, 352, 440, 552, and so on, each about 25% above the last and a multiple of 8 above 64. Sequences of every length from 1 to 512 make 24 shapes. Its results are cut back to the real sizes, so the caller sees what f gives. shapescript.padded(f, "n") pads only n, and is a ValueError naming the operation if n can’t be padded.
• Boolean selections (FAST-4). x[mask] with a mask of one axis has an unknown size ?k. When every operation on the selection has a masked version, it’s kept at x’s size, with the mask saying which positions are selected, so its shape doesn’t depend on the data. A function that returns it, like the README’s nonempty_rows, gives it back kept so to a caller that can keep it too, and compacts it otherwise: Python always gets the compact value, as it would without padding.

Only these vary. A data field still has a fixed shape (E1302), so variable-length examples come into a program as a generic axis of a function called from Python, or are padded to a fixed maximum by the dataset, with a Mask field. Examples of different lengths in one batch (a mask per example) aren’t supported.

Operations across a padded axis

The checker follows a padded axis through the program by its size, as it follows the batch axis across several devices: every value whose type has an axis of that size is padded along it, including in the functions it’s passed to. Each operation across the axis is made exact:

Anything else across the axis has no masked version, and keeps the axis from being padded: sort, argsort, topk, argmax, argmin, gather, scatter and cat over it, cross_entropy over it, indexing or slicing it, a selection’s size used as a number (which would wait for the device), an as that gives it a fixed size, a pattern that groups it with another axis, print or a Python function given a padded value, an extern fn that isn’t elementwise, a parameter without a shape (unless the function only does elementwise maths), scaled_dot_product_attention’s keys, and, for a selection, whose positions aren’t in order, causal. A loop’s batch, or a function’s axis, is padded only if nothing it reaches is one of these; a selection that reaches one is compacted where it’s made.

Why padding is exact

The padded positions hold copies of real ones: the loader repeats the last batch’s own examples, shapescript.padded repeats a sequence’s positions, and a selection’s unselected positions are copies of its first selected one. So every value computed at a padded position is one a real position could have, as finite as a real one’s. Masking uses where, never a multiplication by the mask, so a padded value never reaches a result, and the gradient that reaches a padded position is exactly zero; since that position’s values are finite, the zero is never multiplied by an infinity, which would make a NaN. (Padding with zeros would break this: log of a zero is -inf, and its gradient times zero is NaN.) A masked reduction adds exact zeros and divides by the number of real positions, so the only difference from the unpadded program is the order of additions.

With Train(seed=…), dropout draws its masks for the padded shape, so the real examples’ masks may not be the ones an unpadded step would draw; with Eval there’s no difference.

When padding is on

Loops over batches and selections are padded wherever the checker proves it exact, and functions when wrapped with shapescript.padded. shapescript.fixed_shapes(False) turns all three off at run time, and shapescript.fixed_shapes(True) back on; it returns what it was. Off, the masked operations get no masks and are the plain ones, so the program runs exactly as it would have been compiled without padding. shapescript.shapes_seen() gives the distinct shapes each loop over batches and each padded function has seen, by the loop’s location or the function’s name, to check that an epoch has one shape. Notebook cells and exported programs aren’t padded.

The generated code shows what was rewritten: a loop over batches is _ss_rt.batch_loop(…, True, where), a function keeps the masks of its padded axes in _ss_pads, from its batches or from its caller (a keyword argument _ss_pad, which Python callers leave out), and a masked operation is a call such as _ss_rt.cross_entropy(z, y, 1, _ss_pads, ((4, 0),)).

Padding on several devices

A padded batch is split across the devices like any other, and its mask with it: 64 padded examples are split 22/21/21, and a device may have only padding. The distributed lowerings of a padded batch add up each device’s masked sum and its number of real examples, so a mean, or the batch’s cross_entropy, is over the real examples of the whole batch.

Recorded steps

A step is a function that changes a model or optimizer state, a mut parameter of either, and does device work itself: value_and_grad, update or changing a field. The README’s train_step is one; run, which calls it in a loop, isn’t, since it does no device work itself. The checker records whether each step can be recorded: it can unless it, or a function it calls, loops over a dataset, prints, saves, loads, writes to a log sink, uses Python or calls a Python extern fn (those for torch.… and the prelude’s are device work), waits for the device with if on a tensor, or makes a size that depends on the data that isn’t kept padded. A step that can’t be recorded has hint H3101 naming what prevents it, and runs operation by operation. Hovering a step’s name says whether it’s recorded, and hovering what keeps it from being recorded says so.

A recordable step is compiled as train_step = _ss_rt.recorded(train_step, "train_step", where). On CUDA, when it’s called:

• The first call runs as it is, on a side stream: the warm-up CUDA graphs need, in which the optimizer makes its state and cuBLAS its workspace. Then the step is recorded for that call’s signature.
• A call with a new signature records the step as a CUDA graph: forward, backward, with several devices the gradients’ all-reduce (NCCL collectives can be recorded), and the optimizer’s update. The signature is the shapes, dtypes and devices of the tensors it’s given (batch fields and their pads included), which model and optimizer state (and their tensors), its host values (int, float, bool, strings), and whether each mode is Train or Eval. So with padding an epoch is one graph; without it, the last batch’s shape is a second.
• Later calls copy their tensors into the graph’s own inputs, replay it, and return copies of its results.

What a step does on the host happens around the replay: the learning rate its optimizer’s schedule gives for its update count is written into the device tensor the update reads (AdamW is made with capturable=True on CUDA), and the count goes up after. Each dropout draw of a recorded step has a generator of its own, which the graph reads its seed and offset from on each replay; before each replay it’s seeded with the seed the draw would have had without recording, from the mode’s seed and path (Random draws). So a replayed step draws different masks at every step, and the very masks it would have drawn without recording: results with and without recording are the same, bit for bit on the RTX 5090 it was measured on. A step that reads a padded batch’s real size as a number is recorded for each real size.

If recording fails at run time, as when a step calls Python through a value the checker can’t see into, the step runs operation by operation from then on, with a RuntimeWarning. A step records at most 64 signatures; further ones run as they are. Recording only happens on CUDA: on the CPU, MPS and CPU workers, steps run as they are.

shapescript.record_steps(False) turns recording off, and shapescript.record_steps(True) back on; it returns what it was. It’s on by default, unless $SHAPESCRIPT_RECORD_STEPS is 0, false, off or no. shapescript.recorded_steps() says, for each step by name, how many graphs it recorded and for which shapes, how many calls replayed, ran as they are, or waited for a recording, and how long recording took; with several devices, the first worker’s.

Recording ahead

A loop over batches(n) knows its epoch’s batch shapes before it loads them, from the dataset’s length, the batch size and padding, and announces them when it starts. A step that has run once is recorded for each announced shape it hasn’t seen, in a thread of its own on a stream of its own, while the steps run, so the first step with that shape replays a graph that’s ready instead of recording one. A step announced before its first call is recorded for those shapes after its warm-up. While a recording is under way in the background, replays go on; a step that would run operation by operation waits for it, since a step run as it is beside a recording can spoil it, and recorded_steps() counts such waits.

With several devices each worker is told its part of each announced shape with the next call, and records it after answering, until the controller’s next message comes: NCCL takes collectives from one thread at a time, so a worker doesn’t record in the background.

In the editor

Hovering a loop’s batches says whether the last batch is padded, and if not, which operation prevents it and where. Hovering that operation says what it keeps from being padded. Hovering a function’s name says which of its generic axes shapescript.padded can pad, and hovering a selection whether it’s kept padded.