pytrio.TrainingClient

class TrainingClient:
    def __init__(
        self,
        task_id: str,
        base_model: str,
        lora: LoraRunSpec,
    ):

TrainingClient runs LoRA training. Create it with ServiceClient.create_lora_training_client().

import pytrio as trio

client = trio.ServiceClient()
training_client = client.create_lora_training_client(base_model="Qwen/Qwen3.5-4B")

tokenizer = training_client.get_tokenizer()
tokens = tokenizer.encode("The meaning of life is")
input_tokens = tokens[:-1]
target_tokens = tokens[1:]

data = [
    trio.Datum(
        model_input=trio.ModelInput.from_ints(input_tokens),
        loss_fn_inputs={"target_tokens": target_tokens},
    )
]

# Forward and backward pass
future = training_client.forward_backward(data=data)
output = future.result()

# Optimizer update
training_client.optim_step(trio.AdamParams(learning_rate=1e-4)).result()

Properties

Methods

forward

def forward(
    self,
    data: list[Datum],
    loss_fn: str = "cross_entropy",
    loss_fn_config: dict[str, object] | None = None,
    auto_shift: bool = False,
) -> APIFuture[ForwardBackwardOutput]

Run the forward pass and compute the loss without accumulating gradients.

Parameters

Returns

APIFuture[ForwardBackwardOutput]. Call .result() to read the output. The output contains:

• loss_fn_outputs: Loss function outputs, one item per input sample.
• metrics: Forward pass metrics.

Example

future = training_client.forward(data=data)
output = future.result()
print(output.metrics)

forward_backward

def forward_backward(
    self,
    data: list[Datum],
    loss_fn: str = "cross_entropy",
    loss_fn_config: dict[str, object] | None = None,
    auto_shift: bool = False,
) -> APIFuture[ForwardBackwardOutput]

Run forward and backward, then accumulate gradients.

Parameters

Returns

APIFuture[ForwardBackwardOutput]. Call .result() to read the output. The output contains:

• loss_fn_outputs: Loss function outputs, one item per input sample.
• metrics: Forward pass metrics.

Example

future = training_client.forward_backward(data=data, loss_fn="cross_entropy")
output = future.result()

forward_backward_custom

def forward_backward_custom(
    self,
    data: list[Datum],
    loss_fn: Callable[
        [list[Datum], list["torch.Tensor"]], tuple["torch.Tensor", dict[str, float]]
    ],
) -> APIFuture[ForwardBackwardOutput]

Run forward and backward with a custom PyTorch loss function. Local torch is required.

Parameters

Returns

APIFuture[ForwardBackwardOutput]. Call .result() to read the output. The output contains:

• loss_fn_outputs: Loss function outputs, one item per input sample.
• metrics: Forward pass metrics.

Example

def my_loss(data, logprobs):
    loss = -sum(lp.mean() for lp in logprobs)
    return loss, {"my_loss": loss.item()}

future = training_client.forward_backward_custom(data=data, loss_fn=my_loss)
output = future.result()

optim_step

def optim_step(self, adam_params: AdamParams) -> APIFuture[OptimStepResponse]

Apply one Adam optimizer update from the accumulated gradients, then clear the gradients.

Parameters

Returns

APIFuture[OptimStepResponse]. Call .result() to read optimizer metrics.

Example

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

Checkpoint Name Rules

The name parameter of save_state() and save_weights_for_sampler() follows the same normalization and validation rules. The name becomes part of a storage path, so its length limit derives from filesystem file-name limits:

• The name must match ^[A-Za-z0-9](?:[A-Za-z0-9._-]*[A-Za-z0-9_-])?$: it must start with an ASCII letter or digit, remaining characters may only be ASCII letters, digits, periods, underscores, or hyphens, and it cannot end with a period.
• The SDK first replaces ASCII spaces with hyphens and logs a WARNING containing the original and normalized names.
• If the name is longer than 200 characters after space replacement, the SDK shortens it to the first 180 characters followed by the first 20 hexadecimal characters of the full name's SHA-256 digest, then logs a WARNING.
• Any other invalid name raises PyTrioError with code validation.invalid_checkpoint_name before the request is submitted.

