ZoRRo Train Overview

ZoRRo Train Overview

What is ZoRRo Train?

ZoRRo stands for Zero Redundancy Rollouts. It is a prompt-deduplication optimization that automatically removes redundant prompt computation during reinforcement learning (RL) training, particularly in scenarios like RLHF (reinforcement learning from human feedback) and policy gradient training with PPO or GRPO.

The Problem: Prompt Redundancy in RL

During RL training, a critical computational bottleneck emerges: prompt redundancy.

Typical RL Training Workflow

1. Sampling: Generate multiple responses for the same prompt to explore different strategies
2. Evaluation: Score each response using a reward model
3. Training: Compute policy gradients by processing all prompt-response pairs

This creates significant computational waste:

Prompt A + Response 1  ──┐
Prompt A + Response 2  ──┤
Prompt A + Response 3  ──┼──→  Same prompt processed N times!
    ...                  │
Prompt A + Response N  ──┘

Impact

• In typical RL training, 80–95% of tokens are redundant prompt tokens
• For a 10K-token prompt with 10 responses of 1K tokens each:
  • Total: 110K tokens (10K × 10 + 1K × 10)
  • Unique: 20K tokens (10K prompt + 1K × 10 responses)
  • Redundancy: 82% of computation is wasted on duplicate prompts

Why This Matters

Transformer attention has O(n²) complexity. When the same prompt appears multiple times:

• Memory: N copies of the same prompt embeddings, attention keys, and values
• Compute: N identical attention computations over the same prompt tokens
• Time: Training throughput is severely limited

For long-context RL (prompts with 4K–32K tokens), this redundancy becomes the dominant cost.

How ZoRRo Train Works

The optimization operates in three phases:

1. Prompt Detection & Deduplication

ZoRRo identifies which sequences in a batch share identical prompts and groups them:

Input Batch (8 sequences):
┌─────────────────────────────────┐
│ [Prompt A][Response 1]          │
│ [Prompt A][Response 2]          │  ──→  Deduplicated Batch:
│ [Prompt B][Response 3]          │       [Prompt A][Response 1][Response 2]
│ [Prompt B][Response 4]          │       [Prompt B][Response 3][Response 4]
│     ...                         │
└─────────────────────────────────┘

For each unique prompt, a single concatenated sequence is created with that prompt followed by all its responses. Reconstruction metadata is tracked for reversing the transformation.

2. Optimized Attention

Two attention strategies are available:

Standard QKV Optimization (use_split_attention=False):

• Compute Q, K, V projections on deduplicated batch
• Reconstruct to full batch shape
• Run standard causal attention

Split Attention Optimization (use_split_attention=True, default):

• Compute Q, K, V projections on deduplicated batch
• Split into prompt and response parts
• Run two separate attention calls:
  1. Prompt-to-Prompt: Deduplicated prompts attend to themselves
  2. Response-to-Full: Each response attends to its prompt + itself
• Replicate prompt results and concatenate with response results

The split attention approach saves additional computation by avoiding redundant prompt-to-prompt attention computations.

3. Transparent Reconstruction

• Output logits are automatically reconstructed to match the original batch shape
• Gradients flow correctly through the deduplicated computation
• The rest of the model (embeddings, MLP layers, final projection) sees the expected batch shape

Architecture

The optimization is delivered via monkey-patching a Hugging Face Qwen3 model. Two patchers are provided:

• Qwen3ModelOncePatcher and QwenAttentionOncePatcher — patched once onto the model by the DeepSpeed worker. Called with the full [B, S] batch, it deduplicates shared prompts internally, runs the packed forward/backward, and returns per-response-token logprobs / entropy in the original sample order (no padding). This is the production path.

• Qwen3ModelPatcher and QwenAttentionPatcher — a context-manager harness for the demo and gradient-correctness reference.

