`bridge.training.optim`#

Module Contents#

Functions#

`setup_optimizer`	Set up the optimizer and scheduler.
`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.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.