torch-compile-volta-cp

A V100-friendly torch.compile wrapper + Ulysses/Ring context parallelism for transformer training on NVIDIA Volta (Compute Capability 7.0). Built end-to-end autonomously by the ml-intern Claude Code skill on a 4× Tesla V100-SXM2 32 GB box. Sister project to flash-attn-volta.

Headline: train Qwen2.5-1.5B at 32 K context on 4× V100 (32 GB each) at 4 733 tokens/sec, 24 GB peak per rank — that's 4× the context and +57 % throughput vs the strongest single-GPU baseline (xformers @ seq=8 192 = 3 018 tok/s, OOM-ceiling).

What's inside

torch-compile-volta-cp/
├── tc_volta/
│   ├── __init__.py        # compile / compile_layers / compile_mlp_only — V100-safe torch.compile
│   ├── diagnose.py        # failure-mode probes for every compile mode / dtype
│   └── inductor_ops.py    # registers tc_volta::flash_attn (xformers Cutlass backend on V100)
├── cp_volta/
│   └── __init__.py        # UlyssesAttention + RingAttention + qwen2/qwen3 patches
├── bench/
│   ├── single_gpu.py      # eager / xformers / xformers+compile sweep
│   └── cp_16k.py          # CP fwd+bwd sweep (cp=2/3/4) up to 32 K
├── tests/
│   ├── test_compile_correctness.py   # output + grad parity for every (mode, dtype) cell
│   └── test_cp_grad_parity.py        # cp=1 vs cp=N grad cosine sim
└── BENCH.md  RESEARCH.md  PLAN.md  RESULTS.md  TASK.md  VERIFY.md  TORCH_COMPILE_VOLTA.md

Why does this exist

torch.compile (Inductor + Triton) was designed assuming Ampere+ targets. On a V100 SM 7.0 box:

• reduce-overhead mode produces silently wrong gradients (cos-sim = 0 vs eager) because cudagraph buffer aliasing collides with Volta's stream-ordering semantics.
• Inductor's compiled backward at seq ≥ 4 K uses ~5 GB more memory than eager and goes OOM at seq=8 K.
• flash_attn_volta (our sister project's Triton-2.3 kernel) can't coexist with Inductor on torch 2.0.1 because Inductor needs Triton 2.0 and Triton 2.0's V100 backend has the tt.reduce MMA-layout-unification bug that breaks every FAV shape.

Context parallelism for ≥ 16 K sequences on Volta isn't shipped by anyone — Megatron / DeepSpeed Ulysses / nemo / torch.distributed.tensor.parallel all assume Ampere+ at some layer of the stack.

This repo lifts those gates with the minimal patches you actually need on a Volta + PyTorch-2.0.1 + Triton-2.0 stack: a torch.compile wrapper that auto-applies the V100-safe config knobs, a custom Inductor op that points SDPA at the xformers Cutlass backend (the only attention path that works on V100 today), and Ulysses-style + ring-style context parallelism on top of torch.distributed over NCCL.

Install

git clone https://github.com/AlexWortega/torch-compile-volta-cp
cd torch-compile-volta-cp
pip install --user xformers==0.0.22 transformers==4.46  # versions known to work
# torch 2.0.1+cu117 and triton 2.0.0 should already be installed on a Volta box
python -c "import tc_volta; import cp_volta; print('OK')"

Use — single GPU, torch.compile that actually wins on V100

import torch
from tc_volta import compile as tc_compile, register_flash_attn_volta
from transformers import AutoModelForCausalLM

register_flash_attn_volta()          # routes F.scaled_dot_product_attention through xformers Cutlass

model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B", torch_dtype=torch.float16).cuda()
model = tc_compile(model)            # V100-safe defaults: cudagraphs OFF, dynamo error-suppress ON
# Now model.forward is compiled. Train as usual.

Or, for more granular control:

from tc_volta import compile_layers       # compile each decoder layer in isolation
from tc_volta import compile_mlp_only     # compile only the FFN, keep attention eager

Use — context parallelism (Ulysses) on 2/3/4 GPUs

import os, torch.distributed as dist, torch
from transformers import AutoModelForCausalLM
from cp_volta import patch_for_cp, init_cp_group

dist.init_process_group("nccl")
cp_group = init_cp_group(cp_size=4)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-1.5B",
    torch_dtype=torch.float16,
    attn_implementation="eager",
).cuda(int(os.environ["LOCAL_RANK"]))

patch_for_cp(model, cp_group=cp_group, mode="ulysses")  # or mode="ring"

# Each rank now gets seq[rank * (S/cp) : (rank+1) * (S/cp)] of the global sequence.
# Forward + backward run on the shard; an all-to-all redistributes K/V over heads
# at attention boundaries.

Constraints: n_heads % cp_size == 0 (Ulysses head-shard) and seq % cp_size == 0 (sequence-shard).

Benchmark — single GPU (Qwen2.5-1.5B, V100 32 GB, fp16, fwd+bwd, grad-ckpt)

