"""第 24 章：TinyTransformer 的参数、矩阵乘法预算和状态内存账本。"""


def parameters(width, vocab=35, context=32, layers=2):
    # 输入嵌入与输出投影不共享；注意力投影无偏置，前馈层有偏置。
    return 2*vocab*width + context*width + layers*(12*width**2+9*width) + 2*width


def matrix_flops(width, batch=16, length=32, vocab=35, layers=2):
    # 一乘一加计 2 FLOPs；只统计密集矩阵乘法，不计逐元素运算。
    forward = layers*(24*batch*length*width**2 + 4*batch*length**2*width)
    forward += 2*batch*length*width*vocab
    return 3*forward  # 前向一次，加两个矩阵输入各自的反向。


BUDGET = 400*matrix_flops(32)


def main():
    print(f"每次训练预算上限：{BUDGET:,} 矩阵 FLOPs")
    for width in [32, 64]:
        count, cost = parameters(width), matrix_flops(width)
        steps = BUDGET//cost
        print(f"宽度 {width}：参数={count:,}；每步={cost:,}；更新={steps}")
        print(f"  预算使用={steps*cost/BUDGET:.2%}；"
              f"处理位置={steps*16*32:,}；真实目标={steps*16*27:,}")
        print(f"  假设 FP32 参数/梯度/Adam 状态：{16*count/1024**2:.3f} MiB")
    n, tokens = 7_000_000_000, 140_000_000_000
    print(f"另一个假设例子：7B 参数、140B 训练词元，6ND≈{6*n*tokens:.2e} FLOPs")
    print(f"同一个 7B 模型：仅 BF16 权重≈{2*n/10**9:.1f} GB；"
          f"按每参数 16 字节的训练状态≈{16*n/10**9:.1f} GB")


if __name__ == "__main__":
    main()