Distributed training: a GPT on four GPUs

Intro

0:00 Hello and welcome back to episode twelve, where we're going to train a model about six times bigger than our laptop's, on four GPUs we rent by the hour. Let's begin with a quick look at where we're going to be by the end of the episode.

0:15 By the end of today, our model has about ninety-one million parameters, pretrained for thirty minutes on four GPUs that we rent by the hour, and it writes stories like this one, from the prompt once upon a time.

0:28 On the website of the series, you can find more information about the project. The code is open source on GitHub, tagged at each step of each episode, so that you can follow along, and the docs get a page on training on several GPUs today. And each episode has a page here as well, with its chapters and its code.

0:48 Now, before we get to today's work, let's quickly go over what we did last time. We wrote a byte-level BPE tokenizer in C++, as a Python module, and checked it against a version in plain Python, merge for merge and token for token. Any text now encodes with no unknown tokens, and decodes back as it was. We pretrained our model on its tokens, with the rotary embeddings from episode ten, for ten minutes on my laptop, and measured it in bits per character, which compares models across tokenizers. That model has about fourteen million parameters, trained for ten minutes on my laptop. So today, we're going to rent some GPUs, to train a bigger model on more of the stories.

1:29 Today, we're going to train a model in the shape of GPT-2 small, twelve blocks of twelve heads, with a context of five hundred and twelve tokens. We train it with data parallelism, where each GPU trains a copy of the model, and the copies average their gradients after each step. We write it with PyTorch's DistributedDataParallel, and try it on the laptop before we pay for anything. A bigger model can learn from more stories, so we tokenize all of TinyStories, on each of a machine's cores. Then we rent a machine with four RTX 4090s on Vast.ai, measure how training speeds up from one GPU to two and four, and pretrain for thirty minutes. Last, we compare the new model with the laptop's, in bits per character and in its stories, and save it for Hugging Face. All right, let's get started with how PyTorch trains a model on several GPUs at once.

Data parallelism

2:21 The paper is PyTorch Distributed: Experiences on Accelerating Data Parallel Training, from 2020, by Shen Li and others at Facebook, who built PyTorch's distributed data parallel module. Its abstract describes the technique: each GPU holds a copy of the model and computes gradients on its own data, then the GPUs exchange those gradients at each step, so that the copies stay consistent. Configured well, it scaled almost linearly, up to two hundred and fifty-six GPUs. Figure four shows two processes, each with a copy of a small model. As the backward pass goes from the loss back up the layers, the gradients are copied into buckets, and as soon as a bucket is full, an all-reduce, in red, averages it across the processes, while the layers before it are still computing theirs.

3:08 With data parallelism, each GPU holds a whole copy of the model, and trains it on batches of its own, so four GPUs read four times as many stories in each step. After the backward pass, each GPU has the gradients from its own batch, and an all-reduce averages them, so that all four end up holding the average. Then each copy takes the optimiser's step with that average, and as all four started from identical weights, they stay identical, step after step. The whole model has to fit on one GPU, which ours does easily, while far bigger models split the model across several GPUs instead.

3:45 Each GPU computes g i, the gradient of the loss on its own batch, and the loss is an average over the tokens. So the average of the four gradients is the gradient of the four batches together, and each copy takes the step that one GPU would take with a batch four times as large. With ninety-six stories of five hundred and twelve tokens on each GPU, the batch we use for the long run, a step trains on about two hundred thousand tokens, sixteen times a laptop step, which is forty-eight stories of two hundred and fifty-six tokens.

Pretraining with DDP

