feat: add PrivateUse1 backend extension support - #206
Open
chen2021673 wants to merge 1 commit into
Open
Conversation
- 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
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
1 participant
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