`nemo_rl.models.megatron.hybridep`#

Module Contents#

`uses_hybridep_flex_dispatcher`
`configure_hybridep_packed_input_padding`
`_get_hybridep_aligned_seq_len`
`_get_hybridep_local_pad_multiple`
`_get_packed_seq_boundaries`
`_shard_packed_seq_on_this_cp_rank`
`pad_packed_seq_for_hybridep`	Align packed inputs once, before model collectives can overlap.
`get_packed_seq_padding_mask`	Return a mask whose true entries are physical THD padding tokens.

nemo_rl.models.megatron.hybridep.uses_hybridep_flex_dispatcher( megatron_cfg: Mapping[str, object], ) → bool#

nemo_rl.models.megatron.hybridep.configure_hybridep_packed_input_padding( model_cfg: Any, config: Mapping[str, Any], ) → None#

nemo_rl.models.megatron.hybridep._get_hybridep_aligned_seq_len( local_seq_len: int, multiple: int, device: torch.device, ) → int#

nemo_rl.models.megatron.hybridep._get_hybridep_local_pad_multiple( pad_packed_seq_to_multiple_of: int, cp_size: int, ) → int#

nemo_rl.models.megatron.hybridep._get_packed_seq_boundaries( cu_seqlens_padded: torch.Tensor, ) → list[tuple[int, int]]#

nemo_rl.models.megatron.hybridep._shard_packed_seq_on_this_cp_rank( packed_tensor: torch.Tensor, cu_seqlens_padded: torch.Tensor, *, cp_rank: int, cp_size: int, seq_dim: int = 1, ) → torch.Tensor#

nemo_rl.models.megatron.hybridep.pad_packed_seq_for_hybridep( input_ids: torch.Tensor, input_ids_cp_sharded: torch.Tensor, packed_seq_params: megatron.core.packed_seq_params.PackedSeqParams, cu_seqlens_padded: torch.Tensor, pad_packed_seq_to_multiple_of: int, cp_rank: int, cp_size: int, ) → tuple[torch.Tensor, torch.Tensor, megatron.core.packed_seq_params.PackedSeqParams, torch.Tensor]#: Align packed inputs once, before model collectives can overlap.

nemo_rl.models.megatron.hybridep.get_packed_seq_padding_mask( cu_seqlens: torch.Tensor, cu_seqlens_padded: torch.Tensor, total_tokens: int, ) → torch.Tensor#: Return a mask whose true entries are physical THD padding tokens.