Training Pipeline & Processors

Overview

The Arctic RL training pipeline provides a composable, two-phase architecture for forward passes: post-forward processors compute derived outputs (like logprobs from logits), and loss functions compute scalar losses and metrics. Register processors by name to make them available at runtime via batch configuration.

Training Pipeline & Processors

The training pipeline in Arctic RL orchestrates forward passes, post-forward processing, loss computation, and backward passes in a unified, composable framework. Processors are registered by name into phase-specific registries and invoked at runtime based on a processing dict embedded in each batch.

Overview

Two-Phase Architecture

The pipeline executes in two phases:

Post-Forward Phase

Signature: (model_outputs, meta, device) → dict
Computes derived outputs from model activations (e.g., logprobs from logits, entropy)
Returns only new keys to add to model_outputs — not the full outputs
Keeps wire responses compact by excluding raw logits sent over HTTP

Loss Function Phase

Signature: (model_outputs, meta, device, config) → (loss, metrics)
Computes a scalar loss and an arbitrary metrics dict
Receives both original model outputs and post-processed outputs
Optional; can be skipped for forward-only (no-grad) passes

Batch Layout

When the processing pipeline is active, the batch dict should contain:

{
    "args": (),
    "batch": { ... },           # Passed to model.forward()
    "meta": { ... },            # Extra data: labels, advantages, flags
    "processing": {
        "post": ["name1", "name2"],  # Post-processors to apply
        "loss_fn": "name",            # Loss function to use
        "config": { ... },            # Loss function configuration
    },
}

Registering Processors

Post-Forward Processors

Decorate a function with @register_post_processor to register it as a post-forward processor:

from arctic_platform.rl.processors.pipeline import register_post_processor

@register_post_processor("my_processor")
def my_processor(model_outputs, batch, meta, device):
    """Compute derived outputs from model activations.
    
    Args:
        model_outputs: Dict with model output tensors (logits, logprobs, entropy, loss)
        batch: Dict with batch tensors (unpacked to 1D if packing was used)
        meta: Dict with metadata (labels, advantages, flags)
        device: Target device string (e.g. "cuda:0")
    
    Returns:
        Dict of new keys to add to model_outputs. Do NOT return the full
        model_outputs — only what this processor computed.
    """
    result = {}
    if "logits" in model_outputs:
        # Example: compute logprobs from logits
        result["logprobs"] = torch.log_softmax(model_outputs["logits"], dim=-1)
    return result

Loss Functions

Decorate a function with @register_loss_fn to register it as a loss function:

from arctic_platform.rl.processors.pipeline import register_loss_fn

@register_loss_fn("my_loss")
def my_loss(model_outputs, batch, meta, config, device):
    """Compute loss and metrics.
    
    Args:
        model_outputs: Dict with model outputs and post-processor results
        batch: Dict with batch tensors
        meta: Dict with metadata (updated with post-processor outputs)
        config: Dict with loss function configuration from processing["config"]
        device: Target device string
    
    Returns:
        Tuple of (loss_tensor, metrics_dict)
            loss_tensor: Scalar loss (requires_grad=True in train mode)
            metrics_dict: Dict of {metric_name: float}
    """
    loss = torch.tensor(0.0, device=device)
    metrics = {"loss": float(loss)}
    return loss, metrics

Dotted-Path Import

You can also pass a full dotted-path string as the processor name. The registry will import it on first use:

processing = {
    "post": ["mypackage.module.my_processor"],
    "loss_fn": "mypackage.module.my_loss",
}

The registry caches the imported function for subsequent calls.

Running the Pipeline

Main Entry Point

from arctic_platform.rl.processors.pipeline import run_pipeline

result = run_pipeline(
    engine,                          # DeepSpeed engine
    args=(),                         # Positional args to engine()
    batch=batch_dict,                # Microbatch tensors
    meta=meta_dict,                  # Metadata (labels, advantages, flags)
    processing=processing_dict,      # {"post": [...], "loss_fn": "...", "config": {...}}
    device="cuda:0",
    backward=True,                   # True: backward, False: no-grad, "loss_only": compute loss, no backward
    pack=True,                       # Auto-pack sequences into [1, T] for flash attention
    max_tokens_per_mb=4096,          # Token budget per microbatch
    return_tensors=False,            # If False, convert output tensors to numpy/lists
)

Parameters

engine
DeepSpeed training engine to call for the forward pass.

args
Positional arguments passed directly to engine().

