Test NeMo RL

Ce guide décrit comment tester NeMo RL à l’aide de tests unitaires et fonctionnels, détaillant les étapes pour une exécution locale ou basée sur Docker, la configuration des dépendances et le suivi des métriques pour garantir des tests efficaces et fiables.

Tests Unitaires

uv run tests/unit/prepare_unit_test_assets.py

# Exécuter les tests unitaires en utilisant les GPU locaux
# Configuration 1 : Tests par défaut uniquement - exclut les tests hf_gated et mcore
uv run --group test bash tests/run_unit.sh
# Configuration 2 : Tests par défaut + tests HF gated, en excluant les tests mcore
uv run --group test bash tests/run_unit.sh --hf-gated
# Configuration 3 : UNIQUEMENT les tests mcore, en excluant ceux avec hf_gated
uv run --extra mcore --group test bash tests/run_unit.sh --mcore-only
# Configuration 4 : UNIQUEMENT les tests mcore, en incluant ceux avec hf_gated
uv run --extra mcore --group test bash tests/run_unit.sh --mcore-only --hf-gated

Expérimental : Itération de test local plus rapide avec pytest-testmon

Nous prenons en charge pytest-testmon pour accélérer les exécutions de tests unitaires locaux en réexécutant uniquement les tests impactés. Cela fonctionne pour le code en processus régulier et les workers @ray.remote hors processus via un assistant de sélection léger, uniquement pour les tests.

Utilisation :

# Réexécuter uniquement les tests unitaires impactés
uv run --group test pytest --testmon tests/unit
# Vous pouvez également combiner avec des marqueurs/chemins
uv run --group test pytest --hf-gated --testmon tests/unit/models/policy/test_dtensor_worker.py

À quoi s’attendre :

• Lors de la première exécution dans un nouvel espace de travail, testmon peut exécuter un ensemble plus large (ou désélectionner tout si rien n’a encore été exécuté) pour construire son cache de dépendances.
• Lors des exécutions suivantes, modifier du code non distant limite la sélection aux seuls tests qui importent/utilisent ces modules.
• Modifier du code à l’intérieur d’acteurs @ray.remote déclenche également les tests impactés. Nous maintenons un mappage statique des modules de test aux modules nemo_rl transitifs qu’ils importent et nous recoupons celui-ci avec les fichiers modifiés lorsque --testmon est présent.
• Après une exécution impactée réussie, une deuxième invocation --testmon (sans autre modification) désélectionnera tous les tests.
• L’exécution de pytest avec -k sous_chaîne_nom_test exécutera toujours les tests qui correspondent, même si --testmon est passé.

Limitations et conseils :

• La sélection est basée sur les importations Python et les mtimes des fichiers ; les ressources non Python (YAML/JSON/shell) ne sont pas suivies. Lors de la modification de celles-ci, réexécutez les tests cibles explicitement.
• La sélection consciente des ressources distantes utilise une carte d’importation statique conservative (pas de résolution d’importation dynamique). Si un test charge du code dynamiquement qui n’est pas visible via les importations, vous devrez peut-être l’exécuter explicitement une fois pour amorcer la carte.
• L’assistant est uniquement destiné aux tests et n’altère pas le comportement de la bibliothèque. Il s’active automatiquement lorsque vous passez --testmon.

Actualisation des artefacts de sélection distante

Si vous modifiez la disposition des tests ou refactorisez considérablement les importations, les artefacts de sélection distante peuvent devenir obsolètes. Pour les reconstruire, supprimez les fichiers suivants à la racine du dépôt et réexécutez avec --testmon pour réamorcer :

# À la racine de nemo-rl
rm .nrl_remote_map.json .nrl_remote_state.json

Exécuter des Tests Unitaires dans un Environnement Hermétique

Pour les environnements manquant de dépendances nécessaires (par exemple, gcc, nvcc) ou lorsque la configuration environnementale peut être problématique, les tests peuvent être exécutés dans Docker avec ce script :

CONTAINER=... bash tests/run_unit_in_docker.sh

