Keyboard shortcuts

Press ← or → to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Not yet implemented

Training on one device is implemented: control flow, host scalars, models, mut parameters, gradients, modes and dropout, data records and datasets, logs, sinks and print, optimizers, gradient transforms and schedules, checkpoints, the standard library’s fit and the devices declaration. So is data parallelism on several devices, with a controller and a worker per device, with Python calls checked and extern fns run on each device’s examples, and the options shard params and split H; all of it has been run on worker processes on the CPU (devices cpu[0 1 2]) and, with NCCL, on 3×T4 and 3×A100 GPUs on Modal, where it matches one GPU’s results. Placement and communication show on hover. Steps are recorded as CUDA graphs, with new fixed shapes recorded ahead of time, and on each worker of several GPUs. A step on several GPUs is as fast as PyTorch’s DDP when the dataset fits on the GPUs, but not yet with large batches loaded from the CPU (see Speed on several devices). What isn’t done yet, beyond milestone M10 in IMPLEMENTATION_PLAN.md: distributed lowerings of reductions over the batch other than sums and means; several devices in notebooks; and dropout masks that match one device’s. Also not yet: x to cpu and x to device, and random operations other than dropout. The mask constructors tril, triu, eye and band are planned too. The only built-in optimizers are SGD and AdamW, the only built-in gradient transform is clip_norm, and the only schedules are cosine and warmup.