Training

The Training SDK provides classes for configuring and managing machine learning model training jobs on Baseten. This reference documents the key classes used to define training configurations.

Deploy a `TrainingJob`

To deploy a training job, use the train push command:

truss train push <path_to_config_file>

The following classes are used to configure and deploy training jobs:

TrainingJob

Defines a complete training job configuration.

class TrainingJob:
    image: Image             # Container image configuration
    compute: Compute         # Compute resource configuration
    runtime: Runtime         # Runtime environment configuration

For example, to create a training job that uses a PyTorch base image, 4 CPUs, and 16GB of memory, you set the following:

training_job = TrainingJob(
    image=Image(base_image="pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime"),
    compute=Compute(cpu_count=4, memory="16Gi"),
    runtime=Runtime(
        start_commands=["python train.py"],
        checkpointing_config=CheckpointingConfig(enabled=True),
        cache_config=CacheConfig(enabled=True),
    )
)

TrainingProject

Organizes training jobs and provides project-level configuration.

class TrainingProject:
    name: str           # Project name
    job: TrainingJob   # Training job configuration
    team_name: str       # Name of the team that owns this project

For example, if you want to create a training project for a team named my-team, you set the following:

project = TrainingProject(
    name="llm-fine-tuning",
    job=training_job,
    team_name="my-team"
)

Image

Specifies the container image for the training environment.

class Image:
    base_image: str  # Docker image to use for training
    docker_auth: Optional[DockerAuth] # Authentication details for private docker images

For example, to create an image that uses the pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime base image, you can set the following:

image = Image(base_image="pytorch/pytorch:2.0.0-cuda11.7-cudnn8-runtime")

DockerAuth

Configures authentication for private Docker registries. Ensure that any SecretReference used has been set in your Baseten Workspace. See secrets for more details.

class DockerAuth:
    auth_method: truss_config.DockerAuthType  # Authentication method type
    registry: str                              # Docker registry URL
    aws_iam_docker_auth: Optional[AWSIAMDockerAuth] = None
    gcp_service_account_json_docker_auth: Optional[GCPServiceAccountJSONDockerAuth] = None

AWSIAMDockerAuth

Authenticates with AWS ECR using IAM credentials.

class AWSIAMDockerAuth:
    access_key_secret_ref: SecretReference      # AWS access key ID as a secret reference
    secret_access_key_secret_ref: SecretReference  # AWS secret access key as a secret reference

GCPServiceAccountJSONDockerAuth

Authenticates with Google Container Registry using service account JSON.

class GCPServiceAccountJSONDockerAuth:
    service_account_json_secret_ref: SecretReference  # GCP service account JSON secret reference

For example, to create an image that uses a GCP service account JSON secret, you can set the following:

# Configure image with GCP authentication
image = Image(
    base_image="gcr.io/my-project/my-repo/my-private-image:latest",
    docker_auth=DockerAuth(
        auth_method=truss_config.DockerAuthType.GCP_SERVICE_ACCOUNT_JSON,
        registry="gcr.io",
        gcp_service_account_json_docker_auth=GCPServiceAccountJSONDockerAuth(
            service_account_json_secret_ref=SecretReference(name="gcp_service_account_json")
        )
    )
)

Compute

Specifies compute resources for training jobs.

class Compute:
    node_count: int = 1      # Number of nodes for distributed training
    cpu_count: int = 1       # Number of CPU cores
    memory: str = "2Gi"      # Memory allocation
    accelerator: Optional[AcceleratorSpec] = None  # GPU configuration

For example, to create a compute resource that uses 2 GPUs and 4 CPUs, you can set the following:

# Configure a training job with 2 GPUs and 4 CPUs
compute = Compute(
    accelerator=AcceleratorSpec(accelerator="H100", count=4)
)

Runtime

Defines the runtime environment for training jobs.

class Runtime:
    start_commands: List[str] = []  # Commands to run at job start
    environment_variables: Dict[str, Union[str, SecretReference]] = {}
    enable_cache: bool = False      # Enable caching
    checkpointing_config: CheckpointingConfig = CheckpointingConfig()

