Use EasyCkpt to save and resume foundation model trainings - Platform For AI: Central de Documentação do Alibaba Cloud

Salve checkpoints de treinamento do PyTorch com sobrecarga quase zero e recupere-se de falhas sem recalcular iterações perdidas usando o EasyCkpt com Megatron ou DeepSpeed.

O treinamento de modelos fundamentais em grandes clusters de GPU sofre interrupções frequentes devido a falhas de hardware, problemas no sistema e erros de rede. O checkpointing tradicional leva minutos para concluir em modelos com dezenas a centenas de bilhões de parâmetros e pausa o treinamento durante cada salvamento. Se houver interrupção, será necessário recalcular todas as iterações desde o último checkpoint.

O EasyCkpt usa checkpointing hierárquico assíncrono para salvar o estado do modelo sem pausar o treinamento. A recuperação de falhas ocorre por meio de réplicas de dados armazenadas nas GPUs, o que elimina a necessidade de recálculo.

Arquitetura

O EasyCkpt emprega três técnicas para salvar o estado do modelo sem interromper o treinamento:

Checkpointing hierárquico assíncrono: copia primeiro o estado para a memória e depois grava no disco em segundo plano
Sobreposição entre checkpoint e computação: permite continuar o treinamento enquanto as operações de checkpoint executam de forma assíncrona
Checkpointing assíncrono ciente da rede: utiliza as interfaces disponíveis mesmo diante de falhas parciais

O design aproveita três características das falhas:

As falhas afetam workers específicos. Geralmente, uma falha tem origem em uma ou duas máquinas e impacta apenas alguns workers no job de treinamento distribuído, não o job inteiro.

As falhas afetam componentes específicos. Erros de GPU não afetam operações de CPU e memória. Como os erros ocorrem apenas em interfaces de rede específicas, os nós ainda conseguem se comunicar apesar das falhas parciais.

As falhas afetam réplicas específicas do modelo. O treinamento de modelos fundamentais usa paralelismo 3D (dados, tensor, pipeline) ou Zero Redundancy Optimizer (ZeRO), mantendo várias réplicas de dados nas GPUs. Quando uma GPU falha, o treinamento se recupera a partir das réplicas em outras máquinas.

Pré-requisitos

Python 3.6, 3,8 ou 3,10
PyTorch com Megatron ou DeepSpeed
AIMaster SDK (consulte Instale o AIMaster SDK)

Instale o AIMaster SDK

Execute pip install para sua versão do Python:

# Python 3.6
pip install -U http://odps-release.cn-hangzhou.oss.aliyun-inc.com/aimaster/pai_aimaster-1.2.1-cp36-cp36m-linux_x86_64.whl

# Python 3.8
pip install -U http://odps-release.cn-hangzhou.oss.aliyun-inc.com/aimaster/pai_aimaster-1.2.1-cp38-cp38-linux_x86_64.whl

# Python 3.10
pip install -U http://odps-release.cn-hangzhou.oss.aliyun-inc.com/aimaster/pai_aimaster-1.2.1-cp310-cp310-linux_x86_64.whl

Integração com Megatron

Integre o EasyCkpt modificando seu arquivo de loop de treinamento (por exemplo, training.py) e importando um hook no arquivo do modelo (por exemplo, pretrain_gpt.py).

Modifique o loop de treinamento

Substitua a importação nativa do Megatron load_checkpoint pela versão do EasyCkpt e adicione as chamadas initialize_easyckpt e save_checkpoint_if_needed ao loop de treinamento.

from megatron.core.utils import get_model_config
from megatron import print_rank_0
from megatron import print_rank_last
# Replace the native Megatron load_checkpoint with EasyCkpt's version
# from megatron.checkpointing import load_checkpoint
from megatron.checkpointing import save_checkpoint
from megatron.model import Float16Module
from megatron.model import GPTModel
from megatron.utils import report_memory
from megatron.model.vision.knn_monitor import compute_feature_bank

# Import EasyCkpt functions
from aimaster.python.torch.easyckpt.megatron import (load_checkpoint,
                                                    initialize_easyckpt,
                                                    save_checkpoint_if_needed)

def print_datetime(string):
    """Note that this call will sync across all ranks."""
    timers('interval-time', log_level=0).start(barrier=True)
    print_datetime('before the start of training step')
    report_memory_flag = True

    # Initialize EasyCkpt before the training loop.
    # save_mem_interval=1: copy to memory every iteration
    # save_storage_interval=5: write to storage every 5 iterations
    # max_ckpt_num=5: keep up to 5 checkpoints on disk
    initialize_easyckpt(save_mem_interval=1,
                        save_storage_interval=5,
                        max_ckpt_num=5,
                        log_file_path='./test.log')

    while iteration < args.train_iters:
        if args.profile and \
           iteration == args.profile_step_start and \
                                       args.micro_batch_size * \
                                       get_num_microbatches()

        # Trigger an in-memory checkpoint if conditions are met
        save_checkpoint_if_needed(iteration, model, optimizer, opt_param_scheduler)

        # Logging.
        loss_scale = optimizer.get_loss_scale().item()
        params_norm = None

Importe o hook do EasyCkpt

No arquivo de treinamento do modelo (por exemplo, pretrain_gpt.py), importe o hook do EasyCkpt para Megatron:

from megatron.utils import average_losses_across_data_parallel_group
from megatron.arguments import core_transformer_config_from_args

# Import the EasyCkpt hook to enable async checkpointing
import aimaster.python.torch.easyckpt.megatron.hook

def model_provider(pre_process=True, post_process=True):
    """Build the model."""

Referência da API

O EasyCkpt fornece três funções para integração com Megatron:

`initialize_easyckpt`