4:16 We start a new file, pretrain DDP dot py, next to pretrain dot py, as most of what training does stays as it was. It imports argparse, os and time, and null context, a context manager with no effect, for the CPU, where there's no autocast. Then NumPy and torch, torch distributed as dist, which is how the processes send each other tensors, and DistributedDataParallel, which wraps our model, shortened to DDP. SummaryWriter logs the losses and the samples to TensorBoard, like pretraining's. From our own files, it takes the checkpoints and runs folders and save checkpoint, the data folder and load tokenizer, the GPT and its config. From pretrain dot py, it takes the evaluation interval, the learning rate, estimate loss, get batch and the learning rate's schedule, so that both scripts train with one recipe. And write, from sample dot py, for the stories along the way. Setup runs in each process, and reads three numbers from the environment: its rank, its local rank and the world size. Torchrun, PyTorch's launcher, starts a process per GPU, and tells each its rank, which numbers the processes from zero, along with its GPU on this machine and how many processes there are. If there's a CUDA GPU, the device is the one numbered by the local rank, made the default for this process. Init process group joins the processes into a group, with NCCL, NVIDIA's library for moving data between GPUs. Otherwise, the device is the CPU, with gloo, for trying it out on a machine without NVIDIA GPUs, such as my laptop, as gloo is PyTorch's own backend for CPUs. Main takes the minutes to train as its first argument, like pretraining. The model's size gets flags of its own, and the comment above them says the defaults are GPT-2 small's shape, with a smaller vocabulary, about ninety-one million parameters. So twelve layers of twelve heads, an embedding size of seven hundred and sixty-eight, and a context of five hundred and twelve tokens, twice the laptop's. Batch is the number of sequences per GPU per step, forty-eight by default, and dropout defaults to none, as is usual when pretraining on this much data. Full trains on all of TinyStories, which we're going to tokenize next. Compile compiles the model with torch compile, from episode six, and name names the run's folder. Each process calls setup, and rank zero is the main one, which is going to evaluate, log and save. Each rank seeds PyTorch with thirteen thirty-seven plus its rank, so the ranks draw different batches. With one seed for all of them, they'd train on identical batches, and four GPUs would train no faster than one. The tokenizer is last episode's BPE, the config takes the flags, with rotary embeddings, and raw is the model, on the process's device. Model wraps it in DDP, compiled first if we ask for it, and tells DDP which GPU it's on. DDP gives all ranks rank 0's starting weights, and averages the gradients across them after each backward. The training file is all of TinyStories with dash dash full, and the laptop's otherwise. Splits memory-maps the training and held-out tokens, as sixteen-bit numbers, so they stay on disk until a batch reads them, which matters with about a gigabyte of them. The held-out stories are the ones the laptop's models were measured on, so the losses compare. Rank zero prints the parameters, how many GPUs and which, and how many training tokens there are, so each run's log starts by saying what it trains, and on what. The optimiser is AdamW, with pretraining's settings, decaying the weights of the matrices and leaving the biases and norms alone. The run's folder holds its checkpoints, and the writer exists on rank zero alone, so a single process writes the logs. On NVIDIA GPUs, the forward pass runs under autocast in bfloat16, a sixteen-bit format with float32's range and fewer digits, and its comment gives the reason: bfloat16 for the matrix multiplications, half the memory traffic, a lot more throughput. On the CPU, null context leaves it all in float32. Allow TF32 lets the multiplications that stay in float32 use the GPU's tensor cores, at slightly lower precision. Evaluate works on rank zero, with raw, the model underneath DDP's wrapper, as an evaluation has no gradients to average, and the other ranks don't take part. Under autocast, it estimates the losses, writes forty tokens of a story, and logs and prints them, like pretraining does. Then it saves a checkpoint into the run's folder, over the last one, along with how many GPUs trained it. It saves at each evaluation, as a rented machine can stop in the middle of a run. Barrier then waits until all the ranks reach it. The other ranks wait here while rank 0 evaluates, so they all set off again together. The budget is in seconds, the clock starts, and the first evaluation is at step zero. Elapsed is a tensor of one number, as processes exchange tensors. In the loop, rank zero writes its elapsed time into it, and broadcast copies it from rank zero to all the others. Rank 0's clock decides, for all of them: one learning rate, and one last step. With a clock each, a rank could stop a step before the others, and they'd hang, waiting for its gradients. If the budget is up, the loop stops, on all the ranks at once. The learning rate comes from pretraining's schedule, get batch draws a batch, and the forward pass runs under autocast. Then zero grad, backward, the clipped gradients and the step, as in pretraining. The backward pass is where DDP does its work: as each bucket of gradients is ready, it averages it across the GPUs, so when backward returns, each GPU holds the average, and the optimiser's steps match. Each fifty steps, rank zero works out the speed. Tokens per second for all the GPUs together: each did a batch per step. So it's fifty steps, times the GPUs, times the batch, times the context, over the seconds since the last print. It logs the batch's loss and the speed to TensorBoard, and prints them on a line, with the learning rate. At each evaluation step, we evaluate, and last starts again, so the evaluation's time stays out of the speed. After the loop, there's a last evaluation if the run didn't end on one. Rank zero closes the writer, and counts the tokens, the steps times the GPUs times the batch times the context. It prints the steps, the tokens, the minutes and the speed, and on CUDA, the most memory each GPU held. Destroy process group leaves the group, so each process exits cleanly.

