`nemo_rl.models.megatron.zero_train_gen_mismatch`#

Zero train/generation KL (batch-invariant Megatron): validate the recipe, then enable.

The recipe sets every required value; this module only checks them. Only BF16 (no FP8) is supported.

Module Contents#

Classes#

`ZeroTrainGenValidation`	Collected validation outcome.
`_Tp1LocalLogprob`	At TP size 1, use the log-softmax Megatron Inference uses (bitwise).

Functions#

`validate_zero_train_gen_mismatch`	Run all zero-KL gates; return violations and non-fatal warnings.
`validate_batch_invariant_mode`	Checks for `batch_invariant_mode=True` without enabling MCore kernels.
`allow_installed_flash_attn_4`	Enable MCore’s FA4 path for the pinned `flash-attn-4` b19.
`_use_fused_log_softmax_at_tp1`
`enable_batch_invariant_kernels`	Pin TE FA support and call MCore `enable_batch_invariant_mode`.
`validate_zero_train_gen_kl`	Raise if the zero-KL / batch-invariant config is invalid.
`_validate_required_values`	The recipe must set every value the preset needs; nothing is rewritten.
`_validate_backend`
`_validate_platform`	Blackwell + FA4 only (mcore extra: flash-attn-4, CuteDSL). Hopper/FA3 not supported.
`_validate_megatron_core_features`
`_package_version`
`_first_package_version`
`_validate_packages`
`validate_zero_train_gen_model_provider`	Reject unsupported architectures using the resolved Bridge provider.
`_validate_precision`	Only BF16 without FP8 is supported, on both the train and generation side.

Data#

`MEGATRON_CORE_REQUIRED_CONFIG_FIELDS`
`TRANSFORMER_ENGINE_MIN_VERSION`
`FLASH_ATTN_MIN_VERSION`
`CUTEDSL_MIN_VERSION`
`_ZERO_KL_MEGATRON_REQUIRED`
`_ZERO_KL_GENERATION_REQUIRED`
`_TE_NATIVE_ENV_VARS`

API#

nemo_rl.models.megatron.zero_train_gen_mismatch.MEGATRON_CORE_REQUIRED_CONFIG_FIELDS#: (‘batch_invariant_mode’, ‘batch_invariant_backend’, ‘batch_invariant_collective’, ‘flash_attention_v…

nemo_rl.models.megatron.zero_train_gen_mismatch.TRANSFORMER_ENGINE_MIN_VERSION#: ‘Version(…)’

nemo_rl.models.megatron.zero_train_gen_mismatch.FLASH_ATTN_MIN_VERSION#: ‘Version(…)’

nemo_rl.models.megatron.zero_train_gen_mismatch.CUTEDSL_MIN_VERSION#: ‘Version(…)’

nemo_rl.models.megatron.zero_train_gen_mismatch._ZERO_KL_MEGATRON_REQUIRED: dict[str, Any]#: None

nemo_rl.models.megatron.zero_train_gen_mismatch._ZERO_KL_GENERATION_REQUIRED: dict[str, Any]#: None

nemo_rl.models.megatron.zero_train_gen_mismatch._TE_NATIVE_ENV_VARS#: None

class nemo_rl.models.megatron.zero_train_gen_mismatch.ZeroTrainGenValidation#

Collected validation outcome.

violations: list[str]#: ‘field(…)’

warnings: list[str]#: ‘field(…)’

raise_if_invalid(header: str) → None#

emit_warnings() → None#

nemo_rl.models.megatron.zero_train_gen_mismatch.validate_zero_train_gen_mismatch( config: nemo_rl.models.policy.PolicyConfig, *, check_packages: bool = True, check_platform: bool = True, ) → nemo_rl.models.megatron.zero_train_gen_mismatch.ZeroTrainGenValidation#: Run all zero-KL gates; return violations and non-fatal warnings.

