昇思大模型推理：输出 logits

一、引言

在昇思 MindSpore + MindFormers 大模型推理链路中，logits 是模型最后一层 Transformer 输出、未经过 Softmax 的原始分值向量 ，维度一般为 [batch, seq_len, vocab_size]。每个位置上的 logit 值代表该位置预测为对应词表 token 的原始得分，经过 softmax 后才是 token 概率，再通过采样 / 贪心得到生成 token。

常规model.generate()默认只返回 token 序列，屏蔽中间 logits，适合直接对话场景；而做可控生成、奖励模型对齐、PPL 困惑度计算、采样策略自定义、输出概率诊断、对比不同量化精度差异时，需要直接拿到 logits。

关键点：

1. MindFormers 的generate接口默认走推理生成分支，会内部执行采样；要输出 logits，有两条路径：调用 model.predict（返回完整 logits） 、或者自定义 generate 循环，每一步取出 next-token 的 logit；

2. 大模型 KV Cache 推理时，每步只输出当前新 token 位置的 logits（仅最后一维） ，形状 [batch, vocab_size]，这是增量推理的核心；

3. 显存开销：完整序列 logits 占用很大显存，长序列下不要全量保存，按需取最后一步 logits；

4. 可配合昇腾 Profiling 观察最后一层算子输出，用于精度漂移、量化误差定位。

环境：MindSpore 2.3，MindFormers，CANN，昇腾 NPU；也可在鲲鹏 openEuler CPU 后端运行。

二、开发教程

1. 依赖导入与基础配置

import numpy as np
import logging
from mindspore import set_context, GraphMode, Tensor
from mindformers import AutoModel, AutoTokenizer

logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(message)s")


# 昇腾NPU；鲲鹏CPU改为 device_target="CPU"
set_context(
    mode=GraphMode.GRAPH,
    device_id=0,
    device_target="Ascend"
)

2. 方式一：predict 接口，直接返回完整 logits（非增量，适合单步诊断）

model.predict 不做采样解码，直接返回模型原始输出，包含 logits，适合拿到整段序列所有位置的分值。

def get_full_logits(model_name, ckpt_path, prompt):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(
        model_name,
        load_checkpoint=ckpt_path,
        set_train=False
    )
    model.set_train(False)

    inputs = tokenizer(prompt, return_tensors="ms")
    # predict返回字典，logits键即为原始输出
    output = model.predict(**inputs)
    logits: Tensor = output["logits"]

    logging.info(f"logits shape: {logits.shape}")
    # logits: [B, S, Vocab]
    logits_np = logits.asnumpy()
    # 取最后一个位置的logits，即下一个token预测分值
    last_pos_logits = logits_np[:, -1, :]
    return logits_np, last_pos_logits, tokenizer

3. 方式二：KV Cache 增量推理，每步取出 next logits（推荐，和 generate 对齐）

generate内部的 KV Cache 增量推理，每一步只计算新 token 位置，返回[B, Vocab]的 logits，适合流式自回归生成，可自定义 softmax、top_k、temperature。

def generate_with_logits(model_name, ckpt_path, prompt, max_new_tokens=32, temperature=1.0, top_k=50):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(
        model_name,
        load_checkpoint=ckpt_path,
        set_train=False
    )
    model.set_train(False)

    inputs = tokenizer(prompt, return_tensors="ms")
    input_ids = inputs["input_ids"]
    attention_mask = inputs["attention_mask"]

    past_key_values = None
    generated_ids = input_ids
    logits_record = []  # 保存每一步next token logits

    for _ in range(max_new_tokens):
        # 单步前向，使用KV Cache
        out = model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            use_cache=True
        )
        next_logits = out[0][:, -1, :]  # shape [B, vocab_size]
        past_key_values = out[1]
        logits_record.append(next_logits.asnumpy())

        # 采样逻辑，等价于原生generate
        if temperature > 0:
            next_logits = next_logits / temperature
        # top-k
        vals, idx = next_logits.topk(top_k)
        mask = next_logits < vals[:, -1:].expand_as(next_logits)
        next_logits = next_logits.masked_fill(mask, -1e9)

        # softmax + sample
        probs = next_logits.softmax(axis=-1)
        next_id = probs.argmax(axis=-1, keepdims=True)

        generated_ids = np.concatenate([generated_ids.asnumpy(), next_id.asnumpy()], axis=-1)
        input_ids = Tensor(next_id)
        attention_mask = np.ones_like(input_ids)

    text_out = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    return text_out, logits_record, tokenizer

