`nemo_rl.models.megatron.alignment`#

Token-alignment requirements shared by Megatron data paths.

Module Contents#

`get_fp8_token_alignment`	Return the token-dimension alignment required by FP8 or FP4.
`get_parallel_token_alignment`	Return the token alignment required by context and sequence parallelism.

nemo_rl.models.megatron.alignment.get_fp8_token_alignment( megatron_cfg: collections.abc.Mapping[str, Any], ) → int[source]#: Return the token-dimension alignment required by FP8 or FP4.

nemo_rl.models.megatron.alignment.get_parallel_token_alignment( megatron_cfg: collections.abc.Mapping[str, Any], ) → int[source]#: Return the token alignment required by context and sequence parallelism.