# TaoTrain: Production-Grade LLM Training Framework

**TaoTrain** is a sophisticated PyTorch framework for training large language models at every scale—from experimental pretraining through supervised fine-tuning to reinforcement learning. Unlike fragmented training scripts or heavyweight frameworks, TaoTrain unifies the **entire training pipeline** in a clean, modular codebase that appeals to both ML engineers and software engineers.

## Current Taotern Work

TaoTrain now includes the Taotern comparison architectures used by the current SSM LLM work:

- `taonet`: the attention/MLA baseline.
- `taonet_ssm`: the TaoNet shell with the attention mixer replaced by the Gamma Space Model DPLR SSM.
- `taonet_hybrid`: an alternating attention/SSM TaoNet used for the current best 200M-class candidate.

The current selected deployment-oriented run is `hybrid_ssm_first_199m`, a `199,480,928` parameter model with 16 layers: SSM layers at `0,2,4,6,8,10,12,14` and attention layers at `1,3,5,7,9,11,13,15`. It uses the DPLR SSM core with split two-lane mixing, channel gates, per-channel local shift, and the faster convolution path for long-sequence training.

Remote run `taotern-200m-hybrid-chat-20260512` trains this model on TaoData for a 4B-token base stage and then runs SFT so the final artifact can be loaded as a chat model. The trainable fixes added for this run are:

- Async JSONL iteration keeps polling while tokenization workers are alive instead of ending early after a temporary empty queue.
- Cached JSONL scan metadata is reused safely while recomputing chunk ranges for the active `samples_per_chunk` and `max_samples` settings.

## Why TaoTrain?

- **Complete Unified Pipeline**: Pretraining → SFT → RL in a single, consistent framework. No context switching between different codebases or architectures.
- **Production-Grade Engineering**: Type-safe Pydantic configs, comprehensive checkpointing, AimStack integration, and proper gradient handling—not research code, but a framework you can deploy.
- **Extensibility Without Modification**: Register custom models, optimizers, schedulers, and datasets via decorators. Experiment freely without forking the framework.
- **Developer Experience First**: Interactive TUI for inference, intuitive YAML configurations, async data loading that eliminates I/O bottlenecks, and clear abstractions that make the codebase a pleasure to work with.

## Key Capabilities

| Capability | Details |
|---|---|
| **Multi-Stage Training** | Unified infrastructure for pretraining, SFT, and RL. Share model checkpoints, logging, and evaluation across stages. |
| **Advanced Optimization** | Hybrid Muon + AdamW optimizer: efficient 2D weight updates via SVD-based methods + adaptive learning for 1D parameters. |
| **Modern Architectures** | DeepSeek MLA with grouped query attention (GQA), YaRN context extension, and factorized embeddings—all configurable via YAML. |
| **Production Features** | BF16 mixed precision training, gradient accumulation, proper gradient clipping, checkpoint resumption, and validation loops. |
| **Async Data Pipeline** | Background tokenization with multi-threaded workers. Stream billion-token datasets from JSONL without loading into memory. |
| **Interactive Inference** | TUI chat interface with real-time generation speed metrics and multi-model comparison. |
| **Logging & Monitoring** | AimStack integration tracks loss, metrics, hyperparameters, and git hashes for reproducibility. Visualize training runs in your browser. |

## Getting Started

### Installation

```bash
git clone https://github.com/lobakkang/taoTrain.git
cd taoTrain
pip install -e .
```

### Training Examples

**Pretraining on a custom dataset:**
```bash
train pretrain --config configs/pretrain.yaml
```
Starts from scratch, learns representations from raw text via next-token prediction.

**Supervised Fine-tuning:**
```bash
train sft --config configs/sft.yaml
```
Fine-tune a pretrained model on instruction-response pairs for improved task performance.

**Reinforcement Learning (DPO):**
```bash
train rl --config configs/rl_dpo.yaml
```
Align models with human preferences using Direct Preference Optimization.

**Interactive Chat:**
```bash
tui-chat --model checkpoints/model.pt
```
Launch an interactive TUI to chat with your model and monitor generation metrics in real-time.

### Configuration

All training is configured via YAML with Pydantic validation. Configs are type-safe and automatically validated:

```yaml
# configs/sft.yaml
model:
  architecture_type: "mla"  # DeepSeek MLA with GQA
  hidden_dim: 2048
  num_layers: 24
  num_heads: 32
  d_latent_kv: 1536  # KV compression factor

training:
  num_epochs: 3
  batch_size: 32
  learning_rate: 1e-4
  warmup_ratio: 0.1
  max_grad_norm: 1.0

optimizer:
  optimizer_type: "muon_adamw"  # Hybrid Muon + AdamW
  muon_momentum: 0.95

data:
  dataset_type: "sft_jsonl"  # or "sft_hf" for HuggingFace
  path: "data/sft_training.jsonl"
  
logging:
  log_to_aim: true
  aim_repo: "/tmp/aim_logs"
```

See `configs/` for complete examples.

## Project Architecture

