`nemo_rl.models.policy`#

Subpackages#

nemo_rl.models.policy.workers

Submodules#

Package Contents#

Classes#

`LoRAConfigDisabled`
`LoRAConfig`
`AutomodelBackendConfig`	Configuration for custom MoE implementation backend in Automodel.
`AutomodelFreezeConfig`	Which sub-modules of a multi-modal Automodel to freeze during training.
`AutomodelKwargs`
`AutomodelConfigDisabled`
`MoEParallelizerOptions`	MoE parallelizer config options (mirrors Automodel’s MoEParallelizerConfig).
`AutomodelCheckpointConfig`	Automodel checkpoint settings.
`AutomodelConfig`
`SequencePackingConfigDisabled`
`SequencePackingConfig`
`RewardModelConfig`
`MegatronPeftConfigDisabled`
`MegatronPeftConfig`
`MegatronOptimizerConfig`
`MegatronSchedulerConfig`
`MegatronDDPConfig`
`Fp8Config`
`Fp4Config`	Transformer Engine NVFP4 training configuration.
`MegatronConfigDisabled`
`MegatronCheckpointConfig`	Checkpoint knobs passed through to Megatron Bridge CheckpointConfig.
`MegatronConfig`
`TokenizerConfig`
`PytorchOptimizerConfig`
`SinglePytorchSchedulerConfig`
`SinglePytorchMilestonesConfig`
`DynamicBatchingConfigDisabled`
`DynamicBatchingConfig`
`RouterReplayConfigDisabled`
`RouterReplayConfig`
`OnPolicyDistillationFullTransport`	Resolved full-vocabulary MOPD settings carried to the policy workers.
`PolicyConfig`

Functions#

_patch_transformers_tokenizer_class_set

Undo the transformers block on deepseek_v3 tokenizers.

Data#

`BLOCK_DRAFT_ALGOS`
`DRAFT_ALGOS`
`SchedulerMilestones`

API#

nemo_rl.models.policy._patch_transformers_tokenizer_class_set()#

Undo the transformers block on deepseek_v3 tokenizers.

Root cause: transformers >=5.4 lists “deepseek_v3” in two internal registries – MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS (a set) and TOKENIZER_MAPPING_NAMES (a dict pinning it to “TokenizersBackend”). Together they force the fast tokenizer backend and suppress trust_remote_code, so AutoTokenizer can only load via a local tokenizer.json. Models like Moonlight-16B-A3B ship no tokenizer.json (only tiktoken.model + a remote-code TikTokenTokenizer), so offline loading fails.

Removing both entries restores the trust_remote_code / auto_map path. discard/pop-with-default are no-ops when the entries are absent, so this is safe on any transformers version in the currently-supported range.

Placed here (nemo_rl/models/policy/init.py) so it fires exactly once per process the first time any policy code is imported – covers the driver (via nemo_rl.algorithms.grpo) and every policy worker (Megatron and Automodel both import from nemo_rl.models.policy) without polluting nemo_rl consumers that don’t touch tokenizers.

class nemo_rl.models.policy.LoRAConfigDisabled#

Bases: typing.TypedDict

enabled: Literal[False]#: None

class nemo_rl.models.policy.LoRAConfig#

Bases: typing.TypedDict

enabled: Literal[True]#: None

target_modules: list[str]#: None

exclude_modules: list[str]#: None

match_all_linear: NotRequired[bool]#: None

dim: int#: None

alpha: int#: None

dropout: float#: None

dropout_position: Literal[pre, post]#: None

lora_A_init: str#: None

use_triton: NotRequired[bool]#: None

restore_from: NotRequired[str | None]#: None

class nemo_rl.models.policy.AutomodelBackendConfig#

Bases: typing.TypedDict

Configuration for custom MoE implementation backend in Automodel.

Used when setting the backend in automodel_kwargs in your config. Alternatively, pass force_hf: true in automodel_kwargs to fall back to the HuggingFace implementation.

Initialization

Initialize self. See help(type(self)) for accurate signature.

