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.

Concretely, Trainite generates:

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

Dynamic Configuration

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.

Table of Contents

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

Quickstart

Requires Python >= 3.10.

Option 1 — uv (recommended):

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

# Initialize the project in any workspace
cd ..
trainite init my-experiment --model rope-transformer --dataset string-reverse
cd my-experiment
uv sync
uv run python main.py config.yaml

Option 2 — pip:

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

# Initialize the project in any workspace
cd ..
trainite init my-experiment --model rope-transformer --dataset string-reverse
cd my-experiment
pip install -e .
python main.py config.yaml

Usage

[!NOTE] If you are using the uv workflow, prefix commands with uv run (e.g., uv run trainite init or uv run python main.py config.yaml). If you are using standard pip/venv, run them directly (e.g., trainite init or python main.py config.yaml).

Initialize a Project

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

Run Training

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

Monitor & Iterate

• 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

Built-in Components

Models

Preprocessors

Datasets

Trainers

Configuration

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.

Swapping Components with _target_

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.

Customizing Your Project

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

Examples

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.