The digest is computed from the full name after space replacement, so the same normalized name maps to the same saved name.

save_state

def save_state(
    self,
    name: str,
    ttl_seconds: int | None = None,
    overwrite: bool = False,
) -> APIFuture[SaveWeightsResponse]

Save the LoRA adapter and optimizer state as a train checkpoint for resume training.

Parameters

Returns

APIFuture[SaveWeightsResponse]. Call .result() to read the saved checkpoint path and model name.

• path: Saved checkpoint URI.
• model: Saved model name.

Example

result = training_client.save_state(name="step-100").result()
print(result.path)

load_state

def load_state(self, path: str) -> APIFuture[dict]

Load checkpoint weights produced by save_state() before the first training, optimizer, or save operation. Optimizer state is not restored.

Parameters

Returns

APIFuture[dict] - a completed empty-result future on success, marking that Trainer initialization has completed.

Behavior

The SDK first reads the source checkpoint configuration and compares the base model, LoRA rank, train_mlp, train_attn, and train_unembed. If the structures are incompatible, it raises validation.checkpoint_config_mismatch with all differences in the error details. The seed is not part of the structural compatibility check; restore keeps the seed declared when the target client was created.

After compatibility is confirmed, the SDK atomically binds the checkpoint and checkpoint Actor through Control. It then uses the returned final Actor, JWT, and Trainer capability to create the first ActorTrainingSession, starts the heartbeat, and submits Trainer initialization once. Calling this method after Trainer initialization raises validation.already_initialized.

Example

future = training_client.load_state(checkpoint_uri)
future.result()

load_state_with_optimizer

def load_state_with_optimizer(self, path: str) -> APIFuture[dict]

Load checkpoint weights and optimizer state produced by save_state() before the first training, optimizer, or save operation. Parameters, return value, and compatibility checks are the same as load_state().

Example

future = training_client.load_state_with_optimizer(checkpoint_uri)
future.result()

save_weights_for_sampler

def save_weights_for_sampler(
    self,
    name: str,
    ttl_seconds: int | None = None,
) -> APIFuture[SaveWeightsForSamplerResponse]

Save only the LoRA adapter for later inference or sampling. Optimizer state is not saved.

Parameters

Returns

APIFuture[SaveWeightsForSamplerResponse]. Call .result() to read the saved checkpoint path, model name, and weight size.

• path: Saved weights checkpoint path.
• model: Saved model name.
• size: Weights size in bytes.

Example

result = training_client.save_weights_for_sampler(name="step-100").result()
print(result.path)

create_sampling_client

def create_sampling_client(
    self,
    model_path: str,
) -> SamplingClient

Create a SamplingClient from a LoRA adapter path during training. The client reuses the current training client's base model.

Parameters

Returns

SamplingClient.

Example

sampling_client = training_client.create_sampling_client(
    model_path="/path/to/weights",
)

save_weights_and_get_sampling_client

def save_weights_and_get_sampling_client(self) -> SamplingClient

Save the current model weights to a temporary anonymous archive and immediately return a SamplingClient loaded with those weights. This method does not require an archive name. The saved weights are intended only for temporary inference and sampling during training, such as Agent-RL workflows that continuously sample from the latest policy inside the training loop.

Returns

SamplingClient.

Example

sampling_client = training_client.save_weights_and_get_sampling_client()

get_tokenizer

def get_tokenizer(self)

Return the tokenizer for the current base model. It uses AutoTokenizer from transformers / modelscope.

Example

tokenizer = training_client.get_tokenizer()
tokens = tokenizer.encode("The meaning of life is")

Async Methods

forward_async

async def forward_async(
    self,
    data: list[Datum],
    loss_fn: str = "cross_entropy",
    loss_fn_config: dict[str, object] | None = None,
    auto_shift: bool = False,
) -> APIFuture[ForwardBackwardOutput]

Async version of forward. Parameters are the same.