_target_: str#: None

attn: NotRequired[str]#: None

linear: NotRequired[str]#: None

rms_norm: NotRequired[str]#: None

experts: NotRequired[str]#: None

dispatcher: NotRequired[str]#: None

fake_balanced_gate: NotRequired[bool]#: None

enable_hf_state_dict_adapter: NotRequired[bool]#: None

enable_fsdp_optimizations: NotRequired[bool]#: None

gate_precision: NotRequired[str]#: None

class nemo_rl.models.policy.AutomodelFreezeConfig#

Bases: typing.TypedDict

Which sub-modules of a multi-modal Automodel to freeze during training.

Used when setting freeze_config in automodel_kwargs in your config.

Initialization

Initialize self. See help(type(self)) for accurate signature.

freeze_vision_tower: NotRequired[bool]#: None

freeze_audio_tower: NotRequired[bool]#: None

freeze_language_model: NotRequired[bool]#: None

class nemo_rl.models.policy.AutomodelKwargs#

Bases: typing.TypedDict

use_liger_kernel: NotRequired[bool]#: None

backend: NotRequired[nemo_rl.models.policy.AutomodelBackendConfig]#: None

freeze_config: NotRequired[nemo_rl.models.policy.AutomodelFreezeConfig]#: None

force_hf: NotRequired[bool]#: None

class nemo_rl.models.policy.AutomodelConfigDisabled#

Bases: typing.TypedDict

enabled: Literal[False]#: None

class nemo_rl.models.policy.MoEParallelizerOptions#

Bases: typing.TypedDict

MoE parallelizer config options (mirrors Automodel’s MoEParallelizerConfig).

Initialization

Initialize self. See help(type(self)) for accurate signature.

ignore_router_for_ac: NotRequired[bool]#: None

reshard_after_forward: NotRequired[bool]#: None

lm_head_precision: NotRequired[str | None]#: None

wrap_outer_model: NotRequired[bool]#: None

class nemo_rl.models.policy.AutomodelCheckpointConfig#

Bases: typing.TypedDict

Automodel checkpoint settings.

Initialization

Initialize self. See help(type(self)) for accurate signature.

model_save_format: Literal[torch_save, safetensors] | None#: None

save_consolidated: Literal[false, final, every]#: None

single_rank_consolidation: bool#: None

consolidation_timeout_minutes: int#: None

class nemo_rl.models.policy.AutomodelConfig#

Bases: typing.TypedDict

enabled: Literal[True]#: None

env_vars: NotRequired[dict[str, str] | None]#: None

tensor_parallel_size: int#: None

context_parallel_size: int#: None

expert_parallel_size: NotRequired[int]#: None

dp_replicate_size: NotRequired[int]#: None

sequence_parallel: bool#: None

activation_checkpointing: bool#: None

cpu_offload: bool#: None

custom_parallel_plan: NotRequired[str | None]#: None

defer_fsdp_grad_sync: NotRequired[bool]#: None

moe_parallelizer: NotRequired[nemo_rl.models.policy.MoEParallelizerOptions]#: None

lora_cfg: NotRequired[nemo_rl.models.policy.LoRAConfig | nemo_rl.models.policy.LoRAConfigDisabled]#: None

automodel_kwargs: NotRequired[nemo_rl.models.policy.AutomodelKwargs]#: None

clear_cache_every_n_steps: NotRequired[int | None]#: None

checkpoint: NotRequired[nemo_rl.models.policy.AutomodelCheckpointConfig]#: None

deterministic: NotRequired[bool]#: None

class nemo_rl.models.policy.SequencePackingConfigDisabled#

Bases: typing.TypedDict

enabled: Literal[False]#: None

class nemo_rl.models.policy.SequencePackingConfig#

Bases: typing.TypedDict

enabled: Literal[True]#: None

train_mb_tokens: int#: None

logprob_mb_tokens: NotRequired[int]#: None

algorithm: str#: None

microbatch_order: NotRequired[Literal[packer, largest_first]]#: None

