Passer à la navigation
Documents de conception

Formation des Backends

Afficher en Markdown

NeMo RL prend en charge plusieurs backends de formation pour s’adapter à différentes tailles de modèles et configurations matérielles.

Backends disponibles

  • DTensor (FSDP2) - Formation distribuée de PyTorch avec une efficacité mémoire améliorée.
  • Megatron - Framework de formation à haute performance de NVIDIA pour mettre à l’échelle les grands modèles (>100 milliards de paramètres).

Format de point de contrôle d’entrée pris en charge

À l’heure actuelle, NeMo RL ne prend en charge que les points de contrôle Hugging Face comme entrées pour les scripts de formation. Cela s’applique à la fois au backend DTensor et au backend Megatron.

  • DTensor utilise le point de contrôle Hugging Face pour initialiser le backend de formation et configurer vllm, en veillant à ce que les implémentations de modèles correspondent exactement. Cela est crucial pour la précision.
  • Megatron utilise également le point de contrôle Hugging Face pour configurer vllm, et effectue une conversion unique vers un point de contrôle au format Megatron pour initialiser le backend de formation.

Si vous souhaitez un support direct des points de contrôle Megatron, veuillez partager votre cas d’utilisation sur https://github.com/NVIDIA-NeMo/RL/issues/671.

Sélection du Backend

Le backend de formation est automatiquement déterminé en fonction de vos paramètres de configuration YAML. Voici comment configurer chaque backend.

Backend Megatron

Pour activer la formation Megatron :

  1. Initialisez les sous-modules NeMo et Megatron en exécutant git submodule update --init --recursive
  2. Ajoutez la clé megatron_cfg à votre configuration de politique.
  3. Définissez policy.megatron_cfg.enabled=True.
  4. Reportez-vous à examples/configs/grpo_math_1B_megatron.yaml pour un exemple de configuration complet.

Remarque : Lors de l’utilisation de Megatron, l’optimiseur et le planning du taux d’apprentissage sont configurés via policy.megatron_cfg.optimizer et policy.megatron_cfg.scheduler, respectivement.

Backend DTensor

Pour activer la formation DTensor (FSDP2) :

  1. Définissez policy.dtensor_config.enabled=True.
  2. Reportez-vous à examples/configs/grpo_math_1B.yaml pour un exemple de configuration.

Priorité des Backends

Megatron prend le pas sur DTensor. Si les deux backends sont activés simultanément (policy.megatron_cfg.enabled=True et policy.dtensor_config.enabled=True), le backend Megatron sera utilisé.

Exemples de Configuration

Pour des exemples complets de chaque algorithme et backend, consultez le dossier examples/configs/recipes/llm. Ce répertoire contient des configurations prêtes à l’emploi pour diverses combinaisons prises en charge.

Configuration de Megatron

Le backend Megatron nécessite un répertoire de point de contrôle pour stocker les poids de modèle Hugging Face convertis au format Megatron. Ce répertoire doit être accessible depuis tous les nœuds de votre configuration de formation distribuée.

Priorité des Variables d’Environnement (De la plus haute à la plus basse)

  1. NRL_MEGATRON_CHECKPOINT_DIR - Le chemin du répertoire de point de contrôle personnalisé.
  2. [RECOMMANDÉ] HF_HOME/nemo_rl - Utilise le répertoire de cache Hugging Face, s’il est disponible.
  3. ~/.cache/huggingface/nemo_rl - L’emplacement de repli par défaut.

Exemples de Configuration

# Option 1 : Définir un répertoire de point de contrôle personnalisé
export NRL_MEGATRON_CHECKPOINT_DIR="/shared/nfs/checkpoints/megatron"
# Option 2 : Utiliser le répertoire personnel de HuggingFace (recommandé pour les configurations partagées)
export HF_HOME="/shared/nfs/huggingface"
# Cela utilisera /shared/nfs/huggingface/nemo_rl
# Option 3 : Utiliser la valeur par défaut (aucune variable d'environnement nécessaire)
# Utilise ~/.cache/huggingface/nemo_rl

Meilleures Pratiques

  • Monter le répertoire de point de contrôle : Si vous utilisez Docker, assurez-vous que le chemin du point de contrôle Megatron est couvert par -v/--mount. De même, si vous utilisez SLURM+pyxis, assurez-vous que --container-mounts inclut ce chemin.
  • Utiliser un stockage partagé : Assurez-vous que le répertoire de point de contrôle est accessible depuis tous les nœuds (par exemple, NFS, système de fichiers partagé).
  • Préférer HF_HOME : Si vous avez déjà HF_HOME monté sur plusieurs nœuds, cela réduit le nombre de variables d’environnement à gérer.
  • Espace suffisant : Assurez-vous d’avoir suffisamment d’espace disque pour les points de contrôle de modèle convertis.