mode	seq	tok/s	peak GB	notes
eager (math SDPA)	1 024	3 347	8.97	baseline A
eager (math SDPA)	8 192	1 516	24.0	largest seq before OOM
xformers SDPA	1 024	3 741	8.97	baseline B
xformers SDPA	8 192	3 018	24.0	best non-CP V100 baseline
xformers + tc_volta.compile (round 2)	1 024	4 020	8.97	+7 % over xformers
xformers + tc_volta.compile (round 2)	4 096	3 932	14.3	+9 % (round-1 was −47 % from cache pollution; fixed)
xformers + tc_volta.compile (round 2)	8 192	3 312	24.0	+10 % (round-1 OOM'd; fixed)

Round-2 numbers from a cold GPU. Round-1 cache-pollution diagnosis in BENCH.md §F3 revisit.

Benchmark — 4× V100 (Qwen2.5-1.5B, CP=4 Ulysses, fp16, fwd+bwd, grad-ckpt)

cp	seq	seq/rank	tok/s	peak/rank GB
4	8 192	2 048	9 542	9.6
4	16 384	4 096	7 323	14.3
4	24 576	6 144	5 740	19.2
4	32 768	8 192	4 733	24.0

Grad parity test on (B=1, H=12, N=2304, D=64): cos-sim = 1.000000 on out / dq / dk / dv between cp=1 and cp=4. Reproduces exactly through 50 timed steps.

Headline — round 2

	seq	tok/s	peak/rank GB	tok/s / GPU
Single-GPU xformers	8 192	3 018	24.0	3 018
CP=3 xformers (round 1)	16 128	5 514	17.3	1 838
CP=4 xformers (round 2)	32 768	4 733	24.0	1 183

4× the context for +57 % tok/s vs strongest single-GPU baseline. Per-GPU efficiency falls with cp_size as expected (each rank's local attention is O(N²) and all-to-all volume grows linearly with cp), but the alternative on V100 is "OOM" — eager hits 24 GB at seq=8 K, you simply can't go further on one card.

Going round-1 → round-2 (cp=3 seq=16 128 → cp=4 seq=32 768): +2.0× context for −14 % tok/s, +39 % peak memory. Sub-linear cost for doubling the trainable window.

Failure-mode catalogue (V100 + torch 2.0.1)

Tested every cell of {mode = default | reduce-overhead | max-autotune} × {dtype = fp16 | fp32}. Full table + root causes in TORCH_COMPILE_VOLTA.md.

mode	dtype	output	grad	verdict
default	fp32	match	match (cos 1.0)	OK
default	fp16	match	cos 1.0, ULP-level abs diff	OK
reduce-overhead	fp32	match	garbage (cos 0.0)	DO NOT USE on Volta
reduce-overhead	fp16	match	garbage (cos 0.0)	DO NOT USE on Volta
max-autotune	fp32	match	match (cos 1.0)	OK, slow compile
max-autotune	fp16	match	cos 1.0, abs diff ≤ 5 e-2	OK with caveat

tc_volta.compile() flips torch._inductor.config.triton.cudagraphs = False unconditionally — that's the mitigation for the reduce-overhead zero-gradient bug, and it's the only knob you have to remember.

Reproduce

# Single-GPU sweep (eager vs xformers vs xformers+compile)
for SEQ in 1024 2048 4096 8192; do
  for MODE in eager xformers xformers_compile xformers_mlp_compile; do
    CUDA_VISIBLE_DEVICES=0 python bench/single_gpu.py \
      --model Qwen/Qwen2.5-1.5B --seq $SEQ --mode $MODE --steps 25
  done
done

# CP sweep (cp_size=2/3/4, seq up to 32 768)
for CP in 2 3 4; do
  for SEQ in 8192 16384 24576 32768; do
    DEVS=$(seq -s, 0 $((CP-1)))
    CUDA_VISIBLE_DEVICES=$DEVS python -m torch.distributed.run \
      --nproc-per-node=$CP --master-port=29504 \
      bench/cp_16k.py --model Qwen/Qwen2.5-1.5B --seq $SEQ --steps 30 --no-compile
  done
done

# Headline 32 K run, 50 steps
CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.run \
  --nproc-per-node=4 --master-port=29504 \
  bench/cp_16k.py --model Qwen/Qwen2.5-1.5B --seq 32768 --steps 50 --no-compile

# Correctness + grad parity
python -m pytest tests/test_compile_correctness.py -v
CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.run \
  --nproc-per-node=4 tests/test_cp_grad_parity.py

Raw JSON per row lives in results/.

Known issues / what's NOT here

• flash_attn_volta Triton-2.3 kernel does not fire in the compiled graph. Inductor needs Triton 2.0, Triton 2.0's V100 backend has the tt.reduce MMA-layout-unification bug. The Inductor SDPA custom op falls back to xformers Cutlass instead. Diagnosed in TORCH_COMPILE_VOLTA.md; the fix is either a Triton-2.0-compatible FAV rewrite or torch ≥ 2.1.
• Qwen3-4B requires transformers ≥ 4.51 which requires Python ≥ 3.9. The eva01 box runs Python 3.8 and we measured on Qwen2.5-1.5B. Same patches (patch_qwen3) work where the deps allow.
• reduce-overhead unusable for backward on V100 due to cudagraph-aliasing zero-gradient bug. Root cause not chased; mitigation (cudagraphs=False) is auto-applied.
• Backward kernels (this repo): we inherit attention backward from xformers Cutlass. The sister project flash-attn-volta has hand-rolled Triton dQ/dK/dV kernels that aren't yet wired through this stack (see first bullet).

License

Apache 2.0. Performance optimisations adopted in spirit (no source copy) from xformers' Triton autotune patterns and DeepSpeed Ulysses' all-to-all layout. xformers is BSD-3, DeepSpeed is Apache 2.0.

Credits

• DeepSpeed Ulysses: Jacobs et al., arXiv 2309.14509.
• Ring Attention: Liu et al., arXiv 2310.01889.
• FlashAttention: Tri Dao et al., arXiv 2205.14135.
• xformers: Lefaudeux et al. (Meta), https://github.com/facebookresearch/xformers.
• Sister kernel project: flash-attn-volta.
• Built by ml-intern Claude Code skill.