Skip to content

[WIP]feat: add MACA backend support - #1

Open
chen2021673 wants to merge 10 commits into
masterfrom
feat/maca-backend
Open

chen2021673 wants to merge 10 commits into
masterfrom
feat/maca-backend

Conversation

@chen2021673

@chen2021673 chen2021673 commented Aug 20, 2026

Copy link
Copy Markdown

背景

InfiniTrain 已提供基于 PrivateUse1 的后端扩展接口,但厂商相关的编译器、SDK、运行时、算子和通信实现不适合继续放在框架核心仓库中。

本 PR 初始化 InfiniTrain-Backends,建立独立于 InfiniTrain 核心仓库的 accelerator backend 扩展机制,并以 MACA 作为首个 provider 完成验证。

整体目标是形成如下扩展关系:

InfiniTrain Core
        ↓
PrivateUse1 Extension Interface
        ↓
InfiniTrain-Backends
        ↓
Vendor Backend

作为后续国产加速器接入的统一组织方式。

整体架构

仓库按照 provider 进行隔离:

InfiniTrain-Backends
├── backends/
│   └── maca/
├── examples/
├── tests/
└── third_party/
    └── InfiniTrain

每个厂商后端独立维护在:

backends/<provider>

其中包含该后端自身的构建配置、Backend 注册以及具体实现。

InfiniTrain 通过 submodule 固定一个明确版本,Backend 仓库依赖 InfiniTrain 提供的 PrivateUse1 接口,而不是反向修改框架核心代码。

这样可以保证:

  • InfiniTrain 核心保持与具体硬件厂商无关;
  • 不同厂商 Backend 相互隔离;
  • Backend 可以独立开发、测试和演进;
  • InfiniTrain 与 Backend 之间通过统一接口协作。

Backend 接入逻辑

构建时通过:

INFINITRAIN_BACKEND=<provider>

选择当前使用的 Backend。

一个 build tree 只对应一个 provider。

Backend 的接入主要分为两个阶段。

构建阶段

部分厂商 Backend 不只是链接额外的 runtime library,而是需要使用自己的编译器工具链。

CMake 的 compiler 必须在第一次 project() 之前确定,因此顶层工程会先加载:

backends/<provider>/cmake/pre_project.cmake

完成 provider 相关的编译环境配置,再进入 InfiniTrain 的正常构建流程。

因此 Backend 不只是 InfiniTrain 构建完成后附加的一组库,而是整个构建过程的一部分。

运行时注册阶段

每个 provider 对外提供统一的 Backend 注册入口。

MACA 当前通过:

infini_train::maca::RegisterBackend();

将自身能力注册到 InfiniTrain 的 PrivateUse1 Backend。

注册完成后,上层模型和训练逻辑仍然使用 InfiniTrain 原有抽象:

Model / Autograd / Distributed
              ↓
         Dispatcher
              ↓
        PrivateUse1
              ↓
           MACA

上层代码只需要感知 device = maca,而不需要直接依赖 MACA SDK 或 MACA 特有接口。

MACA Backend

本 PR 将 MACA 支持迁移到新的 Backend 架构下,并适配 InfiniTrain 当前的 PrivateUse1 接口。

MACA Backend 负责承接厂商相关能力,并通过统一注册机制暴露给 InfiniTrain。

整体逻辑为:

InfiniTrain 定义 Backend 接口
        ↓
MACA 实现 Backend 接口
        ↓
启动时注册到 PrivateUse1
        ↓
InfiniTrain Dispatcher 统一调度

核心目标是在 InfiniTrain 中不增加 MACA 专用执行分支,而是通过统一 Backend 接口完成硬件能力接入,从而保持框架逻辑与硬件实现解耦。

测试与开发方式

测试基于 InfiniTrain 的 PrivateUse1 公共测试体系。

Backend 仓库主要负责将指定 provider 注册到测试程序中,而 Tensor、Autograd 等通用行为继续由 InfiniTrain 的 shared tests 进行验证。

