TinyTrain

TinyTrain is a lightweight training library for fine-tuning small language models on consumer hardware.

It is being developed as part of the TinyLM project, whose goal is to build capable open-source language models that can be trained and deployed under strict compute and memory constraints.

TinyTrain focuses on making the training process understandable, controllable, and efficient rather than hiding everything behind a large training framework.

🚧 Early development: TinyTrain is currently experimental and its API may change.

Features

• 🧠 Supervised Fine-Tuning (SFTTrainer)
• ⚙️ Simple, explicit training configuration (SFTConfig)
• 💾 Gradient accumulation for limited VRAM
• 🔥 Automatic mixed-precision training
• 🧩 Gradient checkpointing support
• 📦 Support for Hugging Face Dataset and streaming datasets
• 📝 Chat-style datasets using tokenizer chat templates
• 🎯 Assistant-only loss masking for SFT
• 📚 Tight sequence packing to reduce wasted tokens
• 📊 Training and validation metrics
• 💾 Automatic best-checkpoint saving
• 🔄 Training checkpoint resume
• 🌡️ GPU temperature monitoring and cooldown
• 🛡️ Basic VRAM/OOM protection
• 📈 Logging of loss, perplexity, entropy, token accuracy, learning rate, and gradient norm

Installation

Install the latest release from PyPI:

pip install tinytrain

For development:

git clone https://github.com/Se00n00/TinyTrainer.git
cd TinyTrainer

pip install -e .

Quick Start

from tinytrain import SFTTrainer, SFTConfig

config = SFTConfig(
    total_samples=100_000,
    batch_size=2,
    grad_accum_steps=16,
    learning_rate=2e-5,
    max_length=512,
    checkpoint_dir="checkpoints",
)

trainer = SFTTrainer(
    training_name="my-model",
    current_example=0,
    model=model,
    tokenizer=tokenizer,
    ds=dataset,
    config=config,
)

trainer.train()

The model and tokenizer are intentionally passed into the trainer rather than being hidden behind a high-level abstraction. This keeps the training process explicit and makes TinyTrain easier to adapt to custom small-model architectures.

Dataset Format

TinyTrain supports datasets containing either messages or text.

Chat / SFT format

{
    "messages": [
        {
            "role": "user",
            "content": "What is machine learning?"
        },
        {
            "role": "assistant",
            "content": "Machine learning is a..."
        }
    ]
}

When the tokenizer provides a chat template, TinyTrain uses it to format the conversation and applies loss masking to assistant responses.

Text format

{
    "text": "This is an example training document."
}

Text examples are treated as regular causal-language-model training data.

Configuration

SFTConfig controls the main training parameters:

Parameter	Default	Description
total_samples	required	Total number of examples available
test_train_ratio	0.01	Fraction reserved for validation
batch_size	2	Micro-batch size
grad_accum_steps	16	Gradient accumulation steps
learning_rate	2e-5	Initial learning rate
warmup_steps_ratio	0.10	Warmup fraction
max_length	512	Maximum packed sequence length
checkpoint_dir	checkpoints	Checkpoint output directory
resume	None	Checkpoint path or automatic resume
logging_steps	10	Training logging interval
eval_steps	200	Validation interval
gradient_checkpointing	True	Enable memory-saving checkpointing
weight_decay	0.1	AdamW weight decay
vram_limit_mb	4000	VRAM safety limit
max_temp	75	Maximum GPU temperature
cooldown_temp	60	Temperature at which training resumes

Memory-Constrained Training

TinyTrain is designed around the idea that useful language-model training should not require access to a large GPU cluster.

The trainer supports several techniques intended to make experimentation practical on smaller machines:

Small batch size
      ↓
Gradient accumulation
      ↓
Effective larger batch
      ↓
Gradient checkpointing
      ↓
Lower activation memory
      ↓
VRAM monitoring / recovery

The goal is not to hide hardware limitations, but to make them easier to work with.

Logging

Training and validation metrics are written to CSV files under the checkpoint directory:

checkpoints/
└── my-model/
    ├── model.pt
    ├── logs_train.csv
    └── logs_validation.csv

Training logs currently include:

• Loss
• Perplexity
• Entropy
• Mean token accuracy
• Learning rate
• Gradient norm
• Number of processed tokens/examples

Validation logs contain the corresponding evaluation metrics.

Project Philosophy

TinyTrain is intentionally small and transparent.

The project is being built alongside TinyLM, an open-source effort to explore how capable language models can be trained and deployed under strict parameter, memory, and compute constraints.

The long-term goal is:

Train
  ↓
Evaluate
  ↓
Improve
  ↓
Compress
  ↓
Deploy

with the entire process remaining accessible to developers who do not have access to large-scale compute.

Roadmap

• Initial SFT trainer
• Gradient accumulation
• Mixed precision
• Gradient checkpointing support
• Streaming dataset support
• Sequence packing
• Checkpoint resume
• Training/validation logging
• More robust checkpoint recovery
• Better dataset preprocessing utilities
• More training strategies
• Distributed training support
• Improved evaluation utilities
• Comprehensive test suite
• Stable public API
• Documentation and examples

Contributing

TinyTrain is open source and contributions are welcome.

If you find a bug, have an optimization idea, or want to improve the training experience on limited hardware, open an issue or submit a pull request.

The project is still young, so feedback on the API and architecture is especially valuable.

Related Project

TinyLM — the open-source small language model project built using the ideas and tooling developed around TinyTrain.

The broader goal is to create language models that are small enough to be practical on edge and consumer hardware while remaining useful.

License

TinyTrain is released under the MIT License. See LICENSE for details.