Supported Models

Supported Models

Currently Supported

ZoRRo Train provides optimized prompt deduplication for the following model families:

• qwen3 — Dense Qwen3 models
• qwen3-moe — Mixture of Experts variant
• qwen3-next-moe — Next-generation MoE architecture
• qwen3.6 — Qwen3.6 dense models
• qwen3.6-moe — Qwen3.6 MoE variant

These span dense, MoE (Mixture of Experts), and hybrid (linear + full attention) architectures.

Model Requirements

To use ZoRRo Train with a supported model, ensure you have:

• PyTorch and Transformers installed:

  pip install torch transformers
  

• (Optional) Flash Attention for best performance:

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

Attention Implementation Support

ZoRRo Train is flexible and supports multiple attention implementations:

• Eager attention (attn_implementation="eager") — Standard PyTorch attention
• Flash Attention 2 (attn_implementation="flash_attention_2")
• Flash Attention 3 (attn_implementation="flash_attention_3")

Specify the attention implementation when initializing the actor:

from arctic_platform.rl.zorro_train.actor import DeduplicatedActor

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

Example: Loading a Supported Model

from arctic_platform.rl.zorro_train.actor import DeduplicatedActor

# Initialize with a Qwen3 model
actor = DeduplicatedActor(
    model_name_or_path="Qwen/Qwen3-0.6B",
    device="cuda",
    logits_optimization="none",
    use_split_attention=True,
    attn_implementation="eager",
)

# The model is now ready for deduplicated training
output = actor.forward(batch, temperature=1.0, calculate_entropy=True)

Future Model Support

More model families will be added in the future. If you need support for a specific architecture, refer to the qwen_model_patcher.py implementation as a reference for extending ZoRRo Train to new models.

Testing Model Compatibility

The test suite validates ZoRRo Train across multiple configurations:

# Run all ZoRRo Train tests
pytest tests/zorro_train/

# Test against tiny-random Qwen3 checkpoints (dense, MoE, hybrid)
# with flash_attention_2 and eager attention
pytest tests/zorro_train/test_once_patcher.py

The test suite automatically sweeps:

• Dense, MoE, and hybrid architectures
• flash_attention_2 and eager implementations
• Padded and unpadded batches
• All logits optimization modes (none, memory, compute)