国产AI芯片实战：寒武纪MLU370-X8运行LLaMA-Factory微调全流程深度解析

当技术自主可控成为行业刚需，国产AI芯片的实战能力究竟如何？去年在部署某金融风控模型时，我们首次尝试将训练任务从N卡迁移到寒武纪MLU370-X8平台，整个过程犹如在未知海域航行——官方文档的只言片语、社区经验的零星碎片，每一步都充满试探。本文将完整呈现从环境配置到模型导出的全链路实战经验，重点解析那些官方手册未曾提及的"暗礁"与应对策略。

1. 环境配置：破解定制化生态的密码

寒武纪MLU平台的特殊性在于其完整的定制化软件栈。与通用GPU环境不同，我们需要重新理解整个工具链的协作关系。官方提供的cambricon_pytorch_container:v24.02.1镜像已经预装了PyTorch 2.1.0和部分基础组件，但这仅仅是起点。

1.1 关键组件安装实战

核心依赖库的安装需要特别注意版本匹配问题。以下是经过验证的组件矩阵：

具体操作时，deepspeed的兼容性处理最为关键。由于LLaMA-Factory会检查deepspeed包而非deepspeed_mlu，必须修改dist-info元数据：

bash复制# 修正deepspeed包识别问题
cp -r /torch/venv3/pytorch/lib/python3.10/site-packages/deepspeed_mlu-0.10.1.dist-info/ \
      /torch/venv3/pytorch/lib/python3.10/site-packages/deepspeed-0.10.1.dist-info/
sed -i 's/Name: deepspeed-mlu/Name: deepspeed/g' \
      /torch/venv3/pytorch/lib/python3.10/site-packages/deepspeed-0.10.1.dist-info/METADATA

提示：所有通过torch_gpu2mlu.py转换的库都需要用pip install -e .方式安装，这会保留源码链接便于调试

1.2 环境验证技巧

建立验证环境时，建议按以下顺序检查：

1. 基础算子支持：torch.mlu.is_available()返回值
2. 内存管理：torch.mlu.empty_cache()的实际效果
3. 混合精度：torch.mlu.amp的可用性
4. 分布式通信：torch.distributed在MLU上的初始化

我们在测试中发现，当使用8卡配置时，需要额外设置：

python复制os.environ['CNCL_IB_HCA'] = 'mlx5_0'  # 指定RDMA网卡
os.environ['CNCL_IB_GPU_DIRECT'] = '1'  # 启用GPUDirect

2. LLaMA-Factory的适配改造

原生的LLaMA-Factory框架对MLU支持存在多处盲点，需要针对性改造。从GitHub克隆Mu-L维护的特定分支后，关键的适配点集中在设备管理和内存回收机制。

2.1 设备管理层的修改

框架中所有cuda相关调用都需要替换为设备无关实现。最稳妥的方式是修改src/llmtuner/core/utils.py中的设备检测逻辑：

python复制def get_device() -> torch.device:
    if torch.mlu.is_available():
        return torch.device("mlu")
    elif torch.cuda.is_available():
        return torch.device("cuda")
    else:
        return torch.device("cpu")

2.2 内存回收机制调整

MLU370的显存管理策略与N卡存在显著差异。原始代码中的torch.cuda.ipc_collect()在MLU上不可用，需要注释掉misc.py中的相关行：

python复制def torch_gc():
    gc.collect()
    if torch.mlu.is_available():
        torch.mlu.empty_cache()
        # 注释掉以下行
        # torch.mlu.ipc_collect()  

注意：这个修改会影响长时间训练时的内存累积，建议将--logging_steps调小以便更频繁触发回收

3. 微调实战：参数配置的黄金法则

在MLU370-X8上微调ChatGLM3-6b时，参数配置需要遵循三个特殊原则：

1. 批次尺寸保守原则：相同模型下，MLU的batch_size通常需设为N卡的1/2
2. 梯度累积补偿原则：通过增加gradient_accumulation_steps维持总batch量
3. 序列长度折衷原则：cutoff_len建议不超过1024以避免OOM

以下是我们验证过的两种典型配置：

单卡高效配置：

yaml复制per_device_train_batch_size: 2
gradient_accumulation_steps: 8
cutoff_len: 1024
lr: 5e-5
warmup_ratio: 0.03

8卡分布式配置：

bash复制deepspeed --num_gpus 8 --master_port=9901 src/train_bash.py \
    --deepspeed ds_config.json \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 16 \
    --cutoff_len 768

在具体任务中，医疗问答数据微调显示：

• 收敛速度：MLU370比V100快1.7倍（相同batch_size）
• 显存效率：MLU的24G显存实际可用约21.5G
• 通信开销：8卡AllReduce耗时比NVIDIA NCCL高约15%

4. 模型导出与性能对比

模型导出阶段的主要挑战在于设备类型检查。需要修改peft库的适配器加载逻辑：

python复制# 修改peft/utils/save_and_load.py
def load_adapter_weights(filename, device):
    if str(device).startswith('mlu'):
        device = 'cpu'  # 先加载到CPU再转移到MLU
    return torch.load(filename, map_location=device)

性能对比数据揭示了一些有趣现象：

这个数据反映出MLU370-X8的两个典型特征：

1. 计算密度优势：在适当batch下算力释放充分
2. 显存瓶颈明显：相同模型下支持的batch_size较小

在金融风控实际场景中，MLU370完成20000条样本微调耗时2.1小时，相比同价位GPU集群有15%的成本优势，但调试时间会多出3-5个工作日。这种tradeoff使得它更适合：

• 对数据隐私要求极高的场景
• 已有寒武纪部署环境的机构
• 需要长期稳定运行的生产系统

5. 常见陷阱与进阶技巧

三个月内我们踩过的坑值得专门总结：

典型报错1：RuntimeError: CNCL未初始化

• 原因：未正确设置多卡环境变量
• 解决：

  bash复制export CNCL_IB_HCA=mlx5_0
  export CNCL_IB_GPU_DIRECT=1
  

典型报错2：OOM when using batch_size=1

• 检查点：
  1. 确认已注释ipc_collect调用
  2. 尝试设置torch.mlu.set_per_process_memory_fraction(0.8)
  3. 降低cutoff_len到768

性能调优技巧：

• 在ds_config.json中启用梯度检查点：

  json复制{
    "gradient_checkpointing": {
      "use_reentrant": false,
      "mlu_optimized": true
    }
  }
  

• 使用--fp16而非--bf16（MLU对fp16优化更好）
• 将数据预处理移至CPU：Dataset.set_transform(fn, device='cpu')

6. 生态现状与发展建议

当前寒武纪MLU生态存在三个明显短板：

1. 第三方库适配滞后（如bitsandbytes缺失）
2. 调试工具链不完善（缺乏等效NSight的工具）
3. 社区知识沉淀不足（问题解决依赖官方支持）

但优势同样明显：

• 硬件计算密度实际测试比同代GPU高20-30%
• 自主可控带来的安全合规价值
• 针对中文NLP任务的特定优化

对于考虑迁移的团队，建议采取渐进策略：

1. 先在小规模评估任务验证
2. 建立内部知识库积累经验
3. 与寒武纪技术团队保持紧密沟通
4. 关键业务保留GPU备选方案

在完成六个医疗对话模型的微调后，我们发现当工作流稳定后，MLU平台的性价比优势会逐渐显现。特别是在需要持续增量训练的场景，省去的License费用相当可观。