Planeje uma execução reproduzível do PyTorch

Prove uma etapa completa antes de escalar

Fixe as revisões do modelo e do conjunto de dados, as versões das dependências, a configuração de treinamento, as sementes aleatórias e o digest do contêiner. Execute o carregamento do modelo, uma etapa completa de treinamento, avaliação, salvamento de checkpoint e restauração de checkpoint em um acelerador.

Use o mesmo comprimento de sequência e configurações de micro-lote pretendidos para a execução mais longa. Um piloto minúsculo que omite o estado do otimizador ou usa sequências mais curtas pode ocultar o problema de memória que você está tentando medir.

Meça uma etapa de otimização representativa

Dentro do seu programa de treinamento existente, envolva uma etapa representativa com uma medição de pico de memória. O exemplo assume que o modelo, o lote e o otimizador já foram inicializados na GPU. Inclua a etapa do otimizador porque parte do estado é criada de forma preguiçosa.

torch.cuda.reset_peak_memory_stats()
optimizer.zero_grad(set_to_none=True)
loss = model(**batch).loss
loss.backward()
optimizer.step()
torch.cuda.synchronize()
peak_gib = torch.cuda.max_memory_allocated() / 2**30
print(f"Peak tensor memory: {peak_gib:.2f} GiB")

Isso reporta o pico de alocação de tensores, não a pegada completa do dispositivo. Inspecione também a memória reservada, a sobrecarga de runtime e o relatório de memória do dispositivo do sistema.

Escolha DDP ou sharding pelo motivo certo

O DistributedDataParallel mantém uma réplica do modelo em cada processo e sincroniza gradientes. É útil quando o estado completo de treinamento cabe em cada GPU e você deseja processar mais dados em paralelo; ele não mescla VRAM em um pool contíguo.

O Fully Sharded Data Parallel pode distribuir parâmetros, gradientes e estado do otimizador. Use-o quando a pressão de memória justificar o sharding e considere a comunicação extra e a complexidade de checkpoints distribuídos.

Precisão mista e checkpointing de ativação são escolhas separadas. Valide o comportamento numérico para precisão mais baixa; meça o custo computacional de recomputar ativações quando o checkpointing economiza memória.

Lance um piloto de nó único

O script de treinamento deve inicializar a execução distribuída e atribuir um processo a cada GPU usando as informações de local-rank. Combine a contagem de processos com as GPUs visíveis. O exemplo a seguir é um exemplo de inicialização para um script com reconhecimento de distribuição, não uma implementação completa de treinamento.

torchrun --standalone --nproc-per-node=4 \
  train.py --config configs/pilot.yaml

Execute isso somente em um nó provisionado. A demonstração local do AnchorGPU não executa treinamento nem aloca GPUs remotas.

Selecione o hardware com base no resultado do piloto

Use a A100 como uma linha de base 80 GB CUDA de preço mais baixo e teste a H100 se a carga de trabalho tiver um caminho útil específico para Hopper. A H200 oferece uma configuração de memória CUDA maior. A MI300X oferece mais capacidade por acelerador neste catálogo, desde que a aplicação completa seja qualificada em ROCm.

Para AMD, verifique extensões personalizadas, kernels, versões de pacotes e a imagem PyTorch testada com ROCm como um conjunto completo. A compatibilidade no nível do framework não valida todas as operações opcionais.

Teste a restauração, não apenas o salvamento

Registre o modelo, o otimizador, o agendador, a etapa de treinamento, a configuração e o estado do escalador de precisão mista quando relevante. Salve um checkpoint, encerre o processo e, em seguida, restaure e retome antes de se comprometer com a execução longa.

Use um método de checkpoint apropriado ao estado distribuído. Teste se os artefatos salvos podem ser lidos independentemente do processo original e se a saída da avaliação permanece dentro da tolerância esperada.

Orçamento para todo o pipeline

Inclua downloads, pré-processamento, carregamento de dados, validação, checkpointing, exportação e cópia das saídas para fora do nó. Meça a vazão juntamente com memória, paradas de carregamento de dados, tempo de comunicação, duração do checkpoint e qualidade da validação.

Um prazo de 7 dias é útil para uma passagem limitada de compatibilidade e perfilamento. Escolha uma reserva de 30 dias somente após estimar o trabalho útil e a margem operacional. Mais GPUs ajudam apenas quando a computação adicional supera os limites de comunicação e do pipeline de entrada.