FSDP et DeepSpeed

Accelerate offre une flexibilité en matière de frameworks d’entraînement en intégrant deux outils extrêmement puissants pour l’entraînement distribué, à savoir Pytorch FSDP et Microsoft DeepSpeed. L’objectif de ce tutoriel est de mettre en parallèle ces deux approches, tout en soulignant les différences potentielles, afin de permettre à l’utilisateur de passer d’un framework à l’autre sans accroc.

Pour passer d’un framework à l’autre, nous recommandons de lancer le code accelerate launch en passant le fichier de configuration approprié avec --config_file, ou en passant directement les arguments respectifs pour FSDP et DeepSpeed .

Des exemples de configurations Accelerate peuvent être trouvés ici pour DeepSpeed et FSDP, ou dans le zoo d’exemples sous « Launch Configurations »

Ce tutoriel est destiné uniquement aux scénarios mono-nœud, multi-GPU.

Configuration des fonctionnalités

Les tenseurs du modèle sont répartis sur différents GPU dans le but d’augmenter la taille des modèles ; cette opération est appelée sharding dans FSDP et partitioning dans DeepSpeed. Le sharding FSDP et les étapes de partitioning DeepSpeed ZeRO sont configurés respectivement par --fsdp_sharding_strategy et --zero_stage. En particulier, FULL_SHARD de FSDP correspond à l’étape 3 de DeepSpeed ZeRO ; consultez ce mapping exhaustif entre les configurations de sharding FSDP et de ZeRO DeepSpeed. Le tableau ci-dessous résume et regroupe les configurations similaires :

Pour des descriptions détaillées des éléments ci-dessus, consultez la documentation de lancement d'Accelerate.

Pour accéder à d'autres configurations de DeepSpeed, telles que les paramètres de précision mixte, vous devez passer un --deepspeed_config_file, voir la documentation.

DeepSpeed peut également être configuré via DeepSpeedPlugin, par exemple DeepSpeedPlugin.zero_stage est l’équivalent de --zero_stage, et DeepSpeedPlugin.hf_ds_config peut être utilisé pour passer --deepeed_config_file.

FSDP peut également être configuré via FullyShardedDataParallelPlugin, par exemple, FullyShardedDataParallelPlugin.sharding_strategy est l’équivalent de --fsdp_sharding_strategy.

Point de contrôle

Notez que, bien que FSDP puisse être configuré via --fsdp_state_dict_type pour enregistrer des points de contrôle complets ou fractionnés.

Pour DeepSpeed Zero3, on peut passer --zero3_save_16bit_model true, ce qui consolide commodément le modèle sur un seul rang et l’enregistre ; c’est l’équivalent FSDP de fsdp_state_dict_type: FULL_STATE_DICT.

Pour les grands modèles, la consolidation du modèle sur un seul rang peut être très lente.

Pour un point de contrôle plus rapide, utilisez fsdp_state_dict_type: SHARDED_STATE_DICT pour FSDP, et pour DeepSpeed Zero3 utilisez le script zero_to_fp32.py pour convertir a posteriori les points de contrôle fragmentés.

Offloading

FSDP n'autorise qu'un offload all-or-nothing (c'est-à-dire soit offload des paramètres, des gradients et de l'optimiseur, soit les garder tous sur le GPU), mais DeepSpeed peut offload les paramètres et l'optimiseur de manière différente. De plus, DeepSpeed prend également en charge l'offload vers NVME.

Préchargement

FSDP permet deux configurations de préchargement, --fsdp_forward_prefetch et --fsdp_backward_prefetch, afin d’améliorer le chevauchement des communications et des calculs au prix d’une consommation mémoire supplémentaire, voir la documentation FSDP. Pour DeepSpeed, le préchargement est activé lorsque nécessaire, et son activation dépend de certains hyperparamètres tels que stage3_param_persistence_threshold, stage3_max_reuse_distance, etc., qui peuvent être configurés pour Zero3 ; accelerate peut définir ces hyperparamètres automatiquement si vous ne les avez pas définis explicitement dans le fichier de configuration deepspeed.

Pour FSDP, définissez fsdp_backward_prefetch: BACKWARD_PRE pour améliorer les débits si la mémoire le permet.

