""" ============================================================================= 第 2 讲 · 代码 2：浮点数格式 —— FP32、FP16、BF16、FP8、FP4 ============================================================================= 对应讲义：lectures/lecture02_pytorch_resource_accounting.md 第 3 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 上一个文件告诉我们：张量占的内存 = 元素个数 × 每个元素的字节数。 想省内存、省带宽、算得更快，最直接的办法就是**每个数用更少的比特**。 但比特少了，数字就表示得不准了。这个文件用实验回答： - 一个浮点数是怎么用比特表示的？（符号 / 指数 / 尾数） - 为什么 FP16 会"下溢"（很小的数变成 0）而 BF16 不会？ - 为什么 BF16 连 1 + 0.001 都算不对，却仍然是训练大模型的首选？ - 为什么"累加"必须用 FP32？（混合精度训练的核心原则） - FP8 / FP4 这么粗糙，怎么还能用？（缩放因子，尤其是"分块缩放"） ----------------------------------------------------------------------------- 【整体思路】 ----------------------------------------------------------------------------- 第 1 部分：各格式的规格表（位数、最大值、最小正规数、机器精度） 第 2 部分：亲眼看看比特：把浮点数拆成 符号 | 指数 | 尾数 第 3 部分：范围：上溢（变成 inf）与下溢（变成 0） 第 4 部分：精度：数与数之间的"间隔" 第 5 部分：累加误差 —— 为什么累加要用 FP32 第 6 部分：混合精度：torch.autocast 第 7 部分：FP8 与缩放；FP4 与分块缩放 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 02_floating_point.py """ import torch def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) # ============================================================================= # 第 1 部分：规格表 # ============================================================================= # 【思路】 # 一个浮点数 = (-1)^符号 × 1.尾数 × 2^(指数 - 偏置) # - 指数位越多 → 能表示的范围越大（从很小到很大） # - 尾数位越多 → 精度越高（相邻两个可表示的数挨得越近） # 总位数固定时，指数和尾数是此消彼长的关系。 # # 格式 总位数 符号 指数 尾数 # FP32 32 1 8 23 # FP16 16 1 5 10 # BF16 16 1 8 7 ← 指数位和 FP32 一样多！只砍了尾数 # FP8 E4M3 8 1 4 3 # FP8 E5M2 8 1 5 2 # FP4 E2M1 4 1 2 1 FORMATS = [ # (名字, dtype, 指数位, 尾数位) ("FP32", torch.float32, 8, 23), ("FP16", torch.float16, 5, 10), ("BF16", torch.bfloat16, 8, 7), ("FP8 E4M3", torch.float8_e4m3fn, 4, 3), ("FP8 E5M2", torch.float8_e5m2, 5, 2), ] def demo_spec_table(): """演示 1：用 torch.finfo 查看各格式的规格。""" print_header("演示 1：各种浮点格式的规格") # 【PyTorch 小知识】torch.finfo(dtype) 返回浮点类型的规格： # max ：能表示的最大有限值 # smallest_normal：最小的"正规"正数（再小就进入精度更差的"次正规数"，再小就变成 0） # eps ：1 与"比 1 大的下一个可表示数"之间的差，衡量相对精度 print(f" {'格式':<10}{'位数':>4}{'指数':>5}{'尾数':>5}{'最大值':>14}{'最小正规数':>14}{'eps(相对精度)':>16}") for name, dt, e, m in FORMATS: fi = torch.finfo(dt) print(f" {name:<10}{fi.bits:>5}{e:>6}{m:>6}{fi.max:>15.4g}{fi.smallest_normal:>14.4g}{fi.eps:>16.4g}") print(""" 读表： - FP16 最大只能到 65504，最小正规数约 6e-5 → 范围很窄，容易上溢/下溢。 - BF16 的范围和 FP32 几乎一样（都是 8 位指数），但 eps 约 0.0078 → 只有 2~3 位有效数字。 - FP8 E4M3 精度稍好、范围小（最大 448）；E5M2 范围大（最大 57344）、精度更差。 常见用法：前向传播的权重和激活用 E4M3，反向传播的梯度（数值范围更大）用 E5M2。""") # ============================================================================= # 第 2 部分：亲眼看看比特 # ============================================================================= def float_bits(value, dtype, exp_bits, man_bits): """ 把一个数转成指定浮点格式，返回它的二进制表示，按"符号 | 指数 | 尾数"分开。 【思路】 PyTorch 的 x.view(整数类型) 可以把同一块内存"当作整数来读"（不做任何数值转换）， 这样就能拿到浮点数的原始比特。 【参数】 value : float，要转换的数 dtype : 浮点 dtype，例如 torch.bfloat16 exp_bits : 该格式的指数位数 man_bits : 该格式的尾数位数 【返回】 str，例如 "0 | 01111111 | 0000000" 【边界】 - 超出范围的数会变成 inf（FP16、BF16、E5M2），或 NaN（E4M3fn 没有 inf，溢出时的行为要看具体实现）。 """ x = torch.tensor(value, dtype=torch.float32).to(dtype) total = 1 + exp_bits + man_bits int_dtype = {32: torch.int32, 16: torch.int16, 8: torch.uint8}[total] raw = x.view(int_dtype).item() raw &= (1 << total) - 1 # 把可能的负数（有符号整数）变成无符号的比特串 s = format(raw, f"0{total}b") return f"{s[0]} | {s[1:1 + exp_bits]} | {s[1 + exp_bits:]}" def demo_bits(): """演示 2：同一个数在不同格式下的比特。""" print_header("演示 2：把浮点数拆成 符号 | 指数 | 尾数") for value in [1.0, -2.5, 0.1]: print(f"\n {value}：") for name, dt, e, m in FORMATS: stored = torch.tensor(value).to(dt).float().item() print(f" {name:<9} {float_bits(value, dt, e, m):<40} 实际存下的值 = {stored!r}") print(""" 观察： - BF16 的比特恰好是 FP32 的前 16 位（符号、指数完全相同，尾数只留前 7 位）。 所以 FP32 → BF16 的转换就是"砍掉后 16 位（并四舍五入）"，非常便宜。 - 0.1 在二进制里是无限循环小数，任何格式都存不准；位数越少，误差越大。""") # ============================================================================= # 第 3 部分：范围 —— 上溢与下溢 # ============================================================================= def demo_range(): """演示 3：很小和很大的数。""" print_header("演示 3：范围 —— 下溢（变成 0）与上溢（变成 inf）") for value in [1e-8, 1e-5, 70000.0, 1e20]: row = f" {value:<8g}" for name, dt, _, _ in FORMATS[:3]: stored = torch.tensor(value).to(dt).float().item() row += f" {name}: {stored:<12.4g}" print(row) print(""" - 1e-8 在 FP16 里直接变成了 0（下溢）。梯度常常就是这么小的数： 用 FP16 训练时，很多梯度会"消失"，所以 FP16 训练需要"损失缩放"（loss scaling）： 先把损失乘以一个大数（比如 1024），让梯度变大到 FP16 能表示的范围，更新前再除回去。 - 70000 在 FP16 里变成了 inf（上溢）。激活值或注意力分数偶尔很大时就会出问题。 - BF16 的范围和 FP32 一样，这两个问题都不存在 → 不需要损失缩放。这是 BF16 流行的主要原因。""") # ============================================================================= # 第 4 部分：精度 —— 数与数之间的间隔 # ============================================================================= def next_representable_gap(value, dtype): """ 返回在 dtype 中，value 附近两个相邻可表示数之间的间隔（ULP，unit in the last place）。 【思路】 浮点数不是均匀分布的：离 0 越远，相邻两个数之间的间隔越大。 间隔大约等于 |value| × eps。 用 torch.nextafter 找到"下一个"可表示数，与当前数相减，就是间隔。 【边界】value 超出 dtype 范围时结果无意义。 """ # nextafter 只支持 float32/float64，所以对低精度格式我们用"间隔 = 2^(指数) × eps"来算 fi = torch.finfo(dtype) x = torch.tensor(value, dtype=torch.float32).to(dtype).float() exponent = torch.floor(torch.log2(x.abs())).item() return (2.0 ** exponent) * fi.eps def demo_precision(): """演示 4：精度。""" print_header("演示 4：精度 —— 1 + 0.001 等于几？") for small in [1e-3, 1e-2]: row = f" 1 + {small:<6g}" for name, dt, _, _ in FORMATS[:3]: r = (torch.tensor(1.0, dtype=dt) + torch.tensor(small, dtype=dt)).float().item() row += f" {name}: {r:<12.8g}" print(row) print("\n BF16 中相邻可表示数之间的间隔：") for v in [1.0, 10.0, 100.0, 1000.0]: print(f" 在 {v:>6g} 附近：{next_representable_gap(v, torch.bfloat16):g}") print(""" BF16 在 1 附近的间隔是 1/128 ≈ 0.0078，所以 1 + 0.001 会被四舍五入回 1。 对神经网络来说，单个数算得"不太准"通常没关系（网络对噪声很鲁棒）； 但如果把大量小数"累加"起来，误差会积累成灾难 —— 见下一个演示。""") # ============================================================================= # 第 5 部分：累加误差 # ============================================================================= def naive_accumulate(value, n, dtype): """ 用指定精度的累加器，把 value 加 n 次（一个一个加，每一步都用 dtype 存结果）。 【思路】 模拟"用低精度做累加"：每加一次，结果都被舍入到 dtype 能表示的最近值。 当累加和变大后，相邻可表示数的间隔超过 2 × value，再加 value 就会被舍入回原值， 累加和就"卡住"不动了。 【参数】 value : float，每次加的数 n : int，加的次数（>= 0） dtype : 累加器的数据类型 【返回】 float，最终的累加和 """ acc = torch.tensor(0.0, dtype=dtype) v = torch.tensor(value, dtype=dtype) for _ in range(n): acc = acc + v return acc.float().item() def demo_accumulation(): """演示 5：累加误差，以及为什么累加要用 FP32。""" print_header("演示 5：把 0.01 加 10000 次（正确答案 100）") for name, dt in [("FP32", torch.float32), ("FP16", torch.float16), ("BF16", torch.bfloat16)]: print(f" 用 {name} 逐个累加：{naive_accumulate(0.01, 10000, dt):.4f}") # PyTorch 的 sum 在内部会用更高的精度累加，即使输入是 BF16 x = torch.full((10000,), 0.01, dtype=torch.bfloat16) print(f"\n torch.full((10000,), 0.01, dtype=bfloat16).sum() = {x.sum().float().item():.4f}" f" ← 库函数内部用 FP32 累加，结果接近 100") print(f" （注意：0.01 本身存成 BF16 就已经不准了，存下的是 {torch.tensor(0.01, dtype=torch.bfloat16).item():.6f}，" f"所以乘以 10000 约为 {torch.tensor(0.01, dtype=torch.bfloat16).item() * 10000:.2f}）") print(""" BF16 累加到 4 左右就卡住了：在 [4, 8) 之间，相邻数的间隔是 1/32 ≈ 0.031， 而 0.01 小于间隔的一半，加上去会被四舍五入回原值。FP16 也在 32 左右卡住。 【大模型联系】这就是"混合精度训练"的核心原则： - 矩阵乘法的输入可以是 BF16，但乘积的累加要在 FP32 中进行（GPU 的 Tensor Core 硬件就是这么做的）； - 求和、均值、方差（LayerNorm/RMSNorm）、softmax 的分母，都要用 FP32； - 优化器更新参数时，θ ← θ - lr × 更新量，更新量往往比 θ 小几个数量级， 如果 θ 用 BF16 存，很多更新会直接被舍入掉 → 所以要保留一份 FP32 的"主权重"（master weights）。""") def demo_weight_update(): """演示 5b：参数用 BF16 存时，小的更新会丢失。""" print_header("演示 5b：用 BF16 存参数，小更新会被吞掉") w32 = torch.tensor(1.0, dtype=torch.float32) w16 = torch.tensor(1.0, dtype=torch.bfloat16) update = 1e-4 # 学习率 × 梯度，通常就是这个量级甚至更小 for _ in range(1000): w32 = w32 - update w16 = w16 - torch.tensor(update, dtype=torch.bfloat16) print(f" 从 1.0 开始，减去 1e-4 共 1000 次（正确答案 0.9）：") print(f" FP32 参数：{w32.item():.6f}") print(f" BF16 参数：{w16.float().item():.6f} ← 一次都没减下去") # ============================================================================= # 第 6 部分：混合精度 autocast # ============================================================================= def demo_autocast(): """演示 6：torch.autocast 自动为不同运算选择精度。""" print_header("演示 6：混合精度 —— torch.autocast") torch.manual_seed(0) a = torch.randn(256, 256) b = torch.randn(256, 256) ref = a @ b # FP32 参考答案 # 【PyTorch 小知识】with torch.autocast(device_type, dtype): 在这个代码块里， # PyTorch 会按一张"规则表"自动选择精度： # - "适合低精度"的运算（矩阵乘法、卷积）用 dtype（这里是 BF16）计算； # - "需要高精度"的运算（例如损失函数）提升到 FP32。 # 你不用手动转换类型。 # 【边界】这张规则表因设备而异：在 CUDA 上，softmax、log_softmax、求和等也会被提升到 FP32； # 在 CPU 上，规则表更短（softmax 仍是 BF16，但交叉熵等损失函数会提升到 FP32）。 # 所以写训练代码时，关键的归约运算最好自己显式地用 .float() 转成 FP32，不要完全依赖 autocast。 target = torch.randint(0, 256, (256,)) with torch.autocast(device_type="cpu", dtype=torch.bfloat16): c = a @ b s = torch.softmax(c, dim=-1) loss = torch.nn.functional.cross_entropy(c, target) print(f" 输入 dtype：{a.dtype}") print(f" autocast(CPU) 下 a @ b 的输出 dtype： {c.dtype} ← 矩阵乘法用 BF16") print(f" autocast(CPU) 下 softmax 的输出 dtype： {s.dtype} ← CPU 规则表里 softmax 不提升（CUDA 上会提升为 FP32）") print(f" autocast(CPU) 下 cross_entropy 的输出 dtype： {loss.dtype} ← 损失函数提升到 FP32") rel_err = ((c.float() - ref).norm() / ref.norm()).item() print(f" BF16 矩阵乘法相对 FP32 的相对误差：{rel_err:.2e}（与 BF16 的精度 2^-8 ≈ 3.9e-3 同一量级）") # ============================================================================= # 第 7 部分：FP8 与缩放；FP4 与分块缩放 # ============================================================================= def quantize_fp8(x, scale): """ 模拟 FP8（E4M3）量化：x / scale → 转成 FP8 → 转回 FP32 → × scale。 【思路】 FP8 E4M3 的可表示范围很窄（约 ±448，最小正规数约 0.0156）。 先除以一个缩放因子 scale，把数值"挪"到 FP8 擅长的范围内，量化后再乘回来。 scale 通常取 max(|x|) / 448，让最大的数正好落在 FP8 能表示的最大值附近。 【参数】 x : FP32 张量 scale : float 或可广播到 x 的张量；scale = 1 表示不缩放 【返回】 FP32 张量，量化再反量化后的值（与 x 形状相同） """ return (x / scale).to(torch.float8_e4m3fn).float() * scale FP4_GRID = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]) # FP4 E2M1 能表示的非负值 def quantize_fp4(x, block_size=None): """ 模拟 FP4（E2M1）量化，可选"分块缩放"。 【思路】 FP4 只有 4 个比特：连正负号在内一共只能表示 15 个不同的数： ±{0, 0.5, 1, 1.5, 2, 3, 4, 6}（+0 和 -0 算一个） 这么粗糙的格式想要可用，关键在于缩放因子： - 整体缩放（block_size=None）：整个张量共享一个缩放因子 = max|x| / 6； - 分块缩放（block_size=16）：每 16 个相邻元素共享一个缩放因子（NVIDIA NVFP4 的做法）。 分块越小，每一块的缩放因子越能"贴合"这一块的数值范围，一个异常大的值（离群值） 只会影响它所在的那一小块，而不会拖累整个张量。 【参数】 x : 一维 FP32 张量，长度必须能被 block_size 整除 block_size : None 或 int 【返回】 FP32 张量，量化再反量化后的值 【边界】 - 真实的 NVFP4 中，缩放因子本身还要用 FP8 存储；这里为了简单用 FP32 存缩放因子。 - 全零的块：缩放因子为 0 会导致除以 0，我们把 0 缩放因子替换成 1。 """ if block_size is None: blocks = x.view(1, -1) else: if x.numel() % block_size != 0: raise ValueError("张量长度必须能被 block_size 整除") blocks = x.view(-1, block_size) scale = blocks.abs().amax(dim=1, keepdim=True) / 6.0 scale = torch.where(scale == 0, torch.ones_like(scale), scale) y = blocks / scale # 找到网格中最近的值：比较 |y| 与每个网格点的距离，取最小的 idx = (y.abs().unsqueeze(-1) - FP4_GRID).abs().argmin(dim=-1) q = FP4_GRID[idx] * y.sign() return (q * scale).view_as(x) def relative_error(approx, exact): """相对误差 = ||approx - exact|| / ||exact||。""" return ((approx - exact).norm() / exact.norm()).item() def demo_low_precision_scaling(): """演示 7：FP8 和 FP4 为什么需要缩放因子。""" print_header("演示 7：FP8 / FP4 与缩放") torch.manual_seed(0) # 模拟一组很小的梯度（量级 1e-4） g = torch.randn(4096) * 1e-4 no_scale = quantize_fp8(g, 1.0) scale = g.abs().max().item() / 448.0 with_scale = quantize_fp8(g, scale) print(" FP8 E4M3 量化一组量级为 1e-4 的数：") print(f" 不缩放：{(no_scale == 0).float().mean().item():.0%} 的数变成了 0，相对误差 {relative_error(no_scale, g):.3f}") print(f" 按 max/448 缩放：{(with_scale == 0).float().mean().item():.0%} 的数变成了 0，相对误差 {relative_error(with_scale, g):.3f}") # 带离群值的激活：大部分是正常值，少数几个特别大（大模型激活中很常见） x = torch.randn(4096) x[::512] *= 100.0 # 每 512 个元素放一个 100 倍的离群值 normal = torch.ones_like(x, dtype=torch.bool) normal[::512] = False # 标记"正常值"（非离群值）的位置 q_tensor, q_block = quantize_fp4(x), quantize_fp4(x, 16) print("\n FP4 量化一组带离群值的数（每 512 个元素有一个 100 倍的离群值）：") print(f" 整个张量一个缩放因子：相对误差 {relative_error(q_tensor, x):.3f}，" f"正常值中有 {(q_tensor[normal] == 0).float().mean().item():.0%} 被量化成了 0") print(f" 每 16 个元素一个缩放因子：相对误差 {relative_error(q_block, x):.3f}，" f"正常值中有 {(q_block[normal] == 0).float().mean().item():.0%} 被量化成了 0") print(""" 整体缩放时，缩放因子被离群值"撑大"，正常的数都被挤到 0 附近，大部分直接变成了 0； 分块缩放让离群值只影响它所在的 16 个元素。 【大模型联系】这就是 MXFP8（每 32 个元素一个缩放因子）、NVFP4（每 16 个一个）等新格式的核心思想， 最新的 GPU（如 B200）在硬件上直接支持它们，用来进一步加速训练和推理（第 5 讲）。""") def main(): demo_spec_table() demo_bits() demo_range() demo_precision() demo_accumulation() demo_weight_update() demo_autocast() demo_low_precision_scaling() print("\n全部演示运行完毕。") if __name__ == "__main__": main()