> 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/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://fr.nvidia-localization.ferndocs.com/fr/_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).