fuse_loss: NotRequired[bool]#: None

pair_grouping_key: NotRequired[Literal[pair_index]]#: None

max_sequences_per_bin: NotRequired[int]#: None

class nemo_rl.models.policy.RewardModelConfig#

Bases: typing.TypedDict

enabled: bool#: None

reward_model_type: str#: None

class nemo_rl.models.policy.MegatronPeftConfigDisabled#

Bases: typing.TypedDict

enabled: Literal[False]#: None

class nemo_rl.models.policy.MegatronPeftConfig#

Bases: typing.TypedDict

enabled: Literal[True]#: None

target_modules: list[str]#: None

exclude_modules: list[str]#: None

dim: int#: None

alpha: int#: None

dropout: float#: None

dropout_position: Literal[pre, post]#: None

lora_A_init_method: str#: None

lora_B_init_method: str#: None

a2a_experimental: bool#: None

lora_dtype: str | None#: None

share_expert_adapters: NotRequired[bool]#: None

restore_from: NotRequired[str | None]#: None

class nemo_rl.models.policy.MegatronOptimizerConfig#

Bases: typing.TypedDict

optimizer: str#: None

lr: float#: None

min_lr: float#: None

weight_decay: float#: None

bf16: bool#: None

fp16: bool#: None

params_dtype: str#: None

adam_beta1: float#: None

adam_beta2: float#: None

adam_eps: float#: None

sgd_momentum: float#: None

use_distributed_optimizer: bool#: None

use_precision_aware_optimizer: bool#: None

clip_grad: float#: None

optimizer_cpu_offload: bool#: None

optimizer_offload_fraction: float#: None

overlap_cpu_optimizer_d2h_h2d: NotRequired[bool]#: None

exp_avg_dtype: NotRequired[str]#: None

exp_avg_sq_dtype: NotRequired[str]#: None

store_param_remainders: NotRequired[bool]#: None

class nemo_rl.models.policy.MegatronSchedulerConfig#

Bases: typing.TypedDict

start_weight_decay: float#: None

end_weight_decay: float#: None

weight_decay_incr_style: str#: None

lr_decay_style: str#: None

lr_decay_iters: NotRequired[int | None]#: None

lr_warmup_iters: int#: None

lr_warmup_init: float#: None

class nemo_rl.models.policy.MegatronDDPConfig#

Bases: typing.TypedDict

grad_reduce_in_fp32: bool#: None

overlap_grad_reduce: bool#: None

overlap_param_gather: bool#: None

use_custom_fsdp: bool#: None

data_parallel_sharding_strategy: str#: None

class nemo_rl.models.policy.Fp8Config#

Bases: typing.TypedDict

enabled: bool#: None

fp8: NotRequired[str]#: None

fp8_recipe: NotRequired[str]#: None

fp8_param: NotRequired[bool]#: None

fp8_quantizer_factory: NotRequired[str]#: None

force_clear_fp8_caches: NotRequired[bool]#: None

class nemo_rl.models.policy.Fp4Config#

Bases: pydantic.BaseModel

Transformer Engine NVFP4 training configuration.

enabled: bool#: None

fp4: str | None#: None

fp4_recipe: str#: ‘nvfp4’

fp4_param: bool#: False

class nemo_rl.models.policy.MegatronConfigDisabled#

Bases: typing.TypedDict

enabled: Literal[False]#: None

class nemo_rl.models.policy.MegatronCheckpointConfig#

Bases: typing.TypedDict

Checkpoint knobs passed through to Megatron Bridge CheckpointConfig.

Initialization

Initialize self. See help(type(self)) for accurate signature.

async_save: bool#: None

ckpt_assume_constant_structure: bool#: None

ckpt_fully_parallel_save_process_group: str#: None

ckpt_fully_parallel_load_process_group: str#: None

ckpt_fully_parallel_load_exchange_algo: str#: None

class nemo_rl.models.policy.MegatronConfig#

Bases: typing.TypedDict

enabled: Literal[True]#: None

env_vars: NotRequired[dict[str, str] | None]#: None

