第10章精讲：应用实战（图像处理 / 科学计算 / AI 推理 / 多卡扩展）

一、PyTorch 编程：torch_musa 与 MUSAExtension

PyTorch 是当下被广泛使用的深度学习框架，由 Meta AI 团队开发并开源，以动态计算图（Dynamic Computational Graph）机制著称，支持传统机器学习、深度学习模型训练推理以及科学计算等任务，并配备 TorchScript、TorchDynamo、TorchInductor 等编译加速组件。从学术研究到工业部署，PyTorch 都是最主流的选择之一，逐渐形成了"PyTorch 生态"的概念。

1.1 torch_musa：让 PyTorch 模型无缝迁移到 MUSA

MUSA 作为面向通用计算和 AI 推理训练的统一 GPU 架构，为大模型训练、推理和图像处理等任务提供算力支持。为便于开发者高效使用 MUSA 芯片进行模型开发与部署，摩尔线程开发了 torch_musa——作为 PyTorch 的后端，支持几乎所有 PyTorch 功能。

torch_musa 是基于 PyTorch 的后端扩展模块，旨在将基于 PyTorch 的模型无缝迁移至 MUSA 平台运行。其底层对接 MUSA 驱动和计算库，提供张量运算、自动求导、调度执行等核心能力的支持。用户通过 to("musa") 调用即可将模型迁移到 MUSA GPU 上运行，无须额外修改模型逻辑。torch_musa 支持 PyTorch 的主流张量操作、模块接口、前向与反向计算，并通过持续适配社区主线版本保持与 PyTorch 最新语义的兼容性。在性能方面，torch_musa 结合 MUSA 芯片的异构并行调度机制与高带宽访存架构，已在 CNN、Transformer 等模型中实现与主流平台相当的推理与训练效率。

import torch
import torch.nn as nn
device = torch.device("musa")  # 指定使用 MUSA 设备
model = nn.Linear(512, 256).to(device)
input = torch.randn(8, 512, device=device)
output = model(input)

1.2 MUSAExtension：自定义 MUSA 算子的构建工具

在 PyTorch 框架中，为了提升性能或扩展功能，用户常常需要编写自定义 C++/MUSA 扩展模块。在 CUDA 平台上，torch.utils.cpp_extension.CUDAExtension 已成为事实标准，其封装了复杂的构建流程；为在 MUSA 平台提供一致的开发体验，MUSA 软件栈提供了功能对等的工具模块：MUSAExtension。它是为 PyTorch 生态中自定义算子开发量身定制的构建工具，接口风格与 PyTorch 社区版本完全一致，允许用户通过 setup.py 脚本以模块化方式构建包含 MUSA 后端的 C++ 算子，并自动完成编译器配置、依赖路径设置与链接库注入，开发者无须手动配置编译流程。

from setuptools import setup
from torch_musa.utils.musa_extension import (
    MUSAExtension,
    BuildExtension,
)
setup(
    name='custom_musa_op',
    ext_modules=[
        MUSAExtension(
            name='custom_musa_op',
            sources=[
                'op.cpp',
                'op_kernel.mu'],
        )
    ],
    cmdclass={
        'build_ext': BuildExtension
    }
)

MUSAExtension 使用 g++ / gcc 编译 .cpp 文件，使用 MUSA 提供的 mcc 编译 .mu 文件。若要传入编译选项，可添加 extra_compile_args 字典，通过配置 "cxx" 和 "mcc" 的键值对来控制：

from setuptools import setup
from torch_musa.utils.musa_extension import (
    MUSAExtension,
    BuildExtension,
)
c_flags = {
    "cxx": ["-O3"],  # 增加 O3 优化
    "mcc": ["-O3"],  # 增加 O3 优化
}

setup(
    name='custom_musa_op',
    ext_modules=[
        MUSAExtension(
            name='custom_musa_op',
            sources=[
                'op.cpp',
                'op_kernel.mu'],
            extra_compile_args=c_flags,  # 传入编译选项
        )
    ],
    cmdclass={
        'build_ext': BuildExtension
    }
)

Step-by-Step：验证 torch_musa 迁移与最小算子构建

1. 在装有 torch_musa 的环境中运行以下代码，验证模型可迁移到 MUSA 设备：

   import torch
   import torch.nn as nn
   device = torch.device("musa")
   model = nn.Linear(512, 256).to(device)
   input = torch.randn(8, 512, device=device)
   output = model(input)
   print(output.shape, output.device)

2. 参照 MUSAExtension 示例，编写最小可运行的 setup.py，包含一个 op.cpp 与 op_kernel.mu（可先用占位实现），执行构建：

   python setup.py build_ext --inplace

