How Prompt Deduplication Works

How Prompt Deduplication Works

The Problem: Prompt Redundancy in RL Training

During reinforcement learning training (e.g., PPO), the typical workflow involves:

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 on all prompt-response pairs

This creates significant computational waste because the same prompt is processed multiple times:

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

Impact on Computation

In typical RL training:

• 80–95% of tokens are redundant prompt tokens
• Transformer attention has O(n²) complexity, so processing the same prompt N times costs N² of computation
• For a 10K-token prompt with 10 responses of 1K tokens each:
  • Total: 110K tokens
  • Unique: 20K tokens
  • Redundancy: 82% of computation is wasted on duplicate prompts

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

The Solution: Automatic Prompt Deduplication

ZoRRo Train removes prompt redundancy through three optimization phases:

Phase 1: Prompt Detection & Deduplication

The algorithm identifies which sequences 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]
│     ...                         │
└─────────────────────────────────┘

Key operations:

• Group samples by identical prompts using find_prompt_groups()
• Create a single concatenated sequence where each unique prompt appears once followed by its responses
• Track reconstruction metadata for reversing the transformation

Phase 2: Optimized Attention

Two attention strategies are available:

Standard QKV Optimization (use_split_attention=False)

# Compute Q, K, V on deduplicated batch
dedup_input_ids, _, reconstruction_info = ZoRRoTrain.create_deduplicated_batch(
    input_ids, position_ids, response_length, prompt_groups, unique_prompts
)
# Forward through model
output = model(dedup_input_ids)
# Reconstruct to original batch shape
reconstructed = ZoRRoTrain.reconstruct_sequences(output, reconstruction_info)

Process:

1. Compute Q, K, V projections on the deduplicated batch
2. Reconstruct to full batch shape
3. Run standard causal attention

Split Attention Optimization (use_split_attention=True, default)

This approach saves additional computation by avoiding redundant prompt-to-prompt attention:

1. Prompt-to-Prompt: Deduplicated prompts attend only to themselves (computed once for all copies)
2. Response-to-Full: Each response attends to its unique prompt + itself

The split attention approach is more efficient because prompt-to-prompt computations are only done once per unique prompt, not once per response.

Phase 3: Transparent Reconstruction

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

Core Algorithm: ZoRRoTrain Class

The ZoRRoTrain class provides model-free tensor deduplication utilities:

Finding Prompt Groups

prompt_groups, unique_prompts = ZoRRoTrain.find_prompt_groups(
    input_ids,      # [batch_size, seq_len]
    response_length # length of response portion
)

Returns:

• prompt_groups: List of lists, each containing indices of samples sharing a prompt
• unique_prompts: [num_unique, prompt_len] tensor of unique prompt tokens

Example:

# Input: 4 samples where prompts at [0,1] and [2,3] match
prompt_groups = [[0, 1], [2, 3]]  # Two groups
unique_prompts = [[prompt_A_tokens], [prompt_B_tokens]]

Creating Deduplicated Batch

dedup_input_ids, dedup_position_ids, reconstruction_info = ZoRRoTrain.create_deduplicated_batch(
    input_ids,      # [batch_size, seq_len]
    position_ids,   # [batch_size, seq_len]
    response_length,
    prompt_groups,
    unique_prompts,
    attention_mask=None,  # Optional
    use_unpad=False       # If True, remove padding
)

Output structure:

• dedup_input_ids: [1, total_dedup_tokens] – single packed sequence
• dedup_position_ids: [1, total_dedup_tokens] or [1, total_valid_tokens] if unpacked
• reconstruction_info: Dict containing metadata for reconstruction

Concatenation pattern:

[Prompt_1][Response_1_1][Response_1_2]...[Prompt_2][Response_2_1][Response_2_2]...

Reconstructing Original Batch

reconstructed_hidden = ZoRRoTrain.reconstruct_sequences(
    dedup_hidden,         # [1, total_tokens, hidden_dim]
    reconstruction_info
)

Returns: Original batch shape [batch_size, seq_len, hidden_dim]

The reconstruction uses a Triton-optimized kernel for efficiency:

@triton.jit
def _triton_reconstruct_seq(...):
    # For each output sample, copy its unique prompt once
    # then copy its response tokens

Architecture: Patching Strategy

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

Production Patcher: Qwen3ModelOncePatcher

Used in production and GPU training loops. Installed once and patches the model for the entire training run.

from arctic_platform.rl.zorro_train import Qwen3ModelOncePatcher

patcher = Qwen3ModelOncePatcher(
    model,
    use_split_attention=True,
    attn_implementation="eager"  # or "flash_attention_2"
)
patcher.patch(response_len=512, rollout_n=10, temperature=1.0)

