torchtitan

torchtitan

torchtitan은 PyTorch의 대규모 언어 모델용 분산 학습 프레임워크예요. Transformers 모델을 torchtitan의 분산 학습 인프라에서 그대로 사용하는 방법을 살펴볼게요.

출처: 문서

본문

torchtitan은 PyTorch의 대규모 언어 모델용 분산 학습 프레임워크예요. Fully Sharded Data Parallelism (FSDP), 텐서, 파이프라인, 컨텍스트 병렬화(4D 병렬화)를 지원해요. torchtitan은 torch.compile과 완전히 호환되어 커널 퓨전과 그래프 최적화를 사용할 수 있고, 이 덕분에 메모리 오버헤드가 크게 줄고 학습이 빨라져요.

[!NOTE] 현재는 밀집(dense) 모델만 지원해요.

Transformers 모델을 torchtitan의 분산 학습 인프라에서 직접 사용할 수 있어요.

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.*** = 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()가 주어진 모델의 config를 불러와요. config 값은 HFTransformerModelArgs에서 torchtitan 스타일 args로 복사돼요.
  2. torchtitan의 HFTransformerModel 래퍼가 config의 architecture 필드를 살펴보고 LlamaForCausalLM 같은 해당 모델 클래스를 인스턴스화하고 불러와요.
  3. forward 경로는 torchtitan의 병렬화·최적화 함수를 사용하면서 네이티브 Transformers 컴포넌트를 써요. torchtitan은 Transformers 모델을 아무것도 다시 쓰지 않고도 torchtitan 모델처럼 취급해요.

리소스 (Resources)

더 알아보기 (Learn more)