BMTrain

BMTrain is an efficient large model training toolkit launched by OpenBMB, used for training large models with parameter counts reaching tens of billions, supporting pretraining and fine-tuning. Its core goal is to keep distributed training code nearly as concise as single-machine training.

Purpose BMTrain is aimed at training large models (which may contain tens of billions of parameters), covering pretraining and fine-tuning scenarios, and completing training in a distributed manner.

Main capabilities

• Perform distributed training with a code style close to single-machine training.
• Supports ZeRO optimization: replace torch.nn.Module with bmtrain.DistributedModule, replace torch.nn.Parameter with bmtrain.DistributedParameter, and wrap the transformer module with bmtrain.Block.
• Supports communication optimization: use bmtrain.TransformerBlockList instead of torch.nn.ModuleList, and adjust the module invocation method to reduce communication overhead and overlap communication with computation time.
• Supports all native PyTorch optimizers and loss functions, and provides BMTrain's fused optimizer for mixed precision training.
• Provides common LRScheduler in the bmtrain.lr_scheduler module.
• Provides OptimManager to uniformly manage the optimizer, which executes zero_grad(), backward(), clip_grad_norm(), and step(). During mixed precision training, loss scale can be used to prevent gradient underflow; loss_scale adaptively adjusts according to gradients during training; when mixed precision is not used, it can be set to None (default).
• Supports tensor parallelism (introduced in version 1.0.0), and the code has been refactored.

Usage Installation: pip install bmtrain is recommended; you can also download the source code and run pip install . Installation requires compiling c/cuda source code, which may take from a few minutes to more than ten minutes. It is recommended to compile directly in the training environment to avoid problems caused by environment differences. Usage steps:

1. Initialization: call init_distributed at the beginning of the code, similar to using init_process_group at the beginning of PyTorch's distributed module. Note thatWhen using BMTrain, do not use PyTorch's distributed modules or their related communication functions.
2. Enable ZeRO optimization: perform the above module replacements and wrap the transformer module with bmtrain.Block.
3. Enable communication optimization: use TransformerBlockList and adjust module calls.
4. Start distributed training: use the same launch command as PyTorch's distributed module, where the command involves MASTER_ADDR (master node IP), MASTER_PORT (master node port), NNODES (total number of nodes), GPU_PER_NODE (number of GPUs per node), NODE_RANK (current node rank).

Examples and Training Workflow The official example for training GPT-2 with BMTrain mainly includes: model definition (using DistributedModule and DistributedParameter to enable ZeRO optimization), BMTrain initialization (init_distributed sets the random seed, init_parameters initializes distributed parameters), optimizer and learning rate scheduler initialization, and the training loop (the workflow is similar to an ordinary training loop, with the additional management of optimizers via optim_manager's add_optimizer).

Performance and Supported Models Officially, a 13B-parameter GPT-2 model was trained on 4 servers with 8 V100s each, and per-GPU throughput (samples per GPU per second) was measured, comparing DeepSpeed + Megatron's ZeRO stages and model parallelism configurations. The GPT-2 model structure is: 128 attention heads, 5120 hidden dimensions, 512 sequence length. BMTrain has already migrated and supported most common models in NLP; the list is in the ModelCenter repository.

Applicable Scenarios Research and engineering scenarios that require large-model distributed pretraining or fine-tuning with relatively low code modification cost; training workflows that already use PyTorch and wish to enable ZeRO and communication optimization.

Limitations and Notes

• When using BMTrain, PyTorch distributed modules and their communication functions should not be used at the same time.
• BMTrain modifies PyTorch's underlying implementation; if the program output is abnormal, information can be submitted via an issue.
• Installation requires compiling c/cuda source code, which is time-consuming and is recommended to be done in the training environment.
• Specific performance data and comparison details are subject to the official documentation and repository.

License and Community This package is released under the Apache 2.0 license. You can contact us via QQ group 735930538 or the website https://www.openbmb.cn、微博 http://weibo.cn/OpenBMB、Twitter https://twitter.com/OpenBMB .