The Falcon 180B is a monumental language model pioneered by Abu Dhabi's Technology Innovation Institute (TII). As an extension of the preceding ‘Falcon’ series, particularly Falcon 40B, Falcon 180B amplifies its predecessor's capabilities. This auto-regressive language model employs a refined transformer architecture and was educated using a vast dataset of 3.5 trillion tokens. Boasting a staggering 180 billion parameters, Falcon 180B claims the title of the most expansive open-source Large Language Model in existence, catering to both research and business applications. It proudly occupies the top position on the Hugging Face Leaderboard for open-source LLMs.
Now, let's delve into the process of fine-tuning Falcon 180B. We'll harness the power of DeepSpeed, Hugging Face Transformers, LoRA, and Flash Attention on a set-up with multiple GPUs.
First, What Exactly Is DeepSpeed ZeRO?
Originating from Microsoft Research, ZeRO, which stands for Zero Redundancy Optimizer, is a groundbreaking memory optimization solution tailored for expansive distributed deep learning. It introduces a novel optimizer that drastically curtails the resources demanded by model and data parallelism, while simultaneously amplifying the trainable parameter count. ZeRO has been engineered to eradicate memory overlaps in both data and model parallel training, ensuring minimized communication overhead and maximized computational precision. This ensures a scalable model size that aligns with the number of available devices, all while maintaining peak efficiency.
Incorporated within DeepSpeed, an open-source toolkit from Microsoft Research, ZeRO propels large model training by enhancing scalability, speed, affordability, and user experience. This makes a 100-billion-parameter model training achievable. Designed to complement PyTorch, DeepSpeed provides comprehensive support for all stages of ZeRO, namely stages 1, 2, and 3, and also facilitates CPU and Disk offloading for optimizer states, gradients, and parameters.
What Is LoRA?
LoRA, or Low-Rank Adaptation, is a specialized training method applied in various domains, including the generation of AI-driven art and the optimization of language models. In the realm of AI art, LoRA aids in the fine-tuning of Stable Diffusion models. It allows the system to be trained on distinct concepts, whether they be specific characters or artistic styles, by enacting minimal alterations to the related model file. Notably, models developed using LoRA are much more compact than standard checkpoint models, offering a convenient solution for those limited by storage capacity.
When applied to language models, particularly those as expansive as GPT-3 with its billions of parameters, LoRA offers a refined approach to fine-tuning. Instead of altering the vast number of pre-trained weights, LoRA introduces trainable layers within each transformer block, thereby cutting down on the number of adjustable parameters and the GPU memory they consume. This strategy bypasses the need to compute gradients for the majority of model weights, which expedites the fine-tuning process. With LoRA, models can be fine-tuned using a fraction of the parameters typically required, marking a significant improvement over conventional methods.
In essence, LoRA stands out as a pivotal training approach, beneficial for fine-tuning both AI-generated art models and mammoth language models. It introduces efficiencies by modifying model files with slight changes and providing compact models that are storage-friendly. Furthermore, in the landscape of language models, LoRA enhances the fine-tuning process, making it both faster and more resource-efficient.
What Is Flash Attention?
Flash Attention is an optimization algorithm designed to accelerate the attention mechanism within Transformer-based language models. The attention mechanism, a vital component in Transformers, allows the model to focus on different parts of the input sequence when making predictions. However, this mechanism can be computationally expensive, particularly when processing longer text sequences, leading to high memory costs and slower processing times.
The Flash Attention algorithm addresses this issue through a series of techniques. It utilizes methods such as tiling, which involves dividing the input sequence into smaller segments, and recomputation, which involves recalculating certain values to reduce memory usage. These techniques help to streamline the computations involved in the attention mechanism, thereby enhancing the model's efficiency in handling longer text sequences.
Flash Attention-2 is an improved version of the original algorithm that further optimizes the parallelism and work partitioning of computations. By refining the distribution of tasks and enhancing parallel processing capabilities, Flash Attention-2 achieves a notable 2x speedup compared to its predecessor. This enhancement is especially significant as it enables the algorithm to reach a high processing speed of 230 TFLOPS/s (tera floating-point operations per second) on A100 GPUs, indicating a substantial improvement in computational performance and efficiency.
Steps for Fine-Tuning Falcon 180B with DeepSpeed ZeRO, LoRA & Flash Attention on E2E Cloud
1. We have to first make sure that we have accepted the license tiiuae/falcon-180B to be able to use it. You can accept the license by clicking on the Agree and Access Repository button on the model page at https://huggingface.co/tiiuae/falcon-180B.
2. Launch a GPU node on the E2E Cloud Computing platform.
For this article, we will be using the NVIDIA-A-100 node.
3. Set up the environment.
When you execute this command, Conda will set up a new isolated environment named ‘hf’ with Python 3.10 installed within it. This will allow you to work on projects that specifically require this environment without affecting other Python projects or environments on your system. Additionally, the use of the ‘conda-forge’ channel ensures that packages from this community-driven repository are used for the environment.
Let’s log in into our Hugging Face account by running the following command:
The purpose of this command is to authenticate the user's identity on the Hugging Face Hub and enable them to perform various actions on the Hub.
4. Load and pre-process the dataset.
We'll work with ‘Dolly’, an open-source dataset containing instruction-following records produced by numerous Databricks employees. The dataset falls into various behavioral categories detailed in the InstructGPT paper, such as brainstorming, classification, closed QA, generation, information extraction, open QA, and summarization.
We use the load_dataset() method from the Hugging Face Datasets library.
To fine-tune our model effectively, we are required to transform our organized examples into a set of tasks specified by instructions. We establish a formatting function designed to accept a sample and produce a string containing our formatted instruction.
Let's test our formatting function on a random example.
Also, apart from formatting our samples, we aim to consolidate multiple samples into one sequence to enhance the efficiency of our training.
We create several auxiliary functions to bundle our samples into sequences of a specified length, and subsequently tokenize them.
Once we have processed the datasets, our intention is to store them on disk so that we can utilize the processed dataset later for training purposes.
5. Fine-tune Falcon 180B using DeepSpeed, Hugging Face Transformers, and LoRA with Flash Attention.
The Hugging Face Transformers Trainer seamlessly incorporates DeepSpeed ZeRO, allowing for easy utilization with the provision of a DeepSpeed config file, with the Trainer managing the process. Two DeepSpeed configurations, including CPU offloading, were generated for our experiments.
ds_falcon_180b_z3.json
ds_falcon_180b_z3_offload.json
Alongside the DeepSpeed configuration, we require a training script that incorporates LoRA and integrates flash-attention into our model via the falcon_patch.py utilities. Our run_ds_lora.py script has been designed for this purpose, incorporating the implementation of LoRA with the assistance of peft_utils.py.
run_ds_lora.py
peft_utils.py
falcon_patch.py
Upon confirming the appropriate configuration and training script, we can initiate the training process using torchrun.
Please note that due to our usage of LoRA, we are solely saving the ‘trained’ adapter weights to conserve storage space. If you intend to consolidate the adapters back into the base model and preserve the merged model, you can either include ‘--merge_adapters True’ or employ the merge_adapter_weights.py script.
merge_adapter_weights.py
6. Inference on Falcon 180B after fine-tuning with LoRA.
Here is an example code snippet that shows how to perform inference on Falcon 180B after fine-tuning with LoRA:
Conclusion
In the blog post, we provided a detailed account of our process for refining the Falcon 180B model, utilizing a combination of DeepSpeed, Hugging Face Transformers, and LoRA with Flash Attention on a multi-GPU set-up. The primary components of our approach included:
1. Exploiting DeepSpeed ZeRO, specifically stage 3 (ZeRO-Infinity), to optimize memory usage and enable the training of models with trillions of parameters within constrained GPU memory.
2. Employing Hugging Face Transformers and Datasets for seamless loading and preparation of the text dataset, facilitated by the user-friendly Trainer API.
3. Implementing LoRA, a methodology that efficiently fine-tunes large LLMs by updating a small subset of parameters during each iteration, thereby significantly reducing memory consumption and computational expenses.
4. Incorporating Flash Attention, a highly optimized attention implementation that further reduces the memory footprint.
5. Conducting inference on the model after training the adapter weights.
By combining these methodologies, we successfully conducted fine-tuning on LLMs with over 100B+ parameters, effectively managing resource constraints. This example serves as a comprehensive template for efficiently fine-tuning the largest publicly accessible models.
