Getting Started with ZoRRo Train

Getting Started with ZoRRo Train

What is ZoRRo Train?

ZoRRo stands for Zero Redundancy Rollouts. It is a prompt deduplication optimization for reinforcement learning (RL) training that dramatically reduces redundant computation without changing the mathematical correctness of your model.

The Problem It Solves

In RLHF and RL-based language model training, the same prompt is processed multiple times with different responses. 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  ──┘

Key Statistics:

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

Since Transformer attention has O(n²) complexity, this redundancy becomes the dominant cost for long-context RL (4K-32K token prompts).

How ZoRRo Train Works

The optimization operates in three phases:

Phase 1: Prompt Detection & Deduplication

The system identifies sequences that share identical prompts and creates a deduplicated batch:

Input Batch (8 sequences):
[Prompt A][Response 1]
[Prompt A][Response 2]     ──→  [Prompt A][Response 1][Response 2]
[Prompt B][Response 3]           [Prompt B][Response 3][Response 4]
[Prompt B][Response 4]

Phase 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:
  • Prompt-to-Prompt: Deduplicated prompts attend to themselves
  • Response-to-Full: Each response attends to its prompt + itself

Phase 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 sees the expected batch shape

Installation

Ensure you have the required dependencies:

pip install torch transformers

For Flash Attention (optional, recommended for best performance):

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

Basic Usage

Step 1: Initialize the Actor

from arctic_platform.rl.zorro_train.actor import DeduplicatedActor

actor = DeduplicatedActor(
    model_name_or_path="Qwen/Qwen3-0.6B",
    device="cuda",
    logits_optimization="none",      # "none" | "compute" | "memory"
    use_split_attention=True,        # split attention (default)
    attn_implementation="eager",     # or "flash_attention_2"
)

Constructor Arguments:

Step 2: Create a Batch with Shared Prompts

from arctic_platform.rl.zorro_train.tests import create_dummy_batch

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,
)

The batch should contain:

• input_ids: [batch_size, seq_len]
• position_ids: [batch_size, seq_len]
• attention_mask: [batch_size, seq_len]
• responses: [batch_size, response_len]

Step 3: Run Deduplicated Forward Pass

output = actor.forward(batch, temperature=1.0, calculate_entropy=True)
log_probs = output.logprobs      # shape: [num_valid_response_tokens]
entropy = output.entropy         # shape: [num_valid_response_tokens]

Note: logprobs and entropy are packed 1D tensors containing only valid response tokens in the original sample order (padding removed).

Step 4: Compute Policy Loss and Run Backward

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}")

Supported Models

ZoRRo Train currently supports the following Qwen3 model families:

• qwen3 (dense)
• qwen3-moe (MoE)
• qwen3-next-moe (MoE)
• qwen3.6 (dense)
• qwen3.6-moe (MoE)

Support for additional model architectures is planned for future releases.

Running the Demo

To see ZoRRo Train in action:

python arctic_platform/rl/zorro_train/demo.py

This runs:

1. A gradient/logprob correctness test (deduplicated vs. baseline)
2. An optional performance benchmark

Testing Your Setup

Run the test suite to verify correct installation and functionality:

# Run all ZoRRo Train tests
pytest tests/zorro_train/

# Test the deduplication algorithm (CPU, no model required)
pytest tests/zorro_train/test_dedup.py

# Test the patcher with actual model forward/backward (GPU required)
pytest tests/zorro_train/test_once_patcher.py

# Test sequence-length balancing (CPU)
pytest tests/zorro_train/test_seqlen_balancing.py

The test_once_patcher.py test sweeps across:

• Dense, MoE, and hybrid model architectures
• Both flash_attention_2 and eager attention implementations
• Padded and unpadded batches
• All three logits optimization modes

Performance Expectations

Expected speedups depend on the deduplication ratio:

Actual speedups depend on:

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

Run the performance benchmark:

python arctic_platform/rl/zorro_train/test_perf.py

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

API Overview

DeduplicatedActor

Reference harness for running deduplicated forward and backward passes.

Key Methods:

• forward(micro_batch, temperature=1.0, calculate_entropy=False) — Deduplicated forward pass. Returns ModelOutput with logprobs and optionally entropy as packed 1D tensors (valid response tokens only, original sample order).

• compute_policy_loss_and_backward(micro_batch, temperature=1.0, gradient_accumulation=1) — PPO training step (forward + backward). Returns a metrics dictionary with actor/policy_loss, actor/pg_loss, etc.

• patch(response_len, rollout_n, temperature) — (Re)install Qwen3ModelOncePatcher on the model. Useful for changing parameters between training steps.

• train() / eval() — Set model to training or evaluation mode.

ZoRRoTrain

Static utility class implementing the core deduplication tensor logic. Key methods:

• find_prompt_groups(input_ids, response_length) — Group rows by prompt identity.

• create_deduplicated_batch(input_ids, position_ids, response_length, prompt_groups, unique_prompts, ...) — Pack each unique prompt followed by its responses.

• reconstruct_sequences(dedup_hidden, reconstruction_info) — Reconstruct the full batch from deduplicated output.

Project Structure

arctic_platform/rl/zorro_train/
├── README.md                  # Detailed documentation
├── __init__.py                # Public API exports
├── zorro_train.py             # Core dedup algorithm
├── actor.py                   # DeduplicatedActor reference harness
├── qwen_model_patcher.py      # Qwen3 model-level patching
├── qwen_attention_patcher.py  # Attention-level patching
├── module_patcher.py          # Base patching utilities
├── seqlen_balancing.py        # Sequence-length balancing
├── demo.py                    # Interactive demonstration
├── test_perf.py               # Performance benchmark
└── tests.py                   # Batch builders and helpers

Next Steps

1. Run the Demo — Get a feel for the system:

   python arctic_platform/rl/zorro_train/demo.py
   

2. Review Test Examples — See how the patcher is used in tests/zorro_train/test_once_patcher.py

3. Integrate into Your Workflow — Replace your model's forward pass with DeduplicatedActor.forward() in your RL training loop

4. Benchmark — Measure speedups on your hardware and batch configurations using test_perf.py

Troubleshooting

Issue: "Model not supported"
Solution: ZoRRo Train currently supports Qwen3 models. For other architectures, consider using the base ModuleReconstructionPatcher class as a template.

Issue: Low speedup despite high deduplication ratio
Solution: Ensure you're using flash_attention_2 for best performance. Also check that your GPU is memory-bound (not compute-bound) by profiling with torch.profiler.

Issue: Gradient mismatch between deduplicated and baseline
Solution: Run test_once_patcher.py to verify correctness on your hardware. Most differences are within numerical precision for bfloat16.

Citation

If you use ZoRRo Train in your research, please cite:

@misc{zorro_train_2025,
  title={ZoRRo Train: Zero Redundancy Rollouts for Efficient RL Training},
  author={Snowflake AI Research},
  year={2025},
  howpublished={\url{https://github.com/Snowflake-AI-Research/Arctic-Platform}}
}

License

Apache License 2.0. See LICENSE.