Byte Latent Transformer (BLT)

If you have access to a SLURM cluster, you can use the provided setup script to build the environment.

git clone https://github.com/facebookresearch/blt
cd blt

bash setup/create_env.sh
# or if you have access to a SLURM cluster
sbatch setup/create_env.sh

After completion, activate the environment using the generated name: conda activate blt_<date>

For a reproducible build using a lock file, use uv. This method installs pre-build groups and compiles xformers before syncing the environment.

uv pip install --group pre_build --no-build-isolation
uv pip install --group compile_xformers --no-build-isolation
uv sync
uv run python download_blt_weights.py
uv run python demo.py "A BLT has"

Use setup/download_prepare_hf_data.py to download and prepare datasets like fineweb_edu, fineweb_edu_10bt, or dclm_baseline_1.0.

Arguments:

• <dataset_name>: e.g., fineweb_edu
• <MEMORY>: Memory allocation for terashuf.
• --data_dir: Directory to save data (default ./data).
• --seed: Random seed.
• --nchunks: Number of chunks. If running on fewer than 32 GPUs, it is recommended to set this to 1 or match the number of GPUs (nchunks = NGPUs).

python setup/download_prepare_hf_data.py fineweb_edu <MEMORY> --data_dir ./data --seed 42 --nchunks <NCHUNKS>

To download a specific tokenizer (e.g., llama3), use the setup/download_tokenizer.py script.

python setup/download_tokenizer.py llama3 <SAVE_PATH> --api_key <HUGGINGFACE_TOKEN>

To set up a local development environment using Conda, clone the repository and follow these steps to create a Python 3.12 environment with the necessary dependencies, including nightly PyTorch and a specific version of xformers.

git clone https://github.com/facebookresearch/blt
cd blt
conda create -n blt python=3.12
conda activate blt
pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121
pip install ninja
pip install -v -U git+https://github.com/facebookresearch/xformers.git@de742ec3d64bd83b1184cc043e541f15d270c148
pip install -r requirements.txt

The TrainState class (implementing Stateful) tracks the progress of the training loop and is used for checkpointing. It ensures that training can be resumed from the exact point of interruption.

Key fields:

• step: The current number of optimizer steps taken.
• acc_step: The current number of accumulation steps completed since the last optimizer step.
• scheduler: The current state of the learning rate scheduler.
• data_loader_state: The state of the data loader (e.g., PackTokensState) to ensure data continuity upon restart.

@dataclass
class TrainState(Stateful):
    step: int  # Nb of steps taken by the optimizer
    acc_step: int  # Nb of accumulation steps done since last optimizer step
    scheduler: lr_scheduler.LambdaLR
    data_loader_state: PackTokensState

The training script uses OmegaConf to allow hierarchical configuration via the command line. You can provide a base configuration file and override specific nested attributes using dot notation.

Configuration Merging Order:

1. Default values in TrainArgs dataclass.
2. Values from the provided configuration file (--config <path>).
3. Values provided via command line arguments.

Example Usage: If your config file defines model.dim: 128, you can override it to 64 by passing model.dim=64 as a CLI argument.

To use BLT models in your own code, you can load the LMTransformer, ByteLatentTransformer, and BltTokenizerAndPatcher from their respective HuggingFace repositories. Note that you must first request access to the weights on HuggingFace and log in via huggingface-cli login.

from bytelatent.transformer import LMTransformer
from bytelatent.model.blt import ByteLatentTransformer
from bytelatent.hf import BltTokenizerAndPatcher

entropy_repo = "facebook/blt-entropy"
blt_repo = "facebook/blt-1b"

# Load components
entropy_model = LMTransformer.from_pretrained(entropy_repo)
blt_model = ByteLatentTransformer.from_pretrained(blt_repo)
tok_and_patcher = BltTokenizerAndPatcher.from_pretrained(blt_repo)

# Build tokenizer and patcher
tokenizer = tok_and_patcher.tokenizer_args.build()
patcher = tok_and_patcher.patcher_args.build()

You can load weights directly from the HuggingFace hub using the bytelatent.hf module via the command line.

python -m bytelatent.hf load-transformers --entropy-repo facebook/blt-entropy --blt-repo facebook/blt-1b --prompt "My test prompt" hub