Trainite

Trainite gives you a complete, working training project as your starting point. The goal of generating a training project is to provide a clean, standalone codebase for learning, rapid prototyping, and experimenting without being constrained by framework abstractions. (Note: Code generation is deterministic and template-based, not AI-based).

Built on PyTorch-Ignite, the generated code is yours — read it, modify it, extend it however your research needs.

my-experiment/
├── config.yaml     # YAML configuration (edit hyperparameters here)
├── config.py       # Pydantic validation models (IDE autocomplete)
├── models/         # Model architecture templates
├── datasets/       # Dataset and tokenization templates
├── preprocessors/  # Tokenizer and preprocessing templates
├── trainer.py      # Training loop built on PyTorch-Ignite
├── utils.py        # Config loading and instantiation helpers
├── main.py         # Project entrypoint (runs training)
├── pyproject.toml  # Isolated package dependencies
└── README.md       # Documentation tailored to your selected components

Swap components (like models, optimizers, or datasets) dynamically via config using dotted import paths (_target_):

model:
  _target_: models.rope_transformer.RoPETransformerModel
  hidden_size: 64

optimizer:
  _target_: torch.optim.AdamW
  lr: 0.001

Trainite is explicitly not:

• A replacement for HuggingFace Trainer.
• A tool for training large models on massive corpora.
• A framework that forces you into its runtime abstractions.

• Installation
• Quickstart
• Usage
  • Initialize a Project
  • Run Training
  • Monitor & Iterate
• Built-in Components
• Configuration
• Customizing Your Project
• Examples

Requires Python >= 3.10.

From pip:

pip install trainite

From source (using uv):

git clone https://github.com/pytorch-ignite/trainite.git
cd trainite
uv sync
source .venv/bin/activate

From source (using pip):

git clone https://github.com/pytorch-ignite/trainite.git
cd trainite
pip install -e .

Once Trainite is installed, generate and run your first experiment project:

# 1. Initialize a starter project
trainite init my-experiment --model rope-transformer --dataset string-reverse

# 2. Navigate into the project and install dependencies
cd my-experiment
pip install -e .

# 3. Run training
python main.py config.yaml

You can generate a starter project by passing configuration options as command-line flags:

trainite init my-experiment --model rope-transformer --dataset string-reverse

Or run it interactively (Trainite will walk you through the options step-by-step):

trainite init

Interactive prompt preview:

(Note: Trainite supports multi-model selection; you can select multiple models to include in your starter project.)

? Project directory: my-experiment
? Model(s): rope-transformer
? Dataset: string-reverse
? Trainer: decoder-trainer
? Output directory: outputs
? Run name: rope-transformer__string-reverse

 Generated config.yaml
 Generated models/rope_transformer.py
 Generated datasets/string_reverse.py
 Generated datasets/transformed.py
 Generated trainer.py
 Generated utils.py
 Generated main.py
 Generated config.py
 Generated preprocessors/char_tokenizer.py
 Generated README.md
 Generated pyproject.toml

Your generated project is a standalone application with no runtime dependency on Trainite. Navigate to it and run:

cd my-experiment
python main.py config.yaml

• TensorBoard: tensorboard --logdir outputs
• ClearML: Run clearml-init, then set logger: clearml in config.yaml.
• Edit architecture: Open models/transformer.py and modify the model.
• Edit hyperparameters: Open config.yaml and change learning rates, batch sizes, data splits, etc.

With ClearML enabled, Trainite keeps checkpoints in the run directory and uploads them to ClearML's default file server. Generated projects document how to use another storage URI, defer to ClearML configuration, or keep checkpoints local only.

Training outputs are organized by run:

outputs/
└── rope-transformer__string-reverse/
    └── 20260611_1430/
        ├── output.log          # Training logs
        ├── config.yaml         # Exact config used for this run
        ├── best_checkpoint_*.pt
        ├── last_checkpoint_*.pt
        └── tensorboard/        # TensorBoard event files

All configuration lives in config.yaml. The top-level blocks map to base validation schemas in config.py (e.g. ModelConfig, DatasetConfig). This ensures that your configuration has the correct overall structure and valid keys, while allowing you to pass custom hyperparameters to your local components.

Components (such as models, datasets, optimizers, or collation functions) are specified via the _target_ key. This key holds a dotted import path resolved at runtime, allowing you to swap components directly from config without altering training scripts:

model:
  _target_: models.rope_transformer.RoPETransformerModel
  hidden_size: 64

optimizer:
  _target_: torch.optim.AdamW
  lr: 0.001

For detailed configuration parameters and data splitting options (e.g., automatic splitting vs. explicit splits) specific to your selected components, refer to the generated README.md inside your initialized project.

Since the generated code is yours, you can override at any layer without touching the others:

Working examples are in the examples/ directory:

Each example is a standalone project — cd into it, install dependencies, and run python main.py config.yaml.

For contributor and development docs, see CONTRIBUTING.md.