3. 尝试在 extra_compile_args 中添加 -O3 优化选项，重新构建，并比较（若有可测的算子）执行时间差异。

1.3 本节要点小结

二、CIFAR-10 图像分类实战（AutoDL 全流程）

图像分类是计算机视觉的基本任务，其技术进步推动着自动驾驶、医学影像诊断、安防监控等领域的发展，是验证框架适配与 GPU 加速效果的理想场景。CIFAR-10 数据集是图像分类任务的经典基准，包含 10 个类别的彩色图像，是深度学习入门与实践的理想选择。本节以 CIFAR-10 为例，演示从环境搭建、模型训练到推理优化的全流程，把 torch_musa 的能力转化为实际项目成果，涵盖 CNN、交叉熵损失函数、Adam 优化器、数据增强、PyTorch 数据管道、GPU 加速、模型评估指标、训练与验证等知识点。

本节基于 AutoDL 平台完成图像分类任务，涵盖六大步骤：源代码与数据集准备、环境搭建、模型设计与训练、训练过程监控与结果分析、模型推理与性能测试、超参数调优。通过 SimpleCNN 与 ResNet（残差网络）两种模型的对比训练，既能了解 CNN 的基本原理，又能体会深层网络的性能优势；MUSA GPU 的加速效果会通过训练时间、吞吐量等指标直观呈现。

2.1 步骤一：源代码与数据集准备

教材所需源代码已集成在 AutoDL 平台的专用镜像中，保存在社区镜像里，无须手动下载。镜像中包含以下核心文件：

CIFAR-10 数据集包含 50000 张训练图像和 10000 张测试图像，共 10 个类别（飞机、汽车、鸟、猫、鹿、狗、青蛙、马、船、卡车），图像均为 32×32 像素 RGB 彩色图像，首次运行训练脚本时会自动下载，存储在 ./data 目录下。数据预处理代码如下：

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                         download=True, transform=transform)

该代码将图像转换为张量并做标准化处理，使像素值分布在 [-1, 1]，有助于提升训练稳定性和收敛速度。

Step-by-Step：AutoDL 环境搭建与冒烟检查

1. 访问 AutoDL 官网，注册账户并登录。
2. 进入控制台，选择"容器实例"，单击"租用新实例"按钮。
3. 在实例配置页面依次选择"按量计费" → "摩尔线程专区" → "004 机"（GPU 型号示例：MTT S4000）。
4. 在"镜像"区域选择"社区镜像"，搜索 "cifar10"，选择对应的实训镜像，单击"创建并开机"启动实例。
5. 单击"JupyterLab"进入开发环境，再单击"终端"按钮新建终端（若已直接进入新终端则跳过该步）。
6. 正式训练之前，检查设备是否支持 MUSA：

   cd moore_lesson
   python check_musa.py

若环境配置正确，将输出摩尔线程全功能 GPU 的设备名称和显存大小等信息；若未检测到 GPU，则会显示 CPU 信息。教材示例输出：

================================================
MUSA环境检查
================================================
PyTorch版本: 2.2.0
torch_musa已安装
torch_musa版本: 1.3.0+81caf0a
MUSA设备可用
MUSA设备数量: 1
设备 0: MTT S4000
  显存: 47.9 GB

测试基本MUSA操作...
MUSA张量运算正常
测试张量形状: torch.Size([100, 100])
张量设备: musa:0

性能测试
------------------------------------------------
CPU矩阵乘法测试...
CPU时间: 0.019秒
MUSA矩阵乘法测试...
MUSA时间: 0.003秒
加速比: 7.5x
================================================

check_musa.py 的核心检测逻辑如下——该段代码会自动检测系统是否支持 MUSA 加速，若检测到 GPU 会显示设备名称和显存大小：

def get_device_info():
    """检测可用的计算设备"""
    if torch.musa.is_available():
        device = torch.device('musa')
        device_name = torch.musa.get_device_name(0)
        memory = torch.musa.get_device_properties(0).total_memory // 1024**2
        return f"GPU: {device_name} ({memory}MB)"
    else:
        return "CPU: " + platform.processor()

由于 AutoDL 平台虚拟机默认未安装中文字体，为正确显示中文图表需先安装字体，再验证是否成功：

bash setup_fonts.sh
python font_config.py

2.2 步骤三：模型设计与训练

model.py 构建了 SimpleCNN 和 ResNet 两个模型。SimpleCNN 包含卷积层、池化层和全连接层，适合作为入门模型：