```
src/taoTrain/
├── cli.py                      # Main CLI entry point
├── config.py                   # Pydantic configuration schemas
│
├── core/                       # Base abstractions
│   └── base.py                 # BaseModel, BaseDataset, BaseTrainer
│
├── models/                     # Pluggable architecture system
│   ├── registry.py             # Architecture factory with @register_architecture
│   ├── taonet.py               # SimpleLLM with DeepSeek MLA
│   ├── mla_components.py       # KV compression, GQA, YaRN
│   ├── embeddings.py           # Factorized embeddings
│   └── transformer.py          # Standard Transformer reference
│
├── data/                       # Advanced data pipeline
│   ├── factory.py              # Dataset factory (HF + JSONL backends)
│   ├── async_loader.py         # Async batch iteration (no I/O bottleneck)
│   ├── tokenization_queue.py   # Background multi-threaded tokenization
│   ├── chunk_manager.py        # Stream billion-token JSONL files
│   ├── hf_pretrain.py          # HuggingFace pretraining datasets
│   ├── hf_sft.py               # HuggingFace SFT datasets
│   ├── hf_rl.py                # HuggingFace RL datasets
│   ├── pretrain_jsonl.py       # JSONL pretraining
│   ├── sft_jsonl.py            # JSONL SFT with instructions
│   └── rl_jsonl.py             # JSONL RL with preferences
│
├── training/                   # Unified training infrastructure
│   └── trainer.py              # Trainer + PretrainTrainer, SFTTrainer, RLTrainer
│
├── optimizers/                 # Pluggable optimizer system
│   ├── registry.py             # Optimizer factory with @register_optimizer
│   ├── hybrid_muon_adamw.py    # Composite: Muon (2D) + AdamW (1D)
│   ├── adamw.py                # AdamW with weight decay
│   ├── adam.py                 # Standard Adam
│   └── sgd.py                  # SGD variants
│
├── schedulers/                 # Learning rate schedules
│   ├── registry.py             # LR scheduler factory
│   ├── cosine_warmup.py        # 3-phase: linear warmup → plateau → cosine decay
│   ├── linear_warmup.py        # Linear warmup + constant
│   └── constant.py             # Constant learning rate
│
├── inference/                  # Inference & interaction
│   ├── inferencer.py           # Load & run inference from checkpoints
│   └── tui.py                  # Interactive chat with metrics display
│
├── checkpointing/              # State management
│   └── checkpoint.py           # Save/load model + optimizer + config + metrics
│
├── logging/                    # Experiment tracking
│   └── aim_logger.py           # AimStack integration (loss, metrics, hyperparams)
│
├── benchmarks/                 # Evaluation tools
│   └── runner.py               # Perplexity, speed, and task-specific benchmarks
│
└── utils/
    └── helpers.py              # Utility functions

configs/                        # Example YAML configurations
├── pretrain.yaml               # Pretraining config
├── sft.yaml                    # SFT config
├── rl_dpo.yaml                 # RL/DPO config
└── tokenizer.yaml              # Tokenizer config

tests/                          # Unit & integration tests
└── test_dataset.py
```

## Extensible Architecture: The Registry Pattern

TaoTrain's power lies in its **pluggable design**. Add custom models, optimizers, schedulers, and datasets without modifying the framework.

### Custom Model Architecture

```python
from taoTrain.models import register_architecture, BaseModel
import torch.nn as nn

@register_architecture("custom_moe")
class MixtureOfExperts(BaseModel):
    """Your custom MoE architecture"""
    def __init__(self, config):
        super().__init__(config)
        self.experts = nn.ModuleList([
            nn.Linear(config.hidden_dim, config.hidden_dim)
            for _ in range(config.num_experts)
        ])
        self.router = nn.Linear(config.hidden_dim, config.num_experts)
    
    def forward(self, input_ids, attention_mask=None):
        # Your implementation
        logits = self.compute_logits(input_ids)
        loss = self.compute_loss(logits, labels) if labels is not None else None
        return {"logits": logits, "loss": loss}
```

Then use it in your config:

```yaml
model:
  architecture_type: "custom_moe"
  hidden_dim: 2048
  num_experts: 8
```

### Custom Optimizers & Schedulers

The same pattern works for optimizers and learning rate schedules:

```python
from taoTrain.optimizers import register_optimizer
from torch.optim import Optimizer

@register_optimizer("my_adaptive_optimizer")
class MyAdaptiveOptimizer(Optimizer):
    def step(self, closure=None):
        # Your optimization logic
        pass
```

```python
from taoTrain.schedulers import register_scheduler

@register_scheduler("my_schedule")
def my_schedule(initial_lr, step, total_steps, **kwargs):
    return initial_lr * (1.0 - step / total_steps)  # Linear decay
```

**The key principle**: No framework code needs to change. You register once, it's available everywhere.

### Dataset Backend Flexibility

Define custom datasets (JSONL, HF, streaming, etc.) and let the factory route to them:

```python
from taoTrain.data import register_dataset

@register_dataset("pretrain", "my_backend")
class MyPretrainDataset(BaseDataset):
    def __init__(self, config):
        # Load from your custom backend
        pass
    
    def __getitem__(self, idx):
        return {"input_ids": ..., "attention_mask": ...}
```

