ERNIE-Image 部署文档

Transformer Config:
  - hidden_size: 4096
  - num_layers: 36
  - num_attention_heads: 32
  - ffn_hidden_size: 12288
  - in_channels: 128 (latent)
  - out_channels: 128
  - patch_size: 1
  - text_in_dim: 3072
  - eps: 1e-06
  - qk_layernorm: True
  - rope_axes_dim: [32, 48, 48]
  - rope_theta: 256

parallel_config:
  sequence_parallel_size: 4  # 4个NPU并行
  ulysses_degree: 4          # Ulysses SP算法
  ring_degree: 1
  tensor_parallel_size: 1
  pipeline_parallel_size: 1
  data_parallel_size: 1

# 安装 modelscope
pip install modelscope

# 下载 ERNIE-Image 模型
python3 << 'EOF'
from modelscope import snapshot_download

model_dir = snapshot_download(
    'PaddlePaddle/ERNIE-Image',
    cache_dir='/opt/data/modelscope/hub'
)
print(f"Model downloaded to: {model_dir}")
EOF

/opt/data/modelscope/hub/models/PaddlePaddle/ERNIE-Image/
├── transformer/
│   ├── config.json
│   ├── diffusion_pytorch_model.safetensors
│   └── model.safetensors.index.json
├── vae/
│   ├── config.json
│   └── diffusion_pytorch_model.safetensors
├── text_encoder/
│   ├── config.json
│   ├── model.safetensors
│   └── tokenizer_config.json
├── scheduler/
│   └── scheduler_config.json

docker pull m.daocloud.io/quay.io/ascend/vllm-ascend:v0.19.1rc1-openeuler

docker run -it -u root -d --net=host \
  --privileged   --ipc=host   \
  --device=/dev/davinci_manager \
  --device=/dev/devmm_svm  \
  --device=/dev/hisi_hdc  \
  -v /usr/local/Ascend/driver:/usr/local/Ascend/driver   \
  -v /usr/local/dcmi:/usr/local/dcmi  \
  -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
  -v /usr/local/sbin:/usr/local/sbin  \
  -v /usr/local/Ascend/driver/tools/hccn_tool:/usr/local/Ascend/driver/tools/hccn_tool  \
 -v /opt/data/modelscope:/opt/data/modelscope   --name ernie-image  \
  m.daocloud.io/quay.io/ascend/vllm-ascend:v0.19.1rc1-openeuler    /bin/bash

