Reliability

Reliability#

Checkpoint Support#

Tunix provides robust checkpointing capabilities to save and resume training progress, leveraging Orbax as the backend. This includes saving model parameters (supporting full state or only LoRA parameters for PEFT) and optimizer state.

Checkpointing is managed by CheckpointManager and integrated into PeftTrainer. SFT uses PeftTrainer directly, while RL uses rl.Trainer, a subclass of PeftTrainer, inside of the RLLearner. Therefore, both SFT and RL share the same checkpointing mechanism. Checkpointing and restarting are built-in features that require no special setup beyond configuration. To enable checkpointing, users simply need to set checkpoint_root_directory in SFTConfig or RLConfig; if this path is provided, Tunix automatically saves checkpoints and resumes training from the most recent one if interrupted, restoring model weights, optimizer state, and training step count. By default, checkpointing is disabled if checkpoint_root_directory is not specified. Users can further customize checkpointing behavior via checkpointing_options in the config.

Users customize background preservation behavior granularly using components defined inside checkpoint_options:

  • Save Decision Policies: Dictates when to initiate a checkpoint based on defined steps or intervals. Supported configurations include FixedIntervalPolicy and ContinuousCheckpointingPolicy. The default is ContinuousCheckpointingPolicy(minimum_interval_secs=180) (saves every 180 seconds). See Orbax v1 save_decision_policies.py for the complete interface contracts.

  • Preservation Policies: Sets specifications regarding tracking checkpoints over bounded timelines (e.g., LatestN). The default is LatestN(n=3) (keeps the latest 3 checkpoints). See Orbax v1 preservation_policies.py for the complete interface contracts.

  • Step Name Format: Defines the representation of directory names for step checkpoints. The default is ocp.path.step.standard_name_format() (uses simple integer step names).

  • Asynchronous Processing: Manage asynchronous behavior by specifying:

    • enable_async_checkpointing: Whether to use async checkpointing. Defaults to True. It is recommended to keep this enabled to prevent the main thread from blocking during training runs while checkpoints are written to storage.

    • timeout_secs: The timeout for asynchronous operations. Defaults to 1200 seconds.

Fault Tolerance#

Tunix ensures fault tolerance primarily through its checkpointing mechanism, allowing training to resume after interruptions such as machine restarts or pre-emptions.

Additionally, to prevent out-of-memory (OOM) errors due to excessive HBM usage, Tunix includes an InflightThrottler. This mechanism limits the number of TPU computations that can be scheduled concurrently, as configured by max_inflight_computations in TrainingConfig, thus providing more stable training runs on memory-constrained hardware.

Determinism Guarantee#

Tunix supports deterministic training runs through careful management of random number generation and data handling:

  • Model Initialization: Models can be initialized with a specific random seed (rng_seed or random_seed) to ensure consistent initial weights across runs.

  • Data Shuffling: RL learners accept a data_shuffle_seed parameter, which ensures that dataset shuffling is deterministic.

  • Dropout and Stochastic Layers: JAX and Flax RNG handling ensures that stochastic operations can be made deterministic if RNGs are correctly seeded and managed.

By providing explicit seeds for these components, users can ensure reproducibility of training experiments.