MT-TransformerEngine 入门

MT-TransformerEngine 简介

MT-TransformerEngine 是由摩尔线程 AI-Infra 团队开发的高性能深度学习框架，基于 TransformerEngine 和 torch_musa 构建，为摩尔线程 GPU 提供优化的 FP8 训练支持。与 MT-Megatron 集成后，可以在摩尔线程 GPU 上使用 FP8 训练方案，并支持跨大规模集群的大模型训练。

环境准备

建议使用摩尔线程官方训练镜像。当前示例基于 MTT S5000、MUSA Driver 4.3.7 和 training-suite:v2.1.5-musa-4.3.7。

创建并启动容器

sudo docker create --privileged --env MTHREADS_VISIBLE_DEVICES=all --net host -v /data/:/data/ --name te-train registry.mthreads.com/mcctest/training-suite:v2.1.5-musa-4.3.7 sleep infinity
sudo docker start te-train
sudo docker exec -it te-train bash
service ssh restart

进入容器后使用 mthreads-gmi 检查 GPU，必要时执行 ./mccl_test 验证通信。

安装 MT-TransformerEngine

通过提供的安装脚本编译 MUSA 内核和 C++ 源文件：

bash install.sh

该脚本会编译以下目录中的源文件：

transformer_engine/musa/common
transformer_engine/musa/pytorch/csrc