Async

for i in range(15):
    fwdbwd_future = training_client.forward_backward(...)
    fwdbwd_result = fwdbwd_future.result()

async def main():
    for i in range(15):
        fwdbwd_future = await training_client.forward_backward_async(...)
        fwdbwd_result = await fwdbwd_future

if __name__ == "__main__":
    asyncio.run(main())

import pytrio as trio
import numpy as np
import asyncio

async def main():
    service_client = trio.ServiceClient()

    base_model = "Qwen/Qwen3.5-4B"
    training_client = await service_client.create_lora_training_client_async(
        base_model=base_model,
        rank=32,
    )

    examples = [
        {"input": "what is trio", "output": "TRIO is an AI infrastructure product from EmotionMachine."},
        {"input": "can you explain what trio is", "output": "TRIO helps teams run large-scale LLM post-training from a CPU-only machine."},
        {"input": "tell me about trio", "output": "TRIO provides APIs for LLM training, sampling, and reinforcement learning workflows."},
    ]

    print("Loading tokenizer...")
    tokenizer = training_client.get_tokenizer()
    print("Tokenizer loaded")

    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

        return trio.Datum(
            model_input=trio.ModelInput.from_ints(tokens=tokens[:-1]),
            loss_fn_inputs=dict(weights=weights[1:], target_tokens=tokens[1:])
        )

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

    print("Starting training")
    print_task_queue = []
    for iter in range(15):
        fwdbwd_future = await training_client.forward_backward_async(processed_examples, "cross_entropy")
        optim_future = await training_client.optim_step_async(trio.AdamParams(learning_rate=1e-4))

        async def print_loss_async(fwdbwd_future, optim_future, iter: int):
            fwdbwd_result = await fwdbwd_future
            await optim_future
            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])
            loss = -np.dot(logprobs, weights) / weights.sum()
            print(f"Iter{iter+1} Loss per token: {loss:.4f}")
            return loss
        
        print_task_queue.append(print_loss_async(fwdbwd_future, optim_future, iter))
        
    await asyncio.gather(*print_task_queue)

if __name__ == "__main__":
    asyncio.run(main())