`nemo_rl.models.megatron.setup`#

Module Contents#

Classes#

MoEFloat16Module

Float 16 Module with the ability to keep the expert bias in float32.

Functions#

`_patch_hf_config_double_instantiation`	Patch HF config classes whose post_init fails with Megatron’s recursive instantiation.
`_force_sync_optimizer_fp32_from_model`	Force-sync the distributed optimizer’s FP32 master copies from the BF16 model params.
`_force_sync_model_from_optimizer_fp32`	Restore BF16 compute weights from loaded HybridDeviceOptimizer masters.
`_sync_consumed_samples_on_scheduler_load`	Keep Megatron-Bridge from resetting the LR schedule on resume.
`_resolve_optimizer_dtype_kwargs`	Resolve optimizer dtype strings, including TE’s uint8-backed FP8 moments.
`enable_batch_invariant_mode`	Validate, then enable Megatron-Core batch-invariant kernels before CUDA init.
`destroy_parallel_state`	Safely destroy parallel state and reset async call tracking.
`configure_refit_environment`	Set the refit allocator mode before NCCL caches the value.
`_distributed_timeout_minutes`	Return `megatron_cfg.distributed_timeout_minutes`, or None when unset.
`setup_distributed`	Handle NCCL settings, dtype mapping, and basic config setup.
`validate_and_set_config`
`_canonicalize_hf_config_overrides`	Return a stable JSON string for hf_config_overrides.
`_get_hf_config_overrides_hash`	Return a short stable hash for hf_config_overrides.
`_resolve_iter_dir_from_root`	Resolve the latest iteration directory under `path`.
`_resolve_peft_restore_dir`	Resolve a `megatron_cfg.peft.restore_from` path to an iteration directory.
`_validate_peft_restore_config`	Fail closed if the donor checkpoint’s PEFT config is incompatible.
`_create_peft_warm_start_hook`	Create a pre-wrap hook that warm-starts LoRA adapters from a donor checkpoint.
`validate_model_paths`	Validate and setup model paths.
`validate_megatron_config`	Validate Bridge config, then preserve explicit NeMo-RL prepadding.
`setup_model_config`	Handle all the model configuration logic.
`_validate_model_override_conflicts`	Reject overrides that duplicate first-class NeMo-RL Megatron settings.
`_merge_model_overrides`	Construct a model config with recursively merged user overrides.
`_merge_model_override_value`	Merge one override value, descending into config hierarchies as needed.
`_apply_parallelism_config`	Apply tensor/pipeline/context parallelism configuration.
`_apply_multimodal_config`	Map legacy Omni freeze controls onto canonical provider attributes.
`_apply_moe_config`	Apply Mixture of Experts configuration.
`_apply_mtp_config`	Apply Multi-Token Prediction settings onto the mcore model config.
`_quant_recipe_name`
`_validate_te_precision_config`
`_validate_nvfp4_pertoken_precision_recipe`	Check effective module precision under the outer NVFP4 autocast context.
`_apply_precision_config`	Apply dtype and Transformer Engine precision configuration.
`_apply_performance_config`	Apply performance optimization configuration.
`_validate_optimizer_config`	Validate optimizer configuration.
`_validate_chunking_config`	Validate chunking configuration.
`_create_checkpoint_config`	Create checkpoint configurations.
`_validate_training_config`	Validate training configuration.
`_validate_dtype_config`
`_create_megatron_config`	Create the final Megatron configuration container.
`_create_draft_pre_wrap_hook`	Create the hook that attaches draft weights before mixed-precision/DDP wrapping.
`_patch_bridge_signal_handler_for_worker_threads`	Make Megatron-Bridge’s signal-handler install safe off the main thread.
`build_inference_model`	Build a second, inference-layout model for colocated Megatron refit.
`setup_model_and_optimizer`
`handle_model_import`	Convert and cache the initial model checkpoint if it does not yet exist.
`setup_reference_model_state`	Setup the reference model for inference and return its state dict.
`load_teacher_output_layer_weight`	Load this rank’s shard of a teacher checkpoint’s LM-head weight.
`finalize_megatron_setup`	Finalize the setup with remaining configurations.
`make_policy_like_config`	Adapt a ValueConfig to look like a PolicyConfig for reusing setup functions.

