BMTrain: distributed training for models with tens of billions of parameters

A toolkit for models that do not fit on one device

BMTrain is described as an efficient large model training toolkit that can be used to train large models with tens of billions of parameters, and the sentence that explains the design goal is the last one: it can train models in a distributed manner while keeping the code as simple as stand-alone training.

That simplicity is relative, and the README is careful about where the effort lands. The toolkit does not hide the distributed execution from you; it changes which classes you use, so the model code looks like ordinary PyTorch with three substitutions. The heavy lifting sits in the C and CUDA layer, in the csrc directory, with CMake driving the native build and pybind11 bridging it to Python.

The project comes from OpenBMB, the group behind the BM series of models, and the repository structure is the open source half of that work. There is a bmtrain Python package, the csrc native sources, a cmake directory, docs, tests, and an example directory with benchmark.py, train.py, run.sh, sbatch.sh, a layers folder and a models folder.

The release history also tells you where the project is. The tags are v1.0.1 from 2025-05-13, v1.0.0 from 2024-02-26 and v0.2.3 from 2023-08-17, and the last push to the repository was on 2026-07-07. The 1.0.0 release is described as a code refactoring with tensor parallel support, and the what's new list goes back to a first public beta in February 2022.

init_distributed replaces init_process_group, and PyTorch collectives go away

The first step of using BMTrain is initialisation, and the README draws the analogy with PyTorch directly. Just as using the distributed module of PyTorch requires the use of init_process_group at the beginning of the code, using BMTrain requires the use of init_distributed at the beginning of the code.

The call looks like this:

import bmtrain as bmt
bmt.init_distributed(
    seed=0,
    # ...
)

The seed parameter is there because BMTrain owns the randomness setup for the distributed run, and the remaining arguments are left as an ellipsis in the documentation because they mirror the PyTorch equivalents.

Then comes the constraint that decides whether BMTrain fits your codebase: do not use PyTorch's distributed module and its associated communication functions when using BMTrain. That is not a style preference, it is a requirement, because BMTrain has its own communication layer and the two cannot be mixed.

The practical consequence is that any code you already have that calls torch.distributed directly, whether for gradient all-reduce, for broadcasting initial state, or for a custom collective, has to be rewritten. In a codebase where that logic is spread across the training loop, this is the expensive part of the migration, and the README does not discuss it, so it is worth auditing your own training script before you start.

ZeRO-2 through two class substitutions and a Block wrapper

Enabling ZeRO optimisation is, in the README's phrase, a matter of making some simple replacements to the original model's code. The first substitution is torch.nn.Module becoming bmtrain.DistributedModule, and the second is torch.nn.Parameter becoming bmtrain.DistributedParameter. The transformer blocks then get wrapped with bmtrain.Block.

The before and after of the README's example makes the shape of the change clear:

import torch
import bmtrain as bmt
class MyModule(bmt.DistributedModule): # changed here
    def __init__(self):
        super().__init__()
        self.param = bmt.DistributedParameter(torch.empty(1024)) # changed here
        self.module_list = torch.nn.ModuleList([
            bmt.Block(SomeTransformerBlock(), zero_level=3), # changed here, support 2 and 3 now
            bmt.Block(SomeTransformerBlock(), zero_level=3), # changed here, support 2 and 3 now
            bmt.Block(SomeTransformerBlock(), zero_level=3)  # changed here, support 2 and 3 now
        ])

The zero_level argument is the part that needs reading rather than copying. The comment in the example says it supports 2 and 3 now, and the release notes record that ZeRO-2 optimisation was added in version 0.1.7 in June 2022. So zero_level selects how much of the model is sharded, and the value you pick changes both the memory you save and the communication you pay for.

The forward method is untouched. In the README's example the loop over the module list stays exactly as it was, which is the payoff of the design: the ZeRO machinery lives in DistributedParameter and in the Block wrapper, not in your forward pass.

TransformerBlockList and overlapping communication with compute

The third step is optional and about speed rather than memory. To further reduce the extra overhead of communication and overlap communication with computing time, TransformerBlockList can be used, and enabling it is two more substitutions: torch.nn.ModuleList becomes bmtrain.TransformerBlockList, and the loop `for module in self.module_list: x = module(x, ...)` becomes `x = self.module_list(x, ...)`.

import torch
import bmtrain as bmt
class MyModule(bmt.DistributedModule):
    def __init__(self):
        super().__init__()
        self.param = bmt.DistributedParameter(torch.empty(1024))
        self.module_list = bmt.TransformerBlockList([ # changed here
            bmt.Block(SomeTransformerBlock()),
            bmt.Block(SomeTransformerBlock()),
            bmt.Block(SomeTransformerBlock())
        ])

The mechanism is a scheduler. Rather than a list of blocks that you iterate over one at a time, the list becomes an object that can interleave the collective communication of one layer with the computation of the next, so the network is not idle while gradients are being reduced. That is the standard trick for hiding communication latency, and BMTrain exposes it as a container swap rather than as something you have to write yourself.

