Partitionnement de l'état de l'optimiseur

Le partitionnement de l'état de l'optimiseur est une technique d'économie de mémoire utile qui partitionne l'état de l'optimiseur (l'ensemble de poids qui décrit l'état de l'optimiseur) entre des groupes d'appareils parallèles aux données. Vous pouvez utiliser le partitionnement de l'état de l'optimiseur chaque fois que vous utilisez un optimiseur avec état (tel qu'Adam) ou un optimiseur FP16 (qui stocke les copies FP16 et FP32 des paramètres).

Note

Le sharding d'état de l'optimiseur est disponible PyTorch dans la bibliothèque de parallélisme des SageMaker modèles v1.6.0 et versions ultérieures.

Utilisation du partitionnement de l'état de l'optimiseur

Vous pouvez activer le partitionnement de l'état de l'optimiseur en définissant "shard_optimizer_state": True dans la configuration modelparallel.

Lorsque cette fonction est activée, la bibliothèque partitionne l'ensemble des paramètres du modèle en fonction du degré de parallélisme de données. Les gradients correspondant à la i-ième partition ne sont réduits qu'au i-ième rang parallèle de données. À la fin du premier appel à une fonction de décorateur smp.step, l'optimiseur enveloppé par smp.DistributedOptimizer redéfinit ses paramètres pour qu'ils ne soient limités qu'aux paramètres correspondant à la partition du rang parallèle aux données actuel. Les paramètres redéfinis sont appelés paramètres virtuels et partagent le stockage sous-jacent avec les paramètres d'origine. Lors du premier appel à optimizer.step, les états de l'optimiseur sont créés en fonction de ces paramètres redéfinis, qui sont partitionnés en raison de la partition d'origine. Après la mise à jour de l'optimiseur, l' AllGatheropération (dans le cadre de l'optimizer.stepappel) s'exécute sur les rangs parallèles des données pour obtenir des états de paramètres cohérents.

Astuce

Le partitionnement de l'état de l'optimiseur peut être utile lorsque le degré de parallélisme de données est supérieur à 1 et que le modèle comporte plus d'un milliard de paramètres.

Le degré de parallélisme de données est calculé par (processes_per_host * instance_count / pipeline_parallel_degree) et la fonction smp.dp_size() gère le dimensionnement en arrière-plan.

Configuration d'un SageMaker PyTorch estimateur

mpi_options = {
    "enabled" : True,
    "processes_per_host" : 8,               # 8 processes
    "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none "
}

smp_options = {
    "enabled":True,
    "parameters": {
        "microbatches": 4,
        "pipeline_parallel_degree": 2,    # alias for "partitions"
        "placement_strategy": "cluster",
        "tensor_parallel_degree": 2,      # tp over 2 devices
        "ddp": True,
        "shard_optimizer_state": True
    }
}

Adaptez votre script PyTorch d'entraînement

Voir Adapter votre script PyTorch d'entraînement dans la section Parallélisme de Tensor combiné au parallélisme de pipeline. Aucune modification supplémentaire n'est requise pour le script.