Data#

`_HF_CONFIG_PATCHED`
`_NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT`
`TokenizerType`
`_OPTIMIZER_DTYPE_KEYS`
`_BRIDGE_SIGNAL_HANDLER_PATCHED`

API#

nemo_rl.models.megatron.setup._HF_CONFIG_PATCHED#: False

nemo_rl.models.megatron.setup._NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT#: ‘expanded_sequence_v1’

nemo_rl.models.megatron.setup._patch_hf_config_double_instantiation()#

Patch HF config classes whose post_init fails with Megatron’s recursive instantiation.

Megatron-LM’s instantiate_utils recursively instantiates all nested configs that have a target key. Some HF config classes (e.g. Qwen3OmniMoeTalkerConfig) then try to re-instantiate those nested configs in post_init via ** unpacking, which fails because the value is already an object, not a dict.

This adds isinstance guards so the post_init is a no-op when the nested config is already the correct type.

nemo_rl.models.megatron.setup._force_sync_optimizer_fp32_from_model(optimizer, model)#

Force-sync the distributed optimizer’s FP32 master copies from the BF16 model params.

With HybridDeviceOptimizer (selected by optimizer_cpu_offload=True) three parallel parameter copies exist that all need pretrained values after a fine-tune-style checkpoint load:

shard_fp32_from_float16_groups – per-DP-rank FP32 GPU shard used as the Adam master parameter.
hdo.gpu_params_map_cpu_copy – CPU clones the CPU sub-optimizer steps against (this is what makes “cpu offload” actually offload).
hdo.param_to_fp32_param – an additional FP32 working copy that HybridDeviceOptimizer keeps so it can do its async D2H/H2D dance without aliasing.

Vanilla load_checkpoint only refreshes the BF16 model parameters and then calls reload_model_params() which currently only walks level 1 for HybridDeviceOptimizer. Levels 2 and 3 keep their default (random) init.

Failure mode without this helper (the optimizer_cpu_offload=True + HF -> mcore + finetune=True path): the first optimizer step does Adam on the stale FP32 master, then writes the result into BF16. BF16 now approximately equals random init plus a tiny Adam delta, and every subsequent forward / refit / rollout uses an essentially untrained model. The training loss looks plausible, but RL reward collapses, KL explodes, and the inference engine produces garbage.

This helper propagates BF16 -> all three FP32 levels right after the checkpoint load to avoid that reversion.

nemo_rl.models.megatron.setup._force_sync_model_from_optimizer_fp32(optimizer)#

Restore BF16 compute weights from loaded HybridDeviceOptimizer masters.

On a full checkpoint resume, HybridDeviceOptimizer.load_state_dict restores its FP32 master parameters, but the BF16 parameter shards used by the first forward can still contain the pre-load values. The first optimizer step copies the masters back and masks the problem from subsequent steps, producing a one-step loss/reward discontinuity exactly at the resume boundary.

Copy each loaded FP32 working parameter back to its BF16 shard and then force a synchronous DP parameter all-gather before any reference-policy or training forward. This is the inverse of _force_sync_optimizer_fp32_from_model and is only called for a genuine optimizer-state resume.

nemo_rl.models.megatron.setup.TokenizerType#: ‘TypeVar(…)’

nemo_rl.models.megatron.setup._OPTIMIZER_DTYPE_KEYS#: (‘params_dtype’, ‘main_grads_dtype’, ‘main_params_dtype’, ‘exp_avg_dtype’, ‘exp_avg_sq_dtype’)

nemo_rl.models.megatron.setup._sync_consumed_samples_on_scheduler_load( state: megatron.bridge.training.state.GlobalState, scheduler: Any, ) → collections.abc.Iterator[None]#

Keep Megatron-Bridge from resetting the LR schedule on resume.

