Frameworks de Treinamento Distribuído

PyTorch

• URL: pytorch.org
• Versão estável: 2.x (2024)
• Linguagem: Python + C++/CUDA
• Padrão de facto para pesquisa e produção de LLMs

Features Relevantes para LLMs

• FSDP2 (Fully Sharded Data Parallel): Estratégia ZeRO-3 nativa; prefetching implícito
• torch.compile: Compilação JIT → 30–50% speedup automático
• FlexAttention: API flexível para variantes de atenção
• DTensor: Primitiva de tensor distribuído (base do FSDP2, Tensor Parallel)
• torch.distributed: NCCL, Gloo, MPI backends

JAX (Google)

• URL: github.comgooglejax
• Paradigma: Funcional puro; arrays imutáveis; transformações (grad, jit, vmap, pmap)
• XLA (Accelerated Linear Algebra): Compilador JIT para TPU e GPU
• jit: Compila função para XLA → speedup automático
• vmap: Vectorized map — batching automático de funções escalares
• pmap: Parallel map — distribui computação em múltiplos dispositivos

Ecossistema:

• Flax: Módulos de rede neural (NNX API, Linen)
• Optax: Optimizadores (Adam, AdaFactor, Lion)
• Orbax: Checkpointing distribuído
• Grain: Pipeline de dados para JAX

Quando usar: Pesquisa Google; TPUs; workloads que se beneficiam de transformações funcionais

Keras (François Chollet)

• URL: keras.io · GitHub: keras-team/keras
• Criador: François Chollet (ex-Google, hoje na Anthropic; mesmo criador do ARC-AGI — ver 08-benchmarks/gerais-raciocinio.md)
• História:
  • Keras 1.0 (2015): API alto-nível para Theano (depois TensorFlow). Um dos motivos da adoção massiva de DL fora de pesquisa.
  • Keras 2 (2017): integrado ao TensorFlow como tf.keras (default API do TF 2.x).
  • Keras 3 (2023, "Keras Core"): multi-backend — mesma API roda sobre TensorFlow, JAX e PyTorch. Reset estratégico após "Keras virou só wrapper de TF" perceived weakness.
• Paradigma: Sequential / Functional / Subclassing API. Foco em ergonomia para beginner + production.
• KerasCV / KerasNLP / KerasHub: componentes pré-treinados; keras_hub.models.Llama3CausalLM.from_preset() interface tipo HuggingFace.
• KerasTuner: hyperparameter search.

Estado em LLM-land (2026):

• Não dominante. PyTorch + HuggingFace é a stack hegemônica para LLM training (PEFT, TRL, Axolotl, Megatron-LM). JAX para TPU em Google. Keras 3 multi-backend tenta retomar relevância via "escreva uma vez, escolha backend" mas adoção em LLMs frontier permanece marginal.
• Onde Keras ainda brilha: ensino, prototipagem rápida, modelos clássicos (CNN, RNN, transformers menores), edge deployment via TF Lite, integração com TensorBoard.
• Quando usar: projetos pedagógicos; transição de DL clássico para LLM; quando você quer trocar de backend sem reescrever; integração com ecossistema TensorFlow legacy.

Por que está aqui: referência cultural relevante (Chollet), e Keras 3 multi-backend é um exemplo de design que vale conhecer mesmo se você não usar diariamente.

DeepSpeed (Microsoft)

• URL: github.commicrosoftDeepSpeed
• Uso: ZeRO (123/Infinity), pipeline parallelism, activation checkpointing
• Integrações: PyTorch nativo; HuggingFace accelerate

ZeRO Optimizer (ver também 04-treinamento/pre-treino.md)

DeepSpeed Inference

• Kernel de atenção otimizado
• Quantização INT8 integrada
• Quando usar: Se já usa DeepSpeed para treino; kernels customizados

FSDP — Fully Sharded Data Parallel (PyTorch Nativo)

• Alternativa ao ZeRO-3 sem dependência de DeepSpeed
• FSDP2 (2024): API renovada; fully_shard() por layer; melhor interoperabilidade com torch.compile
• Documentação: pytorch.orgtutorialsintermediate/FSDP_tutorial.html
• Quando usar: Treino de modelos 7B–70B em 2–16 GPUs; padrão para maioria dos times

Megatron-LM (NVIDIA)

• URL: github.comNVIDIAMegatron-LM
• Origem: NVIDIA Research (2019)
• Especialidade: Tensor Parallelism + Pipeline Parallelism + Data Parallelism (3D parallelism)
• Performance: FlashAttention integrado; FP16BF16FP8; mixed precision
• Escala: Modelos 1T+ parâmetros em clusters de 1000+ GPUs

Quando usar:

• Pré-treino de modelos muito grandes (>70B parâmetros)
• Acesso a clusters NVIDIA (DGX, Selene)
• Reprodução de papers NVIDIA

ColossalAI

• URL: github.comhpcaitechColossalAI
• Foco: Auto-parallelism; buscador automático de estratégia de paralelismo
• Features: Sequência Parallelism, Tensor Parallelism, ZeRO alternativo
• Quando usar: Experimentação com estratégias de paralelismo automático

Nanotron (HuggingFace)

• URL: github.comhuggingfacenanotron
• Foco: Pré-treino minimalista e reprodutível
• Design: Simples de entender; Tensor + Pipeline + Data parallelism
• Uso: HuggingFace usa internamente; bom para pesquisa e reprodução

nanoGPT (Karpathy)

• URL: github.comkarpathynanoGPT
• ~300 linhas de Python — implementação minimal de GPT-2 treinável
• Valor: Referência educacional; base para experimentos
• llm.c: Versão em C puro de Karpathy (2024) — GPT-2 em C sem dependências

LightSeq / Liger Kernel

LightSeq (ByteDance)

• Kernels CUDA otimizados para transformer (atenção, layer norm, embeddings)
• 1.5–2× speedup em fine-tuning

Liger Kernel (LinkedIn/Liger)

• URL: github.comlinkedinLiger-Kernel
• Kernels Triton para RMSNorm, RoPE, SwiGLU, CrossEntropy com chunking
• Reduz memória de ativações em 60%; compatível com HuggingFace
• Drop-in: from liger_kernel.transformers import apply_liger_kernel_to_llama

Accelerate (HuggingFace)

• URL: github.comhuggingfaceaccelerate
• Papel: Camada de abstração sobre FSDP, DeepSpeed, TPU, multi-GPU
• Uso: Accelerator() → mesmo código roda em 1 GPU, 8 GPUs, TPU
• Integração: TRL, Axolotl, LLaMA-Factory usam internamente

Stack de Pré-Treino Recomendado (2026)

Monitoring de Treino

• Weights & Biases (wandb): Padrão de facto; curvas de loss, LR, gradients
• TensorBoard: Integrado ao PyTorch; menor overhead
• MLflow: Open-source; experiment tracking; model registry
• Comet ML: Alternativa ao wandb com foco enterprise

Integração mínima:

import wandb
wandb.init(project="kode-pretraining")
wandb.log({"loss": loss, "lr": scheduler.get_last_lr()[0]})