class SimpleCNN(nn.Module):
    """SimpleCNN模型，适用于CIFAR-10图像分类"""

    def __init__(self, num_classes=10):
        super(SimpleCNN, self).__init__()

        # 卷积层
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)

        # 池化层
        self.pool = nn.MaxPool2d(2, 2)

        # 全连接层
        self.fc1 = nn.Linear(128 * 4 * 4, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, num_classes)

        # Dropout层
        self.dropout = nn.Dropout(0.5)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.pool(F.relu(self.conv3(x)))
        x = x.view(-1, 128 * 4 * 4)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.fc3(x)
        return x

python train_cifar10.py --epochs 5 --batch-size 64

ResNet 通过残差连接解决了深层网络训练困难的问题，能学习更复杂的特征表示。本节实现简化版 ResNet：

class ResNetBlock(nn.Module):
    """ResNet 基本块"""

    def __init__(self, in_channels, out_channels, stride=1):
        super(ResNetBlock, self).__init__()

        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3,
                                stride=stride, padding=1, bias=False)
        self.bn1 = nn.BatchNorm2d(out_channels)

        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3,
                                stride=1, padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(out_channels)

        # 残差连接
        self.shortcut = nn.Sequential()
        if stride != 1 or in_channels != out_channels:
            self.shortcut = nn.Sequential(
                nn.Conv2d(in_channels, out_channels, kernel_size=1,
                          stride=stride, bias=False),
                nn.BatchNorm2d(out_channels)
            )

    def forward(self, x):
        out = F.relu(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        out += self.shortcut(x)
        out = F.relu(out)
        return out

python train_cifar10.py --model simple_resnet --epochs 10

2.3 步骤四：训练过程监控与结果分析

训练过程中会实时显示每个 epoch 的损失值变化、训练和验证准确率、GPU 利用率和内存占用、预计剩余时间。以 SimpleCNN 训练 5 轮为例，教材记录的关键日志：

训练完成!
总耗时: 36s
最佳验证准确率: 70.38%
平均每epoch耗时: 7s

正在生成训练曲线...
训练曲线已保存为 training_curves.png
训练脚本执行完成!

训练完成后会自动生成 training_curves.png（训练损失和准确率曲线），模型参数保存在 checkpoints/ 目录下：训练损失曲线随 Epoch 增加而下降（训练/验证损失均下降），准确率曲线随 Epoch 增加而上升（训练/验证准确率均上升）。

2.4 步骤五：模型推理与性能测试

# 如果训练用的是 SimpleCNN 模型
python inference.py --model simple_cnn

# 如果训练用的是 ResNet 模型
python inference.py --model simple_resnet

以 SimpleCNN 为例执行 python inference.py --model simple_cnn，教材示例日志片段：

批处理大小: 64
批处理时间: 8.7018秒
平均每张图像: 11.30毫秒
批次准确率: 71.88% (46/64)

性能指标主要包括推理速度（单张图片处理时间）、批处理性能（不同批次大小下的吞吐量）、内存占用（峰值 GPU 内存使用量）、准确率分析（各类别预测精度对比）。执行性能基准测试：

python inference.py --model simple_cnn --benchmark --visualize

推理结果可视化（inference_results.png）会展示测试样本的"真实类别 / 预测类别 / 置信度"，例如猫→猫 55.7%、船→船 93.8%、飞机→船 53.1%（预测错误示例）、青蛙→青蛙 94.8% 等。

2.5 步骤六：超参数调优

超参数是训练前需要人为设定的参数：学习率（--lr）过大可能导致训练不收敛，过小则收敛速度慢甚至陷入局部最优；批次大小（--batch-size）大批次更稳定、速度快但显存占用高，小批次有助于跳出局部最优但训练波动大；训练轮数（--epochs）过少会欠拟合，过多则可能过拟合。

# 默认参数训练
python train_cifar10.py

# 调大学习率和批次大小，减少训练轮数
python train_cifar10.py --lr 0.01 --batch-size 256 --epochs 10

# 降低学习率，增加训练轮数
python train_cifar10.py --lr 0.0005 --batch-size 64 --epochs 30

Step-by-Step：CIFAR-10 全流程复现实验

1. 按上文步骤完成 AutoDL 账户注册、实例租用（"按量计费" → "摩尔线程专区"）、社区镜像（搜索 "cifar10"）启动。
2. 进入 JupyterLab 终端，执行 python check_musa.py 确认 MUSA 设备可用。
3. 执行 bash setup_fonts.sh 与 python font_config.py 完成中文字体配置。
4. 分别训练 SimpleCNN 与 ResNet 模型：

   python train_cifar10.py --epochs 5 --batch-size 64
   python train_cifar10.py --model simple_resnet --epochs 10

5. 查看 training_curves.png，对比两模型的训练损失曲线与准确率曲线。
6. 分别执行推理与基准测试：

   python inference.py --model simple_cnn
   python inference.py --model simple_cnn --benchmark --visualize

7. 调整超参数，至少完成一组"调大学习率+批大小、减少 epoch"和一组"降低学习率、增加 epoch"的对比实验，记录最终验证准确率差异。

产出物：SimpleCNN 与 ResNet 的准确率/训练时间对比表；不同超参数配置下的收敛曲线对比。

2.6 本节要点小结

三、大模型训练实战：MT-Megatron

大模型训练对算力、通信效率、内存管理提出了极高要求，是检验 MUSA 生态完整性的核心场景。本节基于摩尔线程开源的 MT-Megatron 框架，演示大模型训练的全流程。

3.1 环境搭建

摩尔线程提供的官方 Docker 镜像已集成 MT-Megatron、MT-TransformerEngine 的新版本以及 MT-Deepspeed、Simmumax 等核心工具，无须手动配置复杂依赖。执行安装脚本即可完成环境部署：

cd TransformerEngine/
chmod a+x ./install.sh
./install.sh

MT-Megatron 将在训练期间自动安装。

3.2 数据准备：词元化与二进制序列化

大模型训练的核心流程包括数据准备、训练脚本编写、参数配置与模型拉起。数据准备包括数据收集、数据清洗、数据词元化、数据二进制序列化，其中前两步因训练目标不同而各有差异，本节着重讲后两步。

数据词元化（Tokenization）是训练的前置核心步骤，MT-Megatron 提供了便捷的构建工具：

class args:
    # 指定分词器模型路径和类型
    tokenizer_model = './dataset/llama3_tokenizer/'
    tokenizer_type = 'HuggingFaceTokenizer'
    # 设置进程秩（Rank），默认从 0 开始
    rank = 0
    # 词表大小对齐系数，用于优化张量并行（TP）下的计算性能
    make_vocab_size_divisible_by = 128  # vocab_size_padding, 用于 TP 并行
    # 张量并行度设置
    tensor_model_parallel_size = 1
    vocab_extra_ids = 0

tokenizer = build_tokenizer(args)

通过调用 build_tokenizer 工具，可以实现对分词器的自动化封装，将原始文本转换为模型可理解的词元：

with open(folder_path + file, 'r', encoding='gbk', errors='ignore') as doc:
    sentences = doc.readlines()
doc_ids = []
sentence_lens = []
for sentence in sentences:
    sentence_ids = tokenizer.tokenize(sentence)
    if len(sentence_ids) > 0:
        doc_ids.extend(sentence_ids)
        sentence_lens.append(len(sentence_ids))
if len(doc_ids) > 0:
    doc_ids.append(tokenizer.eod)
    sentence_lens[-1] += 1

词元化之后需要将数据转换为二进制格式（bin）以便后续训练，教材提供了多进程处理示例：

output_bin_files = {}
output_idx_files = {}
builders = {}
output_bin_files[file] = "{}_{}_{}.bin".format(output_prefix, file, level)
output_idx_files[file] = "{}_{}_{}.idx".format(output_prefix, file, level)
builders[file] = indexed_dataset.IndexedDatasetBuilder(
    output_bin_files[file],
    dtype=indexed_dataset.DType.optimal_dtype(tokenizer.vocab_size),
)
builders[file].add_document(doc_ids, sentence_lens)
builders[file].finalize(output_idx_files[file])
return np.sum(sentence_lens)


import os
import multiprocessing
def process_files_in_folder(folder_path='./dataset/raw_data//'):
    # 获取文件夹中的所有文件
    files = os.listdir(folder_path)
    # 过滤出 .txt 文件
    txt_files = [f for f in files if f.endswith('.txt')]
    files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith('.txt')]
    # 使用多进程池处理文件
    with multiprocessing.Pool(processes=multiprocessing.cpu_count()) as pool:
        results = pool.map(preprocess, txt_files)
    # 汇总结果
    total_tokens = np.sum(results)
    pwd = os.getcwd()
    output = []
    for i, file in enumerate(txt_files):
        append_str = "{} {}/dataset/data_{}_document".format(results[i], pwd, file)
        output.append(append_str)
    return output, total_tokens

