1.Overview

2.Resource Accounting

2.1 Memory

2.1.1 Tensor Concept

张量是常见的数据结构的拓展：例如：标量：零维张量；向量：一维张量；矩阵：二维张量。

在机器学习中，之所以使用张量，是因为：1.硬件加速。硬件对张量有专门的设计，适配GPU进行快速计算（谷歌还专门设计了TPU）；2.自动求导。设置了requires_grad = True 时，这个张量会记录自己是怎么被算出来的（即记录父节点）。在后续训练中，只需调用.backward()，梯度就能自动沿着张量传播回来；3.广播机制（Broadcasting）。将一个标量（0维张量）与一个矩阵（2维张量）相加，系统会自动把标量“扩展”成同样大小的矩阵再相加。

2.1.2 Types of floats

float32（fp32）：单精度，一个数字32位，4bytes（较大） 

自己研究了一下一个小数的存储方法，有点意思。例如19303.0890是怎么存成一个float32的，这下懂了为啥会出现 0.1 + 0.2 = 0.300000001这种类似的问题了

float16（fp16）：半精度，动态范围较小只有5位，很容易出现下溢到0或上溢到无穷大的现象

bfloat16（bf16）：

fp8：by Nvidia in 2022

混合精度训练（mix-precision-training）：在训练的pipeline中，例如在前向传播、反向传播、优化器或是梯度累积阶段，不同阶段用到的精度可能不一样。

2.2 Compute

import torch

memory_allocated = torch.cuda.memory_allocated()

z = torch.zeros(32, 32, device="cuda:0")

new_memory_allocated = torch.cuda.memory_allocated()

memory_used = new_memory_allocated - memory_allocated

print(memory_used)

结果输出4096 （32 * 32 * 4， 一个float32占4bytes）

2.2.1 Tensor Storage

一个大小为 4 x 4的矩阵，在内存中是一个一维长数组。张量在 PyTorch 中是指向已分配内存的指针。张量对象本身拥有元数据，指定了如何在数组中找到特定的地址。在这个二维矩阵中，元数据具有两个维度：stride[0] 和 stride[1]，分别代表从各自的维度到下一级要跨越多少个单位（即stride[0] 代表每过一行要跨越4个单位，stride[0] = 4; 而 stride[1] 代表每过一列要跨越1个单位，stride[1] = 1）

2.2.2 Tensor Operations

a. Tensor Slicing

张量的很多操作实际上并不会创建一个新的张量，它们只会创建一个不同的视图，而不会进行复制

import torch


def print_storage(x: torch.Tensor):
    print(x.untyped_storage().data_ptr())


x = torch.tensor([[1., 2, 3], [4, 5, 6]])
y1 = x[0]
y2 = x[:, 1]
y3 = x.view(3, 2)
y4 = x.transpose(1, 0)

但 print_storage() 输出的结果一致，即占用同一块内存空间。

x[0, 0] = 100
print(y4[0, 0])          # tensor(100.)
print(y4[0, 0].item())   # 100.0

所以当修改 x 时，其他 Tensor 中对应元素也会修改。

强制复制张量的操作：

1.显式调用 clone()；

2.当操作要求连续内存而视图不连续时调用.contiguous()：某些操作例如转置，会破坏原本张量的连续性（contiguous）。遍历具有连续性的张量意味着能够平滑地遍历存储中的一维数组；反之则是跳跃访问的。

对一个转置后的视图调用 .view() 改变形状，会报错。.contiguous()的作用是把逻辑顺序散乱的数据重新连续排列。这个时候会发生一次复制，对 x 和 y5 进行 print_storage()，得到的结果不一样。

b. Tensor Elementwise

下列操作会得到新的张量，很好理解——计算结果总是要保存的

x = torch.tensor([1, 4, 9])
assert torch.equal(x.pow(2), torch.tensor([1, 16, 81]))
assert torch.equal(x.sqrt(), torch.tensor([1, 2, 3]))
assert torch.equal(x.rsqrt(), torch.tensor([1, 1 / 2, 1 / 3]))  # i -> 1/sqrt(x_i)

assert torch.equal(x + x, torch.tensor([2, 8, 18]))
assert torch.equal(x * 2, torch.tensor([2, 8, 18]))
assert torch.equal(x / 0.5, torch.tensor([2, 8, 18]))


