acytoo

大模型入门教程3 - 算子加速与原生多模态：用 12.5M 参数同时算中、英、钢琴曲

#算子加速与原生多模态：用 12.5M 参数同时算中、英、钢琴曲

系列导读：前两篇（第一篇：从 Bigram 到 GPT-2、第二篇：用 GPT-2 算钢琴曲）我们用最直白的代码手搓了 GPT-2，并用它分别算出了莎士比亚文本和 MIDI 钢琴曲。但之前的写法偏向教学演示，每一步张量操作拆得太碎，在 GPU 上计算效率很低。本篇我们先切回 Andrej Karpathy 的标准 nanoGPT 实现，通过合并 QKV 矩阵与引入 FlashAttention 解决算子效率问题；随后不改一行网络结构，直接把中文、英文、钢琴曲拼进同一张词表，训练一个仅 1,251 万（12.51M）参数的原生多模态模型，跑通跨模态的“文生音乐”与“听曲写文”。

#1. 算子提速：为什么之前的教学代码跑得慢？

在 第一篇 的 SimpleLLM 里，为了让代码和数学公式一一对应，我们把自注意力拆成了多个独立线性层和显式的矩阵乘法。这种写法逻辑清晰，但在现代 GPU 上却撞上了 显存读写带宽(Memory Wall) 与 CUDA Kernel 启动开销。

对比教学代码，Karpathy 的 nanoGPT 在数学结果完全等价的前提下，核心做了三步底层重构：

  1. QKV 矩阵合并（c_attn）： 不再用三个独立的 nn.Linear(d, d) 分别计算 $Q, K, V$，而是合并为一个 nn.Linear(d, 3 * d)，单次大矩阵乘法算完后再用 .split(n_embd, dim=2) 切开，并直接变形为 (B, nh, T, hs) 的四维张量并行计算所有注意力头。这直接砍掉了 $2/3$ 的小矩阵 Kernel 启动开销，减少了对输入张量的重复读取。
  2. FlashAttention 替代手写注意力： 手写版的 q @ k.T $\to$ masked_fill $\to$ softmax $\to$ @ v，每一步都要把形状为 (B, nh, T, T) 的中间矩阵写回慢速显存（HBM）再重新读出。改用 PyTorch 原生的 F.scaled_dot_product_attention(..., is_causal=True) 后，底层通过分块计算（Tiling）在 GPU 极快的片上缓存（SRAM）里一次性算完因果注意力，完全不在显存中实例化 $T \times T$ 矩阵。
  3. 词表对齐与权重共享（Weight Tying）： 输入端 Embedding 表 wte 与输出层 lm_head 直接指向同一块权重内存；同时将词表大小向上填充（Pad）到 64 的整数倍（如将 GPT-2 的 50257 补齐为 50304），对齐 GPU Tensor Core 的硬件分块边界。

在第一个 Colab 中，我们直接把 HuggingFace 上的 OpenAI gpt2-large（774M）官方预训练权重加载进这套不到 300 行的代码里，在免费 T4 GPU（FP16）上测试续写，验证了架构实现的完全对齐。

#2. 怎么把中文、英文和钢琴曲塞进同一个模型？

第二篇 已经证明：对 Transformer 而言，文字和琴键都只是一维整数编号（Token IDs）。那么要让一个模型同时学会中文、英文和音乐，甚至在它们之间互相翻译，需要设计专门的“音乐编码器”或“语言解码器”吗？

完全不需要。只需把所有模态的离散符号排进同一张词表（vocab_size = 4,928），再用几个特殊控制符标记模态切换：

  • 0 .. 307（MIDI 钢琴事件，308 个 Token）：完全复用第二篇的编码: 0..87 为按下琴键 NOTE_ON，88..175 为松开琴键 NOTE_OFF，176..275 为时间步进 TIME_SHIFT（每步 10ms），276..307 为 32 档按键力度 VELOCITY。
  • 308 .. 311（模态控制符，4 个 Token）：
    • 308（TOK_TEXT）：标记接下来是文本模态；
    • 309（TOK_MUSIC）：标记接下来是音乐模态；
    • 310（TOK_EOS）：序列结束符；
    • 311（TOK_UNK）：未知字符占位符。
  • 312 .. 4887（中英文字符，4,576 个 Token）：将训练语料（《西游记》中文、Little Brother 英文小说、中英文乐评描述）中出现的所有唯一字符按序编号。
  • 4888 .. 4927（对齐补位，40 个 Token）：将原始词表 4,888 向上补齐到 64 的整数倍 4,928。

