Metadata-Version: 2.4
Name: tinytrain
Version: 0.1.0
Summary: A lightweight training library for small language models
Author: Se00n00
License-Expression: MIT
Project-URL: Homepage, https://github.com/Blue-sky-2025/TinyTrain
Project-URL: Repository, https://github.com/Blue-sky-2025/TinyTrain
Project-URL: Issues, https://github.com/Blue-sky-2025/TinyTrain/issues
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch
Requires-Dist: datasets
Dynamic: license-file

# 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:

```bash
pip install tinytrain
```

For development:

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

pip install -e .
```

## Quick Start

```python
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

```python
{
    "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

```python
{
    "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:

```text
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:

```text
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:

```text
Train
  ↓
Evaluate
  ↓
Improve
  ↓
Compress
  ↓
Deploy
```

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

## Roadmap

* [x] Initial SFT trainer
* [x] Gradient accumulation
* [x] Mixed precision
* [x] Gradient checkpointing support
* [x] Streaming dataset support
* [x] Sequence packing
* [x] Checkpoint resume
* [x] 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`](LICENSE) for details.