For example, to create a runtime environment that uses the python train.py start command, a BATCH_SIZE environment variable set to 32, and a WANDB_API_KEY environment variable set to a secret reference named WANDB_KEY, you can set the following:

runtime = Runtime(
    start_commands=["python train.py"],
    environment_variables={
        "BATCH_SIZE": "32",
        "WANDB_API_KEY": SecretReference(name="WANDB_KEY")
    },
    checkpointing_config=CheckpointingConfig(enabled=True)
)

Training Cache

By default, the training cache provides two mount locations:

$BT_PROJECT_CACHE_DIR, which is shared and accessible by all jobs within a Project ID
$BT_TEAM_CACHE_DIR which is shared by jobs that belong to the same Team.

Baseten will export these variables in your job’s environment. The cache storage is separate from ephemeral storage limits of your training job. Training Projects provide storage segregation within the cache. Training jobs within the same project share the same cache, while training jobs in different projects cannot access each other’s data. Example usage:

runtime = Runtime(
    start_commands=["python train.py"],
    environment_variables={
        "BATCH_SIZE": "32",
        "WANDB_API_KEY": SecretReference(name="WANDB_KEY")
    },
    cache_config=CacheConfig(enabled=True)
)

SecretReference

Used to securely reference secrets stored in your Baseten workspace.

class SecretReference:
    name: str  # Name of the secret in your workspace

Example usage:

# Reference a secret named "WANDB_API_KEY" 
secret_ref = SecretReference(name="WANDB_API_KEY")

CheckpointingConfig

Configures model checkpointing behavior during training. Baseten will export the $BT_CHECKPOINT_DIR within the Training Job’s environment. The checkpointing storage is independent of the ephemeral storage of the pod

class CheckpointingConfig:
    enabled: bool = False              # Enable/disable checkpointing
    checkpoint_path: Optional[str] = None  # Custom checkpoint directory path
    volume_size_gib: Optional[int] = None # Custom size of your checkpointing directory

For example, to enable checkpointing with the default path, you can set the following:

# Enable checkpointing with default path
checkpointing = CheckpointingConfig(enabled=True)

Baseten provided environment variables

Baseten automatically provides several environment variables in your training job’s environment to help integrate your code with the Baseten platform.

Environment variables

Environment Variable	Description	Example
`BT_TRAINING_JOB_ID`	ID of the Training Job	`"gvpql31"`
`BT_TRAINING_PROJECT_ID`	ID of the Training Project	`"aghi527"`
`BT_NUM_GPUS`	Number of available GPUs per node	`"4"`
`BT_PROJECT_CACHE_DIR`	Directory shared across Training Jobs within a singular Training Project	`"/root/.cache/user_artifacts"`
`BT_TEAM_CACHE_DIR`	Directory shared across Training Jobs within a singular Team	`/root/.cache/team_artifacts`
`BT_CHECKPOINT_DIR`	Directory where checkpoints are automatically saved during training	`"/mnt/ckpts"`
`BT_LOAD_CHECKPOINT_DIR`	Directory of where loaded checkpoints will be	`"/tmp/loaded_checkpoints"`
`BT_TRAINING_JOB_NAME`	Name your Training Job	`"gpt-oss-20b-lora"`
`BT_TRAINING_PROJECT_NAME`	Name your Training Project	`"gpt-oss-finetunes"`

Multinode Environment Variables

The following environment variables are particularly useful for multinode training jobs:

Environment Variable	Description	Example
`BT_GROUP_SIZE`	Number of nodes in the multinode deployment	`"2"`
`BT_LEADER_ADDR`	Address of the leader node	`"10.0.0.1"`
`BT_NODE_RANK`	Rank of the node	`"0"`

For multinode deployments, any traditionally used port number (e.g. 29500) will work. There is no specific port number required by Baseten.

Deploy Checkpoints as a Model

Deploy checkpoints CLI wizard

The easiest way to deploy your checkpoints is by using the CLI wizard:

truss train deploy_checkpoints --job-id <job id>

