torchtitan

torchtitan é o framework de treinamento distribuído do PyTorch para grandes modelos de linguagem. Ele suporta Fully Sharded Data Parallelism (FSDP), paralelismo de tensor, de pipeline e de contexto (paralelismo 4D). O torchtitan é totalmente compatível com torch.compile, permitindo fusão de kernels e otimizações de grafo que reduzem significativamente a sobrecarga de memória e aceleram o treinamento.

No momento, apenas modelos densos são suportados.

Use um modelo do Transformers diretamente na infraestrutura de treinamento distribuído do 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.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)

Integração com Transformers

1. AutoConfig.from_pretrained() carrega a configuração de um modelo específico. Os valores da configuração são copiados para os argumentos no estilo torchtitan em HFTransformerModelArgs.
2. O wrapper HFTransformerModel do torchtitan verifica o campo architecture na configuração e instancia e carrega a classe de modelo correspondente, como LlamaForCausalLM.
3. O caminho forward usa componentes nativos do Transformers, apoiando-se nos métodos de paralelização e otimização do torchtitan. O torchtitan trata o modelo do Transformers como um modelo do torchtitan, sem necessidade de reescrever nada.

Recursos

• repositório torchtitan