Training

Learn how to write SFT and reinforcement learning training logic locally with PyTRIO and run it on a managed GPU cluster.

You can use PyTRIO to run large-scale LLM post-training from your CPU machine without dealing with infrastructure and environment complexity.

Post-training is usually split into supervised fine-tuning (SFT) and reinforcement learning (RL). PyTRIO's SFT workflow looks like this:

PyTRIO's RL workflow looks like this:

If we abstract the common pieces, data processing (Datum), forward and backward computation (forward_backward), and weight updates (optim_step) are always required. The following sections break down these three modules and show how to use them in SFT and RL.

Data Processing (Datum)

Before passing a dataset to the LLM and the loss function, you need to process it.

Datum is a wrapper for training data. You need to convert data into Datum before passing it to forward_backward. It contains two parts:

model_input: tokens passed to the model.
loss_fn_inputs: parameters passed to the loss function. Different post-training tasks require different parameters:
- For the SFT loss function (cross_entropy), pass:
  - weights: per-token loss weights, as a list of 0s and 1s (0 = ignore, 1 = compute loss). The prompt part is usually 0, and the output part is usually 1.
  - target_tokens: tokens shifted one position to the right, usually tokens[1:].
```
datum = trio.Datum(
    model_input=trio.ModelInput.from_ints(tokens=input_tokens),
    loss_fn_inputs=dict(
        weights=weights,
        target_tokens=target_tokens,
    )
)
```
- For RL loss functions (importance_sampling, ppo), pass:
  - target_tokens: target token IDs. The length must match model_input.
  - logprobs: sampling logprobs from rollout.
  - advantages: reward advantages. To exclude prompt tokens from the loss, set the prompt part of advantages to 0.
```
rl_datum = trio.Datum(
    model_input=trio.ModelInput.from_ints(tokens=input_tokens),
    loss_fn_inputs=dict(
        target_tokens=target_tokens,
        logprobs=sampling_logprobs,
        advantages=advantages,
    )
)
```

After the concept, here is a more practical SFT example. Suppose we have a dataset and want to build Datum objects for SFT:

Prepare a dataset:

examples = [
    {"input": "1+1", "output": "2"},
    {"input": "1+2", "output": "3"},
    {"input": "2*3", "output": "6"},
]

Convert each example into input_tokens, target_tokens, and weights, then place them into Datum. The final result is a processed_examples dataset made of Datum objects:

def process_example(example, tokenizer):
    prompt = f"Formula: {example['input']}\nAnswer:"

    prompt_tokens = tokenizer.encode(prompt, add_special_tokens=True)
    prompt_weights = [0] * len(prompt_tokens)
    completion_tokens = tokenizer.encode(f" {example['output']}\n\n", add_special_tokens=False)
    completion_weights = [1] * len(completion_tokens)

    tokens = prompt_tokens + completion_tokens
    weights = prompt_weights + completion_weights

    input_tokens = tokens[:-1]
    target_tokens = tokens[1:]
    weights = weights[1:]

    return trio.Datum(
        model_input=trio.ModelInput.from_ints(tokens=input_tokens),     
        loss_fn_inputs={                                                
            "weights": np.asarray(weights, dtype=np.float32),           
            "target_tokens": np.asarray(target_tokens, dtype=np.int32), 
        }
    )

processed_examples = [process_example(ex, tokenizer) for ex in examples]

Finally, pass processed_examples to training_client.forward_backward to start training:

import numpy as np

for _ in range(6):
    fwdbwd_future = training_client.forward_backward(processed_examples, "cross_entropy")
    optim_future = training_client.optim_step(trio.AdamParams(learning_rate=1e-4))
    ...

Accumulating Gradients (forward_backward)

After the dataset is processed, pass Datum objects to forward_backward to accumulate gradients.

As the name suggests, forward_backward runs a forward pass through the LLM, then runs the loss function and backward pass to produce gradients for the weight update.

forward_backward takes two parameters:

data: a list of Datum objects. Each item contains input tokens and loss-function parameters.
loss_fn: the loss function used to compute gradients. It can be a built-in function such as cross_entropy, importance_sampling, or ppo, or a custom function. See Loss Functions for the mathematical definitions.

fwdbwd_future = training_client.forward_backward(
    data=[datum],
    loss_fn="cross_entropy"
)
fwdbwd_future = fwdbwd_future.result()

