Skip to content
View as Markdown

TrainingClient

Generated from tinker 0.30.4 at commit 1e5777e. Source links point at that snapshot.

class tinker.TrainingClient(holder, model_seq_id, model_id)

Client for training ML models with forward/backward passes and optimization.

The TrainingClient corresponds to a fine-tuned model that you can train and sample from. You typically get one by calling service_client.create_lora_training_client(). Key methods:

  • forward_backward() - compute gradients for training
  • optim_step() - update model parameters with Adam optimizer
  • save_weights_and_get_sampling_client() - export trained model for inference

Parameters:

Example:

training_client = service_client.create_lora_training_client(base_model="Qwen/Qwen3-8B")
fwdbwd_future = training_client.forward_backward(training_data, "cross_entropy")
optim_future = training_client.optim_step(types.AdamParams(learning_rate=1e-4))
fwdbwd_result = fwdbwd_future.result()  # Wait for gradients
optim_result = optim_future.result()    # Wait for parameter update
sampling_client = training_client.save_weights_and_get_sampling_client("my-model")

get_console_url()

Return the Tinker Console URL for this training run.

Returns: str

forward(data, loss_fn, loss_fn_config=None)

Compute forward pass without gradients.

Parameters:

  • data (List[types.Datum]) – List of training data samples
  • loss_fn (types.LossFnType) – Loss function type (e.g., "cross_entropy")
  • loss_fn_config (Mapping[str, float | str] | None, default: None) – Optional configuration for the loss function

Returns:

  • APIFuture containing the forward pass outputs and loss

Example:

data = [types.Datum(
    model_input=types.ModelInput.from_ints(tokenizer.encode("Hello")),
    loss_fn_inputs={"target_tokens": types.ModelInput.from_ints(tokenizer.encode("world"))}
)]
future = training_client.forward(data, "cross_entropy")
result = await future
print(f"Loss: {result.loss}")

Async variant: forward_async()

forward_backward(data, loss_fn, loss_fn_config=None)

Compute forward pass and backward pass to calculate gradients.

Parameters:

  • data (List[types.Datum]) – List of training data samples
  • loss_fn (types.LossFnType) – Loss function type (e.g., "cross_entropy")
  • loss_fn_config (Mapping[str, float | str] | None, default: None) – Optional configuration for the loss function

Returns:

  • APIFuture containing the forward/backward outputs, loss, and gradients

Example:

data = [types.Datum(
    model_input=types.ModelInput.from_ints(tokenizer.encode("Hello")),
    loss_fn_inputs={"target_tokens": types.ModelInput.from_ints(tokenizer.encode("world"))}
)]

# Compute gradients
fwdbwd_future = training_client.forward_backward(data, "cross_entropy")

# Update parameters
optim_future = training_client.optim_step(
    types.AdamParams(learning_rate=1e-4)
)

fwdbwd_result = await fwdbwd_future
print(f"Loss: {fwdbwd_result.loss}")

Async variant: forward_backward_async()

forward_backward_custom(data, loss_fn, loss_type_input='logprobs')

Compute forward/backward with a custom loss function.

Allows you to define custom loss functions that operate on log probabilities. The custom function receives logprobs and computes loss and gradients.

Parameters:

  • data (List[types.Datum]) – List of training data samples
  • loss_fn (CustomLossFnV1) – Custom loss function that takes (data, logprobs) and returns (loss, metrics)
  • loss_type_input (Literal['logprobs'], default: 'logprobs') – Input space for loss_fn. Currently the only supported value is "logprobs".

Returns:

  • APIFuture containing the forward/backward outputs with custom loss

Example:

def custom_loss(data, logprobs_list):
    # Custom loss computation
    loss = torch.mean(torch.stack([torch.mean(lp) for lp in logprobs_list]))
    metrics = {"custom_metric": loss.item()}
    return loss, metrics

future = training_client.forward_backward_custom(data, custom_loss)
result = future.result()
print(f"Custom loss: {result.loss}")
print(f"Metrics: {result.metrics}")

Async variant: forward_backward_custom_async()

optim_step(optim_params=None, adam_params=None)

Update model parameters using the optimizer selected at model creation.

The Adam optimizer used by tinker is identical to torch.optim.AdamW. Note that unlike PyTorch, Tinker's default weight decay value is 0.0 (no weight decay).

Parameters:

  • optim_params (types.OptimParams | None, default: None) – AdamParams or DimuonParams matching the optimizer selected at model creation.
  • adam_params (types.AdamParams | None, default: None) – Legacy keyword alias for AdamParams; cannot be combined with optim_params.

Returns:

  • APIFuture containing optimizer step response

Example:

# First compute gradients
fwdbwd_future = training_client.forward_backward(data, "cross_entropy")

# Then update parameters
optim_future = training_client.optim_step(
    types.AdamParams(
        learning_rate=1e-4,
        weight_decay=0.01
    )
)

# Wait for both to complete
fwdbwd_result = await fwdbwd_future
optim_result = await optim_future

Async variant: optim_step_async()

save_state(name, ttl_seconds=None, overwrite=False, user_metadata=None)

