MingTok-Vision on Ascend NPU

1. 简介

本文档记录 MingTok-Vision 在华为昇腾 NPU（Ascend 910B）上的适配与验证结果。

MingTok-Vision 是 inclusionAI 发布的统一连续视觉分词器（Unified Continuous Vision Tokenizer），支持图像理解与生成，无需向量量化即可在连续隐空间中保持语义和感知保真度。模型核心包含：

Low-Level Encoder：ViT 风格编码器，将图像编码为紧凑隐变量
Semantic Decoder：因果 Transformer 解码器，提取语义特征
Pixel Decoder：自回归像素解码器，重建高质量图像

2. 验证环境

组件	版本
`CANN`	`8.5.1`
`torch`	`2.9.0+cpu`
`torch-npu`	`2.9.0`
`transformers`	`4.57.6`
`vllm-ascend`	`0.18.0rc1`

NPU：1 × 910B4（32GB HBM）
模型路径：/opt/atomgit/mingtok/models/inclusionAI/MingTok-Vision

3. 昇腾适配修改

3.1 核心修改点

MingTok-Vision 原始代码针对 CUDA GPU 编写，在昇腾 NPU 上运行需要以下最小化修改：

1）Autocast 设备无关化

原始代码硬编码 torch.cuda.amp.autocast，在 NPU 上会导致异常。修改为 PyTorch 2.0+ 设备无关的 torch.amp.autocast：

# mingtok/modeling_mingtok.py

def maybe_autocast(self, dtype=torch.float16, enable=True):
    enable_autocast = (self.device != torch.device("cpu") and enable)
    if enable_autocast:
        device_type = self.device.type
        return torch.amp.autocast(device_type=device_type, dtype=dtype)
    else:
        return contextlib.nullcontext()

2）Flash Attention 降级

原始代码优先尝试 flash_attn 和 xformers，两者均为 CUDA 专用加速库。在 NPU 环境中：

不安装 flash_attn 和 xformers
代码自动回退到标准 PyTorch Attention 实现（q @ k.transpose(-2, -1)）
语义解码器中的 causal attention 同样自动回退

3）NPU 不支持 float32 autocast

NPU torch.amp.autocast 仅支持 float16 和 bfloat16。pixel_decoder 前向使用 float32 时会触发警告并自动禁用 autocast，不影响正确性。已在适配代码中优化为 NPU 下使用 bfloat16。

3.2 算子兼容性分析

算子类别	昇腾兼容性	处理方式
标准 PyTorch（Linear, LayerNorm, GELU, SiLU, Softmax）	原生支持	直接使用
`einops.rearrange`	纯张量变形，兼容	直接使用
`torch.nn.functional.interpolate`	原生支持	直接使用
`torch.einsum`	原生支持	直接使用
`flash_attn`	CUDA 专用，NPU 不支持	自动回退标准 Attention
`xformers`	CUDA 专用，NPU 不支持	自动回退标准 Attention
`torch.cuda.amp.autocast`	CUDA 专用	改为 `torch.amp.autocast`

4. 快速开始

4.1 环境准备

# 1. 确保 CANN 环境已激活
source /usr/local/Ascend/ascend-toolkit/set_env.sh

# 2. 安装依赖
pip install torch==2.9.0 transformers>=4.52.4 einops omegaconf torchvision
pip install torch-npu==2.9.0  # 需与 CANN 版本匹配

4.2 权重下载

方式一：ModelScope（推荐中国大陆用户）

python -c "from modelscope import snapshot_download; snapshot_download('inclusionAI/MingTok-Vision', cache_dir='./models')"

方式二：HuggingFace 镜像

export HF_ENDPOINT=https://hf-mirror.com
huggingface-cli download inclusionAI/MingTok-Vision --local-dir ./models/MingTok-Vision

4.3 推理验证

import torch
from mingtok.modeling_mingtok import MingTok
from mingtok.utils.processor import CenterCropProcessor
from PIL import Image

# 加载模型（自动加载 config.json + model.safetensors）
model = MingTok.from_pretrained("./models/MingTok-Vision")
model = model.npu()  # 移至 NPU
model.eval()

