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)
- DCLM Baseline (
- 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
pip install -e .
Usage
Run Training
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 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:
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
1 maintainer
Metadata
1 maintainer
Release files for blt-trainer 0.3.3
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | Provenance | |
|---|---|---|---|---|
| blt_trainer-0.3.3.tar.gz | 38.9 kB | Without provenance | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| blt_trainer-0.3.3-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 75.8 kB
Release files / blt_trainer-0.3.3.tar.gz
| Download URL | blt_trainer-0.3.3.tar.gz |
|---|---|
| Size | 38.9 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
241ad0d0f080239a5f5667f69f4d59ec930254282dd97bcb97bd07aa53c458a1
|
|
BLAKE2b-256 checksum How to use checksums |
4c7ca33c2afc5b050d605a9f1ed9a89a671339457e824abd020b1264477e9036
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.14.4
|
Release files / blt_trainer-0.3.3-py3-none-any.whl
| Download URL | blt_trainer-0.3.3-py3-none-any.whl |
|---|---|
| Size | 37.0 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
918367b7d581e8d3e782297064a8967920ab1467a766b18d02a66d734a7d7d56
|
|
BLAKE2b-256 checksum How to use checksums |
9281a3bff66e668ec3d3bf7870bf066fae014bc5cf28160ed2dfb276416e98d9
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.14.4
|
Security for blt-trainer 0.3.3
Provenance
Provenance describes where a file came from. On PyPI, provenance is shared via attestations, which provide a verifiable record of the build or publishing details. View details, limitations and caveats.
No provenance found
None of the files in this release were uploaded with provenance metadata.
A note on package security
PyPI is a repository for distributing packages, not a service for auditing them. While third-party security researchers and our staff do their best to detect and remove malicious packages, PyPI cannot guarantee the quality, security or suitability of any given package.