Save model weights to persistent storage.

Parameters:

  • name (str) – Name for the saved checkpoint
  • ttl_seconds (int | None, default: None) – Optional TTL in seconds for the checkpoint, between 1 hour (3600) and 10 years (None = never expires)
  • overwrite (bool, default: False) – If True, overwrite any existing checkpoint with the same name. This replaces the entire existing user_metadata mapping; if user_metadata is not provided, the previous user metadata is deleted.
  • user_metadata (dict[str, str] | None, default: None) – Optional user-provided metadata to attach to the checkpoint

Returns:

  • APIFuture containing the save response with checkpoint path

Example:

# Save after training
save_future = training_client.save_state("checkpoint-001")
result = await save_future
print(f"Saved to: {result.path}")

Async variant: save_state_async()

load_state(path, weights_access_token=None)

Load model weights from a saved checkpoint.

This loads only the model weights, not optimizer state (e.g., Adam momentum). To also restore optimizer state, use load_state_with_optimizer.

Parameters:

  • path (str) – Tinker path to saved weights (e.g., "tinker://run-id/weights/checkpoint-001")
  • weights_access_token (str | None, default: None) – Optional access token for loading checkpoints under a different account.

Returns:

Example:

# Load checkpoint to continue training (weights only, optimizer resets)
load_future = training_client.load_state("tinker://run-id/weights/checkpoint-001")
await load_future
# Continue training from loaded state

Async variant: load_state_async()

load_state_with_optimizer(path, weights_access_token=None)

Load model weights and optimizer state from a checkpoint.

Parameters:

  • path (str) – Tinker path to saved weights (e.g., "tinker://run-id/weights/checkpoint-001")
  • weights_access_token (str | None, default: None) – Optional access token for loading checkpoints under a different account.

Returns:

Example:

# Resume training with optimizer state
load_future = training_client.load_state_with_optimizer(
    "tinker://run-id/weights/checkpoint-001"
)
await load_future
# Continue training with restored optimizer momentum

Async variant: load_state_with_optimizer_async()

save_weights_for_sampler(name, ttl_seconds=None, user_metadata=None)

Save model weights for use with a SamplingClient.

Parameters:

  • name (str) – Name for the saved sampler weights
  • ttl_seconds (int | None, default: None) – Optional TTL in seconds for the checkpoint, between 1 hour (3600) and 10 years (None = never expires)
  • user_metadata (dict[str, str] | None, default: None) – Optional user-provided metadata to attach to the checkpoint

Returns:

  • APIFuture containing the save response with sampler path

Example:

# Save weights for inference
save_future = training_client.save_weights_for_sampler("sampler-001")
result = await save_future
print(f"Sampler weights saved to: {result.path}")

# Use the path to create a sampling client
sampling_client = service_client.create_sampling_client(
    model_path=result.path
)

Async variant: save_weights_for_sampler_async()

save_weights_external(name, ttl_seconds=None)

Save model weights in an external (e.g. HuggingFace) format.

The saved checkpoint is stored in the external-weights bucket and can be managed (listed, deleted) like any other checkpoint.

Parameters:

  • name (str) – Name for the saved external weights
  • ttl_seconds (int | None, default: None) – Optional TTL in seconds for the checkpoint, between 1 hour (3600) and 10 years (None = never expires)

Returns:

  • APIFuture containing the save response with the external weights path

Async variant: save_weights_external_async()

get_info()

Get information about the current model.

Returns:

Example:

info = training_client.get_info()
print(f"Model ID: {info.model_data.model_id}")
print(f"Base model: {info.model_data.model_name}")
print(f"LoRA rank: {info.model_data.lora_rank}")

Async variant: get_info_async()

get_tokenizer()

Get the tokenizer for the current model.

Returns:

  • PreTrainedTokenizer compatible with the model

Example:

tokenizer = training_client.get_tokenizer()
tokens = tokenizer.encode("Hello world")
text = tokenizer.decode(tokens)

create_sampling_client(model_path, retry_config=None, record_stability_info=False)

Create a SamplingClient from saved weights.

Parameters:

Returns:

Example:

sampling_client = training_client.create_sampling_client(
    "tinker://run-id/weights/checkpoint-001"
)
# Use sampling_client for inference

Async variant: create_sampling_client_async()

save_weights_and_get_sampling_client(name=None, retry_config=None, record_stability_info=False)

Save current weights and create a SamplingClient for inference.

Parameters:

  • name (str | None, default: None) – Deprecated, has no effect. Will be removed in a future release.
  • retry_config (RetryConfig | None, default: None) – Optional configuration for retrying failed requests
  • record_stability_info (bool, default: False)

Returns:

Example:

# After training, create a sampling client directly
sampling_client = training_client.save_weights_and_get_sampling_client()

# Now use it for inference
prompt = types.ModelInput.from_ints(tokenizer.encode("Hello"))
params = types.SamplingParams(max_tokens=20)
result = sampling_client.sample(prompt, 1, params).result()

Async variant: save_weights_and_get_sampling_client_async()

Referenced by