3.3 训练脚本与入口函数

参照 MT-Megatron 提供的示例改写训练脚本，指定数据路径的方式（每行为"该文档 token 数量 + 文档路径"，多个数据文件按行拼接）：

DATA_PATH="
234215 /data/lecture/megatron-lm-musa-patch/examples/lecture/dataset/data_10.txt_document
74019  /data/lecture/megatron-lm-musa-patch/examples/lecture/dataset/data_11.txt_document
209570 /data/lecture/megatron-lm-musa-patch/examples/lecture/dataset/data_9.txt_document"

MT-Megatron 默认使用 torchrun 作为训练入口函数，利用 torch 配置分布式训练所需的环境、建立通信组：

torchrun --nproc_per_node=4 train.py

torchrun 能自动进行环境配置，包含启动/守护进程、创建通信 ID、创建通信 rank、容错等能力。

3.4 模型结构与 MoE 参数配置

MT-Megatron 用户可以仅通过修改配置进行高效大模型训练，完整配置参数参见 arguments.py 文件。Transformer/MoE 模型训练流程为 Embedding → 多层 Transformer Block（含 Self Attention 多头机制）→ LM Head → Cross Entropy，MoE 场景下还包含 Router 与多个并行 MLP 专家模块。

模型结构参数（MODEL_ARGS）示例：

