> For clean Markdown of any page, append .md to the page URL. > For a complete documentation index, see https://fr.nvidia-localization.ferndocs.com/fr/nemo/rl/latest/nemo-rl-models-megatron-common/llms.txt. > For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://fr.nvidia-localization.ferndocs.com/_mcp/server. # nemo\_rl.models.megatron.common ## Module Contents ### Fonctions | Nom | Description | | ------------------------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------------------- | | [`_pack_sequences_for_megatron`](#nemorlmodelsmegatroncommonpacksequencesformegatron) | Empaqueter des séquences pour le traitement du modèle Megatron avec une parallélisation de contexte facultative. | | [`_unpack_sequences_from_megatron`](#nemorlmodelsmegatroncommonunpacksequencesfrommegatron) | Désempaqueter des séquences à partir du format de sortie Megatron. | | [`forward_step_arbitrary_loss`](#nemorlmodelsmegatroncommonforwardsteparbitraryloss) | Étape d'entraînement avant avec prise en charge des séquences empaquetées et de la parallélisation de contexte. | | [`broadcast_tensor`](#nemorlmodelsmegatroncommonbroadcasttensor) | Diffuse un tenseur de src\_rank à tous les rangs du groupe à l'aide de broadcast\_object\_list pour les métadonnées. | ### API ```python nemo_rl.models.megatron.common._pack_sequences_for_megatron( input_ids: torch.Tensor, seq_lengths: torch.Tensor, pad_individual_seqs_to_multiple_of: int = 1, pad_packed_seq_to: typing.Optional[int] = None, cp_rank: int = 0, cp_size: int = 1 ) -> tuple[torch.Tensor, megatron.core.packed_seq_params.PackedSeqParams, torch.Tensor, typing.Optional[torch.Tensor]] ``` Empaqueter des séquences pour le traitement du modèle Megatron avec une parallélisation de contexte facultative. Arguments : input\_ids: ID de jetons d'entrée \[batch\_size, seq\_length] seq\_lengths: Longueurs de séquence réelles pour chaque échantillon \[batch\_size] pad\_individual\_seqs\_to\_multiple\_of: Bourrer les séquences individuelles à un multiple de cette valeur pad\_packed\_seq\_to: Bourrer les séquences empaquetées jusqu'à cette valeur (avant CP) cp\_size: Taille de la parallélisation de contexte Retourne : Tuple de : * packed\_input\_ids: Tenseur d'entrée empaqueté \[1, T] * input\_ids\_cp\_sharded: Tenseur d'entrée fragmenté \[cp\_size, T // cp\_size] * packed\_seq\_params: Objet PackedSeqParams * cu\_seqlens: Longueurs de séquence cumulatives * cu\_seqlens\_padded: Longueurs de séquence cumulatives bourrées ```python nemo_rl.models.megatron.common._unpack_sequences_from_megatron( output_tensor: torch.Tensor, seq_lengths: torch.Tensor, cu_seqlens: torch.Tensor, cu_seqlens_padded: typing.Optional[torch.Tensor], original_batch_size: int, original_seq_length: int ) -> torch.Tensor ``` Désempaqueter des séquences à partir du format de sortie Megatron. Arguments : output\_tensor: Tenseur de sortie empaqueté \[1, T, vocab\_size] seq\_lengths: Longueurs de séquence réelles pour chaque échantillon cu\_seqlens: Longueurs de séquence cumulatives cu\_seqlens\_padded: Longueurs de séquence cumulatives bourrées (si CP était utilisé) original\_batch\_size: Taille de batch originale original\_seq\_length: Longueur de séquence maximale originale Retourne : Tenseur de sortie désempaqueté \[batch\_size, seq\_length, vocab\_size] ```python nemo_rl.models.megatron.common.forward_step_arbitrary_loss( state: megatron.bridge.training.state.GlobalState, global_valid_seqs: torch.Tensor, global_valid_toks: torch.Tensor, data_iterator: typing.Iterator[nemo_rl.distributed.batched_data_dict.BatchedDataDict[typing.Any]], model: megatron.core.models.gpt.GPTModel, loss_fn: nemo_rl.algorithms.loss_functions.LossFunction, pack_sequences: bool = False, seq_length_key: typing.Optional[str] = None, pad_individual_seqs_to_multiple_of: int = 1, pad_full_seq_to: typing.Optional[int] = None, cp_normalize: bool = True, policy_cfg: typing.Optional[dict] = None ) ``` Étape d'entraînement avant avec prise en charge des séquences empaquetées et de la parallélisation de contexte. Arguments : state (GlobalState): État global de l'exécution global\_valid\_seqs: Nombre global de séquences valides global\_valid\_toks: Nombre global de jetons valides data\_iterator: Itérateur de données d'entrée model (GPTModel): Le modèle GPT loss\_fn (LossFunction): Fonction de perte à appliquer pack\_sequences (bool): Indique s'il faut empaqueter les séquences pour l'efficacité seq\_length\_key (Optional\[str]): Clé dans data\_dict contenant les longueurs de séquence réelles cp\_normalize (bool): Indique s'il faut normaliser la perte par cp\_size policy\_cfg (Optional\[dict]): Configuration de stratégie contenant les paramètres de génération Notes sur les séquences empaquetées avec parallélisation de contexte (CP) : * Lorsque CP > 1, chaque séquence est bourrée à un multiple de (cp\_size \* 2) * Le facteur 2 assure un équilibrage de charge pour l'attention causale * cu\_seqlens suit les limites de séquence réelles * cu\_seqlens\_padded suit les limites de séquence bourrées pour CP * Nécessite TransformerEngine >= 1.10 pour la prise en charge CP ```python nemo_rl.models.megatron.common.broadcast_tensor( tensor: torch.Tensor | None, src_rank: int, group: torch.distributed.ProcessGroup ) -> torch.Tensor ``` Diffuse un tenseur de src\_rank à tous les rangs du groupe à l'aide de broadcast\_object\_list pour les métadonnées. Gère le cas où le tenseur d'entrée peut être None sur des rangs non sources. Si le tenseur d'entrée est fourni sur des rangs non sources, il doit avoir la forme et le type de données correspondant au tenseur sur le rang source. Arguments : tensor: Le tenseur à diffuser sur le rang source. Peut être None sur des rangs non sources (sera créé avec la forme/dtype correcte). Si ce n'est pas None sur des rangs non sources, il est utilisé comme tampon pour la diffusion et doit correspondre aux métadonnées du tenseur source. src\_rank (int): Le rang global du processus source. group: Le groupe de processus pour la communication. Retourne : torch.Tensor: Le tenseur diffusé. Sur des rangs non sources, ce sera le tenseur reçu de la source. Lève : ValueError: Si le tenseur est None sur le rang source, ou si un tenseur fourni sur un rang non source a une forme/dtype/device incompatible. TypeError: Si la diffusion des métadonnées échoue (par exemple, en raison de problèmes de sérialisation).