model_overrides: NotRequired[dict[str, Any]]#: None

empty_unused_memory_level: int#: None

activation_checkpointing: bool#: None

recompute_granularity: NotRequired[Literal[full, selective]]#: None

recompute_modules: NotRequired[list[str] | None]#: None

tensor_model_parallel_size: int#: None

pipeline_model_parallel_size: int#: None

num_layers_in_first_pipeline_stage: int | None#: None

num_layers_in_last_pipeline_stage: int | None#: None

context_parallel_size: int#: None

radio_force_eval_mode: NotRequired[bool]#: None

radio_force_cpe_eval_mode: NotRequired[bool]#: None

freeze_vision_model: NotRequired[bool]#: None

freeze_vision_projection: NotRequired[bool]#: None

freeze_sound_encoder: NotRequired[bool]#: None

freeze_sound_projection: NotRequired[bool]#: None

recompute_vision: NotRequired[bool]#: None

vision_recompute_granularity: NotRequired[str | None]#: None

vision_recompute_method: NotRequired[str | None]#: None

vision_recompute_num_layers: NotRequired[int | None]#: None

pipeline_dtype: str#: None

sequence_parallel: bool#: None

freeze_moe_router: bool#: None

freeze_vision_encoder: NotRequired[bool]#: None

freeze_vision_projector: NotRequired[bool]#: None

freeze_audio_encoder: NotRequired[bool]#: None

freeze_audio_projector: NotRequired[bool]#: None

moe_router_dtype: str | None#: None

moe_router_load_balancing_type: str | list[str]#: None

moe_router_bias_update_rate: float#: None

moe_permute_fusion: bool#: None

expert_tensor_parallel_size: int#: None

expert_model_parallel_size: int#: None

defer_fp32_logits: NotRequired[bool]#: None

apply_rope_fusion: bool#: None

bias_activation_fusion: bool#: None

force_reconvert_from_hf: NotRequired[bool]#: None

attention_backend: NotRequired[str]#: None

moe_per_layer_logging: bool#: None

moe_enable_deepep: bool#: None

moe_token_dispatcher_type: str#: None

inference_moe_token_dispatcher_type: NotRequired[str]#: None

inference_grouped_gemm_backend: NotRequired[str]#: None

inference_flashinfer_mxfp8_token_capacity: NotRequired[int | None]#: None

moe_router_num_groups: NotRequired[int | None]#: None

moe_router_group_topk: NotRequired[int | None]#: None

transformer_impl: NotRequired[str]#: None

cuda_graph_impl: NotRequired[str]#: None

cuda_graph_modules: NotRequired[str | list[str]]#: None

cuda_graph_warmup_steps: NotRequired[int]#: None

moe_pad_experts_for_cuda_graph_inference: NotRequired[bool]#: None

moe_shared_expert_overlap: bool#: None

fine_grained_activation_offloading: NotRequired[bool]#: None

offload_modules: NotRequired[list[str] | None]#: None

use_gloo_process_groups: NotRequired[bool]#: None

distributed_timeout_minutes: NotRequired[float | None]#: None

moe_grouped_gemm: NotRequired[bool]#: None

moe_flex_dispatcher_backend: NotRequired[str]#: None

moe_hybridep_num_sms: NotRequired[int]#: None

moe_hybridep_prepad_packed_inputs: NotRequired[bool]#: None

hybridep_num_ranks_per_nvlink_domain: NotRequired[int]#: None

hybridep_use_mnnvl: NotRequired[bool]#: None

peft: NotRequired[nemo_rl.models.policy.MegatronPeftConfig | nemo_rl.models.policy.MegatronPeftConfigDisabled]#: None

optimizer: nemo_rl.models.policy.MegatronOptimizerConfig#: None

scheduler: nemo_rl.models.policy.MegatronSchedulerConfig#: None

distributed_data_parallel_config: nemo_rl.models.policy.MegatronDDPConfig#: None

checkpoint: NotRequired[nemo_rl.models.policy.MegatronCheckpointConfig]#: None