Each result contains:

loss_fn_outputs: per-sample logprobs and elementwise_loss.
metrics: computed metrics, including loss_sum, loss_mean, and token_count.

fwdbwd_result.loss_fn_outputs
# [{'logprobs':..., 'elementwise_loss':...}, ...]

fwdbwd_result.metrics
# {'loss_sum': ..., 'loss_mean': ..., 'token_count': ...}

To compute per-token loss, you can use metrics:

loss_sum = fwdbwd_result.metrics['loss_sum']
weights = np.concatenate([example.loss_fn_inputs['weights'].tolist() for example in processed_examples])
print(f"Loss: {loss_sum / weights.sum():.4f}")

You can also compute it from logprobs in loss_fn_outputs:

logprobs = np.concatenate([output['logprobs'].tolist() for output in fwdbwd_result.loss_fn_outputs])
weights = np.concatenate([example.loss_fn_inputs['weights'].tolist() for example in processed_examples])
print(f"Loss per token: {-np.dot(logprobs, weights) / weights.sum():.4f}")

The two calculations are equivalent.

Updating Weights (optim_step)

After gradients are accumulated, use the optimizer to update weights.

After optim_step runs, PyTRIO updates weights according to the gradients accumulated by forward_backward:

optim_future = training_client.optim_step(
    trio.AdamParams(learning_rate=1e-4)
)
optim_future = optim_future.result()

optim_step has one parameter, adam_params, which should be a trio.AdamParams object.

Available trio.AdamParams parameters:

learning_rate: learning rate. Default 1e-4.
beta1: AdamW beta1. Default 0.9.
beta2: AdamW beta2. Default 0.95.
eps: AdamW epsilon. Default 1e-12.
weight_decay: weight decay coefficient. Default 0.
grad_clip_norm: upper bound for gradient clipping norm. Default 0.

Supervised Fine-Tuning (SFT)

The following SFT example teaches the model to answer that PyTRIO is an AI Infra product:

import pytrio as trio
import numpy as np

# 1. Connect to TRIO
service_client = trio.ServiceClient()

# 2. Create a training client
base_model = "Qwen/Qwen3.5-4B"
training_client = service_client.create_lora_training_client(
    base_model=base_model,
    rank=32,
)

# 3. Dataset: teach the LLM what TRIO is
examples = [
    {"input": "what is trio", "output": "trio is emotionmachine's AI Infra products."},
    {"input": "can you explain what trio is", "output": "trio is an AI infra product developed by emotionmachine."},
    {"input": "tell me about trio", "output": "trio is a product from emotionmachine that provides AI Infra capabilities."},
]

# 4. Get the tokenizer
print("Loading tokenizer...")
tokenizer = training_client.get_tokenizer()
print("Tokenizer finish")

# 5. Process the dataset into the format required for training
def process_example(example: dict, tokenizer) -> trio.Datum:
    prompt = f"Question: {example['input']}\nAnswer:"

    prompt_tokens = tokenizer.encode(prompt, add_special_tokens=True)
    prompt_weights = [0] * len(prompt_tokens)

    completion_tokens = tokenizer.encode(f" {example['output']}\n\n", add_special_tokens=False)
    completion_weights = [1] * len(completion_tokens)

    tokens = prompt_tokens + completion_tokens
    weights = prompt_weights + completion_weights

    input_tokens = tokens[:-1]
    target_tokens = tokens[1:]
    weights = weights[1:]

    # Convert to the format required by TRIO training
    return trio.Datum(
        model_input=trio.ModelInput.from_ints(tokens=input_tokens),
        loss_fn_inputs={
            "weights": np.asarray(weights, dtype=np.float32),
            "target_tokens": np.asarray(target_tokens, dtype=np.int32),
        },
    )


processed_examples = [process_example(ex, tokenizer) for ex in examples]

# 6. Train
print("Start Training")
for iter in range(15):
    fwdbwd_future = training_client.forward_backward(processed_examples, "cross_entropy")  # Forward/backward
    optim_future = training_client.optim_step(trio.AdamParams(learning_rate=1e-4))  # Adam optimizer update

    fwdbwd_result = fwdbwd_future.result()
    optim_result = optim_future.result()

    logprobs = np.concatenate([output['logprobs'].tolist() for output in fwdbwd_result.loss_fn_outputs])
    weights = np.concatenate([example.loss_fn_inputs['weights'].tolist() for example in processed_examples])
    print(f"Iter{iter+1} Loss per token: {-np.dot(logprobs, weights) / weights.sum():.4f}")

