{"id":5647703,"date":"2023-04-06T12:45:16","date_gmt":"2023-04-06T16:45:16","guid":{"rendered":"https:\/\/lightning.ai\/pages\/?p=5647703"},"modified":"2023-06-22T13:30:45","modified_gmt":"2023-06-22T17:30:45","slug":"accelerating-llama-with-fabric-a-comprehensive-guide-to-training-and-fine-tuning-llama","status":"publish","type":"post","link":"https:\/\/lightning.ai\/pages\/community\/tutorial\/accelerating-llama-with-fabric-a-comprehensive-guide-to-training-and-fine-tuning-llama\/","title":{"rendered":"Accelerating LLaMA with Fabric: A Comprehensive Guide to Training and Fine-Tuning LLaMA"},"content":{"rendered":"

Takeaways<\/h3> In this tutorial, we will learn how to train and fine-tune LLaMA (Large Language Model Meta AI). Lit-LLaMA, a rewrite of LLaMA, can run inference on an 8 GB consumer GPU. We will also discover how it utilizes Lightning Fabric to accelerate the PyTorch code. <\/div>\n

What is LLaMA \ud83e\udd99<\/h2>\n

LLaMA is a foundational large language model that has been released by Meta AI.<\/p>\n

LLaMA comes in four size variants: 7B, 13B, 33B, and 65B parameters. The paper shows that training smaller foundation models on large enough tokens is desirable, as it requires less computing power and resources. The 65B parameter models have been trained on 1.4 trillion tokens, while the LLaMA 7B model has been trained on 1 trillion tokens.<\/p>\n

Just a few weeks after the release of LLaMA, the open-source community embraced it by creating an optimized version and expanding its use cases. Now, you can fine-tune LLaMA using LoRA (reduces the number of trainable parameters for fine-tuning) and train a chatbot with Stanford Alpaca<\/a>.<\/p>\n

Lightning AI has also joined the trend by providing an open-source, from-scratch rewrite of LLaMA called Lit-LLaMA<\/a>. The main highlight of Lit-LLaMA is that it is released under the Apache 2.0 license, which makes it easier to adopt for other deep learning projects that use similar permissive licenses and also enables commercial use. It has scripts for optimized training and fine-tuning with LoRA<\/a>.<\/p>\n

Lit-LLaMA: simple, optimized, and completely open-source \ud83d\udd25<\/span><\/h2>\n

\"\"<\/p>\n

 <\/p>\n

Lit-LLaMA is a scratch rewrite of LLaMA that uses Lightning Fabric<\/a> for scaling PyTorch code. It focuses on code readability and optimizations to run on consumer GPUs. As of the time of writing this article, you can run Lit-LLaMA on GPUs with 8 GB of memory \ud83e\udd2f.<\/p>\n

Note: Currently you need to download the official Meta AI LLaMA pre-trained model weights for fine-tuning or running inference.<\/em><\/p>\n

Lit-LLaMA supports training, fine-tuning, and generating inference. Let’s discuss each functionality in detail.<\/p>\n

Training LLaMA<\/h3>\n

Note: We won\u2019t go into too many details about training LLaMA from scratch and instead focus more on fine-tuning and inference because the computational need for training is not available to everyone in the community.<\/em><\/p>\n

The repo comes with a simple and readable LLaMA model implementation<\/a> and a training script<\/a> accelerated by Fabric.<\/p>\n

Large language models may not fit into a single GPU. Fully Sharded Data Parallelism (FSDP) is a technique that shards model parameters, gradients, and optimizer states across data parallel workers. Fabric provides a unified API that makes it easy to use FSDP.<\/p>\n

To use FSDP (Fully-Sharded Data Parallel) with Fabric, create an FSDPStrategy<\/code> object by specifying the auto-wrap policy<\/a> and passing it as an argument to the Fabric<\/code> class.<\/p>\n

Fabric helps to automatically place the model and tensors on the correct devices, enabling distributed training, mixed precision, and the ability to select the number of devices to train on.<\/p>\n

\n\n
import lightning as L\n\n
from lightning.fabric.strategies import FSDPStrategy\n\n
import torch\n\n
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy\n\nfrom lit_llama.model import Block, LLaMA, LLaMAConfig\n\n
def main():\n\n
\u00a0 \u00a0 # \u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f Initialize FSDP strategy \u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\n\n
\u00a0 \u00a0 <\/span>auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={Block})\n\n
\u00a0 \u00a0 <\/span>strategy = FSDPStrategy(auto_wrap_policy=auto_wrap_policy, activation_checkpointing=Block)\n\n
\u00a0 \u00a0 # \u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f Initialize Fabric \u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\n\n
\u00a0 \u00a0 # setting for 4 GPUs with bf16 mixed precision and FSDP distributed training strategy\n\n
\u00a0 \u00a0 <\/span>fabric = L.Fabric(accelerator=\"cuda\", devices=4, precision=\"bf16-mixed\", strategy=strategy)\n\n
\u00a0 \u00a0 <\/span>fabric.launch()\n\n
\u00a0 \u00a0 # Load data\n\n
\u00a0 \u00a0 <\/span>train_data, val_data = load_datasets()\n\n
\u00a0 \u00a0 # Load model configs\n\n
\u00a0 \u00a0 <\/span>config = LLaMAConfig.from_name(\"7B\")\n\n
\u00a0 \u00a0 <\/span>config.block_size = block_size\n\n
\u00a0 \u00a0 <\/span>config.vocab_size = 100\u00a0 <\/span># from prepare_shakespeare.py\n\n
\u00a0 \u00a0 # \u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f initialize model \u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\n\n
\u00a0 \u00a0 <\/span>with fabric.device:\n\n
\u00a0 \u00a0 \u00a0 \u00a0 <\/span>model = LLaMA(config)\n\n
\u00a0 \u00a0 # \u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f Setup model and optimizer for distributed training \u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\u26a1\ufe0f\n\n
\u00a0 \u00a0 <\/span>model = fabric.setup_module(model)\n\n
\u00a0 \u00a0 <\/span>optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(beta1, beta2))\n\n
\u00a0 \u00a0 <\/span>optimizer = fabric.setup_optimizers(optimizer)\n\n
\u00a0 \u00a0 <\/span>train(fabric, model, optimizer, train_data, val_data)\n\n<\/code>