Levanter

你无法阻止雷暴，但你可以利用电力；你无法控制风向，但你可以调整帆布，让船按你的意愿前进，无论风从哪个方向吹来。
— Cora L. V. Hatch

Levanter 是一个用于训练大型语言模型（LLMs）和其他基础模型的框架，致力于实现可读性、可扩展性和可重现性：

可读性：Levanter 使用我们的命名张量库 Haliax 来编写易于理解、可组合的深度学习代码，同时保持高性能。
可扩展性：Levanter 可扩展到大型模型，并能够在各种硬件上进行训练，包括 GPU 和 TPU。
可重现性：Levanter 具有逐位确定性，这意味着相同的配置将始终产生相同的结果，即使在面对抢占和恢复的情况下也是如此。

我们使用 JAX、Equinox 和 Haliax 构建了 Levanter。

文档

Levanter 的文档可在 levanter.readthedocs.io 获取。 Haliax 的文档可在 haliax.readthedocs.io 获取。

特性

分布式训练：我们支持在 TPU（不久后还将支持 GPU）上进行分布式训练，包括 FSDP 和张量并行。
兼容性：Levanter 支持与 Hugging Face 生态系统之间的模型导入和导出，包括分词器、数据集和通过 SafeTensors 的模型。
性能：Levanter 的性能可与商业支持的框架（如 MosaicML 的 Composer 或 Google 的 MaxText）相媲美。
缓存式按需数据预处理：我们在线预处理语料库，但会缓存预处理结果，使得恢复训练更快，后续运行甚至更快。一旦缓存的第一部分完成，Levanter 就会开始训练。
优化：Levanter 支持新的 Sophia 优化器，它可能比 Adam 快 2 倍。我们还支持使用 Optax 进行 AdamW 等优化。
日志记录：Levanter 支持多种不同的日志后端，包括 WandB 和 TensorBoard。（添加新的日志后端很容易！）Levanter 甚至提供了在 JAX jit 函数内部记录日志的能力。
可重现性：在 TPU 上，Levanter 具有逐位确定性，这意味着相同的配置将始终产生相同的结果，即使在面对抢占和恢复的情况下也是如此。
分布式检查点：通过 Google 的 TensorStore 库支持分布式检查点。训练甚至可以在不同数量的主机上恢复，尽管目前这会破坏可重现性。

Levanter 由斯坦福基础模型研究中心（CRFM）的研究工程团队创建。你也可以在非官方的 Jax LLM Discord 的 #levanter 频道找到我们。

入门

以下是一些帮助你入门的简单示例。有关各种配置选项的更多信息，请参阅入门指南或深入配置指南。你也可以使用 --help 或查看其他配置文件，以了解所有可用的选项。

安装 Levanter

在根据你的平台安装 JAX 之后，你可以通过以下方式安装 Levanter：

pip install levanter

或者使用 GitHub 上的最新版本：

git clone https://github.com/stanford-crfm/levanter.git
cd levanter
pip install -e .
wandb login  # 可选，我们使用 wandb 进行日志记录

如果你同时开发 Haliax 和 Levanter，可以这样做：

git clone https://github.com/stanford-crfm/levanter.git
cd levanter
pip install -e .
cd ..
git clone https://github.com/stanford-crfm/haliax.git
cd haliax
pip install -e .
cd ../levanter

有关如何安装 Levanter 的更多信息，请参阅安装指南。

如果你使用 TPU，更完整的设置文档可在此处获取。GPU 支持仍在进行中；文档可在此处获取。

训练 GPT2-nano

作为一种"Hello World"示例，以下是如何在小型数据集上训练 GPT-2 "nano" 大小的模型。

python -m levanter.main.train_lm --config_path config/gpt2_nano.yaml

# 或者，如果你没有使用 -e 并且在不同的目录中
python -m levanter.main.train_lm --config_path gpt2_nano

这将在 WikiText-103 数据集上训练一个 GPT2-nano 模型。

在你自己的数据上训练 GPT2-small

你也可以通过更改配置文件中的 dataset 字段来更改数据集。如果你的数据集是 Hugging Face 数据集，你可以使用 data.id 字段来指定它：

python -m levanter.main.train_lm --config_path config/gpt2_small.yaml --data.id openwebtext

# 可选地，你可以指定一个分词器和/或缓存目录，可以是本地的或在 gcs 上的
python -m levanter.main.train_lm --config_path config/gpt2_small.yaml --data.id openwebtext --data.tokenizer "EleutherAI/gpt-neox-20b" --data.cache_dir "gs://path/to/cache/dir"

如果你的数据是一个 URL 列表，你可以使用 data.train_urls 和 data.validation_urls 字段来指定它们。数据 URL 可以是本地文件、gcs 文件或 http(s) URL，或任何 fsspec 支持的格式。 Levanter（实际上是 fsspec）将自动解压 .gz 和 .zstd 文件，可能还支持其他格式。

python -m levanter.main.train_lm --config_path config/gpt2_small.yaml --data.train_urls ["https://path/to/train/data_*.jsonl.gz"] --data.validation_urls ["https://path/to/val/data_*.jsonl.gz"]

自定义配置文件

你可以修改配置文件来更改模型、数据集、训练参数等。以下是 gpt2_small.yaml 文件的内容：

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: 768
    num_heads: 12
    num_layers: 12
    seq_len: 1024
    gradient_checkpointing: true
    scale_attn_by_inverse_layer_idx: true
trainer:
  tracker:
    type: wandb
    project: "levanter"
    tags: [ "openwebtext", "gpt2"]

  mp: p=f32,c=bfloat16
  model_axis_size: 1
  per_device_parallelism: 4

  train_batch_size: 512
optimizer:
  learning_rate: 6E-4
  weight_decay: 0.1
  min_lr_ratio: 0.1

其他架构

目前，我们支持以下架构：

GPT-2
LLama 1 或 2
Backpacks
MosaicML 的 MPT

我们计划在未来添加更多架构。

使用 Llama 1 或 Llama 2 继续预训练

以下是如何在 OpenWebText 数据集上继续预训练 Llama 1 或 Llama 2 模型的示例：

python -m levanter.main.train_lm --config_path config/llama2_7b_continued.yaml

分布式和云训练

在 TPU 云虚拟机上训练

请参阅 TPU 入门指南以获取有关如何设置 TPU 云虚拟机并在那里运行 Levanter 的更多信息。

使用 CUDA 训练

请参阅 CUDA 入门指南以获取有关如何设置 CUDA 环境并在那里运行 Levanter 的更多信息。

贡献

我们欢迎贡献！请查看 CONTRIBUTING.md 以获取更多信息。

许可证

Levanter 在 Apache License 2.0 下授权。完整的许可证文本请参见 LICENSE。