TRL (Transformers Reinforcement Learning)

The GOLDTrainer (located in trl.experimental.gold) is an extension of SFTTrainer designed for knowledge distillation between student and teacher models. Its primary advantage is supporting cross-tokenizer alignment, allowing you to distill knowledge from a teacher to a student even if they use different tokenizers or belong to different model families (e.g., a LLaMA student with a Qwen teacher).

Key features include:

• Cross-tokenizer alignment: It aligns textual spans and merges logits so no completion tokens are lost during distillation.
• Hybrid ULD loss: When uld_use_hybrid_loss is enabled, it uses exact vocabulary matches where possible and falls back to sorted-probability Universal Logit Distillation (ULD) for unmatched tokens.
• GKD Integration: It inherits on-policy/off-policy scheduling from GKDTrainer.

Note: As part of the trl.experimental namespace, the API is subject to change.

Sequence Parallelism (also called Context Parallelism) is a technique to enable training with very long sequences by splitting the sequence dimension across multiple GPUs. This allows training with sequences that exceed the memory capacity of a single GPU.

TRL distinguishes between two types of sequence splitting:

• Context Parallelism (CP): Implemented as Ring Attention using the FSDP2 backend. It splits sequences across GPUs.
• Sequence Parallelism (SP): Implemented as ALST/Ulysses using the DeepSpeed backend. It uses attention head parallelism.

Key Terminology:

• Global sequence length: The full sequence length before splitting. In TRL, max_seq_length (or max_length) refers to this value.
• Micro sequence length: The sequence length processed by each individual GPU after splitting.

Note that these parallelism dimensions multiply. For example, if TP=2 and CP=2, you require 4 GPUs ($2 \times 2$).

The Harmony format extends conversational datasets to support richer structures like reasoning (thinking), tool calls, and metadata.

Key components:

• Developer role: Used for high-level instructions and tool lists (similar to a system prompt).
• Channels: Separates assistant output into analysis (for thinking), final (for content), and commentary (for tool calls).
• Reasoning effort: Can be set to "low", "medium", or "high".
• Model identity: Defines the assistant's persona.

When using tokenizer.apply_chat_template, you can pass reasoning_effort and model_identity to format the messages correctly for models supporting Harmony.

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("openai/gpt-oss-20b")

messages = [
    {"role": "developer", "content": "Use a friendly tone."},
    {"role": "user", "content": "What is the meaning of life?"},
    {"role": "assistant", "thinking": "Deep reflection...", "content": "The final answer is..."},
]

print(
    tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        reasoning_effort="low",
        model_identity="You are HuggingGPT, a large language model trained by Hugging Face.",
    )
)

You can train a single model across multiple environments simultaneously by using a meta-environment class and environment routing via the dataset.

Implementation Steps:

1. Dataset Routing: Add a column (e.g., "env") to your dataset containing the identifier for the target environment for each sample.
2. Meta-Environment Class: Create a class that wraps all possible environments. In its reset(**kwargs) method, use kwargs.get("env") to select and initialize the correct environment.
   • Lazy Initialization: Initialize clients inside reset() rather than __init__() to avoid unnecessary connections.
   • Cleanup: Close the previous environment's client before opening a new one to prevent server capacity errors.
   • Tool Exposure: The meta-class should expose all tools from all sub-environments. The model will learn to use the correct ones based on the system prompt.
3. Per-Environment Reward Functions: Define separate reward functions for each environment. Each function should return None for samples that do not belong to its environment. TRL handles None values using nansum/nanmean, ensuring each sample is only scored by its relevant reward.

Monitoring

When training with multiple environments, monitor per-reward-function metrics (e.g., train/reward_func_0, train/reward_func_1) instead of the combined train/reward, as the combined metric will be noisy due to alternating environments.

# 1. Dataset with routing
dataset = Dataset.from_dict({
    "prompt": ([[...]] * n) + ([[...]] * n),
    "env": ["wordle"] * n + ["catch"] * n,
})

# 2. Meta-environment routing
class MultiEnv:
    def reset(self, **kwargs) -> str | None:
        self.active = kwargs.get("env", "wordle")
        if self.active == "wordle":
            # Initialize wordle client...
            return observation
        elif self.active == "catch":
            # Initialize catch client...
            return observation

# 3. Per-environment rewards
def wordle_reward(environments, **kwargs) -> list[float | None]:
    return [env.reward if env.active == "wordle" else None for env in environments]

def catch_reward(environments, **kwargs) -> list[float | None]:
    return [env.reward if env.active == "catch" else None for env in environments]

For stateful training where the environment manages its own state and rewards, use the environment_factory argument. GRPOTrainer creates one environment instance per rollout and exposes its public methods as tools.

Environment Class Requirements:

• reset(self, **kwargs) -> str | None (Required): Called at the start of each rollout. If it returns a string, that string becomes the user prompt. It can also receive keyword arguments from the train_dataset rows.
• get_reward(self) -> float (Optional): Returns a float representing the reward based on the environment's internal state. Called once per completed rollout.
• Public Methods: Any other public methods defined in the class are automatically exposed to the model as tools.

Note: environment_factory requires transformers>=5.2.0.

import random
from trl import GRPOConfig, GRPOTrainer

class IncrementEnv:
    def reset(self, **kwargs) -> str | None:  # required
        self.counter = 0
        self.target = random.randint(1, 6)
        return f"Increment the counter by {self.target}."

    def get_reward(self) -> float:  # optional
        return float(self.counter == self.target)

    def increment(self, step: int) -> int:  # exposed as tool
        """
        Increment the internal counter.

        Args:
            step: Value to add to the counter.

        Returns:
            The updated counter value.
        """
        self.counter += step
        return self.counter

trainer = GRPOTrainer(
    model="Qwen/Qwen3-0.6B",
    args=GRPOConfig(max_steps=1000, chat_template_kwargs={"enable_thinking": False}),
    environment_factory=IncrementEnv,
)
trainer.train()