经过这层统一编码，无论是单模态续写还是跨模态生成，在模型眼里全都变成了同一套一维整数接龙：

  • 单模态文本续写：[308(TOK_TEXT), 汉字/英文字符序列...]
  • 单模态钢琴生成：[309(TOK_MUSIC), MIDI 事件序列..., 310(TOK_EOS)]
  • 文生音乐（Text $\to$ Music）：[308(TOK_TEXT), 音乐描述字符..., 309(TOK_MUSIC), MIDI 事件序列..., 310(TOK_EOS)]
  • 听曲写文（Music $\to$ Text）：[308(TOK_TEXT), “用中文描述这段音乐：”, 309(TOK_MUSIC), MIDI 事件序列..., 308(TOK_TEXT), 音乐描述字符..., 310(TOK_EOS)]

在跨模态任务中，为了防止长前缀（比如两三百个 MIDI 音符）稀释目标输出的梯度，我们在构造标签 y 时将分界符之前的 Prompt 前缀全部设为 -1（利用 F.cross_entropy 的 ignore_index = -1 跳过损失计算），让反向传播 100% 集中在目标模态的生成上。

#3. 实验结果：12.51M 参数的五种生成模式

我们保持 6 层、6 头、隐藏层维度 384 的迷你配置（将上下文窗口 block_size 扩大到 512 以完整装下音文配对样本），模型总参数量为 12.51M。

在中文《西游记》（65 万 Token）、英文小说（50 万 Token）、12 首肖邦钢琴曲（6.3 万 MIDI Token）以及 488 组“音乐片段—中英文描述”配对数据上混合训练 2,000 步后，同一个权重文件直接跑通了五项任务：

  1. 中文续写：给定 [TOK_TEXT] + "悟空道："，模型自动限制在汉字区间生成《西游记》风格的对话；
  2. 英文续写：给定 [TOK_TEXT] + "Jack looked"，模型自动切换到英文单词与语法分布；
  3. 无条件弹奏肖邦：仅给定单个起始符 [TOK_MUSIC]，模型连续生成的肖邦风格钢琴片段如下：
  1. 文生音乐（Text $\to$ Music）：输入一段中文情绪与节奏描述并以 [TOK_MUSIC] 结尾，自注意力直接跨模态检索前文汉字特征，生成对应风格的钢琴 MIDI 序列；
  2. 听曲写文（Music $\to$ Text）：输入一段验证集中未见过的钢琴 MIDI 序列并以 [TOK_TEXT] 结尾，模型回顾整段音符的高低起伏与时间密度，自动输出对应的中文乐评描述。

#4. 小结：所谓的“多模态理解”，依然是共享空间里的条件概率

把这三篇实验连起来看，原生多模态大模型的底层逻辑其实非常朴素：

  1. 模态切换的本质是子空间路由：模型内部没有任何 if modality == "music" 的分支代码。所有中文字、英文字母和钢琴琴键共享完全相同的 6 层注意力与 MLP 权重。它之所以不会在写英文时突然蹦出钢琴音符，纯粹是因为训练数据决定了 TOK_TEXT 引导的上下文对 MIDI 区间（0..307）的条件概率趋近于零（推理时再配合一层简单的 token_mask 即可完全锁死输出模态）。
  2. 跨模态对齐就是跨区间的 Attention 检索：当执行“文生音乐”或“听曲写文”时，QKV 自注意力机制做的事情和单语言翻译没有任何区别，只不过这次 Query 来自 312..4887 号文字区间，而去检索的 Key 和 Value 位于 0..307 号 MIDI 区间。
  3. 从“能跑”走向现代架构：到这里，我们已经把 2019 年的 GPT-2 架构压榨到了极致。但正如 第一篇 末尾所留下的悬念，这套标准 GPT-2 架构在面对更长上下文与更大参数规模时，依然受困于绝对位置编码 wpe、Dense MLP 的算力绑定以及无 KV Cache 的推理瓶颈，这正是 下一篇（分词器与位置编码） 及后续章节我们要逐一拆解的现代 LLM 架构演进。

相关文章：