Performance Benchmarks

Performance Benchmarks

Expected Speedups

Performance improvements from ZoRRo Train's prompt deduplication depend primarily on the deduplication ratio — the fraction of redundant prompt tokens in your RL training workload.

Typical Scenarios

Key insight: Speedups scale with the number of identical prompts processed together. When each sequence has a unique prompt, deduplication offers negligible benefit.

Performance Factors

Actual speedup depends on several hardware and software factors:

Hardware Characteristics

• Compute vs. memory bottleneck: Speedups are maximized on memory-bound hardware (longer prompts, larger models)
• GPU architecture: Results vary across different GPU generations and families
• Memory bandwidth: Higher bandwidth reduces the relative cost of attention computation

Sequence Lengths

• Longer prompts = more benefit: O(n²) attention cost means 8K–32K token prompts yield the largest improvements
• Short prompts: < 512 tokens rarely justify deduplication overhead

Attention Implementation

• Flash Attention 2/3: Best performance; leverages IO-aware kernels
• Eager attention: Good speedups but slower baseline, so relative improvements are larger
• Split attention mode (use_split_attention=True): Additional savings by avoiding redundant prompt-to-prompt attention

Benchmark Suite

Running Performance Benchmarks

The performance benchmark measures forward and backward pass speedup on large batches with realistic prompt deduplication:

python arctic_platform/rl/zorro_train/test_perf.py

This script:

1. Creates batches with varying deduplication ratios
2. Measures wall-clock time for forward and backward passes
3. Compares deduplicated vs. baseline (non-deduplicated) execution
4. Reports speedup factors and per-token throughput

Redundancy Analysis

For typical RL training, redundancy is substantial:

Example: 10K-token prompt, 10 responses of 1K tokens each

Total tokens:      110K  (10K × 10 + 1K × 10)
Unique tokens:      20K  (10K prompt + 1K × 10 responses)
Redundancy:         82%  of computation wasted on duplicate prompts

At transformer attention's O(n²) complexity, duplicating the prompt 10 times means:

• Memory: 10 copies of prompt embeddings, keys, and values
• Compute: 10 identical attention operations over the same 10K prompt tokens
• Impact: Dominant bottleneck for long-context RL (4K–32K token prompts)

Testing Correctness and Performance

Correctness Tests

Verify gradient and numerical correctness:

pytest tests/zorro_train/test_once_patcher.py

This test sweeps:

• Model architectures: dense Qwen3, MoE, hybrid (linear + full attention)
• Attention implementations: flash_attention_2, eager
• Batch configurations: padded and unpadded
• Logits optimizations: none, memory, compute

Logprobs and gradients are validated against a non-deduplicated reference implementation.

Deduplication Algorithm Tests

CPU-level tests for the core dedup tensor transformations:

pytest tests/zorro_train/test_dedup.py
pytest tests/zorro_train/test_seqlen_balancing.py

Usage Tips for Optimal Performance

1. Maximize Prompt Deduplication Ratio

Structure your RL training to group identical prompts in the same batch:

from arctic_platform.rl.zorro_train.actor import DeduplicatedActor
from arctic_platform.rl.zorro_train.tests import create_dummy_batch

actor = DeduplicatedActor(
    model_name_or_path="Qwen/Qwen3-0.6B",
    device="cuda",
    use_split_attention=True,        # Split attention saves redundant prompt-to-prompt compute
    attn_implementation="flash_attention_2",  # Best performance
)

# Create batch with high dedup ratio (2 unique prompts across 16 sequences)
batch = create_dummy_batch(
    batch_size=16,
    num_unique_prompts=2,  # ← Maximize this ratio
    prompt_len=8192,       # ← Longer prompts = more benefit
    response_len=512,
    device="cuda",
)

output = actor.forward(batch, temperature=1.0)

2. Use Split Attention

Enable split attention mode (default) for additional savings:

actor = DeduplicatedActor(
    model_name_or_path="Qwen/Qwen3-0.6B",
    device="cuda",
    use_split_attention=True,  # ← Splits into prompt-to-prompt + response-to-full
)

Split attention avoids redundant prompt-to-prompt attention computations by running two separate attention calls instead of one full attention.

3. Choose Flash Attention for Long Sequences

For prompts > 2K tokens, use Flash Attention 2 or 3:

actor = DeduplicatedActor(
    model_name_or_path="Qwen/Qwen3-0.6B",
    device="cuda",
    attn_implementation="flash_attention_2",  # Faster for long sequences
)

Verify Flash Attention is installed:

pip install flash-attn --no-build-isolation

4. Monitor Speedup with Profiling

Profile your specific workload to understand the actual speedup:

import time

actor = DeduplicatedActor(
    model_name_or_path="Qwen/Qwen3-0.6B",
    device="cuda",
)

batch = create_dummy_batch(batch_size=16, num_unique_prompts=2, prompt_len=8192, response_len=512)

# Measure deduplicated forward
torch.cuda.reset_peak_memory_stats()
torch.cuda.synchronize()
start = time.perf_counter()

output = actor.forward(batch, temperature=1.0, calculate_entropy=True)

torch.cuda.synchronize()
elapsed = time.perf_counter() - start
peak_mem = torch.cuda.max_memory_allocated() / 1e9

print(f"Forward time: {elapsed*1000:.2f}ms, Peak memory: {peak_mem:.2f}GB")

Performance Limitations

When Deduplication Provides No Benefit

• Each sequence has a unique prompt: Deduplication ratio = 0, overhead outweighs gains
• Very short prompts: < 256 tokens; overhead exceeds savings
• Small batch sizes: < 4 sequences with shared prompts

Future Optimization Opportunities

• Partial prompt overlap: Prefix caching for prompts that share common prefixes (not yet supported)
• Model architecture expansion: Optimization currently supports Qwen3 family; other models (Llama, Mistral) could benefit similarly

References

For detailed implementation information and correctness validation, see:

• Core algorithm: arctic_platform/rl/zorro_train/zorro_train.py — ZoRRoTrain class and ReconstructionInfo
• Production patcher: arctic_platform/rl/zorro_train/qwen_model_patcher.py — Qwen3ModelOncePatcher
• Reference harness: arctic_platform/rl/zorro_train/actor.py — DeduplicatedActor
• Full documentation: See the ZoRRo Train README for architecture details, API reference, and getting started guide