When override_opt_param_scheduler=True, Megatron-Bridge will override the scheduler step with the checkpointed train_state.consumed_train_samples after loading the scheduler from the checkpoint. Because NeMo-RL doesn’t use this counter, older checkpoints have this dummy state set to 0, which then overrides the actual scheduler.num_steps checkpoint value.

To support forward compatibility matching Megatron, NeMo-RL should begin to checkpoint train_state.consumed_train_samples = scheduler.num_steps to identify / sync these two values when saving a checkpoint.

To support backward compatibility, this context manager will do the same after checkpoint load to maintain the checkpointed scheduler.num_steps.

HACK(@cspades): When https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/6321 is merged, we should use checkpoint_state_migration_hook to implement the backward-compatibility logic given the state_dict and train_state.

load_checkpoint(
    state,
    model,
    optimizer,
    ...
    checkpoint_state_migration_hook=checkpoint_compat_hook,
)

nemo_rl.models.megatron.setup._resolve_optimizer_dtype_kwargs( optimizer_cfg: dict[str, Any], ) → dict[str, Any]#: Resolve optimizer dtype strings, including TE’s uint8-backed FP8 moments.

nemo_rl.models.megatron.setup.enable_batch_invariant_mode( config: nemo_rl.models.policy.PolicyConfig, ) → None#: Validate, then enable Megatron-Core batch-invariant kernels before CUDA init.

nemo_rl.models.megatron.setup.destroy_parallel_state()#

Safely destroy parallel state and reset async call tracking.

This function is called during initialization to clean up temporary distributed state from model import operations. Resetting async call tracking ensures that when the main Megatron distributed context is created, all ranks start with consistent call_idx values for async checkpointing.

nemo_rl.models.megatron.setup.configure_refit_environment(config) → None#: Set the refit allocator mode before NCCL caches the value.

nemo_rl.models.megatron.setup._distributed_timeout_minutes( config: nemo_rl.models.policy.PolicyConfig | nemo_rl.models.value.config.ValueConfig, ) → Optional[float]#: Return megatron_cfg.distributed_timeout_minutes, or None when unset.

nemo_rl.models.megatron.setup.setup_distributed(config) → None#: Handle NCCL settings, dtype mapping, and basic config setup.

nemo_rl.models.megatron.setup.validate_and_set_config( config, rank, hf_model_name, pretrained_path, weights_path, optimizer_path, *, skip_weight_load: bool = False, )#

nemo_rl.models.megatron.setup._canonicalize_hf_config_overrides( overrides: dict[str, Any], ) → str#: Return a stable JSON string for hf_config_overrides.

nemo_rl.models.megatron.setup._get_hf_config_overrides_hash(overrides: dict[str, Any]) → str#: Return a short stable hash for hf_config_overrides.

nemo_rl.models.megatron.setup._resolve_iter_dir_from_root( path: str, not_found_msg: str, config_key: str = 'pretrained_checkpoint.path', ) → str#

Resolve the latest iteration directory under path.

Checks latest_checkpointed_iteration.txt first; falls back to scanning for iter_* subdirectories and taking the last one (lexicographic order).

nemo_rl.models.megatron.setup._resolve_peft_restore_dir(restore_from: str) → str#

Resolve a megatron_cfg.peft.restore_from path to an iteration directory.

Accepts either a specific iteration directory (containing run_config.yaml) or a checkpoint root with a latest_checkpointed_iteration.txt tracker file / iter_* subdirectories.

nemo_rl.models.megatron.setup._validate_peft_restore_config( restore_dir: str, peft_cfg: nemo_rl.models.policy.MegatronPeftConfig, ) → None#

Fail closed if the donor checkpoint’s PEFT config is incompatible.

Adapter tensors are loadable only when the rank (dim) matches, and only functionally equivalent when the scaling (alpha) matches, so both must agree with this run’s megatron_cfg.peft before any weights are loaded. The targeted module sets must match as well: the adapter-only distributed load only raises on donor keys missing from the checkpoint, while unrequested donor keys are discarded silently (DCP’s ASSUME_OK_UNEXPECTED skips the mismatch check), so a superset donor would otherwise load with no diagnostics and leave this run’s extra adapters at fresh init.

