MegaTrain 是一个RAM-centric（内存中心）的单机大模型训练框架：它把 100B+ 参数大语言模型的参数和优化器状态存放在主机内存（CPU RAM）中，只把 GPU 当作"临时计算引擎"，逐层把权重流式搬上显存。这样一来，一张 H100 80GB 就能做 GPT-OSS 120B 这类模型的全精度（Full Precision）训练，无需模型切分、无需多机集群。

一、单卡训 120B，卡在哪？

传统方案（DeepSpeed ZeRO-3 等）把参数、梯度、优化器状态切到多卡/多机，通信开销和工程复杂度随规模飙升。而单卡方案的瓶颈非常直白：

• 显存装不下参数：120B 模型仅 BF16 权重就约 240 GB，是 80GB 显存的 3 倍；
• 优化器状态更贵：Adam 需要 FP32 主权重 + 一阶/二阶动量，约为参数量的 12 倍（约 12 GB / 1B 参数）；
• 通信不可接受：CPU↔GPU 之间频繁同步会吃掉所有算力。

MegaTrain 的破题思路一句话概括：既然显存装不下，就让参数住在内存里，GPU 上永远只保留"正在算的那一层"。

二、RAM-centric 核心思想：CPU 是"主人"，GPU 是"引擎"

MegaTrain 的架构角色划分非常清晰：

其中 CPU 侧的 infinity/optimizer.py 定义了ParameterState：每个参数都持有 FP32master、grad、m、v，训练步更新的是内存中的主权重，GPU 上只是"用完即走"的工作副本。

配合DeepSpeed CPUAdam（SIMD 加速的 CPU 端优化器），优化器步比 PyTorch AdamW 快5-7 倍——因为优化器恰好跑在数据所在的 CPU 内存里，省掉了最贵的一次参数搬运。

三、逐层流式流水线：120B 是如何"流过"一张卡的

MegaTrain 的核心实现在 infinity/model/cpu_master.py 的CPUMasterModel类（L337）。它用一条双缓冲异步流水线让"搬参数"和"算参数"完全重叠。

1. 双缓冲（Double Buffering）

GPU 侧按"参数结构组"建立两层缓冲（L523-L627）：

• gpu_flat_buffers：两块固定大小的显存平铺缓冲，搬第 N+1 层权重时，第 N 层还在上面做前向/反向；
• layer_pinned_flats：每层权重预先展平并锁页（pinned）到内存，DMA 传输走最快路径（L446-L456）。

2. 三条 CUDA 流，各司其职

每个 GPU 上下文维护三条流（L568-L571）：

流之间用 CUDA Event 精确同步（weight_ready_events、buffer_free_events等），保证"数据没到不计算、缓冲没空不写入"，全程无阻塞等待。

3. K-slab 梯度池：反传也不等

反向时梯度不能逐层阻塞地搬回内存。MegaTrain 建了一个K 块（默认 12 块）平铺梯度 slab 池（L584-L595）：梯度落到空闲 slab 后立刻释放 GPU 层模板，再由一个独立的CPU 工作线程（L784-L823）把 slab 内容累加到 FP32 主梯度上。"GPU 反传 → 梯度 DMA → CPU 累加"三级完全流水化。

4. 检查点重计算（Gradient Checkpointing）

默认每 4 层重算一次激活（checkpoint_interval: 4），把显存从"存所有层激活"压到"只存一个检查点窗口"，是 120B 能塞进 80GB 的关键开关之一。

流水线总效果：GPU 上任意时刻只有1~2 层权重 + 少量激活，官方实测训练期显存占用仅约4-9 GB 临时开销，其余显存可留给推理引擎或其他组件。

5. 模块速查：流式调度的完整拼图

四、异构架构也通用：混合注意力与 MoE

120B 级模型往往不是"整齐划一"的层。MegaTrain 通过按参数结构对层分组（_group_layers_by_structure）解决：

• 同构层（如 Llama 的每个 Transformer 层）共享同一套 GPU 模板，缓冲复用；
• 混合注意力（Qwen3.5 的 linear + full attention）和MoE 层（Mixtral、DeepSeek、Qwen3-Next）各自成组，互不干扰；
• 模型结构通过 HuggingFace 属性内省自动发现（_discover_model_components），任何AutoModelForCausalLM加载的 decoder-only 模型（含 VLM）零代码改动即可训练。

五、多卡数据并行：4 卡 = 4 倍加速，但不需要 NCCL

MegaTrain 的多卡扩展（infinity/model/mp_worker.py + infinity/model/mp_state.py）采用spawn 进程 + 共享内存：

• 每张卡一个独立 worker 进程，权重从共享内存各自读取，梯度带锁累加回共享区；
• 全程没有 NCCL 集合通信，因此 4× H100 实现了 4× 线性加速；
• 单卡与多卡共用同一套_GPUContext抽象，配置里加一行devices: [0,1,2,3]即可切换（参考 examples/sft/configs/qwen_7b_4gpu.yaml）。

六、不止 SFT：单卡 GRPO 强化学习闭环

MegaTrain 已完整接入 VERL 框架（仓库内 verl/ 目录），可在同一张 H100上跑单卡 GRPO 训练：

三个组件共存且无需重新加载权重：实测 Qwen2.5-7B 单卡 GRPO 约 60s/step（~120 tok/s），Qwen3.5-27B 约 230s/step。ref_in_actor指针交换技巧还能以零成本计算参考模型 log-prob。参考脚本：examples/rl/run_qwen2_5_7b_megatrain.sh、examples/rl/run_qwen3_5_27b_megatrain.sh。

七、3 分钟上手：安装与训练命令

# 克隆并安装（自动构建 CUDA 流水线扩展） git clone https://gitcode.com/gh_mirrors/me/MegaTrain cd MegaTrain pip install -e . # 第一步（官方强烈建议）：用资源计算器算出最优 batch_size python scripts/calc_resource.py # SFT 训练（25+ 份现成 YAML 配置） python examples/sft/train.py --config examples/sft/configs/qwen3_5_27b.yaml # 单卡 GRPO 强化学习 CUDA_VISIBLE_DEVICES=0 bash examples/rl/run_qwen2_5_7b_megatrain.sh

数据侧采用 LlamaFactory 风格的注册表：在 data/dataset_info.json 里登记数据集，YAML 里按名字引用即可，支持 alpaca / sharegpt 格式，详见 data/README.md。

八、关键参数速查表

总结

MegaTrain 把"大模型训练"重新定义成了一道内存带宽与流水线的调度题：

1. RAM-centric：参数/优化器状态常驻主机内存，GPU 只留当前层；
2. 双缓冲 + 三 CUDA 流：搬、算、收梯度三线并行，GPU 不空转；
3. 结构分组 + K-slab 池：混合注意力、MoE、VLM 全部通用；
4. 无 NCCL 多卡 + VERL 单卡 RL：从 SFT 到 GRPO 一条技术栈打通。

对只有单卡却想动 100B+ 参数的人来说，这套"流式搬运"设计目前是全精度训练下绕不开的参考实现。