Fine-tuning Llama 2 70B using PyTorch FSDP

Hugging Face has detailed a methodology for fine-tuning the Llama 2 70B model using PyTorch Fully Sharded Data Parallelism (FSDP), leveraging the Hugging Face Transformers, Accelerate, and TRL libraries. This approach enables the training of massive models across multi-node, multi-GPU setups by sharding optimizer states, gradients, and parameters across devices.

Overcoming CPU RAM Bottlenecks during Model Loading

Loading a Llama 2 70B model typically requires significant CPU RAM; if every process on a node loads the model, it can require approximately 2TB of CPU RAM (70B parameters * 4 bytes * 8 GPUs). To prevent out-of-memory (OOM) errors, Hugging Face utilizes a specific initialization strategy implemented in transformers and accelerate:

  1. Meta Device Initialization: The model is created on all ranks using the meta device, meaning it is initialized without weights.
  2. Rank 0 Loading: The state dict is loaded only on rank 0.
  3. Empty Parameter Allocation: All other ranks create empty parameters on the meta device using torch.empty().
  4. State Broadcasting: By setting sync_module_states=True, FSDP broadcasts the weights from rank 0 to all other ranks before training begins.

This method ensures that only one process per node loads the pre-trained model into CPU RAM, drastically reducing the memory footprint during the setup phase.

Efficient Checkpointing with Sharded State Dicts

Saving full intermediate checkpoints using FULL_STATE_DICT with CPU offloading on rank 0 often leads to NCCL Timeout errors and significant delays. To resolve this, Hugging Face recommends using SHARDED_STATE_DICT in the FSDP configuration.

  • Intermediate Checkpoints: SHARDED_STATE_DICT saves shards per GPU separately, allowing for faster saving and resuming of training.
  • Final Model Export: To obtain a standard model state dict for deployment, the state dict type is switched to FULL_STATE_DICT only at the end of training before calling trainer.save_model().

Optimizing VRAM and Training Speed

To reduce compute costs and increase training speed, the implementation employs two primary techniques: Gradient Checkpointing and Flash Attention.

Flash Attention

Standard attention mechanisms are often memory-bound due to redundant High Bandwidth Memory (HBM) reads and writes during elementwise operations (masking, softmax, and dropout). Flash Attention optimizes this via:

  • Kernel Fusion: It keeps intermediate steps in SRAM and writes the final result back to HBM only once.
  • Tiling: The NxN softmax/scores computation is chunked into blocks to fit within SRAM limits, utilizing an online softmax algorithm.
  • Recomputation: The backward pass recomputes necessary values instead of storing the entire NxN matrix from the forward pass, significantly reducing memory consumption.

Gradient Checkpointing

Gradient checkpointing is enabled to further reduce VRAM usage, allowing for larger batch sizes or longer sequence lengths during the fine-tuning of the 70B parameter model.

Implementation and Hardware Specifications

Hardware Configuration

The fine-tuning was performed using the following hardware:

  • Nodes: 2 nodes (minimum 1 required).
  • GPUs: 8 A100 (80GB) GPUs per node.
  • Interconnects: NVLink (intra-node) and Elastic Fabric Adapter (inter-node).
  • System RAM: 1TB per node.
  • CPU: 96 cores per node.

Training Execution

The training process utilized the accelerate launch command with a FULL_SHARD strategy and TRANSFORMER_BASED_WRAP auto wrap policy. Mixed precision training was enabled using bf16. For single-node setups with 8 A100 80GB GPUs, the paged_adamw_32bit optimizer from bitsandbytes is recommended to manage memory.

Fine-tuning was completed in approximately 13.5 hours using the meta-llama/Llama-2-70b-chat-hf model and the smangrul/code-chat-assistant-v1 dataset.

Sources