torchtitan

torchtitan 是 PyTorch 面向大语言模型的分布式训练框架。它支持完全分片数据并行（FSDP）、张量并行、流水线并行和上下文并行（4D 并行）。torchtitan 与 torch.compile 完全兼容，能够实现内核融合和图优化，从而显著降低内存开销并加快训练速度。

目前仅支持稠密模型。

直接在 torchtitan 的分布式训练基础设施中使用 Transformers 模型。

import torch
from torchtitan.config.job_config import JobConfig
from torchtitan.experiments.transformers_modeling_backend.job_config import (
    HFTransformers,
)
from torchtitan.experiments.transformers_modeling_backend.model.args import (
    TitanDenseModelArgs,
    HFTransformerModelArgs,
)
from torchtitan.experiments.transformers_modeling_backend.model.model import (
    HFTransformerModel,
)

job_config = JobConfig()

job_config.hf_transformers = HFTransformers(model="Qwen/Qwen2.5-7B")

titan_args = TitanDenseModelArgs()
model_args = HFTransformerModelArgs(titan_dense_args=titan_args).update_from_config(
    job_config
)

model = HFTransformerModel(model_args)

Transformers 集成

AutoConfig.from_pretrained() 用于加载指定模型的配置。配置值会被复制到 HFTransformerModelArgs 中的 torchtitan 风格参数中。
torchtitan 的 HFTransformerModel 封装会扫描配置中的 architecture 字段，并实例化及加载对应的模型类，例如 LlamaForCausalLM。
forward 路径使用 Transformers 的原生组件，同时借助 torchtitan 的并行化和优化方法。torchtitan 将 Transformers 模型视为 torchtitan 模型，无需重写任何内容。

资源

torchtitan 仓库

Update on GitHub