torchtitan

torchtitan은 대규모 언어 모델을 위한 PyTorch의 분산 학습 프레임워크입니다. Fully Sharded Data Parallelism(FSDP), 텐서, 파이프라인, 컨텍스트 병렬성(4D 병렬성)을 지원합니다. 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 통합

AutoConfig.from_pretrained()는 주어진 모델의 config를 로드합니다. config 값은 HFTransformerModelArgs의 torchtitan 스타일 args로 복사됩니다.
torchtitan의 HFTransformerModel 래퍼는 config의 architecture 필드를 스캔하여 LlamaForCausalLM과 같은 해당 모델 클래스를 인스턴스화하고 로드합니다.
forward 경로는 네이티브 Transformers 구성 요소를 사용하되, torchtitan의 병렬화 및 최적화 기법에 의존합니다. torchtitan은 Transformers 모델을 재작성할 필요 없이 torchtitan 모델로 취급합니다.

리소스

torchtitan 저장소

Update on GitHub