Training Systems

This section contains documentation for training systems.

Modules

Callbacks

• base
• checkpoint
• early_stopping
• logging
• profiling
• visualization

Distributed

• data_parallel
• mesh
• model_parallel

Optimizers

• adafactor
• adamw
• lion
• optax_wrappers

Rl

• dpo
• ppo
• reinforce

Schedulers

• cosine
• exponential
• factory
• linear
• scheduler

Trainers

• autoregressive_trainer
• diffusion_trainer
• energy_trainer
• flow_trainer
• gan_trainer
• vae_trainer

Training

• base
• gradient_accumulation
• mixed_precision
• tracking
• trainer
• utils

Statistics

• Total modules: 35
• Total classes: 2
• Total functions: 13