CUDA Streams et record_stream

L'allocateur avec cache ne réutilise un bloc que sur le stream qui l'a alloué. Si un tensor est utilisé sur un autre stream, alors avant sa libération vous avez besoin soit (a) d'une synchronisation de ce stream vers le stream de création (un événement, ou creation_stream.wait_stream(side_stream)), soit (b) record_stream() / CUDACachingAllocator::recordStream(). (a) est déterministe. (b) ne l'est pas : lors d'allocations ultérieures, l'allocateur sonde les streams enregistrés, donc si la mémoire est réutilisée dépend de la progression du GPU, et l'utilisation mémoire peut exploser quand le CPU prend du retard. Préférez (a) quand c'est possible.

Lors de l'écriture ou la révision de code qui en a besoin, déterminez lequel des cas suivants s'applique :

• Évitable : le code effectuant l'utilisation multi-stream contrôle aussi quand le tensor est libéré, ou peut maintenir une référence jusqu'à un point qu'il contrôle. Enregistrez un événement après la dernière utilisation sur side-stream et faites en sorte que le stream de création l'attende avant de relâcher la référence. N'utilisez pas record_stream. FSDP2 fonctionne ainsi (voir docs/source/distributed.fsdp.fully_shard.md).
• Inévitable sous l'API actuelle : le tensor va à du code utilisateur arbitraire, qui le libère à un moment inconnu, et rien que nous contrôlons ne s'exécute entre deux points où nous pourrions insérer la synchronisation. Ici record_stream est le bon outil et ne devrait pas être signalé.
• Mais nous concevons les APIs. Si l'API est nouvelle ou peut encore changer, demandez-vous si une forme différente nous donnerait un point de synchronisation. Exemples : retourner un handle dont la méthode wait() fait la synchronisation et libère les références internes ; allouer les sorties sur le stream de l'appelant et synchroniser le side stream avant de retourner ; garder la propriété des buffers en interne. c10d montre les deux côtés : les collectives gardent leurs tensors vivants dans l'objet Work jusqu'à ce que wait() synchronise les streams, sans recordStream, mais isend ne demande pas wait(), donc il revient à recordStream (voir ProcessGroupNCCL.cpp). Pesez le coût de conception par rapport au bénéfice de déterminisme, et pointez l'option plutôt que d'y insister.
• Difficile à éviter : parfois l'éviter nécessiterait une refonte. Par exemple, le moteur autograd déplace les gradients entre nœuds dont les streams viennent de la forward pass (torch/csrc/autograd/input_buffer.cpp). Le moteur a le graphe et pourrait retarder les libérations jusqu'à un point où il synchronise, mais les hooks utilisateur arbitraires peuvent voir et conserver les gradients, donc il ne sait pas quand la dernière utilisation se produit. Ne traitez pas de tels cas comme réglés : dites quelles informations manquent pour un point de synchronisation et ce qu'il faudrait pour les obtenir.

Quand record_stream est le bon appel, laissez un commentaire expliquant pourquoi le point de libération ne peut pas être contrôlé. Voir aussi la docstring de record_stream dans torch/_tensor_docs.py.