{"id":13451288,"url":"https://github.com/OpenBMB/BMTrain","last_synced_at":"2025-03-23T18:32:12.694Z","repository":{"id":37867786,"uuid":"433663107","full_name":"OpenBMB/BMTrain","owner":"OpenBMB","description":"Efficient Training (including pre-training and fine-tuning) for Big Models","archived":false,"fork":false,"pushed_at":"2024-07-22T03:28:12.000Z","size":2475,"stargazers_count":558,"open_issues_count":13,"forks_count":77,"subscribers_count":11,"default_branch":"main","last_synced_at":"2024-10-28T18:15:02.596Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"apache-2.0","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/OpenBMB.png","metadata":{"files":{"readme":"README-ZH.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2021-12-01T02:58:58.000Z","updated_at":"2024-10-27T05:28:21.000Z","dependencies_parsed_at":"2023-11-24T09:24:29.280Z","dependency_job_id":"3839c46b-26a0-49b0-b407-2cd17ab966ab","html_url":"https://github.com/OpenBMB/BMTrain","commit_stats":null,"previous_names":[],"tags_count":16,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/OpenBMB%2FBMTrain","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/OpenBMB%2FBMTrain/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/OpenBMB%2FBMTrain/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/OpenBMB%2FBMTrain/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/OpenBMB","download_url":"https://codeload.github.com/OpenBMB/BMTrain/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":245149529,"owners_count":20568917,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-07-31T07:00:51.189Z","updated_at":"2025-03-23T18:32:07.669Z","avatar_url":"https://github.com/OpenBMB.png","language":"Python","funding_links":[],"categories":["Python","LLM Training Frameworks","预训练模型","LLM训练框架","Training","Fine-tuning \u0026 Quantization (18)"],"sub_categories":["LLM 评估工具"],"readme":"\u003cdiv align=\"center\"\u003e\n\n\u003ch1\u003e\u003cimg src=\"docs/logo.png\" height=\"28px\" /\u003e BMTrain\u003c/h1\u003e\n\n**大模型高效训练工具包**\n\n\u003cp align=\"center\"\u003e\n  \u003ca href=\"#总览\"\u003e总览\u003c/a\u003e • \u003ca href=\"#文档\"\u003e文档\u003c/a\u003e • \u003ca href=\"#安装\"\u003e安装\u003c/a\u003e • \u003ca href=\"#使用说明\"\u003e使用说明\u003c/a\u003e • \u003ca href=\"#性能\"\u003e性能\u003c/a\u003e • \u003ca href=\"./README.md\" target=\"_blank\"\u003eEnglish\u003c/a\u003e\n\u003cbr\u003e\n\u003c/p\u003e\n\n\u003cp align=\"center\"\u003e\n\n\u003ca href='https://bmtrain.readthedocs.io/en/latest/?badge=latest'\u003e\n    \u003cimg src='https://readthedocs.org/projects/bmtrain/badge/?version=latest' alt='Documentation Status' /\u003e\n\u003c/a\u003e\n\n\u003ca href=\"https://github.com/OpenBMB/BMTrain/releases\"\u003e\n    \u003cimg alt=\"GitHub release (latest by date including pre-releases)\" src=\"https://img.shields.io/github/v/release/OpenBMB/BMTrain?include_prereleases\"\u003e\n\u003c/a\u003e\n\n\u003ca href=\"https://github.com/OpenBMB/BMTrain/blob/main/LICENSE\"\u003e\n    \u003cimg alt=\"GitHub\" src=\"https://img.shields.io/github/license/OpenBMB/BMTrain\"\u003e\n\u003c/a\u003e\n\n\u003c/p\u003e\n\n\u003c/div\u003e\n\n## 最新动态\n- 2022/06/14 **BMTrain** [0.1.7](https://github.com/OpenBMB/BMTrain/releases/tag/0.1.7) 发布。支持了ZeRO-2优化!\n- 2022/03/30 **BMTrain** [0.1.2](https://github.com/OpenBMB/BMTrain/releases/tag/0.1.2) 发布。适配了[OpenPrompt](https://github.com/thunlp/OpenPrompt)和 [OpenDelta](https://github.com/thunlp/OpenDelta)工具包。\n- 2022/03/16 **BMTrain** [0.1.1](https://github.com/OpenBMB/BMTrain/releases/tag/0.1.1) 公开发布了第一个稳定版本，修复了 beta 版本中的一些问题。\n- 2022/02/11 **BMTrain** [0.0.15](https://github.com/OpenBMB/BMTrain/releases/tag/0.0.15) 公开发布了第一个 beta 版本。\n\n\u003cdiv id=\"总览\"\u003e\u003c/div\u003e\n\n## 总览\n\nBMTrain 是一个高效的大模型训练工具包，可以用于训练数百亿参数的大模型。BMTrain 可以在分布式训练模型的同时，能够保持代码的简洁性。\n\n\u003cdiv id=\"文档\"\u003e\u003c/div\u003e\n\n## 文档\n我们的[文档](https://bmtrain.readthedocs.io/en/latest/index.html)提供了关于工具包的更多信息。\n\n\u003cdiv id=\"安装\"\u003e\u003c/div\u003e\n\n## 安装\n\n- 用 pip 安装（推荐）: ``pip install bmtrain``\n\n- 从源代码安装: 下载工具包，然后运行 ``pip install .`` (setup.py的安装方式将会在未来被setuptools弃用)\n\n安装 BMTrain 可能需要花费数分钟的时间，因为在安装时需要编译 c/cuda 源代码。\n我们推荐直接在训练环境中编译 BMTrain，以避免不同环境带来的潜在问题。\n\n\n\u003cdiv id=\"使用说明\"\u003e\u003c/div\u003e\n\n## 使用说明\n\n### 步骤 1: 启用 BMTrain\n\n首先，你需要在代码开头初始化 BMTrain。正如在使用 PyTorch 的分布式训练模块需要在代码开头使用 **init_process_group** 一样，使用 BMTrain 需要在代码开头使用 **init_distributed**。\n\n```python\nimport bmtrain as bmt\nbmt.init_distributed(\n    seed=0,\n    zero_level=3,   # 目前支持2和3\n    # ...\n)\n```\n\n**注意：** 使用 BMTrain 时请不要使用 PyTorch 自带的 `distributed` 模块，包括 `torch.distributed.init_process_group` 以及相关通信函数。\n\n### 步骤 2: 使用 ZeRO 优化\n\n使用ZeRO优化需要对模型代码进行简单替换：\n\n* `torch.nn.Module` -\u003e `bmtrain.DistributedModule`\n* `torch.nn.Parameter` -\u003e `bmtrain.DistributedParameter`\n\n并在 transformer 模块上使用 `bmtrain.CheckpointBlock`。\n\n下面是一个例子：\n\n**原始代码**\n\n```python\nimport torch\nclass MyModule(torch.nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.param = torch.nn.Parameter(torch.empty(1024))\n        self.module_list = torch.nn.ModuleList([\n            SomeTransformerBlock(),\n            SomeTransformerBlock(),\n            SomeTransformerBlock()\n        ])\n    \n    def forward(self):\n        x = self.param\n        for module in self.module_list:\n            x = module(x, 1, 2, 3)\n        return x\n\n```\n\n**替换后代码**\n\n```python\nimport torch\nimport bmtrain as bmt\nclass MyModule(bmt.DistributedModule): # 修改这里\n    def __init__(self):\n        super().__init__()\n        self.param = bmt.DistributedParameter(torch.empty(1024)) # 修改这里\n        self.module_list = torch.nn.ModuleList([\n            bmt.CheckpointBlock(SomeTransformerBlock()), # 修改这里\n            bmt.CheckpointBlock(SomeTransformerBlock()), # 修改这里\n            bmt.CheckpointBlock(SomeTransformerBlock())  # 修改这里\n        ])\n    \n    def forward(self):\n        x = self.param\n        for module in self.module_list:\n            x = module(x, 1, 2, 3)\n        return x\n    \n```\n\n### 步骤 3: 通信优化\n\n为了进一步缩短通信额外开销，将通信与运算时间重叠，可以使用 `TransformerBlockList` 来进一步优化。\n\n在使用时需要对代码进行简单替换：\n\n* `torch.nn.ModuleList` -\u003e `bmtrain.TransformerBlockList`\n* `for module in self.module_list: x = module(x, ...)` -\u003e `x = self.module_list(x, ...)`\n\n**原始代码**\n\n```python\nimport torch\nimport bmtrain as bmt\nclass MyModule(bmt.DistributedModule):\n    def __init__(self):\n        super().__init__()\n        self.param = bmt.DistributedParameter(torch.empty(1024))\n        self.module_list = torch.nn.ModuleList([\n            bmt.CheckpointBlock(SomeTransformerBlock()),\n            bmt.CheckpointBlock(SomeTransformerBlock()),\n            bmt.CheckpointBlock(SomeTransformerBlock())\n        ])\n    \n    def forward(self):\n        x = self.param\n        for module in self.module_list:\n            x = module(x, 1, 2, 3)\n        return x\n    \n```\n\n**替换后代码**\n\n```python\nimport torch\nimport bmtrain as bmt\nclass MyModule(bmt.DistributedModule):\n    def __init__(self):\n        super().__init__()\n        self.param = bmt.DistributedParameter(torch.empty(1024))\n        self.module_list = bmt.TransformerBlockList([ # 修改这里\n            bmt.CheckpointBlock(SomeTransformerBlock()),\n            bmt.CheckpointBlock(SomeTransformerBlock()),\n            bmt.CheckpointBlock(SomeTransformerBlock())\n        ])\n    \n    def forward(self):\n        x = self.param\n        x = self.module_list(x, 1, 2, 3) # 修改这里\n        return x\n    \n```\n\n### 步骤 4: 运行分布式训练代码\n\nBMTrain 使用 PyTorch 原生的分布式训练启动器，你可以根据 PyTorch 版本选择下列命令中的一个。\n\n* `${MASTER_ADDR}` 为主节点的 IP 地址\n* `${MASTER_PORT}` 为主节点的端口\n* `${NNODES}` 为节点数量\n* `${GPU_PER_NODE}` 为每个节点的 GPU 数量\n* `${NODE_RANK}` 为本节点的 rank\n\n#### torch.distributed.launch\n```shell\n$ python3 -m torch.distributed.launch --master_addr ${MASTER_ADDR} --master_port ${MASTER_PORT} --nproc_per_node ${GPU_PER_NODE} --nnodes ${NNODES} --node_rank ${NODE_RANK} train.py\n```\n\n#### torchrun\n\n```shell\n$ torchrun --nnodes=${NNODES} --nproc_per_node=${GPU_PER_NODE} --rdzv_id=1 --rdzv_backend=c10d --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} train.py\n```\n\n更多信息请参考 PyTorch [官方文档](https://pytorch.org/docs/stable/distributed.html#launch-utility)。\n\n## 样例\n\n我们提供了一个使用 BMTrain 训练 GPT-2 的[样例](https://github.com/OpenBMB/BMTrain/tree/main/example)。\n代码主要包含以下几个部分。\n\n### 第 1 部分: 模型定义\n\n```\n├── layers\n│   ├── attention.py\n│   ├── embedding.py\n│   ├── feedforward.py\n│   ├── __init__.py\n│   ├── layernorm.py\n│   └── linear.py\n└── models\n    ├── gpt.py\n    └── __init__.py\n```\n\n上面是代码的目录结构。\n\n我们定义了 GPT-2 需要的所有模型层，并使用 BMTrain 的 `DistributedModule` 和 `DistributedParameter` 来启用 ZeRO 优化。\n\n### 第 2 部分: 初始化 BMTrain\n\n```python\nbmtrain.init_distributed(seed=0)\n\nmodel = GPT(\n    num_layers=8,\n    vocab_size=10240, \n    dim_model=2560,\n    dim_head=80,\n    num_heads=32,\n    dim_ff=8192,\n    max_distance=1024,\n    bias=True,\n    dtype=torch.half\n)\n\nbmtrain.init_parameters(model) # 或者使用`bmtrain.load`加载checkpoint\n\n# ... 其他初始化（例如数据集） ...\n```\n\n`bmtrain.init_distributed(seed=0)` 用于初始化分布式训练环境，并设置随机数种子便于复现。\n\n`bmtrain.init_parameters(model)` 用于初始化模型的分布式参数。\n\n### 第 3 部分: 初始化优化器和学习率调整策略\n\n```python\nloss_func = torch.nn.CrossEntropyLoss(ignore_index=-100)\noptimizer = bmtrain.optim.AdamOffloadOptimizer(model.parameters(), weight_decay=1e-2)\nlr_scheduler = bmtrain.lr_scheduler.Noam(optimizer, start_lr=1e-3, warmup_iter=40, end_iter=1000, num_iter=0)\n```\n\nBMTrain 支持**所有** PyTorch 原生的优化器和损失函数，同时你也可以使用 BMTrain 提供的融合（fused）优化器用于混合精度训练。\n\n此外，在 `bmtrain.lr_scheduler` 中 BMTrain 也提供了常见的学习率调整策略。\n\n### 第 4 部分: 训练\n\n```python\n# 新建优化器管理器实例\noptim_manager = bmtrain.optim.OptimManager(loss_scale=1024)\n# 将所有的 optimzer 及(可选)其对应的 lr_scheduler 收入优化器管理器管理。\noptim_manager.add_optimizer(optimizer, lr_scheduler)\n# 可以再次调用 add_optimizer 加入其他优化器\n\nfor iteration in range(1000):\n    # ... 为每个rank加载数据 ...\n\n    # 前向传播并计算梯度\n    pos = torch.arange(enc_input.size(1)).long().cuda().repeat(enc_input.size(0), 1)\n    logits = model(\n        enc_input,\n        pos,\n        pos \u003c enc_length[:, None]\n    )\n    batch, seq_len, vocab_out_size = logits.size()\n\n    loss = loss_func(logits.view(batch * seq_len, vocab_out_size), targets.view(batch * seq_len))\n    \n    global_loss = bmtrain.sum_loss(loss).item() # 聚合所有rank上的损失, 仅用于输出训练日志\n\n    # 梯度清零\n    optim_manager.zero_grad() # 为每个 optimizer 调用 zero_grad\n\n    # 损失缩放和反向传播\n    optim_manager.backward(loss)\n\n    # 梯度裁剪\n    grad_norm = optim_manager.clip_grad_norm(optimizer.param_groups, max_norm=1.0)\n\n    # 更新参数\n    optim_manager.step()\n\n    # ... 保存checkpoint、打印日志 ...\n```\n\n这部分代码略有些长，但写起来就像常见的训练代码一样，你不需要为分布式训练调整太多的代码。\n\n你可以根据代码中的注释来了解各部分代码的作用。\n\n唯一需要说明的是 `optim_manager`。在使用 BMTrain 后，优化器的部分相关操作需要有一些细节上的调整。我们在 `optim_manager` 帮你实现了这些细节, 你只需要通过 `add_optimizer` 将优化器和学习率调整策略收入 `optim_manager` 管理，并由 `optim_manger` 代为执行 `zero_grad()`, `backward()`, `clip_grad_norm()` 和 `step()` 等操作。\n\n如果你没有使用混合精度训练，你可以不用损失缩放，只需要将 `OptimManger(loss_scale=None)` 构造函数中 `loss_scale` 置为 None 即可, 这也是 `OptimManager` 的默认构造参数。\n\n如果你使用了混合精度训练，**损失缩放**是混合精度训练中的一项常用技术，我们在 `optim_manager.backward(loss)` 帮你对 `loss` 进行了放缩，用于避免梯度下溢。只需要将 `OptimManger` 构造函数中 `loss_scale` 置为一个浮点数即可。 `loss_scale` 会在训练过程中根据梯度进行自适应的调整。\n\n\u003cdiv id=\"性能\"\u003e\u003c/div\u003e\n\n## 性能\n\n我们训练了一个有130亿参数的 GPT-2 模型，使用了4台服务器，每台服务器有8张V100显卡。我们测试了训练过程中每个GPU的吞吐量（每个GPU每秒处理的样本数），结果见下表。\n\n模型结构：\n* 40层\n* 128个注意力头\n* 5120的隐藏层维数\n* 512的序列长度\n\n\n| batch size  | 8     | 16    | 24    | 32    |\n|-------------|-------|-------|:------|:------|\n| BMTrain     | 24.15 | 26.94 | 29.42 | 28.28 |\n| ZeRO3(mp=1) | 14.88 | 21.69 | 24.38 | -     |\n| ZeRO3(mp=4) | 15.51 | -     | -     | -     |\n| ZeRO3(mp=8) | 15.51 | -     | -     | -     |\n| ZeRO2(mp=1) | -     | -     | -     | -     |\n| ZeRO2(mp=4) | 22.85 | -     | -     | -     |\n| ZeRO2(mp=8) | 21.33 | -     | -     | -     |\n\n**ZeROa(mp=b)** 表示 DeepSpeed + Megatron ZeRO stage a 和 model parallelism = b。\n\n表格中的 **-** 表示超出显存。\n\n## 模型支持\n\n我们已经将大多数常见的 NLP 模型移植到了 BMTrain 中。你可以在 [ModelCenter](https://github.com/OpenBMB/ModelCenter) 项目中找到支持模型的列表。\n\n## 开源社区\n欢迎贡献者参照我们的[贡献指南](https://github.com/OpenBMB/BMTrain/blob/master/CONTRIBUTING.md)贡献相关代码。\n\n您也可以在其他平台与我们沟通交流：\n- QQ群: 735930538\n- 官方网站: https://www.openbmb.org\n- 微博: http://weibo.cn/OpenBMB\n- Twitter: https://twitter.com/OpenBMB\n\n## 开源许可\n\n该工具包使用[Apache 2.0](https://github.com/OpenBMB/BMTrain/blob/main/LICENSE)开源许可证。\n\n## 其他说明\n\n`BMTrain` 工具包对 PyTorch 进行了底层修改，如果你的程序输出了意料之外的结果，可以在 issue 中提交相关信息。\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FOpenBMB%2FBMTrain","html_url":"https://awesome.ecosyste.ms/projects/github.com%2FOpenBMB%2FBMTrain","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FOpenBMB%2FBMTrain/lists"}