Use in config:

```yaml
data:
  dataset_type: "pretrain"
  backend_type: "my_backend"  # Routes to MyPretrainDataset
```

## Why TaoTrain Framework?

### Async Data Loading: No I/O Bottleneck

Most training frameworks load and tokenize data on the main training thread, blocking compute. TaoTrain's **multi-threaded tokenization pipeline**:

- Tokenizes data in background workers while your GPU trains
- Supports streaming billion-token JSONL files without loading into memory
- Intelligent chunking (by file size or sample count)
- Metadata caching to avoid rescanning

**Result**: 10-100x faster data iteration on large datasets.

### Type-Safe Configuration

Forget YAML parsing errors or mysterious config bugs. TaoTrain uses **Pydantic dataclasses** for configuration:

- Automatic type validation: mistyped `learning_rate: "1e-4"` becomes an error, not silent failure
- Serialization: configs are part of checkpoints, ensuring reproducibility
- IDE support: autocomplete and type hints for all config fields
- Defaults: sensible defaults for all parameters

### Benchmarking & Metrics

Track what matters:

- **Perplexity**: Language modeling quality on held-out data
- **Generation Speed**: Tokens-per-second (useful for TUI or deployment)
- **Task-Specific Accuracy**: Evaluate on downstream tasks
- **Training Metrics**: Loss curves, gradient norms, effective batch size

All logged to AimStack with git hashes for reproducibility.

## Logging with AimStack

Automatically track and visualize experiments:

```bash
aim up --host 0.0.0.0
```

Then open `http://localhost:43800` to see:

- **Loss curves** per training step
- **Hyperparameters** (learning rate, batch size, model architecture)
- **Git hashes** for reproducibility
- **Custom metrics** (perplexity, validation accuracy, generation speed)
- **Compare runs**: Side-by-side experiment comparison

## Advanced Features

### Checkpointing with Resumption

TaoTrain saves complete training state:

```python
checkpoint = {
    "step": 12500,
    "model_state": model.state_dict(),
    "optimizer_state": optimizer.state_dict(),
    "config": config,  # Full config as Pydantic object
    "metrics": metrics_tracker.to_dict(),
}
```

Resume training from any checkpoint without loss of state. Keep last N checkpoints automatically.

### Mixed Precision Training (BF16)

```yaml
training:
  use_bfloat16: true
  gradient_accumulation_steps: 4
```

- BF16 via `torch.autocast` for ~2x speedup with minimal accuracy loss
- Proper gradient scaling and clipping
- Compatible with all optimizers and architectures

### 3-Phase Learning Rate Schedule

```yaml
scheduler:
  scheduler_type: "cosine_warmup"
  warmup_ratio: 0.1          # 10% of training steps
  steady_ratio: 0.5          # 50% at steady rate
  min_lr_ratio: 0.1          # Final LR = 0.1 × initial_lr
  num_cycles: 1
```

This schedule:
1. **Linear warmup** (0 → 1) over 10% of steps
2. **Steady plateau** at full LR over 50% of steps
3. **Cosine decay** (1 → 0.1) over remaining 40% of steps

Better convergence than simple cosine or linear decay.

### Gradient Accumulation & Clipping

Simulate larger batch sizes with gradient accumulation:

```yaml
training:
  batch_size: 32
  gradient_accumulation_steps: 4  # Effective batch = 128
  max_grad_norm: 1.0               # Gradient clipping
```

## Contributing

Contributions are welcome! TaoTrain is designed to make contributions easy:

1. **Add a model**: Implement `BaseModel` and `@register_architecture("name")`
2. **Add an optimizer**: Implement `torch.optim.Optimizer` and `@register_optimizer("name")`
3. **Add a dataset**: Implement `BaseDataset` and `@register_dataset(mode, backend_type)`
4. **Improve the core**: Submit PRs to `training/`, `data/`, `logging/`, etc.

Ensure new code includes:
- Type hints throughout
- Pydantic configs for new parameters
- Unit tests in `tests/`
- Documentation in docstrings and README

## Current Scope & Roadmap

### ✅ Currently Supported

- **Single GPU / single node** training
- **Pretraining, SFT, and RL training** stages
- **HuggingFace and JSONL** data backends
- **BF16 mixed precision** training
- **Checkpoint saving/loading** with resumption
- **Interactive inference** via TUI
- **Benchmarking** (perplexity, speed)
- **Pluggable architectures, optimizers, schedulers, datasets**

### 🚀 Roadmap (Future)

- **Distributed training** (DDP, FSDP) for multi-GPU/multi-node scaling
- **Quantization** support (INT8, QLoRA)
- **Advanced evaluation** (BLEU, ROUGE, custom tasks)
- **Streaming inference** with KV cache
- **Speculative decoding** for faster generation
- **Integration with popular model hubs** (Hugging Face Hub upload/download)

---

## Getting Help

- **Questions?** Open an issue on GitHub
- **Want to contribute?** See `CONTRIBUTING.md` (coming soon)
- **Found a bug?** Report it with a minimal reproduction script

## License

MIT