Le CONTAINER requis peut être construit en suivant les instructions dans la documentation Docker.

Suivre les Métriques dans les Tests Unitaires

Les tests unitaires peuvent également enregistrer des métriques dans une fixture. La fixture s’appelle tracker et dispose de l’API suivante :

# Suivre une métrique arbitraire (doit être sérialisable en json)
tracker.track(metric_name, metric_value)
# Enregistrer la mémoire maximale sur l'ensemble du cluster. Acceptable pour les tests car ils sont exécutés en série.
tracker.log_max_mem(metric_name)
# Renvoie la mémoire maximale. Utile si vous mesurez des changements de mémoire.
tracker.get_max_mem()

L’inclusion de la fixture tracker suit également le temps écoulé pour le test implicitement.

Voici un exemple de test :

def test_exponentiate(tracker):
    starting_mem = tracker.get_max_mem()
    base = 2
    exponent = 4
    result = base ** exponent
    tracker.track("result", result)
    tracker.log_max_mem("memory_after_exponentiating")
    change_in_mem = tracker.get_max_mem() - starting_mem
    tracker.track("change_in_mem", change_in_mem)
    assert result == 16

Qui produirait ce fichier dans tests/unit/unit_results.json :

{
  "exit_status": 0,
  "git_commit": "f1062bd3fd95fc64443e2d9ee4a35fc654ba897e",
  "start_time": "2025-03-24 23:34:12",
  "metrics": {
    "test_hf_ray_policy::test_lm_policy_generation": {
      "avg_prob_mult_error": 1.0000039339065552,
      "mean_lps": -1.5399343967437744,
      "_elapsed": 17.323044061660767
    }
  },
  "gpu_types": [
    "NVIDIA H100 80GB HBM3"
  ],
  "coverage": 24.55897613282601
}

jq -r '[.start_time, .git_commit, .metrics["test_hf_ray_policy::test_lm_policy_generation"].avg_prob_mult_error] | @tsv' tests/unit/unit_results/*
# Exemple de sortie :
#2025-03-24 23:35:39     778d288bb5d2edfd3eec4d07bb7dffffad5ef21b        1.0000039339065552
#2025-03-24 23:36:37     778d288bb5d2edfd3eec4d07bb7dffffad5ef21b        1.0000039339065552
#2025-03-24 23:37:37     778d288bb5d2edfd3eec4d07bb7dffffad5ef21b        1.0000039339065552
#2025-03-24 23:38:14     778d288bb5d2edfd3eec4d07bb7dffffad5ef21b        1.0000039339065552
#2025-03-24 23:38:50     778d288bb5d2edfd3eec4d07bb7dffffad5ef21b        1.0000039339065552

Tests Fonctionnels

Les tests fonctionnels se trouvent sous tests/functional/.

# Exécuter le test fonctionnel pour sft
uv run bash tests/functional/sft.sh

À la fin de chaque test fonctionnel, les vérifications de métriques seront imprimées ainsi que leur réussite ou échec. Voici un exemple :

                              Vérifications des Métriques
┏━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━┓
┃ Statut ┃ Vérification                   ┃ Valeur            ┃ Message ┃
┡━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━┩
│ PASS   │ data["train/loss"]["9"] < 1500 │ 817.4517822265625 │         │
└────────┴────────────────────────────────┴───────────────────┴─────────┘

Exécuter des Tests Fonctionnels dans un Environnement Hermétique

Pour les environnements manquant de dépendances nécessaires (par exemple, gcc, nvcc) ou lorsque la configuration environnementale peut être problématique, les tests peuvent être exécutés dans Docker avec ce script :

CONTAINER=... bash run_functional_in_docker.sh functional/sft.sh

Vérification Statique des Types avec MyPy

La vérification statique des types peut être effectuée sans ressources GPU :

uv run --group test mypy {programme}.py

Par exemple,

uv run --group test mypy examples/run_grpo_math.py
uv run --group test mypy examples/run_sft.py

mypy.ini contrôle la configuration de mypy.