ZoRRo Train API Reference

DeduplicatedActor(
    model_name_or_path: str,
    device: str = "cuda",
    logits_optimization: str = "none",
    use_split_attention: bool = True,
    attn_implementation: str = "eager",
    world_size: int = 1,
    max_token_len: int = 4096,
    dtype: torch.dtype = torch.bfloat16,
)

forward(
    micro_batch: Dict[str, torch.Tensor],
    temperature: float = 1.0,
    calculate_entropy: bool = False,
) -> ModelOutput

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=4096,
    response_len=512,
    device="cuda",
    include_training_fields=True,
)

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]

compute_policy_loss_and_backward(
    micro_batch: Dict[str, torch.Tensor],
    temperature: float = 1.0,
    gradient_accumulation: int = 1,
) -> Dict[str, float]

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

patch(
    response_len: int,
    rollout_n: int,
    temperature: float,
) -> Qwen3ModelOncePatcher

train() -> None
eval() -> None

@staticmethod
find_prompt_groups(
    input_ids: torch.Tensor,
    response_length: int,
) -> Tuple[List[List[int]], torch.Tensor]

from arctic_platform.rl.zorro_train import ZoRRoTrain

input_ids = torch.tensor([
    [1, 2, 3, 4, 5, 6],  # Prompt A + Response 1
    [1, 2, 3, 7, 8, 9],  # Prompt A + Response 2
    [10, 11, 12, 13, 14, 15],  # Prompt B + Response 3
])

prompt_groups, unique_prompts = ZoRRoTrain.find_prompt_groups(
    input_ids, response_length=3
)
# prompt_groups: [[0, 1], [2]]
# unique_prompts: [[1, 2, 3], [10, 11, 12]]

@staticmethod
create_deduplicated_batch(
    input_ids: torch.Tensor,
    position_ids: torch.Tensor,
    response_length: int,
    prompt_groups: List[List[int]],
    unique_prompts: torch.Tensor,
    attention_mask: Optional[torch.Tensor] = None,
    use_unpad: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor, Dict]

prompt_groups, unique_prompts = ZoRRoTrain.find_prompt_groups(input_ids, 512)

dedup_ids, dedup_pos, recon_info = ZoRRoTrain.create_deduplicated_batch(
    input_ids=input_ids,
    position_ids=position_ids,
    response_length=512,
    prompt_groups=prompt_groups,
    unique_prompts=unique_prompts,
    attention_mask=attention_mask,
    use_unpad=True,
)

@staticmethod
reconstruct_sequences(
    dedup_hidden: torch.Tensor,
    reconstruction_info: Dict,
) -> torch.Tensor

@staticmethod
deduplicate_sequences(
    full_hidden: torch.Tensor,
    reconstruction_info: Dict,
) -> torch.Tensor

@staticmethod
extract_unpadded_responses_from_deduped_packed_ids(
    packed_ids: torch.Tensor,
    reconstruction_info: Dict,
    offset: int = 0,
) -> torch.Tensor

@staticmethod
responses_in_orig_sample_order(
    packed_responses: torch.Tensor,
    reconstruction_info: Dict,
) -> torch.Tensor

@staticmethod
attention_mask_to_flash_attn_params(
    attention_mask: torch.Tensor,
    device: torch.device = None,
) -> Dict[str, torch.Tensor]

# Self-attention with 3 sequences: [5, 3, 4] tokens each
flash_params = ZoRRoTrain.attention_mask_to_flash_attn_params(attention_mask)
# cu_seqlens_q: [0, 5, 8, 12]
# cu_seqlens_k: [0, 5, 8, 12]
# max_seqlen_q: 5
# num_sequences: 3

class ReconstructionInfo(dict):
    """Ensures object identity persists; forces use of update() to overwrite contents."""
    pass

from arctic_platform.rl.zorro_train import ReconstructionInfo

recon_info = ReconstructionInfo(a=1, b=2)
old_id = id(recon_info)

recon_info.update(a=2, b=3)  # Overwrites contents, same object
new_id = id(recon_info)

assert old_id == new_id  # Object identity preserved

def packed_ppo_policy_loss(
    log_prob: torch.Tensor,
    old_log_prob: torch.Tensor,
    advantages: torch.Tensor,
    ref_log_prob: Optional[torch.Tensor] = None,
    clip_ratio: float = 0.2,
    kl_loss_coef: float = 0.001,
) -> Tuple[torch.Tensor, Dict[str, float]]

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

# 1. Initialize actor (installs Qwen3ModelOncePatcher on first forward)
actor = DeduplicatedActor(
    model_name_or_path="Qwen/Qwen3-0.6B",
    device="cuda",
    logits_optimization="none",
    use_split_attention=True,
    attn_implementation="eager",
)

# 2. Create batch with shared prompts
batch = create_dummy_batch(
    batch_size=8,
    num_unique_prompts=2,  # Only 2 unique prompts for 8 samples
    prompt_len=4096,
    response_len=512,
    device="cuda",
    include_training_fields=True,
    add_padding=True,
)

# 3. Forward pass (automatically deduplicates)
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]

# 4. 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}")

# All ZoRRo Train tests
pytest tests/zorro_train/

# Specific test suites
pytest tests/zorro_train/test_dedup.py             # CPU: dedup algorithm round-trips
pytest tests/zorro_train/test_once_patcher.py      # GPU: forward/backward correctness
pytest tests/zorro_train/test_seqlen_balancing.py  # CPU: sequence-length balancing

# Performance benchmark
python arctic_platform/rl/zorro_train/test_perf.py