nanotron

环境要求

• Python 3.11（README 使用 --python 3.11 创建虚拟环境）
• uv（用于创建虚拟环境和安装依赖）
• 带 CUDA 的 PyTorch（README 安装的是 cu124 版本，需 NVIDIA GPU 环境）
• 已训练好的 checkpoint 目录（运行生成脚本时需要）

安装与启动步骤

1. 1创建虚拟环境

   用 uv 创建名为 nanotron 的 Python 3.11 虚拟环境并激活，同时升级 pip。HF 集群用户建议设置 UV_LINK_MODE=copy。

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

   复制

2. 2安装 PyTorch

   从 PyTorch 官方 cu124 索引安装 PyTorch，默认适配 CUDA 12.4 环境。

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

   复制

3. 3安装核心依赖

   在项目根目录以可编辑模式安装 nanotron 本体及核心依赖。

   uv pip install -e .

   复制

4. 4安装示例依赖

   要运行示例脚本，再安装数据集、模型加载、日志等可选依赖。

   uv pip install datasets transformers datatrove[io] numba wandb

   复制

5. 5运行文本生成

   用 torchrun 从已有 checkpoint 启动生成脚本；1 张卡时 tp、pp 都设为 1。

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

   复制

关键配置

如何确认成功

torchrun 命令正常退出且无报错，并输出模型生成的文本结果。

常见问题

注意事项

• 所有命令需先激活虚拟环境（source nanotron/bin/activate）后再执行。
• 命令中的 checkpoints/{checkpoint_number}/ 要替换成你实际的 checkpoint 目录。
• --tp、--pp 的取值需与可用 GPU 数量匹配，否则可能启动失败。
• README 只给出了生成脚本 run_generate.py 的示例，其余训练脚本请自行查阅仓库文档。