⚡️ Nanotron

安装 • 快速入门 • 功能 • 基准测试 • 贡献

预训练模型，轻松搞定

Nanotron 是一个用于预训练 Transformer 模型的库。它提供了一个简单而灵活的 API，用于在自定义数据集上预训练模型。Nanotron 的设计宗旨是易用、快速且可扩展。其构建遵循以下原则：

• 简单性：Nanotron 设计简洁易用，提供了一个简单灵活的 API，方便用户在自定义数据集上进行模型预训练。
• 性能：Nanotron 在速度和可扩展性方面进行了优化，采用最新技术，使模型训练更快、更高效。

📚 查看我们的 超大规模实践手册 —— 一份使用 Nanotron 高效扩展 LLM 训练的全面指南！

📝 AI 生成的文档，感谢 DeepWiki

安装

要运行该项目中的代码，首先使用例如 uv 创建一个 Python 虚拟环境：

uv venv nanotron --python 3.11 && source nanotron/bin/activate && uv pip install --upgrade pip

[!提示] 对于 Hugging Face 集群用户，请将 export UV_LINK_MODE=copy 添加到您的 .bashrc 文件中，以抑制 uv 的缓存警告。

接下来，安装 PyTorch：

uv pip install torch --index-url https://download.pytorch.org/whl/cu124

然后安装核心依赖项：

uv pip install -e .

要运行示例脚本，还需安装以下依赖项：

uv pip install datasets transformers datatrove[io] numba wandb
# 融合内核
uv pip install ninja triton "flash-attn>=2.5.0" --no-build-isolation

接下来，按照如下方式登录您的 Hugging Face 和 Weights & Biases 账户：

huggingface-cli login
wandb login

最后，检查您的系统是否已安装 Git LFS，以便能够加载和推送模型/数据集到 Hugging Face Hub：

git-lfs --version

如果未安装，请运行：

sudo apt-get install git-lfs

快速入门

训练一个小型 Llama 模型

以下命令将在单节点的 8 张 H100 显卡上大约 10 分钟内训练一个小型 Llama 模型：

CUDA_DEVICE_MAX_CONNECTIONS=1 torchrun --nproc_per_node=8 run_train.py --config-file examples/config_tiny_llama.yaml

模型将保存在配置文件中指定的 checkpoints 目录下。

[!注意] 您可以使用 examples/config_tiny_llama.py 生成您自己的训练配置

有关如何训练您的第一个模型的详细说明，请参阅我们的 您的第一个训练指南。如需使用 Slurm 进行多节点训练，请参阅我们的 多节点训练指南。

从您的检查点运行生成

torchrun --nproc_per_node=1 run_generate.py --ckpt-path checkpoints/{checkpoint_number}/ --tp 1 --pp 1

增加 --tp（张量并行）的值以利用多 GPU 加速生成，并为非常大的模型使用更大的 --pp（流水线并行）值。

使用 VSCode 调试

要使用 VSCode 调试，请将以下配置添加到您的 launch.json 文件中：

{
    "name": "run_train.py",
    "type": "python",
    "request": "launch",
    "program": "torchrun", // 或者通过运行 `which torchrun` 获取 torchrun 的完整路径
    "console": "integratedTerminal",
    "justMyCode": false,
    "args": [
        "--nproc_per_node=2",
        "run_train.py",
        "--config-file=examples/config_tiny_llama.yaml", // 或者使用 examples/config_tiny_llama.py 生成您自己的配置
    ],
    "env": {
        // "NANOTRON_BENCHMARK": "1", // 启用此选项可在几个步骤内对您的训练进行基准测试
        "CUDA_DEVICE_MAX_CONNECTIONS": "1",
        "WANDB_MODE": "disabled",
    }
},

[!注意] 更多信息请参阅 在多 GPU 上调试 Nanotron 示例

自定义示例

您可以在 /examples 目录中找到更多示例：

示例	描述
custom-dataloader	将自定义数据加载器接入 Nanotron
datatrove	使用 datatrove 库加载数据
doremi	使用 DoReMi 加速训练
mamba	训练一个 Mamba 示例模型
moe	训练一个专家混合（MoE）示例模型
mup	使用谱 µTransfer 扩展模型规模
examples/config_tiny_llama_with_s3_upload.yaml	用于自动将检查点上传至 S3

我们正在努力添加更多示例！欢迎提交 PR 添加您自己的示例。🚀

基准测试

我们对 Nanotron 在不同模型大小和配置下的性能进行了广泛的基准测试。完整的基准测试数据、配置和日志都可在我们的 ultrascale-playbook-data 仓库中找到。

上图展示了我们在 Nanotron v0.5 中针对不同模型大小和节点数量所发现的最佳配置，突出了最优的 MFU（模型 FLOPS 利用率）和内存使用情况。这些是我们通过全面基准测试过程确定的最高效训练设置。敬请期待更多优化方案的推出！🚀

