Metadata-Version: 2.4
Name: pytrio
Version: 0.2.9
Summary: training is on
Author: Emotion Machine
License-Expression: MIT
Keywords: training,machine learning,post train
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: click>=8.2.0
Requires-Dist: httpx>=0.28.1
Requires-Dist: jinja2>=3.1.0
Requires-Dist: modelscope>=1.15.0
Requires-Dist: msgpack>=1.1.0
Requires-Dist: numpy>=1.26.0
Requires-Dist: pydantic>=2.12.5
Requires-Dist: tomli>=2.3.0
Requires-Dist: transformers>=4.57.3
Provides-Extra: full

<div align="center">

<p>
  <img src="https://pytrio.com/assets/trio-logo.svg" alt="PyTRIO" width="280">
</p>

### Connect to WiFi. Train large language models.

<p>
  <a href="https://pypi.org/project/pytrio/"><img src="https://img.shields.io/pypi/v/pytrio?color=3776AB&label=PyPI&logo=pypi&logoColor=white" alt="PyPI version"></a>
  <a href="https://pypi.org/project/pytrio/"><img src="https://img.shields.io/badge/python-3.10%2B-3776AB?logo=python&logoColor=white" alt="Python 3.10+"></a>
  <a href="https://pypi.org/project/pytrio/"><img src="https://img.shields.io/pypi/l/pytrio?color=22c55e" alt="License"></a>
</p>

**A clean Python SDK for model training, inference, and post-training experiments.**

<p>
  <a href="https://pytrio.com/">Website</a> ·
  <a href="https://docs.pytrio.com/">Docs</a>
</p>

</div>

## ✨ 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](https://pytrio.com/):

```bash
pip install pytrio
trio login
```

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

```python
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:

```python
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

| Goal | Start here |
| --- | --- |
| Connect to the training engine | `trio.ServiceClient` |
| Create and update a LoRA policy | `trio.TrainingClient` |
| Generate rollouts or compute logprobs | `trio.SamplingClient` |
| Define token, text, or image input | `trio.ModelInput` |
| Control rollout generation | `trio.SamplingParams` |
| Save, resume, and download checkpoints | `trio.RestClient` |

## 🔌 OpenAI-compatible inference

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

```python
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)
```

<div align="center">

### Ready to start training?

[**Read the official docs →**](https://docs.pytrio.com/)

**Training Is On.** ⚡

</div>