┌──────────────────────────────────────────────────────┐
│  ZoRRoTrain                                          |
│  ┌────────────────────────────────────────────────┐  │
│  │  1. Find prompt groups                         │  │
│  │  2. Create deduplicated batch                  │  │
│  │  ┌──────────────────────────────────────────┐  │  │
│  │  │  Qwen3ModelOncePatcher                   │  │  │
│  │  │  ┌────────────────────────────────────┐  │  │  │
│  │  │  │  QwenAttentionOncePatcher          │  │  │  │
│  │  │  │  - Intercepts attention forward()  │  │  │  │
│  │  │  │  - Applies deduplication logic     │  │  │  │
│  │  │  │  - Reconstructs outputs            │  │  │  │
│  │  │  └────────────────────────────────────┘  │  │  │
│  │  │  Model Forward Pass                      │  │  │
│  │  └──────────────────────────────────────────┘  │  │
│  │  3. Reconstruct logits to original batch       │  │
│  └────────────────────────────────────────────────┘  │
│  Backward Pass (with patching active)                │
└──────────────────────────────────────────────────────┘

Key Features

✅ Mathematically Correct — Gradients match baseline implementation (within numerical precision)

✅ Transparent — Works as a drop-in replacement for standard forward/backward

✅ Flexible — Supports multiple attention implementations (eager, Flash Attention 2/3)

✅ Gradient Checkpointing — Compatible with activation checkpointing

✅ Mixed Precision — Optimized for bfloat16 training

Supported Models

ZoRRo Train currently supports these Qwen model families:

• qwen3
• qwen3-moe
• qwen3-next-moe
• qwen3.6
• qwen3.6-moe

These span dense, MoE, and hybrid (linear + full attention) architectures. More models will be added in the future.

Performance Expectations

Expected speedups depend on the deduplication ratio:

Note: Actual speedups depend on:

• Hardware (GPU compute vs memory bottleneck)
• Attention implementation (Flash Attention gives best results)
• Sequence lengths (longer prompts = more benefit)

Integration

ZoRRo Train is installed transparently by the DeepSpeed training/log-prob engines and toggled per run via the RL config (zorro_train.enable).

For framework integrations, see the Arctic RL documentation.

Quick Reference

Core Classes

from arctic_platform.rl.zorro_train import (
    ZoRRoTrain,
    Qwen3ModelOncePatcher,
    QwenAttentionOncePatcher,
    ReconstructionInfo,
)
from arctic_platform.rl.zorro_train.actor import DeduplicatedActor

Basic Usage Example

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

# Initialize the actor
actor = DeduplicatedActor(
    model_name_or_path="Qwen/Qwen3-0.6B",
    device="cuda",
    logits_optimization="none",      # "none" | "compute" | "memory"
    use_split_attention=True,        # split attention (prompt-to-prompt + response-to-full)
    attn_implementation="eager",     # or "flash_attention_2"
)

# Create a batch with shared prompts
batch = create_dummy_batch(
    batch_size=8,
    num_unique_prompts=2,  # 8 sequences, but only 2 unique prompts
    prompt_len=4096,
    response_len=512,
    device="cuda",
    include_training_fields=True,
    add_padding=True,
)

# Forward pass (automatically deduplicates)
output = actor.forward(batch, temperature=1.0, calculate_entropy=True)
log_probs, entropy = output.logprobs, output.entropy

# Training step with backward pass
actor.model.train()
metrics = actor.compute_policy_loss_and_backward(
    batch,
    temperature=1.0,
    gradient_accumulation=1
)

print(f"Policy loss: {metrics['actor/policy_loss']:.4f}")

Testing

Correctness tests live in the top-level test suite:

pytest tests/zorro_train/                          # everything below

pytest tests/zorro_train/test_dedup.py             # CPU: dedup-algorithm round-trips
pytest tests/zorro_train/test_once_patcher.py      # GPU: forward/backward vs reference
pytest tests/zorro_train/test_seqlen_balancing.py  # CPU: sequence-length balancing

Run the demo:

python arctic_platform/rl/zorro_train/demo.py

Performance benchmark:

python arctic_platform/rl/zorro_train/test_perf.py

Limitations & Future Work

Current Limitations:

• Best speedups require identical prompts; partial overlap is not exploited

Future Directions:

• Support for other model architectures (Llama, Mistral, etc.)
• Prefix caching for partially overlapping prompts

References

For the complete design, deduplication/attention internals, and detailed API reference, see arctic_platform/rl/zorro_train/README.md.