| from __future__ import annotations |
|
|
| import inspect |
| import math |
|
|
| import torch |
|
|
|
|
| def compute_elf_lr(lr: float, blr: float, effective_batch_size: int) -> float: |
| if lr is not None and lr > 0: |
| return float(lr) |
| return float(blr) * float(effective_batch_size) / 256.0 |
|
|
|
|
| def create_learning_rate_fn( |
| num_train_steps: int, |
| num_warmup_steps: int, |
| learning_rate: float, |
| schedule: str = "constant", |
| min_lr: float = 0.0, |
| ): |
| """ELF pytorch_elf linear warmup schedule, returning absolute LR.""" |
| alpha = (min_lr / learning_rate) if learning_rate > 0 else 0.0 |
|
|
| def fn(step: int) -> float: |
| step = int(step) |
| if num_warmup_steps > 0 and step < num_warmup_steps: |
| return learning_rate * step / max(1, num_warmup_steps) |
| if schedule == "cosine": |
| progress = (step - num_warmup_steps) / max(1, num_train_steps - num_warmup_steps) |
| progress = min(max(progress, 0.0), 1.0) |
| cosine = 0.5 * (1.0 + math.cos(math.pi * progress)) |
| return learning_rate * (alpha + (1.0 - alpha) * cosine) |
| return learning_rate |
|
|
| return fn |
|
|
|
|
| def build_optimizer(args, model: torch.nn.Module, world_size: int) -> torch.optim.Optimizer: |
| effective_batch = int(args.batch_size) * int(world_size) * int(args.grad_accum) |
| args.effective_batch_size = effective_batch |
| args.lr = compute_elf_lr(float(args.lr), float(args.blr), effective_batch) |
| if args.optimizer == "muon": |
| from utils.muon_utils import muon_with_aux_adam |
|
|
| return muon_with_aux_adam(model, lr=args.lr) |
| params = [p for p in model.parameters() if p.requires_grad] |
| fused = "fused" in inspect.signature(torch.optim.AdamW).parameters and torch.cuda.is_available() |
| return torch.optim.AdamW( |
| params, |
| lr=args.lr, |
| weight_decay=args.weight_decay, |
| betas=(args.adam_beta1, args.adam_beta2), |
| eps=args.adam_eps, |
| fused=fused, |
| ) |
|
|
|
|
| def build_lr_scheduler(args, optimizer: torch.optim.Optimizer): |
| lr_fn = create_learning_rate_fn( |
| num_train_steps=args.steps, |
| num_warmup_steps=args.warmup_steps, |
| learning_rate=args.lr, |
| schedule=args.lr_schedule, |
| min_lr=args.min_lr, |
| ) |
| base_lr = optimizer.param_groups[0]["lr"] |
| return torch.optim.lr_scheduler.LambdaLR( |
| optimizer, |
| lr_lambda=lambda step: lr_fn(step) / max(base_lr, 1e-12), |
| ) |
|
|
