feat: add PrivateUse1 backend extension support - #206
Open
chen2021673 wants to merge 1 commit into
Open
Conversation
JYMiracle305
self-requested a review
August 17, 2026 07:17
| auto hook = std::make_unique<infini_train::autograd::AllReducePostAccumulateHook>( | ||
| function::ReduceOpType::kAvg, ddp_pg_); | ||
| const auto reduce_op | ||
| = ddp_config.average_in_collective ? function::ReduceOpType::kAvg : function::ReduceOpType::kSum; |
Contributor
There was a problem hiding this comment.
关于 reduce_op 类型和架构后端实现不是强相关的，是不是单独提一个PR
| } else { | ||
| bucket.work = ddp_pg->AllReduce(bucket.contents, function::ReduceOpType::kAvg, true); | ||
| const auto reduce_op | ||
| = ddp_config_.average_in_collective ? function::ReduceOpType::kAvg : function::ReduceOpType::kSum; |
| @@ -143,18 +145,11 @@ void DeviceGuardImplRegistry::Register(Device::DeviceType type, std::unique_ptr< | |||
| LOG(FATAL) << std::format("DeviceGuardImpl for type {} already registrered", static_cast<int>(type)); | |||
| } | |||
|
|
|||
Contributor
There was a problem hiding this comment.
删除单加速器后端限制后，Tensor::To(Device) 中原有的跨后端复制路径就可以覆盖到了， tensor.cc::161 存在一个问题，第二步 H2D 复制根据 buffer_device获取impl，本来应该使用目标 device来获取tmpl。这里comment作记录，可以另外PR修复，加单元测例覆盖一下。
| # ------------------------------------------------------------------------------ | ||
|
|
||
| add_library(infini_train STATIC ${SRC}) | ||
| add_library(InfiniTrain::infini_train ALIAS infini_train) |
Contributor
There was a problem hiding this comment.
新增的 InfiniTrain::infini_train alias 是不是给外部 provider 直接链接使用的？目前 runtime、CCL 和 kernel 都依赖静态注册，而保证这些注册代码不被链接器裁掉的 --whole-archive 只加在 link_infini_train_exe() 里。如果外部工程直接 target_link_libraries(... InfiniTrain::infini_train)，没有调用 link_infini_train_exe()，运行时报 runtime 或 kernel 未注册
| void RegisterFakeRuntime() { | ||
| CHECK_EQ(core::GetPrivateUse1BackendName(), "fake"); | ||
| CHECK_EQ(Device(Device::DeviceType::kPrivateUse1, 0).ToString(), "Device(fake, 0)"); | ||
| INFINI_TRAIN_REGISTER_DEVICE_GUARD_IMPL(Device::DeviceType::kPrivateUse1, FakePrivateUse1GuardImpl) |
Contributor
There was a problem hiding this comment.
warning: unused variable ‘__infini_train_device_guard_registered__COUNTER__’ [-Wunused-variable]
236 | static const bool __infini_train_device_guard_registered##__COUNTER__ = []() { \
| ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
这里会有一个告警，因为之前加了-Wunused 编译选项，是不是在 Register 宏里加一下[[maybe_unused]]。
同时也发现了一个问题，在宏里 __COUNTER__直接接触 ##，没有展开成数字，在此记录一下，后续另提PR修改。
| @@ -143,18 +145,11 @@ void DeviceGuardImplRegistry::Register(Device::DeviceType type, std::unique_ptr< | |||
| LOG(FATAL) << std::format("DeviceGuardImpl for type {} already registrered", static_cast<int>(type)); | |||
- add a provider-neutral PrivateUse1 device type and registration API - validate runtime, kernel, and optional CCL backend registrations - initialize external device runtimes lazily on first use - support provider names in device parsing and display - require explicit autocast dtype for PrivateUse1 devices - allow examples to register an external backend before flag parsing - honor average_in_collective consistently across DDP paths - expose embeddable CMake targets and add fake backend tests
chen2021673
force-pushed
the
refactor/privateuse1-backend
branch
from
August 20, 2026 08:38
2d8e754 to
0c5953b
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
背景
InfiniTrain 原有设备体系只包含 CPU 和 CUDA。接入新的后端时，需要在核心框架中增加厂商专属的
DeviceType、runtime、CCL、kernel 以及模型侧判断，导致核心代码与具体厂商耦合。本 PR 引入通用的
DeviceType::kPrivateUse1扩展槽位，使外部厂商后端可以在不修改 InfiniTrain 核心设备枚举和模型逻辑的情况下注册自己的实现。MACA 后端将作为独立仓库使用该接口接入。主要改动
PrivateUse1 注册接口
新增
PrivateUse1BackendRegistration和RegisterPrivateUse1Backend()，统一编排以下注册流程：macaDeviceGuardImplCclImpl保留并复用现有三个注册宏：
REGISTER_KERNELINFINI_TRAIN_REGISTER_DEVICE_GUARD_IMPLINFINI_TRAIN_REGISTER_CCL_IMPLPrivateUse1 后端至少需要提供以下基础算子：
CastFillNoOpForwardNoOpBackward设备名称解析
新增统一的
Device::ParseType()：cpu映射到kCPUcuda映射到kCUDAprivateuse1映射到kPrivateUse1maca，映射到kPrivateUse1Device::ToString()同样使用注册后的厂商名称展示设备。Example 外部后端入口
GPT2、LLaMA3 和 Mixtral example 支持在解析 --device 前调用外部 backend registrar，同时不直接依赖任何厂商头文件。
外部仓库可以注入：
Test