Tried on the laptop

10:51 Before renting anything, we try it on the laptop, where gloo runs it on the CPU. Torchrun starts the processes, with nnodes one, a single machine, and nproc per node two, two processes on it. Master addr is the address where they meet, this machine, as torchrun's standalone option hung on my Mac. The model is tiny, two layers of two heads and sixty-four channels, for a quarter of a minute. Rank zero prints the parameters, two processes on the CPU, and the first evaluation, which both processes wait for at the barrier. The steps go by, and the run ends on a last evaluation, so the process group, the broadcast of the clock and the barriers all work.

All of TinyStories

11:34 So far, data dot py downloads the first eighty megabytes of TinyStories, which is plenty for a model of fourteen million parameters. A model six times bigger can learn from all of it, two point two gigabytes of stories, so data dot py gets a dash dash full option. It imports Pool from multiprocessing, Python's way of running a function in several processes at once. Download takes the most bytes to read as an argument now, with eighty megabytes as before, and none to read the whole file. Encode part encodes a piece of the text, from bytes, into an array of sixteen-bit ids. Start worker runs once in each of the pool's processes, and loads the tokenizer there, into a global, so each process loads it once, and not for each piece. Build full starts with the comment: all of TinyStories' training stories, 2.2 GB, for GPUs that can read them, with the merges learned on the first 80 MB. That way, the tokenizer is the laptop model's, and the held-out stories are the ones it was measured on. It encodes the held-out stories with those merges, downloading them if they're missing, into valid BPE dot bin, as before. Then it downloads the whole training file, if it isn't there yet, and reads all of it into memory, as bytes. Reading bytes leaves the decoding to each process, for its own piece, so the main process doesn't decode two gigabytes on its own first. Pieces of about 32 MB, each ending after a story, so they encode as the whole file would, a piece per core. Bounds collects where each piece ends: thirty-two megabytes on, the next end of text, so no story is cut in half. Then the pool, with a process per core, each started by start worker, and imap hands it the pieces, and gives back their ids in order, as they're done. Each piece's ids are appended to the file as raw bytes, and at the end, we print the tokens, the bytes per token, and the seconds it took on how many cores. In main, dash dash full builds the full data and returns.

13:38 A rented machine starts out empty, so a new bash script, remote dot sh, gets it ready from the laptop. Its header comment lists what that means: the code at HEAD, the BPE merges learned here, the C++ tokenizer built, and all of TinyStories, tokenized. We run it with the name of a host in our ssh config, as a rented machine's address changes from one rental to the next. Set dash e u o pipefail stops the script at the first command that fails, and host keeps the name it's given. Git archive packs the files of the commit, and they go over ssh to be unpacked into a folder there. Scp copies the merges, as those merges are what make the rented machine's tokens match the laptop's. Then, on the machine, g plus plus gets installed if it's missing, to build our C++. The image's Python is the system's, and pip won't install into it without break system packages. We pass that flag, as the machine is destroyed after the run, and pip installs NumPy, TensorBoard and pybind11. Setup dot py builds the tokenizer, and data dot py dash dash full makes the data.

