`nemo_rl.models.megatron.batch_invariant`#

Module Contents#

Functions#

batch_invariant_token_multiple

Smallest token multiple compatible with batch-invariant MCore inference.

Data#

MCORE_TOKEN_ROUNDER

API#

nemo_rl.models.megatron.batch_invariant.MCORE_TOKEN_ROUNDER#: 64

nemo_rl.models.megatron.batch_invariant.batch_invariant_token_multiple(tp_size: int) → int[source]#: Smallest token multiple compatible with batch-invariant MCore inference.