gradient_accumulation_fusion: NotRequired[bool]#: None

use_fused_weighted_squared_relu: NotRequired[bool]#: None

use_fused_linear_logprobs: NotRequired[bool]#: None

fused_linear_logprobs_chunk_size: NotRequired[int]#: None

fp32_lm_head: NotRequired[bool]#: None

mtp_num_layers: NotRequired[int]#: None

mtp_loss_scaling_factor: NotRequired[float]#: None

train_iters: NotRequired[int]#: None

mtp_use_repeated_layer: NotRequired[bool]#: None

mtp_detach_heads: NotRequired[bool]#: None

clear_memory_caches_before_refit: NotRequired[bool]#: None

fp8_cfg: NotRequired[nemo_rl.models.policy.Fp8Config]#: None

fp4_cfg: NotRequired[nemo_rl.models.policy.Fp4Config]#: None

first_last_layers_bf16: NotRequired[bool]#: None

num_layers_at_start_in_bf16: NotRequired[int]#: None

num_layers_at_end_in_bf16: NotRequired[int]#: None

te_precision_config_file: NotRequired[str]#: None

freeze_config: NotRequired[dict[str, Any]]#: None

batch_invariant_mode: NotRequired[bool]#: None

batch_invariant_backend: NotRequired[Literal[deepgemm, te_native, triton]]#: None

batch_invariant_collective: NotRequired[Literal[multimem, ordered]]#: None

flash_attention_version: NotRequired[Literal[2, 3, 4] | None]#: None

zero_train_gen_mismatch: NotRequired[bool]#: None

nemo_rl.models.policy.BLOCK_DRAFT_ALGOS#: (‘dspark’, ‘dflash’)

nemo_rl.models.policy.DRAFT_ALGOS#: (‘eagle3’,)

class nemo_rl.models.policy.TokenizerConfig#

Bases: typing.TypedDict

name: str#: None

chat_template: NotRequired[str | None]#: None

chat_template_kwargs: NotRequired[dict[str, Any] | None]#: None

tokenizer_kwargs: NotRequired[dict[str, Any] | None]#: None

audio: NotRequired[dict[str, Any]]#: None

video: NotRequired[dict[str, Any]]#: None

use_processor: NotRequired[bool]#: None

use_fastokens: NotRequired[bool]#: None

class nemo_rl.models.policy.PytorchOptimizerConfig#

Bases: typing.TypedDict

name: str#: None

kwargs: dict[str, Any]#: None

class nemo_rl.models.policy.SinglePytorchSchedulerConfig#

Bases: typing.TypedDict

name: str#: None

kwargs: dict[str, Any]#: None

class nemo_rl.models.policy.SinglePytorchMilestonesConfig#

Bases: typing.TypedDict

milestones: list[int]#: None

nemo_rl.models.policy.SchedulerMilestones#: None

class nemo_rl.models.policy.DynamicBatchingConfigDisabled#

Bases: typing.TypedDict

enabled: Literal[False]#: None

class nemo_rl.models.policy.DynamicBatchingConfig#

Bases: typing.TypedDict

enabled: Literal[True]#: None

train_mb_tokens: int#: None

logprob_mb_tokens: NotRequired[int]#: None

sequence_length_round: int#: None

class nemo_rl.models.policy.RouterReplayConfigDisabled#

Bases: typing.TypedDict

enabled: Literal[False]#: None

class nemo_rl.models.policy.RouterReplayConfig#

Bases: typing.TypedDict

enabled: Literal[True]#: None

class nemo_rl.models.policy.OnPolicyDistillationFullTransport#

Bases: typing.TypedDict

Resolved full-vocabulary MOPD settings carried to the policy workers.

A model_dump of OnPolicyDistillationFullConfig plus the resolved payload_field and teacher_index_field. That BaseModel remains the authoritative schema and the only place defaults are declared, so readers must take these keys as required rather than supplying their own fallbacks.

Initialization

Initialize self. See help(type(self)) for accurate signature.