Back to blog

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%。

二、核心思想

问题定义

训练大语言模型面临的核心挑战包括:

  1. 计算效率瓶颈:现有训练框架(如 Megatron-LM、DeepSpeed)存在内核效率不足、通信开销大、GPU 利用率低等问题。
  2. 内存限制:模型参数、梯度、优化器状态和激活值的内存需求随模型规模线性增长,单 GPU 无法容纳大模型。
  3. 扩展性不足:从单 GPU 到多节点集群的无缝扩展是实际训练中的关键需求。
  4. 灵活性受限:许多框架对模型架构、数据格式、损失函数的支持有限,难以适应快速迭代的研究需求。

解决方案概述

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):

  1. 将 _weight_shard 复制到 _weight_buffer 的本地切片
  2. ZeRO Stage 3:在数据并行组上执行 all_gather(_weight_buffer)
  3. 双缓冲(num_weight_buffers=2):阶段 i+1 的 all-gather 与阶段 i 的计算异步执行

反向过程(reduce_gradients):

  1. 序列并行时:在 tensor-and-sequence-data 组上 all-reduce 序列并行梯度贡献
  2. reduce_scatter(_grad_buffer -> _grad_shard) 在数据并行组上(平均归约)
  3. 双缓冲(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_rankTP all-reduce
data相同 tensor_rank, pipeline_rankZeRO 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_rankZeRO + TP 组合
tensor_and_sequence_data相同 pipeline_rank, batch_data_rankSequence-TP 激活值
model_and_sequence_data相同 batch_data_rank跨流水线序列梯度

Transformer Block 结构

每个 Transformer Block 由以下组件组成:

组件配置类功能
Mixer (Attention)AttentionConfig多头/分组查询注意力 + RoPE
MLPMLPConfig / MoEMLPConfig前馈网络(稠密或 MoE)
NormalizationNormalizationConfigLayerNorm / RMSNorm
DropoutfloatDropout 率

支持的注意力机制:

  • 标准多头注意力(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-EntropyLanguageModelLabelEntropyLossConfig标准语言建模
Label EntropyLanguageModelLabelEntropyLossConfig标签熵正则化
Z-LossLanguageModelZLossConfig防止过度自信
DPOLanguageModelDPOLossConfig直接偏好优化
GRPOLanguageModelGRPOLossConfig群体相对策略优化
GSPOLanguageModelGSPOLossConfig群体序列策略优化
DistillationLanguageModelDistillationLossConfig知识蒸馏

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 格式:

格式名称对应模型
llamaLlama 3.1 8B, Llama 3.x 系列
qwen2Qwen 2.5 7B 系列
mistralMistral 7B
mixtralMixtral 8x7B (MoE)
mtp_llamaMulti-Token Prediction Llama
diffusion_dreamDiffusion Dream
diffusion_llamaDiffusion Llama
apriel_hybrid_ssmApriel Hybrid SSM
apriel2Apriel2
gemma4Gemma 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 遵循严格的设计原则:

  1. 泛化而非特化:新功能应扩展现有抽象,而非创建并行结构。如果 Attention 不覆盖新变体,扩展其配置而非引入 MyModelAttention。

  2. 复用优先:在编写新辅助函数、配置形状或一次性代码路径前,检查框架是否已提供。

  3. 未使用时零开销:新功能在禁用路径上无可测量成本:无新内核启动、GPU 同步点或更慢的代码路径。通过配置标志短路。

  4. 无死代码:不添加不值得维护的模块、类、抽象或配置选项。

  5. 信任内部边界:在系统边界(用户输入、外部 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 size8k tokens
Total batch size256k tokens
并行策略3D 并行(DP + TP + PP)

支持的训练场景

Fast-LLM 的文档提供了以下训练 recipe:

  1. Llama 3.1 8B 从头训练:100K 迭代,cosine 学习率调度,ZeRO Stage 2
  2. Qwen 2.5 7B 从头训练:类似配置,支持 8192 上下文长度
  3. SmolLM2-135M 快速实验:单 GPU 快速验证
  4. MoE 上循环(Upcycle):从 Llama-3B 构建 MoE 模型
  5. 指令微调:使用 DPO/GRPO 进行对齐训练
  6. 继续训练:从预训练检查点继续训练

与其他框架的对比

虽然 Fast-LLM 没有发布正式的基准测试论文,但从其设计和实际使用来看:

  • vs Megatron-LM:Fast-LLM 提供更简洁的配置系统和更好的 HuggingFace 集成;Triton 内核比 Megatron 的 CUDA 内核更易维护
  • vs DeepSpeed:Fast-LLM 将并行策略、内存优化和训练循环集成在统一框架中,而非作为独立库
  • vs FSDP:Fast-LLM 在 PyTorch FSDP 基础上增加了流水线并行、张量并行和自定义内核

七、总结

核心贡献

  1. 高性能训练框架:基于 Triton 的自定义内核实现,在实际训练中达到 10,000+ tokens/s/GPU 的吞吐量
  2. 完整的分布式训练体系:支持 DP + TP + PP + SDP 四种并行策略的自由组合,配合 ZeRO Stage 1/2/3 内存优化
  3. 统一的模型架构抽象:用可配置的组件支持所有常见 GPT 类架构,降低维护和扩展成本
  4. 生产级工程质量:Apache 2.0 开源许可、完善的文档、Docker 支持、Slurm/Kubernetes 集群部署
  5. 实际验证:用于训练 StarCoder2-3B(3B 参数,3T token),证明了框架的可靠性和效率

技术影响

  • 降低大模型训练门槛:通过简化配置和提供预构建 Docker 镜像,使中小团队也能进行大规模训练
  • 推动 Triton 生态:展示了 Triton 在训练框架中的广泛应用,为社区提供了可参考的实现
  • 促进开源协作:Apache 2.0 许可和透明的开发流程鼓励社区贡献

局限性

  1. 缺少正式基准论文:没有发表学术论文进行系统性的性能对比,难以与其他框架进行严格的基准比较
  2. 主要支持 GPT 类架构:虽然通过 SSM 层支持了 Mamba 等架构,但主要设计面向 Transformer decoder-only 模型
  3. Expert Parallelism 尚未实现:MoE 的专家并行目前标注为 TODO
  4. Sequence-Tensor Parallelism 与 MoE 不兼容:MoE 层目前不支持 sequence-tensor parallelism
  5. 部分功能仍在开发中:自定义数据加载器、损失函数等可扩展性功能标注为 “in progress”

未来方向

  • 实现 Expert Parallelism 以支持更大规模的 MoE 模型
  • 扩展多模态模型支持(视觉-语言、代码等)
  • 增加更多模型架构的原生支持
  • 完善文档和教程

八、参考资源

官方资源

关键文档

成功案例

相关项目

依赖库

  • PyTorch >= 2.9.0
  • CUDA >= 12.1
  • Python >= 3.12
  • FlashAttention(默认启用)
  • NVIDIA Apex(提供 fast_layer_norm 等内核)