🚂 Loooooooong Sequence Lengths

Optimizations for extremely long sequence lengths in Megatron-DeepSpeed as part of the DeepSpeed4Science project.

Sam Foreman 2024-02-12

• DeepSpeed4Science (09/2023)
  • New Features
  • New optimizations
  • Initial Results
• Installation
  • Using install.sh
  • Step-by-Step
• Running
• ZeRO Offloading

Figure 1: This work was done as part of the DeepSpeed4Science project, in collaboration with Microsoft.

The new Megatron-DeepSpeed release contains a variety of improvements / optimizations to enable pre-training Transformer based architectures with significantly longer sequences than was previously possible.

DeepSpeed4Science (09/2023)

New Features

• Enabled Megatron-LM’s sequence parallel.

• Enabled rotary positional embedding.

• Enabled FlashAttention v1 and v2.

• Enabled new fused kernels from NVIDIA.

New optimizations

• Enabled attention map memory optimization, where we first generated attention mask on CPU memory and then moved it into GPU memory to avoid out-of-memory errors when training with very large sequence lengths.

• Position embedding partitioning, where we split weights of position encoding across all GPUs when enabling sequence parallel to further reduce the memory footprint.

Initial Results

gpus = ('32', '64', '128')

colors = {
    'Old Megatron-DS': '#FF5252',
    'Megatron-LM': '#76b900',
    'New Megatron-DS':  '#1A8FFF',
}

data = {
    '25B': {
        'Old Megatron-DS': np.array([36, 42, 42]),
        'Megatron-LM': np.array([26, 48, 52]),
        'New Megatron-DS': np.array([192, 448, 512]),
    },
    '33B': {
        'Old Megatron-DS': np.array([28, 32, 32]),
        'Megatron-LM': np.array([14, 46, 52]),
        'New Megatron-DS': np.array([128, 384, 448]),
    },
}

x = np.arange(len(gpus))
width = 0.25
multiplier = 0

outdir = Path(os.getcwd()).joinpath('assets')
outdir.mkdir(exist_ok=True, parents=True)

improvement = {}
for idx, (model_size, d) in enumerate(data.items()):
    multiplier = 0
    figure, axes = plt.subplots(figsize=(7.5, 4))
    fig = plt.gcf()
    ax = plt.gca()
    for label, value in d.items():
        offset = width * multiplier
        rects = ax.barh(
          x + offset,
          value,
          width,
          label=label,
          color=colors[label],
          alpha=0.8
        )
        ax.bar_label(
          rects,
          padding=3,
          color=colors[label],
          family='monospace',
          weight='bold'
        )
        multiplier += 1
    ax.set_ylabel(
        'GPUs',
        fontsize=18,
        family='sans-serif',
        loc='center',
    )
    ax.set_yticks(x + width, gpus)
    plt.figtext(
        0.005, 0.93, f"{model_size}", fontsize=24, fontweight='bold', ha='left'
    )
    ax.set_xlabel(
        'Sequence Length (k)', fontsize=18, loc='center'
    )
    ax.legend(
        bbox_to_anchor=(0.005, 1.04, 0.99, .098),
        alignment='center',
        edgecolor="#83838320",
        frameon=True,
        ncols=3,
        fontsize=13,
        mode="expand",
        borderaxespad=0.01
    )
    save_figure(fname=f'{model_size}', outdir=outdir)
    _ = plt.show()

Installation

Using install.sh

git clone https://github.com/ramanthanlab/GenSLM/
cd GenSLM/examples/long-sequences/
./install.sh

Step-by-Step

For completeness, we describe below the steps for installing and building each of the dependencies.

1. Clone GitHub repo:

   git clone https://github.com/ramanthanlab/GenSLM

2. Load conda module:

   • ThetaGPU:

     # ThetaGPU:
     if [[ "$(hostname)==theta*" ]]; then
         export MACHINE="ThetaGPU"
         export CONDA_DATE="2023-01-10"
         module load conda/2023-01-11
         conda activate base
     fi

   • Polaris:

     # Polaris:
     if [[ "$(hostname)==x3*" ]]; then
         export MACHINE="Polaris"
         export CONDA_DATE="2023-01-10"
         module load conda/2023-01-10-unstable
         conda activate base
     fi

3. Setup Virtual Environment³:

   cd ./genslm/examples/long-sequences
   # create a new virtual environment
   mkdir -p "venvs/${MACHINE}/${CONDA_DATE}"
   python3 -m venv "venvs/${MACHINE}/${CONDA_DATE}" --system-site-packages
   source "venvs/${MACHINE}/${CONDA_DATE}/bin/activate"

4. Create a new folder (genslm/examples/long-sequences/deps/${MACHINE}) where we’ll installing dependencies locally:

   mkdir -p "deps/${MACHINE}"
   cd "deps/${MACHINE}"

Dependencies

We provide below the details needed to install each of the required dependencies.

Running

The ALCF/ directory contains shell scripts for setting up the environment and specifying the options to be used when launching.

Various options can be specified dynamically at runtime by setting them in your environment, e.g.:

MODEL_SIZE_KEY="GPT25B" SEQ_LEN=128000 USE_FLASH_ATTN=1 MICRO_BATCH=1 GAS=1 SP_TYPE="megatron" ZERO_STAGE=1 ./ALCF/train-gpt3.sh

Explicitly:

• ALCF/train-gpt3.sh: Main entry point for training
  • This script will automatically source the rest of the required ALCF/*.sh scripts below
• ALCF/models.sh: Contains some example model architectures for GPT3-style models
• ALCF/args.sh: Logic for parsing / setting up runtime options for Megatron and DeepSpeed
• ALCF/setup.sh: Locate and activate virtual environment to be used, ensure MPI variables are set properly
• ALCF/launch.sh: Identify available resources and build the command to be executed
  • i.e. figure out how many: {nodes, GPUs per node, GPUs total}, to pass to mpi{run,exec}
  • then, use this to build mpiexec <mpiexec-args> python3 pretrain_gpt.py

ZeRO Offloading

🚀 W&B Report: Looooooooong Sequences

These newly introduced optimizations, in combination with ZeRO-Offload allows us to go even further.

By employing ZeRO-Offloading, we are able to free up additional memory which can be used for even longer sequences.

Though work is still ongoing, this is a promising direction that will allow us to consider significantly larger genomes than previously possible.

We use Weights & Biases to track these experiments, and have aggregated our initial results in the W&B Report below.

We can evaluate the performance of our model by looking at two different metrics for throughput: samples_per_sec and TFLOPS.

Explicitly, we see that we are able to scale up to significantly longer sequences (420k / 128k ~ 3.3x) with only a minimal impact on throughput performance (81 / 105 ~ 77\%)⁴.