Untied-Ulysses

An efficient implementation of Untied Ulysses — a sequence parallelism method for training LLMs on extremely long contexts (up to 5M+ tokens).

Installation

# From repo root
uv venv && source .venv/bin/activate

cd torchtitan
uv pip install torch==2.9
uv pip install flash-attn==2.8.3 --no-build-isolation --no-cache
uv sync --no-build-isolation --active

cd ../long-context-attention
uv pip install -e .

Optional: FlashAttention 3 (Hopper GPUs)

Required to reproduce exact tokens/second numbers from the paper.

git clone https://github.com/Dao-AILab/flash-attention
cd flash-attention/hopper
uv pip install -U setuptools wheel packaging ninja
MAX_JOBS=128 uv pip install . --no-build-isolation

Usage

Scripts

Scripts to reproduce Table 3 and Figure 4 from the paper are in torchtitan/exps/llama3_8b/ and torchtitan/exps/llama3_70b/.

Quick Start

cd torchtitan/exps/llama3_8b

# Run single experiment
CONTEXT_LENGTH=131072 bash run_single.sh

# Run all experiments
bash run_all.sh

Multinode (8B, 2 nodes)

Run the same command on both nodes simultaneously:

cd torchtitan/exps/llama3_8b

# Run single experiment
NNODES=2 RDZV_ENDPOINT="master-hostname:29500" bash run_single_multinode.sh

# Run all experiments
NNODES=2 RDZV_ENDPOINT="master-hostname:29500" bash run_all_multinode.sh

Configuration

Attention Methods

For Hopper GPUs, use _fa3_ variants instead of _fa2_.