MODEL_ARGS=(
    --num-layers $NUM_LAYERS
    --hidden-size $HIDDEN_SIZE
    --num-attention-heads 128
    --seq-length $SEQ_LENGTH
    --max-position-embeddings $MAX_POSITION_EMBEDDINGS
    --norm-epsilon 1e-6
    --attention-dropout 0.0
    --hidden-dropout 0.0
    --disable-bias-linear
    --vocab-size $VOCAB_SIZE
    --ffn-hidden-size $FFN_HIDDEN_SIZE
    --position-embedding-type rope
    --no-position-embedding
    --swiglu
    --normalization RMSNorm
    --untie-embeddings-and-output-weights
)

其中 --num-attention-heads 128 指定注意力头数、--position-embedding-type rope 使用旋转位置编码、--swiglu 使用 SwiGLU 激活函数代替 ReLU/GELU、--normalization RMSNorm 使用 RMSNorm 归一化以加快训练、--disable-bias-linear 禁用线性层偏置项以减少参数量。

MoE 参数（MOE_ARGS）常用配置：

MOE_ARGS=(
    --num-experts $NUM_EXPERTS
    --expert-model-parallel-size $EP_SIZE
    --moe-token-dispatcher-type alltoall
    --moe-router-score-function softmax
    --moe-router-num-groups $EP_SIZE
    --moe-router-group-topk $MOE_ROUTER_GROUP_TOPK
    --moe-router-load-balancing-type seq_aux_loss
    --moe-router-topk $MOE_TOP_K
    --moe-router-pre-softmax
    --moe-router-topk-scaling-factor $MOE_ROUTER_TOPK_SCALING_FACTOR
    --moe-aux-loss-coeff 3e-3
    --moe-expert-capacity-factor $MOE_EXPERT_CAPACITY_FACTOR
    --moe-device-level-capacity
    --moe-device-level-aux-loss-coeff 5e-2
    --moe-comm-aux-loss-coeff 2e-2
    --moe-ffn-hidden-size $MOE_FFN_HIDDEN_SIZE
    --moe-shared-expert-intermediate-size $MOE_SHARED_EXPERT_INTERMEDIATE_SIZE
    --moe-layer-freq "$MOE_LAYER_FREQ"
    --moe-grouped-gemm
    --moe-permute-fusion
)

--moe-router-topk 指定每个 token 分配给 top-k 个专家（通常为 1 或 2）；--moe-aux-loss-coeff 是 token 级负载均衡辅助损失权重，防止 token 全部涌向少数专家；--moe-grouped-gemm 启用 Grouped GEMM 优化，将多个小矩阵乘法合并以提升计算效率。

训练相关参数（TRAINING_ARGS）涉及并行策略与训练技术：

TRAINING_ARGS=(
    --seed 42
    --micro-batch-size $MICRO_BS
    --global-batch-size $GLOBAL_BS
    --train-iters $TRAIN_STEPS
    --init-method-std 0.086
    --use-mcore-models
    --no-rope-fusion
    --use-distributed-optimizer
    --use-flash-attn
    --sequence-parallel
    --recompute-granularity full
    --recompute-method block
    --recompute-num-layers $RECOMPUTE_LAYERS
    --distributed-backend nccl
    --num-layers-per-virtual-pipeline-stage 2
    --multi-latent-attention
    --qk-layernorm
    --wait-after-ckpt 0
    --overlap-grad-reduce
    --moe-shared-expert-overlap
)

--use-flash-attn 启用 Flash Attention 大幅降低内存占用并提升速度；--recompute-granularity full 配合 --recompute-method block 对整个 Transformer 块做激活重计算以节省显存；--distributed-backend nccl 设置分布式通信后端，在 MUSA 上实际是 MCCL，沿用 NCCL 名称是为了方便 CUDA 用户迁移。