Renting four GPUs

14:47 We rent from Vast.ai, a marketplace where data centres and people rent out their GPUs by the hour, through its command line tool, vastai. Search offers asks for machines with four RTX 4090s, a reliability over ninety-eight percent, and a fast download, for the data. Dash o d p h sorts them by the price per hour, with eighty gigabytes of disk, and jq turns the JSON that comes back into a line per offer. The cheapest, in Japan, at about one point six seven dollars an hour, got stuck downloading its image when I tried it before recording. So we take one of the two in Norway, at about two point nine four dollars an hour.

15:29 Create instance rents it, with PyTorch's own image, which comes with CUDA and PyTorch installed, eighty gigabytes of disk, and ssh straight to the machine, labelled mvgpt. It answers with the new instance's id, along with an API key for the instance, which I've hidden.

15:45 A few minutes later, show instances has it running, and I've given its address to ssh under the name vast. On the machine, nvidia smi shows its four RTX 4090s, each with twenty-four gigabytes of memory, all idle. And topo dash m shows how they're connected to each other. P H B and node both mean the path between two GPUs goes over PCIe, and through a bridge on the CPU's side. N V would mean NVLink, NVIDIA's direct link between GPUs, which this machine doesn't have, so averaging the gradients is going to cost more here than on a data centre's training machine.

16:22 Remote dot sh vast gets it ready, and its last step makes the data. The held-out stories come first, about twenty-eight thousand of them, as before. Then all of the training stories, about two million seven hundred thousand of them in two point two gigabytes, become about five hundred and fifty million tokens, at about four point one bytes per token, in about fourteen seconds on the machine's sixty-four cores.

One, two and four GPUs

16:47 Before the long run, we measure how training speeds up, with two minutes on one GPU, then two, then four, all compiled. Torchrun's standalone option runs on a single machine, and nproc per node is how many GPUs to use. The model has ninety-one million parameters, and after its first fifty steps, which include compiling it, one GPU trains at about one hundred and forty thousand tokens a second. Two GPUs do about two hundred and thirty thousand, one point seven times one, and four do about three hundred and ninety thousand, two point nine times one. So four GPUs do about seventy-two percent of four times one GPU's work. Each step, they average about three hundred and sixty-five megabytes of gradients, ninety-one million numbers in thirty-two bits, over PCIe, and with forty-eight stories per GPU, a step's arithmetic is short enough that the averaging takes a noticeable share of the step. In those two minutes, four GPUs reach a held-out loss of about one point nine two, against about two point zero five on one.

17:50 With twice the batch on each GPU, ninety-six stories, each step does twice the arithmetic for one averaging of the gradients. Four GPUs then train at about four hundred and seventy thousand tokens a second, using about nineteen point three gigabytes of each GPU's twenty-four. We keep that batch for the long run, which comes next.

18:10 Side by side, each GPU we add speeds training up, though by less than one GPU's worth, as the averaging over PCIe takes its share of each step. With twice the batch, four GPUs get to about three point five times what one GPU does at the smaller batch.

Thirty minutes on four GPUs

18:25 Then thirty minutes on the four GPUs, at ninety-six stories each. Nohup keeps torchrun going if the connection drops, its output goes to train dot log, and echo writes its process id to train dot pid. Tail dash f follows the log as it grows, until that process exits. It trains at about four hundred and seventy thousand tokens a second, and the held-out loss goes from about nine point two at step zero to about one point nine nine after two hundred and fifty steps. At the end, it has taken forty-one hundred steps, about eight hundred and ten million tokens, which is one point five times the training stories, in thirty point one minutes. With the evaluations, that's about four hundred and fifty thousand tokens a second, and each GPU held about nineteen point three gigabytes at most.

