Fine-tuning Llama 2 70B using PyTorch FSDP

Introduction

In this blog post, we will look at how to fine-tune Llama 2 70B using PyTorch FSDP and related best practices. We will be leveraging Hugging Face Transformers, Accelerate and TRL. We will also learn how to use Accelerate with SLURM.

Fully Sharded Data Parallelism (FSDP) is a paradigm in which the optimizer states, gradients and parameters are sharded across devices. During the forward pass, each FSDP unit performs an all-gather operation to get the complete weights, computation is performed followed by discarding the shards from other devices. After the forward pass, the loss is computed followed by the backward pass. In the backward pass, each FSDP unit performs an all-gather operation to get the complete weights, with computation performed to get the local gradients. These local gradients are averaged and sharded across the devices via a reduce-scatter operation so that each device can update the parameters of its shard. For more information on what PyTorch FSDP is, please refer to this blog post: Accelerate Large Model Training using PyTorch Fully Sharded Data Parallel.

(Source: link)

Hardware Used

Number of nodes: 2. Minimum required is 1.
Number of GPUs per node: 8
GPU type: A100
GPU memory: 80GB
intra-node connection: NVLink
RAM per node: 1TB
CPU cores per node: 96
inter-node connection: Elastic Fabric Adapter

Challenges with fine-tuning LLaMa 70B

We encountered three main challenges when trying to fine-tune LLaMa 70B with FSDP:

Let’s look at how to solve the above challenges and fine-tune a 70B model!

Before we get started, here's all the required resources to reproduce our results:

Pre-requisites

Fine-Tuning

Addressing Challenge 1

PRs huggingface/transformers#25107 and huggingface/accelerate#1777 solve the first challenge and requires no code changes from user side. It does the following:

Addressing Challenge 2

Let’s create the accelerate config via below command:

accelerate config --config_file "fsdp_config.yaml"

Addressing Challenge 3

(Source: link)

(Source: link)