Testing & Correctness Verification

Testing & Correctness Verification

Overview

ZoRRo Train includes a complete test suite to verify:

• Correctness of the deduplication algorithm (CPU)
• Gradient correctness and numerical precision (GPU)
• Sequence-length balancing (CPU)
• Performance benchmarking (GPU)

All tests live in tests/zorro_train/ at the repository root.

Test Files

test_dedup.py — Core Deduplication Algorithm (CPU)

Verifies the deduplication and reconstruction logic without requiring a model or GPU.

What it tests:

• find_prompt_groups(): Correctly identifies sequences with identical prompts
• create_deduplicated_batch(): Packs unique prompts and responses correctly
• reconstruct_sequences(): Round-trip correctness (deduplicated → original batch shape)
• deduplicate_sequences(): Inverse operations match

Run:

pytest tests/zorro_train/test_dedup.py

test_once_patcher.py — Forward & Backward Correctness (GPU)

The primary correctness test for the production patcher (Qwen3ModelOncePatcher). Compares deduplicated forward/backward against a reference baseline implementation.

What it tests:

• Forward pass: Log probabilities and entropy match the baseline (within numerical precision)
• Backward pass: Gradients flow correctly through deduplicated computation
• Multiple architectures: Dense, MoE, and hybrid Qwen3 models
• Attention implementations: Both eager and flash_attention_2
• Batch configurations: Padded and unpadded batches
• Logits optimization modes: none, memory, and compute

Test sweep:

• Tiny-random Qwen3 checkpoints (faster than full models)
• Multiple batch sizes and deduplication ratios
• Both use_split_attention=True and use_split_attention=False

Run:

pytest tests/zorro_train/test_once_patcher.py

Example test case:

# Tests that deduplicated forward/backward matches baseline
# across different model families, attention implementations, and batch types
test_qwen3_model_once_patcher(
    model_name="tiny-random-Qwen3",
    attn_impl="flash_attention_2",
    use_padded_batch=True,
    use_split_attention=True,
    logits_optimization="compute"
)

test_seqlen_balancing.py — Sequence-Length Balancing (CPU)

Verifies the sequence-length balancing logic for distributing work across micro-batches.

What it tests:

• Sequences are grouped by similar length
• Micro-batches are created without exceeding token limits
• All sequences are assigned to exactly one micro-batch

Run:

pytest tests/zorro_train/test_seqlen_balancing.py

Performance Benchmarking

test_perf.py — Dedicated Benchmark

Measures forward and backward pass speedup of deduplication on realistic batches.

What it measures:

• Forward pass duration (baseline vs. deduplicated)
• Backward pass duration (baseline vs. deduplicated)
• Overall speedup factor
• Memory usage (optional)

Run:

python arctic_platform/rl/zorro_train/test_perf.py

Output example:

Forward pass:   1.2s (deduplicated) vs 3.5s (baseline) → 2.9x speedup
Backward pass:  0.8s (deduplicated) vs 2.1s (baseline) → 2.6x speedup
Overall:        2.8x speedup

Benchmark parameters (configurable in the script):

• Batch size: 16–32
• Number of unique prompts: 1–4
• Prompt length: 4K–8K tokens
• Response length: 512–1K tokens

Running All Tests

Quick verification (CPU only, ~30 seconds):

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

Full test suite (including GPU correctness):

pytest tests/zorro_train/

GPU correctness only (with verbose output):

pytest tests/zorro_train/test_once_patcher.py -v

With coverage report:

pytest tests/zorro_train/ --cov=arctic_platform.rl.zorro_train --cov-report=html

Gradient Correctness Reference

The DeduplicatedActor class (in arctic_platform/rl/zorro_train/actor.py) includes a built-in gradient correctness check:

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

batch = create_dummy_batch(
    batch_size=8,
    num_unique_prompts=2,
    prompt_len=1024,
    response_len=256,
    device="cuda",
)

# Forward pass
output_dedup = actor.forward(batch, temperature=1.0)

# Compare with baseline (inside test harness)
# Gradients should match within numerical precision (typically 1e-5 relative error)