batch
Dict with tensors belonging to this microbatch (e.g., input_ids, attention_mask, position_ids).

meta
Extra batch data — flags and tensors that shouldn't have been sharded (e.g., labels, advantages, pad_token_id). Available to the model, post-processors, and loss functions.

processing
Dict with keys:

"post" (optional): List of post-processor names to apply
"loss_fn" (required for backward): Loss function name
"config" (optional): Loss function configuration dict

device
Target device string (e.g., "cuda:0").

backward

True: Compute loss and call engine.backward(loss)
False: Forward-only, no-grad pass
"loss_only": Compute loss in train mode, return raw loss tensor with requires_grad=True, caller responsible for scaling and backward

pack
If True (default), automatically handle sequence packing for flash attention:

Splits batch into microbatches by token budget
Packs each microbatch from [B_mb, S] to [1, T]
Runs the pipeline on each
Unpacks and concatenates results

Packing is fully transparent: caller passes padded [B, S] tensors and receives [B, S] outputs. Requires attention_mask in batch.

Auto-detects already-packed input: if cu_seqlens is present (added by pack_for_dss), skips to avoid double-packing. Safe to pass pack=True unconditionally.

max_tokens_per_mb
Token budget per microbatch when pack=True. Sequences are grouped by first-fit-decreasing algorithm so no microbatch exceeds this limit. Default: 4096.

return_tensors
If False, convert output tensors to numpy/lists for serialization. If True, return PyTorch tensors.

Return Values

Forward-only (no loss):

{
    "batch": {
        "logprobs": tensor,  # Post-processor outputs
        ...
    },
    "metrics": {},
}

With loss:

{
    "avg_loss": float,
    "metrics": {
        "metric_name": float,
        ...
    },
    "batch": {  # Included if post-processors ran
        "logprobs": tensor,
        ...
    },
}

With backward="loss_only":

{
    "avg_loss": float,
    "metrics": {...},
    "batch": {...},
    "loss_tensor": tensor,  # Undetached, requires_grad=True
}

Sequence Packing

When pack=True, the pipeline automatically:

Splits the batch into microbatches by token budget (first-fit-decreasing)
Packs each microbatch from [B, S] to [1, T] for flash attention
Runs the full pipeline on each packed microbatch
Unpacks results back to [B, S] shape
Aggregates losses and metrics across microbatches:
- Token-level aggregation (default): weights by loss_mask.sum()
- Sequence-level aggregation: weights by batch size if configured

Auto-Detection

The pipeline auto-detects already-packed input by checking for cu_seqlens in the batch or meta. If present, packing is skipped. This makes pack=True safe to pass unconditionally — double-packing is automatically avoided.

Example: Custom Processor

from arctic_platform.rl.processors.pipeline import register_post_processor, register_loss_fn
import torch

@register_post_processor("compute_logprobs")
def compute_logprobs(model_outputs, batch, meta, device):
    """Compute logprobs from model logits."""
    logits = model_outputs["logits"]
    logprobs = torch.log_softmax(logits, dim=-1)
    return {"logprobs": logprobs}

@register_loss_fn("simple_loss")
def simple_loss(model_outputs, batch, meta, config, device):
    """Simple cross-entropy loss."""
    logprobs = model_outputs["logprobs"]
    labels = meta["labels"]
    
    loss = -logprobs.gather(-1, labels.unsqueeze(-1)).squeeze(-1).mean()
    metrics = {
        "loss": float(loss),
        "accuracy": float((logprobs.argmax(-1) == labels).float().mean()),
    }
    return loss, metrics

# Usage:
processing = {
    "post": ["compute_logprobs"],
    "loss_fn": "simple_loss",
    "config": {},
}

result = run_pipeline(
    engine, args, batch, meta, processing, device,
    backward=True,
)

Entropy Optimization

The pipeline automatically gates entropy computation when it cannot affect the loss:

If actor_config["entropy_coeff"] == 0.0, entropy is not computed
Only applies to loss computation paths (not forward-only)
Saves computation since entropy requires full-vocab softmax

Memory and Performance

The pipeline includes:

Memory profiling: see_memory_usage() checkpoints before/after forward, post-processing, and backward
Timing instrumentation: Per-phase timers (forward, post-forward, loss, backward) when ENABLE_TIMERS=True
Profiling support: Integrates with C, Torch, or no-op profilers via PROFILER_TYPE

Metrics are always returned as non-tensor types (float/int) for serialization.