这样不同 Backend 可以复用同一套接口语义和测试标准,减少厂商后端重复维护测试逻辑。

开发场景下,默认使用仓库中固定版本的 InfiniTrain submodule;同时支持通过 INFINITRAIN_SOURCE_DIR 指向本地 InfiniTrain 工作区,方便框架接口与 Backend 联调。

当前约束

当前设计下:

  • 一个 build tree 只选择一个 provider;
  • 一个进程最多注册一个 PrivateUse1 provider;
  • provider 通过统一接口接入 InfiniTrain;
  • 厂商相关实现保持在 InfiniTrain 核心仓库之外。

本 PR 以 MACA 验证上述 Backend 架构,为后续其他 accelerator provider 接入提供统一的目录结构、构建流程和运行时注册模式。

- align clang-format configuration with InfiniTrain
- format existing MACA backend sources
- add GitHub Actions format checking
- support .maca files in the formatting script
Comment thread backends/maca/src/kernels/register_maca_kernels.cc Outdated
@chen2021673

Copy link
Copy Markdown
Author
img_v3_0214t_f6ab87d7-e0cd-4652-8c6e-dc15de21822g

- run upstream PrivateUse1 suites through BUILD_TEST
- reject incompatible CUDA builds and remove the local allocator test
- register BF16 autocast and fix pthread detection
- correct elementwise launches, embedding indexing, linear reductions, and layer norm gradients
- remove redundant tensor initialization and strengthen dtype and shape checks
- scope ballot and shuffle operations to 32-lane logical warps
- use matching CUB reduction storage for logical warp groups
- remove invalid host/device attributes from the infinity constant
Use tiled FP32 multiplication with fixed-order FP64 accumulation in linear backward input.
Hoist kernel, runtime, and MCCL registrations to file scope, remove the
centralized MACA kernel registrar, and rely on the whole-archive executable
link interface to retain registration-only objects.
@chen2021673

chen2021673 commented Sep 18, 2026

Copy link
Copy Markdown
Author

精度测试(threshold-fp32 1e-5),将 MACA 结果与 CUDA(g3025)运行结果比较:
image
image
主要精度差来源于 LinearBackward 浮点加法规约顺序带来的误差。修复后精度能达标但性能会大幅下降,因此现阶段删除该修复。当 threshold-fp32 1e-4 时大部分测例可通过:
image

性能测试:
image
大部分性能落后,但部分测例差距在20%以内,具体见https://gxtctab8no8.feishu.cn/wiki/OL5BwuWZtixckNkyKeWcZTtwnVh

Comment thread README.md
infini_train::Device device(type, 0);
```

`RegisterBackend()` installs the process-wide PrivateUse1 name, MACA kernels,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

现在 RegisterBackend() 只注册 meta data 了,需要更新下 README。

}
}

void MacaGuardImpl::ResetMemPoolHighWatermarks(Device device) const {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

https://github.com/InfiniTensor/InfiniTrain/blob/e661492a8a741de303123e319ee4f063be7a8d09/infini_train/src/core/runtime/maca/maca_guard_impl.cc#L764
我看 InfiniTrain 之前适配的代码里 memInfo 这块是实现了的,是后来发现还是有问题吗?

@@ -0,0 +1,721 @@
{
"variables": {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

没有 xxx_LORA_WEIGHTS_DIR,确认下是否跟进了 InfiniTrain 对于 lora 精度的修复呢?


#include "common/common_maca.h"
#include "infini_train/include/core/runtime/device_guard.h"
#include "infini_train/include/dispatcher.h"

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

建议按依赖层次对 include 分组,例如将 infini_train 头文件单独分组,backend 内部头文件单独分组放在最后,方便阅读和维护。
我看似乎有的文件做了单独分组,有的没做,建议统一一下。

#include <cstdint>
#include <memory>

#include "common/kernel_helper.cuh"

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

同上

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants