Chapter 6. Configuring model checkpointing for distributed training with Kubeflow Trainer V2

How you distribute GPUs across nodes affects storage usage, checkpoint download times, and training performance. When using S3 storage, each pod independently downloads the model and checkpoints to its local emptyDir volume. Minimizing the number of pods reduces the total storage consumed and the number of redundant downloads.

Consider the following two configurations for training with six GPUs:

trainer = TransformersTrainer(
    func=train_fn,
    num_nodes=6,
    resources_per_node={
        "nvidia.com/gpu": 1,
        "memory": "64Gi",
        "cpu": "4",
    },
    output_dir="s3://my-bucket/llama3-fine-tune",
)

trainer = TransformersTrainer(
    func=train_fn,
    num_nodes=2,
    resources_per_node={
        "nvidia.com/gpu": 3,
        "memory": "192Gi",
        "cpu": "12",
    },
    output_dir="s3://my-bucket/llama3-fine-tune",
)

The second configuration is more efficient for the following reasons:

When possible, maximize the number of GPUs per node to reduce the total number of pods in your training job.

When using S3 checkpointing, each pod writes checkpoints to local emptyDir storage before uploading them to S3. Storage usage fluctuates during training, with temporary spikes during checkpoint consolidation. Provision local storage to handle peak usage rather than average usage, to avoid pod eviction due to storage pressure.

Storage usage follows this general pattern:

The following tables show storage measurements from internal benchmarks. Your actual storage requirements can vary depending on your model, precision, training strategy, and dataset. Use these figures as a starting point for capacity planning and validate with with testing against your own workload.

Estimate per-pod local storage requirements for training workloads using a formula that accounts for model cache, checkpoint downloads, consolidation peaks, and final model storage. Example calculations are provided for DDP, FSDP, and DeepSpeed ZeRO-3 strategies.

Use the following formula to estimate per-pod local storage requirements:

Per-pod storage = base_cache + checkpoint_download + (N x checkpoint_size) + consolidation_peak + final_model + safety_buffer

where:

Sequential cleanup automatically removes old checkpoints as new ones are controlled by save_total_limit parameter, resulting in the lower storage estimates above. Full calculations show worst case if cleanup does not occur.

During checkpoint saves, the underlying training framework temporarily requires additional storage to consolidate model state before writing the final checkpoint files. This temporary spike is the most common cause of pod eviction due to storage pressure.

DeepSpeed ZeRO-3 has the largest consolidation peaks. A checkpoint that is only 1.6 GB in its final form can require up to 68 GB of temporary storage during the save operation. If you provision storage based on the steady-state checkpoint size alone, pods might be evicted during checkpoint operations.

The PeriodicCheckpointConfig class controls how often Kubeflow Trainer saves checkpoints during training and how many recent checkpoints it retains in local storage.

Pass this configuration to your TransformersTrainer:

from kubeflow.trainer import TransformersTrainer, PeriodicCheckpointConfig

checkpoint_config = PeriodicCheckpointConfig(
    save_strategy="steps",   # or "epoch"
    save_steps=50,           # Save every 50 steps
    save_total_limit=2,      # Keep only the 2 most recent checkpoints locally
)

trainer = TransformersTrainer(
    func=train_fn,
    num_nodes=2,
    resources_per_node={"nvidia.com/gpu": 2},
    output_dir="s3://my-bucket/llama3-fine-tune",
    data_connection_name="my-s3-checkpoint-storage",
    periodic_checkpoint_config=checkpoint_config,
)

When you configure periodic checkpoints, consider the following guidelines:

Monitor local storage consumption on training pods during training runs to detect storage pressure before pods are evicted. Use oc exec with standard Linux commands (df and du) to inspect storage usage on individual pods.

# Check storage usage on a training pod
$ oc exec <pod-name> -- df -h /mnt/kubeflow-checkpoints

# Detailed breakdown by directory
$ oc exec <pod-name> -- du -h /mnt/kubeflow-checkpoints | sort -h | tail -20

Compare the storage characteristics of FSDP, DDP, and DeepSpeed ZeRO-3 distributed training strategies to plan checkpoint storage capacity and understand performance tradeoffs. Each distributed training strategy has different storage characteristics :