The demo script (demo.py) also runs this check automatically:

python arctic_platform/rl/zorro_train/demo.py

This will:

1. Run a gradient correctness test (deduplicated vs. baseline)
2. Report relative error in log probabilities and gradients
3. Optionally benchmark performance

Key Test Assertions

The test suite verifies:

Supported Configurations

The test suite covers:

Model families:

• Qwen3 (dense)
• Qwen3-MoE
• Qwen3-Next-MoE
• Qwen3.6 (dense)
• Qwen3.6-MoE

Attention implementations:

• eager (reference PyTorch attention)
• flash_attention_2 (production optimized)

Batch types:

• Padded sequences (with attention_mask)
• Unpadded sequences (for maximum efficiency)

Logits optimization modes:

• none (baseline, no optimization)
• compute (CPU-optimized logits computation)
• memory (distributed logits, requires process group)

Debugging Failed Tests

If test_once_patcher.py fails:

1. Check gradient error is not due to numerical precision:

   pytest tests/zorro_train/test_once_patcher.py -v -k "test_qwen3_model_once_patcher"
   

   Look for relative errors > 1e-4; smaller errors are acceptable due to floating-point arithmetic.

2. Isolate to a specific model/attention combination:

   pytest tests/zorro_train/test_once_patcher.py -v -k "eager" # eager attention only
   pytest tests/zorro_train/test_once_patcher.py -v -k "flash" # flash attention only
   

3. Test with unpadded batches (simpler case):

   pytest tests/zorro_train/test_once_patcher.py -v -k "unpadded"
   

4. Run the demo for manual inspection:

   python arctic_platform/rl/zorro_train/demo.py
   

   This prints detailed gradient comparisons.

If test_dedup.py fails:

1. Run with verbose output to see intermediate values:

   pytest tests/zorro_train/test_dedup.py -v -s
   

2. Check for off-by-one errors in position indices:

   • Verify prompt_groups correctly identifies identical prompts
   • Verify reconstruction_info metadata is correct

If test_perf.py is slow:

1. Reduce batch size or sequence lengths in the script
2. Use smaller models (e.g., Qwen3-0.6B)
3. Run only forward pass (comment out backward pass)

Test Utilities

Shared test helpers are in arctic_platform/rl/zorro_train/tests.py:

create_dummy_batch(...)

Creates a batch with controlled prompt deduplication for testing.

batch = create_dummy_batch(
    batch_size=8,                        # Total sequences
    num_unique_prompts=2,                # How many unique prompts
    prompt_len=1024,                     # Prompt tokens
    response_len=256,                    # Response tokens
    device="cuda",
    include_training_fields=True,        # Add attention_mask, position_ids, etc.
    add_padding=True,                    # Pad to max length
)

Returns a batch dict with:

• input_ids: [batch_size, max_len]
• attention_mask: [batch_size, max_len]
• position_ids: [batch_size, max_len]
• response_length: scalar (same for all sequences in batch)

compare_gradients(...)

Compares gradients between deduplicated and baseline implementations.

from arctic_platform.rl.zorro_train.tests import compare_gradients

relative_error = compare_gradients(
    dedup_grads,
    baseline_grads,
    atol=1e-5,
    rtol=1e-4,
)

Continuous Integration

In CI/CD pipelines, run:

# Fast checks (CPU only, ~1 minute)
pytest tests/zorro_train/test_dedup.py tests/zorro_train/test_seqlen_balancing.py

# Full suite (with GPU, ~5 minutes for one model/attention combo)
pytest tests/zorro_train/ -m "gpu" --timeout=300

To mark tests for CI environments:

import pytest
@pytest.mark.gpu  # GPU-only test
def test_once_patcher():
    ...

Expected Results

All tests should pass with:

• Log probability error: < 1e-5 relative
• Entropy error: < 1e-5 relative
• Gradient error: < 1e-4 relative
• Dedup round-trip error: 0 (exact match, no floating-point error)

Larger errors may indicate:

• Incorrect deduplication logic
• Gradient checkpoint incompatibility
• Numerical instability in attention implementation
• Attention mask not applied correctly