Use EasyCkpt to save and resume foundation model trainings - Platform For AI

Enregistrez des points de contrôle d'entraînement PyTorch avec une surcharge quasi nulle et reprenez l'entraînement après une panne sans recalculer les itérations perdues grâce à EasyCkpt utilisé avec Megatron ou DeepSpeed.

L'entraînement de modèles de base sur de grands clusters GPU subit fréquemment des interruptions dues à des pannes matérielles, des problèmes système et des erreurs réseau. La création traditionnelle de points de contrôle prend plusieurs minutes pour les modèles comportant des dizaines à des centaines de milliards de paramètres, ce qui suspend l'entraînement à chaque sauvegarde. En cas d'interruption, toutes les itérations effectuées depuis le dernier point de contrôle doivent être recalculées.

EasyCkpt utilise la création asynchrone de points de contrôle hiérarchiques pour enregistrer l'état du modèle sans interrompre l'entraînement. Il permet de récupérer après une panne en exploitant les réplicas de données stockés sur les GPU, éliminant ainsi le besoin de recalcul.

Architecture

EasyCkpt s'appuie sur trois techniques pour enregistrer l'état du modèle sans interrompre l'entraînement :

La création asynchrone de points de contrôle hiérarchiques copie d'abord l'état en mémoire, puis l'écrit sur le disque en arrière-plan.
Le chevauchement entre le calcul et la création de points de contrôle permet de poursuivre l'entraînement pendant que les opérations de point de contrôle s'exécutent de manière asynchrone.
La création asynchrone de points de contrôle consciente du réseau exploite les interfaces réseau disponibles malgré les pannes partielles.

La conception tire parti de trois caractéristiques des pannes :

Les pannes affectent des workers spécifiques. Une panne provient généralement d'une ou deux machines, n'impactant que quelques workers dans la tâche d'entraînement distribuée, et non l'ensemble de la tâche.

Les pannes affectent des composants spécifiques. Les erreurs GPU n'affectent pas les opérations CPU et mémoire. Les erreurs ne se produisent que sur des interfaces réseau spécifiques, de sorte que les nœuds peuvent toujours communiquer malgré les pannes partielles.

Les pannes affectent des réplicas de modèles spécifiques. L'entraînement de modèles de base utilise le parallélisme 3D (données, tenseur, pipeline) ou Zero Redundancy Optimizer (ZeRO), maintenant plusieurs réplicas de données sur les GPU. Lorsqu'un GPU tombe en panne, l'entraînement reprend à partir des réplicas situés sur d'autres machines.

Prérequis

Python 3.6, 3,8 ou 3,10
PyTorch avec Megatron ou DeepSpeed
AIMaster SDK (voir Installation du SDK AIMaster)

Installation du SDK AIMaster

Exécutez pip install pour votre version de 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

Intégration avec Megatron

Intégrez EasyCkpt en modifiant votre fichier de boucle d'entraînement (par exemple, training.py) et en important un hook dans votre fichier de modèle (par exemple, pretrain_gpt.py).

Modification de la boucle d'entraînement

Remplacez l'importation native load_checkpoint de Megatron par la version d'EasyCkpt, puis ajoutez les appels initialize_easyckpt et save_checkpoint_if_needed à la boucle d'entraînement.

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

Importation du hook EasyCkpt