y = torch.ones(3, 3).triu()  # triu()   # 创建上三角矩阵
assert torch.equal(y, torch.tensor([
    [1, 1, 1],
    [0, 1, 1],
    [0, 0, 1]],
))

c. Tensor Matmul

对于一个 16 x 32 和一个 32 x 2的矩阵，得到的结果是一个 16 x 2 的结果。

但在深度学习中，常见的张量形式是一个四维矩阵：w = [a， b， c，d]，其中 a 为 batch，b 为 sequence。c 和 d 为数据。而对 w 的乘法，是将每个 batch 、每个 sequence 都进行一次矩阵的乘法，即：

x = torch.ones(4, 8, 16, 32)
w = torch.ones(32, 2)
y = x @ w    
# @ 只会对张量的“最后两个维度”执行矩阵乘法
# 前面的所有维度都当作“批量维度（Batch）”原样保留。
assert y.size() == torch.Size([4, 8, 16, 2])

2.2.3 Tensor Einops

为了张量代码的可读性，在操作张量时我们命名张量的所有维度，而不像以前那样本质上只依赖于索引

a. jaxtyping basics

# old way : 通过注释等记录
x = torch.ones(2, 2, 1, 3)    # batch seq heads hidden

# eniops way
x: Float[torch.Tensor, "batch seq heads hidden"] = torch.ones(2, 2, 1, 3)

b. einops einsum

本质上还是矩阵乘法

x: Float[torch.Tensor, "batch seq1 hidden"] = torch.ones(2, 3, 4)  
y: Float[torch.Tensor, "batch seq2 hidden"] = torch.ones(2, 3, 4)


z = x @ y.transpose(-2, -1)   # 将倒数第一维度和倒数第二维度转置   
z = einsum(x, y, "batch seq1 hidden, batch seq2 hidden -> batch seq1 seq2")
# 两种方式等价


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

c. einops reduce

对张量的某个维度做“归约（聚合）”操作，比如求和、求平均、取最大值等

import torch
from einops import reduce

x = torch.ones(2, 3, 4)

# 对 hidden 维度求和，消掉它
y = reduce(x, "batch seq hidden -> batch seq", "sum")
print(y.shape)  # torch.Size([2, 3])
print(y)

d. einops rearrange

x: Float[torch.Tensor, "batch seq total_hidden"] = torch.ones(2, 3, 8)
x = rearrange(x, "... (heads hidden1) -> ... heads hidden1", heads=2)

# 得到的结果维度为（2,3,2,4）   from （2,3,8）

2.2.4 Tensor Operation Flops

张量操作本质上是浮点操作（floating-point operation，FLOP），而浮点操作是指任何涉及浮点数的运算，比如加法或乘法

区分 FLOPs 以及 FLOPS。前者小写 s 说明是复数形式，代表浮点操作次数，衡量的是完成的计算总量；后者大写 S 说明是 per second，代表每秒浮点运算次数，衡量的是硬件的速度 。

Linear Model：

假设一个矩阵 M 大小为（B，D），乘以一个矩阵（ D, K）将 D 维向量映射到一个 K 维向量，则：

M = torch.ones(B, D, device=device)
# 要得到的矩阵N 形状为 B x K ，则线性操作矩阵T形状为 D x K
T = torch.randn(D, K, device=device)
N = M @ T

要进行的 FLOPs  = 2 * B * D * K （对于 N 中每个元素，都要计算： D 次乘法以及 D - 1 次加法，共 2D - 1 次操作；而 N 中有 B * K个元素，故 FLOPs  = 2 * B * D * K ）

其他操作的浮点运算次数通常与矩阵或张量的大小呈线性关系。在深度学习中，对于足够大的矩阵，矩阵乘法是所有操作中最耗时的

推广到 Transformer——

B： the number of data points

D * K : the number of parameters

FLOPs for forward pass is (2 * tokens * parameters) (忽略了序列长度和其他因素)

得到了浮点运算的次数，计算实际运行的时间：