# 7. Sampling and evaluation
print("Start Sampling")
sampling_base_client = service_client.create_sampling_client(base_model=base_model)
sampling_sft_client = training_client.save_weights_and_get_sampling_client()

prompt = trio.ModelInput.from_ints(tokenizer.encode("Question: what is trio\nAnswer:"))
params = trio.SamplingParams(max_tokens=20, temperature=0.0)

future_base = sampling_base_client.sample(prompt=prompt, sampling_params=params, num_samples=1)
result_base = future_base.result()
future_sft = sampling_sft_client.sample(prompt=prompt, sampling_params=params, num_samples=1)
result_sft = future_sft.result()

print("Base Responses:")
print(f"{repr(result_base.sequences[0].text)}")

print("SFT Responses:")
print(f"{repr(result_sft.sequences[0].text)}")

Reinforcement Learning (RL)

The following RL example trains the model to answer math questions in the required format:

import re
import pytrio as trio
import numpy as np

# 1. Connect to TRIO
service_client = trio.ServiceClient()

# 2. Create a training client
base_model = "Qwen/Qwen3.5-4B"
training_client = service_client.create_lora_training_client(
    base_model=base_model,
    rank=32,
)

# 3. Dataset: simple math questions
dataset = [
    ("What is 2 + 3?", 5),
    ("What is 7 - 4?", 3),
    ("What is 6 * 8?", 48),
    ("What is 12 / 3?", 4),
    ("Solve for x: x + 5 = 9", 4),
    ("Solve for x: 2x = 10", 5),
    ("What is 3 squared?", 9),
    ("What is the square root of 81?", 9),
    ("What is 15 + 27?", 42),
    ("What is 100 - 58?", 42),
]

eval_dataset = [
    ("Solve for x: x + 7 = 12", 5),
    ("What is 9 * 7?", 63),
    ("What is 81 / 9?", 9),
    ("What is 14 + 28?", 42),
]

# 4. Get the tokenizer
print("Loading tokenizer...")
tokenizer = training_client.get_tokenizer()
print("Tokenizer finish")

# 6. Parse numeric answers from model output
def parse_number(text: str):
    match = re.fullmatch(r"-?\d+(?:\.\d+)?", text.strip())
    return float(match.group()) if match else None

# 7. Reward function
def compute_reward(text: str, gold: float) -> float:
    pred = parse_number(text)
    if pred is None:
        return -1.0
    if abs(pred - gold) < 1e-6:
        return 2.0
    return -0.5

# 8. Convert to a numpy array for loss statistics
def to_np(x):
    values = x.tolist() if hasattr(x, "tolist") else x
    return np.array([0.0 if value is None else value for value in values], dtype=float)

# 9. Convert one rollout result into the Datum format required by TRIO training
def process_rollout(prompt_tokens, completion_tokens, completion_logprobs, reward_value):
    tokens = prompt_tokens + completion_tokens

    completion_logprobs = [0.0 if value is None else float(value) for value in completion_logprobs]
    old_logprobs = ([0.0] * len(prompt_tokens) + completion_logprobs)[:len(tokens)]
    old_logprobs += [0.0] * (len(tokens) - len(old_logprobs))

    input_tokens = tokens[:-1]
    target_tokens = tokens[1:]
    old_logprobs = old_logprobs[1:]
    advantages = ([0.0] * len(prompt_tokens) + [reward_value] * len(completion_tokens))[1:]

    return trio.Datum(
        model_input=trio.ModelInput.from_ints(tokens=input_tokens),
        loss_fn_inputs=dict(
            target_tokens=target_tokens,
            logprobs=old_logprobs,
            advantages=advantages,
        ),
    )

# 10. RL training
print("Start RL Training")

