FP8 pour NeMo-RL

Ce module fournit une suite d’outils pour permettre la quantification FP8 pour les grands modèles de langage. Ce module est encore en développement. Actuellement, nous prenons en charge la génération FP8, en utilisant FP8 de style Deepseek (mise à l’échelle de sous-canal).

NeMo-RL applique des corrections à plusieurs fonctions vLLM pour permettre des générations FP8 pour l’apprentissage par renforcement. La fonction init_fp8 corrige les composants vLLM clés lors de l’initialisation :

1. RayDistributedExecutor : Pour l’inférence multi-GPU, l’exécuteur est corrigé pour s’assurer que chaque processus worker applique les mêmes correctifs FP8 avant l’initialisation du modèle.
2. Utilitaires de Quantification : Les fonctions au sein de vllm.model_executor.layers.quantization sont remplacées par des versions qui prennent en charge la mise à l’échelle de puissance de 2 et d’autres fonctionnalités personnalisées.
3. Chargement des Poids : Une fonction personnalisée load_weights gère la quantification à la volée des poids du modèle d’un format de précision supérieure vers FP8 avec les facteurs de mise à l’échelle corrects.

Utilisation

Les générations FP8 sont recommandées d’être configurées avec les paramètres suivants :

 loss_fn:
     # l'échantillonnage d'importance aide à améliorer la stabilité
     use_importance_sampling_correction: true
 policy:
     generation:
         vllm_cfg:
             precision: 'fp8'
             # DeepGemm est beaucoup plus performant que les noyaux de sous-canal FP8 cutlass par défaut de vLLM
             use_deep_gemm: true
             # Garder les première et dernières trois couches en bf16 réduit l'erreur multi-token sans 
             # effet significatif sur les performances
             num_last_layers_in_bf16: 3
             num_first_layers_in_bf16: 1
             # Utiliser des facteurs de mise à l'échelle FP32. L'arrondi des facteurs de mise à l'échelle au 
             # plus proche pow2 peut améliorer la fidélité de la quantification, mais cette fonctionnalité 
             # est encore en recherche.
             use_weight_pow2_scale: False
             use_activation_pow2_scale: False

Précision

Nous observons sur la recette Llama 8b une perte de précision d’environ 5% avec les générations FP8. La convergence est toujours en cours de recherche active et les générations FP8 doivent être utilisées avec prudence. Nous recherchons des moyens de réduire l’écart de précision et d’améliorer encore les performances.