pytrio 0.2.9

✨ What is PyTRIO?

PyTRIO lets you focus on the parts of post-training that matter—your data and algorithms—while the engine handles distributed execution, scheduling, fault tolerance, and GPU infrastructure.

Write your training loop on a CPU machine, choose a base model with one string, and keep full control of your loss function, optimizer, rollout strategy, and experiment logic.

• 🧪 Built for post-training: Run SFT, RL, preference optimization, and custom objectives
• ⚡ Async research: Keep working locally while training runs remotely
• 🚀 Managed scale: Train across GPUs without managing CUDA or clusters
• 🔁 Fast iteration: Sample fresh weights, checkpoint, and resume anytime
• 🔌 Application ready: Serve trained weights through OpenAI-compatible APIs

🚀 Train, then sample

Install the SDK and sign in with an API key from the Trio console:

pip install pytrio
trio login

The same training client can update LoRA weights and turn the latest policy into a sampler immediately:

import pytrio as trio

# Connect to the PyTRIO training engine.
service = trio.ServiceClient()
trainer = service.create_lora_training_client(
    base_model="Qwen/Qwen3.5-4B",
    rank=32,
)
tokenizer = trainer.get_tokenizer()

# Prepare one supervised training example.
tokens = tokenizer.encode(
    "Question: what is Trio?\nAnswer: a model training platform."
)
batch = [
    trio.Datum(
        model_input=trio.ModelInput.from_ints(tokens[:-1]),
        loss_fn_inputs={
            "target_tokens": tokens[1:],
            "weights": [1.0] * (len(tokens) - 1),
        },
    )
]

# Every step of the training loop stays in your hands.
for step in range(10):
    trainer.forward_backward(batch, "cross_entropy").result()
    trainer.optim_step(trio.AdamParams(learning_rate=1e-4)).result()

# Save the latest policy and sample from it immediately.
sampler = trainer.save_weights_and_get_sampling_client()
prompt_tokens = tokenizer.encode("Question: what is Trio?\nAnswer:")
result = sampler.sample(
    prompt=trio.ModelInput.from_ints(prompt_tokens),
    sampling_params=trio.SamplingParams(max_tokens=32, temperature=0.0),
    num_samples=1,
).result()

print(result.sequences[0].text)

🎲 Rollouts in, policy updates out

PyTRIO keeps rollout generation and policy updates in one workflow. Sample multiple trajectories from the latest policy, score them with your reward function, then train directly from their tokens and logprobs:

question = "What is 6 * 8?"
prompt_tokens = tokenizer.encode(
    f"Question: {question}\nReturn only the final numeric answer.\nAnswer:"
)

# Generate a group of rollouts from the latest policy.
rollouts = sampler.sample(
    prompt=trio.ModelInput.from_ints(prompt_tokens),
    sampling_params=trio.SamplingParams(max_tokens=8, temperature=0.7),
    num_samples=4,
).result()

training_data = []
for sequence in rollouts.sequences:
    reward = 2.0 if sequence.text.strip() == "48" else -1.0
    completion_tokens = list(sequence.tokens)
    tokens = prompt_tokens + completion_tokens
    old_logprobs = (
        [0.0] * len(prompt_tokens)
        + [0.0 if value is None else float(value) for value in sequence.logprobs]
    )
    advantages = [0.0] * len(prompt_tokens) + [reward] * len(completion_tokens)

    training_data.append(
        trio.Datum(
            model_input=trio.ModelInput.from_ints(tokens[:-1]),
            loss_fn_inputs={
                "target_tokens": tokens[1:],
                "logprobs": old_logprobs[1:],
                "advantages": advantages[1:],
            },
        )
    )

# Update the policy from rewarded rollouts.
trainer.forward_backward(
    training_data,
    loss_fn="importance_sampling",
).result()
trainer.optim_step(
    trio.AdamParams(learning_rate=1e-5),
).result()

Refresh the sampler from the latest weights on every iteration to build an on-policy loop. The same primitives support SFT, PPO-style objectives, custom losses, checkpointing, and evaluation without hiding the algorithm behind a black box.

🧭 API primitives

🔌 OpenAI-compatible inference

Move a trained policy into an application with a model identifier and the standard OpenAI client:

from openai import OpenAI

client = OpenAI(
    base_url="https://pytrio.com/api/openai/v1",
    api_key="TRIO_API_KEY",
)
response = client.chat.completions.create(
    model="trio://your-model/your-version",
    messages=[{"role": "user", "content": "Hello!"}],
)
print(response.choices[0].message.content)