bridge.recipes.utils.optimizer_utils

Module Contents

Functions

distributed_fused_adam_with_cosine_annealing

Creates a distributed fused Adam optimizer with cosine annealing scheduler.

API

bridge.recipes.utils.optimizer_utils.distributed_fused_adam_with_cosine_annealing(

precision: str = 'bf16-mixed',

lr_warmup_iters: int = 2000,

lr_decay_iters: int = 2000,

adam_beta1: float = 0.9,

adam_beta2: float = 0.95,

adam_eps: float = 1e-05,

weight_decay: float = 0.1,

max_lr: float = 0.0001,

min_lr: Optional[float] = None,

clip_grad: float = 1.0,

) → tuple[megatron.bridge.training.config.OptimizerConfig, megatron.bridge.training.config.SchedulerConfig]

Creates a distributed fused Adam optimizer with cosine annealing scheduler.