Llama 7B 迁移至 MindSpore 实战指南：避坑与优化

⚠️ 注意：有的实现把 cos/sin 的 layout 写反了；decode 阶段 pos 要累加（pos_offset += 1），别反复从 0 开始。

04 权重转换：从 HuggingFace → MindSpore .ckpt

HuggingFace 的 Llama 权重是多个 pytorch_model-*.bin。思路：用 torch.load 拿 state_dict，做键名映射，再 mindspore.save_checkpoint。

1、键名映射表（示例）

HuggingFace（常见） → MindSpore（示例命名）：

2、转换脚本（最小可用）

import os, torch, mindspore as ms
from mindspore import save_checkpoint, Tensor

def map_key(hf_key: str):
    key = hf_key
    key = key.replace("model.embed_tokens.weight", "tok_embeddings.embedding_table")
    key = key.replace("model.norm.weight", "final_norm.gamma")
    key = key.replace("lm_head.weight", "lm_head.weight")
    key = key.replace("model.layers.", "blocks.")
    key = key.replace(".self_attn.q_proj.", ".attn.wq.")
    key = key.replace(".self_attn.k_proj.", ".attn.wk.")
    key = key.replace(".self_attn.v_proj.", ".attn.wv.")
    key = key.replace(".self_attn.o_proj.", ".attn.wo.")
    key = key.replace(".mlp.gate_proj.", ".mlp.w1.")
    key = key.replace(".mlp.down_proj.", ".mlp.w2.")
    key = key.replace(".mlp.up_proj.", ".mlp.w3.")
    key = key.replace(".input_layernorm.weight", ".ln1.gamma")
    key = key.replace(".post_attention_layernorm.weight", ".ln2.gamma")
    return key

def torch_to_mindspore_ckpt(hf_dir, ms_ckpt_path, dtype=ms.float16):
    # 1) 收集所有 shard
    sd = {}
    for name in sorted(os.listdir(hf_dir)):
        if name.startswith("pytorch_model-") and name.endswith(".bin"):
            part = torch.load(os.path.join(hf_dir, name), map_location="cpu")
            sd.update(part)
        elif name == "pytorch_model.bin":
            sd.update(torch.load(os.path.join(hf_dir, name), map_location="cpu"))
    # 2) 键名映射 + 类型转换
    ms_params = []
    for k, v in sd.items():
        ms_k = map_key(k)
        if "rope.freqs" in ms_k:
            continue
        np_v = v.numpy()
        ms_params.append({"name": ms_k, "data": Tensor(np_v).astype(dtype)})
    save_checkpoint(ms_params, ms_ckpt_path)
    print(f"Saved MindSpore ckpt to: {ms_ckpt_path}")

# 用法：
# torch_to_mindspore_ckpt("/path/to/llama-hf", "llama7b_ms.ckpt", dtype=ms.float16)

⚠️ 注意：LayerNorm 在 Llama 是无 bias，MindSpore 里如果你 LayerNorm 定义带 beta，要么删掉，要么初始为 0 并在图里不使用；否则数值会'飘'。

05 Llama 前向与 KV Cache（prefill + decode）

1、Attention mask 语义

• **训练：**通常是 [bs, 1, seq, seq] 或 [bs, seq] 的下三角 + padding mask。
• **推理：**prefill 阶段 mask 仍按下三角；decode 阶段仅对新 token 做与历史的点积，mask 形状变小。

建议统一为 float mask，填充不可见位置为 -1e4（或和你 softmax 实现一致的 -inf），避免 dtype 乱战。

2、简化版 KV Cache

class KvCache:
    def __init__(self, n_layer, n_head, max_batch, max_len, head_dim, dtype=ms.float16):
        self.k = [ms.numpy.zeros((max_batch, n_head, max_len, head_dim), dtype=dtype) for _ in range(n_layer)]
        self.v = [ms.numpy.zeros((max_batch, n_head, max_len, head_dim), dtype=dtype) for _ in range(n_layer)]
        self.pos = 0 # 当前 decode 写入位置

    def update(self, layer_idx, k_new, v_new):
        # [bs, head, 1, dim]
        p = self.pos
        self.k[layer_idx][:, :, p:p+1, :] = k_new
        self.v[layer_idx][:, :, p:p+1, :] = v_new

    def step(self):
        self.pos += 1

