Jean Zay : Chargement de bases de données pour l’apprentissage distribué d’un modèle

Cette documentation présente la gestion des datasets et dataloaders en PyTorch et TensorFlow, soit toute la chaîne qui permet de transmettre les données par batch au modèle lors de l’apprentissage d’un réseau de neurones.

Nous nous positionnons pour cette documentation dans un contexte adapté à Jean Zay :

  • avec de larges modèles et des bases de données volumineuses (images, images 3D, vidéos, texte, audio, etc). Le chargement des données se fait sur les CPU, puis celles-ci sont transférées batch par batch aux GPU pour l’apprentissage ou l’inférence du modèle.
  • avec une distribution de type parallélisme de données. Si vous n’utilisez pas ce type de parallélisme, vous n’aurez qu’à ignorer l’étape de distribution.

La chaîne de chargement et pré-traitement des données d'entrée est schématisée ci-dessous :  Chaîne complète de chargement de données pour l’apprentissage distribué d’un modèle

1/ Le Data Set est l’ensemble de données, disponible en général sous la forme d’un ensemble structuré de fichiers.

2/ L’étape de Structure Discovery permet de découvrir la structure de l’ensemble de données. Celui-ci n’est pas encore chargé en mémoire à ce stade. On récupère par exemple la liste des chemins vers les fichiers, ou plus généralement la liste des pointeurs de données associées à leurs étiquettes pour le calcul de la fonction de perte dans le cas d’un apprentissage supervisé.

3/ Le Shuffle est l’étape de permutation aléatoire de l’ensemble de données. Elle est nécessaire pour l’apprentissage et est effectuée au début de chaque epoch.
Les pointeurs de données sont stockées dans un buffer mémoire avant d’être permutées. Idéalement, la taille de buffer doit être suffisante pour stocker l’ensemble des données. Cependant, les contraintes mémoires ne le permettent pas la plupart du temps. Une taille de buffer inférieure, pouvant contenir ~1 000 ou ~10 000 pointeurs, est alors considérée comme acceptable et sans impact sur la qualité de l'entraînement.
Remarque : pour l’inférence ou de la validation, l’étape du Shuffle n’est pas effectuée.

4/ La Distribution est l’étape de parallélisation de données qui permet de distribuer les données sur les différents GPU. Idéalement, la Distribution se fait après le Shuffle. En pratique, elle est plutôt faite avant le Shuffle (ou après le Batch) pour éviter les communications inter-processus. Cela est considéré comme acceptable et sans impact perceptible sur la qualité de l’entraînement.

5/ Le Batch est l’étape d’échantillonnage, i.e. de découpage de la base données en batches de taille fixe. Une boucle itérative est ensuite lancée sur l’ensemble de ces batches.
Remarque : avec le parallélisme de données, il faudra distinguer “la taille de batch” qui est définie lors du pré-traitement des données d'entrée, de “la taille de mini-batch” qui correspond à la portion de données diffusée sur chaque GPU lors de l'entraînement.

6/ Le Load and Transform est l’étape où le batch est chargé en mémoire, décodé et transformé (pré-traitement). C’est l’étape qui demande le plus de temps de calcul CPU et est donc un point de congestion évident.
Certaines opérations de transformation sont déterministes, c’est-à-dire qu’elles genèrent la même donnée transformée à chaque epoch, alors que d’autres sont liées à la Data Augmentation et sont aléatoires.
Il est possible d’effectuer les opérations déterministes en amont de la boucle sur les epochs afin d’éviter de refaire les calculs à chaque itération. Pour cela, on peut stocker les données transformées dans un cache mémoire (dans ce cas, le stockage se fait lors de la première epoch, qui sera plus longues que les suivantes). On peut aussi créer un fichier d’entrée personnalisé de type TFRecords avant l’apprentissage.
Les opérations aléatoires doivent être effectuées à chaque itération. Une optimisation possible est de transférer ces calculs sur le GPU. Des solutions, comme DALI par exemple, proposent de faire cela dans le cadre de la Data Augmention d'images. Les GPU sont bien plus rapides que les CPU pour ce genre de traitement.
Remarque : notre expérience sur Jean Zay nous montre qu’en optimisant le multiprocessing et le Prefetch (voir point suivant), les CPU sont finalement rarement congestionnés pour le pré-traitement de données. À notre connaissance, le transfert des opérations sur GPU n'est donc pas indispensable sur Jean Zay.

7/ Le Prefetch permet de créer une file d’attente de type FIFO (First In, First Out) pour le chargement des batches sur GPU. Cette fonctionnalité permet de générer du recouvrement transferts/calcul et accélère largement la chaîne de chargement de données lorsque les GPU sont plus sollicités que les CPU. Lorsque l’on utilise les GPU sur Jean Zay, le Prefetch est largement conseillé.

Mise en pratique