FSDP Training

A PyTorch implementation of distributed language model training using Fully Sharded Data Parallel (FSDP) for efficient training of Transformer models at scale.

Overview

This project implements a GPT-style Transformer model trained using PyTorch's Fully Sharded Data Parallel (FSDP) framework. FSDP enables training large models by sharding model parameters, gradients, and optimizer states across multiple GPUs, reducing memory requirements while maintaining training efficiency.

Requirements

Install the required dependencies:

pip install -r requirements.txt

Usage

Single Node Training

For single-node multi-GPU training:

torchrun --nproc_per_node=2 main.py \
  --train_steps=1000 \
  --d_model=512 \
  --num_heads=8 \
  --num_layers=6 \
  --learning_rate=1e-4 \
  --batch_size=32 \
  --seq_length=256

Multi-Node Training

For distributed training across multiple nodes using SLURM:

sbatch submit_2gpu.sub

Grid Search

Run hyperparameter grid search:

sbatch sbatch_grid_search.sub

Command Line Arguments

Argument	Default	Description
--train_steps	1000	Number of training steps
--d_model	256	Model dimension
--num_heads	4	Number of attention heads
--num_layers	4	Number of transformer layers
--learning_rate	1e-4	Learning rate
--batch_size	64	Batch size per GPU
--seq_length	256	Sequence length
--vocab_size	50257	Vocabulary size (GPT-2)
--max_len	256	Maximum sequence length
--save_checkpoint_dir	None	Directory to save checkpoints
--load_checkpoint_dir	None	Directory to load checkpoints
--log_train_loss_freq	100	Training loss logging frequency
--log_valid_loss_freq	100	Validation loss logging frequency

SLURM Scripts

Basic 2-GPU Training (submit_2gpu.sub)

• Single node with 2 A100 GPUs
• 20-minute time limit
• Basic training configuration

Grid Search (sbatch_grid_search.sub)

• Multi-node training (2 nodes)
• Hyperparameter sweep over learning rates
• Array job for parallel experiments
• 5-hour time limit

Load/Save Operations

• sbatch_load.sub: Resume training from checkpoint
• sbatch_save.sub: Save model checkpoint

Data

The project uses the C4 dataset (Common Crawl) for training.

Data Path Configuration

Update the dataset paths in get_dataloader() function to point to your local dataset:

# Training data
hf_dataset = load_from_disk("/path/to/your/c4/train")

# Validation data  
hf_dataset = load_from_disk("/path/to/your/c4/validation")

Monitoring

Wandb Integration

The training automatically logs metrics to Weights & Biases:

• Training loss per rank
• Validation loss per rank
• Learning rate schedule
• Final validation loss

Initialize your wandb project:

wandb login
# Update project name in main.py
wandb.init(name="your-run-name", project="your-project")