# Now all forward passes use deduplication internally
output = model(input_ids)

Reference Patcher: Qwen3ModelPatcher

A context-manager harness for demos and correctness testing:

from arctic_platform.rl.zorro_train.qwen_model_patcher import Qwen3ModelPatcher

with Qwen3ModelPatcher(model, use_split_attention=True):
    output = model(input_ids)  # Runs deduplicated
# Patching automatically removed

Lower-Level: QwenAttentionPatcher

Operates at the attention layer for fine-grained control:

from arctic_platform.rl.zorro_train.qwen_attention_patcher import QwenAttentionPatcher

patcher = QwenAttentionPatcher(
    attention_layer,
    use_split_attention=True
)
patcher.patch()

Working with Packed/Unpacked Sequences

ZoRRo Train handles both formats:

Padded Format (with attention mask)

# Input: [batch_size, seq_len] with padding
# attention_mask: 1 for valid tokens, 0 for padding

dedup_ids, dedup_pos, info = ZoRRoTrain.create_deduplicated_batch(
    input_ids,
    position_ids,
    response_length,
    prompt_groups,
    unique_prompts,
    attention_mask=attention_mask,  # Include mask
    use_unpad=False
)

Unpacked Format (variable lengths, no padding)

# Removes padding during deduplication
dedup_ids, dedup_pos, info = ZoRRoTrain.create_deduplicated_batch(
    input_ids,
    position_ids,
    response_length,
    prompt_groups,
    unique_prompts,
    attention_mask=attention_mask,
    use_unpad=True  # Enable unpacking
)

When unpacked, the returned reconstruction_info contains padding boundaries so gradients and outputs can be correctly mapped back.

Flash Attention Integration

For non-square attention masks (e.g., with KV caching), convert to Flash Attention parameters:

flash_params = ZoRRoTrain.attention_mask_to_flash_attn_params(
    attention_mask  # [total_q, total_kv] sparse block-diagonal
)

# Now use in Flash Attention kernels
output = flash_attention_forward(
    q, k, v,
    cu_seqlens_q=flash_params["cu_seqlens_q"],
    cu_seqlens_k=flash_params["cu_seqlens_k"],
    max_seqlen_q=flash_params["max_seqlen_q"],
    max_seqlen_k=flash_params["max_seqlen_k"]
)

Correctness & Gradients

ZoRRo Train maintains mathematical correctness:

• Gradients computed through deduplicated batch are numerically equivalent to baseline (within floating-point precision)
• The deduplication is transparent to the model—no API changes required
• Activation checkpointing and mixed precision (bfloat16) are fully supported

Gradient flow:

Input → Dedup → Forward → Attention (deduplicated) → Reconstruct → Output
         ↓                                              ↑
         └──────────── Backward (gradients flow through both) ────┘

Performance Expectations

Speedup depends on the deduplication ratio:

Factors affecting speedup:

• Deduplication ratio: More identical prompts → more savings
• Attention implementation: Flash Attention gives best results
• Sequence lengths: Longer prompts = more benefit (due to O(n²) attention)
• Hardware: GPU memory vs. compute bound

Supported Models

Currently ZoRRo Train supports:

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

These span dense, MoE, and hybrid architectures. More models will be added in the future.

Example: Using DeduplicatedActor

For a complete reference implementation, use the DeduplicatedActor harness:

from arctic_platform.rl.zorro_train.actor import DeduplicatedActor

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

# Forward pass (auto-deduplicated)
output = actor.forward(
    batch,
    temperature=1.0,
    calculate_entropy=True
)
log_probs = output.logprobs    # [num_valid_response_tokens]
entropy = output.entropy        # [num_valid_response_tokens]

# Training step
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}")

Output logprobs and entropy are packed 1D tensors in the original sample order, containing only valid (non-padded) response tokens.

Debugging & Analysis

Batch Statistics

Analyze a normal (non-deduplicated) batch:

from arctic_platform.rl.zorro_train.zorro_train import analyze_normal_batch_via_attention_mask

analyze_normal_batch_via_attention_mask(
    input_ids,
    attention_mask,
    padded_response_len=4096
)

Output example:

batch stats: prompt lens: 15.9Ki resp count=16/min=0.4Ki/mean=0.6Ki/max=1.0Ki 
             non-pad-tokens=26.6K/82.5%

This shows:

• Prompt length
• Response counts (min/mean/max)
• Total non-padded tokens and percentage

Key Limitations & Future Work

Current Limitations:

• Best speedups require identical prompts (partial overlap not yet exploited)
• Implemented for Qwen models only (other architectures coming soon)

Future Directions:

• Support for other model families (Llama, Mistral, etc.)
• Prefix caching for partially overlapping prompts
• Dynamic deduplication detection without manual group specification