如何用单卡训100B+ 大模型

小卡训大模，穷人要翻身！ 今天，我们正式发布MegaTrain训练框架，可以单卡、全精度、全参数、训练100B+大模型，且不掉速。 config遵循Llama Factory 格式，开箱即用。 唯一需要注意的是，批次（batchsize）越大跑得越快，各位可能得改改习惯。 希望能帮助各位摆脱GPU Poor 的困扰。 Github: MegaTrain github.com/DLYuanGod/MegaTrain MegaTrain: Full Precision Training of 100B+ Parameter Large Language Models on a Single GPU arxiv.org/abs/2604.05091

单张H200，训gpt-oss 120b。30% MFU是没怎么优化的速度，需要的话到时候改改kernel。 实话说，我们搞AI的人，小时候一定都逃了不少课。 学数学的时候，除了一句“假设模型满足什么什么条件”，其他严谨的证明、形式逻辑的抽象，都没有学会，以至于落笔，往往皆是“满纸荒唐言，一串不等式。试问审稿人，却道reject依旧“。 学CS 的时候，逃课的少了些，但像经典的如《计算机架构》之类的课，也大多是稀里糊涂混过去的。比如为什么要有不同浮点数格式，为什么要有存储层级，如何设计并行算法，我们都不关心。我们惯用python——tensor一建，矩阵一乘，AGI就有了，谁管那底层具体怎么实现的。 譬如训练大模型。 大多数人，因为吃过CUDA Out of Memory 的苦，还是知道自己的GPU 的显存是有多大。 倘若有心算一算，可以知道，哦，我训的是7B模型，2倍参数（bf16）、2倍梯度（bf16）、外加adam的动量和方差（fp32*2），至少需要至少12倍显存，96GB。外加零零总总的activation 等，刚好填满H200，再大就训不了了。批次大小也只能几千词元。 然而，HBM之上，还有SRAM（L2/L3 Cache），这点Tri Dao在flash attention里给我们补过课。HBM 之下，有各种CPU memory （DDR/LPDDR），还有固态硬盘存储（NVMe SSD），能干什么就不清楚了，只当是无用的配件，放在一旁。 我自也是当中的一员，常因数学出身为借口，不好好学习计算机知识（我数学学得也不好其实）。 然而每日因借不到卡，做不了大模型的后训练，苦恼不已。又遇上Qwen 出神入化的蒸馏，4B 小模型不训极好用，一训就崩，tool call都不会了。我也只好每日骂两句OpenAI的有钱人，说那么多Blackwell，都不给我用，怎能叫Open。 于是小规模跑一跑，或微调或GRPO，凑篇文章投投稿。 可审稿人意见下来，一个个的，都说自己的太小。 此时此刻，那种瘀堵的心情，真想喊声：“却不是特地来消遣洒家？“ 但转念一想，自己当审稿人时，也喜欢说别人的太小。 毕竟，后训练研究的是涌现能力。模型小了，试验就是没意义，有些涌现能力就是跑不出来。 所以，看到拒稿那一刻，几个月心血付诸东流。除了委屈生气，又当如何？ 然而，这一切烦恼都因不理解memory hierarchy。其本是计算机里最重要的大道理之一，小时候上课都讲过。 人类设计过各种各样的存储。速度由快到慢、价格由贵到贱、存储由小到大，实现不同，各有所用。 为了充分利用这些不同的存储，人们往往将他们搭成一个hierarchy（层级）。

