TorchTitan - Preentrenamiento de LLM nativo de PyTorch
TorchTitan es una plataforma de código abierto nativa de PyTorch para preentrenamiento de modelos generativos (LLM) a gran escala. Es una implementación de referencia de espacio en blanco de las técnicas de entrenamiento distribuido más recientes de PyTorch, componiendo paralelismo 4D — datos, tensor, pipeline y paralelo de contexto — de manera modular, sobre FSDP2, torch.compile y punto de control distribuido, con escalado elástico. Es el punto de partida para equipos que desean entrenar previamente modelos usando escalado nativo de PyTorch en lugar de una pila de terceros.
El preentrenamiento distribuido requiere muchos recursos. Valida configuraciones en un modelo/escala pequeña antes de comprometerse en una ejecución multi-nodo grande.
Instalación
| Método | Comando |
|---|
| Desde fuente | git clone https://github.com/pytorch/torchtitan && cd torchtitan && pip install -e . |
| Requerimientos | PyTorch reciente (nightly a menudo recomendado), GPUs CUDA |
| Conjuntos de datos | Usa conjuntos de datos de HF / tokenizadores |
| Verificar | ejecuta una configuración de depuración en 1 GPU |
Dimensiones de paralelismo (4D)
| Dimensión | Divide |
|---|
| Paralelo de datos (FSDP2) | El lote, parámetros de particionamiento/optimizador/gradientes |
| Paralelo de tensor | Capas individuales entre GPUs |
| Paralelo de pipeline | El modelo en etapas entre GPUs |
| Paralelo de contexto | Secuencias largas entre GPUs |
Estos se componen: por ejemplo, dp × tp × pp × cp mapea un modelo sobre una malla de GPU grande.
Configuración
TorchTitan es impulsado por archivos de configuración TOML más anulaciones de línea de comandos.
# Inicia el entrenamiento con una configuración (multi-GPU a través de torchrun)
CONFIG_FILE=./train_configs/llama3_8b.toml ./run_train.sh
| Sección de configuración | Controla |
|---|
[model] | Arquitectura, tamaño, sabor |
[training] | Tamaño de lote, pasos, longitud de secuencia, LR |
[parallelism] | data_parallel_shard_degree, tensor_parallel_degree, pipeline_parallel_degree, context_parallel_degree |
[optimizer] | Optimizador + hiperparámetros |
[checkpoint] | Frecuencia/rutas de punto de control distribuido |
[activation_checkpoint] | Política de recomputación |
Características principales
| Característica | Beneficio |
|---|
| FSDP2 | Paralelo de datos completamente cortante moderno |
torch.compile | Compilación gráfica para velocidad |
| Punto de control distribuido asincrónico | Guarda/reanuda ejecuciones enormes de manera eficiente |
| Float8 / precisión baja | Entrenamiento más rápido en hardware compatible |
| Escalado elástico | Adaptar a la disponibilidad de GPU cambiante |
| Inicialización de dispositivo Meta | Construye modelos enormes sin OOM en la inicialización |
Lanzamiento de múltiples nodos
# torchrun entre nodos (conceptual)
torchrun --nnodes=4 --nproc_per_node=8 \
--rdzv_backend=c10d --rdzv_endpoint=$MASTER:29500 \
-m torchtitan.train --job.config_file llama3_70b.toml
| Anulación | Ejemplo |
|---|
| Grados de paralelismo | --parallelism.tensor_parallel_degree 8 |
| Pasos | --training.steps 10000 |
| Longitud de secuencia | --training.seq_len 8192 |
| Compilar | --training.compile |
Monitoreo y puntos de control
| Aspecto | Cómo |
|---|
| Logging | Integración TensorBoard / W&B |
| Métricas | Pérdida, tokens/seg, MFU (utilización de FLOPs de modelo) |
| Puntos de control | Punto de control distribuido (DCP) guardar/cargar |
| Reanudar | Apunta la configuración en un paso de punto de control |
TorchTitan vs otros marcos de entrenamiento
| Aspecto | TorchTitan | Megatron-LM | Nanotron |
|---|
| Origen | Equipo PyTorch | NVIDIA | Hugging Face |
| Estilo | Nativo de PyTorch, composable | Completo, maduro | Minimalista 3D |
| Paralelismo | 4D (dp/tp/pp/cp) | 3D+ (extenso) | 3D |
| Mejor para | Escalado de referencia nativo de PyTorch | Escala máxima MoE/densa | Entrenamiento 3D pequeño hackable |
Recursos