19:14 From step two hundred and fifty, the held-out loss keeps falling all the way, more slowly as it goes, and ends at about one point zero seven. The training loss stays close to it, as the model sees each story about one point five times, so it learns how the stories are written, without memorising them.

19:32 Scp copies the model back to the laptop, then the runs' TensorBoard logs. And vastai destroy instance gives the machine back, so it stops costing money. The session, from renting the machine to destroying it, came to about two point six two dollars.

What thirty minutes bought

19:53 Back on the laptop, bits dot py from last episode measures both models on the first two million characters of the held-out stories: the laptop's model from last episode, and the one from the four GPUs. Both read about four point zero characters per token, with no unknown tokens, and both give each story back as it was. The laptop's model needs about zero point seven six bits per character, and the bigger one about zero point three nine, about fifty-one percent of that. They share the tokenizer, so their losses per token compare too.

20:25 Six times the parameters, twice the context, and thirty minutes on four GPUs took the bits per character from about zero point seven six for the laptop's model to about zero point three nine.

20:36 From one seed, the laptop's model writes this, and the model from the four GPUs writes this. In the laptop's story, Tim pretends he was a fast, his toys go inside to play, and his friend swims in the yard. The bigger model's story stays with Tim cleaning his room, from his mother asking him to the hug she gives him at the end.

On Hugging Face

20:58 To share the model, hub dot py from episode five gets a second architecture to save it as. It imports json, for the model card's example, then Regex and BPE from Hugging Face's tokenizers library, which we're going to build the tokenizer from. From transformers, it takes GPT-NeoX's config and its model for writing text, and bytes to unicode, from GPT-2's tokenizer code, and from our own files, the BPE tokenizer. Transformers and tokenizers came with episode five's requirements, so those don't change. Episode five saved our model as a GPT-2, which learns a table of positions, so a model with rotary embeddings needs another architecture. To neox is a new function, which converts it. With rotary embeddings, the model is a GPT-NeoX: pre-LayerNorm blocks with biases, rotary embeddings on all the channels of each head, in the rotate half form, and GELU computed with erf, as ours is. GPT-NeoX is EleutherAI's architecture, which transformers and Transformers.js both run. Our rotary embeddings turn each channel together with the one half a head further on, the first with the thirty-third, and that's what rotate half means, so the weights carry over as they are. Head size is the number of channels in each head, sixty-four here. The config takes our sizes, the MLP four times as wide, with rotary on all of each head's channels, at a base of ten thousand, like ours. Use parallel residual is false, as GPT-NeoX normally adds attention and the MLP to the residual stream side by side, while our blocks add them one after the other. The biases are on, dropout is off, the head is tied to the embeddings, and end of text and padding get their ids. Ours is our model's state dict, and state maps our weights to GPT-NeoX's names, the token embeddings and the final norm first. Then, in each layer, b and h are the two prefixes, ours and theirs. GPT-NeoX projects q, k and v with one layer, each head's query, key and value side by side. So for the weights and then the biases, we view each projection by head, stack the three along a new dimension, so each head's three sit together, and reshape them into one layer, three times as wide. The other layers get new names, from a list of pairs, theirs then ours: the attention's output, the MLP's two layers, and the two norms, each with its weight and its bias. GPT-NeoX names the MLP's layers dense h to four h and dense four h to h, after their shapes, from the embedding size to four times it and back. Then we make the GPT-NeoX model and load the state into it, with strict false, which lists what's missing instead of failing. The head is tied to embed in, as ours is to embed tokens, so it comes with it. So the assert lets the head's weight be missing, and no other weight. To H F BPE writes our tokenizer as a tokenizer dot json, which transformers and Transformers.js read. A space or a new line can't sit in a vocabulary file as it is, so byte-level BPE uses GPT-2's table, which gives each of the two hundred and fifty-six bytes a printable character. A token is then the string of its bytes' characters. Chars is that table, a dictionary from each byte to its character. Fast is Hugging Face's BPE model, made from that vocabulary, each string with its id, and our merges, as pairs of strings, in the order we learned them. Our chunks, from bpe reference dot py's regex, on characters instead of bytes: anything past ASCII is a letter, as the bytes of a multi-byte character are. Ours splits bytes, while Hugging Face's pre-tokenizer splits text, so a character above ASCII stands in for its bytes. The pre-tokenizer splits the text into those chunks, then maps each chunk's bytes to their characters, and the decoder maps them back. Padding and end of text are added as tokens of their own, and the fast tokenizer is wrapped for transformers, with its padding, start and end tokens. To H F tokenizer now hands a BPE tokenizer to it, and keeps the word-level one as before. Describe writes what the model card says about a model, as the Hub shows the card of the latest push. Size is a phrase for the blocks, the heads and the embedding size, which both kinds of model share. For a model with rotary embeddings, the architecture says how many parameters it has, counted from the model, in millions, so the card stays true for whichever model we push, and that it's saved as a GPT-NeoX. Otherwise, it's saved as a GPT-2, as in episode five, with its table of learned positions, and the card says how many there are. For a BPE tokenizer, the vocabulary is the two hundred and fifty-six bytes, the merges, and padding and end of text, and its limitations say any text can be encoded, and decoded back as it was. A word-level model keeps episode five's wording, about words outside the vocabulary becoming the unknown token, and the spaces that decoding puts around punctuation, and describe returns all of it, with the context window and the tag for the card's metadata. Push takes those, along with the prompt for the card's example, written with json dumps, so quotes and new lines in it stay valid Python. In main, the description mentions GPT-NeoX, and prompt sets the example's prompt, once upon a time by default, as a pretrained model carries on a story. A model with rotary embeddings is saved with to neox, and the others as a GPT-2, and push gets the model's description. The model card's template takes those values: the architecture's tag in its metadata, the architecture and the vocabulary, the example's prompt, the context window, and the limitations, along with a line on the new revision, for the model of ninety-one million parameters, pretrained on four GPUs.