nemo_rl.models.megatron.setup._create_peft_warm_start_hook( megatron_cfg: megatron.bridge.training.config.ConfigContainer, state: megatron.bridge.training.state.GlobalState, restore_dir: str, ) → Callable[[list[megatron.core.transformer.MegatronModule]], list[megatron.core.transformer.MegatronModule]]#

Create a pre-wrap hook that warm-starts LoRA adapters from a donor checkpoint.

The hook must run immediately after the PEFT pre-wrap hook (which loads the pretrained base weights and attaches fresh adapters). It routes the donor load through Megatron-Bridge’s PEFT-resume path: loading from checkpoint.load with PEFT configured and finetune=False filters the generated sharded state dict to adapter keys and drops to a non-strict load, so an adapter-only donor checkpoint restores cleanly onto the fresh base with distributed resharding handled by the torch_dist loader.

nemo_rl.models.megatron.setup.validate_model_paths( config: nemo_rl.models.policy.PolicyConfig, ) → tuple[str, str, bool]#

Validate and setup model paths.

Returns:

hf_model_name is the HuggingFace model name / path used for architecture config resolution and tokenizer setup.
pretrained_path is the path of the checkpoint that will be used as the pretrained starting point. For megatron_bridge format this is resolved to the specific iteration directory containing run_config.yaml. For megatron_lm format this is resolved to the specific iteration directory (via latest_checkpointed_iteration.txt or by scanning iter_* subdirs if a root dir is provided, since the bridge does not resolve iterations itself). For the default HF path this is the Megatron-Bridge cache directory.
pt_checkpoint_exists is True when the checkpoint at pretrained_path is already present and does not need to be created.

Return type:

A (hf_model_name, pretrained_path, pt_checkpoint_exists) tuple where

nemo_rl.models.megatron.setup.validate_megatron_config( megatron_cfg: Any, config: collections.abc.Mapping[str, Any], ) → None#: Validate Bridge config, then preserve explicit NeMo-RL prepadding.

nemo_rl.models.megatron.setup.setup_model_config( config: nemo_rl.models.policy.PolicyConfig, rank, dtype, hf_model_name: str, pretrained_path: str, weights_path: Optional[str] = None, optimizer_path: Optional[str] = None, *, skip_weight_load: bool = False, ) → tuple[megatron.bridge.training.config.ConfigContainer, Any]#

Handle all the model configuration logic.

Parameters:

config – Policy config.
rank – Global rank (used in error messages).
dtype – Training dtype.
hf_model_name – HF model id (or local path).
pretrained_path – Path to the pretrained Megatron checkpoint.
weights_path – Path to save/load training weights.
optimizer_path – Path to the optimizer state (None if not resuming).
skip_weight_load – This policy never loads the checkpoint (weights arrive via refit).

nemo_rl.models.megatron.setup._validate_model_override_conflicts( megatron_cfg: collections.abc.Mapping[str, Any], overrides: dict[str, Any], ) → None#: Reject overrides that duplicate first-class NeMo-RL Megatron settings.

nemo_rl.models.megatron.setup._merge_model_overrides( target: Any, overrides: dict[str, Any], path: str = 'policy.megatron_cfg.model_overrides', ) → Any#

Construct a model config with recursively merged user overrides.

Dataclass-backed config objects are reconstructed with dataclasses.replace so the returned provider is the canonical object later stored and serialized by Megatron Bridge’s ConfigContainer. Nested mappings and config objects are copied before updates; the input provider is never mutated.

Parameters:

target – Model provider, nested config object, or mapping to merge.
overrides – YAML-derived override hierarchy.
path – User-facing config path used in error messages.

Returns:

A new config object or mapping containing the merged values.

Raises:

AttributeError – If an override does not match an object attribute.

nemo_rl.models.megatron.setup._merge_model_override_value( current_value: Any, override_value: Any, path: str, ) → Any#: Merge one override value, descending into config hierarchies as needed.

nemo_rl.models.megatron.setup._apply_parallelism_config( model_cfg: Any, config: nemo_rl.models.policy.PolicyConfig, ) → None#: Apply tensor/pipeline/context parallelism configuration.

