1. Tensor

torch.tensor类是大模型推理的基础，利用torch搭建模型网络，

实际上

• data
• parameter
• gradient
• optimizer state
• activation

这些内容的存储形式，都是tensor

张量可以是任意维度，维度也被称为rank

import torch

x = torch.zeros(4)        # rank 1：vector
x = torch.zeros(4, 8)     # rank 2：matrix
x = torch.zeros(4, 8, 2)  # rank 3

Transformer 中常见 rank-4 tensor：分别表示批处理大小，序列长度，多头注意力头数，隐藏层维度

B = 32  # batch size
S = 16  # sequence length
H = 16  # number of heads
D = 64  # hidden dimension per head
x = torch.zeros(B, S, H, D)

tensor在底层占用的内存也可以估算，numel接口可以获取张量的元素个数，element_size可以获取每个元素的内存消耗，相乘则是整个张量消耗的内存。

这里数据类型为float32，每个元素4字节，元素个数32，总的内存占用是128B

def get_memory_usage(x: torch.Tensor) -> int:
    return x.numel() * x.element_size()

x = torch.zeros(4, 8)
assert x.dtype == torch.float32
assert x.numel() == 32
assert x.element_size() == 4
assert get_memory_usage(x) == 128

2. 数据类型

LLM推理中常见的数据类型有fp32,fp16,bf16

2.1 fp32

最常见的就是fp32，如下图，每个占4B，32个bit，8bit存指数，23bit存有效数字，1bit存符号。数值范围和精度都比较可靠，问题是太消耗内存了，而AI不是科学计算，对精度要求没有那么高，牺牲性能换精度并不划算，所以fp32是最基础的类型，但不是追求极限性能场景使用的类型。

2.2 fp16

另一个常见的是fp16，和fp32类似，只是内存消耗变小了，只有2B，5为指数，10位精度。

五位指数位能表示的数值范围变小了很多，可能出现数值下溢，如下代码的断言实际上会通过，这可能导致训练中出现梯度消失。

x = torch.tensor([1e-8], dtype=torch.float16)
assert x == 0  # underflow

我们具体分析一下，fp16的计算规则如下，S是符号位，E是指数位，M是尾数。公式如下，在指数位不全为0时，会给尾数额外+1。这里的尾数M表示的是一个二进制小数，也就是从高位到低位表示的分别是$1/2,1/2^2,1/2^3...$
$$\text{实际值}=(-1{)}^S\times (1+M)\times 2^{E-15}$$

如果想表示很小的数，那么在指数E全为0时，指数E硬编码为-14，并且位数不再+1，公式如下

$$\text{实际值}=(-1{)}^S\times (0+M)\times 2^{-14}$$
那么能表示的最小非零数是E全为0，M只有最低位为1，表示的数值为
$$2^{-10}\times 2^{-14}=2^{-24}$$
这约等于$5\times 10^{-8}$，因此如果赋值为$1\times 10^{-8}$，无法表示这么小的数，会被映射到0。

2.3 bf16

为了解决fp16的下溢问题，谷歌(Google Brain)提出了bf(Brain Float-Point)16类型，增加了指数位，减少了尾数位

增加后的指数位长度和fp32一样，可以表示的最大值，最小值也接近，具体来说可以表示$[9.18\times 10^{-41},3.4\times 10^{38}]$的范围，有效解决了fp16梯度消失的问题，但同样有内存优势。至于损失的尾数精度，大模型对精度并不敏感。

2.4 混合精度

只用 fp32 很耗 memory；只用 fp16，甚至只用 bf16，又可能数值不稳定。mixed precision training 的常见方案是：

• parameter、activation、gradient 使用 bf16。
• optimizer state 使用 fp32。

PyTorch 的 Automatic Mixed Precision（AMP）会在相对安全的操作上自动降低精度，例如 matmul 可以使用 bf16，而 exp 等敏感操作可保留更高精度。

with torch.amp.autocast("cuda", dtype=torch.bfloat16):
    x = torch.zeros(4, 8, device="cuda")

2.5 fp8 与 fp4

H100开始，新支持两种 FP8：E4M3 的范围约为 $[-448,448]$，E5M2 的范围约为 $[-57344,57344]$。更低精度可以显著减少 memory 和提高吞吐，但需要 scale 等机制维持数值范围。

NVIDIA 在 2025 年推出 NVFP4，每个值仅 4 bits。每个 block 使用独立 scale factor，因此整体 dynamic range 可以较大，但相邻值不能任意跨数量级变化。部分细节由 NVIDIA library 在用户代码之外处理。

2.6 把 Tensor 放到 GPU

PyTorch 默认在 CPU memory 中创建 tensor。要利用 GPU 的大规模并行能力，必须把数据移入 GPU memory：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

x = torch.zeros(32, 32)
x = x.to(device)

# 也可以直接在目标 device 上创建
with torch.device(device):
    x = torch.zeros(32, 32)

3. 使用 einops 思考 Tensor

传统 PyTorch 常用负数下标表示维度，例如这里转置交换最后一个维度和倒数第二个维度：

x = torch.ones(2, 2, 3)      # batch, seq, hidden
y = torch.ones(2, 2, 3)      # batch, seq, hidden
z = x @ y.transpose(-2, -1)  # batch, seq, seq

代码很短，但 -2、-1 容易写错，也难以审查。einops 库用名称标识维度，更适合推导复杂 tensor operation。它的核心思想是用显式的命名来指代各个维度的操作。

3.1 einsum

einsum 是带有的命名的矩阵乘法接口，规则是，结果表达式中未出现的维度会进行矩阵乘法，例如下面就是对两个张量的hidden维度求和

from einops import einsum

x = torch.ones(3, 4)  # seq1, hidden
y = torch.ones(4, 3)  # hidden, seq2

z = einsum(x, y, "seq1 hidden, hidden seq2 -> seq1 seq2")

如果增加batch维度，则如下，这里可以用省略号表示任意个广播维度，表达式只具体写进行矩阵乘法运算的维度。

x = torch.ones(2, 3, 4)  # batch, seq1, hidden
y = torch.ones(2, 3, 4)  # batch, seq2, hidden

z = einsum(
    x,
    y,
    "batch seq1 hidden, batch seq2 hidden -> batch seq1 seq2",
)

# ... 表示任意数量的广播维度
z = einsum(x, y, "... seq1 hidden, ... seq2 hidden -> ... seq1 seq2")

3.2 reduce

reduce 对指定维度执行 sum、mean、max、min 等归约，同样可以广播，进行的操作类型通过一个字符串传入

from einops import reduce

x = torch.ones(2, 3, 4)  # batch, seq, hidden
y = reduce(x, "... hidden -> ...", "sum")

3.3 rearrange

一个扁平维度有时实际包含多个逻辑维度，例如 total_hidden = heads * hidden1，可以用括号表示把这这两个维度拆开，可以合并。拆分时需要传入拆分的必要维度信息，生育信息可以自动推导，比如已知输入的heads*hidden1维度，再传入heads=2即可推导出hidden1维度的大小。

from einops import rearrange

x = torch.ones(3, 8)  # seq, total_hidden
w = torch.ones(4, 4)  # hidden1, hidden2

x = rearrange(x, "... (heads hidden1) -> ... heads hidden1", heads=2)
x = einsum(x, w, "... hidden1, hidden1 hidden2 -> ... hidden2")
x = rearrange(x, "... heads hidden2 -> ... (heads hidden2)")

eniops库的价值在于，shape推导是显式写在代码里的，不需要阅读代码时再人脑里推算，并且维度进行命名，不需要手动对照操作的是哪个维度。