TrainingClient sends operations to a Loops trainer endpoint in either availability mode. Long-running operations use a submit-and-retrieve protocol: the submit fires immediately on the calling thread (so validation errors surface at call time) and .result() long-polls the server until the operation finishes. You can submit multiple operations before awaiting any of them. Construct one with ServiceClient.create_lora_training_client.
Run one training step and save a checkpoint. Each long-running call is submit-then-.result():
Every long-running server operation on
ServiceClient, TrainingClient, and SamplingClient (for example, forward_backward, sample, create_lora_training_client) has an await-able *_async counterpart for callers running their own event loop. The async variants accept the same arguments as their synchronous names. Simpler blocking calls like health, ensure_ready, get_tokenizer, and close (whose async form is aclose) have no *_async twin.ForwardBackwardFuture
Run a forward and backward pass over
data (a list of Datum objects) using the specified loss function. Returns a ForwardBackwardFuture; call .result() to block until the pass completes and retrieve the loss.loss_fn defaults to "cross_entropy". The accepted values are cross_entropy, importance_sampling, ppo, dppo, cispo, dro, and dpo. See Loss functions for each loss’s data shape, usage example, and loss_fn_config defaults.Set top_k_logprobs above 0 to include the trainer’s highest-probability token IDs and their log-probabilities at each scored position. The default value of 0 skips the top-k computation and does not add these tensors to the response.The accepted loss_fn_inputs shapes depend on the loss. Cross-entropy accepts [N] targets or [N, K] targets with matching weights; other built-in losses use one value per token position. See Loss functions.ForwardBackwardFuture
Run a forward and backward pass with a custom PyTorch loss function.
loss_fn receives (data, logprobs_list) and returns (loss, metrics). Each item in logprobs_list matches the corresponding datum’s target_tokens shape, including [N, K] targets.The only supported loss_type_input is "logprobs". Install baseten-loops[torch] to include PyTorch. See Custom loss functions for the callback contract and example.ForwardBackwardFuture
Run a forward pass without gradient computation. Same inputs and output shape as
forward_backward, but the gradient buffer is left untouched, so it is safe to interleave with gradient accumulation steps.Inspect trainer-selected tokens
Passtop_k_logprobs=K to forward or forward_backward. Each datum’s entry in loss_fn_outputs then contains two additional TensorData values:
topk_token_idshas shape[sequence_length, K]and dtypeint64.topk_logprobshas the same shape and dtypefloat32, ordered from highest to lowest probability at each position.
logprobs value reports the trainer log-probability for the target token at each position. The top-k values report the alternatives selected by the trainer from the same distribution. Inactive loss positions contain token ID -1 and log-probability 0.0.
top_k_logprobs is available on the native baseten.loops.TrainingClient. It is not part of the Tinker-compatible tinker.TrainingClient signature.
APIFuture[OptimStepResponse]
Apply the accumulated gradients using the Adam optimizer configured by
AdamParams. Call this after one or more forward_backward calls. The result’s metrics include the gradient norm and mean loss for the update; see OptimStepResponse for every key.APIFuture[SaveWeightsResponse]
Save a training checkpoint under
name. In the dedicated mode, when a weight sync URI is configured server-side, save_state also publishes the LoRA adapter so a polling sampler can hot-swap to the new weights.APIFuture[SaveWeightsResponse]
Save the LoRA adapter for sampling under
name without returning a SamplingClient. Use save_weights_and_get_sampling_client to publish weights and get a client that selects them.SamplingClient
Publish the LoRA adapter and return a
SamplingClient that selects the published weights. The first call creates the run’s paired sampler if needed; later calls reuse it. The call blocks until the publish completes.In the serverless mode, each call saves a distinct checkpoint and pins the returned client to it. Use the client’s checkpoint_name when downloading that save; it can differ from the requested name. Later publishes don’t change the weights served by that client.In the dedicated mode, the returned client sends X-Min-Policy-Version on each sample() call and can serve the published version or a later one. Replicas below that version return 408, and the SDK retries with backoff up to ready_timeout. The client polls deployment status before the first sample and raises SamplerTerminatedError if the deployment fails.Set record_stability_info=True to record every sample from the returned client for router replay. Each returned SampledSequence then carries a sequence_id to pass to Datum.with_provenance().APIFuture[LoadWeightsResponse]
Load weights from a
bt://loops:<run_id>/weights/<checkpoint> URI into this trainer. Use to resume training from a checkpoint.APIFuture[LoadWeightsResponse]
Same as
load_state but also restores Adam moments. Use when you want bit-exact resumption.list[Checkpoint]
List checkpoints for the run bound to this client. Requires that this client was constructed using
ServiceClient.create_lora_training_client (which populates the necessary session and run IDs automatically). Returns a list of Checkpoint.CheckpointFilesResponse
Return presigned URLs for every file in an S3-backed checkpoint folder. Pass either
checkpoint_id (a Checkpoint.id) or name, resolved against this client’s run; target (a CheckpointTarget) picks one when a name was saved for both a trainer and a sampler checkpoint. Volume-backed checkpoints can’t be fetched by URL; use download_checkpoint() instead. Same semantics as ServiceClient.get_checkpoint_archive_url.CheckpointDownload
Download a checkpoint to a local directory, whether its files live in S3 or a Baseten volume. Pass either
checkpoint_id (a Checkpoint.id) or name, resolved against this client’s run; target (a CheckpointTarget) picks one when a name was saved for both a trainer and a sampler checkpoint.dest_dir must be empty for an S3-backed checkpoint. A volume-backed checkpoint lands in a subdirectory under dest_dir, and that landing directory must be empty. The returned CheckpointDownload reports the actual destination, file count, and bytes written.SamplingClient
Return a
SamplingClient bound to the run’s paired sampler, serving the weights at model_path (a bt://loops:<run_id>/sampler_weights/<checkpoint> URI). In the serverless mode, the client stays pinned to that checkpoint. In the dedicated mode, it serves the checkpoint until the trainer publishes a newer version. The first call creates the paired sampler if needed; later calls reuse it. The checkpoint must use the run’s LoRA configuration. Before returning, the call validates the checkpoint against the run and raises httpx.HTTPStatusError if validation fails. Distinct from ServiceClient.create_sampling_client, which creates a standalone sampler with no weight sync. record_stability_info works as in save_weights_and_get_sampling_client.PreTrainedTokenizer
Return the Hugging Face
PreTrainedTokenizer for the base model. Cached after the first load.GetInfoResponse
Return the model configuration for this training session (base model name, LoRA rank, and max sequence length) without a server round-trip.
str | None
Property. The run ID this client is bound to. Use this when filtering checkpoints or making HTTP API calls against the same run.
str | None
Property. The custom base checkpoint the run started from, pinned to a commit as
hf://<org>/<repo>@<sha>, or None when the run trains from base_model’s weights. Set it with base_checkpoint on create_lora_training_client.int
Property. The number of
optim_step calls applied to a trainer deployment so far. Each access issues a GET /policy_version round-trip, so read it deliberately rather than in a tight loop. Serverless trainers raise NotImplementedError. For a published checkpoint, use save_weights_for_sampler(...).result().version, or call save_weights_and_get_sampling_client(...) to sample the saved weights.APIFuture[InitTrainerServerResponse]
Reset trainer state to a fresh LoRA adapter at
lora_rank. Use to start a new adapter on an existing trainer without reprovisioning.None
Check the trainer’s
/health endpoint. Returns None on success and raises if the trainer is unreachable or unhealthy.None
Close the client’s HTTP connections and finish any active Weights & Biases run. In
async code, call aclose() instead.An open TrainingClient keeps a dedicated run active between training calls. The client stops doing so when you call close(), when it’s garbage-collected, when your process exits, or after it can’t reach the run for about 20 minutes. Closing the client doesn’t deactivate the run: a dedicated run ends after inactivity, 15 minutes by default, and bills for GPU time until then. Deactivate the run to end it without waiting.None
Async counterpart to
close() for callers running their own event loop.