26:37 Saved this way, the folder has GPT-NeoX's config and weights, and the tokenizer's files. The checks load it the way anyone else would, through transformers, and the token ids match ours on about forty-eight thousand tokens, with the logits within about two point five in a hundred thousand of ours. Exported to ONNX for the browser, the model is about three hundred and seventy megabytes, or about ninety-two in eight bits, and on a window of held-out text, both give about PyTorch's loss, one point six one. I've pushed it to the Hub, as the revision ep012 dash pretrain, and it's one of the two models in the chat on the series' website, where it carries on a story from the beginning you give it.

27:19 Last, the docs get a page on training on several GPUs, with what pretrain DDP dot py changes from pretraining, the full data and the rented machine, how training scaled, and the thirty-minute run against the laptop's model. It goes at the end of the navigation, after the BPE tokenizer.

Wrap-up

27:37 So, today we pretrained with data parallelism, where each GPU trains a copy of the model on batches of its own, and DistributedDataParallel averages the gradients during the backward pass. We tokenized all of TinyStories, about five hundred and fifty million tokens, on a rented machine's sixty-four cores. Four RTX 4090s, with no NVLink between them, trained about two point nine times as fast as one, and about three point five times with twice the batch. In thirty minutes, the model of ninety-one million parameters read about eight hundred and ten million tokens, and went from about zero point seven six bits per character for the laptop's model to about zero point three nine, for about two point six two dollars. And it's on Hugging Face, saved as a GPT-NeoX, with a byte-level BPE tokenizer.

28:26 If you rent GPUs for your own training runs, let me know in the comments where you rent them, and how far a few dollars got you. And if you'd like to follow along, subscribe to the channel. I hope you find this useful, and thanks for watching.