Package: levanter
LevanterYou could not prevent a thunderstorm, but you could use the electricity; you could not direct the wind, but you could trim your sail so as to propel your vessel as you pleased, no matter which way the wind blew.— Cora L. V. HatchLevanter is a framework for training large language models (LLMs) and other foundation models that strives for legibility, scalability, and reproducibility:Legible: Levanter uses our named tensor libraryHaliaxto write easy-to-follow, composable deep learning code, while still being high performance.Scalable: Levanter scales to large models, and to be able to train on a variety of hardware, including GPUs and TPUs.Reproducible: Levanter is bitwise deterministic, meaning that the same configuration will always produce the same results, even in the face of preemption and resumption.We built Levanter withJAX,Equinox,
andHaliax.FeaturesDistributed Training: We support distributed training on TPUs (and soon, GPUs), including FSDP and tensor parallelism.Compatibility: Levanter supports importing and exporting models to/from the Hugging Face ecosystem, including tokenizers, datasets, and models viaSafeTensors.Performance: Levanter's performance rivals commercially-backed frameworks like MosaicML's Composer or Google's MaxText.Reproducibility: Levanter is bitwise deterministic, meaning that the same configuration will always produce the same results, even in the face of preemption and resumption.Cached On-Demand Data Preprocessing: We preprocess corpora online, but we cache the results of preprocessing so
that resumes are much faster and so that subsequent runs are even faster. As soon as the first part of the cache is complete, Levanter will start training.Logging: Logging is done withWandB, complete with a fancy online visualization of the validation set during training.Distributed Checkpointing: Distributed checkpointing is supported via Google'sTensorStorelibrary. Training can even be resumed on a different number of hosts, though this breaks reproducibility for now.Optimization: Levanter usesOptaxfor optimization. Our new optimizer,Sophia, is coming to Levanter soon!Levanter was created byStanford's Center for Research on Foundation Models (CRFM)'s research engineering team. (We're hiring!) You can also find us in the #levanter channel on the unofficialJax LLM DiscordGetting StartedHere is a small set of examples to get you started. For more information about the various configuration options,
please see theGetting Startedguide or theIn-Depth Configuration Guide.
You can also use--helpor poke around other configs to see all the options available to you.DocumentationLevanter's documentation is available atlevanter.readthedocs.io.
Haliax's documentation is available athaliax.readthedocs.io.Installing LevanterAfterinstalling JAXwith the appropriate configuration
for your platform, you can install Levanter with:pipinstalllevanteror using the latest version from GitHub:gitclonehttps://github.com/stanford-crfm/levanter.gitcdlevanter
pipinstall-e.
wandblogin# optional, we use wandb for loggingIf you're developing Haliax and Levanter at the same time, you can do something like.gitclonehttps://github.com/stanford-crfm/levanter.gitcdlevanter
pipinstall-e.cd..
gitclonehttps://github.com/stanford-crfm/haliax.gitcdhaliax
pipinstall-e.cd../levanterPlease refer to theInstallation Guidefor more information on how to install Levanter.If you're using a TPU, more complete documentation for setting that up is availablehere. GPU support is still in-progress; documentation is availablehere.Training a GPT2-nanoAs a kind of hello world, here's how you can train a GPT-2 "nano"-sized model on a small dataset.python-mlevanter.main.train_lm--config_pathconfig/gpt2_nano.yaml# alternatively, if you didn't use -e and are in a different directorypython-mlevanter.main.train_lm--config_pathgpt2_nanoThis will train a GPT2-nano model on theWikiText-103dataset.Training a GPT2-small on your own dataYou can also change the dataset by changing thedatasetfield in the config file.
If your dataset is aHugging Face dataset, you can use thedata.idfield to specify it:python-mlevanter.main.train_lm--config_pathconfig/gpt2_small.yaml--data.idopenwebtext# optionally, you may specify a tokenizer and/or a cache directory, which may be local or on gcspython-mlevanter.main.train_lm--config_pathconfig/gpt2_small.yaml--data.idopenwebtext--data.tokenizer"EleutherAI/gpt-neox-20b"--data.cache_dir"gs://path/to/cache/dir"If instead your data is a list of URLs, you can use thedata.train_urlsanddata.validation_urlsfields to specify them.
Data URLS can be local files, gcs files, or http(s) URLs, or anything thatfsspecsupports.
Levanter (really, fsspec) will automatically uncompress.gzand.zstdfiles, and probably other formats too.python-mlevanter.main.train_lm--config_pathconfig/gpt2_small.yaml--data.train_urls["https://path/to/train/data_*.jsonl.gz"]--data.validation_urls["https://path/to/val/data_*.jsonl.gz"]Customizing a Config FileYou can modify the config file to change the model, the dataset, the training parameters, and more. Here's
thegpt2_small.yamlfile:data:train_urls:-"gs://pubmed-mosaic/openwebtext-sharded/openwebtext_train.{1..128}-of-128.jsonl.gz"validation_urls:-"gs://pubmed-mosaic/openwebtext-sharded/openwebtext_val.{1..8}-of-8.jsonl.gz"cache_dir:"gs://pubmed-mosaic/tokenized/openwebtext/"model:gpt2:hidden_dim:768num_heads:12num_layers:12seq_len:1024gradient_checkpointing:truescale_attn_by_inverse_layer_idx:truetrainer:wandb:project:"levanter"tags:["openwebtext","gpt2"]mp:p=f32,c=bfloat16model_axis_size:1per_device_parallelism:4train_batch_size:512optimizer:learning_rate:6E-4weight_decay:0.1min_lr_ratio:0.1Other ArchitecturesCurrently, we support the following architectures:GPT-2LLama 1 or 2BackpacksMosaicML'sMPTWe plan to add more in the future.Continued Pretraining with Llama 1 or Llama 2Here's an example of how to continue pretraining a Llama 1 or Llama 2 model on the OpenWebText dataset:python-mlevanter.main.train_lm--config_pathconfig/llama2_7b_continued.yamlDistributed and Cloud TrainingTraining on a TPU Cloud VMPlease see theTPU Getting Startedguide for more information on how to set up a TPU Cloud VM and run Levanter there.Training with CUDAPlease see theCUDA Getting Startedguide for more information on how to set up a CUDA environment and run Levanter there.ContributingWe welcome contributions! Please seeCONTRIBUTING.mdfor more information.LicenseLevanter is licensed under the Apache License, Version 2.0. SeeLICENSEfor the full license text.