`bridge.training.optim`#

Module Contents#

Functions#

`setup_optimizer`	Set up the optimizer and scheduler.
`_optimizer_state_is_fp32_adam`	Return whether a state dict contains only FP32 Adam moment tensors.
`_get_te_fused_adam_class`	Return Transformer Engine’s FusedAdam class when it is available.
`memory_efficient_fp32_optimizer_state_loading`	Avoid redundant TE FusedAdam state copies during checkpoint loading.
`sync_hybrid_device_optimizer_fp32_master_copies`	Refresh `HybridDeviceOptimizer` FP32 master copies from BF16 model parameters.
`_get_scheduler`	Get the optimizer parameter scheduler.

Data#

G_LOGGER

API#

bridge.training.optim.G_LOGGER#: ‘getLogger(…)’

bridge.training.optim.setup_optimizer( optimizer_config: megatron.core.optimizer.OptimizerConfig, scheduler_config: megatron.bridge.training.config.SchedulerConfig, model: Union[megatron.core.transformer.module.MegatronModule, list[megatron.core.transformer.module.MegatronModule]], use_gloo_process_groups: bool = False, pg_collection: Optional[megatron.core.process_groups_config.ProcessGroupCollection] = None, optimizer_config_override_provider: Optional[megatron.bridge.training.config.OptimizerConfigOverrideProvider] = None, ) → tuple[megatron.core.optimizer.MegatronOptimizer, megatron.core.optimizer_param_scheduler.OptimizerParamScheduler]#

Set up the optimizer and scheduler.

Parameters:

optimizer_config – Configuration for the optimizer
scheduler_config – Configuration for the scheduler
model – The model to optimize
use_gloo_process_groups – Whether to use Gloo process groups
pg_collection – Optional process group collection for distributed training

Returns:

tuple containing the optimizer and scheduler

bridge.training.optim._optimizer_state_is_fp32_adam( state_dict: collections.abc.Mapping[str, object], ) → bool#: Return whether a state dict contains only FP32 Adam moment tensors.

bridge.training.optim._get_te_fused_adam_class() → type[torch.optim.Optimizer] | None#: Return Transformer Engine’s FusedAdam class when it is available.

bridge.training.optim.memory_efficient_fp32_optimizer_state_loading( optimizer: megatron.core.optimizer.MegatronOptimizer | None, ) → collections.abc.Iterator[int]#

Avoid redundant TE FusedAdam state copies during checkpoint loading.

Megatron-Core’s distributed optimizer creates FP32 Adam moment tensors as its loading scaffold. Transformer Engine’s loader replaces those tensors with another FP32 copy even though no conversion is needed. This context temporarily uses PyTorch’s base loader only for compatible standard distributed FusedAdam instances, allowing them to adopt the scaffold tensors directly.

Parameters:: optimizer – Optimizer participating in checkpoint loading.
Yields:: Number of compatible FusedAdam instances using the base loader.

bridge.training.optim.sync_hybrid_device_optimizer_fp32_master_copies( optimizer: megatron.core.optimizer.MegatronOptimizer | None, ) → bool#

Refresh HybridDeviceOptimizer FP32 master copies from BF16 model parameters.

Workaround for an upstream Megatron-Core gap: when a checkpoint is loaded into the BF16 model parameters, reload_model_params() only refreshes the level-1 FP32 GPU shards inside HybridDeviceOptimizer. The level-2 CPU clones (gpu_params_map_cpu_copy) and the level-3 FP32 working copy (param_to_fp32_param) keep their default initialisation values.

Without this helper, the first optimizer step on a fine-tuning run that combines optimizer_cpu_offload=True + distributed optimizer + BF16 runs Adam on stale FP32 masters and writes the result back into the BF16 model, effectively reverting the loaded weights to fresh nn.Module random init. Training loss looks plausible at step 1 and collapses at step 2 because the model is no longer the one loaded from the checkpoint.

Mirrors the workaround in NVIDIA-NeMo/RL PR #2372. Once mcore’s reload_model_params() walks all three FP32 levels, this helper can be removed from both Bridge and RL.

Parameters:: optimizer – The Megatron optimizer returned by :func:setup_optimizer. No-op when None or when no sub-optimizer wraps a HybridDeviceOptimizer (i.e. CPU offload is not enabled).
Returns:: True when at least one HybridDeviceOptimizer sub-optimizer was synced; False otherwise.

bridge.training.optim._get_scheduler( optimizer_config: megatron.core.optimizer.OptimizerConfig, scheduler_config: megatron.bridge.training.config.SchedulerConfig, optimizer: megatron.core.optimizer.MegatronOptimizer, ) → megatron.core.optimizer_param_scheduler.OptimizerParamScheduler#

Get the optimizer parameter scheduler.