4. 主函数调用示例

if __name__ == "__main__":
    model_cfg = "qwen_1_8b"
    ckpt = "./qwen1.8b.ckpt"
    prompt = "解释什么是logits"

    # 1. 获取完整序列logits
    full_logits, last_logits, tok = get_full_logits(model_cfg, ckpt, prompt)
    print("最后位置logits形状", last_logits.shape)

    # 2. 自回归生成，每步记录next logits
    gen_text, logit_list, tok2 = generate_with_logits(model_cfg, ckpt, prompt, max_new_tokens=20)
    print("生成文本：", gen_text)
    print(f"共采集 {len(logit_list)} 组next-token logits，每组shape {logit_list[0].shape}")

5. 辅助函数：由 logits 算概率、求 PPL

def calc_ppl(logits_np, label_ids):
    """
    logits: [B,S,V]
    label_ids: [B,S]
    返回序列PPL
    """
    B, S, V = logits_np.shape
    exp_log = np.exp(logits_np - np.max(logits_np, axis=-1, keepdims=True))
    prob = exp_log / np.sum(exp_log, axis=-1, keepdims=True)
    # 取出标签对应概率
    label_prob = np.take_along_axis(prob, label_ids[..., None], axis=-1)
    nll = -np.log(label_prob + 1e-10)
    ppl = np.exp(np.mean(nll))
    return ppl

三、问题总结

1. 直接调用 model.generate 拿不到 logits generate是高层封装，内部做完采样、丢弃原始 logits，返回 token id。必须走底层model()前向或者model.predict。

2. KV Cache 增量推理时 logits 形状不对，误以为是 B,S,V

   开启use_cache=True后，每轮只输入最新 token，输出仅新位置的 logits，是[B, vocab_size]，不是完整序列。

3. 长序列 predict 显存 OOM predict输出全序列 logits，[B,seq_len,vocab]，vocab 通常上万，序列 4k 时显存暴涨。诊断优先用增量方式，只保存每步 next logits。

4. logits 数值极大 / 极小，softmax 溢出

   直接 softmax 会数值不稳定，代码里做logits - max_val的 log-sum-exp 平移，防止 exp 溢出。

5. MindFormers 版本差异，输出字典 key 不一致

   老版本部分模型前向返回 tuple 而不是 dict，logits 在out[0]；新版本 predict 返回 dict，key 为logits，报错时打印model.output_names查看。

6. 量化后 logits 偏移，和 FP16 基线对比差异大

   做精度评估时，用这套代码分别跑 FP16、INT8，比对 logits 的 KL 散度，定位量化带来的分布漂移，比单纯看生成文本更敏感。

四、流程

整套流程可分为 4 步：

1. 模型加载：加载 MindFormers 大模型，关闭训练模式，固定昇腾 / 鲲鹏后端；

2. 选择前向模式：诊断全序列分布用predict；流式生成、低显存采集用use_cache=True的单步循环；

3. 采集 logits：按需保存原始分值，不要默认全量存储；基于 logits 自定义采样、计算 PPL、KL 散度；

4. 精度回归：同一 prompt 固定输入，保存基线 logits，新版本 / 量化后自动比对，用于性能 + 精度基准测试。

输出 logits 是大模型可观测性的基础：不仅可以自定义解码策略，还可以做奖励信号、不确定性评估、量化精度诊断，常和前面的昇思基准测试脚本组合使用。