## add by wei (Line 123)
def apply_rotary_emb_npu(x_in: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
    rot_dim = freqs_cis.shape[-1]
    x, x_pass = x_in[..., :rot_dim], x_in[..., rot_dim:]
    cos_ = torch.cos(freqs_cis).to(x.dtype)
    sin_ = torch.sin(freqs_cis).to(x.dtype)

    # Use NPU optimized rotary embedding
    out = rotary_position_embedding(x, cos_, sin_, rotated_mode='rotated_half')

    return torch.cat((out, x_pass), dim=-1)

# Auto-select NPU or CPU implementation
if freqs_cis is not None:
    if is_torch_npu_available():
        query = apply_rotary_emb_npu(query, freqs_cis)
        key = apply_rotary_emb_npu(key, freqs_cis)
    else:
        query = apply_rotary_emb(query, freqs_cis)
        key = apply_rotary_emb(key, freqs_cis)

## add by wei (Line 3311)
def is_supported_laser_attention(head_dim, q_seqlen, kv_seqlen):
    MAX_DIM = 128
    MIN_SEQLEN_SELF = 4000
    MIN_SEQLEN_CROSS = 118404
    MAX_SEQLEN_CROSS = 119056

    if head_dim > MAX_DIM:
        return False
    if q_seqlen == kv_seqlen:
        return q_seqlen >= MIN_SEQLEN_SELF
    else:
        return (MIN_SEQLEN_CROSS <= q_seqlen <= MAX_SEQLEN_CROSS) and \
            (MIN_SEQLEN_CROSS <= kv_seqlen <= MAX_SEQLEN_CROSS)


def _broadcast_attn_mask_npu(query, key, attn_mask):
    if attn_mask is not None:
        if attn_mask.ndim == 2 and attn_mask.shape[0] == query.shape[0] and attn_mask.shape[1] == key.shape[1]:
            batch_size, seq_len_q, seq_len_kv = attn_mask.shape[0], query.shape[1], key.shape[1]
            attn_mask = attn_mask.unsqueeze(1).expand(batch_size, seq_len_q, seq_len_kv).unsqueeze(1).contiguous()
        elif attn_mask.ndim == 4 and attn_mask.shape[1:3] == (1, 1):
            attn_mask = attn_mask.expand(-1, -1, query.shape[1], -1).contiguous()

    return attn_mask

# Step 1: Gather hidden_states
B, S_local, C = hidden_states.shape

# Convert to sequence-first for gather
hidden_states = hidden_states.transpose(0, 1)  # [B, S, C] -> [S, B, C]
hidden_states_full = sp_gather(hidden_states, dim=0)  # Gather along sequence
hidden_states_full = hidden_states_full.transpose(0, 1)  # Back to [B, S_full, C]

# Step 2: Gather rotary embedding (same process)
cos = cos.transpose(0, 1)
cos_full = sp_gather(cos, dim=0)
cos_full = cos_full.transpose(0, 1)

# Step 3: Create full attention mask
mask_full = torch.ones((B, 1, S_full, S_full), ...)

# Step 4: Compute attention on full sequence
output_full = processor(attn, hidden_states_full, mask_full, rotary_full)

# Step 5: Scatter back
output_full = output_full.transpose(0, 1)
output_local = sp_shard(output_full, dim=0)
output_local = output_local.transpose(0, 1)

# 1. 修改 diffusers (需要root权限)
sudo cp patches/diffusers_transformer_ernie_image.patch \
  /usr/local/python3.11.14/lib/python3.11/site-packages/diffusers/models/transformers/
  
sudo patch -p0 < patches/diffusers_transformer_ernie_image.patch

sudo cp patches/diffusers_attention_dispatch.patch \
  /usr/local/python3.11.14/lib/python3.11/site-packages/diffusers/models/
  
sudo patch -p0 < patches/diffusers_attention_dispatch.patch

# 2. 修改 vllm-omni
cd /vllm-workspace/vllm-omni/vllm_omni/diffusion/models/ernie_image/

patch -p0 < patches/vllm_omni_pipeline_ernie_image.patch
patch -p0 < patches/vllm_omni_ernie_image_transformer.patch
patch -p0 < patches/vllm_omni_ulysses_attention.patch

cd /vllm-workspace/ernie-image/bin

# 单卡配置文件（需自行创建）
vllm serve /opt/data/modelscope/hub/models/PaddlePaddle/ERNIE-Image \
  --config config/ernie_stage_single.yaml \
  --port 8000 \
  --dtype bfloat16

cd /vllm-workspace/ernie-image/bin

# 使用SP=4配置启动
bash start_ernie.sh

#!/bin/bash
echo "Starting ERNIE-Image (v0.19.0rc1 + diffusers adapter)..."
echo "Model: /opt/data/modelscope/hub/models/PaddlePaddle/ERNIE-Image"
echo "Port: 8000"
echo "Config: /vllm-workspace/ernie-image/config/ernie_stage_sp4_custom.yaml"

vllm serve /opt/data/modelscope/hub/models/PaddlePaddle/ERNIE-Image \
  --config /vllm-workspace/ernie-image/config/ernie_stage_sp4_custom.yaml \
  --port 8000 \
  --dtype bfloat16

# 检查服务状态
curl http://localhost:8000/v1/models

# 测试生成
curl -X POST http://localhost:8000/v1/images/generations \
  -H "Content-Type: application/json" \
  -d '{
    "prompt": "A beautiful sunset over the ocean",
    "size": "1024x1024",
    "num_inference_steps": 50,
    "guidance_scale": 4.0
  }'

# 旧
class ErnieImageTransformer2DModel(nn.Module):

# 新
class ErnieImageTransformer2DModel(CachedTransformer):

无 Cache-DiT（50 步全部计算）：
  Step 0: 全部 N 层计算 → hidden_0
  Step 1: 全部 N 层计算 → hidden_1
  ...
  Step 49: 全部 N 层计算 → hidden_49
  总计算量: 50 × N 层

有 Cache-DiT（缓存部分步的部分层）：
  Step 0-3: 全部 N 层计算               ← warmup，始终全算
  Step 4-6: 残差差值 < threshold → 缓存 ← 只算前几层，其余复用
  Step 7:   连续缓存达上限，强制全算    ← max_continuous_cached_steps
  ...
  总计算量: 大幅减少（取决于缓存比例）

