torchtitan

torchtitan は、大規模言語モデル向けの PyTorch の分散学習フレームワークです。Fully Sharded Data Parallelism (FSDP)、テンソル並列、パイプライン並列、コンテキスト並列（4D 並列）をサポートしています。torchtitan は torch.compile と完全に互換性があり、カーネル融合やグラフ最適化を有効化することで、メモリオーバーヘッドを大幅に削減し、学習を高速化します。

現時点では、Denseモデルのみがサポートされています。

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 の統合

1. AutoConfig.from_pretrained() loads the config for a given model. The config values are copied into torchtitan style args in HFTransformerModelArgs.
2. torchtitan の HFTransformerModel ラッパーは、config の architecture フィールドをスキャンし、LlamaForCausalLM のような対応するモデルクラスをインスタンス化して読み込みます。
3. forward パスでは、ネイティブの Transformers コンポーネントを使用しつつ、torchtitan の並列化および最適化手法に依存しています。torchtitan は、何を書き直す必要もなく、Transformers モデルを torchtitan モデルとして扱います。

リソース

• torchtitanリポジトリ