def time_matmul(a: torch.Tensor, b: torch.Tensor) -> float:
    """Return the number of seconds required to perform `a @ b`."""

    # Wait until previous CUDA threads are done
    if torch.cuda.is_available():
        torch.cuda.synchronize()
        # PyTorch 在 GPU 上运行张量运算时，是异步的。
        # Python 会立刻把指令 a @ b 发给 GPU，然后不等 GPU 算完，Python 就继续执行下一行代码了。
        # torch.cuda.synchronize() 的作用就是强制 Python 停下来等待, 直到 GPU 上的所有线程全部执行完毕.
        # 保证时间计算准确

    def run():
        # Perform the operation
        a @ b

        # Wait until CUDA threads are done
        if torch.cuda.is_available():
            torch.cuda.synchronize()

    # Time the operation `num_trials` times
    num_trials = 5
    total_time = timeit.timeit(run, number=num_trials)

    return total_time / num_trials

# B = 16384 D = 32768 K = 8192
actual_num_flops = 2 * B * D * K
actual_time = time_matmul(M, T)  # @inspect actual_time
actual_flop_per_sec = actual_num_flops / actual_time

print(actual_num_flops)       # 8796093022208        8.80e12
print(actual_time)            # 1.2285193999996409
print(actual_flop_per_sec)    # 7159913813498.241    7.16e12

模型浮点运算利用率 (Model Flops Utilization，MFU)，即实际浮点次数除以承诺的浮点运算次数

mfu = actual_flop_per_sec / promised_flop_per_sec

查了一下3070的 promised_flop_per_sec = 20.31TFLOPs（但没详细写明是什么float数据格式）

默认利用的float类型为float32， 将其换成为bfloat16，得到的结果：

M = M.to(torch.bfloat16)
T = T.to(torch.bfloat16)

# actual_num_flops = 8.80e12 不变
# actual_time = 0.42
# actual_flop_per_sec = 21.06 TFLOPs

每秒浮点运算次数（FLOP/s）取决于硬件以及数据类型。5090 > 4090，bfloat16 > float32

2.2.5 Gradients FLOPs

整个计算过程中除了计算矩阵乘法之外，还有一部分用于计算梯度

Example:

M = torch.ones(B, D, device=device)
T1 = torch.randn(D, D, device=device, requires_grad=True)
T2 = torch.randn(D, K, device=device, requires_grad=True)
# 一个两层线性模型
# M --T1--> N1 --T2--> N2 -> loss

N1 = M @ T1
N2 = N1 @ T2
loss = N2.pow(2).mean()

前向传播计算量：2 * B * D * D + 2 * B * D * K

计算梯度过程：

forward pass :   2 x ( data points) x ( parameters) FLOPs

backward pass: 4 x ( data points) x ( parameters) FLOPs

total: 6 x ( data points) x ( parameters) FLOPs

3.Model

3.1 Module Parameters

import torch
from torch import nn

input_dim = 16384
output_dim = 32

w = nn.Parameter(torch.randn(input_dim, output_dim))
assert isinstance(w, torch.Tensor)
# 检查 w 是不是 torch.Tensor 的一个实例
# Parameter 在行为上完全是一个张量，可以对它做 w @ x、w.view()、w.transpose() 等所有张量操作

assert type(w.data) == torch.Tensor
# w.data 的类型是精确的 torch.Tensor（而非子类），不带 Parameter 的“身份标记”
# 在训练过程中，应通过 w.data 来修改底层数据（比如做权重裁剪、初始化），而不是重新赋值 w 本身。 
# 否则，w 会丢失 Parameter 的身份，优化器将不再更新它

# nn.Parameter 相较于普通张量 torch.Tensor
# 前者默认 requires_grad = True
# 会被自动添加到模型的 .parameters() 迭代器

parameter initialization

发现问题：

（x的形状为（intput_dim），却能和形状为（input_dim，output_dim）的w相乘，得到形状为（output_dim）的output，是PyTorch默认把一维张量当作行向量来参与矩阵乘法）

得到的结果较大，这是因为数值增长的量级本质上与隐藏维度的平方根成正比

输入数据x的标准差是 1，经过一个全连接层（y = x @ w）后，输出 y 的方差会变成 input_dim × 1，所以 y 的标准差会变成 sqrt(input_dim)，输出结果与标准差量级相同，发生数值发散的爆炸现象，导致模型训练不稳定

解决方法：