Dans votre fichier d'entraînement de modèle (par exemple, pretrain_gpt.py), importez le hook Megatron d'EasyCkpt :

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."""

Référence API

EasyCkpt fournit trois fonctions pour l'intégration avec Megatron :

`initialize_easyckpt`

Initialise EasyCkpt. Appelez cette fonction avant la boucle d'entraînement.

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

Paramètre	Type	Par défaut	Description
`save_mem_interval`	int	`1`	Fréquence d'itération pour la copie de l'état du modèle en mémoire
`save_storage_interval`	int	`1`	Fréquence d'itération pour l'écriture asynchrone des points de contrôle sur le stockage persistant
`max_ckpt_num`	int	`1`	Nombre maximal de points de contrôle à conserver sur le disque. Les anciens points de contrôle sont supprimés automatiquement
`log_file_path`	str	`None`	Chemin vers un fichier journal pour les logs des opérations de point de contrôle

`load_checkpoint`

Étend la fonction native load_checkpoint() de Megatron avec un paramètre concat pour la prise en charge de l'optimiseur distribué.

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

Paramètre	Type	Par défaut	Description
`model`	Model	Obligatoire	L'instance du modèle Megatron
`optimizer`	Optimizer	Obligatoire	L'instance de l'optimiseur
`opt_param_scheduler`	Scheduler	Obligatoire	Le planificateur de paramètres de l'optimiseur
`load_arg`	str	`'load'`	Nom de l'argument pour le répertoire de chargement
`strict`	bool	`True`	Indique s'il faut appliquer une correspondance stricte du dictionnaire d'état
`concat`	bool	`False`	Définissez sur `True` lors du changement du nombre d'instances ou de la fusion des paramètres de l'optimiseur distribué

Remarque

Pour Megatron 2305 ou 2306 avec l'optimiseur distribué activé, définissez concat=True dans load_checkpoint() lors du changement du nombre d'instances pendant le chargement ou de la fusion des paramètres de l'optimiseur distribué.

Pour Megatron 2304, remplacez le load_checkpoint natif par la version d'EasyCkpt.

`save_checkpoint_if_needed`

Déclenche un point de contrôle en mémoire. Appelez cette fonction à l'intérieur de la boucle d'entraînement.

save_checkpoint_if_needed(iteration, model, optimizer, opt_param_scheduler)

Paramètre	Type	Par défaut	Description
`iteration`	int	Obligatoire	Numéro d'itération d'entraînement actuel
`model`	Model	Obligatoire	L'instance du modèle Megatron
`optimizer`	Optimizer	Obligatoire	L'instance de l'optimiseur
`opt_param_scheduler`	Scheduler	Obligatoire	Le planificateur de paramètres de l'optimiseur

Tous les paramètres sont des variables Megatron existantes dans la boucle d'entraînement.

Intégration avec DeepSpeed

Pour les tâches DeepSpeed utilisant Transformers Trainer, EasyCkpt s'intègre via une classe TrainerWrapper.

Configuration des paramètres de point de contrôle

EasyCkpt réutilise les paramètres de point de contrôle standard de Transformers. Ajoutez-les à votre script de lancement :

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

Cet exemple enregistre un point de contrôle toutes les 2 étapes et conserve un maximum de 2 points de contrôle.

Encapsulation du Trainer

Importez TrainerWrapper depuis le module Transformers d'EasyCkpt. Encapsulez votre instance Trainer et activez la reprise à partir du point de contrôle :

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()

Référence API

`TrainerWrapper`

Encapsule une instance Trainer de Transformers avec des capacités de création asynchrone de points de contrôle.

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

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

Paramètres de point de contrôle

Ces paramètres suivent les définitions standard TrainingArguments de Transformers :

Paramètre	Type	Valeurs valides	Description
`save_strategy`	str	`"no"`, `"epoch"`, `"steps"`	Moment de la sauvegarde des points de contrôle. `no` : désactive la sauvegarde. `epoch` : sauvegarde à la fin de chaque époque. `steps` : sauvegarde à l'intervalle spécifié par `save_steps`
`save_steps`	int	Entier positif	Intervalle d'étapes pour la sauvegarde des points de contrôle. S'applique uniquement lorsque `save_strategy="steps"`
`save_total_limit`	int	Entier positif	Nombre maximal de points de contrôle conservés. Les anciens points de contrôle sont supprimés automatiquement lorsque cette limite est atteinte

Avertissement

Lorsque save_total_limit est activé, les points de contrôle obsolètes sont supprimés automatiquement. Sauvegardez toutes les données stockées dans ces dossiers avant leur suppression. Pour plus de détails, consultez la documentation officielle de Transformers.

Sécurité des données

EasyCkpt lit et écrit des données uniquement dans les répertoires de stockage que vous spécifiez. Il peut supprimer les anciens points de contrôle pour respecter le nombre maximal de points de contrôle.

PAI applique les garanties suivantes :

Toutes les opérations sont limitées aux répertoires de sauvegarde et de chargement configurés.
Toutes les opérations de sauvegarde et de suppression sont journalisées.

Important

Ne stockez pas d'autres données dans le répertoire de sauvegarde ou de chargement du modèle. EasyCkpt risque de ne pas fonctionner correctement et vous êtes responsable de toute perte de données qui en résulterait.