Inicializa o EasyCkpt. Chame esta função antes do loop de treinamento.

initialize_easyckpt(save_mem_interval, save_storage_interval, max_ckpt_num, log_file_path=None)

Parâmetro	Tipo	Padrão	Descrição
`save_mem_interval`	int	`1`	Frequência de iterações para copiar o estado do modelo para a memória
`save_storage_interval`	int	`1`	Frequência de iterações para gravar checkpoints no armazenamento persistente de forma assíncrona
`max_ckpt_num`	int	`1`	Número máximo de checkpoints mantidos no disco. O sistema exclui automaticamente os mais antigos
`log_file_path`	str	`None`	Caminho para o arquivo de log das operações de checkpoint

`load_checkpoint`

Estende a função nativa do Megatron load_checkpoint() com o parâmetro concat para oferecer suporte ao otimizador distribuído.

load_checkpoint(model, optimizer, opt_param_scheduler, load_arg='load', strict=True, concat=False)

Parâmetro	Tipo	Padrão	Descrição
`model`	Model	Obrigatório	Instância do modelo Megatron
`optimizer`	Optimizer	Obrigatório	Instância do otimizador
`opt_param_scheduler`	Scheduler	Obrigatório	Agendador de parâmetros do otimizador
`load_arg`	str	`'load'`	Nome do argumento para o diretório de carregamento
`strict`	bool	`True`	Define se deve impor correspondência estrita do dicionário de estado
`concat`	bool	`False`	Defina como `True` ao alterar a contagem de instâncias ou mesclar parâmetros do otimizador distribuído

Nota

Para Megatron 2305 ou 2306 com otimizador distribuído ativado, defina concat=True em load_checkpoint() ao alterar a contagem de instâncias durante o carregamento ou ao mesclar parâmetros do otimizador distribuído.

Para Megatron 2304, substitua o load_checkpoint nativo pela versão do EasyCkpt.

`save_checkpoint_if_needed`

Aciona um checkpoint na memória. Chame esta função dentro do loop de treinamento.

save_checkpoint_if_needed(iteration, model, optimizer, opt_param_scheduler)

Parâmetro	Tipo	Padrão	Descrição
`iteration`	int	Obrigatório	Número da iteração atual do treinamento
`model`	Model	Obrigatório	Instância do modelo Megatron
`optimizer`	Optimizer	Obrigatório	Instância do otimizador
`opt_param_scheduler`	Scheduler	Obrigatório	Agendador de parâmetros do otimizador

Todos os parâmetros são variáveis existentes do Megatron no loop de treinamento.

Integração com DeepSpeed

Para tarefas DeepSpeed que utilizam Transformers Trainer, o EasyCkpt se integra por meio da classe TrainerWrapper.

Configure os parâmetros de checkpoint

O EasyCkpt reutiliza os parâmetros padrão de checkpoint do Transformers. Adicione-os ao script de inicialização:

--max_steps=10 \
--block_size=2048 \
--num_train_examples=100000 \
--gradient_checkpointing=false \
--save_strategy="steps" \
--save_steps="2" \
--save_total_limit="2"

Este exemplo salva um checkpoint a cada 2 etapas e retém no máximo 2 checkpoints.

Envolva o Trainer

Importe TrainerWrapper do módulo Transformers do EasyCkpt. Envolva a instância Trainer e ative a retomada de checkpoint:

import datasets
import transformers

# Import the EasyCkpt wrapper for Transformers Trainer
from aimaster.python.torch.easyckpt.transformers import TrainerWrapper

logger = logging.getLogger(__name__)

    tokenizer=tokenizer,
    data_collator=transformers.default_data_collator,
  )
  # Wrap the Trainer with EasyCkpt to enable async checkpointing
  trainer = TrainerWrapper(trainer)
  # Resume from the latest checkpoint automatically
  trainer.train(resume_from_checkpoint=True)

if __name__ == "__main__":
  main()

Referência da API

`TrainerWrapper`

Envolve uma instância Trainer do Transformers com recursos de checkpointing assíncrono.

from aimaster.python.torch.easyckpt.transformers import TrainerWrapper

trainer = TrainerWrapper(trainer)
trainer.train(resume_from_checkpoint=True)

Parâmetros de checkpoint

Estes parâmetros seguem as definições padrão de TrainingArguments do Transformers:

Parâmetro	Tipo	Valores válidos	Descrição
`save_strategy`	str	`"no"`, `"epoch"`, `"steps"`	Define quando salvar checkpoints. `no`: desativa o salvamento. `epoch`: salva ao final de cada época. `steps`: salva no intervalo especificado por `save_steps`
`save_steps`	int	Inteiro positivo	Intervalo de etapas para salvar checkpoints. Aplica-se apenas quando `save_strategy="steps"`
`save_total_limit`	int	Inteiro positivo	Número máximo de checkpoints retidos. O sistema exclui automaticamente os mais antigos ao atingir esse limite

Aviso

Quando save_total_limit está ativado, o sistema exclui automaticamente checkpoints desatualizados. Faça backup de quaisquer dados armazenados nessas pastas antes da exclusão. Para mais detalhes, consulte a documentação oficial do Transformers.

Segurança de dados

O EasyCkpt lê e grava dados apenas nos diretórios de armazenamento especificados. Ele pode excluir checkpoints mais antigos para respeitar o limite máximo de checkpoints.

O PAI impõe as seguintes garantias:

Todas as operações ficam restritas aos diretórios de salvamento e carregamento configurados.
Todas as operações de salvamento e exclusão são registradas em log.

Importante

Não armazene outros dados no diretório de salvamento ou carregamento do modelo. O EasyCkpt pode não funcionar corretamente, e você será responsável por qualquer perda de dados resultante.