TorchTitan - Cheatsheet de Pré-treinamento de LLM Nativo PyTorch

TorchTitan é uma plataforma de código aberto nativa PyTorch para pré-treinamento de modelo generativo (LLM) em larga escala. É uma implementação de referência de sala limpa das técnicas de treinamento distribuído mais recentes do PyTorch, compondo paralelismo 4D — dados, tensor, pipeline e paralelo de contexto — de forma modular, em cima de FSDP2, torch.compile e ponto de verificação distribuído, com dimensionamento elástico. É o ponto de partida ideal para equipes que desejam fazer pré-treinar modelos usando dimensionamento PyTorch nativo em vez de uma pilha de terceiros.

Pré-treinamento distribuído consome muitos recursos. Valide configs em um modelo pequeno/escala antes de se comprometer com uma execução grande de vários nós.

Instalação

Método	Comando
From source	`git clone https://github.com/pytorch/torchtitan && cd torchtitan && pip install -e .`
Requirements	PyTorch recente (nightly frequentemente recomendado), GPUs CUDA
Datasets	Usa datasets HF / tokenizers
Verify	execute um config de depuração em 1 GPU

Dimensões de Paralelismo (4D)

Dimensão	Divide
Data Parallel (FSDP2)	O lote, fragmentação de params/otimizador/grads
Tensor Parallel	Camadas individuais em GPUs
Pipeline Parallel	O modelo em estágios em GPUs
Context Parallel	Sequências longas em GPUs

Esses compostos: ex. dp × tp × pp × cp mapeia um modelo em uma grande malha de GPU.

Configuração

TorchTitan é acionado por arquivos de config TOML mais overrides de linha de comando.

# Launch training with a config (multi-GPU via torchrun)
CONFIG_FILE=./train_configs/llama3_8b.toml ./run_train.sh

Seção de Config	Controla
`[model]`	Arquitetura, tamanho, sabor
`[training]`	Tamanho do lote, passos, comprimento da sequência, LR
`[parallelism]`	`data_parallel_shard_degree`, `tensor_parallel_degree`, `pipeline_parallel_degree`, `context_parallel_degree`
`[optimizer]`	Otimizador + hiperparâmetros
`[checkpoint]`	Frequência/caminhos de ponto de verificação distribuído
`[activation_checkpoint]`	Política de recompilação

Recursos-Chave

Recurso	Benefício
FSDP2	Paralelismo de dados totalmente fragmentado moderno
`torch.compile`	Compilação gráfica para velocidade
Async distributed checkpointing	Salvar/retomar execuções enormes com eficiência
Float8 / low precision	Treinamento mais rápido em hardware suportado
Elastic scaling	Adapte-se à disponibilidade de GPU em mudança
Meta-device init	Construir modelos enormes sem OOM na inicialização

Iniciando Multi-Node

# torchrun across nodes (conceptual)
torchrun --nnodes=4 --nproc_per_node=8 \
  --rdzv_backend=c10d --rdzv_endpoint=$MASTER:29500 \
  -m torchtitan.train --job.config_file llama3_70b.toml

Override	Exemplo
Parallel degrees	`--parallelism.tensor_parallel_degree 8`
Steps	`--training.steps 10000`
Seq length	`--training.seq_len 8192`
Compile	`--training.compile`

Monitoramento & Pontos de Verificação

Aspecto	Como
Logging	Integração TensorBoard / W&B
Métricas	Perda, tokens/seg, MFU (utilização de FLOPs do modelo)
Checkpoints	Ponto de verificação distribuído (DCP) salvar/carregar
Resume	Aponte o config para uma etapa de ponto de verificação salva

TorchTitan vs Outros Frameworks de Treinamento

Aspecto	TorchTitan	Megatron-LM	Nanotron
Origem	Equipe PyTorch	NVIDIA	Hugging Face
Estilo	Nativo PyTorch, composável	Recursos-ricos, maduro	Minimalista 3D
Paralelismo	4D (dp/tp/pp/cp)	3D+ (extenso)	3D
Melhor para	Dimensionamento de referência nativo PyTorch	Treinamento de MoE/denso em escala máxima	Treinamento 3D pequeno e hackeável