Metadata-Version: 2.5
Name: blt-trainer
Version: 0.1.6
Summary: Pure functional 96-layer Byte Latent Transformer trainer with Muon and weighted datamix
Requires-Python: >=3.10
Requires-Dist: datasets>=2.16.0
Requires-Dist: numpy>=1.24.0
Requires-Dist: torch>=2.2.0
Requires-Dist: zstandard>=0.22.0
Description-Content-Type: text/markdown

# BLT Trainer

A pure functional PyTorch implementation of a 96-layer **Byte Latent Transformer (BLT)** with Muon + ScheduleFreeAdamW optimization and weighted multi-corpus datamix.

## Features

- **Byte Latent Architecture**: Vectorized entropy/surprise-based dynamic byte patching and cross-attention pooling.
- **Targeted Shard Downloading**: Massive corpora (such as DCLM at 4TB+ and OpenWebMath at 16GB+) only download the necessary training shards (~400MB total) instead of the entire multi-terabyte repositories. Datasets are cached locally using Apache Arrow to prevent Hugging Face API rate-limiting during training.
- **Weighted Multi-Corpus Sampling**:
  - DCLM Baseline (`mlfoundations/dclm-baseline-1.0`): ~100M token shard
  - FLAN Collection (`Muennighoff/flan`): Representative reasoning & QA shards
  - Wikipedia QA (`microsoft/wiki_qa`): Full split (~20MB)
  - OpenWebMath Reasoning (`open-web-math/open-web-math`): Initial shard (~130M tokens)
  - Ubuntu Dialogue QA (`sedthh/ubuntu_dialogue_qa`): Full split (~30MB)
- **Advanced Logging**: Real-time throughput reporting with **TPS (Tokens Per Second)**, **Tokens Processed**, step latency, and ETA.
- **Dual Optimizer Partitioning**: Muon optimizer for 2D hidden projection matrices and Schedule-Free AdamW for 1D embeddings and norm weights.
- **Cut Cross-Entropy**: Chunked sequence loss computation to eliminate peak VRAM spikes from materializing large vocabulary logits.

## Installation

```bash
pip install -e .
```

## Usage

### Run Training

```bash
blt-train --batch-size 32 --block-size 512 --steps 10000 --device cuda
```

### CLI Options

| Flag | Type | Default | Description |
|---|---|---|---|
| `--steps` | int | `10000` | Total training steps |
| `--batch-size` | int | `8` | Batch size per step |
| `--block-size` | int | `512` | Byte context length |
| `--streaming` | flag | `False` | Stream from Hugging Face instead of downloading locally (default: False, downloaded) |
| `--cache-dir` | str | `None` | Custom directory to cache downloaded datasets |
| `--dclm-shards` | int | `1` | Number of DCLM shards to download (default: 1, ~100M tokens) |
| `--download-full-datasets` | flag | `False` | Download entire corpora (warning: DCLM is 4TB+) |
| `--log-interval` | int | `10` | Frequency of progress logs in steps |
| `--device` | str | `cuda` (if available) else `cpu` | Execution device |
| `--compile` | flag | `False` | Enable `torch.compile` kernel fusion across forward & optimizer passes |
| `--compile-mode` | str | `default` | Optimization mode (`default`, `reduce-overhead`, `max-autotune`) |
| `--lr-muon` | float | `0.02` | Learning rate for Muon optimizer |
| `--lr-adam` | float | `1e-3` | Learning rate for Schedule-Free AdamW |

### Logging Output

Logs include step progress, loss, **Tokens Per Second (TPS)**, **Tokens Processed**, step latency, and estimated time of arrival (ETA):

```text
Step 00010/10000 (  0.1%) | Loss: 6.4215 | TPS:  14,250.3 | Tokens: 40,960 (40.96K) | Step Time:   35.9ms | ETA: 00:05:55 | VRAM: 3.45GB
Step 00020/10000 (  0.2%) | Loss: 6.1023 | TPS:  14,310.0 | Tokens: 81,920 (81.92K) | Step Time:   35.7ms | ETA: 00:05:52 | VRAM: 3.45GB
```
