qlabs-eng / dust

# DUST
This is a minimal implementation of our research: 
Train a transformer using forward evaluations and SGD, without backpropagation.
This minimal implementation preserves the paper's estimator and tuned defaults
while omitting execution optimizations used in the full experiments.
`model.py` contains the model, and `dust.py` the algorithm and training loop. A simple
backpropagation SGD baseline uses the same model and data.
## Installation
Requires Linux, Python 3.10+, and an NVIDIA GPU with bfloat16 support.
The following installs PyTorch with CUDA 12.8; use a compatible NVIDIA driver.
 [code] 
## Data
Prepare the 1M and 10M training sets:
 [code] 
This streams FineWeb from Hugging Face and saves `data/train_1M.pt`,
`data/train_10M.pt`, and shared `data/val.pt` and `data/test.pt` files.
Use `--tokens 1M` or `--tokens 10M` to prepare only one training budget.
Each held-out split contains 544 sequences of 2,048 prediction tokens.
The fixed BPE-4096 tokenizer is included; retraining is optional.
Its training recipe and script are in `tokenizer/`.
## DUST
Train with the default population of **16,384** on eight GPUs:
 [code] 
For a smaller, single-GPU run:
 [code] 
Supported populations are **256, 1024, 4096, and 16384**, using respectively
**1, 1, 4, and 8 GPUs** in the paper configurations. Each selects its tuned
per-layer draw allocation, credit decay and optimizer settings. Population counts direct-loss
draws across all GPUs; head and local attention draws are additional. Every GPU
processes the same batch and contributes different perturbations.
Both token budgets use an 8-layer, 512-wide transformer, sequence length 2,048,
and batch size 8. Add `--steps 2` for a short check using the selected dataset.
## Backpropagation baseline
Run SGD backpropagation on one GPU:
 [code] 
## Outputs
Each run saves its configuration, metrics, a validation-selected `best.pt`, and
`result.json` containing the selected checkpoint's test loss. Use a new output
directory for each run. Data, checkpoints, and logs are stored locally.
## License
Released under the MIT License.