""" ============================================================================= 第 2 讲 · 代码 1：张量与内存 —— 一个张量到底占多少内存？存在哪里？ ============================================================================= 对应讲义：lectures/lecture02_pytorch_resource_accounting.md 第 2 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 深度学习里的一切数据——输入、参数、梯度、优化器状态、中间激活——都是"张量"（tensor）。 要做资源核算，第一步就是搞清楚： 1. 一个张量占多少字节？ → 元素个数 × 每个元素的字节数 2. 张量在内存里是怎么排布的？ → 存储（storage）+ 步长（stride） 3. 哪些操作是"免费"的（视图，不复制数据），哪些会真的复制数据？ 4. 张量在 CPU 还是 GPU 上？搬来搬去有什么代价？ ----------------------------------------------------------------------------- 【整体思路】 ----------------------------------------------------------------------------- 第 1 部分：创建张量；形状、数据类型、设备 第 2 部分：内存 = numel × element_size；用 "meta" 设备做"不花内存的核算" 第 3 部分：存储与步长：张量在内存里其实是一条"一维长带" 第 4 部分：视图 vs 复制：转置、切片、view、reshape、contiguous 第 5 部分：设备：CPU / GPU（CUDA）/ Apple 芯片（MPS） ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 01_tensors_and_memory.py 需要安装 PyTorch（见 code/lecture02/README.md）。 ----------------------------------------------------------------------------- 【注释约定】（全课代码通用） ----------------------------------------------------------------------------- 【思路】【参数】【返回】【边界】【Python 小知识】【大模型联系】 另外本讲开始会出现【PyTorch 小知识】，讲解 PyTorch 的用法。 """ # ============================================================================= # 第 0 部分：导入 # ============================================================================= # 【Python 小知识】 # torch 就是 PyTorch。第一次 import 可能要花一两秒，这是正常的。 import time import torch def print_header(title): """打印一个带分隔线的标题。""" print() print("=" * 70) print(f" {title}") print("=" * 70) def get_device(): """ 自动选择可用的最快设备。 【返回】 torch.device：有 NVIDIA GPU 时是 "cuda"；在 Apple 芯片的 Mac 上是 "mps"；否则 "cpu"。 【PyTorch 小知识】 torch.cuda.is_available() ：是否有可用的 NVIDIA GPU（以及驱动） torch.backends.mps.is_available() ：是否能用 Apple 芯片的 GPU（Metal Performance Shaders） """ if torch.cuda.is_available(): return torch.device("cuda") if torch.backends.mps.is_available(): return torch.device("mps") return torch.device("cpu") def tensor_memory_bytes(x): """ 计算一个张量的数据占多少字节。 【思路】 内存 = 元素个数 × 每个元素的字节数。 这是资源核算最基本的公式，后面所有的"显存估算"都建立在它之上。 【参数】 x : torch.Tensor（任何形状、任何数据类型、任何设备，包括 "meta"） 【返回】 int，字节数。 【边界】 - 这里只算"张量自己能看到的元素"。如果 x 是一个大张量的视图（比如切片）， 它底层共享的存储可能更大，见第 4 部分。 - 不包括 PyTorch 为这个张量记录的元信息（形状、步长等），那部分很小，可以忽略。 【PyTorch 小知识】 x.numel() ：元素总个数（number of elements），等于形状各维相乘 x.element_size() ：每个元素的字节数，float32 是 4，bfloat16 是 2 """ return x.numel() * x.element_size() def human_bytes(n): """ 把字节数转成易读的字符串，例如 2415919104 → '2.42 GB'。 【思路】这里用十进制单位：1 KB = 1000 字节，1 GB = 10^9 字节（硬件厂商标注显存时常用）。 操作系统有时用二进制单位（1 GiB = 2^30 字节 ≈ 1.074 GB），两者差约 7%。 """ for unit in ["B", "KB", "MB", "GB", "TB"]: if abs(n) < 1000: return f"{n:.2f} {unit}" n /= 1000 return f"{n:.2f} PB" # ============================================================================= # 第 1 部分：创建张量 # ============================================================================= def demo_create(): """演示 1：创建张量，查看形状、数据类型、设备。""" print_header("演示 1：创建张量") # 【PyTorch 小知识】几种常用的创建方式 a = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) # 从 Python 列表创建 b = torch.zeros(4, 8) # 全 0 c = torch.ones(2, 3, dtype=torch.bfloat16) # 全 1，指定数据类型 d = torch.randn(2, 16, 64) # 标准正态分布随机数 e = torch.arange(6) # 0, 1, ..., 5（整数） for name, t in [("a", a), ("b", b), ("c", c), ("d", d), ("e", e)]: print(f" {name}: shape={tuple(t.shape)}, 维数(rank)={t.dim()}, dtype={t.dtype}, device={t.device}") print(""" "张量"就是多维数组。维数（rank）： 0 维 = 标量，1 维 = 向量，2 维 = 矩阵，3 维及以上就叫张量。 在 Transformer 里最常见的是 4 维张量，比如注意力里的 (batch, heads, seq, head_dim)。 注意 torch.arange 默认产生整数（int64），torch.randn 默认产生 float32。""") # ============================================================================= # 第 2 部分：内存核算 # ============================================================================= def demo_memory(): """演示 2：计算张量占用的内存，并用 meta 设备做"零成本"核算。""" print_header("演示 2：内存 = 元素个数 × 每个元素的字节数") x = torch.zeros(4, 8) print(f" torch.zeros(4, 8)：{x.numel()} 个元素 × {x.element_size()} 字节 = {tensor_memory_bytes(x)} 字节") print("\n 不同数据类型的单个元素大小：") for dt in [torch.float64, torch.float32, torch.bfloat16, torch.float16, torch.float8_e4m3fn, torch.int64, torch.int32, torch.uint16, torch.bool]: size = torch.empty(0, dtype=dt).element_size() print(f" {str(dt):<22} {size} 字节") # 大模型的一个真实张量：GPT-3 某一层前馈网络的第一个权重矩阵，形状 12288 × 49152 # 【PyTorch 小知识】device="meta" 创建一个"只有形状和数据类型、没有真实数据"的张量。 # 它不占内存，但可以正常查询 numel / element_size，非常适合做资源核算。 # 直接在 CPU 上创建 12288 × 49152 的 float32 张量需要约 2.4 GB 内存，没必要。 w = torch.empty(12288, 49152, device="meta") print(f"\n GPT-3 前馈层权重 12288 × 49152（float32）：{human_bytes(tensor_memory_bytes(w))}") w16 = torch.empty(12288, 49152, device="meta", dtype=torch.bfloat16) print(f" 同一个矩阵用 bfloat16 存：{human_bytes(tensor_memory_bytes(w16))} ← 精度减半，内存也减半") print(" GPT-3 有 96 层，每层有 2 个这样大小的前馈矩阵……光这部分参数就要：", human_bytes(96 * 2 * tensor_memory_bytes(w)), "（float32）") # ============================================================================= # 第 3 部分：存储与步长 # ============================================================================= # 【思路】 # 内存本身是"一维"的：一个个字节排成一条长带。 # 一个 3×4 的矩阵在内存里其实是 12 个数排成一排（默认"按行优先"，即先放第 0 行，再放第 1 行……）。 # # 张量 = 底层存储（storage，一维长带）+ 如何解读它的元信息： # shape（形状） ：(3, 4) # stride（步长） ：(4, 1) —— 沿第 0 维走一步，要在长带上跳 4 个元素；沿第 1 维走一步，跳 1 个 # offset（偏移） ：从长带的第几个元素开始 # # 访问 x[i, j] 的位置 = offset + i * stride[0] + j * stride[1] # # 【大模型联系】 # 理解步长后，你就明白为什么"转置"几乎不花时间（只是交换了步长）， # 也明白为什么转置后的张量有时必须 .contiguous()（真正复制一份）才能继续某些操作。 # 第 5、6 讲写 GPU 核函数时，你要亲手用步长计算每个元素的地址。 def demo_storage_and_stride(): """演示 3：存储、形状、步长。""" print_header("演示 3：张量在内存里是一条一维长带") x = torch.arange(12).view(3, 4) print(f" x = torch.arange(12).view(3, 4)：\n{x}") print(f" x.shape = {tuple(x.shape)}, x.stride() = {x.stride()}") i, j = 1, 2 flat_index = i * x.stride(0) + j * x.stride(1) # 【PyTorch 小知识】x.flatten() 把张量拉平成一维；对连续张量，它的顺序就是内存里的顺序 print(f" x[{i}, {j}] = {x[i, j].item()}，在长带上的位置 = {i}×{x.stride(0)} + {j}×{x.stride(1)} = {flat_index}，" f"长带上该位置的值 = {x.flatten()[flat_index].item()}") assert x[i, j].item() == x.flatten()[flat_index].item() # 三维张量的步长 y = torch.zeros(2, 3, 4) print(f"\n torch.zeros(2, 3, 4).stride() = {y.stride()} ← 最后一维步长为 1，前一维 = 后面各维大小之积") # ============================================================================= # 第 4 部分：视图 vs 复制 # ============================================================================= def same_storage(a, b): """ 判断两个张量是否共享同一块底层存储。 【PyTorch 小知识】 t.untyped_storage().data_ptr() 返回底层存储的起始内存地址。 两个张量的这个地址相同，说明它们是同一块内存的不同"视图"。 """ return a.untyped_storage().data_ptr() == b.untyped_storage().data_ptr() def demo_views_and_copies(): """演示 4：哪些操作只是换个视角看同一块内存，哪些会真的复制数据。""" print_header("演示 4：视图（免费）vs 复制（要花内存和时间）") x = torch.arange(12, dtype=torch.float32).view(3, 4) # (a) 切片是视图 row = x[1] # 第 1 行 col = x[:, 1] # 第 1 列 print(f" x[1] 是视图吗？{same_storage(x, row)}；x[:, 1] 是视图吗？{same_storage(x, col)}") print(f" x[:, 1] 的步长 = {col.stride()}（在长带上每隔 4 个取一个）") # 视图共享数据：改视图，原张量也会变！ row[0] = 100.0 print(f" 执行 row[0] = 100 之后，x[1, 0] = {x[1, 0].item()} ← 原张量也被改了（常见 bug 来源）") x[1, 0] = 4.0 # 改回来 # (b) 转置是视图：只交换了 shape 和 stride xt = x.T print(f"\n x.T：shape={tuple(xt.shape)}, stride={xt.stride()}，共享存储={same_storage(x, xt)}，" f"是否连续={xt.is_contiguous()}") # (c) view 要求内存连续；不连续时会报错 try: xt.view(-1) except RuntimeError as e: print(f" 【边界】x.T.view(-1) 报错：{str(e)[:60]}...") # (d) reshape：能用视图就用视图，否则复制 r1 = x.reshape(-1) r2 = xt.reshape(-1) print(f" x.reshape(-1) 共享存储：{same_storage(x, r1)}；x.T.reshape(-1) 共享存储：{same_storage(x, r2)}（悄悄复制了）") # (e) contiguous：不连续时复制成连续的 xc = xt.contiguous() print(f" x.T.contiguous()：stride={xc.stride()}，是否连续={xc.is_contiguous()}，共享存储={same_storage(x, xc)}") # (f) 视图"免费"，复制要花时间：用一个大张量计时 big = torch.randn(4096, 4096) t0 = time.perf_counter() for _ in range(100): _ = big.T # 视图 t_view = (time.perf_counter() - t0) / 100 t0 = time.perf_counter() for _ in range(10): _ = big.T.contiguous() # 复制 t_copy = (time.perf_counter() - t0) / 10 print(f"\n 4096×4096 张量：转置（视图）约 {t_view * 1e6:.1f} 微秒；转置并复制约 {t_copy * 1e3:.1f} 毫秒") print(" → 视图与张量大小无关；复制要读写全部 64 MB 数据。") print("\n 【大模型联系】显存核算时，视图不额外占显存；但 .contiguous()、.clone()、以及大多数运算的输出都会。") # ============================================================================= # 第 5 部分：设备 # ============================================================================= def demo_devices(): """演示 5：CPU 与加速器之间搬运数据。""" print_header("演示 5：设备（CPU / CUDA / MPS）") device = get_device() print(f" 本机最快的可用设备：{device}") x = torch.randn(1024, 1024) # 默认在 CPU 内存里 print(f" torch.randn(...) 默认在：{x.device}") if device.type == "cpu": print(" 没有检测到 GPU。本讲的所有代码都能在 CPU 上运行，只是某些计时演示会慢一些。") return # 【PyTorch 小知识】x.to(device) 把数据复制到指定设备，返回新张量（原张量不变） y = x.to(device) print(f" x.to('{device.type}') 之后：{y.device}") # GPU 上的运算是"异步"的：Python 发出命令后立刻返回，GPU 在后台执行。 # 计时前必须"同步"（等 GPU 真的算完），否则测到的只是"发命令"的时间。第 6 讲会详细讲。 sync = torch.cuda.synchronize if device.type == "cuda" else torch.mps.synchronize sync() t0 = time.perf_counter() for _ in range(10): _ = x.to(device) sync() t_move = (time.perf_counter() - t0) / 10 print(f" 把 4 MB 的张量从 CPU 复制到 {device.type}：每次约 {t_move * 1e3:.2f} 毫秒") print(" 【大模型联系】训练时要让数据和参数一直留在 GPU 上，尽量减少 CPU↔GPU 之间的搬运。") print(" （在 NVIDIA 机器上，CPU 与 GPU 之间通过 PCIe 连接，带宽比 GPU 显存带宽低一两个数量级。") print(" Apple 芯片的 CPU 和 GPU 共用同一块内存，搬运代价小得多，但仍然不是零。）") def main(): torch.manual_seed(0) demo_create() demo_memory() demo_storage_and_stride() demo_views_and_copies() demo_devices() print("\n全部演示运行完毕。") if __name__ == "__main__": main()