算法相关参数还包括正则化（--weight-decay 0.1、--adam-beta1 0.9、--adam-beta2 0.95、--clip-grad 1.0）、学习率（--lr-decay-style cosine 余弦衰减、--lr-warmup-iters 预热步数避免训练不稳定）、模型并行（--tensor-model-parallel-size、--pipeline-model-parallel-size）、混合精度（--bf16、--attention-softmax-in-fp32）等分组，均通过 arguments.py 统一管理。训练可拉起后，通过 wandb 或 TensorBoard 做可视化监控，教材展示了 deepseek_236B 切层模型的 wandb 训练 loss 曲线示例。

3.5 MT-Megatron 代码结构：ModuleSpec 解耦"层构造"与"层执行"

MT-Megatron 的代码结构围绕模块化、可扩展设计，核心包括模型定义、数据处理、训练配置三大模块，重要代码模块包括 Torch Module（参数管理、前向计算、自动求导）、Language Module（TE AttentionBackend 指定、tp 交叉熵计算）、Megatron Module（分片 Parameter 管理）、Custom Module（Parallel State、FaultTolerance、TransformerEngine 等）。

开发者通过 CustomModel 类定义模型实体，Megatron 引入 ModuleSpec 概念（如 transformer_layer_spec），将"层如何构造"与"层如何执行"解耦：

model = CustomModel(
    config=config,
    transformer_layer_spec=transformer_layer_spec,
    vocab_size=args.padded_vocab_size,
    max_sequence_length=args.max_position_embeddings,
    pre_process=pre_process,
    post_process=post_process,
    fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
    parallel_output=True,
    share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
    position_embedding_type=args.position_embedding_type,
    rotary_percent=args.rotary_percent,
    rotary_base=args.rotary_base,
    rope_scaling=args.use_rope_scaling
)

构造函数内部首先初始化 Embedding 层，再堆叠 TransformerBlock 构建计算主体：

self.embedding = LanguageModelEmbedding(
    config=self.config,
    vocab_size=self.vocab_size,
    max_sequence_length=self.max_sequence_length,
    position_embedding_type=position_embedding_type,
    scatter_to_sequence_parallel=scatter_embedding_sequence_parallel,
)

self.decoder = TransformerBlock(
    config=self.config,
    spec=transformer_layer_spec,
    pre_process=self.pre_process,
    post_process=self.post_process,
)

ModuleSpec 方法用于封装模型参数和配置，其使用逻辑是通过 build_layer 遍历 layer_specs 逐层构建：

def _build_layers(self):
    def build_layer(layer_spec, layer_number):
        return build_module(layer_spec, config=self.config, layer_number=layer_number)

    self.layers = torch.nn.ModuleList(
        [
            build_layer(layer_spec, i + 1)
            for i, layer_spec in enumerate(self.submodules.layer_specs)
        ]
    )

def build_module(spec_or_module, *args, **kwargs):
    if isinstance(spec_or_module, types.FunctionType):
        return spec_or_module
    if isinstance(spec_or_module, ModuleSpec) and isinstance(
        spec_or_module.module, types.FunctionType
    ):
        return spec_or_module.module

MoE 场景下，Router 的初始化与前向传播示例：

# 初始化 Router
self.router = TopKRouter(config=self.config)

moe_module_spec = ModuleSpec(
    module=MoELayer,
    submodules=MoESubmodules(experts=experts, shared_experts=shared_experts)
)

def forward(self, input: torch.Tensor):
    """
    Router 的前向传播。
    参数:
        input (torch.Tensor): 输入张量。
    """
    input = self.apply_input_jitter(input)
    logits = self.gating(input)
    scores, routing_map = self.routing(logits)
    return scores, routing_map

Step-by-Step：MT-Megatron 环境与配置演练（概念性练习）

该练习依赖官方 Docker 镜像与多 GPU 集群环境，若无对应硬件资源，可作为"配置阅读与理解"练习。

1. 在提供 MT-Megatron 官方镜像的机器上执行环境安装脚本：

   cd TransformerEngine/
   chmod a+x ./install.sh
   ./install.sh

2. 参照本节示例编写 tokenizer 构建脚本，对一小批文本文件完成词元化与 .bin/.idx 序列化。
3. 参照 MODEL_ARGS、TRAINING_ARGS、REGULARIZATION_ARGS 等参数模板，为一个假设的模型规模（如 6 层、hidden-size 512，仅用于练习，非生产配置）填写完整的启动参数列表。
4. 用 torchrun --nproc_per_node=GPU数量 train.py 的形式写出对应的启动命令。
5. 用文字描述 Transformer 与 MoE 模型的前向传播流程，标注 Router、Shared MLP、MLP 专家模块的位置。

