Fast-LLM: 高效大语言模型训练框架
ServiceNow Research 开源的大语言模型分布式训练库,基于 PyTorch 和 Triton 构建,支持 3D 并行、ZeRO、MoE 等特性
Fast-LLM: 高效大语言模型训练框架
一、项目概述
| 项目 | 内容 |
|---|---|
| 名称 | Fast-LLM |
| 全称 | Fast-LLM: Accelerating your LLM training to full speed |
| 作者 | Joel Lamy Poirier, Max Tian, Raymond Li, Charles Guille-Escuret, Luke Nitish Kumar, Denis Kocetkov, Torsten Scholak |
| 机构 | ServiceNow AI Research |
| 代码 | https://github.com/ServiceNow/Fast-LLM |
| 文档 | https://servicenow.github.io/Fast-LLM |
| 首次发布 | 2024-10-19 |
| 许可 | Apache 2.0 |
| 语言 | Python 3.12+, C++, Triton |
| Star 数 | 314 (截至 2026-06-04) |
| 依赖 | PyTorch, Triton, NVIDIA Apex, FlashAttention |
Fast-LLM 是由 ServiceNow AI Research 开发的开源大语言模型训练框架。该项目并非一篇传统学术论文,而是一个面向生产环境的工程框架,但其在分布式训练、内存优化和内核效率方面的技术实现具有重要的研究和工程价值。Fast-LLM 已在实际项目中得到验证——用于训练 StarCoder2-3B 模型(30 亿参数,3 万亿 token),在 NVIDIA DGX A100-80GB 集群上实现了每 GPU 10,000 tokens/s 的吞吐量,相比其他框架训练时间减少 20%。
二、核心思想
问题定义
训练大语言模型面临的核心挑战包括:
- 计算效率瓶颈:现有训练框架(如 Megatron-LM、DeepSpeed)存在内核效率不足、通信开销大、GPU 利用率低等问题。
- 内存限制:模型参数、梯度、优化器状态和激活值的内存需求随模型规模线性增长,单 GPU 无法容纳大模型。
- 扩展性不足:从单 GPU 到多节点集群的无缝扩展是实际训练中的关键需求。
- 灵活性受限:许多框架对模型架构、数据格式、损失函数的支持有限,难以适应快速迭代的研究需求。
解决方案概述
Fast-LLM 采用以下核心策略解决上述问题:
- 基于 Triton 的自定义内核:使用 Triton 编写高性能 GPU 内核(Adam 优化器、稀疏线性层、损失函数等),替代传统 CUDA 实现,兼顾开发效率和运行性能。
- 完整的 3D 并行体系:支持数据并行(DP)、张量并行(TP)、流水线并行(PP)以及序列数据并行(SDP),可自由组合以适应不同规模的训练任务。
- ZeRO 风格内存优化:实现 ZeRO Stage 1/2/3 的权重、梯度、优化器状态分片,并通过双缓冲(double-buffering)机制实现通信与计算的重叠。
- 统一的模型架构抽象:通过可配置的 Transformer Block 设计,用统一的类层次结构支持所有常见的 GPT 类架构(Llama、Mistral、Mixtral、Qwen2 等)。
- Dropless MoE 实现:提供高效的 Mixture-of-Experts 实现,支持 top-k 和 Sinkhorn 路由,无需 token 丢弃。
三、技术架构
整体框架结构
Fast-LLM 的代码组织遵循清晰的模块化设计:
fast_llm/
├── config.py # 核心配置系统(冻结 dataclass + 动态类型分发)
├── tensor.py # 张量抽象
├── cli.py # 命令行入口
├── core/ # 核心分布式操作
│ ├── distributed.py # 进程组管理
│ ├── kernels.py # 底层内核
│ └── ops.py # 通信操作(all-reduce, all-gather 等)
├── engine/ # 训练引擎(模型无关)
│ ├── base_model/ # BaseModel 抽象层
│ ├── distributed/ # 分布式配置与进程组
│ ├── multi_stage/ # 多阶段模型(ZeRO/FSDP 实现)
│ ├── schedule/ # 流水线调度(DAG 执行)
│ ├── optimizer/ # 优化器
│ ├── checkpoint/ # 检查点保存/加载/转换
│ ├── training/ # Trainer 基类
│ └── inference/ # 推理引擎
├── layers/ # 可复用层组件
│ ├── common/ # Linear, Normalization, PEFT (LoRA)
│ ├── attention/ # 多头/分组查询注意力, RoPE
│ ├── decoder/ # TransformerBlock, MLP, MoE
│ ├── language_model/ # Embedding, LM Head, 损失函数
│ ├── ssm/ # 状态空间模型 (Mamba, GDN, KDA)
│ └── vision/ # 视觉编码器(多模态支持)
├── models/ # 具体模型实现
│ ├── gpt/ # GPT 模型族(主模型)
│ └── multimodal/ # 视觉-语言模型
├── functional/ # 底层函数操作
│ ├── triton/ # Triton 内核集合
│ │ ├── adam.py # Triton Adam 优化器
│ │ ├── entropy_loss.py # 融合熵损失
│ │ ├── grpo_loss.py # GRPO 损失
│ │ ├── gspo_loss.py # GSPO 损失
│ │ ├── mlp.py # MoE MLP 内核
│ │ ├── sparse_linear.py # 稀疏线性层(MoE 用)
│ │ ├── sparse_copy.py # 稀疏拷贝操作
│ │ ├── normalization.py # 归一化
│ │ ├── rotary.py # 旋转位置编码
│ │ └── z_loss.py # Z-loss 正则化
│ ├── entropy_loss.py
│ ├── linear.py
│ └── utils.py
└── data/ # 数据管道
├── dataset/ # Memmap, 混合, 拼接, FIM, 随机, 流式数据集
├── data/gpt/ # GPT 特定数据管道
└── preparation/ # 离线数据预处理工具
配置系统
Fast-LLM 的配置系统是其核心基础设施之一。所有配置都是使用 @config_class() 装饰器的冻结 dataclass,继承自 Config 基类。字段使用 Field(default=..., desc=..., hint=FieldHint.X) 定义,其中 FieldHint 控制序列化行为和验证级别:
| FieldHint | 含义 | 用途 |
|---|---|---|
architecture | 定义模型结构 | 跨检查点比较 |
core | 始终需要显式指定 | 关键训练参数 |
optional/performance/stability/feature/expert | 可选调优参数 | 不同层次的配置 |
derived | 从其他字段计算 | 不参与序列化 |
动态类型分发使用 @config_class(dynamic_type={BaseClass: "name"}),实现 YAML 配置中的 type: 键选择子类。例如 fast-llm train gpt 命令通过两级动态分发:train 选择训练子命令,gpt 选择 GPTModelConfig。
核心公式与算法
1. 分布式并行度计算
Fast-LLM 的四种并行策略通过以下关系组合:
data_parallel = world_size / (tensor_parallel * pipeline_parallel)
batch_data_parallel = data_parallel / sequence_data_parallel
Rank 分配算法:
tensor_rank = rank % tensor_parallel
data_rank = (rank // tensor_parallel) % data_parallel
pipeline_rank = rank // (tensor_parallel * data_parallel)
batch_data_rank = data_rank // sequence_data_parallel
sequence_data_rank = data_rank % sequence_data_parallel
当 pipeline_first=True 时,data_rank 和 pipeline_rank 交换以优化 NUMA 局部性。
2. 张量并行线性层
OutputParallelLinear(列切分):
- 权重形状:
[output_dim / tensor_parallel, input_dim] - 前向:
Y_local = X @ W_local,输出保持分区,无通信 - 若
sequence_parallel=True:输入先 all-gather,然后矩阵乘法 - 反向:梯度 all-reduce(或使用 sequence-TP 时 reduce-scatter)
- 用途:Q/K/V 投影、MLP gate/up 投影
InputParallelLinear(行切分):
- 权重形状:
[output_dim, input_dim / tensor_parallel] - 前向:
Y_local = X_local @ W_local,然后 all-reduce 输出 - 若
sequence_parallel=True:输出 reduce-scatter 而非 all-reduce - 用途:注意力输出投影、MLP down 投影
3. Sequence-Tensor Parallelism
标准 TP 在层间复制完整序列。Sequence-tensor parallelism 在层间保持激活值沿序列维度分布:
OutputParallelLinear: all-gather(序列分块输入) -> 完整序列 x 部分输出列
注意力/逐元素操作: 在完整序列切片上操作
InputParallelLinear: 矩阵乘 -> reduce-scatter(输出) -> 每个 rank 持有 seq_len/TP 行
总通信量(all-gather + reduce-scatter)等于单次 all-reduce,无额外带宽开销。收益是层间激活张量更小。
4. ZeRO/FSDP 内存管理
FSDP 类为每个流水线阶段维护平坦缓冲区:
_weight_shard : [total_params / data_parallel] # 本地分片,常驻内存
_weight_buffer : [total_params] # 完整权重,按需重建(ZeRO-3)
_grad_shard : [total_params / data_parallel] # 本地梯度分片
_grad_buffer : [total_params] # 完整梯度累积缓冲
每个参数是适当缓冲区切片的视图,前向/反向过程中无需拷贝。分片填充到 SHARD_PAD_TO_MULTIPLE(32)的倍数以确保对齐通信。
前向过程(restore_parameters):
- 将
_weight_shard复制到_weight_buffer的本地切片 - ZeRO Stage 3:在数据并行组上执行
all_gather(_weight_buffer) - 双缓冲(
num_weight_buffers=2):阶段 i+1 的 all-gather 与阶段 i 的计算异步执行
反向过程(reduce_gradients):
- 序列并行时:在 tensor-and-sequence-data 组上 all-reduce 序列并行梯度贡献
reduce_scatter(_grad_buffer -> _grad_shard)在数据并行组上(平均归约)- 双缓冲(
num_grad_buffers=2):阶段 i 的 reduce-scatter 与阶段 i-1 的反向重叠
5. 流水线调度
Fast-LLM 使用广度优先调度(breadth-first scheduling)隐藏流水线气泡:同时保持多个 micro-batch 在飞,使每个阶段始终有工作可做。
schedule:
micro_batch_splits: 1 # 沿序列维度细分 micro-batch
breadth_first_micro_batches: 4 # 跨阶段交错的 micro-batch 数
depth_first_micro_batches: 1 # 单阶段内梯度累积步数
调度器构建 ScheduleNode 操作(前向、反向、send、recv、优化器步骤)的 DAG,并在三个 CUDA 流(计算、发送、接收)上执行。流水线通信使用 PyTorch isend/irecv 实现与计算的重叠。
6. 确定性训练与种子管理
Distributed.set_step(step, phase) 在每个前向/反向开始时调用,重新种子化两个 per-rank 生成器:
pp_generator:种子由(step, phase, tensor_rank, data_rank)决定——确保 dropout 在同一 TP 组的流水线 rank 间一致tp_generator:种子由(step, phase, pipeline_rank, data_rank)决定——确保 TP rank 采样相同的 dropout 掩码
大素数偏移确保所有 rank 组合的种子不同,保证确定性训练和精确检查点恢复。
核心组件详解
进程组管理
Distributed 类从 DistributedConfig 构建 NCCL(或 Gloo 用于 CPU)进程组。通过 ProcessGroupPool 去重——如果两个并行维度恰好覆盖相同的 rank 集合,它们共享同一个底层 ProcessGroup。
| 进程组名 | 成员 | 主要用途 |
|---|---|---|
world | 所有 rank | 全局屏障 |
tensor | 相同 data_rank, pipeline_rank | TP all-reduce |
data | 相同 tensor_rank, pipeline_rank | ZeRO reduce-scatter / all-gather |
pipeline | 相同 tensor_rank, data_rank | 流水线 send/recv |
sequence_data | 相同 tensor_rank, pipeline_rank, batch_data_rank | 序列并行归约 |
batch_data | 相同 tensor_rank, pipeline_rank, sequence_data_rank | 非序列数据归约 |
tensor_and_data | 相同 pipeline_rank | ZeRO + TP 组合 |
tensor_and_sequence_data | 相同 pipeline_rank, batch_data_rank | Sequence-TP 激活值 |
model_and_sequence_data | 相同 batch_data_rank | 跨流水线序列梯度 |
Transformer Block 结构
每个 Transformer Block 由以下组件组成:
| 组件 | 配置类 | 功能 |
|---|---|---|
| Mixer (Attention) | AttentionConfig | 多头/分组查询注意力 + RoPE |
| MLP | MLPConfig / MoEMLPConfig | 前馈网络(稠密或 MoE) |
| Normalization | NormalizationConfig | LayerNorm / RMSNorm |
| Dropout | float | Dropout 率 |
支持的注意力机制:
- 标准多头注意力(MHA)
- 分组查询注意力(GQA)
- 多查询注意力(MQA)
- 滑动窗口注意力(Mistral 风格)
支持的 MLP 变体:
- 稠密 MLP(gated/non-gated, SiLU/GeLU/ReLU 激活)
- Mixture-of-Experts(top-k / Sinkhorn 路由,dropless 实现)
- Hybrid MoE(部分层稠密 + 部分层 MoE)
损失函数体系
Fast-LLM 提供丰富的损失函数支持:
| 损失类型 | 配置类 | 用途 |
|---|---|---|
| Cross-Entropy | LanguageModelLabelEntropyLossConfig | 标准语言建模 |
| Label Entropy | LanguageModelLabelEntropyLossConfig | 标签熵正则化 |
| Z-Loss | LanguageModelZLossConfig | 防止过度自信 |
| DPO | LanguageModelDPOLossConfig | 直接偏好优化 |
| GRPO | LanguageModelGRPOLossConfig | 群体相对策略优化 |
| GSPO | LanguageModelGSPOLossConfig | 群体序列策略优化 |
| Distillation | LanguageModelDistillationLossConfig | 知识蒸馏 |
MoE 实现细节
Mixture-of-Experts 实现遵循 Megatron-LM 的设计,但增加了自定义路由:
class MixtureOfExpertMLP[ConfigType: MoEMLPConfig](MLPBase[ConfigType]):
"""
MoE 实现,支持 topk 和 sinkhorn 路由
使用 dropless 策略——无 token 丢弃
"""
关键配置参数:
experts:专家数量unshared_experts:非共享专家数router:路由器配置(LinearConfig)router_normalization:可选的路由器输入归一化router_scale:可学习的 per-feature 缩放router_per_expert_scale:可学习的 per-expert 缩放implementation:实现方式(auto/dropless/looped)RoutingType:路由类型(topk/sinkhorn)
稀疏线性操作通过 Triton 内核 sparse_linear.py 实现,使用 sparse_copy.py 中的稀疏映射机制高效处理 MoE 中的 token-专家分配。
数据管道
数据管道支持多种数据集格式:
- Memmap 数据集:内存映射的 tokenized 序列,支持随机访问
- 混合数据集:多个数据集的加权混合
- 拼接数据集:多个数据集的顺序拼接
- FIM(Fill-in-the-Middle):代码补全训练目标
- 随机数据集:用于测试和基准
- 流式数据集:流式加载大型数据集
数据预处理工具支持从 HuggingFace Hub 下载数据集并转换为 memmap 格式,支持多进程并行处理。
训练流程
支持的模型架构
通过统一的 GPT 模型类,Fast-LLM 支持以下 HuggingFace 格式:
| 格式名称 | 对应模型 |
|---|---|
llama | Llama 3.1 8B, Llama 3.x 系列 |
qwen2 | Qwen 2.5 7B 系列 |
mistral | Mistral 7B |
mixtral | Mixtral 8x7B (MoE) |
mtp_llama | Multi-Token Prediction Llama |
diffusion_dream | Diffusion Dream |
diffusion_llama | Diffusion Llama |
apriel_hybrid_ssm | Apriel Hybrid SSM |
apriel2 | Apriel2 |
gemma4 | Gemma 4 |
训练配置示例
以 Mistral-7B 为例的训练配置:
training:
train_iters: 100
num_workers: 8
logs:
interval: 10
data:
datasets:
training:
type: random
micro_batch_size: 8192
maximum_document_length: 4096
optimizer:
learning_rate:
base: 1.0e-05
decay_style: constant
warmup_iterations: 0
weight_decay: 0.1
beta_1: 0.9
beta_2: 0.95
model:
base_model:
embeddings:
vocab_size: 32000
dropout: 0.0
decoder:
block:
mixer:
type: attention
rotary:
type: default
theta: 10000
heads: 32
head_groups: 8 # GQA: 32 heads / 8 groups = 4 KV heads
head_size: 128
add_linear_biases: false
window_size: 4096
dropout: 0.0
mlp:
intermediate_size: 14336
add_linear_biases: false
gated: true
activation: silu
normalization:
type: rms_norm
epsilon: 1.0e-05
dropout: 0.0
num_blocks: 32
head:
normalization:
type: rms_norm
epsilon: 1.0e-05
hidden_size: 4096
tied_embedding_weight: false
multi_stage:
zero_stage: 2
distributed:
compute_dtype: bf16
seed: 984059
检查点转换
Fast-LLM 支持多种检查点格式:
- Fast-LLM 分布式格式:原生分布式检查点,支持 ZeRO 分片
- SafeTensors 格式:安全的张量存储格式
- HuggingFace 格式:与 HuggingFace Transformers 生态无缝集成
转换通过 CheckpointFormat 和 WeightConverter 体系实现,支持自动的配置映射和权重转换(包括分片合并/拆分、转置等操作)。
四、核心创新
| 创新点 | 说明 | 技术依据 |
|---|---|---|
| Triton 内核优化 | 使用 Triton 编写 Adam 优化器、稀疏线性层、融合损失函数等关键内核,兼顾开发效率和运行性能 | Triton 编程模型比 CUDA 更易维护,同时通过自动调优达到接近手写 CUDA 的性能 |
| Sequence-Tensor Parallelism | 在张量并行基础上,层间保持激活值沿序列维度分布,减少激活内存 TP 倍,通信量不变 | all-gather + reduce-scatter 总量 = 单次 all-reduce |
| Dropless MoE | 无 token 丢弃的 MoE 实现,使用稀疏映射和 Triton 内核高效处理不均匀的专家负载 | 避免了传统 MoE 中 token 丢弃导致的信息损失 |
| 双缓冲通信重叠 | 通过 num_grad_buffers 和 num_weight_buffers 实现梯度/权重通信与计算的异步重叠 | 将通信延迟隐藏在计算中,提升 GPU 利用率 |
| 统一模型抽象 | 通过可配置的 Block 设计和动态类型分发,用单一代码库支持所有 GPT 类架构 | 减少代码重复,降低维护成本 |
| 确定性训练 | 通过精心设计的种子管理机制,保证跨 rank 和跨重启的完全确定性 | 支持精确的检查点恢复和可复现实验 |
| 广度优先流水线调度 | 同时保持多个 micro-batch 在飞,最小化流水线气泡 | bubble fraction 趋近于零(当 breadth_first_micro_batches >= pipeline_parallel 时) |
| 灵活的损失函数框架 | 支持 CE、DPO、GRPO、GSPO、蒸馏等多种损失,均通过统一配置系统管理 | 通过 registry 模式实现动态损失类型选择 |
五、代码实现分析
设计原则
Fast-LLM 遵循严格的设计原则:
-
泛化而非特化:新功能应扩展现有抽象,而非创建并行结构。如果
Attention不覆盖新变体,扩展其配置而非引入MyModelAttention。 -
复用优先:在编写新辅助函数、配置形状或一次性代码路径前,检查框架是否已提供。
-
未使用时零开销:新功能在禁用路径上无可测量成本:无新内核启动、GPU 同步点或更慢的代码路径。通过配置标志短路。
-
无死代码:不添加不值得维护的模块、类、抽象或配置选项。
-
信任内部边界:在系统边界(用户输入、外部 API、文件格式)验证;信任内部调用者和框架不变量。
代码风格
- 无注释默认:只在 “why” 不明显时添加注释
- 无缩写:使用
batch_size而非bs - 类型标注:始终为公共接口添加类型标注,使用现代语法(
X | Y,list[T]) - 断言:使用
fast_llm.utils.Assert命名空间(Assert.eq,Assert.geq等) - 日志:每个模块
logger = logging.getLogger(__name__) - 路径:使用
pathlib.Path而非os.path - Python 版本:3.12+
关键实现细节
Triton Adam 优化器
@triton_jit()
def triton_adam_kernel(
params_ptr, grads_ptr, exp_avgs_ptr, exp_avg_sqs_ptr,
noop_flag_ptr, scale_ptr,
step_size, beta1, beta2, bias_correction, decay_factor, epsilon,
numel: tl_constexpr, block_size: tl_constexpr,
):
# 每个 block 处理 block_size 个参数
# 支持 noop_flag 跳过(用于流水线调度)
# 融合了梯度缩放、动量更新、权重衰减和参数更新
Triton Adam 实现比 Apex 更简单快速,但目前缺少多张量版本。
稀疏线性层(MoE)
@triton_autotune(
configs=autotune_configs,
key=["row_dim", "col_dim", "inner_dim"],
)
@triton_jit()
def sparse_linear_kernel(...):
# 支持多种 block 大小配置的自动调优
# 处理 MoE 中 token-专家的稀疏映射
稀疏线性层通过 SparseMap 数据结构管理 token 到专家的映射,支持高效的稀疏矩阵乘法。
C++ 扩展
fast_llm/csrc/data.cpp 提供 C++ 实现的数据处理内核,用于高性能的数据加载和预处理。
外部模型集成
fast_llm_external_models/ 包含与外部模型的集成:
apriel2/:Apriel2 模型支持apriel_hybrid_ssm/:混合 SSM 模型diffusion_dream/:Diffusion Dream 模型diffusion_llama/:Diffusion Llama 模型llava_hybrid/:LLaVA 混合模型mtp_llama/:Multi-Token Prediction Llama
六、实验结果
StarCoder2 训练案例
Fast-LLM 最重要的实际验证是训练 StarCoder2-3B 模型:
| 指标 | 数值 |
|---|---|
| 模型规模 | 3B 参数 |
| 训练数据 | 3 万亿 token(The Stack V2) |
| 集群配置 | NVIDIA DGX SuperCloud, DGX A100-80GB 节点 |
| 吞吐量 | 10,000 tokens/s/GPU |
| 训练时间缩减 | 比其他框架减少 20% |
| 上下文窗口 | 16,384 token(使用 GQA) |
| 训练目标 | 标准因果语言建模 + FIM |
Mistral-7B 基准
在 Mistral-7B 训练基准中:
| 配置 | 数值 |
|---|---|
| 集群 | 4 节点 x 8 H100 GPU = 32 GPU |
| 吞吐量 | 9,800 tokens/s/H100 |
| Micro-batch size | 8k tokens |
| Total batch size | 256k tokens |
| 并行策略 | 3D 并行(DP + TP + PP) |
支持的训练场景
Fast-LLM 的文档提供了以下训练 recipe:
- Llama 3.1 8B 从头训练:100K 迭代,cosine 学习率调度,ZeRO Stage 2
- Qwen 2.5 7B 从头训练:类似配置,支持 8192 上下文长度
- SmolLM2-135M 快速实验:单 GPU 快速验证
- MoE 上循环(Upcycle):从 Llama-3B 构建 MoE 模型
- 指令微调:使用 DPO/GRPO 进行对齐训练
- 继续训练:从预训练检查点继续训练
与其他框架的对比
虽然 Fast-LLM 没有发布正式的基准测试论文,但从其设计和实际使用来看:
- vs Megatron-LM:Fast-LLM 提供更简洁的配置系统和更好的 HuggingFace 集成;Triton 内核比 Megatron 的 CUDA 内核更易维护
- vs DeepSpeed:Fast-LLM 将并行策略、内存优化和训练循环集成在统一框架中,而非作为独立库
- vs FSDP:Fast-LLM 在 PyTorch FSDP 基础上增加了流水线并行、张量并行和自定义内核
七、总结
核心贡献
- 高性能训练框架:基于 Triton 的自定义内核实现,在实际训练中达到 10,000+ tokens/s/GPU 的吞吐量
- 完整的分布式训练体系:支持 DP + TP + PP + SDP 四种并行策略的自由组合,配合 ZeRO Stage 1/2/3 内存优化
- 统一的模型架构抽象:用可配置的组件支持所有常见 GPT 类架构,降低维护和扩展成本
- 生产级工程质量:Apache 2.0 开源许可、完善的文档、Docker 支持、Slurm/Kubernetes 集群部署
- 实际验证:用于训练 StarCoder2-3B(3B 参数,3T token),证明了框架的可靠性和效率
技术影响
- 降低大模型训练门槛:通过简化配置和提供预构建 Docker 镜像,使中小团队也能进行大规模训练
- 推动 Triton 生态:展示了 Triton 在训练框架中的广泛应用,为社区提供了可参考的实现
- 促进开源协作:Apache 2.0 许可和透明的开发流程鼓励社区贡献
局限性
- 缺少正式基准论文:没有发表学术论文进行系统性的性能对比,难以与其他框架进行严格的基准比较
- 主要支持 GPT 类架构:虽然通过 SSM 层支持了 Mamba 等架构,但主要设计面向 Transformer decoder-only 模型
- Expert Parallelism 尚未实现:MoE 的专家并行目前标注为 TODO
- Sequence-Tensor Parallelism 与 MoE 不兼容:MoE 层目前不支持 sequence-tensor parallelism
- 部分功能仍在开发中:自定义数据加载器、损失函数等可扩展性功能标注为 “in progress”
未来方向
- 实现 Expert Parallelism 以支持更大规模的 MoE 模型
- 扩展多模态模型支持(视觉-语言、代码等)
- 增加更多模型架构的原生支持
- 完善文档和教程
八、参考资源
官方资源
- GitHub 仓库:https://github.com/ServiceNow/Fast-LLM
- 官方文档:https://servicenow.github.io/Fast-LLM
- Docker 镜像:https://github.com/ServiceNow/Fast-LLM/pkgs/container/fast-llm
- ServiceNow AI Research:https://www.servicenow.com/research/
关键文档
- 并行策略指南:https://servicenow.github.io/Fast-LLM/user_guide/parallelism/
- 多阶段训练指南:https://servicenow.github.io/Fast-LLM/user_guide/multi-stage/
- 快速入门:https://servicenow.github.io/Fast-LLM/quick-start/
- 训练 Recipe:https://servicenow.github.io/Fast-LLM/recipes/train/
- 模型转换参考:https://servicenow.github.io/Fast-LLM/developer_guide/conversion/
- 并行内部实现:https://servicenow.github.io/Fast-LLM/developer_guide/parallelism/
成功案例
相关项目
- Megatron-LM:https://github.com/NVIDIA/Megatron-LM(Fast-LLM 的设计参考之一)
- DeepSpeed:https://github.com/microsoft/DeepSpeed
- PyTorch FSDP:PyTorch 原生的全分片数据并行
- FlashAttention:https://github.com/Dao-AILab/flash-attention
- Triton:https://triton-lang.org
- NVIDIA Apex:https://github.com/NVIDIA/apex
依赖库
- PyTorch >= 2.9.0
- CUDA >= 12.1
- Python >= 3.12
- FlashAttention(默认启用)
- NVIDIA Apex(提供 fast_layer_norm 等内核)