[ROLL]

TL;DR: Let’s have a quick loot at how to involve FSDP2 into ROLL RL framework!

As of version 0.2.0, ROLL has successfully integrated FSDP2 (Fully Sharded Data Parallel 2) as one of its training backends. (Release Note).

In this blog post, I'll explain how a specific training backend is implemented in ROLL for RL training, using FSDP2 as a detailed example to walk through the architecture and design patterns.

Background: from Zero Redundancy Optimizer to FSDP2

Motivation: Zero Redundancy Optimizer

In the "early days" of distributed training, the standard approach was Distributed Data Parallel (DDP). The workflow of DDP could be summarized as:

  1. Replicate: Copy the entire model to each GPU
  2. Distribute: Feed each model replica a unique batch of data
  3. Compute: Calculate loss and perform backward pass on each process
  4. Synchronize: Reduce gradients across all processes (i.e., all model replicas on different GPUs)
  5. Update: Update model parameters with the synchronized gradients

However, as models continued to grow in size, a critical bottleneck emerged: a single GPU could no longer fit an entire model. With models growing, the memory walls of individual GPUs became the barriers.

This raised a natural question: Instead of keeping everything redundant (replicating model weights, optimizer states, and gradients across every GPU) could we shard (split) each of these components uniformly across all data parallel processes (i.e., GPUs)?

This is where the Zero Redundancy Optimizer (ZeRO) comes in. If we have $N$ GPUs, we can shard the model weights $P_p$, optimizer states $P_{os}$, and gradients $P_g$ across all $N$ GPUs, and then gather the necessary components on demand during training. Therefore, instead of each GPU holding the full memory footprint, each GPU only maintains $\frac{1}{N}$ of each component.

https://arxiv.org/abs/1910.02054

https://arxiv.org/abs/1910.02054

Implementation: FSDP1

FSDP is the pytorch-native implementation of ZeRO. The training workflow of FSDP follows the ZeRO-3 strategy, implementing a orchestrated gather-compute-discard pattern for each layer:

Forward Pass:
  For each layer:
    1. All-gather: Collect full parameters from all GPUs
    2. Compute: Execute forward computation
    3. Discard: Free gathered parameters (keep only local shard)

Backward Pass:
  For each layer (reverse order):
    1. All-gather: Collect full parameters from all GPUs
    2. Compute: Calculate gradients
    3. Reduce-scatter: Aggregate and distribute gradients
    4. Discard: Free gathered parameters
    
Optimizer Step:
  Each GPU updates only its local parameter shard

FSDP implements the gather-compute-discard workflow using pytorch hook, allowing transparently inject parameter gathering and freeing without modifying model code.