Stage - Fp8 Format Integration In The Aidge Deep Learning Framework H/F CEA

Rejoignez-nous en Stage !

CEA Tech Corporate from CEA Tech on Vimeo.
En tant que stagiaire au CEA, vous aurez l'opportunité de travailler au sein d'un environnement de recherche de renommée mondiale. Nos équipes sont composées d'experts passionnés et dédiés, offrant un cadre propice à l'apprentissage et à la collaboration. Vous aurez accès à des équipements de pointe et à des ressources de recherche de premier ordre pour mener à bien vos missions.
Description du poste :
Les calculs des réseaux de neurones sont historiquement faits dans le format 32 bits classique. L'essor récent de ces réseaux et leur complexité grandissante ont conduit à explorer d'autres formats de données, à même de réduire les temps de calcul, la mémoire utilisée et la consommation électrique. Les derniers développements dans ce domaine ont conduit à l'utilisation du format FP8 [1], où un nombre est encodé sur 8 bits. En pratique, le format FP8 se décompose en deux encodages : E4M3 (4 bits d'exposant, 3 de mantisse), utilisé pour l'inférence ; et E5M2, utilisé pour l'apprentissage.
Le but de ce stage est d'implémenter ce format numérique dans la plateforme logiciel open-source Aidge, développée par les équipes du CEA et dédiée à la conception, la manipulation et l'export de réseaux de neurones, en particulier sur des cibles embarquées [2].

Les attendus du stage sont donc :
La définition des nouveaux formats FP8 dans Aidge, qui se fera en choisissant une implémentation déjà existante et en l'intégrant dans la plateforme.
L'implémentation de noyaux de calculs (sur CPU, uniquement en inférence) utilisant ce format.
L'implémentation de noyaux de calculs CUDA (à la fois pour l'entraînement et l'inférence) utilisant le FP8.
L'entraînement de bout en bout d'un réseau de neurones en FP8 avec Aidge
(Bonus) Si le temps le permet, des sujets annexes peuvent être abordés, principalement les arrondis stochastiques et la prise en compte du facteur de scaling dans les noyaux de calculs.
En fonction des attendus de votre école, les sujets pourront être discutés.

[1] : Micikevicius, Paulius, Dusan Stosic, Neil Burgess, et al. FP8 Formats for Deep Learning. arXiv:2209.05433. Preprint, arXiv, September 29, 2022. https://doi.org/10.48550/arXiv.2209.05433.
[2] : https://eclipse.dev/aidge/