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.
_cpu_staging_get_unscaled_state	Wrap a TE state accessor without depending on its call signature.
memory_efficient_precision_aware_optimizer_state_checkpointing	Stage expanded precision-aware Adam checkpoint tensors on CPU.
memory_efficient_fp32_optimizer_state_loading	Avoid redundant TE FusedAdam state copies during checkpoint loading.
hybrid_optimizer_state_loading	Restore native FP32 HybridDeviceOptimizer metadata omitted from parameter payloads.
sync_hybrid_device_optimizer_fp32_master_copies	Synchronize HybridDeviceOptimizer working copies after checkpoint loading.
_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._cpu_staging_get_unscaled_state(

fallback: collections.abc.Callable[..., object],

) → collections.abc.Callable[..., torch.Tensor]

Wrap a TE state accessor without depending on its call signature.

bridge.training.optim.memory_efficient_precision_aware_optimizer_state_checkpointing(

optimizer: megatron.core.optimizer.MegatronOptimizer | None,

*,

enabled: bool,

) → collections.abc.Iterator[int]

Stage expanded precision-aware Adam checkpoint tensors on CPU.

Transformer Engine’s precision-aware FusedAdam exposes portable checkpoint state by expanding compressed moments to FP32. Its default state_dict path retains those expansions on GPU until the complete state dictionary is built for saving or load scaffolding, which can exceed device memory even when training fits. This context preserves the same unscaled checkpoint values and dtypes, but moves each tensor to CPU immediately so only one expansion is live on GPU at a time.

Parameters:

• optimizer – Optimizer participating in checkpointing.

• enabled – Whether to stage compatible optimizer state on CPU.

Yields:

Number of compatible FusedAdam instances using CPU staging.

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.hybrid_optimizer_state_loading(

optimizer: megatron.core.optimizer.MegatronOptimizer | None,

) → collections.abc.Iterator[None]

Restore native FP32 HybridDeviceOptimizer metadata omitted from parameter payloads.

Core restores Adam counters in common group metadata, while its parameter loader expects every existing state tensor in the payload. For a hybrid optimizer this includes a GPU step tensor introduced during loading and, in fully reshardable checkpoints, the CPU working-copy master_param alias. Supply the saved group counter and checkpoint FP32 parameter for those fields before calling Core’s ordinary tensor-copy implementation.

Parameters:

optimizer – Optimizer whose native checkpoint state is being loaded.

bridge.training.optim.sync_hybrid_device_optimizer_fp32_master_copies(

optimizer: megatron.core.optimizer.MegatronOptimizer | None,

) → bool

Synchronize HybridDeviceOptimizer working copies after checkpoint loading.

A freshly constructed optimizer has no state when only model weights were loaded. Refresh its FP32 working copies from those model weights. When optimizer state was restored, preserve its full-precision masters and synchronize the CPU/GPU sub-optimizers with that state instead. Rebuilding restored masters from BF16 weights discards their low bits and changes the next update. The pinned Core dp_reshardable loader also replaces state tensors without refreshing HybridDeviceOptimizer’s working copies.

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.

The model-only branch mirrors NVIDIA-NeMo/RL PR #2372. This helper can be removed once Core handles model-only master refresh, full-state rebinding, and the restored sub-optimizer step counters.

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.

Parameters:

• optimizer_config – Configuration for the optimizer

• scheduler_config – Configuration for the scheduler

• optimizer – The optimizer to schedule

Returns:

The optimizer parameter scheduler