diff --git a/docs/source_en/Components/Checkpoint Engine/CheckpointEngine.md b/docs/source_en/Components/Checkpoint Engine/CheckpointEngine.md index 1a7c39bfa..a7ca33a81 100644 --- a/docs/source_en/Components/Checkpoint Engine/CheckpointEngine.md +++ b/docs/source_en/Components/Checkpoint Engine/CheckpointEngine.md @@ -39,7 +39,7 @@ class CheckpointEngine(ABC): ## Available Checkpoint Engines -Twinkle provides two checkpoint engine implementations: +Twinkle provides three checkpoint engine implementations: ### NCCLCheckpointEngine @@ -61,10 +61,21 @@ A checkpoint engine that uses HCCL for weight transfer between Ascend NPUs. See: [HCCLCheckpointEngine](HCCLCheckpointEngine.md) +### XCCLCheckpointEngine + +A checkpoint engine that transfers weights over BKCL (XCCL) on Kunlunxin XPU. + +- XPU Support: Drop-in for NCCLCheckpointEngine on the cuda-alike XPU route +- Relay Fallback: Handles duplicate local device indices via a stateless `ProcessGroupXCCL` plus a socket relay +- Compatible Interface: Inherits `NCCLCheckpointEngine`, reusing bucketing and metadata handshake + +See: [XCCLCheckpointEngine](XCCLCheckpointEngine.md) + ## How to Choose - **NCCLCheckpointEngine**: Suitable for GPU environments, provides the highest transfer performance - **HCCLCheckpointEngine**: Suitable for Ascend NPU environments +- **XCCLCheckpointEngine**: Suitable for Kunlunxin XPU environments (selected automatically) > Checkpoint engine is a key component of RLHF training infrastructure, ensuring that trainers and samplers use consistent model weights. > Currently, synchronization is divided into two cases based on merge_and_sync=True/False. When set to True, the LoRA is merged into the base model and then synchronized. diff --git a/docs/source_en/Components/Checkpoint Engine/XCCLCheckpointEngine.md b/docs/source_en/Components/Checkpoint Engine/XCCLCheckpointEngine.md new file mode 100644 index 000000000..6bbaf1c61 --- /dev/null +++ b/docs/source_en/Components/Checkpoint Engine/XCCLCheckpointEngine.md @@ -0,0 +1,37 @@ +# XCCLCheckpointEngine + +A checkpoint engine for Kunlunxin XPU that transfers weights over BKCL (mounted into PyTorch as the `nccl`/XCCL backend). + +## Usage Example + +```python +from twinkle.checkpoint_engine import XCCLCheckpointEngine + +engine = XCCLCheckpointEngine(bucket_size=512<<20) +# Usage is the same as NCCLCheckpointEngine +``` + +On XPU this engine is selected automatically by `CheckpointEngineManager` / `CheckpointEngineMixin`, so you normally do not construct it directly. + +## Features + +- **Drop-in for NCCL**: Inherits `NCCLCheckpointEngine`; bucketing, ZMQ metadata handshake, and double buffering are reused unchanged. +- **XCCL group + relay fallback**: Replaces the direct `ProcessGroupNCCL` construction (which deadlocks when two ranks share a local device index) with a stateless `ProcessGroupXCCL`, plus a relay path for colliding ranks. +- **Store-based init barrier**: The readiness barrier runs through the TCPStore instead of a device collective. + +## Why a Separate Engine + +BKCL identifies a rank's device by its **local device index** and does not translate `CUDA_VISIBLE_DEVICES` to physical cards. In a separated trainer/sampler deployment on a single host, both sides see their device as local `0`, producing duplicate `(host, index)` pairs. Directly building `ProcessGroupNCCL` then fails to form the ring and the first broadcast deadlocks (600s WorkXCCL timeout). + +`XCCLCheckpointEngine` exchanges `(hostname, local_device_index)` through the store, keeps colliding ranks out of the XCCL group, and serves them through the lowest-rank XCCL member over direct TCP sockets (relay). Non-colliding ranks use a flat stateless `ProcessGroupXCCL`. + +## Use Cases + +- Weight synchronization (train ↔ sample) on Kunlunxin XPU +- Colocate / separated RL (GRPO) with full-weight sync (`merge_and_sync=True`) + +## Limitations + +- Multi-host (inter/intra-node) broadcast is implemented but not yet validated; only single-host relay mode is verified. + +> On Kunlunxin XPU, use `merge_and_sync=True` for weight sync; LoRA sampling is currently blocked by vendor kernel dtype limitations (see the [XPU Support](../../Usage%20Guide/XPU-Support.md) guide). diff --git a/docs/source_en/Components/Checkpoint Engine/index.rst b/docs/source_en/Components/Checkpoint Engine/index.rst index bcd188429..cec1df80a 100644 --- a/docs/source_en/Components/Checkpoint Engine/index.rst +++ b/docs/source_en/Components/Checkpoint Engine/index.rst @@ -6,3 +6,4 @@ Checkpoint Engine CheckpointEngine.md NCCLCheckpointEngine.md HCCLCheckpointEngine.md + XCCLCheckpointEngine.md diff --git a/docs/source_en/Usage Guide/Installation.md b/docs/source_en/Usage Guide/Installation.md index 5dd5342e5..a9d4d7ede 100644 --- a/docs/source_en/Usage Guide/Installation.md +++ b/docs/source_en/Usage Guide/Installation.md @@ -53,5 +53,6 @@ sh INSTALL_MEGATRON.sh | GPU A10/A100/H100/RTX series | | | GPU T4/V100 | Does not support bfloat16, Flash-Attention | | Ascend NPU | Some operators not supported | +| Kunlunxin XPU (P800) | Supported via the cuda-alike route; some operators fp16-only, see [XPU Support](XPU-Support.md) | | PPU | Supported | | CPU | Supports partial components like dataset, dataloader | diff --git a/docs/source_en/Usage Guide/XPU-Support.md b/docs/source_en/Usage Guide/XPU-Support.md new file mode 100644 index 000000000..c5d4afbf1 --- /dev/null +++ b/docs/source_en/Usage Guide/XPU-Support.md @@ -0,0 +1,126 @@ +# XPU (Kunlunxin) Quick Start Guide + +This document describes how to install and use the Twinkle framework on Kunlunxin XPU (P800 series). + +## How XPU Support Works + +Unlike Ascend NPU, Kunlunxin XPU does **not** introduce a separate device type in Twinkle. It relies on `torch_xmlir`, which rewrites `torch.cuda` symbols at import time so that: + +- `torch.cuda.is_available()` returns `True` and devices report as `cuda:N` +- `libbkcl.so` (BKCL) is mounted as PyTorch's `nccl` backend (XCCL) +- `CUDA_VISIBLE_DEVICES` is taken over by the XPU runtime + +Because of this "cuda-alike" route, most of Twinkle's GPU code path works unchanged. Only a few CUDA-specific assumptions that the symbol rewrite does not cover need adaptation. The `XPU` platform therefore **inherits from `GPU`** (`src/twinkle/utils/platforms/xpu.py`) rather than being a fully independent device backend like NPU. + +## Environment Requirements + +Adaptation was validated inside the vendor container image (`kylin_v11-vllm_021_swift:*`). The verified component matrix: + +| Component | Version | Notes | +|---|---|---| +| OS / Python | Kylin V11 / 3.10.14 | Meets Twinkle `>=3.10` | +| torch | 2.9.0 | Symbol-rewritten by `torch_xmlir` | +| vllm / vllm_kunlun | 0.21.0 / 0.21.0.dev0 (also revalidated on 0.25.1) | `KunlunPlatform: device_type="cuda", dist_backend="nccl"` | +| megatron-core / transformers / triton | 0.16.1 / 5.9.0 / 3.5.0 | triton runs on the xmlir backend | +| XPU kernels | xpu_flash_attn, xpu_fla, kunlun_ops, xspeedgate_ops, xformers | Provided by the vendor container | + +**Notes**: +- The Kunlunxin driver, `torch_xmlir`, and the XPU-enabled `torch` / `vllm_kunlun` builds are provided by Kunlunxin (typically via the vendor container image). Twinkle does not install them. +- `vllm_kunlun` and its dependencies come from Kunlunxin; some kernels currently have dtype restrictions (see [Known Limitations](#known-limitations)). + +## Supported Hardware + +- Kunlunxin P800 (OAM), verified on an 8-card host + +## Installation Steps + +### 1. Prepare the XPU environment + +Use the Kunlunxin-provided container image (or a host with the Kunlunxin driver, `torch_xmlir`, and the XPU builds of `torch` / `vllm` / `vllm_kunlun` installed). Twinkle does not manage these components. + +### 2. Install Twinkle + +Install from source with `--no-deps` so pip does not attempt to replace the vendor's `torch` / `vllm` builds: + +```bash +git clone https://github.com/modelscope/twinkle.git +cd twinkle +pip install -e . --no-deps +pip install pyzmq +# If the container is missing debug/runtime helpers: +pip install h5py prettytable func_timeout ray redis +``` + +### 3. Verify Installation + +Create test script `verify_xpu.py`: + +```python +import torch +import twinkle # triggers ensure_xpu_compat() +from twinkle.utils.platforms import Platform + +print(f"PyTorch version: {torch.__version__}") +print(f"CUDA (XPU) available: {torch.cuda.is_available()}") +print(f"Device count: {torch.cuda.device_count()}") +print(f"Detected platform: {Platform.get_platform().__name__}") + +if torch.cuda.is_available(): + x = torch.randn(3, 3).cuda() + y = torch.randn(3, 3).cuda() + print(f"XPU computation test passed: {(x + y).shape}") +``` + +Run verification: + +```bash +python verify_xpu.py +``` + +A successful run reports `Detected platform: XPU`, `CUDA (XPU) available: True`, and the correct device count. Platform detection keys off the presence of `xpu-smi` on `PATH`. + +## Verified Capabilities + +The following have been validated end-to-end on Kunlunxin P800 (single-host, 8-card): + +| Capability | Backend | Status | Notes | +|---|---|---|---| +| FSDP2 LoRA SFT | transformers | ✅ Verified | Single-card, loss converges as on GPU | +| vLLM sampling (TP=1) | vLLM | ✅ Verified | Requires `enforce_eager=True`; text models | +| vLLM sampling (TP=2) | vLLM | ✅ Verified | Tensor parallel | +| GRPO (colocate) | native_fsdp | ✅ Verified | Full-weight sync (`merge_and_sync=True`); relay ~0.7 GB/s | +| Megatron LoRA (TP=2) | Megatron | ✅ Verified | Zero Twinkle code change; cuda-alike path | +| Weight sync (train↔sample) | XCCL | ✅ Verified | Via `XCCLCheckpointEngine`, single-host relay | + +## Known Limitations + +These are constraints of the current Kunlunxin XPU stack (vendor kernels / `vllm_kunlun`), not of Twinkle logic: + +- **LoRA sampling is not usable**: `vLLMSampler` with `enable_lora=True` cannot start. The underlying LoRA kernels (`bgmv_shrink_cluster`, `sgmv_expand_sdnn`, `sgmv_expand_slice`) are fp16-only on the 0.21 stack. Use full-weight sync (`merge_and_sync=True`) for RL. (On the 0.25.1 stack the vendor added a bf16 adapter layer, but other LoRA-path issues remain.) +- **fp16 required for some models**: GDN FLA kernels are fp16-only; Qwen3.5 inference must run in float16. +- **CUDA graph disabled**: Use `enforce_eager=True` for vLLM; graph capture on XPU is not yet validated. +- **Multimodal (Qwen-VL) blocked**: The ViT SDPA kernel triggers a hardware-level `kl3ChannelCheckErrors` / `noc idle timeout` that requires a physical card reset. Use text-only models. +- **Single-host only**: `XCCLCheckpointEngine` multi-host (inter/intra-node) broadcast is implemented but not yet validated; only single-host relay mode is verified. +- **`logprobs` on vllm_kunlun 0.25.1**: The logprobs path can return NaN/uninitialized top-logprob token ids, causing tokenizer `OverflowError`. Reported to the vendor. + +## Platform Internals (for reference) + +The XPU adaptation is concentrated in a few files: + +- `src/twinkle/utils/platforms/xpu.py` — `XPU` platform (inherits `GPU`); `ensure_xpu_compat()` neutralizes the native Intel-XPU stub in `torch.xpu`; a device-UUID fallback chain (`current_platform` → `xpu-smi -q` Bus Id → `xpu-smi -L` UUID → sha1) for vLLM. +- `src/twinkle/utils/platforms/base.py` — detects XPU via `xpu-smi` on `PATH`. +- `src/twinkle/infra/_ray/ray_helper.py` — registers the cuda-alike GPU count through `ray.init(num_gpus=N)` (Ray cannot autodetect XPUs). Remote workers started via `ray start` must pass `--num-gpus`. +- `src/twinkle/checkpoint_engine/xpu_checkpoint_engine.py` — `XCCLCheckpointEngine` (see the Checkpoint Engine docs). + +## Reference Resources + +- [vLLM-Kunlun](https://github.com/baidu/vLLM-Kunlun) +- [Twinkle GitHub](https://github.com/modelscope/twinkle) +- [Twinkle Documentation](https://twinkle.readthedocs.io/) + +## Getting Help + +If you encounter issues during use: + +1. **Submit an Issue**: [Twinkle GitHub Issues](https://github.com/modelscope/twinkle/issues) +2. **Vendor stack issues** (LoRA kernels, logprobs, multimodal): report to Kunlunxin / vLLM-Kunlun. diff --git a/docs/source_en/index.rst b/docs/source_en/index.rst index 0a00efdd1..d34dbced6 100644 --- a/docs/source_en/index.rst +++ b/docs/source_en/index.rst @@ -13,6 +13,7 @@ Twinkle DOCUMENTATION Usage Guide/Installation.md Usage Guide/Server and Client/index.rst Usage Guide/NPU-Support.md + Usage Guide/XPU-Support.md Usage Guide/Train-as-a-Service.md Usage Guide/Agentic-RL-Deployment-and-Training.md Usage Guide/Agentic-Evaluator.md diff --git a/docs/source_zh/index.rst b/docs/source_zh/index.rst index 96ac5a05a..da399ecfd 100644 --- a/docs/source_zh/index.rst +++ b/docs/source_zh/index.rst @@ -13,6 +13,7 @@ Twinkle DOCUMENTATION 使用指引/安装.md 使用指引/服务端和客户端/index.rst 使用指引/NPU的支持.md + 使用指引/XPU的支持.md 使用指引/训练服务.md 使用指引/Agentic RL部署与训练.md 使用指引/Agentic评测.md diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/XPU\347\232\204\346\224\257\346\214\201.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/XPU\347\232\204\346\224\257\346\214\201.md" new file mode 100644 index 000000000..0df44c935 --- /dev/null +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/XPU\347\232\204\346\224\257\346\214\201.md" @@ -0,0 +1,126 @@ +# XPU(昆仑芯)开箱指南 + +本文档介绍如何在昆仑芯 XPU(P800 系列)环境下安装和使用 Twinkle 框架。 + +## XPU 支持原理 + +与昇腾 NPU 不同,昆仑芯 XPU 在 Twinkle 中**不引入独立的设备类型**,而是依赖 `torch_xmlir` 在 import 阶段对 `torch.cuda` 符号进行重写: + +- `torch.cuda.is_available()` 返回 `True`,设备表现为 `cuda:N` +- `libbkcl.so`(BKCL)以 `nccl` 后端(XCCL)挂载到 PyTorch +- `CUDA_VISIBLE_DEVICES` 由 XPU runtime 接管 + +得益于这种 "cuda-alike" 路线,Twinkle 的 GPU 代码路径绝大部分可原生复用,仅需适配少数符号重写未覆盖的 CUDA 专属假设。因此 `XPU` 平台**继承自 `GPU`**(`src/twinkle/utils/platforms/xpu.py`),而非像 NPU 那样是完全独立的设备后端。 + +## 环境要求 + +适配在厂商容器镜像(`kylin_v11-vllm_021_swift:*`)内验证,已确认的组件矩阵如下: + +| 组件 | 版本 | 说明 | +|---|---|---| +| OS / Python | Kylin V11 / 3.10.14 | 满足 Twinkle `>=3.10` | +| torch | 2.9.0 | 被 `torch_xmlir` 符号重写 | +| vllm / vllm_kunlun | 0.21.0 / 0.21.0.dev0(0.25.1 亦复验通过) | `KunlunPlatform: device_type="cuda", dist_backend="nccl"` | +| megatron-core / transformers / triton | 0.16.1 / 5.9.0 / 3.5.0 | triton 走 xmlir 后端 | +| XPU kernels | xpu_flash_attn、xpu_fla、kunlun_ops、xspeedgate_ops、xformers | 由厂商容器提供 | + +**说明**: +- 昆仑芯驱动、`torch_xmlir` 以及 XPU 版的 `torch` / `vllm_kunlun` 由昆仑芯提供(通常通过厂商容器镜像),Twinkle 不负责安装。 +- `vllm_kunlun` 及其依赖来自昆仑芯;部分算子当前存在 dtype 限制(见[已知限制](#已知限制))。 + +## 支持的硬件 + +- 昆仑芯 P800(OAM),已在单机 8 卡环境验证 + +## 安装步骤 + +### 1. 准备 XPU 环境 + +使用昆仑芯提供的容器镜像(或已安装昆仑芯驱动、`torch_xmlir` 及 XPU 版 `torch` / `vllm` / `vllm_kunlun` 的宿主机)。这些组件不由 Twinkle 管理。 + +### 2. 安装 Twinkle + +从源码以 `--no-deps` 方式安装,避免 pip 覆盖厂商的 `torch` / `vllm` 构建: + +```bash +git clone https://github.com/modelscope/twinkle.git +cd twinkle +pip install -e . --no-deps +pip install pyzmq +# 若容器缺少调试/运行时辅助依赖: +pip install h5py prettytable func_timeout ray redis +``` + +### 3. 验证安装 + +创建测试脚本 `verify_xpu.py`: + +```python +import torch +import twinkle # 触发 ensure_xpu_compat() +from twinkle.utils.platforms import Platform + +print(f"PyTorch version: {torch.__version__}") +print(f"CUDA (XPU) available: {torch.cuda.is_available()}") +print(f"Device count: {torch.cuda.device_count()}") +print(f"Detected platform: {Platform.get_platform().__name__}") + +if torch.cuda.is_available(): + x = torch.randn(3, 3).cuda() + y = torch.randn(3, 3).cuda() + print(f"XPU computation test passed: {(x + y).shape}") +``` + +运行验证: + +```bash +python verify_xpu.py +``` + +运行成功时应输出 `Detected platform: XPU`、`CUDA (XPU) available: True` 以及正确的设备数量。平台检测依据 `PATH` 中是否存在 `xpu-smi`。 + +## 已验证能力 + +以下能力已在昆仑芯 P800(单机 8 卡)上端到端验证: + +| 能力 | 后端 | 状态 | 说明 | +|---|---|---|---| +| FSDP2 LoRA SFT | transformers | ✅ 已验证 | 单卡,loss 收敛与 GPU 一致 | +| vLLM 采样(TP=1) | vLLM | ✅ 已验证 | 需 `enforce_eager=True`;纯文本模型 | +| vLLM 采样(TP=2) | vLLM | ✅ 已验证 | 张量并行 | +| GRPO(colocate) | native_fsdp | ✅ 已验证 | 全量权重同步(`merge_and_sync=True`);relay ~0.7 GB/s | +| Megatron LoRA(TP=2) | Megatron | ✅ 已验证 | Twinkle 侧零改动;cuda-alike 路径 | +| 权重同步(训练↔采样) | XCCL | ✅ 已验证 | 经 `XCCLCheckpointEngine`,单机 relay | + +## 已知限制 + +以下为当前昆仑芯 XPU 栈(厂商 kernel / `vllm_kunlun`)的限制,而非 Twinkle 逻辑问题: + +- **LoRA 采样不可用**:`enable_lora=True` 的 `vLLMSampler` 无法启动。底层 LoRA 算子(`bgmv_shrink_cluster`、`sgmv_expand_sdnn`、`sgmv_expand_slice`)在 0.21 栈上仅支持 fp16。RL 请使用全量权重同步(`merge_and_sync=True`)。(0.25.1 栈厂商新增了 bf16 适配层,但 LoRA 路径仍有其他问题。) +- **部分模型需 fp16**:GDN FLA 算子仅支持 fp16,Qwen3.5 推理必须用 float16。 +- **CUDA graph 关闭**:vLLM 需 `enforce_eager=True`;XPU 上的 graph 捕获尚未验证。 +- **多模态(Qwen-VL)阻塞**:ViT SDPA 算子触发硬件级 `kl3ChannelCheckErrors` / `noc idle timeout`,需物理重置卡。请使用纯文本模型。 +- **仅支持单机**:`XCCLCheckpointEngine` 的多机(节点间/节点内)广播已实现但尚未验证,当前仅验证单机 relay 模式。 +- **vllm_kunlun 0.25.1 的 `logprobs`**:logprobs 路径可能返回 NaN/未初始化的 top-logprob token id,导致 tokenizer `OverflowError`。已反馈厂商。 + +## 平台内部实现(参考) + +XPU 适配集中在少数文件: + +- `src/twinkle/utils/platforms/xpu.py` — `XPU` 平台(继承 `GPU`);`ensure_xpu_compat()` 桩化 `torch.xpu` 中的原生 Intel-XPU 桩;为 vLLM 提供设备 UUID 降级链(`current_platform` → `xpu-smi -q` Bus Id → `xpu-smi -L` UUID → sha1)。 +- `src/twinkle/utils/platforms/base.py` — 通过 `PATH` 中的 `xpu-smi` 检测 XPU。 +- `src/twinkle/infra/_ray/ray_helper.py` — 经 `ray.init(num_gpus=N)` 注册 cuda-alike 的 GPU 数量(Ray 无法自动发现 XPU)。经 `ray start` 启动的远程 worker 需手动传 `--num-gpus`。 +- `src/twinkle/checkpoint_engine/xpu_checkpoint_engine.py` — `XCCLCheckpointEngine`(详见检查点引擎文档)。 + +## 参考资源 + +- [vLLM-Kunlun](https://github.com/baidu/vLLM-Kunlun) +- [Twinkle GitHub](https://github.com/modelscope/twinkle) +- [Twinkle 文档](https://twinkle.readthedocs.io/) + +## 获取帮助 + +如果您在使用过程中遇到问题: + +1. **提交 Issue**:[Twinkle GitHub Issues](https://github.com/modelscope/twinkle/issues) +2. **厂商栈问题**(LoRA 算子、logprobs、多模态):请反馈昆仑芯 / vLLM-Kunlun。 diff --git "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\256\211\350\243\205.md" "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\256\211\350\243\205.md" index dd9d9a18a..6d9fd9468 100644 --- "a/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\256\211\350\243\205.md" +++ "b/docs/source_zh/\344\275\277\347\224\250\346\214\207\345\274\225/\345\256\211\350\243\205.md" @@ -53,5 +53,6 @@ sh INSTALL_MEGATRON.sh | GPU A10/A100/H100/RTX系列等 | | | GPU T4/V100等 | 不支持bfloat16、Flash-Attention | | Ascend NPU | 部分算子不支持 | +| 昆仑芯 XPU(P800) | 经 cuda-alike 路线支持;部分算子仅 fp16,详见 [XPU 的支持](XPU的支持.md) | | PPU | 支持 | | CPU | 支持dataset、dataloader等部分组件 | diff --git "a/docs/source_zh/\347\273\204\344\273\266/\346\243\200\346\237\245\347\202\271\345\274\225\346\223\216/CheckpointEngine.md" "b/docs/source_zh/\347\273\204\344\273\266/\346\243\200\346\237\245\347\202\271\345\274\225\346\223\216/CheckpointEngine.md" index 338be10db..0d2064036 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/\346\243\200\346\237\245\347\202\271\345\274\225\346\223\216/CheckpointEngine.md" +++ "b/docs/source_zh/\347\273\204\344\273\266/\346\243\200\346\237\245\347\202\271\345\274\225\346\223\216/CheckpointEngine.md" @@ -39,7 +39,7 @@ class CheckpointEngine(ABC): ## 可用的检查点引擎 -Twinkle 提供了两种检查点引擎实现: +Twinkle 提供了三种检查点引擎实现: ### NCCLCheckpointEngine @@ -61,10 +61,21 @@ Twinkle 提供了两种检查点引擎实现: 详见: [HCCLCheckpointEngine](HCCLCheckpointEngine.md) +### XCCLCheckpointEngine + +使用 BKCL(XCCL)在昆仑芯 XPU 上进行权重传输的检查点引擎。 + +- XPU 支持: 在 cuda-alike XPU 路线上作为 NCCLCheckpointEngine 的直替 +- relay 兜底: 通过 stateless `ProcessGroupXCCL` 加 socket 中继处理重复的 local device index +- 兼容接口: 继承 `NCCLCheckpointEngine`,复用分桶与元数据握手 + +详见: [XCCLCheckpointEngine](XCCLCheckpointEngine.md) + ## 如何选择 - **NCCLCheckpointEngine**: 适用于 GPU 环境,提供最高的传输性能 - **HCCLCheckpointEngine**: 适用于昇腾 NPU 环境 +- **XCCLCheckpointEngine**: 适用于昆仑芯 XPU 环境(自动选择) > 检查点引擎是 RLHF 训练基础设施的关键组件,确保训练器和采样器使用一致的模型权重。 > 目前的同步分为merge_and_sync=True/False两种情况,为True时将lora合并仅基模并同步,为False时仅同步lora权重。另外,多租户直接附加lora文件到vLLM上,在merge_and_sync=False,或使用多租户时, diff --git "a/docs/source_zh/\347\273\204\344\273\266/\346\243\200\346\237\245\347\202\271\345\274\225\346\223\216/XCCLCheckpointEngine.md" "b/docs/source_zh/\347\273\204\344\273\266/\346\243\200\346\237\245\347\202\271\345\274\225\346\223\216/XCCLCheckpointEngine.md" new file mode 100644 index 000000000..5cbe080c5 --- /dev/null +++ "b/docs/source_zh/\347\273\204\344\273\266/\346\243\200\346\237\245\347\202\271\345\274\225\346\223\216/XCCLCheckpointEngine.md" @@ -0,0 +1,37 @@ +# XCCLCheckpointEngine + +面向昆仑芯 XPU 的检查点引擎,经 BKCL(以 `nccl`/XCCL 后端挂载到 PyTorch)进行权重传输。 + +## 使用示例 + +```python +from twinkle.checkpoint_engine import XCCLCheckpointEngine + +engine = XCCLCheckpointEngine(bucket_size=512<<20) +# 使用方式与 NCCLCheckpointEngine 相同 +``` + +在 XPU 上该引擎由 `CheckpointEngineManager` / `CheckpointEngineMixin` 自动选择,通常无需手动构造。 + +## 特性 + +- **NCCL 直替**:继承 `NCCLCheckpointEngine`,分桶、ZMQ 元数据握手、双缓冲等逻辑原样复用。 +- **XCCL 组 + relay 兜底**:将直接构造 `ProcessGroupNCCL`(当两个 rank 共享同一 local device index 时会死锁)替换为 stateless 的 `ProcessGroupXCCL`,并对冲突 rank 提供 relay 路径。 +- **基于 store 的初始化 barrier**:就绪 barrier 走 TCPStore,而非设备集合通信。 + +## 为何需要独立引擎 + +BKCL 按 **local device index** 识别 rank 设备,且不做 `CUDA_VISIBLE_DEVICES` → 物理卡的翻译。在单机上分离部署 trainer/sampler 时,两侧可见设备均为 local `0`,产生重复的 `(host, index)` 对。此时直接构造 `ProcessGroupNCCL` 无法成环,首次 broadcast 死锁(600s WorkXCCL 超时)。 + +`XCCLCheckpointEngine` 经 store 交换 `(hostname, local_device_index)`,将冲突 rank 排除出 XCCL 组,并由最小 rank 的 XCCL 成员经直连 TCP socket 中继(relay)。无冲突时使用平铺的 stateless `ProcessGroupXCCL`。 + +## 适用场景 + +- 昆仑芯 XPU 上的权重同步(训练 ↔ 采样) +- colocate / 分离式 RL(GRPO)下的全量权重同步(`merge_and_sync=True`) + +## 限制 + +- 多机(节点间/节点内)广播已实现但尚未验证,当前仅验证单机 relay 模式。 + +> 在昆仑芯 XPU 上权重同步请使用 `merge_and_sync=True`;LoRA 采样当前受厂商算子 dtype 限制而阻塞(详见 [XPU 的支持](../../使用指引/XPU的支持.md) 指南)。 diff --git "a/docs/source_zh/\347\273\204\344\273\266/\346\243\200\346\237\245\347\202\271\345\274\225\346\223\216/index.rst" "b/docs/source_zh/\347\273\204\344\273\266/\346\243\200\346\237\245\347\202\271\345\274\225\346\223\216/index.rst" index 996ddf2bc..b0cb9c113 100644 --- "a/docs/source_zh/\347\273\204\344\273\266/\346\243\200\346\237\245\347\202\271\345\274\225\346\223\216/index.rst" +++ "b/docs/source_zh/\347\273\204\344\273\266/\346\243\200\346\237\245\347\202\271\345\274\225\346\223\216/index.rst" @@ -6,3 +6,4 @@ CheckpointEngine.md NCCLCheckpointEngine.md HCCLCheckpointEngine.md + XCCLCheckpointEngine.md diff --git a/src/twinkle/checkpoint_engine/__init__.py b/src/twinkle/checkpoint_engine/__init__.py index 85febef20..3b9efebac 100644 --- a/src/twinkle/checkpoint_engine/__init__.py +++ b/src/twinkle/checkpoint_engine/__init__.py @@ -19,6 +19,7 @@ from .mixin import CheckpointEngineMixin # Import backend implementations to register them from .nccl_checkpoint_engine import NCCLCheckpointEngine +from .xpu_checkpoint_engine import XCCLCheckpointEngine __all__ = [ 'CheckpointEngine', @@ -26,5 +27,6 @@ 'CheckpointEngineManager', 'NCCLCheckpointEngine', 'HCCLCheckpointEngine', + 'XCCLCheckpointEngine', 'TensorMeta', ] diff --git a/src/twinkle/checkpoint_engine/manager.py b/src/twinkle/checkpoint_engine/manager.py index ba8d23bc5..8cfef1cc8 100644 --- a/src/twinkle/checkpoint_engine/manager.py +++ b/src/twinkle/checkpoint_engine/manager.py @@ -67,10 +67,14 @@ def __init__( @staticmethod def decide_backend_engine(platform: Optional[str] = None) -> 'CheckpointEngine': - if Platform.get_platform(platform).__name__ == 'GPU': + platform_name = Platform.get_platform(platform).__name__ + if platform_name in ('GPU', 'XPU'): + if platform_name == 'XPU': + from twinkle.checkpoint_engine import XCCLCheckpointEngine + return XCCLCheckpointEngine from twinkle.checkpoint_engine import NCCLCheckpointEngine return NCCLCheckpointEngine - elif Platform.get_platform(platform).__name__ == 'NPU': + elif platform_name == 'NPU': from twinkle.checkpoint_engine import HCCLCheckpointEngine return HCCLCheckpointEngine else: diff --git a/src/twinkle/checkpoint_engine/mixin.py b/src/twinkle/checkpoint_engine/mixin.py index 8dc15c926..ed3db6ff0 100644 --- a/src/twinkle/checkpoint_engine/mixin.py +++ b/src/twinkle/checkpoint_engine/mixin.py @@ -12,10 +12,14 @@ class CheckpointEngineMixin: def _get_or_create_checkpoint_engine(self) -> 'CheckpointEngine': """Get or create the checkpoint engine instance (lazy singleton).""" if self._checkpoint_engine is None: - if Platform.get_platform().__name__ == 'GPU': + platform_name = Platform.get_platform().__name__ + if platform_name == 'GPU': from twinkle.checkpoint_engine import NCCLCheckpointEngine self._checkpoint_engine = NCCLCheckpointEngine(self._bucket_size) - elif Platform.get_platform().__name__ == 'NPU': + elif platform_name == 'XPU': + from twinkle.checkpoint_engine import XCCLCheckpointEngine + self._checkpoint_engine = XCCLCheckpointEngine(self._bucket_size) + elif platform_name == 'NPU': from twinkle.checkpoint_engine import HCCLCheckpointEngine # Reusing HCCL communicator across sync steps avoids frequent diff --git a/src/twinkle/checkpoint_engine/xpu_checkpoint_engine.py b/src/twinkle/checkpoint_engine/xpu_checkpoint_engine.py new file mode 100644 index 000000000..576a2e4e5 --- /dev/null +++ b/src/twinkle/checkpoint_engine/xpu_checkpoint_engine.py @@ -0,0 +1,349 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +# Adapted from https://github.com/baidu/vLLM-Kunlun/blob/v0.21.0-dev/vllm_kunlun/distributed/py_kunlun_communicator.py +import pickle +import socket +import struct +import time +from datetime import timedelta +from typing import List, Optional, Sequence + +import torch +import torch.distributed as dist + +from twinkle import get_logger +from .nccl_checkpoint_engine import MasterMetadata, NCCLCheckpointEngine + +logger = get_logger() + +_RELAY_SETUP_TIMEOUT = 300.0 +_RELAY_IO_TIMEOUT = 600.0 +_STORE_WAIT_TIMEOUT = timedelta(seconds=300) + + +def _store_wait(store, keys: List[str], timeout: timedelta = _STORE_WAIT_TIMEOUT) -> None: + try: + store.wait(keys, timeout) + except TypeError: + # Older/newer signatures disagree on the timeout type. + store.wait(keys, timeout.total_seconds()) + + +def _recv_exact(sock, n: int) -> bytearray: + """Read exactly ``n`` bytes from ``sock`` into a fresh bytearray.""" + buf = bytearray(n) + mv = memoryview(buf) + got = 0 + while got < n: + k = sock.recv_into(mv[got:]) + if k == 0: + raise ConnectionError('relay socket closed by peer') + got += k + return buf + + +def _sock_send_tensor(sock, tensor: torch.Tensor) -> None: + """Stream ``tensor`` as raw bytes: [meta_len][meta][payload_len][payload].""" + meta = pickle.dumps((tensor.dtype, tuple(tensor.shape))) + cpu = tensor.detach() + if cpu.device.type != 'cpu': + cpu = cpu.cpu() + if not cpu.is_contiguous(): + cpu = cpu.contiguous() + cpu = cpu.reshape(-1) + sock.sendall(struct.pack('!Q', len(meta))) + sock.sendall(meta) + nbytes = cpu.numel() * cpu.element_size() + sock.sendall(struct.pack('!Q', nbytes)) + sock.sendall(cpu.view(torch.uint8).numpy()) + + +def _sock_recv_tensor(sock, device: torch.device) -> torch.Tensor: + """Receive a tensor sent by ``_sock_send_tensor`` onto ``device``.""" + (mlen,) = struct.unpack('!Q', _recv_exact(sock, 8)) + dtype, shape = pickle.loads(bytes(_recv_exact(sock, mlen))) + (nbytes,) = struct.unpack('!Q', _recv_exact(sock, 8)) + buf = _recv_exact(sock, nbytes) + flat = torch.frombuffer(buf, dtype=torch.uint8) + return flat.view(dtype).reshape(shape).to(device) + + +class _DoneWork: + """Already-completed work object (our collectives are synchronous).""" + + def wait(self): + return None + + +def _build_xccl_pg(store, name: str, ranks: Sequence[int], my_rank: int, device: torch.device): + """Assemble a stateless ProcessGroupXCCL over ``ranks`` (global ids). + + Mirrors ms-swift's ``_build_pg``: a plain ``ProcessGroup`` wrapper with a + ``PrefixStore`` per group and a registered CUSTOM XCCL backend. + """ + from torch._C._distributed_c10d import ProcessGroup + from torch.distributed import PrefixStore, ProcessGroupXCCL + + n = len(ranks) + local_rank = ranks.index(my_rank) + pstore = PrefixStore(name, store) + + pg = ProcessGroup(pstore, local_rank, n) + options = ProcessGroupXCCL.Options() + if hasattr(options, '_timeout'): + options._timeout = _RELAY_IO_TIMEOUT + backend = ProcessGroupXCCL(pstore, local_rank, n, options) + backend._set_sequence_number_for_group() + + backend_type = ProcessGroup.BackendType.CUSTOM + pg._set_default_backend(backend_type) + pg._register_backend(device, backend_type, backend) + return pg + + +class _FlatPG: + """Adapter: raw ``broadcast([tensor], opts)`` API over a ProcessGroupXCCL wrapper.""" + + def __init__(self, pg): + self.pg = pg + + def broadcast(self, tensors, opts=None, *args, **kwargs): + tensor = tensors[0] if isinstance(tensors, (list, tuple)) else tensors + src = getattr(opts, 'rootRank', 0) if opts is not None else 0 + self.pg.broadcast(tensor, src).wait() + return _DoneWork() + + +class _RelayPG: + """Relay-mode collective facade (single host, duplicate local device indices). + + Colliding ranks are kept out of the XCCL group and served through the + lowest-rank XCCL member over direct TCP sockets; the members themselves + keep the device-to-device XCCL path. ``src`` is always rank 0 (trainer + master) in the checkpoint engine's broadcast topology. + """ + + def __init__(self, rank: int, world_size: int, members: List[int], + excluded: List[int], pg, device: torch.device): + self.rank = rank + self.world_size = world_size + self.members = members + self.excluded = excluded + self.leader = members[0] if members else 0 + self.pg = pg # ProcessGroupXCCL wrapper among members (None when degenerate) + self.device = device + self._listen = None + self._conns = {} # leader: {excluded rank: socket} + self._sock = None # excluded: socket to the leader + + # -- direct-socket rendezvous (store-coordinated, tiny payloads) -- + def setup_sockets(self, store, prefix: str) -> None: + addr_key = f'{prefix}_relay_addr' + if not self.excluded: + return + if self.rank == self.leader: + srv = socket.socket() + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(('', 0)) + srv.listen(len(self.excluded)) + srv.settimeout(15) + store.set(addr_key, pickle.dumps((socket.gethostname(), srv.getsockname()[1]))) + conns = {} + deadline = time.time() + _RELAY_SETUP_TIMEOUT + while len(conns) < len(self.excluded): + try: + conn, _ = srv.accept() + except socket.timeout: + if time.time() > deadline: + raise RuntimeError('relay socket rendezvous timed out') + continue + conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + conn.settimeout(_RELAY_IO_TIMEOUT) + (peer,) = struct.unpack('!I', bytes(_recv_exact(conn, 4))) + conns[peer] = conn + self._listen = srv + self._conns = conns + elif self.rank in self.excluded: + host, port = pickle.loads(store.get(addr_key)) + deadline = time.time() + 60 + while True: + try: + s = socket.create_connection((host, port), timeout=10) + break + except OSError: + if time.time() > deadline: + raise RuntimeError('relay socket connect timed out') + time.sleep(0.5) + s.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + s.settimeout(_RELAY_IO_TIMEOUT) + s.sendall(struct.pack('!I', self.rank)) + self._sock = s + + def close(self) -> None: + if self._sock is not None: + try: + self._sock.close() + except OSError: + pass + self._sock = None + for c in self._conns.values(): + try: + c.close() + except OSError: + pass + self._conns = {} + if self._listen is not None: + try: + self._listen.close() + except OSError: + pass + self._listen = None + + # -- collective facade -- + def broadcast(self, tensors, opts=None, *args, **kwargs): + tensor = tensors[0] if isinstance(tensors, (list, tuple)) else tensors + src = getattr(opts, 'rootRank', 0) if opts is not None else 0 + if tensor.device != self.device: + tensor = tensor.to(self.device) + if not tensor.is_contiguous(): + tensor = tensor.contiguous() + + if self.rank in self.excluded: + # Excluded from XCCL: TCP only (src pushes, or leader forwards). + if self.rank == src: + _sock_send_tensor(self._sock, tensor) + # else: the sender (src or leader) pushes to us over the socket. + else: + tensor.copy_(_sock_recv_tensor(self._sock, self.device)) + return _DoneWork() + + # This rank is an XCCL member. + if src in self.excluded: + if self.rank == self.leader: + conn = self._conns.get(src) + tensor.copy_(_sock_recv_tensor(conn, self.device)) + if self.pg is not None: + self.pg.broadcast(tensor, 0).wait() # leader == members[0] + if self.rank == self.leader: + for x in self.excluded: + if x != src: + _sock_send_tensor(self._conns[x], tensor) + else: + if self.rank == src: + for x in self.excluded: + _sock_send_tensor(self._conns[x], tensor) + if self.pg is not None: + self.pg.broadcast(tensor, self.members.index(src)).wait() + return _DoneWork() + + +class XCCLCheckpointEngine(NCCLCheckpointEngine): + """NCCL-checkpoint-engine drop-in for Kunlunxin XPU (BKCL as 'nccl'). + + Replaces the direct ``ProcessGroupNCCL`` construction (which deadlocks + when two ranks share a local device index, see module docstring) with a + stateless ``ProcessGroupXCCL`` plus a relay fallback. Everything else + (bucketing, ZMQ metadata, double buffering) is inherited unchanged. + """ + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self._relay: Optional[_RelayPG] = None + + def _init_xccl_pg(self) -> None: + store = self._store + prefix = self.group_name + + # Exchange (hostname, local device index) per rank through the store + # (a round-trip, not a device collective -- safe on every rank). + devidx = torch.cuda.current_device() + store.set(f'{prefix}_dev_{self.rank}', f'{socket.gethostname()}|{devidx}') + keys = [f'{prefix}_dev_{r}' for r in range(self.world_size)] + _store_wait(store, keys) + info = [store.get(k).decode().split('|') for k in keys] + devs = [(h, int(d)) for h, d in info] + + device = torch.device('cuda', devidx) + + # Single host assumed for the flat case; membership by first + # occurrence of each (host, local device index) pair. + first_owner = {} + excluded = set() + for r, hd in enumerate(devs): + if hd in first_owner: + excluded.add(r) + else: + first_owner[hd] = r + members = [r for r in range(self.world_size) if r not in excluded] + + if not excluded: + self._pg = _FlatPG(_build_xccl_pg(store, f'{prefix}_ws', + list(range(self.world_size)), self.rank, device)) + logger.info(f'XCCLCheckpointEngine: flat XCCL group, world_size={self.world_size}') + else: + pg = None + if len(members) > 1: + pg = _build_xccl_pg(store, f'{prefix}_relay', members, self.rank, device) + self._relay = _RelayPG(self.rank, self.world_size, members, sorted(excluded), pg, device) + self._relay.setup_sockets(store, prefix) + self._pg = self._relay + logger.info(f'XCCLCheckpointEngine: relay mode, members={members}, excluded={sorted(excluded)}') + + def _store_barrier(self) -> None: + prefix = self.group_name + store = self._store + store.set(f'{prefix}_bar_{self.rank}', '1') + _store_wait(store, [f'{prefix}_bar_{r}' for r in range(self.world_size)]) + + def init_process_group(self, rank: int, world_size: int, master_metadata: MasterMetadata): + """Initialize the weight-sync process group (XPU flavour). + + Same rendezvous as the NCCL engine (dedicated TCPStore hosted by the + master), but the XCCL group construction/relay decision happens here + and the readiness barrier runs through the store instead of a device + collective. + """ + # Non-participating trainer ranks: record rank and return. + if rank < 0: + self.rank = rank + self.world_size = world_size + self._group_initialized = True + return + + # Fast path: group already initialized, skip all setup. + if self._group_initialized and not self.rebuild_group: + return + + if self._pg is None: + self.rank = rank + self.world_size = world_size + + is_store_master = (rank == 0) + self._store = dist.TCPStore( + host_name=master_metadata.nccl_store_host, + port=master_metadata.nccl_store_port, + world_size=world_size, + is_master=is_store_master, + wait_for_workers=True, + ) + self._init_xccl_pg() + else: + assert self.rank == rank, f'rank {rank} != self.rank {self.rank}' + assert self.world_size == world_size, ( + f'world_size {world_size} != self.world_size {self.world_size}') + + # Receivers connect to master's ZMQ PUB server. + if self.rank > 0 and self.socket is None: + self._connect_zmq_client(master_metadata) + + # Store-based readiness barrier (no device collective during init). + self._store_barrier() + + self._group_initialized = True + logger.info(f'init_process_group: rank={self.rank}, world_size={self.world_size}') + + def finalize(self): + """Tear down relay sockets before the base cleanup.""" + if self._relay is not None: + self._relay.close() + self._relay = None + super().finalize() diff --git a/src/twinkle/infra/_ray/ray_helper.py b/src/twinkle/infra/_ray/ray_helper.py index ffd4e1a42..f06132782 100644 --- a/src/twinkle/infra/_ray/ray_helper.py +++ b/src/twinkle/infra/_ray/ray_helper.py @@ -28,18 +28,32 @@ def _get_ray_custom_resources(device_groups: List[DeviceGroup]) -> Dict[str, flo # ResourceManager supports one accelerator type per run. Only NPU # currently needs an explicit Ray custom-resource registration. - if device_types != {'NPU'}: - return {} + if device_types == {'NPU'}: + try: + import torch - try: - import torch + npu = getattr(torch, 'npu', None) + npu_count = npu.device_count() if npu is not None and npu.is_available() else 0 + except (ImportError, AttributeError, RuntimeError): + return {} + + return {'NPU': float(npu_count)} if npu_count > 0 else {} + + # Kunlunxin XPU: Ray cannot autodetect the GPUs (no nvidia-smi on the + # XPU runtime), so register the cuda-alike GPU count explicitly for + # local ray.init. Remote workers started via `ray start` must set + # --num-gpus manually. + if device_types == {'GPU'} and Platform.get_platform().__name__ == 'XPU': + try: + import torch + + gpu_count = torch.cuda.device_count() if torch.cuda.is_available() else 0 + except (ImportError, AttributeError, RuntimeError): + return {} - npu = getattr(torch, 'npu', None) - npu_count = npu.device_count() if npu is not None and npu.is_available() else 0 - except (ImportError, AttributeError, RuntimeError): - return {} + return {'GPU': float(gpu_count)} if gpu_count > 0 else {} - return {'NPU': float(npu_count)} if npu_count > 0 else {} + return {} @staticmethod def init_registry(): @@ -97,7 +111,11 @@ def initialize(nproc_per_node: int, ncpu_proc_per_node: int, device_groups: List RayHelper.device_groups = device_groups if not RayHelper.ray_inited(): resources = RayHelper._get_ray_custom_resources(device_groups) - ray.init(ignore_reinit_error=True, resources=resources or None) + # XPU registers the cuda-alike 'GPU' count, but 'GPU' is one of Ray's + # default resources and ray.init() rejects it via `resources=`; route + # it through `num_gpus` instead. + num_gpus = int(resources.pop('GPU')) if 'GPU' in resources else None + ray.init(ignore_reinit_error=True, resources=resources or None, num_gpus=num_gpus) if RayHelper.resource_manager is None: # Resource manager initializes only once in the pipeline process. diff --git a/src/twinkle/utils/platforms/__init__.py b/src/twinkle/utils/platforms/__init__.py index ea327f840..02d7c2298 100644 --- a/src/twinkle/utils/platforms/__init__.py +++ b/src/twinkle/utils/platforms/__init__.py @@ -2,3 +2,6 @@ from .gpu import GPU from .mps import MPS, is_mps_available from .npu import NPU, ensure_hccl_socket_env, ensure_npu_backend +from .xpu import XPU, ensure_xpu_compat + +ensure_xpu_compat() diff --git a/src/twinkle/utils/platforms/base.py b/src/twinkle/utils/platforms/base.py index 483c725a3..49d666d3a 100644 --- a/src/twinkle/utils/platforms/base.py +++ b/src/twinkle/utils/platforms/base.py @@ -17,7 +17,7 @@ def device_prefix(platform: str = None) -> str: @staticmethod def get_platform_names() -> List[str]: - return ['GPU', 'NPU', 'MPS'] + return ['GPU', 'NPU', 'XPU', 'MPS'] @staticmethod def get_platform(platform: str = None) -> Type['Platform']: @@ -27,6 +27,10 @@ def get_platform(platform: str = None) -> Type['Platform']: from .npu import NPU, ensure_npu_backend ensure_npu_backend() return NPU + elif shutil.which('xpu-smi'): + from .xpu import XPU, ensure_xpu_compat + ensure_xpu_compat() + return XPU elif shutil.which('nvidia-smi'): from .gpu import GPU return GPU @@ -43,6 +47,10 @@ def get_platform(platform: str = None) -> Type['Platform']: from .npu import NPU, ensure_npu_backend ensure_npu_backend() return NPU + elif platform.upper() == 'XPU': + from .xpu import XPU, ensure_xpu_compat + ensure_xpu_compat() + return XPU elif platform.upper() == 'MPS': from .mps import MPS return MPS diff --git a/src/twinkle/utils/platforms/xpu.py b/src/twinkle/utils/platforms/xpu.py new file mode 100644 index 000000000..f45ff7df0 --- /dev/null +++ b/src/twinkle/utils/platforms/xpu.py @@ -0,0 +1,112 @@ +# Copyright (c) ModelScope Contributors. All rights reserved. +import hashlib +import os +import re +import socket +import subprocess +from typing import Optional + +from .gpu import GPU + + +def _get_xpu_bus_id_from_xpu_smi(device_id: int) -> Optional[str]: + """Get XPU Bus-Id from `xpu-smi -q` output. + """ + try: + output = subprocess.check_output( + ['xpu-smi', '-q'], + text=True, + stderr=subprocess.STDOUT, + timeout=5, + ) + except Exception: + return None + + pattern = re.compile( + r'^XPU\s+(\d+).*?Bus Id\s*:\s*([0-9A-Fa-f]{4}:[0-9A-Fa-f]{2}:[0-9A-Fa-f]{2}\.[0-9A-Fa-f])', + re.MULTILINE | re.DOTALL, + ) + for match in pattern.finditer(output): + if int(match.group(1)) == device_id: + return match.group(2).lower() + return None + + +def _get_xpu_uuid_from_xpu_smi(device_id: int) -> Optional[str]: + """Get XPU UUID from `xpu-smi -L` output. + """ + try: + output = subprocess.check_output( + ['xpu-smi', '-L'], + text=True, + stderr=subprocess.STDOUT, + timeout=5, + ) + except Exception: + return None + + pattern = re.compile(r'^XPU\s+(\d+):\s+\S.*\(UUID:\s*(\S+)\)', re.MULTILINE) + for match in pattern.finditer(output): + if int(match.group(1)) == device_id: + return match.group(2) + return None + + +def ensure_xpu_compat() -> None: + """Neutralize the native Intel-XPU stub shipped inside torch. + """ + import torch + + xp = getattr(torch, 'xpu', None) + if xp is None: + return + try: + if xp.is_available(): + return # a real (Intel) XPU runtime is present; nothing to stub + except Exception: + return + if not getattr(xp.get_device_name, '_twinkle_xpu_stub', False): + def _get_device_name(*args, **kwargs): # noqa: E306 + return 'Kunlunxin XPU' + _get_device_name._twinkle_xpu_stub = True + xp.get_device_name = _get_device_name + + +class XPU(GPU): + + @staticmethod + def visible_device_env(): + # Kunlunxin XPU runtime takes over CUDA_VISIBLE_DEVICES semantics. + return 'CUDA_VISIBLE_DEVICES' + + @staticmethod + def device_prefix(): + # torch_xmlir symbol-rewrites torch.cuda to XPU, so tensors live on + # 'cuda:N' devices from PyTorch's point of view. + return 'cuda' + + @staticmethod + def get_local_device(idx, **kwargs) -> str: + return f'cuda:{idx}' + + @staticmethod + def device_backend(platform: str = None): + # BKCL is mounted into torch as the 'nccl' backend (XCCL). + return 'nccl' + + @staticmethod + def get_vllm_device_uuid(device_id: int = 0) -> str: + from vllm.platforms import current_platform + try: + return current_platform.get_device_uuid(device_id) + except NotImplementedError: + bus_id = _get_xpu_bus_id_from_xpu_smi(device_id) + if bus_id: + return bus_id + xpu_uuid = _get_xpu_uuid_from_xpu_smi(device_id) + if xpu_uuid: + return xpu_uuid + # Deterministic fallback so both sides compute the same socket name. + visible = os.environ.get(XPU.visible_device_env()) + raw = f'{socket.gethostname()}:{visible}:{device_id}' + return hashlib.sha1(raw.encode('utf-8')).hexdigest()[:16]