w = nn.Parameter(torch.randn(input_dim, output_dim) / np.sqrt(input_dim))
# 将 w 除以输入维度的平方根
# Xavier initialization

得到的结果集中在类似 Normal(0, 1) 的正态分布附近。但由于正态分布曲线永远不会碰到横轴，理论上会出现及其庞大的数据（特别是在大模型参数动辄几十亿数量级的情况下）。故更安全的方法将尾部截断，例如将范围限制在 [-3, 3]

3.2 Custom Model

class Cruncher(nn.Module):
    def __init__(self, dim: int, num_layers: int):
        super().__init__()
        self.layers = nn.ModuleList([
            Linear(dim, dim)
            for i in range(num_layers)
            # Linear(dim, dim)：把输入张量的最后一个维度从 dim 映射到 dim，也就是做一次线性变换
        ])
        # nn.ModuleList 是 PyTorch 里一个专门用来存放多个子模块（如 Linear、Conv2d）的容器
        # 当把层放进 nn.ModuleList 时，PyTorch 会自动“记住”这些层，让优化器能找到它们的参数

        self.final = Linear(dim, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Apply linear layers
        B, D = x.size()
        for layer in self.layers:
            x = layer(x)

        # Apply final head
        x = self.final(x)
        assert x.size() == torch.Size([B, 1])

        # Remove the last dimension
        x = x.squeeze(-1)
        assert x.size() == torch.Size([B])

        return x


D = 64  # Dimension
num_layers = 2
model = Cruncher(dim=D, num_layers=num_layers)
device = get_device()
model = model.to(device)

B = 8  # Batch size
x = torch.randn(B, D, device=device)
y = model(x)
assert y.size() == torch.Size([B])

（当 num_layers = 2 时）这个模型一共有 3 层，两层 Linear 和一个 final 层。前两层的参数量都为 D * D，最后一层的参数量为 D * 1。

3.3 Note about Randomness

# Torch
seed = 0
torch.manual_seed(seed)

# NumPy
import numpy as np
np.random.seed(seed)

# Python
import random
random.seed(seed)

将 Torch、NumPy 和 Python 引入随机性的地方全部“锁死”，让代码每次运行时产生的随机数序列完全一样，增加代码的可复现性。

3.4 Data Loading

data = np.memmap("data.npy", dtype=np.int32)
# "data.npy"：这是一个已存在硬盘上的文件
# np.memmap 的返回值：data 是一个内存映射对象
# 它看起来、用起来和普通 NumPy 数组几乎一模一样（支持索引、切片、算术运算）
# 但数据本身还留在硬盘上，没有加载到内存里。
assert np.array_equal(data, orig_data)

使用 np.memmap() ，适合处理超大规模数据集，内存无法直接 load 整个数据集时。当尝试访问数据时，会按需加载文件

原理： 虚拟内存 + 缺页异常 + 手动管理内存

3.5 Optimizar

class AdaGrad(torch.optim.Optimizer):    # 父类
    def __init__(self, params: Iterable[nn.Parameter], lr: float = 0.01):
        super(AdaGrad, self).__init__(params, dict(lr=lr))
        # 调用 PyTorch 优化器的基类构造函数，把参数和学习率注册进去，这样 self.param_groups 里就有了所有参数的信息

    def step(self):
        for group in self.param_groups:     # self.param_groups 在父类定义，
            lr = group["lr"]
            for p in group["params"]:
                # Optimizer state
                state = self.state[p]
                # 一个字典，用来记录优化器的“状态”，记录每个参数的历史信息
                # 在 AdaGrad 里，它记录的是历史梯度平方的累加和
                grad = p.grad.data  # 在 loss.backward() 反向传播过程中得到

                # Get squared gradients g2 = sum_{i<t} g_i^2
                g2 = state.get("g2", torch.zeros_like(grad))   # g2 为历史梯度平方的累加和

                # Update optimizer state
                g2 += torch.square(grad)    # 把当前梯度的平方累加进去
                state["g2"] = g2
                # g2 的值只会增加，不会减少

                # Update parameters
                p.data -= lr * grad / torch.sqrt(g2 + 1e-5)
                # 如果某个参数的历史梯度平方和 g2 很大（说明这个参数一直在剧烈更新），那么 grad / sqrt(g2) 就会变小，相当于学习率被自动缩小。

optimizer = AdaGrad(model.parameters(), lr=0.01)

x = torch.randn(B, D, device=get_device())
y = torch.tensor([4., 5.], device=get_device())
pred_y = model(x)
loss = F.mse_loss(input=pred_y, target=y)
loss.backward()

optimizer.step()
# 在这里优化器真正发挥作用
state = model.state_dict()
# model.state_dict() 返回一个 Python 字典，里面包含模型所有参数张量的副本
# 注意和优化器里的 state 区分
optimizer.zero_grad(set_to_none=True)
# 释放内存

优化器状态的内存需求：

参数数量 num_parameters ： D * D * num_layers + D

激活值数量 num_activations（前向传播过程中，每一层计算出的中间输出张量）：对于每一层、每一个数据点、每一个维度，都必须存储激活值。即结果为 B * D * num_layers

梯度数量 num_gradients ：梯度数量与参数数量相同，反向传播为每一个参数计算出一个梯度

优化器状态数量 num_optimizer_states ：对于AdaGrad，存储的状态为每个参数对应一个g2，即数量同样与参数数量相同

故总内存量（假设数据格式为Float32，即4个字节）total_memory = 4 * ( num_parameters + num_activations + num_gradients + num_optimizer_states )

FLOPs = 6 * B * num_parameters

3.6 Train Loop

import torch
from torch import nn
from torch.nn import Linear
from torch.optim import SGD
import torch.nn.functional as F

class Cruncher(nn.Module):
    def __init__(self, dim: int, num_layers: int):
        super().__init__()
        self.layers = nn.ModuleList([
            Linear(dim, dim)
            for i in range(num_layers)
        ])
       
        self.final = Linear(dim, 1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # Apply linear layers
        B, D = x.size()
        for layer in self.layers:
            x = layer(x)

        # Apply final head
        x = self.final(x)
        assert x.size() == torch.Size([B, 1])

        # Remove the last dimension
        x = x.squeeze(-1)
        assert x.size() == torch.Size([B])

        return x


def train(name: str, get_batch,
          D: int, num_layers: int,
          B: int, num_train_steps: int, lr: float):
    device = "cuda:0"
    model = Cruncher(dim=D, num_layers=0).to(device)
    # num_layers = 0 说明只有最后的 final 层，输出形状 (B, 1)
    optimizer = SGD(model.parameters(), lr=0.01)

    for t in range(num_train_steps):
        # Get data
        x, y = get_batch(B=B)

        # Forward (compute loss)
        pred_y = model(x)
        loss = F.mse_loss(pred_y, y)

        # Backward (compute gradients)
        loss.backward()

        # Update parameters
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)
        if t % 2 == 0:  # 每2步打印一次
            print(f"Step {t}: loss = {loss.item():.6f}")

    print("Final model weight:")
    print(model.final.weight)
    print("True weight:")
    print(true_w)


D = 16
true_w = torch.arange(D, dtype=torch.float32, device="cuda:0")


def get_batch(B: int) -> tuple[torch.Tensor, torch.Tensor]:
    x = torch.randn(B, D).to("cuda:0")
    true_y = x @ true_w
    return (x, true_y)
# 获取数据对，x 以及 true_y, 用来与训练出的结果 pred_y 对比


train("simple", get_batch, D=D, num_layers=0, B=4, num_train_steps=100, lr=0.01)

num_train_steps 分别设为 10 和 100 跑了两次，结果如下：

3.7 Checkpointing

为了防止训练过程中崩溃，训练结果全部丢失。在训练中途就可以保存训练中间结果。

model = Cruncher(dim=64, num_layers=3).to(get_device())
optimizer = AdaGrad(model.parameters(), lr=0.01)

checkpoint = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
}
torch.save(checkpoint, "model_checkpoint.pt")

loaded_checkpoint = torch.load("model_checkpoint.pt")
# 加载

3.8 Mix Precision Trainig

精度更高，则更准确、更稳定，但成本更高；低精度则相反

a. 对第五个前向传播使用较低的精度，但对其余部分使用 float32

b. PyTorch提供了自动进行混合精度训练（automatic mixed precision，AMP）的工具