Follow the setup wizard to choose your hardware and the checkpoint you’d like to deploy. Baseten will automatically recognize checkpoints for full finetunes and LoRAs for LLMs and Whisper models. Note that FSDP checkpoints aren’t supported by deploy_checkpoints today and must be manually configured in the truss config. Once you’ve completed the wizard, Baseten will generate a truss and deploy a published model according to the specs provided.

Deploy checkpoints with static configuration

If you’d like to keep a static config of your checkpoint deploy, you can create a python config file defining the configuration you’d like to reference:

truss train deploy_checkpoints <path_to_config_file>

DeployCheckpointsRuntime

Configures the runtime environment for deployed checkpoints.

class DeployCheckpointsRuntime:
    environment_variables: Dict[str, Union[str, SecretReference]] = {}

Checkpoint

Represents metadata for a saved model checkpoint.

class Checkpoint:
    training_job_id: str   # ID of the training job
    id: str               # Checkpoint ID
    name: str            # Checkpoint name 
    lora_rank: Optional[int] = None  # LoRA rank if applicable. Auto-detected if not specified.

CheckpointList

Manages a collection of checkpoints and their download configuration.

class CheckpointList:
    download_folder: str = "training_checkpoints"  # Local download location upon deployment
    base_model_id: Optional[str] = None           # Base model identifier. Auto-detected if not specified.
    checkpoints: List[Checkpoint] = []            # List of checkpoints

DeployCheckpointsConfig

Specifies configuration for deploying trained model checkpoints.

class DeployCheckpointsConfig:
    checkpoint_details: Optional[CheckpointList] = None  # Checkpoints to deploy
    model_name: Optional[str] = None                    # Name for the deployed model
    deployment_name: Optional[str] = None               # Name for the deployment
    runtime: Optional[DeployCheckpointsRuntime] = None  # Runtime configuration
    compute: Optional[Compute] = None                   # Compute resources

For example, to deploy a checkpoint with the name checkpoint_1 from a training job with the id gvpql31, you can set the following:

deploy_config = DeployCheckpointsConfig(
    model_name="fine-tuned-llm",
    deployment_name="production-llm",
    checkpoint_details=CheckpointList(
        checkpoints=[
            Checkpoint(
                training_job_id="gvpql31",
                id="checkpoint_1",
                name="checkpoint_1"
            )
        ]
    ),
    compute=Compute(
        accelerator=AcceleratorSpec(accelerator="H100", count=1)
    )
)

Reference

Inference API

Management API

CLI reference

SDK reference

Deploy a `TrainingJob`

TrainingJob

TrainingProject

Image

DockerAuth

AWSIAMDockerAuth

GCPServiceAccountJSONDockerAuth

Compute

Runtime

Training Cache

SecretReference

CheckpointingConfig

Baseten provided environment variables

Environment variables

Multinode Environment Variables

Deploy Checkpoints as a Model

Deploy checkpoints CLI wizard

Deploy checkpoints with static configuration

DeployCheckpointsRuntime

Checkpoint

CheckpointList

DeployCheckpointsConfig

Reference

Inference API

Management API

CLI reference

SDK reference

​Deploy a TrainingJob

​TrainingJob

​TrainingProject

​Image

​DockerAuth

​AWSIAMDockerAuth

​GCPServiceAccountJSONDockerAuth

​Compute

​Runtime

​Training Cache

​SecretReference

​CheckpointingConfig

​Baseten provided environment variables

​Environment variables

​Multinode Environment Variables

​Deploy Checkpoints as a Model

​Deploy checkpoints CLI wizard

​Deploy checkpoints with static configuration

​DeployCheckpointsRuntime

​Checkpoint

​CheckpointList

​DeployCheckpointsConfig

Deploy a `TrainingJob`

TrainingJob

TrainingProject

Image

DockerAuth

AWSIAMDockerAuth

GCPServiceAccountJSONDockerAuth

Compute

Runtime

Training Cache

SecretReference

CheckpointingConfig

Baseten provided environment variables

Environment variables

Multinode Environment Variables

Deploy Checkpoints as a Model

Deploy checkpoints CLI wizard

Deploy checkpoints with static configuration

DeployCheckpointsRuntime

Checkpoint

CheckpointList

DeployCheckpointsConfig