有关这些基准测试的详细分析及最佳实践，请参阅我们的 超大规模实践手册。

功能

我们目前支持以下功能：

• 3D 并行（数据并行 + 模型并行 + 流水线并行）
• 面向 MoE 的专家并行
• 流水线并行的 AFAB 和 1F1B 调度策略
• 提供显式的数据并行和流水线并行 API，便于调试
• ZeRO-1 优化器
• FP32 梯度累积
• 参数共享/分片
• 针对大模型的自定义模块检查点保存
• 用于扩展神经网络的谱 µTransfer 参数化方法
• Mamba 示例
• 基于 CUDA 事件的计时功能，可精确测量 GPU 性能

我们的路线图还包括：

• FP8 训练
• ZeRO-3 优化器（即 FSDP）
• torch.compile 支持
• Ring Attention
• 交错式 1f1b 调度策略

致谢

我们衷心感谢所有从事大语言模型研究的同仁，尤其是那些公开分享研究成果、为我们提供巨大启发的团队：NVIDIA 的 Megatron-LM/apex、微软的 DeepSpeed、HazyResearch 的 flash-attn 等。

Nanotron 快速上手指南

Nanotron 是由 Hugging Face 推出的一个用于预训练 Transformer 模型的开源库。它旨在提供简单灵活的 API，同时针对速度和可扩展性进行了深度优化，支持在自定义数据集上高效训练大语言模型（LLM）。

环境准备

在开始之前，请确保您的系统满足以下要求：

• 操作系统: Linux (推荐)
• Python 版本: 3.11
• GPU: 支持 CUDA 的 NVIDIA GPU (示例基于 H100，其他型号需调整配置)
• 必要工具:
  • git 和 git-lfs (用于加载和推送模型/数据集)
  • uv (推荐的 Python 虚拟环境和包管理工具)

检查并安装 Git LFS：

git-lfs --version
# 如果未安装，运行：
sudo apt-get install git-lfs

安装步骤

推荐使用 uv 进行环境管理和依赖安装，以获得更快的速度和更好的兼容性。

1. 创建虚拟环境

uv venv nanotron --python 3.11 && source nanotron/bin/activate && uv pip install --upgrade pip

提示: 如果您在 Hugging Face 集群环境中使用，建议在 .bashrc 中添加 export UV_LINK_MODE=copy 以抑制缓存警告。

2. 安装 PyTorch

根据官方推荐安装 CUDA 12.4 版本的 PyTorch：

uv pip install torch --index-url https://download.pytorch.org/whl/cu124

国内加速建议: 如果下载速度慢，可尝试使用清华或阿里镜像源替换 index-url，例如：--index-url https://pypi.tuna.tsinghua.edu.cn/simple (需确保 PyTorch 版本兼容性，若官方源不可用再考虑镜像)。

3. 安装核心依赖

uv pip install -e .

4. 安装示例运行所需依赖

如果您打算运行示例脚本或进行完整开发，请安装以下额外依赖：

uv pip install datasets transformers datatrove[io] numba wandb
# 安装融合内核 (Fused kernels)
uv pip install ninja triton "flash-attn>=2.5.0" --no-build-isolation

5. 登录账户

为了使用 Hugging Face Hub 和记录实验数据，请登录相关账户：

huggingface-cli login
wandb login

基本使用

训练一个微型 Llama 模型

以下命令将在单个节点（8 张 H100 GPU）上训练一个微型 Llama 模型，预计耗时约 10 分钟。

CUDA_DEVICE_MAX_CONNECTIONS=1 torchrun --nproc_per_node=8 run_train.py --config-file examples/config_tiny_llama.yaml

• 说明:
  • --nproc_per_node=8: 指定使用的 GPU 数量，请根据您的实际硬件调整。
  • --config-file: 指定配置文件。您可以参考 examples/config_tiny_llama.py 生成自定义的训练配置。
  • 训练完成的模型将保存在配置文件中指定的 checkpoints 目录下。

从检查点生成文本

训练完成后，可以使用以下命令加载检查点进行推理生成：

torchrun --nproc_per_node=1 run_generate.py --ckpt-path checkpoints/{checkpoint_number}/ --tp 1 --pp 1

• 参数调整:
  • --tp (Tensor Parallelism): 增加此值可利用多张 GPU 加速生成过程。
  • --pp (Pipeline Parallelism): 对于超大模型，可增加此值以启用流水线并行。

更多示例

项目 /examples 目录下提供了丰富的进阶示例，包括：

• custom-dataloader: 接入自定义数据加载器
• moe: 训练混合专家模型 (MoE)
• mamba: 训练 Mamba 架构模型
• doremi: 使用 DoReMi 加速训练

详细的多节点训练（Slurm）和调试指南请参阅项目文档中的 docs/your-first-training.md 和 docs/multi-node-training.md。