feat: add PrivateUse1 backend extension support - #206
Open
chen2021673 wants to merge 6 commits into
Open
chen2021673 wants to merge 6 commits into
chen2021673 wants to merge 6 commits into
Conversation
JYMiracle305
self-requested a review
August 17, 2026 07:17
chen2021673
force-pushed
the
refactor/privateuse1-backend
branch
3 times, most recently
from
August 27, 2026 05:39
4b1042c to
66cc91e
Compare
chen2021673
changed the base branch from
master
to
fix/backend-independent-correctness
August 31, 2026 08:18
chen2021673
force-pushed
the
refactor/privateuse1-backend
branch
from
August 31, 2026 09:30
66cc91e to
54dc847
Compare
chen2021673
force-pushed
the
refactor/privateuse1-backend
branch
from
September 2, 2026 09:31
54dc847 to
e8a25ee
Compare
JYMiracle305
approved these changes
Sep 3, 2026
Base automatically changed from
fix/backend-independent-correctness
to
master
September 11, 2026 06:46
chen2021673
force-pushed
the
refactor/privateuse1-backend
branch
from
September 14, 2026 02:36
d9fe648 to
9b4a33c
Compare
kilinchange
requested changes
Sep 16, 2026
| DEFINE_bool(overfit_single_batch, true, "overfit just one batch of data"); | ||
| // memory management | ||
| DEFINE_string(device, "cuda", "device type (cpu/cuda), useless if using parallel training mode"); | ||
| DEFINE_string(device, "cuda", "device type, useless if using parallel training mode"); |
Collaborator
There was a problem hiding this comment.
cpu/cuda/privateuse1/<registered privateuse1 backend name>
改成这样吧，避免用户对大小写或合法取值产生误解。
| // rely solely on file-scope static initialization because an unreferenced static archive member may be discarded. | ||
| using PrivateUse1RegistrationCallback = void (*)(); | ||
|
|
||
| struct PrivateUse1BackendRegistration { |
Collaborator
There was a problem hiding this comment.
如果走静态注册的话，这部分是不是不需要了？
Contributor
Author
There was a problem hiding this comment.
name和default_autocast_dtype应该还是需要，三个回调函数不需要了，我删掉。
| #ifdef INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS | ||
| // FIXME(cx): MACA needs synchronization before training teardown. | ||
| if (INFINITRAIN_EXAMPLE_EXTERNAL_BACKEND_WORKAROUNDS && device.type() == Device::DeviceType::kPrivateUse1) { | ||
| impl->SynchronizeDevice(device); |
Collaborator
There was a problem hiding this comment.
这里麻烦再确认下 main.cc 里这两处 workrounds 是否仍然有必要：
- 沐曦平台改成多进程启动后，是否还会复现原问题；
- 沐曦在多线程分布式启动模式下，确认这两处是否仍然必要（考虑到框架经过一段时间迭代，可能有些 bug 被修复了）。
Contributor
Author
There was a problem hiding this comment.
- 已验证多线程问题依然存在，多进程不会复现原问题，但需要一处新的修复；
- 因为多线程分布式启动模式下问题依然存在，所以两处修复仍有必要，目前在多进程启动模式下单独关闭这两处 workrounds 修复及异步 malloc/free 修复，来观察性能数据。
|
|
||
| TEST(PrivateUse1BackendValidationTest, RejectsNonAsciiBackendNames) { | ||
| auto registration = MinimalRegistration(); | ||
| registration.name = "invalid-name"; |
Collaborator
There was a problem hiding this comment.
这里测试的并不是 non-ASCII，而是不支持的 backend name 字符（如 -），建议改下测试名。
Contributor
Author
There was a problem hiding this comment.
改为 RejectsUnsupportedBackendNameCharacters
- 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
- separate test suite declaration from CPU, CUDA, and provider instantiation - support provider-injected registration, linkage, target names, and CTest labels - add provider-defined default autocast dtype and backend-neutral test helpers - generalize accelerator copy tests and remove the ineffective CUDA optimizer test - fix Exp/Add backward, CUDA bias reduction, and DDP device validation - scope runtime workarounds to the registered MACA backend - document external backend test integration and usage
- group CCL initialization only for multiple local communicators
Register provider metadata explicitly while leaving runtime, kernel, and CCL implementations in their statically initialized registries. Validate the runtime and required kernels after metadata registration.
chen2021673
force-pushed
the
refactor/privateuse1-backend
branch
2 times, most recently
from
September 17, 2026 03:18
99b1438 to
e552299
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
3 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。接入新后端时，需要在核心框架中增加厂商专属的设备枚举、runtime、CCL、kernel、测试和模型入口判断，导致核心代码与具体厂商耦合。
本 PR 引入通用的
DeviceType::kPrivateUse1扩展槽位。外部 Provider 可以注册自己的运行时和算子实现，MACA 后端则通过独立仓库接入，公共 CMake 中不包含 MACA SDK 或厂商构建选项。主要改动
PrivateUse1 注册接口
新增
PrivateUse1BackendRegistration和RegisterPrivateUse1Backend()，统一完成：macaDeviceGuardImplCclImplProvider 名称仅允许小写 ASCII 字母、数字和下划线，且不能占用
cpu、cuda。继续复用
REGISTER_KERNEL、INFINI_TRAIN_REGISTER_DEVICE_GUARD_IMPL和INFINI_TRAIN_REGISTER_CCL_IMPL。设备解析与模型入口
新增统一的
Device::ParseType()：cpu映射到kCPUcuda映射到kCUDAprivateuse1映射到kPrivateUse1kPrivateUse1Device::ToString()使用注册后的厂商名称。AutocastGuard根据 Provider 注册信息取得 PrivateUse1 的默认计算类型，CPU 和 CUDA 的现有行为保持不变。GPT-2、Llama 3 和 Mixtral 使用
Device::ParseType()解析设备。外部工程可通过编译定义注入 Provider 头文件和注册入口，并在 gflags 校验--device前完成注册。并行模型使用用户选择的 accelerator backend，不再固定为 CUDA。运行脚本新增
DEVICE_BACKEND，默认值仍为cuda。GPT-2 和 Llama 3 暂时保留少量仅在 Provider 名称为
maca时启用的同步和进程退出 workaround，并保留 FIXME；这些逻辑不会影响其他 PrivateUse1 Provider。构建与静态注册
新增并明确以下 CMake 接口：
InfiniTrain::infini_train：供库和 Provider 使用的核心接口InfiniTrain::cpu_kernels：CPU kernel targetInfiniTrain::infini_train_executable：最终可执行文件的完整链接接口最终可执行文件通过 archive group 和
--whole-archive保留 DeviceGuard、CCL 和 kernel 的静态注册对象。Provider 可以通过自己的 executable interface 或EXTRA_ARCHIVES加入 Provider 注册 archive。InfiniTrain 作为 submodule 使用时不再构建自身 examples 和 tools，并隔离 glog 的测试选项，避免污染上层工程。
测试复用
公共测试的 suite 声明与设备实例化分离，每个测试二进制只实例化一个设备。CMake 为目标注入设备类型、GTest 前缀和
DEVICE_INDEX，其中设备序号默认是0。新增
infini_train_add_privateuse1_test_suites()，允许外部 Provider 复用全量公共测试：USE_CUDA=OFF，冲突时直接报错test_*_<BACKEND_NAME>Provider 测试test_main在测试环境初始化前显式调用ONLY_CUDA，copy 等测试改为通用 accelerator 测试Provider 测试使用厂商名作为 CTest label，例如 MACA 使用
ctest -L maca；不提供ctest -L privateuse1标签。新增无需真实硬件的 fake PrivateUse1 测试，覆盖注册校验、名称解析、默认 autocast dtype、延迟 runtime 初始化和基础 kernel 调度。
Test