nemo_rl.models.megatron.setup._apply_multimodal_config( model_cfg: Any, config: nemo_rl.models.policy.PolicyConfig, ) → None#: Map legacy Omni freeze controls onto canonical provider attributes.

nemo_rl.models.megatron.setup._apply_moe_config( model_cfg: Any, config: nemo_rl.models.policy.PolicyConfig, ) → None#: Apply Mixture of Experts configuration.

nemo_rl.models.megatron.setup._apply_mtp_config( model_cfg: Any, config: nemo_rl.models.policy.PolicyConfig, ) → None#: Apply Multi-Token Prediction settings onto the mcore model config.

nemo_rl.models.megatron.setup._quant_recipe_name(recipe: Any) → str | None#

nemo_rl.models.megatron.setup._validate_te_precision_config( quant_recipe: Any, fp8_cfg: collections.abc.Mapping[str, Any] | None, ) → None#

nemo_rl.models.megatron.setup._validate_nvfp4_pertoken_precision_recipe( quant_recipe: megatron.core.quantization.quant_config.RecipeConfig, ) → None#: Check effective module precision under the outer NVFP4 autocast context.

nemo_rl.models.megatron.setup._apply_precision_config( model_cfg: Any, config: nemo_rl.models.policy.PolicyConfig, dtype: torch.dtype, ) → None#: Apply dtype and Transformer Engine precision configuration.

nemo_rl.models.megatron.setup._apply_performance_config( model_cfg: Any, config: nemo_rl.models.policy.PolicyConfig, ) → None#: Apply performance optimization configuration.

nemo_rl.models.megatron.setup._validate_optimizer_config( config: nemo_rl.models.policy.PolicyConfig, ) → None#: Validate optimizer configuration.

nemo_rl.models.megatron.setup._validate_chunking_config( config: nemo_rl.models.policy.PolicyConfig, ) → None#: Validate chunking configuration.

nemo_rl.models.megatron.setup._create_checkpoint_config( pretrained_path: Optional[str], weights_path: Optional[str], optimizer_path: Optional[str], load_main_params_from_ckpt: bool = False, ckpt_cfg: Optional[dict[str, Any]] = None, ) → megatron.bridge.training.config.CheckpointConfig#

Create checkpoint configurations.

Parameters:

pretrained_path – Path to the pretrained checkpoint.
weights_path – Path to save/load training weights.
optimizer_path – Path to the optimizer state (None if not resuming optimizer).
load_main_params_from_ckpt – Load optimizer main params from the checkpoint.
ckpt_cfg – MegatronCheckpointConfig dict from YAML (megatron_cfg.checkpoint). Every knob (async_save, ckpt_assume_constant_structure, and the parallel-IO fields) is forwarded only when explicitly set in YAML — no call-site default. When a field (or the whole block) is absent, Megatron Bridge’s own CheckpointConfig default applies, so async_save falls back to synchronous save for configs that don’t set it.

nemo_rl.models.megatron.setup._validate_training_config( config: nemo_rl.models.policy.PolicyConfig, model_cfg: Any, ) → None#: Validate training configuration.

nemo_rl.models.megatron.setup._validate_dtype_config( dtype: torch.dtype, model_cfg: Any, optimizer_cfg: Any, ) → None#

nemo_rl.models.megatron.setup._create_megatron_config( model_cfg: Any, checkpoint_config: megatron.bridge.training.config.CheckpointConfig, config: nemo_rl.models.policy.PolicyConfig, hf_model_name: str, dtype: torch.dtype, fp8_param_enabled: bool = False, ) → megatron.bridge.training.config.ConfigContainer#: Create the final Megatron configuration container.

nemo_rl.models.megatron.setup._create_draft_pre_wrap_hook( policy_cfg: nemo_rl.models.policy.PolicyConfig, megatron_cfg: megatron.bridge.training.config.ConfigContainer, state: megatron.bridge.training.state.GlobalState, *, preload_policy_from_pretrained: bool, ) → Callable[[list[megatron.core.transformer.MegatronModule]], list[megatron.core.transformer.MegatronModule]]#: Create the hook that attaches draft weights before mixed-precision/DDP wrapping.