产出物：一份完整的（练习用）训练启动参数清单 + 前向传播流程文字描述。

3.6 本节要点小结

四、HPC 应用：HPL/HPCG 基准与 GROMACS GPU 加速

HPC 通过并行计算、专用硬件和优化算法聚合计算资源解决复杂大规模问题，应用于气象预测、生命科学、工业仿真等领域。其典型场景包括气象预测（高精度网格化大气-海洋耦合模型）、物理（大型强子对撞机实验数据分析、暗物质分布模拟）、生命科学（基因测序、蛋白质折叠预测）、工业领域（汽车碰撞测试、飞机气动设计与有限元分析）。当前 HPC 基础设施的主要硬件形态是 CPU+GPU 异构融合，主要负载形态是 AI-HPC 超智融合。摩尔线程针对 HPC 场景提供了功能完善的基础数学库、分布式线性代数求解器，以及 GROMACS、AMBER 等 HPC 常见仿真软件及 MCCL 通信库。

4.1 HPL 与 HPCG 基准测试

HPC 硬件效能通常通过基准测试衡量，如 HPL 与 HPCG（High Performance Conjugate Gradients）。HPL 通过求解 N 阶线性方程组衡量浮点计算性能；HPCG 聚焦大规模稀疏线性系统求解，更贴近实际应用的访存与通信特征，通常作为 HPL 的补充。

HPL 是早期 Linpack 评测程序的并行版本，支持大规模并行超级计算系统，其报告的每秒浮点运算次数（FLOPS）是 Top500 排名依据。基于 Linpack 的 Top500 排名始于 1993 年，每年发表两次。2007 年 11 月及之前，排名前 10 的超级计算机全部由同构 CPU 提供算力；2008 年 6 月，Top500 首台性能超过 1P FLOPS 的超级计算机 Roadrunner 采用了 CPU+加速芯片的异构架构；2009 年 11 月排名第 5 的 Tianhe-1 使用了 CPU+GPU 异构架构，此后 CPU+加速器架构成为主流趋势；2019 年 6 月排名前 10 中有 7 台系统使用 CPU+加速器架构。由于 HPL 主要计算集中在 DGEMM（双精度矩阵乘），具有很高的计算访存比，属于计算瓶颈应用，因此 CPU+GPU 等异构加速架构将持续成为支撑大规模科学计算的主流方案。

HPL 通过迭代法求解 N 阶线性方程组：先用带行交换的高斯消元法对系数矩阵做 LU 分解，再进行三角回代求解。LU 分解分为 FACT（Panel factorization）、BCAST（Panel broadcast）、SWAP（Row swapping）、UPDATE（Trailing submatrix update）四个步骤。

HPCG 是一种新的超级计算机排名衡量标准，主要包含 SpMV（稀疏矩阵向量乘）、SymGS（对称高斯-赛德尔平滑器）、Global Dot Product（全局点乘）、Vector Update（向量更新）、Multigrid Preconditioner（多重网格预处理器）几个部分。

4.2 GROMACS 简介与 GPU 加速原理

GROMACS 是一个开源、多功能的分子动力学软件包，能针对具有数亿粒子的系统求解牛顿动力学方程，主要为蛋白质、脂质和核酸等具有复杂键合相互作用的生物化学分子设计，但由于计算非键合相互作用力速度极快，也被广泛用于聚合物和流体动力学等非生物系统研究。GROMACS 支持的算法非常全面，加速手段包括 CPU SIMD 指令集、CUDA、OpenCL、SYCL 等；可同时利用 CPU 和 GPU 资源并包含静态/动态负载均衡策略；使用文本配置文件操控计算任务；支持 MPI 分布式计算；使用 C++17 编写，易移植；可通过 MDModules 接口扩展。

现代 GROMACS 在多核 CPU 基础上引入 GPU 加速，主要思路是将计算量最大的"非键相互作用"和网格化电荷求和（Particle-Mesh Ewald，PME）等模块卸载到 GPU 执行，获得数倍到数十倍的性能提升：

