[VVycccc]

背景

InfiniTrain-Backends 已通过 MACA(PR #1)确立了 backends/<provider> 的厂商后端扩展机制。本 PR 沿用同一架构,以寒武纪 MLU 作为第二个 provider 接入,验证该机制的通用性,并补齐国产加速器的第二个实现参考。

InfiniTrain Core
        ↓
PrivateUse1 Extension Interface
        ↓
InfiniTrain-Backends
        ↓
MLU Backend (CNRT / CNNL / CNCL / BANG)

当前状态 [WIP]

已完成:

• Backend 注册:provider 名 mlu,经 PrivateUse1BackendRegistration 注册,上层 --device mlu 即可启用;
• DeviceGuardImpl(CNRT 7.x):设备管理 / stream(队列)/ 事件(notifier)/ 同步与异步内存拷贝 / 懒初始化 / 内存分配释放 / BLAS 句柄(CNNL,handle 绑定 queue);
• 构建链路:host/core 代码全部使用系统 g++ 编译(设备侧 kernel 编译留给后续 BANG),INFINITRAIN_BACKEND=mlu 一键构建;
• 链接验证:MLU590 环境(CNRT 7.x / CNNL 2.2.14 / CNCL 1.30.8)下 libinfini_train_backend_mlu.a 与 gpt2 / llama3 / mixtral 三个 example 可执行文件全部链接成功;
• 注册冒烟:启动后注册通过,框架自检按预期停在首个缺失的 required kernel(Cast),验证注册路径正确。

下一步(按依赖顺序):

1. 实现 required kernels:Cast / Fill / NoOpForward / NoOpBackward(参考 cast.maca 等 ~150 行);
2. 按框架自检提示迭代补全 CNNL/BANG 算子,直至 gpt2 训练跑通;
3. 复用 InfiniTrain shared tests 全量验证;
4. CNCL 多卡分布式验证。

与 MACA Backend 的实现差异

MLU 侧软件栈与 MACA 不同构,实现上做了相应适配:

• CNRT 无异步 malloc → 沿用 MACA PR 的同步分配 workaround 模式;
• GEMM 无独立 BLAS 库 → 直接使用 CNNL(handle 与 queue 绑定,cnnlSetQueue);
• 事件实现基于 CNRT notifier(place / wait / query / duration);
• 设备侧算子计划采用 BANG C(.mlu),host 侧 dispatch 用系统工具链编译;
• 错误处理:cnrtGetErrorStr / cnnlGetErrorString,检查宏区分 CNRT/CNNL。

Baseline 设备验证(配套 InfiniTrain-Test 仓库日志)

后端尚未跑通模型前,先以 torch_mlu 层面的 baseline 验证了 MLU590 设备与软件栈的可用性(日志与脚本提交在 InfiniTrain-Test 仓库 baseline/{pytorch,megatron}/logs/mlu590_wangyichen_node2*/):

• 44 个 PyTorch run(training 12 / LoRA 12 / lr_scheduler 20)+ 4 个 Megatron Mixtral MoE run,全部通过;
• 跨卡规约精确:CNCL allreduce(AVG)模式下 8 卡 DDP 轨迹与单卡逐位一致(三平台均如此);
• RNG 兼容:torch_mlu 设备侧 RNG 流与 CUDA 兼容,Megatron 随机 init 场景 step1 loss 与 A100 差仅 3e-6;
• 精度对齐:官方 llmc 数据下 fp32 step1 loss 与 A100 差 ≤ 2e-6(多数逐位一致),bf16 全部在阈值内;10 步训练轨迹的 fp32 发散属跨厂商算子内规约顺序差异(归因实验排除跨卡规约;现有 CUDA↔MACA 基线同口径偏差比 MLU 大 2–4 倍)。

当前约束

与 MACA PR 相同:一个 build tree 只选择一个 provider;一个进程最多注册一个 PrivateUse1 provider。