⚠️ 注意：别在 decode 阶段每步都 concat，就地写入 slice，Ascend 的内存移动不白嫖。

06 训练与微调（LoRA/全参）

LoRA 在 MindSpore 的一个常见实现：给线性层包一个 A/B 低秩旁路，前向时加上 x @ A @ B * alpha/r。

建议把 LoRA 的参数单独分组，禁用 weight decay；并只在 target 模块（q_proj, v_proj, o_proj, w1/w3）上挂。

def wrap_lora(linear, r=16, alpha=32):
    in_f, out_f = linear.in_channels, linear.out_channels
    A = ms.Parameter(ms.ops.zeros((in_f, r), ms.float16))
    B = ms.Parameter(ms.ops.zeros((r, out_f), ms.float16))
    scale = alpha / r

    def forward(x):
        base = linear(x)
        lora = ms.ops.matmul(ms.ops.matmul(x, A), B) * scale
        return base + lora

    linear.forward = forward
    return linear

⚠️ 注意：MindSpore Graph 下如果你'猴子补丁'forward，要确保图能稳定跟住；更稳的做法是写一个 LoraLinear(Cell) 包起来。

07 性能小记（不玄学）

• **GRAPH_MODE + O2 混合精度：**不解释。
• **大 batch prefill：**把多条输入拼长些，prefill 吞吐会好不少（当然别 OOM）。
• **KV Cache 扁平化：**把 [bs, head, t, dim] 按设备最友好的内存布局摆放（这块我没深挖，简单就地 slice 已经够用）。
• **避免 Python 回环：**decode loop 尽量把张量操作留在图里，减少 host 参与。
• **检查算子降级：**图编译日志里搜 'fallback/host' 之类关键词，别让关键算子跑到 CPU 端。

08 端到端推理样例（极简）

import mindspore as ms
from mindspore import Tensor
import numpy as np

net = build_llama_from_ckpt("llama7b_ms.ckpt") # 你的加载逻辑
net.set_train(False) # 假装我们已经有 tokenizer
prompt_ids = np.array([[1, 42, 123, 456]]) # <s> ...
attn_mask = np.ones_like(prompt_ids) # prefill
logits, cache, pos = net(Tensor(prompt_ids, ms.int32), Tensor(attn_mask, ms.float32), cache=None, pos=0)

# decode N 步
generated = []
x = np.array([[50256]]) # 假设上一步采样出的 token
for _ in range(32):
    l, cache, pos = net(Tensor(x, ms.int32), Tensor(np.ones_like(x)), cache=cache, pos=pos)
    next_id = int(ms.ops.argmax(l[0, -1], axis=-1).asnumpy())
    generated.append(next_id)
    x = np.array([[next_id]])

⚠️ 注意：很多人把 pos 写死，导致 RoPE 永远用到第 0 行，性能和数值全飞。prefill 后 pos 应等于上下文长度，decode 逐步 +1。

09 常见报错对照（以防手忙脚乱）

• **Shape 不一致：**尤其 attention_mask，MindSpore 的广播规则和你在 PyTorch 的'侥幸成功'未必一致，显式 reshape 保命。
• **LayerNorm gamma/beta：**权重名映射遗漏，或 beta 多出来。
• **溢出：**fp16 的 matmul 穿了，loss scale 或者切到 bf16。
• **图编译卡慢：**第一次长一些正常，第二次还慢，看看是否每次都在重建图（输入 shape 乱飘）。

10 小结

迁 Llama 到 MindSpore 没有想象中那么可怕，难点集中在键名映射、RoPE 位移、KV Cache 写法三件事。 一旦跑通，Ascend 上的吞吐和能效都挺能打。别追求一步封神，先上一个'能打'的版本，再迭代优化。 最后，再次提醒自己：少写骚代码，别给图编译添堵。有时候'朴素写法'反而更快更稳（这点我已经被现实教育过两次，脸疼）。