for iter in range(15):
    sampler = training_client.save_weights_and_get_sampling_client()
    processed_examples = []
    rewards = []
    correct = 0
    total = 0

    for question, gold in dataset:
        prompt_tokens = tokenizer.encode(f"Question: {question}\nReturn only the final numeric answer.\nAnswer:", add_special_tokens=True)

        future_sample = sampler.sample(
            prompt=trio.ModelInput.from_ints(prompt_tokens),
            sampling_params=trio.SamplingParams(max_tokens=8, temperature=0.7),
            num_samples=4,
        )
        sample_result = future_sample.result()

        for sequence in sample_result.sequences:
            reward_value = compute_reward(sequence.text, float(gold))
            pred = parse_number(sequence.text)

            rewards.append(reward_value)
            total += 1
            correct += pred is not None and abs(pred - gold) < 1e-6

            completion_tokens = list(sequence.tokens)

            if completion_tokens:
                processed_examples.append(
                    process_rollout(
                        prompt_tokens=prompt_tokens,
                        completion_tokens=completion_tokens,
                        completion_logprobs=sequence.logprobs,
                        reward_value=reward_value,
                    )
                )

    print(
        f"Iter{iter+1} | Reward: {np.mean(rewards):.4f} | "
        f"Acc: {correct / max(total, 1):.4f} | Samples: {len(processed_examples)}"
    )

    fwdbwd_future = training_client.forward_backward(processed_examples, "importance_sampling")
    optim_future = training_client.optim_step(trio.AdamParams(learning_rate=1e-5))

    fwdbwd_result = fwdbwd_future.result()
    optim_result = optim_future.result()

    logprobs = np.concatenate([to_np(output["logprobs"]) for output in fwdbwd_result.loss_fn_outputs])
    old_logprobs = np.concatenate([to_np(example.loss_fn_inputs["logprobs"]) for example in processed_examples])
    advantages = np.concatenate([to_np(example.loss_fn_inputs["advantages"]) for example in processed_examples])

    mask = advantages != 0
    loss = -np.sum(np.exp(logprobs[mask] - old_logprobs[mask]) * advantages[mask]) / mask.sum()
    print(f"Iter{iter+1} IS Loss: {loss:.4f}\n")

# 11. Sampling and evaluation
print("Start Evaluation")

sampling_base_client = service_client.create_sampling_client(base_model=base_model)
sampling_rl_client = training_client.save_weights_and_get_sampling_client()

for question, gold in eval_dataset:
    prompt = trio.ModelInput.from_ints(
        tokenizer.encode(f"Question: {question}\nReturn only the final numeric answer.\nAnswer:", add_special_tokens=True)
    )

    future_base = sampling_base_client.sample(prompt=prompt, sampling_params=trio.SamplingParams(max_tokens=8, temperature=0.0), num_samples=1)
    future_rl = sampling_rl_client.sample(prompt=prompt, sampling_params=trio.SamplingParams(max_tokens=8, temperature=0.0), num_samples=1)

    result_base = future_base.result()
    result_rl = future_rl.result()

    base_text = result_base.sequences[0].text.strip()
    rl_text = result_rl.sequences[0].text.strip()

    print("=" * 60)
    print(f"Q: {question} | Gold: {gold}")
    print(f"Base: {repr(base_text)} -> {parse_number(base_text)}")
    print(f"RL:   {repr(rl_text)} -> {parse_number(rl_text)}")

Training results:

Iter1 | Reward: -0.5375 | Acc: 0.1500 | Samples: 40
Iter1 IS Loss: 0.7964

Iter2 | Reward: -0.5375 | Acc: 0.1500 | Samples: 40
Iter2 IS Loss: 0.7938

...

Iter15 | Reward: 1.3250 | Acc: 0.7750 | Samples: 40
Iter15 IS Loss: -0.7395

Start Evaluation
============================================================
Q: Solve for x: x + 7 = 12 | Gold: 5
Base: '5\n\nQuestion: Solve for x' -> None
RL:   '5' -> 5.0
============================================================
Q: What is 9 * 7? | Gold: 63
Base: '63. 63.' -> None
RL:   '63' -> 63.0
============================================================
Q: What is 81 / 9? | Gold: 9
Base: '9\n\nQuestion: What is' -> None
RL:   '9' -> 9.0
============================================================
Q: What is 14 + 28? | Gold: 42
Base: '42.' -> None
RL:   '42' -> 42.0

After RL training, the model better follows the required output format.

Multimodal Training

See the multimodal documentation for details on multimodal training.