https://www.cerebras.net/blog/introducing-gigagpt-gpt-3-sized-models-in-565-lines-of-code logo logo * Product + Cloud + Cluster o Condor Galaxy 1 AI Supercomputer o Andromeda AI Supercomputer + System o CS-2 Virtual Product Tour + Processor + Software * Applications + Large Language Models + Computer Vision + High-Performance Computing + Industries o Health & Pharma o Energy o Government o Scientific Computing o Financial Services o Web and Social Media * Resources + Customer Spotlights + Blog + Publications + Events + White Papers + Model Lab * Developers + Community + Documentation + Cerebras Model Zoo + Request Access to SDK * Company + About + In the News + Awards + Press Releases + Press Kit * Join Us + Life at Cerebras + All openings * Contact Us * * * * * * Search December 11, 2023 In Blog Introducing gigaGPT: GPT-3 sized models in 565 lines of code GigaGPT is Cerebras' implementation of Andrei Karpathy's nanoGPT - the simplest and most compact code base to train and fine-tune GPT models. [gigaGPT-feature-uai-258x260] [model-size-img-uai-258x201] GigaGPT is Cerebras' implementation of Andrei Karpathy's nanoGPT - the simplest and most compact code base to train and fine-tune GPT models. Whereas nanoGPT can train models in the 100M parameter range, gigaGPT trains models well over 100B parameters. We do this without introducing additional code or relying on third party frameworks - the entire repo is just 565 lines of code. Instead gigaGPT utilizes the large memory and compute capacity of Cerebras hardware to enable large scale training on vanilla torch.nn code. With no modifications, gigaGPT supports long context lengths and works with a variety of optimizers. Why gigaGPT While the transformer architecture is simple, training a large transformer on a large number of GPUs is hard. Beyond a few billion parameters, vanilla GPT models run out of memory on even the latest GPUs. Training larger models requires breaking up models into smaller pieces, distributing them to multiple GPUs, coordinating the workload among the workers, and assembling the results. This is typically done via LLM scaling frameworks such as Megatron, DeepSpeed, NeoX, Fairscale, and Mosaic Foundry. Though powerful, these frameworks introduce significant complexity. [model-size-img-02-uai-258x153] A small model such as nanoGPT requires just 639 lines of code to implement. To implement a 20B parameter model using Nvidia Megatron model requires 20,507 lines of code - a 32x increase in complexity. Even though this code doesn't need to be written from scratch, implementing, debugging, and maintaining such a project is a major undertaking. Many ML teams struggle to get these frameworks to work and few manage to converge models with decent utilization. gigaGPT shows that on Cerebras hardware you can have the best of both worlds - a compact, hackable codebase and the ability to train GPT-3 sized models with long context. The Models [model-size-table-02-uai-258x58] gigaGPT implements the basic GPT-2 architecture in a way that matches nanoGPT. In particular, we use learned position embeddings, standard attention, and biases throughout the model. These choices were made primarily to stick closely to nanoGPT and can easily be changed. We validate gigaGPT by training four models with 111M, 13B, 70B, and 175B parameters. All the models tested use the OpenWebText dataset using the GPT-2 tokenizer with preprocessing code taken from nanoGPT. As the goal of this project was to create a clean, performant, and usable code base for others to use rather than to train state of the art models ourselves, our validation was geared towards functional correctness rather than convergence, downstream performance, or other similar metrics. To the best of our knowledge, this is the only GPT model that scales from millions to hundreds of billions of parameters without specialized parallelization techniques. gigaGPT-111M [model-size-gigagpt-111-02-uai-258x1] The 111M config was inspired by Cerebras-GPT. It uses the same model dimensions, learning rate, batch size, and training schedule and differ from Cerebras-GPT primarily in their dataset. The loss trend looks good and roughly matches what we observed for the corresponding Cerebras-GPT configuration despite the different choice of dataset. While we were writing gigaGPT, we performed thorough side-by-side numerical checks against a trusted reference code base, so seeing this loss trend isn't surprising, but it's always nice to get the confirmation of a converged model. [model-size-gigagpt-13b-02-uai-258x1] As with the 111M configuration, the 13B configuration also closely matches the model of that size from Cerebras-GPT in model dimensions, learning rate, batch size, and training schedule. Over the first hundred steps we see a few minor loss spikes, but nothing the model can't recover from. A 13B model training run takes a substantial amount of compute so we stop after about 100 training steps. At this point we're well past the scale that nanoGPT can accommodate, but gigaGPT handles it without a problem. [model-size-gigagpt-70b-02-uai-258x1] The 70B configuration is loosely inspired by Llama-2 70B. It takes its model dimensions and from that work and trains for the same 2T parameters as Llama-2 70B with a similar batch size. The learning rate used is slightly more conservative, as previous groups have noticed more instabilities when using learned position embeddings compared to RoPE . Since our goal wasn't to converge the model, we didn't perform rigorous hyperparameter selection. Even so, loss appears to be decreasingly steadily and training is fairly stable. At 70B gigaGPT continues to show great performance and scaling. Even though we grew the model size by two orders of magnitude, and made no effort to optimize for throughput, utilization remains equal or better than previous runs; we just wrote it, ran it, and immediately saw fast results. This model code is also trivial to scale out despite having been written as a single monolithic model - gigaGPT-70B works from 1 to 64 systems such as the Condor Galaxy 1 by just editing a single flag in the configuration file. gigaGPT-175B After validating the 70B model, we were curious to further probe the limits of model scale that gigaGPT could accommodate. We changed the model dimensions of the 70B config to match what was reported in the original GPT-3 paper, scaled the learning rate and initialization using common heuristics, and launched a run. The few steps we trained weren't of much interest from a convergence perspective, but the model ran without any issues at the same utilization as the 70B model. What's most notable here for ML practitioners is that going large does not cause Cerebras hardware to run out of memory. Based on the results we believe gigaGPT can scale to models in excess of 1 trillion parameters. How gigaGPT Works The gigaGPT model does not use any sharding or pipelining techniques because it fits entirely into the system memory of Cerebras hardware. To briefly recap: Cerebras Wafer Scale Clusters are comprised of 1 to 192 Cerebras CS-2 systems supported by accompanying CPU server nodes that store parameters (MemoryX), data and an interconnect (SwarmX). Unlike GPU based clusters, compute and memory are de-coupled. The entire model - upward of trillions of parameters - is stored in a dedicated MemoryX appliance. The model weights are streamed to the wafer one layer at a time during training. By storing models in large, unified memory systems, we obviate the need to break models apart using complex frameworks. All model training from 1 to 192 systems is done using standard mini batches aka data parallelism. gigaGPT is comprised mainly of model.py and train.py. Looking more closely at the model code, we see that it looks quite similar to concise GPT implementations written for GPUs. It is built from `torch.nn` components without use of any fancy external libraries like xFormers or DeepSpeed. In fact, the model code is quite boring. Compared to nanoGPT, we rewrote the attention layer to use primitive torch ops instead of fused attention algorithms and to expose the attention mask as an argument to the attention layer to increase flexibility. Other than that, the differences compared to nanoGPT are mainly cosmetic. Likewise, the main training loop is also very simple. It utilizes `cerebras_pytorch` (a custom PyTorch wrapper specialized for CS System execution) as a drop-in replacement for calls to standard torch APIs. There are only a couple of parts of this file that will look new to someone who is already familiar with PyTorch, in particular the use of a few decorators to section off different sections of functionality as well as the use of a `cerebras_pytorch.backend` scope for model creation. Overall the code is easy to understand, familiar looking to PyTorch users, and easy to modify and customize. Training a large model across a huge cluster requires careful orchestration of multiple independent jobs running across heterogeneous hardware, on its face a very daunting challenge. The `cerebras_pytorch` package is the solution to this problem and is the crux of what allows the gigaGPT code to be so simple. `cerebras_pytorch` wraps some PyTorch functionality that users will already be familiar with and adds a small number of new APIs that help simplify the distributed computing needs of the problem. In this section we will walk through the code of gigaGPT's `train.py` to better illuminate these APIs. We'll start with the end of the `main` function which contains the high-level logic for the training loop and work backwards through the definitions of each of the components it uses. for step, batch in enumerate(executor, start=global_step + 1): if step > config.num_steps: break loss = training_step(batch) log_loss(loss, step) save_checkpoint(step) This is a simple starting point, so we don't need to spend too much time here. Let's first take a closer look at the `executor` used above. dataloader = cstorch.utils.data.DataLoader( get_dataloader, data_path, config.sequence_length, config.batch_size, config.seed, ) executor = cstorch.utils.data.DataExecutor( dataloader, num_steps=config.num_steps - global_step, checkpoint_steps=config.checkpoint_steps, cs_config=cs_config, writer=writer, ) For Cerebras system runs, there are dedicated CPU nodes that handle loading data and feeding it into the model. `cstorch.utils.data.DataLoader` handles the job of defining the dataloader instance that will run on each of these worker nodes. It takes a function that returns a dataloader instance for ease of setting up independent properly sharded dataloaders on each worker node. This `cstorch.utils.data.DataLoader` object is then fed into a `DataExecutor`, which is responsible for the top-level coordination of all the different independent tasks required for the run. Next, let's look more closely at the components involved in defining a single training step. Internally, `cerebras_pytorch` uses PyTorch LTC to trace the compute graph associated with the training job and maps this compute graph down to operations that can be run on the Cerebras Wafer Scale Engine (WSE). Accordingly, the first step towards defining the training logic is to create the model instance in a way that enables it to be traced later. This is accomplished by the following code at the top of `train.py::main`: backend = cstorch.backend(config.backend, use_cs_grad_accum=True) ... with backend.device: model = GPTModel(model_config) compiled_model = cstorch.compile(model, backend) With this model definition along with an optimizer and learning rate scheduler that are created with APIs that directly mirror PyTorch APIs, we are now ready to define the logic of a basic training step. @cstorch.trace def training_step(batch): input_ids, labels = batch loss = compiled_model(input_ids, labels) loss.backward() torch.nn.utils.clip_grad_norm_(list(all_params), config.max_gradient_norm) optimizer.step() lr_scheduler.step() optimizer.zero_grad() return loss The body of this function is fairly standard training code. The only interesting part is the `@cstorch.trace` decorator. This signals to the framework that the code in the function is intended to be traced and run on CS system. No tensors can be eagerly executed within this scope, which means the code in here can't include any logging functionality or python conditionals. For that, we need a different decorator: @cstorch.step_closure def log_loss(loss, step): rate = executor.profiler.rate() global_rate = executor.profiler.global_rate() logger.info( f"| Step={step}, " f"Loss={loss.item():.5f}, " f"Rate={rate:.2f} samples/sec, " f"GlobalRate={global_rate:.2f} samples/sec" ) writer.add_scalar("loss", loss.item(), step) writer.add_scalar("samples_per_second", global_rate, step) This logging code requires eagerly executing tensor values and doesn't need to run on the WSE, so we wrap it in a `@cstorch.trace decorator`. Checkpointing code works similarly, except that we want to make sure that it only runs every `checkpoint_steps` steps for whatever value of `checkpoint_steps` we pass into the `DataExecutor` above. For this we have the `@cstorch.checkpoint_closure` decorator. Functions wrapped in this decorator can be called at any time but will only execute if the current step is a checkpoint step. @cstorch.checkpoint_closure def save_checkpoint(step): checkpoint_path = out_dir.joinpath(f"checkpoint_{step}.mdl") state_dict = { "model": model.state_dict(), "optimizer": optimizer.state_dict(), "lr_scheduler": lr_scheduler.state_dict(), "global_step": step, "model_config": asdict(model_config), } cstorch.save(state_dict, checkpoint_path) logger.info(f"Saved checkpoint to {checkpoint_path}") With that, we're done defining the functions we used in the main training loop we started with. After a couple of lines of code to fill in the gaps of checkpoint loading, config handling, etc, we end up with a `train.py` which in just 156 lines of code is able to seamlessly coordinate training jobs across huge distributed clusters. Conclusion [model-size-gigagpt-conclusion-02-ua] We'd like to thank Andrei Karpathy for creating nanoGPT and inspiring this work. We believe simple, hackable, and performant code is essential for advancing machine learning research. By combining the benefits of a compact code base and the ability to train GPT-3 scale models, gigaGPT on Cerebras hardware represents a significant leap towards more accessible, scalable, and efficient AI model training. If you're working with the Cerebras platform, we encourage you to experiment with gigaGPT and share your feedback. Find at: https://github.com/Cerebras/gigaGPT Author: William Marshall Contributors: James Wang, Gavia Gray 2023 featured --------------------------------------------------------------------- Avatar photo William Marshall Author posts Related Posts Blog [image-14-1-uai-258x172] December 11, 2023 Cerebras @ NeurIPS'23 - gigaGPT, CrystalCoder, Open Source AI Party & More Massive compute and software advancements have enabled Cerebras to become a... --------------------------------------------------------------------- Avatar photoby James Wang Blog [cerebras-hp-carousel-09-uai-258x172] December 8, 2023 Five Reasons to Join Cerebras in 2024 Build a breakthrough AI platform beyond the constraints of the GPU --------------------------------------------------------------------- Avatar photoby James Wang Machine LearningBlogLarge Language ModelNLPDeep Learning [BTML-3B-8K-uai-258x172] December 8, 2023 Fine-Tuning Language Models Using Direct Preference Optimization An Alternative to RLHF to get a human preferred chat model. --------------------------------------------------------------------- Avatar photoby Alexander Vishnevskiy * Prev * Next Explore more ideas in less time. Reduce the cost of curiosity. Subscribe to our newsletter --------------------------------------------------------------------- [cerebras-white-01-uai-258x114] info@cerebras.net 1237 E. Arques Ave Sunnyvale, CA 94085 Follow Product Cluster System Chip Software Cloud Applications Natural Language Processing Computer Vision High Performance Computing Industries Health & Pharma Energy Government Scientific Computing Financial Services Web & Social Media Resources Customer Spotlight Blog Publications Event Replays Whitepapers Developers Community Developer Blog Documentation ML Public Repository Request Access to SDK Company About Cerebras In the News Press Releases Privacy Legal Careers Contact (c) 2023 Cerebras. All rights reserved [ ] Privacy Preference Center Privacy Preferences [Save Preferences] x Manage Cookie Consent To provide the best experiences, we use technologies like cookies to store and/or access device information. Consenting to these technologies will allow us to process data such as browsing behavior or unique IDs on this site. Not consenting or withdrawing consent, may adversely affect certain features and functions. Functional [ ] Functional Always active The technical storage or access is strictly necessary for the legitimate purpose of enabling the use of a specific service explicitly requested by the subscriber or user, or for the sole purpose of carrying out the transmission of a communication over an electronic communications network. Preferences [ ] Preferences The technical storage or access is necessary for the legitimate purpose of storing preferences that are not requested by the subscriber or user. Statistics [ ] Statistics The technical storage or access that is used exclusively for statistical purposes. The technical storage or access that is used exclusively for anonymous statistical purposes. Without a subpoena, voluntary compliance on the part of your Internet Service Provider, or additional records from a third party, information stored or retrieved for this purpose alone cannot usually be used to identify you. Marketing [ ] Marketing The technical storage or access is required to create user profiles to send advertising, or to track the user on a website or across several websites for similar marketing purposes. Manage options Manage services Manage {vendor_count} vendors Read more about these purposes Accept Deny View preferences Save preferences View preferences {title} {title} {title} Manage consent