`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`
`DTensorConfigDisabled`
`MoEParallelizerOptions`	MoE parallelizer config options (mirrors Automodel’s MoEParallelizerConfig).
`DTensorConfig`
`SequencePackingConfigDisabled`
`SequencePackingConfig`
`RewardModelConfig`
`MegatronPeftConfigDisabled`
`MegatronPeftConfig`
`MegatronOptimizerConfig`
`MegatronSchedulerConfig`
`MegatronDDPConfig`
`Fp8Config`
`MegatronConfigDisabled`
`MegatronCheckpointConfig`	Checkpoint knobs passed through to Megatron Bridge CheckpointConfig.
`MegatronConfig`
`DraftConfigDisabled`	Configuration shape for the disabled draft-model training path.
`DraftConfig`	Configuration for Eagle draft-model training alongside the policy model.
`TokenizerConfig`
`PytorchOptimizerConfig`
`SinglePytorchSchedulerConfig`
`SinglePytorchMilestonesConfig`
`DynamicBatchingConfigDisabled`
`DynamicBatchingConfig`
`RouterReplayConfigDisabled`
`RouterReplayConfig`
`PolicyConfig`

Functions#

_patch_transformers_tokenizer_class_set

Undo the transformers block on deepseek_v3 tokenizers.

Data#

SchedulerMilestones

API#

nemo_rl.models.policy._patch_transformers_tokenizer_class_set()#

Undo the transformers block on deepseek_v3 tokenizers.

Root cause: transformers 5.4-5.11 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 / DTensor / DTensor v2 all 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

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

enable_deepep: NotRequired[bool]#: 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.DTensorConfigDisabled#

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.DTensorConfig#

Bases: typing.TypedDict

enabled: Literal[True]#: None

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

_v2: NotRequired[bool]#: 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

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

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

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

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

force_clear_fp8_caches: NotRequired[bool]#: None

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

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

pipeline_dtype: str#: None

sequence_parallel: bool#: None

freeze_moe_router: 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

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

moe_pad_experts_for_cuda_graph_inference: NotRequired[bool]#: None

moe_shared_expert_overlap: bool#: None

use_gloo_process_groups: NotRequired[bool]#: None

moe_grouped_gemm: NotRequired[bool]#: None

moe_flex_dispatcher_backend: NotRequired[str]#: None

moe_hybridep_num_sms: NotRequired[int]#: 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

mtp_num_layers: NotRequired[int]#: None

mtp_loss_scaling_factor: NotRequired[float]#: 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

class nemo_rl.models.policy.DraftConfigDisabled#

Bases: typing.TypedDict

Configuration shape for the disabled draft-model training path.

Initialization

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

enabled: Literal[False]#: None

class nemo_rl.models.policy.DraftConfig#

Bases: typing.TypedDict

Configuration for Eagle draft-model training alongside the policy model.

Initialization

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

enabled: Literal[True]#: None

model_name: NotRequired[str | None]#: None

loss_weight: NotRequired[float]#: None

num_layers: NotRequired[int | None]#: None

aux_layer_indices: NotRequired[list[int] | None]#: None

class nemo_rl.models.policy.TokenizerConfig#

Bases: typing.TypedDict

name: str#: None

chat_template: NotRequired[str]#: None

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

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

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

use_processor: 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.PolicyConfig#

Bases: typing.TypedDict

model_name: str#: None

tokenizer: nemo_rl.models.policy.TokenizerConfig#: None

train_global_batch_size: int#: None

train_micro_batch_size: int#: None

logprob_batch_size: NotRequired[int]#: None

logprob_chunk_size: NotRequired[int | None]#: None

generation: NotRequired[nemo_rl.models.generation.interfaces.GenerationConfig]#: None

generation_batch_size: NotRequired[int]#: None

precision: str#: None

reward_model_cfg: NotRequired[nemo_rl.models.policy.RewardModelConfig]#: None

dtensor_cfg: nemo_rl.models.policy.DTensorConfig | nemo_rl.models.policy.DTensorConfigDisabled#: None

megatron_cfg: NotRequired[nemo_rl.models.policy.MegatronConfig | nemo_rl.models.policy.MegatronConfigDisabled]#: None

draft: NotRequired[nemo_rl.models.policy.DraftConfig | nemo_rl.models.policy.DraftConfigDisabled]#: None

pretrained_checkpoint: NotRequired[nemo_rl.utils.checkpoint.PretrainedCheckpointConfig]#: None

router_replay: NotRequired[nemo_rl.models.policy.RouterReplayConfig | nemo_rl.models.policy.RouterReplayConfigDisabled]#: None

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

dynamic_batching: nemo_rl.models.policy.DynamicBatchingConfig | nemo_rl.models.policy.DynamicBatchingConfigDisabled#: None

sequence_packing: NotRequired[nemo_rl.models.policy.SequencePackingConfig | nemo_rl.models.policy.SequencePackingConfigDisabled]#: None

make_sequence_length_divisible_by: int#: None

max_total_sequence_length: int#: None

max_grad_norm: NotRequired[float | int | None]#: None

refit_buffer_size_gb: NotRequired[float]#: None

optimizer: NotRequired[nemo_rl.models.policy.PytorchOptimizerConfig | None]#: None

scheduler: NotRequired[list[nemo_rl.models.policy.SinglePytorchSchedulerConfig | nemo_rl.models.policy.SinglePytorchMilestonesConfig] | nemo_rl.models.policy.SchedulerMilestones | None]#: None

quant_cfg: NotRequired[str | None]#: None

quant_calib_data: NotRequired[str | None]#: None

quant_calib_size: NotRequired[int | None]#: None

quant_batch_size: NotRequired[int | None]#: None

quant_sequence_length: NotRequired[int | None]#: None

disable_modelopt_layer_spec: NotRequired[bool]#: None

nemo_rl.models.policy#

Subpackages#

Submodules#

Package Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.models.policy`#