nemo_rl.models.megatron.setup._BRIDGE_SIGNAL_HANDLER_PATCHED#: False

nemo_rl.models.megatron.setup._patch_bridge_signal_handler_for_worker_threads() → None#

Make Megatron-Bridge’s signal-handler install safe off the main thread.

See https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/4375

TODO: Remove this hotfix once Megatron-Bridge is bumped.

nemo_rl.models.megatron.setup.build_inference_model( policy_cfg: nemo_rl.models.policy.PolicyConfig, megatron_cfg: megatron.bridge.training.config.ConfigContainer, initial_model_provider: megatron.bridge.models.model_provider.ModelProviderMixin, ) → megatron.core.transformer.MegatronModule#

Build a second, inference-layout model for colocated Megatron refit.

The returned model is resident on GPU; its weights are uninitialized until the first reshard.

Parameters:

policy_cfg – The inference config
megatron_cfg – The training config
initial_model_provider – Pre-wrap provider snapshot taken by setup_model_and_optimizer.

Returns:

The inference model module (single element; not DDP-wrapped, no optimizer).

nemo_rl.models.megatron.setup.setup_model_and_optimizer( policy_cfg: nemo_rl.models.policy.PolicyConfig, megatron_cfg: megatron.bridge.training.config.ConfigContainer, load_optimizer: bool = True, get_embedding_ranks=None, get_position_embedding_ranks=None, pre_load_checkpoint_hook: Optional[Callable] = None, additional_pre_wrap_hooks: Optional[list[Callable]] = None, load_weights: bool = True, )#

nemo_rl.models.megatron.setup.handle_model_import( config: nemo_rl.models.policy.PolicyConfig, hf_model_name: str, pretrained_path: str, pt_checkpoint_exists: bool, model_post_wrap_hook: Optional[Callable] = None, transformer_layer_spec: Optional[Any] = None, mamba_stack_spec: Optional[Any] = None, ) → None#

Convert and cache the initial model checkpoint if it does not yet exist.

Behaviour depends on policy.pretrained_checkpoint.format:

"megatron_bridge": The checkpoint is already in the correct format; no conversion is performed.
"megatron_lm": Megatron-Bridge can load torch_dist MLM checkpoints directly (the bridge falls back to extracting config from the state dict when run_config.yaml is absent), so no conversion is performed.
No pretrained_checkpoint (default): The HuggingFace model identified by hf_model_name is converted to Megatron-Bridge format (existing behaviour).

The force_reconvert_from_hf flag forces the HF conversion to run again even if the output already exists. It has no effect for megatron_bridge or megatron_lm formats.

Parameters:

config – Policy config used for pretrained_checkpoint, hf_config_overrides, and megatron_cfg.
hf_model_name – HF model id (or local path) to import.
pretrained_path – Output directory for the Megatron checkpoint.
pt_checkpoint_exists – Whether a Megatron checkpoint already exists at pretrained_path. If True and force_reconvert_from_hf is False, the import is skipped.
model_post_wrap_hook –
Optional callable forwarded to

func:

import_model_from_hf_name. Invoked on each Megatron model chunk after it is built (and before DDP wrapping).
transformer_layer_spec – Optional Megatron ModuleSpec (or callable returning one) overriding the default layer spec from the model provider.
mamba_stack_spec – Optional Megatron ModuleSpec (or callable returning one) overriding the default stack spec from Mamba model providers.

nemo_rl.models.megatron.setup.setup_reference_model_state( config: nemo_rl.models.policy.PolicyConfig, megatron_cfg: megatron.bridge.training.config.ConfigContainer, pretrained_path: str, pre_load_checkpoint_hook: Optional[Callable] = None, ) → dict#: Setup the reference model for inference and return its state dict.

nemo_rl.models.megatron.setup.load_teacher_output_layer_weight( *, teacher_pretrained_path: str, local_vocab_size: Optional[int], dtype: Optional[torch.dtype], ) → Optional[torch.Tensor]#