nemo_rl.models.megatron.zero_train_gen_mismatch.validate_batch_invariant_mode( config: nemo_rl.models.policy.PolicyConfig, ) → nemo_rl.models.megatron.zero_train_gen_mismatch.ZeroTrainGenValidation#: Checks for batch_invariant_mode=True without enabling MCore kernels.

nemo_rl.models.megatron.zero_train_gen_mismatch.allow_installed_flash_attn_4() → None#

Enable MCore’s FA4 path for the pinned flash-attn-4 b19.

MCore gates HAVE_FA4 on >=4.0.0b20, which needs a newer apache-tvm-ffi than we pin. The call signature is identical.

class nemo_rl.models.megatron.zero_train_gen_mismatch._Tp1LocalLogprob(original: Any, *, chunked: bool)#

At TP size 1, use the log-softmax Megatron Inference uses (bitwise).

Initialization

apply(*args: Any) → Any#

nemo_rl.models.megatron.zero_train_gen_mismatch._use_fused_log_softmax_at_tp1() → None#

nemo_rl.models.megatron.zero_train_gen_mismatch.enable_batch_invariant_kernels( config: nemo_rl.models.policy.PolicyConfig, ) → None#: Pin TE FA support and call MCore enable_batch_invariant_mode.

nemo_rl.models.megatron.zero_train_gen_mismatch.validate_zero_train_gen_kl( config: nemo_rl.models.policy.PolicyConfig, *, check_environment: bool, ) → None#

Raise if the zero-KL / batch-invariant config is invalid.

check_environment=False runs only the config checks, so the driver can call it before Ray allocates GPUs; workers also check the installed packages and the GPU.

nemo_rl.models.megatron.zero_train_gen_mismatch._validate_required_values( config: nemo_rl.models.policy.PolicyConfig, out: nemo_rl.models.megatron.zero_train_gen_mismatch.ZeroTrainGenValidation, ) → None#: The recipe must set every value the preset needs; nothing is rewritten.

nemo_rl.models.megatron.zero_train_gen_mismatch._validate_backend( config: nemo_rl.models.policy.PolicyConfig, out: nemo_rl.models.megatron.zero_train_gen_mismatch.ZeroTrainGenValidation, ) → None#

nemo_rl.models.megatron.zero_train_gen_mismatch._validate_platform( out: nemo_rl.models.megatron.zero_train_gen_mismatch.ZeroTrainGenValidation, *, check_device: bool, ) → None#: Blackwell + FA4 only (mcore extra: flash-attn-4, CuteDSL). Hopper/FA3 not supported.

nemo_rl.models.megatron.zero_train_gen_mismatch._validate_megatron_core_features( out: nemo_rl.models.megatron.zero_train_gen_mismatch.ZeroTrainGenValidation, ) → None#

nemo_rl.models.megatron.zero_train_gen_mismatch._package_version(dist_name: str) → packaging.version.Version | None#

nemo_rl.models.megatron.zero_train_gen_mismatch._first_package_version( dist_names: tuple[str, ...], ) → packaging.version.Version | None#

nemo_rl.models.megatron.zero_train_gen_mismatch._validate_packages( out: nemo_rl.models.megatron.zero_train_gen_mismatch.ZeroTrainGenValidation, ) → None#

nemo_rl.models.megatron.zero_train_gen_mismatch.validate_zero_train_gen_model_provider( config: nemo_rl.models.policy.PolicyConfig, model_cfg: Any, ) → None#: Reject unsupported architectures using the resolved Bridge provider.

nemo_rl.models.megatron.zero_train_gen_mismatch._validate_precision( config: nemo_rl.models.policy.PolicyConfig, out: nemo_rl.models.megatron.zero_train_gen_mismatch.ZeroTrainGenValidation, ) → None#: Only BF16 without FP8 is supported, on both the train and generation side.

nemo_rl.models.megatron.zero_train_gen_mismatch#