添加图片注释，不超过 140 字（可选） 譬如GPU。在每个运算单元里有register（存储器），不过几KB。 周围有SRAM（静态随机存储器），不过几十至一百MB，带宽却能到做到几十TB/s。价格也贵，每GB几千刀 再往外，便是显存了，大家都较为熟悉。 常见的有GDDR，几十GB，带宽1TB/s，涨价前2刀/GB就能买到，游戏卡上配的就是这个。 或者用HBM，可达100GB以上。其本质上是3D堆叠的GDDR，叠了12-16层不等。这个要贵得多，25GB/s。 主机侧的内存，有人也会说cpu内存， 容量可达数TB，带宽却骤降一个数量级。 再往外，还有NVMe SSD，便是所谓固态硬盘。容量以几十TB计，一GB只需几美分。 如何使用这些内存呢？我们一般按照使用信息的频率来安排。 就好像一个人如何布置自己的办公室？ 正在用的东西，拿在手上。 常用的东西，放在桌上，随时伸手可取。 再用得少些的东西，就会放到书柜或者架子上，取用需要走几步路，但仍算方便。 而至于那些几年都不一定用一次的东西，就要打包放仓库了。 科普完memory hierarchy（存储层次结构），终于可以介绍MegaTrain了。 倘若理解了memory hierarchy的伟大思想，MegaTrain的一切设计都是显然的。 要说明，所有的代码都是我的合作者zhengqing写的，我只负责给他打下手+宣传之类。他是我们组的传奇人物，出身一本却代码超群。有机会真应该采访一下他的传奇故事。 首先，如前文所说，大模型训练的内存有三种。persistent storage（持续存储），包括参数、梯度、和优化器状态；激活值，即为了反向传播要存储的中间状态；以及其他一些内存使用。 persistent storage，其随模型变大而变大，合参数量的12倍，这是我们主要处理的内存。 对于模型的参数、梯度、和优化器状态，我们将它全部放在主机内存中，而只把GPU当作一个临时的计算引擎，或者也可理解为高一层的cache，而不是全部信息的存储区。 然后，参数不用一直放卡上，而是用到哪层传哪层。 在这个过程中，产生的梯度也是，从后往前，算出一层下传一层。 至于优化器状态，我们遵循DeepSpeed，根本不上传，而是完全用CPU去运算。 为何呢？首先，优化器状态占了8倍内存，完全没必要触发一次I/O。 若使用adam，其梯度更新不过一串加减乘除，而不涉及矩阵运算。 当然，聪明的小朋友们要问了，muon怎么办？ 实话说没想好。倘若无momentum的muon还好说，直接下传前做完Newton Schwartz就好。用momentum，还得多加两倍的gpu上传momentum。

图为一个block内的前传后传重算设计。下标表示block内的第几层。F为前算，W是权重，G为梯度。开始时，GPU有了第0层的权重，算的同时第一层的权重开始上传，以此类推。等前算算完，开始反向传播时，我们再次上传第一层的权重，以完成block内的重算，然后重算结束开始往下以此传播第3210层的梯度。 另一个问题是，CPU-GPU的通信速度慢很多（比如H200用PCIE好像也就128GB/s），远远慢于HBM的速度（4.8TB/s）。I/O的成本不会很大吗？ 正如前文所说，persistent storage只随参数量变化，而跟批次大小无关。 那么，如果我把批次大小干到特别大，比如几十万上百万词元。那么，平摊到每个词元的I/O成本就很小了。 因为此设计，我们顺道实现了单卡训512k长文本，7B模型。正常fsdp+cp啥的至少得64张卡吧。

添加图片注释，不超过 140 字（可选） 这些做完，我们来处理激活值。 我们决定不下传激活值，毕竟这个数字随批次大小增长。但实际上传与不传都可以，我们只是没优化好，传了会掉速。 那怎么限制其增长呢？就是用非常激进的重算 我们基本是几层一重算，这样内存使用就会少非常多。当然，这么激进的重算也会有额外的计算成本，稍微慢一点，但我想大家都是穷逼不会介意的。 剩下所有的优化就主要是为了通信和计算做重叠了。 比如zhengqing写了很多新的memory management的做法，在cpu和gpu里各划出一片区域。 GPU里叫buffer，里面最多存两层，这样这层算的时候下一层已经在传了，不需要等。 CPU叫gradient slab，也是为了梯度下传提前预留空间，GPU可以Direct Memory Access（DMA）。 另一个改变是，正常pytorch计算需要autograd graph，即记录整个反向传播的路径，很复杂，也会导致scheduling的各种不便利。 我们改成只保留一层的template，然后权重传上来现场绑定，毕竟大部分模型每层长得都差不多。（其实这段我没看懂） 本文介绍了MegaTrain的设计思路。 我们未来会计划持续维护MegaTrain这个库，使得所有的本地训练都能脱离内存的困扰。欢迎大家提需求和意见！ 同时，我们诚挚地欢迎大家来一起为这个开源库添砖加瓦，有兴趣可以私信。 以下是我们准备在近期内支持的项目： RL/DPO 单机多卡（我们不一定准备支持多机） MoE专门优化 Diffusion/视频图像生成 Muon/FlashOptimizer 谢谢大家！