Load this rank’s shard of a teacher checkpoint’s LM-head weight.

Full-vocabulary MOPD ships the teacher’s hidden states and projects them on the student side, which needs the teacher’s output_layer.weight. Only that one tensor is read: instantiating a second Megatron model just to reach it is far more fragile, since provider and config objects accumulate runtime-only distributed state during live training.

The request is built at the student’s tensor-parallel rank and size. Re-sharding a teacher saved at a different tensor-parallel width is dist_checkpointing’s ordinary by-offset load and needs no special flag. allow_shape_mismatch=True covers a different case: a teacher whose padded vocabulary differs from the student’s. Megatron pads to a multiple of make_vocab_size_divisible_by * tp_size, so the two widths diverge whenever the parallel sizes do. Under that flag mcore zero-initializes and partially loads instead of raising, so rows above the narrower of the two padded widths come back as zeros. Those are pad slots rather than real vocabulary entries, so the objective is unaffected in practice – but this is a silent fallback, not a re-sharding mechanism.

Resolving the checkpoint iteration from a checkpoint root goes through Megatron-Bridge’s read_train_state, a broadcast_object_list over the whole student world, so under student pipeline parallelism the stages that own no output_layer must still call this function. They pass local_vocab_size=None to request nothing: the load then runs with an empty sharded state dict and returns only the checkpoint’s common state. dist_checkpointing.load itself is rank-local here – validate_access_integrity=False skips its only all-gather.

Parameters:

teacher_pretrained_path – Megatron checkpoint root of the teacher.
local_vocab_size – This rank’s vocabulary shard width, or None to request nothing.
dtype – Dtype to materialize the shard in. Unused, and expected to be None, when local_vocab_size is None.

Returns:

The [local_vocab_size, hidden_size] weight shard on CPU, or None when this rank requested nothing.

Raises:

FileNotFoundError – If the checkpoint root holds no readable iteration.
KeyError – If neither an output-layer nor a tied-embedding weight exists.
TypeError – If the loaded checkpoint entry is not a torch.Tensor.
ValueError – If the checkpoint tensor is not a rank-2 matrix, or if exactly one of local_vocab_size / dtype is None.

nemo_rl.models.megatron.setup.finalize_megatron_setup( config: nemo_rl.models.policy.PolicyConfig, megatron_cfg: megatron.bridge.training.config.ConfigContainer, hf_model_name: str, worker_sharding_annotations: nemo_rl.distributed.named_sharding.NamedSharding, model, optimizer, ) → tuple#

Finalize the setup with remaining configurations.

Returns:: Tuple of (megatron_tokenizer, megatron_bridge, should_disable_forward_pre_hook, dp_size)

class nemo_rl.models.megatron.setup.MoEFloat16Module( config: megatron.core.transformer.transformer_config.TransformerConfig, module: torch.nn.Module, )#

Bases: megatron.core.transformer.module.Float16Module

Float 16 Module with the ability to keep the expert bias in float32.

.. attribute:: config

Transformer config

Type:: TransformerConfig

.. attribute:: fp16

Specifies if the model runs in fp16 mode

Type:: bool

.. attribute:: bf16

Specifies if the model runs in bf16 mode

Type:: bool
Parameters:: config (TransformerConfig) – The transformer config used to initalize the model

Initialization

re_enable_float32_expert_bias() → None#

Ensure MoE router expert bias stays in float32 for numerical stability.

Walks the wrapped module to find MoE routers and invokes the _maintain_float32_expert_bias() helper which recreates or casts the expert bias tensors to float32 as required by Megatron-LM.

nemo_rl.models.megatron.setup.make_policy_like_config( config: nemo_rl.models.value.config.ValueConfig, ) → dict#

Adapt a ValueConfig to look like a PolicyConfig for reusing setup functions.

The Megatron setup functions expect PolicyConfig fields. This builds a compatible dict from the ValueConfig with the same shape as a PolicyConfig.

The output is deterministic for a given input — callers should cache the result rather than rebuilding on every call.

nemo_rl.models.megatron.setup#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.models.megatron.setup`#