Passer à la navigation

nemo_rl.distributed.model_utils

Afficher en Markdown

Contenu du Module

Classes

NomDescription
DistributedLogprobFonction autograd personnalisée pour calculer les log-probabilités dans un contexte distribué.
ChunkedDistributedLogprobFonction autograd personnalisée pour calculer les log-probabilités dans un contexte distribué.
AllGatherCPTensorAucune

Fonctions

NomDescription
_compute_distributed_log_softmaxCalculer un log softmax distribué stable à travers les workers parallèles de tenseurs.
dtensor_from_parallel_logits_to_logprobsObtenir des log-probabilités à partir de logits de vocabulaire TP+CP partitionnés.
from_parallel_logits_to_logprobsObtenir des log-probabilités à partir de logits de vocabulaire TP+CP partitionnés.
from_parallel_logits_to_logprobs_packed_sequencesObtenir des log-probabilités à partir de logits de vocabulaire TP partitionnés pour des séquences empaquetées.
_get_tokens_on_this_cp_rankObtenir les tokens sur ce rang de parallélisme contextuel.
allgather_cp_sharded_tensorAucune
get_logprobs_from_vocab_parallel_logitsCalcule les log-probabilités à partir de logits de vocabulaire parallèle.

API

[Le reste du document reste inchangé, seul le texte visible a été traduit]