• 非键相互作用加速：将范德华力和静电力的邻域搜索、力及能量计算等核心循环封装为 GPU Kernel，利用 GPU 大规模并行能力同时对多个粒子对进行计算，并通过 warp-shuffle/共享存储器优化数据交换与缓存。
• 网格化电荷求和：分为实空间和倒空间两部分，倒空间部分包括 FFT。GROMACS 支持将 FFT 及插值阶段卸载到 GPU，通过 muFFT 计算库完成，大幅减少 CPU-GPU 通信开销。
• 异步重叠与负载均衡：在多 GPU 多节点场景中，GROMACS 异步地在 CPU 和 GPU 之间交替执行——例如 GPU 计算非键相互作用时，CPU 处理键相互作用、约束算法等，并通过动态域分割（Dynamic Domain Decomposition）自动调整每个 GPU/CPU 的工作负载。

4.3 GROMACS 移植到 MUSA 与优化

得益于 GROMACS 良好的软件工程架构以及 MUSA 生态对 CUDA 的高度兼容性，移植非常简单：在不破坏目录结构的基础上复制 CUDA 源代码、更改文件名，将 CUDA API 调用替换为对应的 MUSA API 调用，再在 CMakeLists 脚本中添加针对 MUSA 的编译选项即可。

在配置 CMake 时开启 GPU 支持：

cmake .. -DGMX_GPU=MUSA -DGMX_MUSA_TARGET_COMPUTE=31 -DGMX_MPI=ON

常用运行命令 gmx mdrun 加上 -gpu_id 选项指定可用 GPU：

gmx mdrun -s topol.tpr -deffnm sim -ntomp 4 -npme 1 -nb gpu -pme gpu -bonded cpu -gpu_id 0,1

其中 -nb gpu 指定非键相互作用在 GPU 上计算，-pme gpu 指定网格化电荷求和在 GPU 上执行，键相互作用（-bonded cpu）仍在 CPU 上处理。

摩尔线程对 GROMACS MUSA 版本进行了深度优化：

• 热点 Kernel 优化：通过第 9 章介绍的 Moore Perf System 抓取 GROMACS 运行过程中的性能数据，定位热点 Kernel 并优化。方法包括——① 确保热点 Kernel 不存在寄存器溢出：调整 __launch_bounds__ 参数、优化 C++ 模板参数在 Kernel 内部的使用，并通过 -resource-usage 参数在编译时确定寄存器使用情况；② 在 Kernel 内部善用循环展开：减少分支指令、提高计算指令比例，某些情况下还能避免寄存器溢出，但复杂循环不一定能得到加速，需要根据具体场景测试。
• Runtime API 优化：基于 Moore Perf System 发现某些数据集上 MUSA Runtime API 调用延迟较高，例如访存操作。针对 musaMemset（针对全局内存空间填充，用于初始化）调用延迟较高的情况，将其替换为一个 GPU Kernel，并合并相邻 GPU Kernel 调用以减少调用次数、降低启动延迟；对于页锁定内存的 musaMemcpyAsync，同样可用 GPU Kernel 直接访问页锁定内存来替换，前提是保证与 GPU 传输数据的主机内存都是页锁定内存，相邻调用同样可以合并。

优化效果：GROMACS MUSA 在 STMV、BenchPEP-h 等知名数据集以及国内重点实验室的数据集上，可测得不亚于国外知名 GPU 厂商解决方案的性能。

Step-by-Step：GROMACS MUSA 移植思路梳理（概念性练习）

1. 描述将一个 CUDA 版 HPC 软件移植到 MUSA 的通用步骤（结合 GROMACS 案例）。
2. 写出在 CMake 配置阶段开启 MUSA GPU 支持所需的关键选项：

   cmake .. -DGMX_GPU=MUSA -DGMX_MUSA_TARGET_COMPUTE=31 -DGMX_MPI=ON

3. 写出一条 gmx mdrun 命令，要求：使用 4 个 OpenMP 线程、1 个 PME 进程、非键相互作用在 GPU 0 和 GPU 1 上计算、PME 在 GPU 上计算、键相互作用在 CPU 上计算。
4. 结合第 9 章 Moore Perf System 的使用方法，描述如何定位 GROMACS 运行中的热点 Kernel，并说明避免寄存器溢出的两种具体手段。

产出物：一份 GROMACS MUSA 移植与优化的操作清单。

4.4 本节要点小结

五、版本与兼容性

六、本课与相邻课程

• m11-l10（第 9 章）：性能分析工具——本课 GROMACS 优化案例直接复用 Moore Perf System 定位热点 Kernel 的方法
• m11-l4（第 4 章）：软件架构——本课 torch_musa/MUSAExtension 底层依赖第 4 章介绍的驱动、运行时库与 muDNN/MCCL
• m11-l5（第 8 章）：性能优化与扩展——本课 GROMACS 的寄存器溢出规避、循环展开手段与第 8 章的 Occupancy 优化方法一脉相承
• 本站 m4-l1：torch_musa 入门——与本课 PyTorch 编程一节互为补充