ADu2021/fastcurl-curriculum-rl-efficient-reasoning

Core Concept

FastCuRL addresses the training inefficiency of large reasoning models by jointly optimizing context length and training data complexity through a curriculum learning framework. The key insight is that controlling context length and selecting data based on problem complexity can significantly improve RL training efficiency while generating more concise Chain-of-Thought (CoT) outputs. The approach uses a cyclical compress-extend strategy that iteratively refines reasoning outputs.

Architecture Overview

FastCuRL integrates three main components:

• Group Relative Policy Optimization (GRPO): A resource-efficient RL algorithm that eliminates the need for a critic model by computing advantages from group-level scores rather than individual baseline scores
• Complexity-Aware Data Selection: Divides training data into three categories (L1, L2, L3) based on input prompt length correlation with output complexity
• Stage-wise Context Scaling: Alternates between compress phases (reducing context length) and extend phases (increasing context length) across multiple training stages

Implementation Steps

1. Group Relative Policy Optimization (GRPO) Algorithm

The GRPO objective maximizes policy improvements with KL divergence regularization and optional entropy bonus:

# Simplified GRPO loss computation
import torch
import torch.nn.functional as F

def compute_grpo_loss(policy_logits, old_policy_logits, rewards,
                      kl_coefficient=0.02, entropy_coeff=0.01, epsilon=0.2):
    """
    Compute Group Relative Policy Optimization loss.

    Args:
        policy_logits: logits from current policy
        old_policy_logits: logits from reference policy
        rewards: shape [batch, group_size]
        kl_coefficient: KL penalty weight
        entropy_coeff: entropy bonus weight
        epsilon: clipping threshold
    """
    # Normalize advantages across group
    mean_reward = rewards.mean(dim=-1, keepdim=True)
    std_reward = rewards.std(dim=-1, keepdim=True) + 1e-8
    advantages = (rewards - mean_reward) / std_reward

    # Compute log probability ratio
    log_ratio = policy_logits - old_policy_logits
    ratio = torch.exp(log_ratio)

    # Clipped surrogate objective
    surr1 = ratio * advantages
    surr2 = torch.clamp(ratio, 1-epsilon, 1+epsilon) * advantages
    policy_loss = -torch.min(surr1, surr2).mean()

    # KL divergence penalty
    kl_loss = kl_coefficient * ((ratio - 1) - log_ratio).mean()

    # Optional entropy bonus for exploration
    entropy = -torch.sum(torch.softmax(policy_logits, -1) *
                         torch.log_softmax(policy_logits, -1), dim=-1)
    entropy_loss = -entropy_coeff * entropy.mean()

    return policy_loss + kl_loss + entropy_loss

2. Complexity-Aware Data Selection

Segment the training dataset by input prompt length to match problem complexity with context requirements:

def partition_dataset_by_complexity(dataset,
                                   complexity_thresholds=None):
    """
    Partition dataset into complexity levels based on prompt length.

    Args:
        dataset: list of (prompt, response) tuples
        complexity_thresholds: percentile cutoffs [25, 75] for L1/L2/L3

    Returns:
        dict with keys 'L1' (short), 'L2' (medium), 'L3' (long)
    """
    prompt_lengths = [len(item['prompt'].split()) for item in dataset]

    if complexity_thresholds is None:
        # Use percentile-based splitting
        p25 = np.percentile(prompt_lengths, 25)
        p75 = np.percentile(prompt_lengths, 75)
        complexity_thresholds = [p25, p75]

    partitions = {'L1': [], 'L2': [], 'L3': []}

    for item, length in zip(dataset, prompt_lengths):
        if length <= complexity_thresholds[0]:
            partitions['L1'].append(item)
        elif length <= complexity_thresholds[1]:
            partitions['L2'].append(item)
        else:
            partitions['L3'].append(item)

    return partitions

3. Stage-wise Context Scaling Training Loop

Implement the multi-stage training with alternating context lengths:

def train_with_stage_scaling(model, train_partitions, num_stages=5,
                            context_lengths=None, epochs_per_stage=1):
    """
    Train model using stage-wise context scaling curriculum.

    Args:
        model: the reasoning model to train
        train_partitions: dict with 'L1', 'L2', 'L3' datasets
        num_stages: number of training stages
        context_lengths: list of context window sizes per stage
        epochs_per_stage: epochs per stage (typically 1 for efficiency)
    """
    if context_lengths is None:
        # Example: compress-extend cycle
        context_lengths = [8192, 8192, 16384, 16384, 24576]

    stage_configs = [
        # Stage 1: L1 (short) at 8K context
        {'dataset': train_partitions['L1'], 'context': context_lengths[0]},
        # Stage 2: L2 (medium) at 8K context
        {'dataset': train_partitions['L2'], 'context': context_lengths[1]},
        # Stage 3: L2 (medium) at 16K context (extend phase)
        {'dataset': train_partitions['L2'], 'context': context_lengths[2]},
        # Stage 4: L3 (long) at 16K context
        {'dataset': train_partitions['L3'], 'context': context_lengths[3]},
        # Stage 5: L3 (long) at 24K context
        {'dataset': train_partitions['L3'], 'context': context_lengths[4]},
    ]

    for stage_idx, config in enumerate(stage_configs):
        print(f"Training Stage {stage_idx + 1}: "
              f"dataset={list(train_partitions.keys())[stage_idx % 3]}, "
              f"context={config['context']}")

        # Set model context window
        model.config.max_position_embeddings = config['context']

        # Train on this stage's dataset
        for epoch in range(epochs_per_stage):
            for batch in DataLoader(config['dataset'], batch_size=8):
                # Forward pass with GRPO
                loss = compute_grpo_loss(batch, model)
                loss.backward()
                optimizer.step()
                optimizer.zero_grad()

Practical Guidance

When to Use FastCuRL

• Training reasoning models on mathematical or complex problem-solving tasks
• Resource-constrained environments where reducing training steps is critical
• When you need both quality reasoning and concise CoT outputs
• Fine-tuning R1-distilled models (e.g., DeepSeek-R1-Distill-Qwen)

When NOT to Use

• Tasks requiring long, detailed explanations as a feature
• Limited access to compute for initial dataset partitioning and analysis
• When training data doesn't show prompt-length to output-length correlation
• Few-shot scenarios with very limited training data

Hyperparameters & Configuration

• KL coefficient (β): Typically 0.02; controls policy divergence from reference
• Entropy coefficient (α): Around 0.01 when needed to prevent entropy collapse
• Epsilon (clip threshold): Standard PPO value of 0.2
• Context lengths: Start with 8K→8K→16K→16K→24K sequence for compress-extend
• Epochs per stage: Keep at 1 for efficiency in low-resource settings
• Group size (G): Number of rollout responses per prompt; typically 4-8

Common Pitfalls

• Not analyzing data distribution: Failing to verify input-output length correlation before applying complexity-based partitioning
• Entropy collapse: Occurs when context window is too large for current dataset; add entropy bonus term to mitigate
• Stage ordering: Placing complex data (L3) too early can cause training instability; start with simpler data
• Context thrashing: Rapidly changing context lengths without stable intermediate stages can prevent convergence
• Ignoring clipping rates: High percentage of clipped outputs indicates context is too small; monitor clipping ratio per stage

Reference

• DeepSeek-AI. 2025. DeepSeek-R1: Incentivizing Reasoning Capability.
• Shao et al. 2024. Group Relative Policy Optimization (GRPO).
• Luo et al. 2025. DeepScaleR: Iterative Context Scaling for Efficient LLM Training.
• Official FastCuRL repository: https://github.com/nick7nlp/FastCuRL