Metadata-Version: 2.5
Name: blt-trainer
Version: 0.3.3
Summary: Pure functional 96-layer Byte Latent Transformer trainer with Muon, DiLoCo HF Buckets, Open-SLM benchmarks, and Axiomic curriculum
Requires-Python: >=3.10
Requires-Dist: datasets>=2.16.0
Requires-Dist: huggingface-hub>=0.20.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 (scale token count per step) |
| `--max-grad-norm` | float | `None` | Optional gradient clipping threshold |
| `--grad-explosion-threshold` | float | `50.0` | Gradient norm threshold to trigger console explosion alert |
| `--eval-interval` | int | `0` | Evaluate Open-SLM-Leaderboard benchmarks every N steps (`0` = disabled) |
| `--eval-limit` | int | `100` | Max examples per benchmark during evaluation (`0` for full dataset) |
| `--eval-benchmarks` | str | `all` | Comma-separated benchmarks (`hellaswag,arc,arcChall,piqa,arithmark3`) |
| `--eval-save-path` | str | `./eval_history.jsonl` | File path to append evaluation metrics JSON |
| `--num-workers` | int | `1` | Total disjoint GPU workers in cluster (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-bucket` | str | `None` | Hugging Face Storage Bucket ID for outer sync (high API rate limit) |
| `--diloco-repo` | str | `None` | Alias for `--diloco-bucket` |
| `--diloco-storage` | str | `bucket` | Storage backend: `bucket` (native S3/object storage) or `repo` (Git repo) |
| `--diloco-repo-type` | str | `model` | Hugging Face repo type (`model` or `dataset`) if using repo mode |
| `--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) |

### Open-SLM-Leaderboard Intelligence Index

Tracks the official [Open-SLM-Leaderboard](https://huggingface.co/spaces/AxiomicLabs/Open_SLM_Leaderboard) Intelligence Index (Int Index):

$$\text{Chance Normalized Score} = \frac{\text{Raw Acc} - \text{Chance}}{100 - \text{Chance}} \times 100$$

$$\text{Int Index} = \frac{\text{HellaSwag}_{\text{norm}} + \text{Combined ARC}_{\text{norm}} + \text{PIQA}_{\text{norm}} + 0.65 \times \text{ArithMark3}_{\text{norm}}}{3.65}$$

- **HellaSwag**: Weight 1.0, Random baseline 25%
- **Combined ARC**: Weight 1.0, Random baseline 25% (mean of ARC-Easy & ARC-Challenge)
- **PIQA**: Weight 1.0, Random baseline 50%
- **ArithMark-3**: Weight 0.65, Random baseline 25%

### Logging Output

Logs include step progress, loss, **Gradient Norm**, **Int Index**, **Tokens Per Second (TPS)**, **Tokens Processed**, step latency, and ETA:

```text
Step 00050/20000 (  0.2%) | Loss: 2.7459 | Grad: 0.4215 | IntIdx: 21.07 | TPS: 135,559.1 | LR: 0.10x (warmup) | Tokens: 3.28M | Step Time: 483.4ms | ETA: 02:49:18 | VRAM: 16.90GB
```
