⚡️ Nanotron

理念 • 核心功能 • 安装 • 使用 • 贡献

该库的目的是提供易于使用的分布式原语，以便使用3D并行性有效地训练各种模型。有关库的内部设计或3D并行性的更多信息，请参阅 [docs.md] 和 [3d_parallelism.md]。

理念

• 让它变得更快。至少与其他开源版本一样快。
• 让它更简约。我们实际上不需要支持所有技术和所有版本的3D并行性。重要的是我们能够有效地使用“最好的”。
• 将所有事情都明确而不是透明。随着我们不断前进，当事情顺利时，让事情透明是有益的，但如果一个人不了解所用技术的含义，那么这将是一个糟糕的调试体验。为了减轻这一点，我们选择明确说明其做事的方式

核心功能

我们支持以下功能

• 3D并行性，包括一前一后管道引擎
• ZeRO-1优化器
• FP32梯度累积
• 参数绑定/分片

安装

需求

• Python >= 3.10
• PyTorch >= 2.0.0
• Flash-Attention >= 2.5.0

安装（在新的环境中）

pip install torch
pip install packaging; pip install "flash-attn>=2.5.0"  --no-build-isolation
pip install nanotron

还需要： pip install transformers datasets python-etcd tensorboardX

我们还支持一组您可以安装的变体，您可以使用 pip install -e [$FLAVOR] 安装

• dev：如果您正在开发 nanotron，则使用它。它特别安装了我们的linter机制。在此之上，您还必须运行 pre-commit install。
• test：我们使用pytest来运行测试套件。为了并行运行测试，它将安装pytest-xdist，您可以通过运行pytest -n 12 tests（12为并行测试的数量）来利用它。

快速示例

在/examples目录中，您可以找到一些示例配置文件和一个运行脚本。

您可以使用以下命令运行示例训练：

torchrun --nproc_per_node=8 run_train.py --config-file examples/debug_run_train.yaml

并且使用以下命令运行示例生成：

torchrun --nproc_per_node=8 run_generation.py --ckpt-path checkpoints/text/4

开发指南

如果您打算在nanotron上进行开发，我们建议您安装dev版本：pip install -e ".[dev]"

我们使用pre-commit在每个提交上运行一系列回调，主要是规范化代码，以便使代码库保持一致性。请运行pre-commit install。

对于代码风格检查

pre-commit install
pre-commit run --config .pre-commit-config.yaml --all-files

我们希望添加的功能

• 支持torch.compile
• 更多优化的内核
• 支持Zero3
• 其他PP调度（如Interleaved 1f1b...）
• 环形注意力/序列并行性
• 3D并行MoEs
• 支持更多架构（Mamba...）
• ...

致谢

我们感谢所有从事LLMs工作的人，特别是那些公开分享其工作的人，我们从他们那里获得了巨大的灵感：Nvidia的Megatron-LM/apex，Microsoft的DeepSpeed，HazyResearch的flash-attn