Spaces:
Running on Zero
Running on Zero
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: OpenMDW-1.1 | |
| import collections | |
| import functools | |
| import itertools | |
| import math | |
| from copy import deepcopy | |
| from typing import Any, Optional | |
| import attrs | |
| import torch | |
| import torch.nn as nn | |
| from torch.distributed.checkpoint.state_dict import StateDictOptions, get_optimizer_state_dict, set_optimizer_state_dict | |
| from torch.distributed.checkpoint.stateful import Stateful | |
| from torch.optim.lr_scheduler import LambdaLR | |
| from cosmos_framework.utils.config import make_freezable | |
| from cosmos_framework.utils.lazy_config import LazyDict | |
| from cosmos_framework.utils import log | |
| from cosmos_framework.utils.reasoner.fused_adam import FusedAdam | |
| class OptimizerConfig: | |
| name: str = "FusedAdam" | |
| lr: float = 2e-6 | |
| init_lr: float = 1e-7 | |
| end_lr: float = 1e-6 | |
| fused: bool = False | |
| weight_decay: float = 0.1 | |
| betas: tuple[float, float] = (0.9, 0.95) | |
| lr_multiplier: LazyDict = LazyDict(dict(vision_encoder=0.1, mm_projector=1.0, llm=1.0)) | |
| # model freeze config | |
| freeze_vision_encoder: bool = False | |
| freeze_mm_projector: bool = False | |
| freeze_llm: bool = False | |
| freeze_llm_moe_gates: bool = False | |
| def _optimizer_cls(params: list[nn.Parameter], optimizer_kwargs: dict[str, Any], name: str): | |
| if name.lower() == "adam": | |
| optimizer = torch.optim.Adam(params, **optimizer_kwargs) | |
| elif name.lower() == "adamw": | |
| optimizer = torch.optim.AdamW(params, **optimizer_kwargs) | |
| elif name.lower() == "fusedadam": | |
| optimizer = FusedAdam( | |
| params, | |
| lr=optimizer_kwargs["lr"], | |
| weight_decay=optimizer_kwargs["weight_decay"], | |
| betas=optimizer_kwargs["betas"], | |
| capturable=True, | |
| master_weights=True, | |
| ) | |
| else: | |
| raise NotImplementedError(f"Optimizer {name} not added.") | |
| return optimizer | |
| class OptimizersContainer(Stateful): | |
| """Util for calling step/zero_grad on multiple optimizers needed for virtual pipeline stages | |
| and saving/loading optimizer state_dict at checkpoint. | |
| """ | |
| def __init__( | |
| self, | |
| model_parts: list[nn.Module], | |
| optimizer_kwargs: dict[str, Any], | |
| name: str, | |
| lr_multiplier: dict[str, float], | |
| model_part_names: list[str], | |
| ) -> None: | |
| assert len(model_parts) == len(model_part_names), "model_parts and model_part_names must have the same length" | |
| self.model_parts = model_parts | |
| self.optimizers = [[] for _ in self.model_parts] | |
| self.model_part_names = model_part_names | |
| for model_id, model in enumerate(self.model_parts): | |
| optimizer_kwargs_copy = deepcopy(optimizer_kwargs) | |
| optimizer_kwargs_copy["lr"] *= lr_multiplier[model_part_names[model_id]] | |
| log.info( | |
| f"model_id: {model_id} | model_part_names: {model_part_names[model_id]} | lr_multiplier: {lr_multiplier[model_part_names[model_id]]} | lr: {optimizer_kwargs_copy['lr']}" | |
| ) | |
| if optimizer_kwargs_copy["fused"]: | |
| # Group the parameters by device mesh to do optimizer fusion. | |
| parameters_by_mesh = collections.defaultdict(list) | |
| for p in model.parameters(): | |
| if p.requires_grad: | |
| device_mesh = p.device_mesh if hasattr(p, "device_mesh") else "default" | |
| parameters_by_mesh[device_mesh].append(p) | |
| for params in parameters_by_mesh.values(): | |
| optimizer = _optimizer_cls(params, optimizer_kwargs_copy, name) | |
| self.optimizers[model_id].append(optimizer) | |
| else: | |
| for p in model.parameters(): | |
| if p.requires_grad: | |
| optimizer = _optimizer_cls([p], optimizer_kwargs_copy, name) | |
| self.optimizers[model_id].append(optimizer) | |
| def __iter__(self) -> torch.optim.Optimizer: | |
| return iter(itertools.chain(*self.optimizers)) | |
| def step(self) -> None: | |
| for optimizer in itertools.chain(*self.optimizers): | |
| optimizer.step() | |
| def zero_grad(self, set_to_none: bool = False) -> None: | |
| for optimizer in itertools.chain(*self.optimizers): | |
| optimizer.zero_grad(set_to_none=set_to_none) | |
| def state_dict(self) -> dict[str, Any]: | |
| result = {} | |
| for idx, (model_part, optimizer) in enumerate(zip(self.model_parts, self.optimizers)): | |
| sd = get_optimizer_state_dict( | |
| model=model_part, | |
| optimizers=optimizer, | |
| options=StateDictOptions(flatten_optimizer_state_dict=True), | |
| ) | |
| prefix = f"optimizer_{idx}/" | |
| result.update({f"{prefix}{k}": v for k, v in sd.items()}) | |
| return result | |
| def load_state_dict(self, state_dict: dict[str, Any]) -> None: | |
| for idx, (model, optimizers) in enumerate(zip(self.model_parts, self.optimizers)): | |
| prefix = f"optimizer_{idx}/" | |
| optimizer_state = { | |
| k[len(prefix) :]: v # Remove prefix | |
| for k, v in state_dict.items() | |
| if k.startswith(prefix) | |
| } | |
| set_optimizer_state_dict( | |
| model=model, | |
| optimizers=optimizers, | |
| optim_state_dict=optimizer_state, | |
| options=StateDictOptions(flatten_optimizer_state_dict=True), | |
| ) | |
| # consider split between PP and non-PP | |
| def build_optimizers( | |
| model_parts: list[nn.Module], | |
| config: OptimizerConfig, | |
| model_part_names: list[str], | |
| ) -> OptimizersContainer: | |
| """Wrap one optimizer per model part in an OptimizersContainer which provides a single | |
| step() and zero_grad() method for all the child optimizers. | |
| """ | |
| lr_multiplier = config.lr_multiplier | |
| for part_name in model_part_names: | |
| assert part_name in lr_multiplier, f"lr_multiplier must have the key {part_name}" | |
| name = config.name | |
| lr = config.lr | |
| fused = config.fused | |
| optimizer_kwargs = { | |
| "lr": lr, | |
| "betas": config.betas, | |
| "weight_decay": config.weight_decay, | |
| "fused": fused, | |
| "foreach": not fused, | |
| } | |
| return OptimizersContainer(model_parts, optimizer_kwargs, name, lr_multiplier, model_part_names) | |
| class SchedulersContainer(Stateful): | |
| """Util for calling step on multiple learning rate schedulers needed for virtual pipeline stages""" | |
| def __init__(self, optimizers: OptimizersContainer, lr_lambda) -> None: | |
| self.schedulers = [] | |
| for optimizer in optimizers: | |
| self.schedulers.append(LambdaLR(optimizer, lr_lambda=lr_lambda)) | |
| def step(self) -> None: | |
| for id, scheduler in enumerate(self.schedulers): | |
| scheduler.step() | |
| def state_dict(self) -> dict[str, Any]: | |
| # Currently, we have one scheduler per optimizer. However, when using MultiSchedule PP or optimizer-in-backward, | |
| # there are multiple optimizers and schedulers, but the scheduler state_dict remains the same for all. | |
| # Therefore, we only save the first one and later load it for all. | |
| assert len(self.schedulers) > 0, "Must have at least one scheduler to save state_dict" | |
| return self.schedulers[0].state_dict() | |
| def load_state_dict(self, state_dict: dict[str, Any]) -> None: | |
| # Load the same state_dict for all schedulers. The key value we're concerned with in scheduler.state_dict() is `last_epoch`, | |
| # which is an integer that will be automatically copied. As long as `training.steps` and `training.warmup_iters` remain | |
| # unchanged when resuming from a checkpoint, this approach is safe. We call `.copy()` here to ensure extra safety. | |
| last_epoch = state_dict["last_epoch"] # Extract last known epoch | |
| _step_count = state_dict["_step_count"] | |
| log.info(f"Resuming schedulers by stepping them to last_epoch: {last_epoch}; _step_count: {_step_count}") | |
| # Manually step all schedulers to match the saved state -- this is a workaround for the inherited issue in the state dict saving (only saved the first scheduler) | |
| # But we have different learning rate for each scheduler, so we need to step them separately instead of loading the state dict | |
| # The benefit of this approach is that we can resume from a checkpoint even if the learning rate is changed | |
| for idx, scheduler in enumerate(self.schedulers): | |
| for step in range(_step_count): | |
| scheduler.step() # Step forward to match previous training state | |
| log.info(f"Scheduler {idx + 1}/{len(self.schedulers)} stepped {_step_count} times.") | |
| log.info(f"Updated learning rate: {scheduler.get_last_lr()}") | |
| def get_last_lr(self) -> list[float]: | |
| return [scheduler.get_last_lr() for scheduler in self.schedulers] | |
| def linear_warmup_linear_decay(warmup_iters: int, decay_steps: int, current_step: int) -> float: | |
| """Computes linear warmup followed by linear decay. | |
| Per LambdaLR requirement, this is accomplished by returning | |
| a multiplicative factor to adjust the learning rate to | |
| create the desired schedule. | |
| """ | |
| if current_step < warmup_iters: | |
| # linear warmup | |
| # 0-indexed step, hence + 1 adjustments | |
| current_step += 1 | |
| curr_adjustment = float(current_step / (warmup_iters + 1)) | |
| else: | |
| # linear decay | |
| normalized_step = decay_steps - (current_step - warmup_iters) | |
| curr_adjustment = 1 - (decay_steps - normalized_step) / decay_steps | |
| return curr_adjustment | |
| def linear_warmup(warmup_iters: int, current_step: int) -> float: | |
| """Computes linear warmup only | |
| Per LambdaLR requirement, this is accomplished by returning | |
| a multiplicative factor to adjust the learning rate to | |
| create the desired schedule. | |
| """ | |
| if current_step < warmup_iters: | |
| # linear warmup | |
| # 0-indexed step, hence + 1 adjustments | |
| current_step += 1 | |
| curr_adjustment = float(current_step / (warmup_iters + 1)) | |
| else: | |
| curr_adjustment = 1 | |
| return curr_adjustment | |
| def linear_warmup_cosine_cooldown( | |
| warmup_iters: int, cooldown_steps: int, current_step: int, base_lr: float, init_lr: float, end_lr: float | |
| ) -> float: | |
| """This scheduler will warmup the learning rate from init_lr to base_lr for warmup_iters, | |
| then decay the learning rate from base_lr to end_lr for cooldown_steps. After cooldown_steps + warmup_iters, | |
| the learning rate will be set to end_lr. | |
| Per LambdaLR requirement, this is accomplished by returning | |
| a multiplicative factor to adjust the learning rate to | |
| create the desired schedule. | |
| Args: | |
| warmup_iters (int): The number of steps to warmup the learning rate. | |
| cooldown_steps (int): The number of steps to decay the learning rate. | |
| current_step (int): The current step. | |
| base_lr (float): The base learning rate. | |
| init_lr (float): The initial learning rate before warmup. | |
| end_lr (float): The final learning rate after cooldown. | |
| Returns: | |
| float: The multiplicative factor to adjust the learning rate. | |
| """ | |
| total_steps = warmup_iters + cooldown_steps | |
| # Normalize | |
| init_multiplier = init_lr / base_lr | |
| end_multiplier = end_lr / base_lr | |
| if current_step <= warmup_iters: | |
| progress = float(current_step / warmup_iters) | |
| return init_multiplier + (1.0 - init_multiplier) * progress | |
| elif current_step <= total_steps: | |
| progress = (current_step - warmup_iters) / cooldown_steps | |
| return end_multiplier + 0.5 * (1.0 - end_multiplier) * (1 + math.cos(math.pi * progress)) | |
| else: | |
| return end_multiplier | |
| def build_lr_schedulers( | |
| optimizers: OptimizersContainer, | |
| name: str, | |
| lr: float, | |
| warmup_iters: int, | |
| lr_decay_iters: Optional[int] = None, | |
| init_lr: Optional[float] = None, | |
| end_lr: Optional[float] = None, | |
| ) -> SchedulersContainer: | |
| decay_steps = float(max(1, lr_decay_iters - warmup_iters)) if lr_decay_iters is not None else None | |
| if name == "warmup_cosine_lr": | |
| assert init_lr is not None and end_lr is not None, "init_lr and end_lr must be provided for warmup_cosine_lr" | |
| assert lr_decay_iters is not None, "lr_decay_iters must be provided for warmup_cosine_lr" | |
| lr_lambda = functools.partial( | |
| linear_warmup_cosine_cooldown, | |
| warmup_iters, | |
| decay_steps, | |
| base_lr=lr, | |
| init_lr=init_lr, | |
| end_lr=end_lr, | |
| ) | |
| elif name == "lambdalinear": | |
| assert lr_decay_iters is not None, "lr_decay_iters must be provided for lambdalinear" | |
| lr_lambda = functools.partial(linear_warmup_linear_decay, warmup_iters, decay_steps) | |
| else: | |
| lr_lambda = functools.partial(linear_warmup, warmup_iters) | |
| return SchedulersContainer(optimizers, lr_lambda) | |