Metadata-Version: 2.5
Name: blt-trainer
Version: 0.3.0
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 |
| `--lr-scheduler` | str | `wsd` | LR scheduler: `wsd` (Warmup-Stable-Decay) or `none` |
| `--warmup-ratio` | float | `0.02` | Warmup fraction of total steps (default: 2%) |
| `--decay-ratio` | float | `0.15` | Decay fraction of total steps (default: 15%) |
| `--min-lr-ratio` | float | `0.05` | Minimum LR floor at end of decay phase |
| `--decay-type` | str | `cosine` | Decay curve: `cosine` or `linear` |
| `--curriculum` | str | `axiomic` | Data curriculum: `axiomic` (GPT-X2.5-135M progressive curriculum) or `static` |
| `--precision` | str | `bf16` | Precision: `bf16` (Blackwell Tensor Cores) or `fp32` |
| `--grad-accum` | int | `1` | Gradient accumulation steps |
| `--num-workers` | int | `1` | Total disjoint GPU workers (e.g. 6) |
| `--worker-rank` | int | `0` | Rank of this GPU worker (0 to `num-workers - 1`) |
| `--diloco` | flag | `False` | Enable DiLoCo distributed outer optimizer |
| `--diloco-repo` | str | `None` | Hugging Face repository ID / bucket for outer model sync |
| `--diloco-repo-type` | str | `model` | Hugging Face repo type (`model` or `dataset`) |
| `--diloco-interval` | int | `500` | Inner local training steps between outer synchronizations |
| `--diloco-lr` | float | `0.7` | Outer SGD learning rate |
| `--diloco-momentum` | float | `0.9` | Outer Nesterov momentum |
| `--diloco-dir` | str | `./diloco_state` | Local directory for outer checkpoints and staging |
| `--hf-token` | str | `None` | Hugging Face auth token (or set `HF_TOKEN` environment variable) |

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