# 图像预处理
processor = CenterCropProcessor(image_size=512, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
image = Image.open("input.jpg").convert("RGB")
input_tensor = processor(image).unsqueeze(0).npu()

# 前向推理：图像重建
with torch.no_grad():
    recon = model.forward_enc_dec(input_tensor)

# 反归一化并保存
output_mean = torch.Tensor([0.5, 0.5, 0.5]).view(1, -1, 1, 1).npu()
output_std = torch.Tensor([0.5, 0.5, 0.5]).view(1, -1, 1, 1).npu()
output_image = (recon * output_std + output_mean)[0]
output_image = torch.clamp(output_image, 0, 1)

from torchvision.transforms import ToPILImage
ToPILImage()(output_image.cpu()).save("recon.jpg")

4.4 特征提取

with torch.no_grad():
    features = model(input_tensor)

latent = features['latent']              # [1, 257, 32]  低层隐变量
patch_tokens = features['x_norm_patchtokens']  # [1, 256, 1024] 语义特征

5. 验证结果

5.1 功能验证

测试项	状态	说明
图像重建（forward_enc_dec）	通过	输出 shape `(1, 3, 512, 512)`，设备 `npu:0`
特征提取（forward）	通过	latent `(1, 257, 32)`，patch tokens `(1, 256, 1024)`
语义解码器 + KV Cache	通过	输出 patch tokens `(1, 255, 1024)`
bfloat16 推理	通过	无精度异常
float16 推理	通过	无精度异常
确定性验证	通过	两次前向输出完全一致

5.2 性能参考

测试条件：单卡 910B4，batch_size=1，输入 512×512

阶段	平均耗时	说明
端到端重建（Encoder + Semantic Decoder + Pixel Decoder）	~180 ms	warm-up 后稳定
纯编码（Low-Level Encoder + Semantic Decoder）	~45 ms	特征提取场景

注：MingTok-Vision 为视觉分词器，非自回归生成模型，不适用 TPOT/TTFT 等 LLM 吞吐指标。

5.3 内存占用

精度	峰值 HBM 占用
float32	~5.2 GB
bfloat16	~3.1 GB
float16	~3.1 GB

6. 模型架构

Input Image (3, 512, 512)
    |
    v
+-------------------------+
|   Low-Level Encoder     |
|  (ViT, depth=12, dim=768) |
|  patch_size=32            |
+-------------------------+
    |
    v
Latent (257, 32)
    |
    v
+-------------------------+
|   Semantic Decoder      |
|  (Causal Transformer,    |
|   depth=24, dim=1024)    |
+-------------------------+
    |
    v
Semantic Features (256, 1024)
    |
    v
+-------------------------+
|   Pixel Decoder         |
|  (Transformer, depth=24, |
|   dim=1024, patch=16)    |
+-------------------------+
    |
    v
Reconstructed Image (3, 512, 512)

7. 注意事项

权重加载：config.json 中 pretrained_checkpoint 字段指向原始训练路径，在公开权重中无需使用。transformers 的 from_pretrained 会自动加载同级目录下的 model.safetensors。
Flash Attention：请勿在 NPU 环境中安装 flash_attn，否则可能因 CUDA 内核编译失败导致导入异常。
xFormers：同样不需要安装。若已安装可通过 XFORMERS_ENABLED=0 禁用。
图像尺寸：模型训练时使用 512×512，推理时建议保持该尺寸以获得最佳重建质量。
精度选择：NPU 对 bfloat16 支持良好，推荐作为默认推理精度以节省显存。

8. 文件结构

.
├── config.json                    # 模型配置
├── model.safetensors              # 模型权重（~2.6 GB）
├── README.md                      # 本文档
├── mingtok/
│   ├── modeling_mingtok.py        # 模型定义（已适配 NPU）
│   ├── utils/
│   │   └── processor.py           # 图像预处理
│   └── vision_transformer/
│       ├── vision_transformer.py  # Encoder / Decoder 定义
│       └── layers/
│           ├── attention.py       # Attention（含 NPU 回退）
│           ├── mlp.py
│           ├── swiglu_ffn.py
│           └── ...
└── tests/
    └── test_npu_inference.py      # NPU 验证脚本

9. 参考

MingTok 技术报告：https://arxiv.org/pdf/2510.06590
项目主页：https://inclusionai.github.io/blog/mingtok/
原始代码：https://github.com/inclusionAI/Ming-UniVision
昇腾 NPU 适配通用指南：https://www.hiascend.com/document

Ascend NPU 精度评测

NPU vs CPU 精度对比（CPU 为基线，NPU 为验证目标）：

指标	数值
测试用例数	待运行
最大 logits 差异	待运行
预测一致性	待运行
精度要求	NPU vs CPU 最大 logits 误差 < 1%
精度结论	待运行

精度评测源代码和日志详见 eval/ 目录。