01 · 模型最后一层为什么也能占很多内存
Transformer产生隐藏向量H[T,D]，输出权重W[V,D]把每个位置映射到V个词表分数：Z=H Wᵀ，形状[T,V]。logit是未经过Softmax的分数；交叉熵再用目标token和整行分数计算Loss。T可表示展平后的batch×序列位置。
自设T=4096、V=128000：FP16完整logits占1000 MiB，若另物化FP32副本则再占2000 MiB。Softmax、梯度、隐藏向量和权重尚未计入；实际是否同时驻留取决于实现，不把这些数字当某框架实测峰值。
依据：Cut Your Losses in Large-Vocabulary Language Models · 原始论文 ↗ · PyTorch · CrossEntropyLoss语义与reduction（main） ↗
02 · 推理：只给真正要采样的位置做投影
普通自回归Prefill后，通常只用每个请求最后有效位置的隐藏向量预测第一个生成token。若所有提示位置都做输出投影，会产生许多本次采样不需要的分数。可以先选这些行再做LM head；Transformer中间层仍需计算提示词。
普通单请求提示词：hidden shape [4096,D]
h_last = hidden[最后有效位置]
logits_next = h_last @ W.T # [V]
# 不能误取padding最后一行；packed batch按请求边界选行。本例单行FP16 logits约0.244 MiB，4096行完整表为1000 MiB。只减少LM head的行数和中间表，不代表整个Prefill加速4096倍。Decode一般本来就只有新位置，收益不能重复计算。
提示词logprob、困惑度、打分、投机验证、多位置预测等任务可能需要更多行；请求需要的分数必须保留。只选目标词也不能直接替代普通采样，因为top-k/top-p或Softmax归一化仍依赖其他候选。
03 · 训练：Loss不是只看目标token的分数
先看无类别权重、无label smoothing的一行交叉熵：L=log(Σ_v exp(z_v))−z_y，其中y是正确目标。即使只输出一个标量Loss，分母仍依赖全部词表。
自设z=[0,ln2,ln3,ln4]，目标y=2（从0编号）。目标概率3/10，Loss=ln(10)−ln(3)=ln(10/3)≈1.203973。只看目标分数，或者只在目标所在片段归一化，都会算错。
逐行稳定形式：
m = max(z)
L = m + log(sum(exp(z-m))) - z[target]
# 不必先生成并保存完整Softmax概率表。04 · 两种分块：分token行，或分词表列
| 方法 | 减少哪个中间表 | 保留什么 |
|---|---|---|
| 按token行分块 | 每次C×V logits，而非T×V | 每行仍遍历完整词表；累积Loss/梯度。 |
| 按词表列分块并融合 | 每次小块投影并更新统计，避免全局完整logits | 全词表log-sum-exp与目标分数。 |
| 只融合已生成logits的交叉熵 | 省部分Softmax/临时缓冲 | 已有T×V logits仍存在，不能等同于融合Linear+Loss。 |
两种分块可以组合。把每块Loss相加后才统一按有效目标数归一化；不是每块取mean再无权平均。普通自动微分即使Python循环逐块前向，也可能把各块中间值都保留到统一backward，无法得到预想峰值，需要配套反向和生命周期设计。
同一个完整词表，临时表可以分多小？
T=4096，V=128000，每logit 2B；仅计一张临时logits表。小块处理后释放并保留必要统计，梯度/权重/隐藏向量等另计；本演示不模拟自动微分或设备内核。演示只计一张FP16 logits临时表；不计输入、权重、统计和梯度，也不假设换分块后计算量变少。
依据：Cut Your Losses in Large-Vocabulary Language Models · 原始论文 ↗
05 · 不存完整词表，怎样算正确分母
为每行维护已见最大值m与指数和ℓ。新块最大值为mb时，mnew=max(m,mb)；旧和乘exp(m−mnew)，再加新块的Σexp(z−mnew)。最终log-sum-exp=m+logℓ。目标所在块顺便取出z_y。
m = -inf; ell = 0; target_logit = 未记录
for each vocabulary tile:
z_tile = H_tile @ W_tile.T
new_m = max(m, max(z_tile))
ell = ell * exp(m-new_m) + sum(exp(z_tile-new_m))
m = new_m
若target在本块：记录对应z
loss = m + log(ell) - target_logit
# 第一块直接初始化，避免-inf-(-inf)。上例按两列分块：块A=[0,ln2]，局部指数总质量3；块B=[ln3,ln4]，质量7。目标在B，正确分母仍为10，不能只用7；误算会得到ln(7/3)≈0.847298。机制与在线Softmax相通，但这里汇总的是词表归一化，不是Attention的Value输出。
依据：Cut Your Losses in Large-Vocabulary Language Models · 原始论文 ↗
06 · 反向仍需要完整词表的贡献
无权重标准交叉熵对logit的梯度是p_v−1[v=y]，所以非目标类别通常也有梯度。给定保存的log-sum-exp，反向可重新生成当前logits块，算对应概率和梯度，再累积隐藏向量与分类权重梯度。
重算当前块z，p = exp(z - saved_logsumexp)
dZ = p - one_hot(target落在本块)
dZ *= Loss归一化及上游梯度
dH += dZ @ W_tile
dW_tile += dZ.T @ H
# dH跨词表块求和；dW按全部有效token累积。这是用部分重算换取中间表容量，仍可能增加投影计算。Apple CCE实现还有梯度过滤等吞吐策略，应与“不物化完整logits”的内存思想分开：过滤微小梯度不等于保留所有梯度项。仓库提供关闭过滤的对照实现，实际支持和数值行为按安装版本核对。
07 · 标签、忽略项和分布式分母不能改错
因果训练通常让位置i预测下一个token。先确定谁负责shift，避免数据预移一次、Loss接口又移一次。ignore_index表示该位置不贡献Loss；Packing和Padding的边界要沿用训练目标，不能因省logits而改成另一套标签。
本篇算例按有效目标数取mean。带类别权重、label smoothing、softcap、Z-loss或蒸馏全分布时，要实现相应定义，不能只套LSE−目标分数的最简公式。所有目标都被忽略时，分母为零，应明确训练框架的处理约定。
词表并行把列放在多个rank：先得到全局最大值，再求全局指数和和目标分数；dH也要按词表分片贡献归约。每卡对自己的词表算独立交叉熵再平均，无法得到完整词表Loss。
依据：PyTorch · CrossEntropyLoss语义与reduction（main） ↗ · Apple · Cut Cross-Entropy原作者实现与数值路径（main） ↗
08 · 小脚本同时核对Loss和梯度
下载脚本比较完整logits与按token/词表双向分块的交叉熵，覆盖非整除块、ignored目标和非均匀行块。输出dH、dW、bias梯度，并用有限差分验证；另展示只用目标词表块分母的错误结果。CPU没有自动微分保留问题，不测试设备峰值和速度。
先记录真实需求：采样行、打分行或训练标签
建立完整logits + 标准Loss基线
核对shift、mask、reduction、额外Loss项
运行融合/分块候选并验证前后向
测峰值显存、投影与重算、完整训练step
# 接口若不再返回logits，下游指标和调试流程也要适配。CCE仓库是Linear+交叉熵融合的实现阅读入口，并提供词表并行和不同数值路径。不是所有loss函数、设备和模型包装都能直接替换。昇腾或其他后端应检查是否存在同定义的融合/分块实现，不能把一个CUDA包当跨平台开关。
依据：Apple · Cut Cross-Entropy原作者实现与数值路径（main） ↗ · vLLM Ascend · 支持功能矩阵（main） ↗
09 · 少占内存和更快，分别报告
| 检查 | 重点 |
|---|---|
| 语义 | 所需logits没有被省掉；Loss、标签、忽略项及全局归一化一致。 |
| 反向 | 隐藏向量、分类权重与bias梯度；混合精度、过滤策略与数值误差。 |
| 内存 | 完整/分块logits、重算缓冲、持久统计、梯度和实际峰值。 |
| 速度 | 前向、反向、额外投影及调度总和；再看完整训练step或服务延迟。 |
- 训练只监督一个目标token，为什么仍要遍历其他词表项？
- 分块的mean为什么不能不加权平均？
- FlashAttention和融合Linear+Loss都减少中间表，两者归一化的轴有什么不同？
流程位置与关联方法
这个方法可能出现在多个位置；计算线包含模型计算与算子实现，执行线包含运行时、内存和通信。
资料与核查范围
原理依据以原始论文和官方文档为主；算例与时间线是教学推导。文档核查日期：2026-10-10。引用 main / latest 文档时，实际接口以安装版本为准。
- Cut Your Losses in Large-Vocabulary Language Models · 原始论文 ↗
- Apple · Cut Cross-Entropy原作者实现与数值路径（main） ↗
- PyTorch · CrossEntropyLoss语义与reduction（main） ↗
- vLLM Ascend · 支持功能矩阵（main） ↗
代码为接入片段或教学伪代码；本页没有声称完成 NVIDIA / 昇腾硬件实测。