self.sp_size = _get_sequence_parallel_world_size_or_one()

if current_omni_platform.is_npu():
    self._sp_plan = {}                          # 清空自动分片计划
    if self.sp_size > 1:
        for layer in self.layers:
            attn = layer.self_attention
            attn.attn.skip_sequence_parallel = True   # 跳过 vLLM 内置 SP
            ulysses_processor = ErnieImageUlyssesAttnProcessorV2()
            attn.set_processor(ulysses_processor)     # 安装自定义处理器

if current_omni_platform.is_npu():
    # 1. 直接读 config.json（需要传入 parallel_config + od_config）
    transformer_config_path = Path(model) / "transformer" / "config.json"
    tf_config_dict = json.load(transformer_config_path)
    self.transformer = ErnieImageTransformer2DModel(
        parallel_config=self.parallel_config,
        od_config=od_config,
        **tf_config_dict,
    ).to(self._execution_device)

    # 2. __init__ 中直接加载权重（DO NOT remove unless verified with benchmark）
    #    实测：__init__ 加载 ~27s，延迟到 load_weights() ~34s（慢 21%）
    transformer_path = Path(model) / "transformer"
    for weight_file in transformer_path.glob("*.safetensors"):
        ...
    self.transformer.load_weights(weights_dict.items())

if current_omni_platform.is_npu():
    x, S_local, ... = self._prepare_npu(hidden_states, text_bth, text_lens)
else:
    x, S_local, ... = self._prepare_gpu(hidden_states, text_bth, text_lens, N_img)

# GPU: pad freqs output (unified_prepare)
freqs_cos = torch.cat([freqs_cos, torch.zeros(...)], dim=...)  # ← RoPE 之后补零

# NPU: pad grid_yx input (_prepare_npu)
pad_grid = torch.zeros((pad_size, 2))  # ← RoPE 之前补零坐标 (0,0)
grid_yx_padded = torch.cat([grid_yx, pad_grid], dim=0)

# ErnieImageAttention 新增
self.processor = None

def set_processor(self, processor):
    self.processor = processor

def forward(self, ...):
    if self.processor is not None:
        return self.processor(self, hidden_states, attention_mask, image_rotary_emb, **kwargs)
    # Default GPU attention logic (原有代码不变)
    ...

输入: hidden_states [B, S_local, C]（本地序列片段）

Step 1: to_q/k/v 投影
  → query/key/value: [B, S_local, H, D]

Step 2: AllToAll gather (scatter heads, gather sequence)
  SeqAllToAll4D(scatter_idx=2, gather_idx=1)
  → [B, S_global, H/P, D]（全局序列，部分 heads）

Step 3: 全局 RoPE
  freqs 如是本地片段 → sp_gather 拼成全局
  → 正确的全局位置编码

Step 4: 全局 Attention
  mask: 从 4D mask 中提取全局掩码
  → attn(query, key, value, metadata)

Step 5: AllToAll scatter (gather heads, scatter sequence)
  SeqAllToAll4D(scatter_idx=1, gather_idx=2)
  → [B, S_local, H, D]（回到本地序列，全部 heads）

Step 6: to_out 输出投影
  → [B, S_local, C]

# 旧：只支持 [B, S, H, D]
cos_ = freqs_cos.unsqueeze(2).to(x.dtype)

# 新：同时支持 [S, B, H, D]（NPU AllToAll 后可能出现的格式）
if x.dim() == 4 and x.shape[0] != freqs_cos.shape[0]:
    cos_ = freqs_cos.transpose(0, 1).unsqueeze(2).to(x.dtype)  # freqs [B,S,D] → [S,B,1,D]
else:
    cos_ = freqs_cos.unsqueeze(2).to(x.dtype)                    # freqs [B,S,D] → [B,S,1,D]

def _vae_decode(self, latents, device, output_type="pil"):
    if current_omni_platform.is_npu():
        sp_size = getattr(self.parallel_config, "sequence_parallel_size", None) or 1
        if sp_size > 1:
            sp_rank = get_sequence_parallel_rank()
            dist.broadcast(latents, src=0, group=get_sp_group().device_group)
            if sp_rank != 0:
                return DiffusionOutput(output=[], peak_memory_mb=0.0)
    # 原有的 BN 校正 + unpatchify + VAE decode
    ...