Returns

APIFuture[ForwardBackwardOutput]. Call .result() or await future to read the output. The output contains:

• loss_fn_outputs: Loss function outputs, one item per input sample.
• metrics: Forward pass metrics.

future = await training_client.forward_async(data=data)
output = await future

forward_backward_async

async def forward_backward_async(
    self,
    data: list[Datum],
    loss_fn: str = "cross_entropy",
    loss_fn_config: dict[str, object] | None = None,
    auto_shift: bool = False,
) -> APIFuture[ForwardBackwardOutput]

Async version of forward_backward. Parameters are the same.

Returns

APIFuture[ForwardBackwardOutput]. Call .result() or await future to read the output. The output contains:

• loss_fn_outputs: Loss function outputs, one item per input sample.
• metrics: Forward pass metrics.

future = await training_client.forward_backward_async(data=data)
output = await future

forward_backward_custom_async

async def forward_backward_custom_async(
    self,
    data: list[Datum],
    loss_fn: Callable[
        [list[Datum], list["torch.Tensor"]],
        tuple["torch.Tensor", dict[str, float]]
    ],
) -> APIFuture[ForwardBackwardOutput]

Async version of forward_backward_custom. Parameters are the same.

Returns

APIFuture[ForwardBackwardOutput]. Call .result() or await future to read the output. The output contains:

• loss_fn_outputs: Loss function outputs, one item per input sample.
• metrics: Forward pass metrics.

future = await training_client.forward_backward_custom_async(data=data, loss_fn=my_loss)
output = await future

optim_step_async

async def optim_step_async(self, adam_params: AdamParams) -> APIFuture[OptimStepResponse]

Async version of optim_step. Parameters are the same.

Returns

APIFuture[OptimStepResponse]. Call .result() or await future to read optimizer metrics.

future = await training_client.optim_step_async(AdamParams(learning_rate=1e-4))
await future

save_state_async

async def save_state_async(
    self,
    name: str,
    ttl_seconds: int | None = None,
    overwrite: bool = False,
) -> APIFuture[SaveWeightsResponse]

Async version of save_state. Parameters are the same.

Returns

APIFuture[SaveWeightsResponse]. Call .result() or await future to read the saved checkpoint path and model name.

• path: Saved checkpoint URI.
• model: Saved model name.

future = await training_client.save_state_async(name="step-100")
result = await future

load_state_async

async def load_state_async(self, path: str) -> APIFuture[dict]

Async version of load_state. It must be called before the first training, optimizer, or save operation.

future = await training_client.load_state_async(checkpoint_uri)
await future

load_state_with_optimizer_async

async def load_state_with_optimizer_async(self, path: str) -> APIFuture[dict]

Async version of load_state_with_optimizer. It must be called before the first training, optimizer, or save operation.

future = await training_client.load_state_with_optimizer_async(checkpoint_uri)
await future

save_weights_for_sampler_async

async def save_weights_for_sampler_async(
    self,
    name: str,
    ttl_seconds: int | None = None,
) -> APIFuture[SaveWeightsForSamplerResponse]

Async version of save_weights_for_sampler. Parameters are the same.

Returns

APIFuture[SaveWeightsForSamplerResponse]. Call .result() or await future to read the saved checkpoint path, model name, and weight size.

• path: Saved weights checkpoint path.
• model: Saved model name.
• size: Weights size in bytes.

future = await training_client.save_weights_for_sampler_async(name="step-100")
result = await future

create_sampling_client_async

async def create_sampling_client_async(
    self,
    model_path: str,
) -> SamplingClient

Async version of create_sampling_client. Parameters are the same.

Returns

SamplingClient.

sampling_client = await training_client.create_sampling_client_async(
    model_path="/path/to/weights",
)

save_weights_and_get_sampling_client_async

async def save_weights_and_get_sampling_client_async(self) -> SamplingClient

Async version of save_weights_and_get_sampling_client. Parameters are the same.

Returns

SamplingClient.

sampling_client = await training_client.save_weights_and_get_sampling_client_async()