Gros plan macro sur un circuit intégré complexe et lumineux, évoquant une haute performance de calcul.

GitHub déploie NCCL Extensions pour accélérer les communications en IA

Ce nouveau dépôt de GitHub rassemble des schémas de communication conçus pour les cas d’usage en intelligence artificielle, s’appuyant sur les API hôtes et de dispositif de NCCL (NVIDIA Collective Communication Library). Il vise, par exemple, à accélérer le transfert de tenseurs dans des architectures complexes comme les modèles Mixture-of-Experts (MoE) ou les entraînements par renforcement.

Les développeurs peuvent désormais accéder à des primitives optimisées pour le routage de jetons MoE, construites sur l’API Device de NCCL, intégrant des opérations LSA et GIN. On trouve également une bibliothèque expérimentale permettant le re-sharding d’un tenseur entre deux groupes de processus GPU disjoints, réalisable en une seule opération sans copie de données, grâce à l’API Window de NCCL.

Pour ceux qui préfèrent travailler en Python, l’annonce de github.com inclut un paquet nccl-extensions fournissant des enveloppes Pythoniques pour nccl_ep et nccl_m2n. L’extension NCCL EP permet un entraînement MoE sans perte de données (dropless) en JAX, tandis que NCCL M2N est présenté comme un collectif sensible à la topologie et au découpage des tenseurs pour le re-sharding distribué.

Vue aérienne d'un vaste parc de serveurs haute performance, illuminé par des lignes lumineuses de données.

Les défis de la communication en entraînement massif

Le contexte de ces outils est celui des modèles de plus en plus gigantesques, où la gestion des poids devient un goulet d’étranglement. Dans le cadre des systèmes d’apprentissage par renforcement (RL), par exemple, lorsqu’on opère en mode asynchrone avec des ressources GPU non colocalisées pour l’inférence et l’entraînement, le transfert des poids entre ces deux pôles est une étape critique, le « refit ». Ce processus exige une communication haute performance à travers des configurations de parallélisme très diverses.

Les implémentations traditionnelles, qui utilisent souvent un modèle de rassemblement puis de diffusion (gather-and-broadcast), impliquent de transférer l’intégralité du paramètre, même si une unité d’inférence n’a besoin que d’un seul shard. Ceci génère une redondance de communication considérable. Pour les modèles actuels, le volume de données échangé peut atteindre, selon les estimations, dix téraoctets par refit.

Optimiser le transfert de poids sans perte de données

L’introduction de nccl-extensions répond à ce besoin de finesse dans le transfert. En particulier, la fonctionnalité NCCL M2N, qui implémente le collectif Reshard, permet de reformuler le refit comme un problème de re-sharding d’un DTensor, indépendamment de la manière dont les poids sont divisés entre les instances d’entraînement et d’inférence. En substance, au lieu de tout envoyer, on ne transfère que le minimum théorique de données nécessaire.

Les performances de ces systèmes dépendent fortement de l’orchestration des communications. D’ailleurs, on observe dans d’autres travaux sur le RL que la manière de former le groupe de communication pour le broadcast est décisive : un groupe mal construit peut rapidement se faire étrangler par une bande passante d’interconnexion unique, même si la bande passante globale est élevée.

Le dépôt fournit des briques matérielles et logicielles pour que les équipes puissent construire des systèmes où les appels collectifs sont optimisés. Les utilisateurs peuvent télécharger des roues Python précompilées pour Linux, en spécifiant leur version CUDA (par exemple, nccl-extensions[cu12]), ou construire les extensions directement à partir du code source en utilisant les Makefile dédiés pour nccl_ep et nccl_m2n.

Sources