Launching is the last piece, and it is the friendly part: BMTrain uses the same launch command as the distributed module of PyTorch, and you can choose one of them depending on your version of PyTorch. The environment variables are the familiar ones, with ${MASTER_ADDR} meaning the IP address of the master node and ${MASTER_PORT} meaning the port of the master node. So the operational part of a distributed launch, the rendezvous and the process count, does not need relearning.

The install compiles CUDA and C, so it takes minutes

Installation is two commands, and the README is honest that the second one is slow.

From pip, which is the recommended route, the command is `pip install bmtrain`. From source code you download the package and run `pip install .`.

The build step in the shipped container shows the source route in practice:

RUN python3 setup.py install

The reason for the delay is stated plainly: installing BMTrain may take a few to ten minutes, as it requires compiling the c/cuda source code at the time of installation. The recommendation is to compile BMTrain directly in the training environment to avoid potential problems caused by the different environments.

That recommendation is the important part. A wheel built somewhere else can be linked against a different CUDA and NCCL stack than the one your cluster has, and BMTrain has its own native layer sitting between your model and the GPUs, so a mismatch shows up as a runtime failure inside the training loop rather than as an import error.

The build requirements are declared in the project metadata: setuptools, pybind11, nvidia-nccl-cu11 at 2.14.3 or later, and cmake above 3.27.0. The setup.py file shows how those are used, with a CMakeBuild extension class that shells out to CMake, honours a CMAKE_GENERATOR environment variable so Conda-Build can override the generator, detects whether ninja is available, and passes the C++ standard, the library output directory and the Python executable through to the native build.

The shipped Dockerfile pins CUDA 10.2, torch 1.10 and AVX512

The Dockerfile in the repository is the fastest way to see what the project was developed against, and it is dated.

It starts FROM nvidia/cuda:10.2-devel and installs the build tools, then installs a pinned torch:

RUN pip3 install torch==1.10.0 -i https://pypi.tuna.tsinghua.edu.cn/simple
RUN pip3 install numpy -i https://pypi.tuna.tsinghua.edu.cn/simple

The index is a Tsinghua mirror, which tells you where the development happened. Then the environment is tuned for the hardware:

ENV TORCH_CUDA_ARCH_LIST=6.1;7.0;7.5
ENV BMT_AVX512=1

The first line restricts the compiled GPU architectures to 6.1, 7.0 and 7.5, which correspond to Pascal, Volta and Turing. The second turns on an AVX512 code path in BMTrain itself, so the CPU side of the training can use those instructions.

After that the image installs other_requirements.txt, copies the source in, runs python3 setup.py install and finally adds the example directory.

Read this as a development reference rather than a deployment recipe. CUDA 10.2 and torch 1.10 are several major versions behind what a current cluster runs, and the GPU architecture list excludes anything newer than Turing. If you want to reproduce a published result, matching this environment is the point. If you want to train on modern hardware, the file tells you exactly which knobs to change, and you should expect to change all three.

Apache-2.0, a quiet release cadence and a Chinese README

The licence is Apache-2.0, which is the permissive licence with an explicit patent grant, and the project is a good fit for commercial use on those terms. Your legal team will still want to look at the licence file at the root of the repository rather than take the summary from the metadata.

The maintenance picture needs reading carefully. The repository is not archived and the last push was on 2026-07-07, so the code is being worked on. The most recent release, v1.0.1, is from 2025-05-13, and the one before it, v1.0.0, is from 2024-02-26. So the commits move faster than the tags, which is the pattern of a project that maintains a branch and publishes when there is a reason to.

The what's new section in the README is the change log in miniature, and it shows the shape of the project over time: the first public beta in February 2022, the first stable version in March 2022, adaptation to OpenPrompt and OpenDelta in March 2022, ZeRO-2 in June 2022, 0.2.0 in December 2022, 0.2.3 in August 2023, and 1.0.0 in February 2024 with the code refactoring and tensor parallel support. Each of those entries links to a separate update log document under docs, which is where the detail lives.

The documentation itself is hosted on readthedocs, and the README links a Simplified Chinese version, which is a reasonable indicator of where the primary audience for the documentation sits. There is also a Release.txt file at the root, and separate requirement files for the documentation and for the other dependencies, so the environment is pinned in more than one place.

Editorial conclusion

Adopt BMTrain if you are training a model large enough that optimiser state, gradients or parameters do not fit on one device, and you want the ZeRO and communication overlap without writing the distributed layer yourself. Do not adopt it for a model that fits comfortably in memory, because the integration cost is real: your module classes change, PyTorch's own collective calls become off limits, and the install compiles native code. Verify first that your PyTorch version, CUDA version and the NCCL build the package pulls in are compatible, that the model code is structured so transformer blocks can be wrapped with bmtrain.Block, and that the pinned versions in the shipped Dockerfile, CUDA 10.2 and torch 1.10, are older than what your cluster actually runs.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. OpenBMB/BMTrain on GitHub
4. README
5. Releases