Olmo-core 3 将重点放在开放、可扩展的大规模混合专家模型(Mixture of Experts,MoE)训练基础设施上。相比稠密模型,MoE 并不是简单地“多放几个专家网络”:路由、专家并行、负载均衡、通信拓扑和分布式检查点都会进入训练主路径。训练框架如果不能把这些问题拆开处理,模型规模越大,故障定位和性能调优就越困难。
由于现有摘要没有给出 Olmo-core 3 的具体 API,下面不会假设其配置字段或调用方式,而是从工程角度说明这类基础设施需要解决什么,以及团队可以如何准备接入和验证。
MoE 的扩展瓶颈不只在显存
典型 MoE 层会先由路由器为每个 token 选择少量专家,再把 token 分发给对应专家计算。模型可以拥有很大的总参数量,但单个 token 只激活其中一部分参数。
这种稀疏激活带来了几个新的系统问题:
- 路由不均衡:少数专家收到过多 token,形成热点;其他专家利用率很低。
- 容量溢出:专家容量不足时,部分 token 可能被丢弃、重路由或延迟处理。
- 跨设备通信:专家位于不同 GPU 或节点时,token 分发通常涉及高代价的集体通信。
- 并行维度组合:数据并行、张量并行、流水线并行和专家并行需要共同工作。
- 恢复复杂度:检查点不仅包含模型参数,还可能涉及优化器状态、路由器状态、随机数状态和数据迭代位置。
因此,“可扩展”不能只用最大参数量衡量。更有意义的指标包括单步耗时、有效 token 吞吐量、通信占比、专家负载偏差、溢出率以及故障恢复时间。
训练基础设施应把策略与执行解耦
一个易于演进的 MoE 训练栈,通常需要把以下层次分开:
- 模型定义:专家结构、专家数量、路由器类型和每个 token 的激活专家数。
- 并行规划:每个维度如何切分,专家如何映射到设备与节点。
- 训练策略:精度格式、梯度累积、优化器、学习率和辅助负载均衡损失。
- 运行时能力:通信、激活检查点、混合精度、编译和容错。
- 可观测性:吞吐量、显存、专家利用率、路由熵、丢弃 token 数量和通信等待时间。
这种拆分的价值在于:模型研究者可以调整路由策略,而平台工程师可以替换通信或检查点实现,双方不必反复改动同一段训练循环。
配置也应当显式表达并行拓扑,而不是把关键参数散落在启动脚本中。例如,可以把团队自己的实验配置设计成下面这样。注意:这只是便于迁移的通用示例,不是对 Olmo-core 3 配置格式的描述。
# moe-smoke.yaml
model:
hidden_size: 1024
num_experts: 8
top_k: 2
parallelism:
data_parallel: 2
tensor_parallel: 1
pipeline_parallel: 1
expert_parallel: 2
training:
global_batch_tokens: 65536
precision: bf16
checkpoint_every_steps: 500
monitoring:
log_expert_load: true
log_router_entropy: true
log_token_overflow: true
用一个可运行的预检程序验证拓扑
在启动昂贵训练任务之前,可以先验证配置约束和分布式通信。下面的脚本不实现真正的 MoE 训练,也不调用 Olmo-core 3;它是一个可以直接运行的通用预检工具,用于检查并行规模、路由容量估算和 PyTorch 进程组。
将前面的 YAML 保存为 moe-smoke.yaml,再创建 preflight.py:
import os
import sys
import torch
import torch.distributed as dist
import yaml
def load_config(path: str) -> dict:
with open(path, "r", encoding="utf-8") as f:
return yaml.safe_load(f)
def main() -> None:
config_path = sys.argv[1] if len(sys.argv) > 1 else "moe-smoke.yaml"
cfg = load_config(config_path)
model = cfg["model"]
parallel = cfg["parallelism"]
training = cfg["training"]
num_experts = int(model["num_experts"])
top_k = int(model["top_k"])
expert_parallel = int(parallel["expert_parallel"])
data_parallel = int(parallel["data_parallel"])
global_batch_tokens = int(training["global_batch_tokens"])
if not 1 <= top_k <= num_experts:
raise ValueError("top_k 必须在 1 和 num_experts 之间")
if num_experts % expert_parallel != 0:
raise ValueError("num_experts 必须能被 expert_parallel 整除")
dist.init_process_group(backend="gloo")
rank = dist.get_rank()
world_size = dist.get_world_size()
expected_world_size = (
data_parallel
* int(parallel["tensor_parallel"])
* int(parallel["pipeline_parallel"])
* expert_parallel
)
if world_size != expected_world_size:
raise RuntimeError(
f"world_size={world_size},但配置要求 {expected_world_size} 个进程"
)
# 理想均衡情况下,每个专家每步接收的 token 分配数。
assignments = global_batch_tokens * top_k
assignments_per_expert = assignments / num_experts
probe = torch.tensor([rank + 1.0])
dist.all_reduce(probe, op=dist.ReduceOp.SUM)
if rank == 0:
experts_per_group = num_experts // expert_parallel
print(f"world_size: {world_size}")
print(f"experts per expert-parallel shard: {experts_per_group}")
print(f"ideal assignments per expert: {assignments_per_expert:.0f}")
print(f"all-reduce checksum: {probe.item():.0f}")
print("preflight passed")
dist.destroy_process_group()
if __name__ == "__main__":
main()
安装依赖并按配置中的 4 个进程启动:
python -m pip install torch pyyaml
torchrun \
--standalone \
--nproc-per-node=4 \
preflight.py moe-smoke.yaml
如果机器无法承载 4 个本地进程,可以把 YAML 中的 data_parallel 和 expert_parallel 都改成 1,然后使用 --nproc-per-node=1。接入真实框架时,应把这类检查放在作业启动阶段,让错误在占用集群资源之前暴露出来。
上线前真正值得检查的事项
评估 Olmo-core 3 或其他 MoE 训练基础设施时,不要只验证模型是否能够启动。可以按以下清单做小规模基准:
- 固定模型与 token 数量,对比单卡、单节点和多节点的扩展效率。
- 记录每个专家收到的 token 数,而不只观察平均辅助损失。
- 分离计算时间、路由时间、跨设备分发时间和检查点时间。
- 人为终止一个训练作业,验证恢复后的步数、数据位置和损失曲线是否连续。
- 检查拓扑约束是否在启动前失败,而不是运行数小时后才报错。
- 评估检查点大小、保存频率和恢复时间之间的权衡。
- 确认训练配置、代码版本、容器镜像和数据版本都可追踪。
开放的训练基础设施最重要的价值,不只是代码可见,而是关键行为能够检查、替换和复现。对于准备采用 Olmo-core 3 的团队,稳妥路径是先从小型 MoE 和单节点通信基线开始,再逐步增加专家并行规模、节点数量与容错要求。只有在路由指标、通信成本和恢复流程都可观测之后,扩大模型规模才不会把系统问题一并放大。