What Is XTuner V1
XTuner V1 is a new-generation large model training engine open-sourced by Shanghai Artificial Intelligence Laboratory. Developed based on PyTorch FSDP, it systematically optimizes memory, communication, and load issues in ultra-large model training, supports training of MoE models at the 1T parameter scale, and can achieve training throughput on models at the 200B scale that surpasses traditional 3D parallelism solutions. XTuner V1 collaborates with Huawei's Ascend technical team to further improve training efficiency and achieve higher model floating-point operation utilization. XTuner V1 provides academia and industry with a high-performance, low-barrier, and easily scalable large model training solution.
Main Features of XTuner V1
- Efficient Training of Ultra-Large Models: Supports training of MoE models at the 1T parameter scale. In mixture-of-experts models at the 200B+ scale, training throughput surpasses traditional 3D parallel training solutions.
- Optimized Memory Usage: Through the automatic Chunk Loss mechanism and Async Checkpointing Swap technology, it effectively reduces peak memory usage and enables training of 200B-parameter-scale MoE models with 64K-length sequences without relying on sequence parallelism technology.
- Masking Communication Overhead: Uses extreme memory optimization to increase the maximum sequence length per iteration and increase the computation time per layer to mask the communication time of parameter aggregation; through Intra-Node Domino-EP, it reduces the communication volume of parameter aggregation and masks the additional communication overhead brought by expert parallelism.
- Achieving DP Load Balancing: Sorts the already concatenated sequences within every n steps so that the maximum subsequence lengths of different DPs are close during each computation, reducing computational bubbles caused by variable-length attention.
- Adaptation to Multiple Hardware Platforms: In cooperation with Huawei's Ascend technical team, it performs deep optimization on Ascend A3 NPU supernodes and fully utilizes supernode hardware characteristics. Despite theoretical compute power being nearly 20% behind NVIDIA H800, it can achieve training throughput surpassing H800 by nearly 5% and MFU surpassing by more than 20%.
Technical Principles of XTuner V1
-
Developed Based on PyTorch FSDP: PyTorch FSDP (Fully Shard Data Parallel) is a data parallelism strategy that evenly partitions model parameters onto each GPU and saves memory by aggregating parameters in advance and re-sharding them. XTuner V1 performs deep optimization on the basis of FSDP to solve its problem of high communication volume.
-
Memory Optimization：
-
Automatic Chunk Loss Mechanism: For the computation graph when calculating the loss function, an automatic Chunk Loss mechanism supporting multiple training scenarios and multiple hardware platforms is developed to effectively reduce peak memory usage.
-
Async Checkpointing Swap: At the start of the model forward computation, the activation values that need to be saved for recomputation are moved from Device to Host to free up memory; during backpropagation, the activation values are moved back from Host to Device in advance, and memory is released after backpropagation ends, further optimizing memory usage.
-
-
Communication Masking：
-
Increasing Computation Time to Mask Communication Time: Through extreme memory optimization, the maximum sequence length per iteration is increased, increasing the computation time per layer so that computation time exceeds communication time, masking communication overhead and avoiding computational bubbles.
-
Intra-Node Domino-EP: For training scenarios limited by memory or communication bandwidth, Intra-Node Domino-EP reduces the communication volume of aggregating parameters at each layer, masking the additional communication overhead introduced by expert parallelism.
-
-
DP Load Balancing: During large model training, multiple sentences are concatenated to a fixed length, and a variable-length attention mechanism is used during computation. XTuner V1 sorts the already concatenated sequences within every n steps, so that the longest subsequence length of different DPs in each computation is close, reducing the 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 repository：https://github.com/InternLM/xtuner
Application scenarios of XTuner V1
-
Natural Language Processing (NLP) field: Used to train extremely large-scale language models, such as GPT, BERT, etc., to improve the model's language understanding and generation capabilities, and applied to tasks such as machine translation, text generation, and sentiment analysis.
-
Computer Vision (CV) field: In tasks such as image recognition and object detection, train large-scale vision models, such as ResNet, Transformer, etc., to improve the model's accuracy and efficiency.
-
Multimodal learning: Combine language and visual information to train multimodal models, such as CLIP, etc., used in tasks such as image caption generation and visual question answering, to improve the model's ability to understand complex scenes.
-
Reinforcement learning: In long-sequence reinforcement learning tasks, such as robot control, game AI, etc., XTuner V1 can efficiently process long-sequence data, accelerate model training, and improve the efficiency of policy learning.