Nanotron - Cheatsheet de Treinamento de LLM Paralelo Minimalista 3D
Nanotron (por Hugging Face) é um framework minimalista para pré-treinamento de modelos de linguagem grande com paralelismo 3D — paralelos de dados, tensor e pipeline — implementado em Python deliberadamente legível. Seu objetivo é ser pequeno e hackeável o suficiente para realmente entender, enquanto ainda distribui treinamento em muitas GPUs. É uma escolha forte quando você deseja aprender ou personalizar a mecânica do pré-treinamento distribuído sem a complexidade dos maiores frameworks.
Pré-treinamento consome muito computador; começar pequeno. Valide um config em um modelo pequeno e um par de GPUs antes de escalar.
Instalação
| Método | Comando |
|---|
| From source | git clone https://github.com/huggingface/nanotron && cd nanotron && pip install -e . |
| Flash attention | pip install flash-attn (recomendado) |
| Requirements | PyTorch + GPUs CUDA |
| Verify | execute um config de exemplo em 1–2 GPUs |
Paralelismo 3D
| Dimensão | Divide | Chave de Config |
|---|
| Data Parallel (DP) | O lote em réplicas | dp |
| Tensor Parallel (TP) | Camadas em GPUs | tp |
| Pipeline Parallel (PP) | O modelo em estágios | pp |
O produto dp × tp × pp deve ser igual à contagem total de GPUs.
Configuração (YAML)
O treinamento do Nanotron é definido por um config YAML descrevendo o modelo, dados, paralelismo e otimizador.
parallelism:
dp: 2
tp: 2
pp: 2
pp_engine: 1f1b # pipeline schedule
model:
model_config:
hidden_size: 2048
num_hidden_layers: 24
num_attention_heads: 16
optimizer:
learning_rate_scheduler:
learning_rate: 3.0e-4
optimizer_factory:
adam_beta1: 0.9
adam_beta2: 0.95
tokens:
sequence_length: 2048
micro_batch_size: 4
batch_accumulation_per_replica: 8
Iniciando Treinamento
# Single node, multiple GPUs
torchrun --nproc_per_node=8 run_train.py --config-file config.yaml
# Generate a starter config
python examples/config_tiny_llama.py
| Chave | Significado |
|---|
run_train.py | Ponto de entrada de treinamento principal |
--config-file | O config YAML de treinamento |
pp_engine | Cronograma de pipeline (ex. 1F1B, AFAB) |
micro_batch_size | Micro lote por etapa |
batch_accumulation_per_replica | Acumulação de gradiente |
Dados
| Opção | Nota |
|---|
| Pretokenized datasets | Aponte o config para dados tokenizados |
| HF datasets | Stream/carregue via Hugging Face |
| Nanoset | Formato de conjunto de dados eficiente do Nanotron |
| Tokenizer | Definir via config do tokenizer |
Pontos de Verificação e Retomar
| Aspecto | Como |
|---|
| Save frequency | checkpoints.checkpoint_interval |
| Save path | checkpoints.checkpoints_path |
| Resume | --config-file + aponte para um ponto de verificação salvo |
| Sharded | Pontos de verificação são cientes de paralelismo |
Monitoramento
| Métrica | Significado |
|---|
| loss / lm_loss | Perda de treinamento |
| tokens/s | Rendimento |
| grad norm | Magnitude de gradiente (estabilidade) |
| Logging | Integração W&B / TensorBoard |
Nanotron vs Outros Frameworks de Treinamento
| Aspecto | Nanotron | TorchTitan | Megatron-LM |
|---|
| Filosofia | Mínimo, hackeável | Referência nativa PyTorch | Recursos-ricos, escala máxima |
| Paralelismo | 3D (dp/tp/pp) | 4D (adiciona contexto) | 3D+ extenso |
| Legibilidade | Alta (aprenda com isso) | Alta | Mais baixa (grande base de código) |
| Melhor para | Compreender/personalizar 3D | Dimensionamento nativo PyTorch | Treinamento em escala de fronteira |
Recursos