Chargement du modèle

Bien que FSDP nécessite l’option explicite --fsdp_cpu_ram_efficient_loading true pour activer le chargement efficace du modèle, transformers active la fonctionnalité similaire dès que DeepSpeed Zero3 est utilisé.

Pour FSDP, chaque fois que --fsdp_cpu_ram_efficient_loading true est défini, accelerate définit automatiquement sync_module_states sur true. Pour un chargement efficace en RAM, les poids sont chargés uniquement sur un seul rang, ce qui nécessite sync_module_states pour diffuser les poids aux autres rangs.

Modèle

FSDP nécessite un --fsdp_auto_wrap_policy explicite pour que l'algorithme décide comment planifier les opérations all-gather et reduce-scatter. Mais pour DeepSpeed, cela est transparent pour l'utilisateur.

Pour FSDP, il suffit de définir fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP. Avec les dernières versions de transformers, nous faisons de notre mieux pour déterminer le fsdp_transformer_layer_cls_to_wrap approprié pour les modèles HF transformers. Toutefois, si vous rencontrez une erreur à ce sujet, veuillez le préciser.

Paramètres d'invocation

FSDP nécessite un drapeau explicite --fsdp_use_orig_params lors de l'utilisation de torch.compile, voir la documentation pytorch. Pour DeepSpeed, cela est transparent pour l'utilisateur.

Pour FSDP, lors de l'utilisation de torch.compile, veuillez définir fsdp_use_orig_params: True.

Entraînement

Deepspeed nécessite les drapeaux explicites --gradient_accumulation_steps et --gradient_clipping. Pour FSDP, cela est transparent pour l'utilisateur.

Lors de l'utilisation de DeepSpeed, définissez gradient_accumulation_steps: "auto" et gradient_clipping: "auto" pour récupérer automatiquement les valeurs définies dans Accelerator ou TrainingArguments (si vous utilisez transformers).

Sur les différences dans la gestion de la précision des données

Pour discuter de la manière dont la précision des données est gérée dans FSDP et DeepSpeed, il est instructif de d’abord présenter un aperçu du traitement des paramètres du modèle dans ces frameworks. Avant que les paramètres du modèle / de l’optimiseur ne soient distribués sur les GPU, une préparation des paramètres est nécessaire pour les « aplatir » d’abord en torch.Tensor unidimensionnels. L’implémentation de FSDP / DeepSpeed diffère quant au dtype dans lequel ces paramètres « aplatis » sont stockés, et cela a des répercussions sur la manière dont torch.Optimizer alloue leurs dtypes. Le tableau ci-dessous résume les processus des deux frameworks ; la colonne « Local » indique le processus qui se produit au niveau de chaque GPU, de sorte que tout surcoût mémoire dû à l’augmentation de la précision doit être compris comme étant amorti par le nombre de GPU utilisés.

En règle générale, pour un entraînement stable avec la précision mixte automatique, tous les paramètres entraînables doivent être en torch.float32.

Par conséquent, lors de l’utilisation de DeepSpeed avec un petit nombre de GPU, gardez à l’esprit les surcoûts mémoire potentiellement importants dus à la conversion vers un type de données plus large effectuée lors de la préparation.

Avec FSDP, en l'absence de précision mixte, il est possible d'exécuter le torch.Optimizer en basse précision torch_dtype, ce qui peut être utile lors de l'utilisation d'un petit nombre de GPU.

Avec la précision mixte, FSDP et DeepSpeed effectuent une montée de précision à l’étape de préparation du modèle (cf. tableau ci-dessus). Mais notez que FSDP enregistre ensuite les points de contrôle dans la précision montée ; DeepSpeed peut toujours enregistrer des points de contrôle en basse précision si --zero3_save_16bit_model est spécifié.

Pour clarifier le tableau ci-dessus, considérez les exemples concrets ci-dessous ; l'étape préliminaire et l'étape effective de l'optimiseur sont combinées pour la concision. Avec FSDP, il est possible de fonctionner dans les deux modes présentés ci-dessous, mais DeepSpeed ne peut fonctionner que dans un seul.