XTuner V1 - A new generation of large model training engine open sourced by Shanghai AI Lab

What is XTuner V1?

XTuner V1 is a new-generation large-model training engine open-sourced by the Shanghai Artificial Intelligence Laboratory. Developed based on PyTorch FSDP, it systematically optimizes memory, communication, and load issues in training ultra-large models. It supports training MoE models with 1T parameters and can achieve training throughput exceeding traditional 3D parallel solutions on models with a scale of 200B. XTuner V1 collaborates with Huawei Ascend technology team to further improve training efficiency and achieve higher model floating-point operation utilization. XTuner V1 provides a high-performance, low-barrier, and easily scalable large-model training solution for academia and industry.

Main functions of XTuner V1

• Efficient training of ultra-large modelsIt supports training MoE models with 1T parameters and achieves training throughput that surpasses traditional 3D parallel training schemes in hybrid expert models with more than 200B parameters.
• Optimize video memory usageBy using an automatic chunk loss mechanism and Async Checkpointing Swap technology, the peak memory usage is effectively reduced, and a 64K-length sequence can be trained on a 200B-level MoE model without relying on sequence parallel technology.
• Covering up communication overheadThe maximum sequence length of a single iteration is increased by using extreme memory optimization, which increases the computation time per layer to mask the communication time of parameter aggregation; the communication volume of parameter aggregation is reduced by Intra-Node Domino-EP to mask the additional communication overhead brought by expert parallelism.
• Achieve DP load balancingSort the concatenated sequences within each n steps to make the longest subsequence length of different DPs similar in each calculation, reducing computational bubbles caused by variable length attention.
• Compatible with various hardwareIn collaboration with Huawei's Ascend technology team, we conducted in-depth optimization on the Ascend A3 NPU supernode, making full use of the supernode's hardware features. Despite theoretically lagging behind the NVIDIA H800 by nearly 20% in computing power, we were able to achieve training throughput that surpassed the H800 by nearly 5% and MFU that surpassed it by more than 20%.

The technical principles of XTuner V1

• Developed based on PyTorch FSDPPyTorch FSDP (Fully Shard Data Parallel) is a data parallelism strategy that evenly distributes model parameters across each GPU, saving GPU memory by pre-aggregating and re-splitting parameters. XTuner V1 deeply optimizes FSDP to address its high communication volume issue.
• Video memory optimization:
  • Automatic Chunk Loss MechanismFor the computation graph used in calculating the loss function, an automatic Chunk Loss mechanism was developed that supports multiple training scenarios and various hardware, effectively reducing peak memory usage.
  • Async Checkpointing SwapAt the start of forward computation, the activation values that need to be saved during recomputation are moved from Device to Host, and the GPU memory is released. During backpropagation, the activation values are moved back from Host to Device in advance, and the GPU memory is released after backpropagation is completed, further optimizing GPU memory usage.
• Communication cover:
  • Increasing computation time masks communication timeBy optimizing the memory to the extreme, the maximum sequence length of a single iteration is increased, and the computation time of each layer is increased, so that the computation time is greater than the communication time, thus masking the communication overhead and avoiding computational bubble.
  • Intra-Node Domino-EPFor training scenarios where GPU memory or communication bandwidth is limited, Intra-Node Domino-EP reduces the communication volume of each layer's aggregation parameters, thus masking the additional communication overhead caused by introducing expert parallelism.
• DP load balancingDuring large model training, multiple sentences are concatenated into a fixed length, and a variable-length attention mechanism is used for computation. XTuner V1 sorts the concatenated sequences within every n steps, making the longest subsequence length of different dynamic programming (DP) operations similar in each computation, reducing computational bubbles caused by variable-length attention and improving training efficiency.

XTuner V1 project address

• Project official website: https://xtuner.readthedocs.io/zh-cn/latest/
• GitHub repositoryhttps://github.com/InternLM/xtuner

Application scenarios of XTuner V1

• Natural Language Processing (NLP)It is used to train ultra-large-scale language models, such as GPT and BERT, to improve the model's language understanding and generation capabilities, and is applied to tasks such as machine translation, text generation, and sentiment analysis.
• Computer vision (CV) fieldIn tasks such as image recognition and object detection, large-scale visual models, such as ResNet and Transformer, are trained to improve the accuracy and efficiency of the models.
• Multimodal learningBy combining language and visual information, multimodal models, such as CLIP, can be trained and used in tasks such as image description generation and visual question answering to improve the model's ability to understand complex scenes.
• reinforcement learningIn long-sequence reinforcement learning tasks, such as robot control and game AI, XTuner V1 can efficiently process long-sequence data, accelerate model training, and improve the efficiency of policy learning.