# npu-adapter

本项目主要是面向大模型昇腾适配过程中，针对场景的torch层算子不兼容，从GPU->NPU适配显存卸载，长序列切分等等能力的昇腾支持，加速整体迁移效率。当前覆盖面暂未包含所有场景，欢迎贡献代码，逐步的沉淀各个场景适配过程中的通用能力，构建更好用的昇腾生态。

## 项目信息

- **版本**: 0.1.0
- **Python要求**: >=3.11
- **主页**: https://gitcode.com/ascend-sact/npu-adapter

## 依赖项

- torch>=2.8.0
- torch_npu>=2.8.0
- yunchang>=0.6.3.post1

## 安装

```bash
pip install npu-adapter
```

## 快速开始

```python
import npu_adapter

# 适配到加速设备
npu_adapter.adapt_to_accelerator_device()

# 初始化分布式训练
npu_adapter.init_distributed_adapter()
```

## 功能模块

### 1. 设备检测与适配

#### has_npu()
检测NPU是否可用。

```python
from npu_adapter import has_npu

if has_npu():
    print("NPU is available")
else:
    print("NPU is not available")
```

#### adapt_to_accelerator_device()
适配到加速设备，打印当前支持的设备。

```python
from npu_adapter import adapt_to_accelerator_device

adapt_to_accelerator_device()
# 输出: load adapter(accelerator device, now support to run on Ascend NPU
# 或: load adapter(accelerator device, now support to run on Nvidia GPU
```

#### init_distributed_adapter()
初始化分布式训练适配，根据设备类型选择不同的后端。

```python
from npu_adapter import init_distributed_adapter

init_distributed_adapter()
# NPU: backend="cpu:gloo,npu:hccl"
# GPU: backend="nccl"
```

### 2. 确定性模式

#### deterministic_on()
启用确定性算法和随机种子设置。

```python
from npu_adapter import deterministic_on

deterministic_on()
```

### 3. 长序列注意力

#### get_longcontext_attention()
获取长序列注意力模块，基于[yunchang](https://github.com/feifeibear/long-context-attention)实现。

<p align="center">
  <img src="https://raw.gitcode.com/user-images/assets/8950477/770c8915-4535-4ecc-8a29-4d40d53a1d82/cb61d353-eb0f-4ce5-94dc-cd8ad12886c0_1038x613.png" alt="长序列注意力" width="500">
</p>

```python
from npu_adapter import get_longcontext_attention

# 自动根据设备类型选择合适的实现
long_context_attn = get_longcontext_attention()
# NPU: LongContextAttention(ring_impl_type="basic_npu", attn_type=AttnType.NPU)
# GPU: LongContextAttention(ring_impl_type="basic", attn_type=AttnType.FA3)
```

### 4. 编译器后端

#### get_compiler_backend()
获取编译器后端，支持昇腾的融合算子自动使能。

```python
import torch
from npu_adapter import get_compiler_backend

# NPU: torchair.get_npu_backend()
# GPU: "inductor"
model = torch.compile(model, backend=get_compiler_backend())
```

对于MindieSDBackend调用方式：
在入口脚本中，将transformer模块整体进行compile，可以通过如下方式使能：
```python
pipe = FluxPipeline.from_pretrained(...)
transformer = torch.compile(pipe.transformer, backend=MindieSDBackend())
setattr(pipe, "transformer", transformer)
```
也可以针对单个Module针对性使用：
```python
@torch.compile(backend=MindieSDBackend())
class FluxSingleTransformerBlock(nn.Module):
```
或者对forward函数使用：
```python
class FluxSingleTransformerBlock(nn.Module):
    @torch.compile(backend=MindieSDBackend())
    def forward(...):
```

### 5. 内存格式

#### contiguous_for_channels_last_3d_memory_format()
适配连续内存格式，NPU仅支持contiguous格式，GPU支持channels_last_3d格式。

```python
from npu_adapter import contiguous_for_channels_last_3d_memory_format

tensor = torch.randn(2, 3, 4, 5, 6)
result = contiguous_for_channels_last_3d_memory_format(tensor)
```

### 6. 激活函数

#### adapter_gelu()
GELU激活函数适配，NPU使用fast_gelu优化实现。

```python
from npu_adapter import adapter_gelu

x = torch.randn(10, 10)
result = adapter_gelu(x)
# NPU: torch_npu.fast_gelu(x)
# GPU: 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
```

### 7. 归一化

#### adapter_norm()
RMS归一化适配（类方法）。

```python
from npu_adapter import adapter_norm

class MyModule:
    Module:
        def __init__(self):
            self.weight = torch.randn(10)
            self.eps = 1e-5
        
        def forward(self, x):
            return adapter_norm(self, x)
```

#### adpater_rmsnorm()
RMS归一化适配（函数形式）。

```python
from npu_adapter import adpater_rmsnorm

x = torch.randn(10, 10)
weight = torch.randn(10)
eps = 1e-5
result = adpater_rmsnorm(x, weight, eps)
```

### 8. 旋转位置编码

#### apply_rotary_pos_emb_adapter()
应用旋转位置编码到查询和键张量，自动根据设备选择实现。

```python
from npu_adapter import apply_rotary_pos_emb_adapter

q = torch.randn(2, 4, 8, 16)
k = torch.randn(2, 4, 8, 16)
cos = torch.randn(2, 8, 16)
sin = torch.randn(2, 8, 16)

q_rotated, k_rotated = apply_rotary_pos_emb_adapter(q, k, cos, sin)
```

## 测试

项目包含完整的测试套件，覆盖NPU和GPU场景。

### 运行所有测试

```bash
python test/test_adapter.py
```

### 使用unittest模块运行

```bash
python -m unittest test.test_adapter -v
```

### 使用测试运行器

```bash
python test/run_tests.py
```

### 测试覆盖范围

- **工具函数测试**: NPU/GPU检测、设备适配、分布式初始化、确定性模式
- **长序列注意力测试**: NPU/GPU场景的长序列注意力获取
- **编译器后端测试**: NPU/GPU场景的编译器后端获取
- **内存格式操作测试**: 连续内存格式适配
- **张量操作测试**: GELU激活、RMS归一化
- **旋转位置编码测试**: 旋转位置嵌入应用
- **集成测试**: 完整工作流测试

## 贡献指南

欢迎贡献代码！请遵循以下步骤：

1. Fork本仓库
2. 创建特性分支 (`git checkout -b feature/AmazingFeature`)
3. 提交更改 (`git commit -m 'Add some AmazingFeature'`)
4. 推送到分支 (`git push origin feature/AmazingFeature`)
5. 开启Pull Request

## 许可证

本项目采用开源许可证，具体请查看LICENSE文件。

## 联系方式

- 项目主页: https://gitcode.com/ascend-sact/npu-adapter

## 致谢

- [yunchang](https://github.com/feifeibear/long-context-attention): 长序列上下文Attention计算
- [MindIE-SD](https://gitcode.com/Ascend/MindIE-SD): 昇腾推理后端
