# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: OpenMDW-1.1 import collections import copy from typing import Any, Iterator, NamedTuple import torch import torch.nn as nn from omegaconf import DictConfig, ListConfig, OmegaConf 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, LRScheduler from cosmos_framework.utils.functional.lr_scheduler import LambdaLinearScheduler, LambdaWarmUpCosineScheduler, WSDScheduler from cosmos_framework.utils import log # Hybrid orthogonalizing optimizers (Muon / Dion2) that own their parameter # categorization and run their own collective communication. They must be built # as a single optimizer instance over all selected params (no per-device-mesh # split) and need ``categorize_params`` called after construction. _AUX_ADAMW_OPTIMIZERS = ("muonwithauxadamw", "dion2withauxadamw") class ParamMetadata(NamedTuple): lr: float enable_weight_decay: bool def _convert_omegaconf_to_python(obj: Any) -> Any: """Convert OmegaConf types to plain Python types. This is needed because PyTorch's checkpoint utilities don't handle OmegaConf types like ListConfig and DictConfig. """ if isinstance(obj, (ListConfig, DictConfig)): return OmegaConf.to_container(obj, resolve=True) return obj def _optimizer_cls( params: list[nn.Parameter] | list[dict[str, Any]], optimizer_type: str, **optimizer_kwargs: Any, ) -> torch.optim.Optimizer: """Dispatch to the requested optimizer class with the right kwargs. Supported ``optimizer_type`` values (case-insensitive): - ``"adam"`` / ``"adamw"``: forwarded to ``torch.optim.Adam`` / ``torch.optim.AdamW``. ``fused`` (if present in ``optimizer_kwargs``) flows through and selects the fused CUDA kernel. - ``"fusedadam"``: NVIDIA's :class:`cosmos_framework.utils.generator.fused_adam.FusedAdam`. It is fused by construction and rejects a ``fused`` kwarg, so any ``fused`` entry is popped before instantiation. We also force ``capturable=True`` and ``master_weights=True`` because those are the only settings exercised in our distributed training stack. Raises ``NotImplementedError`` for any other ``optimizer_type``. """ if optimizer_type.lower() == "adam": optimizer = torch.optim.Adam(params, **optimizer_kwargs) elif optimizer_type.lower() == "adamw": optimizer = torch.optim.AdamW(params, **optimizer_kwargs) elif optimizer_type.lower() == "fusedadam": from cosmos_framework.utils.generator.fused_adam import FusedAdam # FusedAdam is fused by construction and does not accept a ``fused`` kwarg. optimizer_kwargs.pop("fused", None) # Force ``capturable`` / ``master_weights`` on -- the only configuration # exercised in our distributed-training stack. Overwrite in-place # rather than passing as positional keywords, otherwise a caller that # also sets either flag would trigger a duplicate-kwarg ``TypeError``. optimizer_kwargs["capturable"] = True optimizer_kwargs["master_weights"] = True optimizer = FusedAdam(params, **optimizer_kwargs) elif optimizer_type.lower() == "muonwithauxadamw": from cosmos_framework.utils.generator.muon_with_aux_adamw import MuonWithAuxAdamW # Muon's AdamW side is the TE-fused kernel; it is fused by construction and # absorbs ``fused`` via **kwargs, but we pop it here to be explicit. We force # capturable + master_weights to match FusedAdam's mixed-precision setup. optimizer_kwargs.pop("fused", None) optimizer_kwargs["capturable"] = True optimizer_kwargs["master_weights"] = True optimizer = MuonWithAuxAdamW(params, **optimizer_kwargs) elif optimizer_type.lower() == "dion2withauxadamw": from cosmos_framework.utils.generator.dion2_with_aux_adamw import Dion2WithAuxAdamW optimizer_kwargs.pop("fused", None) optimizer_kwargs["capturable"] = True optimizer_kwargs["master_weights"] = True optimizer = Dion2WithAuxAdamW(params, **optimizer_kwargs) else: raise NotImplementedError(f"Optimizer {optimizer_type} not found.") return optimizer def _build_params_with_metadata( model: nn.Module, keys_to_select: list[str], lr_multipliers: dict[str, float], base_lr: float, disable_weight_decay_for_1d_params: bool, ) -> list[tuple[nn.Parameter, ParamMetadata]]: """Filter trainable parameters and tag each with its effective LR and weight-decay flag. Walks ``model.named_parameters()`` once and produces one ``(param, ParamMetadata)`` entry per parameter that survives the ``keys_to_select`` filter and belong to regular model. Has the side effect of freezing (``requires_grad=False``) every parameter that does NOT match ``keys_to_select`` when that list is non-empty. Selection rule: - If ``keys_to_select`` is empty, every parameter is kept. - Otherwise, a parameter is kept iff its dotted name contains at least one entry of ``keys_to_select`` as a substring. Effective LR: ``ParamMetadata.lr = base_lr * matched_multiplier`` where ``matched_multiplier`` is the value associated with the first ``lr_multipliers`` pattern whose key occurs as a substring in the parameter name (iteration order of the dict is significant); defaults to ``1.0`` if no pattern matches. Weight-decay flag: ``ParamMetadata.enable_weight_decay`` is ``False`` exactly when ``disable_weight_decay_for_1d_params`` is True AND the parameter has fewer than two dimensions (the standard heuristic for norm weights, biases, and other 1-D tensors). Otherwise it is ``True``. The flag is consumed downstream by :func:`_build_optimizer_internal`, which materializes the ``weight_decay=0.0`` override on the corresponding param group. Args: model: Module whose ``named_parameters()`` to scan. Only parameters under the ``net.`` subtree are considered; EMA-network parameters (``net_ema.*``) are intentionally skipped. keys_to_select: Substrings used to allowlist parameter names. Empty list = train everything. lr_multipliers: Ordered mapping from name-substring to LR multiplier. First match wins per parameter. base_lr: Base learning rate; each parameter's effective LR is ``base_lr * matched_multiplier``. disable_weight_decay_for_1d_params: When ``True``, 1-D parameters are tagged ``enable_weight_decay=False``. Returns: List of ``(nn.Parameter, ParamMetadata)`` pairs covering every kept parameter. Frozen parameters are not included. """ # Optimize only the regular ``net.`` subtree; the sibling ``net_ema`` # subtree is intentionally skipped. Materialize once so we can ``len()`` # the total below -- ``named_parameters()`` is a generator. net_params = dict(model.net.named_parameters()) param_dict = {pn: p for pn, p in net_params.items() if p.requires_grad} params_with_metadata: list[tuple[nn.Parameter, ParamMetadata]] = [] for pn, p in param_dict.items(): if len(keys_to_select) > 0 and not any(key in pn for key in keys_to_select): p.requires_grad = False continue # Find the matching multiplier for the parameter. matched_mult = 1.0 for pattern, mult in lr_multipliers.items(): if pattern in pn: matched_mult = mult break if disable_weight_decay_for_1d_params and p.dim() < 2: enable_weight_decay = False else: enable_weight_decay = True params_with_metadata.append( ( p, ParamMetadata( lr=base_lr * matched_mult, enable_weight_decay=enable_weight_decay, ), ) ) log.info( f"Total tensors: {len(net_params)}, " f"trainable tensors: {len(param_dict)}, " f"selected tensors: {len(params_with_metadata)}" ) return params_with_metadata def _build_optimizer_internal( params_with_metadata: list[tuple[nn.Parameter, ParamMetadata]], optimizer_type: str, **optimizer_kwargs: Any, ) -> torch.optim.Optimizer: """Bucket per-parameter metadata into PyTorch param groups and instantiate the optimizer. Parameters sharing the same ``ParamMetadata`` (i.e. identical effective LR *and* identical weight-decay enablement) are collapsed into a single ``torch.optim`` param group. Each resulting group carries an explicit ``"lr"``; groups whose metadata has ``enable_weight_decay=False`` also receive an explicit ``"weight_decay": 0.0`` override that suppresses the optimizer-wide ``weight_decay`` kwarg for those parameters only. Args: params_with_metadata: Output of :func:`_build_params_with_metadata` — one ``(param, ParamMetadata)`` entry per trainable parameter. optimizer_type: Optimizer kind string accepted by :func:`_optimizer_cls` (``"adam"``, ``"adamw"``, or ``"fusedadam"``). **optimizer_kwargs: Forwarded to the optimizer constructor. ``"lr"`` is ignored at the group level (each group sets its own LR) but must still be present because PyTorch optimizers require a default. Returns: A ``torch.optim.Optimizer`` configured with one param group per distinct ``ParamMetadata`` value found in ``params_with_metadata``. """ params_by_metadata: dict[ParamMetadata, list[nn.Parameter]] = collections.defaultdict(list) for param, metadata in params_with_metadata: params_by_metadata[metadata].append(param) param_groups: list[dict[str, Any]] = [] for metadata, params in sorted(params_by_metadata.items()): log.info( f"Param group (lr={metadata.lr}, WD={metadata.enable_weight_decay}): " f"{len(params):,} tensors, {sum(p.numel() for p in params):,} elements" ) param_group: dict[str, Any] = {"params": params, "lr": metadata.lr} if not metadata.enable_weight_decay: param_group["weight_decay"] = 0.0 param_groups.append(param_group) return _optimizer_cls(param_groups, optimizer_type, **optimizer_kwargs) class OptimizersContainer(Stateful): """ Utility for calling step/zero_grad on multiple optimizers, and saving/loading optimizer state_dict at checkpoint. """ def __init__( self, model: nn.Module, optimizer_type: str, **optimizer_kwargs: Any, ) -> None: """Build per-device-mesh optimizers and expose them as a single Stateful container. Parameters are first filtered + tagged with metadata via :func:`_build_params_with_metadata` and then sub-grouped by ``param.device_mesh.mesh_dim_names`` (or ``"default"`` when no mesh is attached). One ``torch.optim`` instance is created per resulting mesh bucket so fused kernels see only same-mesh DTensors — required for FSDP2-aware optimizer state-dict resharding. Args: model: Model to build the optimizer(s) for. ``model.parameters()`` provides the parameter universe and is also passed through to ``get_optimizer_state_dict`` / ``set_optimizer_state_dict`` for checkpointing. optimizer_type: Optimizer kind string accepted by :func:`_optimizer_cls` (``"adam"``, ``"adamw"``, ``"fusedadam"``). **optimizer_kwargs: Forwarded to the optimizer constructor, with the following container-level kwargs intercepted first: - ``fused`` (required): Must be truthy. Forwarded to ``torch.optim.AdamW`` / ``Adam`` so they use the fused CUDA kernel; ignored for ``"fusedadam"`` which is fused by construction. Eager optimizers are not supported. - ``lr`` (required): Base learning rate. Each param group's effective LR is ``lr * matched_multiplier``; ``lr`` is still required because PyTorch optimizers expect a default. - ``keys_to_select`` (optional, default ``[]``): List of substrings; only params whose names contain at least one are trained. Empty list = train every parameter. - ``lr_multipliers`` (optional, default ``{}``): Ordered dict mapping param-name substrings to LR multipliers. First match wins; unmatched selected params use multiplier ``1.0``. - ``disable_weight_decay_for_1d_params`` (optional, default ``False``): When true, parameters with ``dim() < 2`` (norm weights, biases, etc.) get ``weight_decay=0.0`` in their param group regardless of the optimizer-wide ``weight_decay``. """ self.model = model self.optimizers: list[torch.optim.Optimizer] = [] # Pop only the container-level kwargs. ``fused`` and ``lr`` are # intentionally NOT popped: ``fused`` must flow through to # ``torch.optim.AdamW`` / ``Adam`` so the fused CUDA kernel is actually # selected, and ``lr`` is needed both here (as the base for per-group # multipliers) and by the underlying optimizer (as the group-level # default). keys_to_select = optimizer_kwargs.pop("keys_to_select", []) lr_multipliers: dict[str, float] = optimizer_kwargs.pop("lr_multipliers", {}) disable_weight_decay_for_1d_params = optimizer_kwargs.pop("disable_weight_decay_for_1d_params", False) if not optimizer_kwargs.get("fused", False): raise ValueError("Optimizers with fused=False are not supported; pass fused=True in optimizer_kwargs.") if "lr" not in optimizer_kwargs: raise ValueError("`lr` is required in optimizer_kwargs (used as the base for per-group LR multipliers).") base_lr = optimizer_kwargs["lr"] params_with_metadata = _build_params_with_metadata( model, keys_to_select=keys_to_select, lr_multipliers=lr_multipliers, base_lr=base_lr, disable_weight_decay_for_1d_params=disable_weight_decay_for_1d_params, ) if optimizer_type.lower() in _AUX_ADAMW_OPTIMIZERS: # Muon / Dion2 categorize parameters into an orthogonalized-matrix group # and an auxiliary AdamW group internally, and run their own collective # communication across the FULL device mesh. They must therefore be a # single optimizer instance over all selected params (no per-mesh split). # # The factory's param grouping is preserved: ``_build_optimizer_internal`` # still produces the same (lr, weight_decay) param groups (from # ``lr_multipliers`` / ``disable_weight_decay_for_1d_params``), and # Muon/Dion2 read lr/weight_decay per group — degenerating to a single # global lr/wd when there is only one group (the reference behavior). optimizer = _build_optimizer_internal( params_with_metadata, optimizer_type, **optimizer_kwargs, ) # Categorize against the trainable network only, mirroring # ``_build_params_with_metadata`` which optimizes the ``net`` subtree # and skips ``net_ema``. optimizer.categorize_params(model.net) self.optimizers.append(optimizer) else: # Sub-group by device mesh so fused optimizers operate on same-mesh params. mesh_groups: dict[str, list[tuple[nn.Parameter, ParamMetadata]]] = collections.defaultdict(list) for param, metadata in params_with_metadata: if hasattr(param, "device_mesh"): # ``mesh_dim_names`` is ``tuple[str, ...] | None`` on DeviceMesh — # fall back to ``default`` when names weren't assigned. names = param.device_mesh.mesh_dim_names mesh_key = "-".join(names) if names else "default" else: mesh_key = "default" mesh_groups[mesh_key].append((param, metadata)) # Create one optimizer per mesh, each with per-LR,weight-decay param groups. for mesh_key, mesh_params in mesh_groups.items(): log.info(f"Building optimizer for mesh '{mesh_key}'") optimizer = _build_optimizer_internal( mesh_params, optimizer_type, **optimizer_kwargs, ) self.optimizers.append(optimizer) log.info(f"Created {len(self.optimizers)} optimizers") def __iter__(self) -> Iterator[torch.optim.Optimizer]: return iter(self.optimizers) def __len__(self) -> int: return len(self.optimizers) def step(self) -> None: # NOTE: This container does not forward a ``torch.amp.GradScaler`` to the # inner optimizers, and training here is bf16 with the scaler disabled # (``grad_scaler_args={"enabled": False}``), so step() always runs unscaled. # MuonWithAuxAdamW / Dion2WithAuxAdamW are intentionally bf16-only (their # AMP grad-scaler support was removed to keep the Newton-Schulz path simple); # only FusedAdam still carries a scaler hook. Enabling fp16 + dynamic loss # scaling end-to-end would require re-adding scaler passthrough here and in # the Muon/Dion2 step paths, which is deliberately not supported. for optimizer in self.optimizers: optimizer.step() def zero_grad(self, set_to_none: bool = True) -> None: # Default matches ``torch.optim.Optimizer.zero_grad`` (set_to_none=True # since PyTorch 1.7); zero-in-place is strictly slower and uses extra # memory, so callers must opt-in explicitly. for optimizer in self.optimizers: optimizer.zero_grad(set_to_none=set_to_none) def state_dict(self) -> dict[str, Any]: return get_optimizer_state_dict( model=self.model, optimizers=self.optimizers, options=StateDictOptions(flatten_optimizer_state_dict=True), ) def load_state_dict(self, state_dict: dict[str, Any]) -> None: set_optimizer_state_dict( model=self.model, optimizers=self.optimizers, optim_state_dict=state_dict, options=StateDictOptions(flatten_optimizer_state_dict=True), ) def build_optimizer( model: nn.Module, optimizer_type: str, **optimizer_kwargs: Any, ) -> "OptimizersContainer": """Build an :class:`OptimizersContainer` for a model. Returns the container — NOT a raw ``torch.optim.Optimizer`` — because the underlying implementation creates one optimizer per device mesh. Callers that need PyTorch ``param_groups`` must iterate ``container.optimizers``. Before instantiation, every kwarg is normalized via :func:`_convert_omegaconf_to_python` so that OmegaConf containers (``ListConfig`` / ``DictConfig``) become plain Python types — PyTorch's checkpoint utilities ``_copy_state_dict`` does not handle OmegaConf types. Args: model: The model to build an optimizer for. Trainable-parameter selection happens inside the container; parameters not matched by ``keys_to_select`` (when non-empty) are frozen as a side effect. optimizer_type: Optimizer kind string accepted by :func:`_optimizer_cls` (``"adam"``, ``"adamw"``, or ``"fusedadam"``). **optimizer_kwargs: Forwarded to :class:`OptimizersContainer`, which intercepts a few container-level kwargs and forwards the rest to the underlying ``torch.optim`` constructor. Container-level kwargs: - ``fused`` (required, truthy): Use the fused CUDA kernel. Eager optimizers are intentionally not supported. - ``lr`` (required): Base learning rate. Each param group's effective LR is ``lr * matched_multiplier``. - ``keys_to_select`` (optional, default ``[]``): Substrings used to allowlist parameter names; empty list = train everything. - ``lr_multipliers`` (optional, default ``{}``): Ordered dict mapping parameter-name substrings to LR multipliers. E.g. ``{"sound2llm": 5.0, "llm2sound": 5.0}`` gives those params 5x the base LR. First match wins; unmatched selected params use multiplier ``1.0``. - ``disable_weight_decay_for_1d_params`` (optional, default ``False``): If true, one-dimensional parameters such as norm weights and biases get ``weight_decay=0.0`` in their param group regardless of the optimizer-wide ``weight_decay``. Returns: An :class:`OptimizersContainer` wrapping one ``torch.optim.Optimizer`` per device mesh discovered on the model's parameters. """ # Convert OmegaConf types to plain Python types to avoid issues with checkpoint saving. # PyTorch's _copy_state_dict doesn't handle OmegaConf types like ListConfig. optimizer_kwargs = {k: _convert_omegaconf_to_python(v) for k, v in optimizer_kwargs.items()} return OptimizersContainer(model, optimizer_type, **optimizer_kwargs) def _lr_scheduler_cls( lr_scheduler_type: str, **lr_scheduler_kwargs: Any, ) -> LambdaLinearScheduler | LambdaWarmUpCosineScheduler | WSDScheduler: """Instantiate a lambda-style scheduler whose ``.schedule(step)`` returns an LR multiplier. Both returned classes expose a ``schedule(step) -> float`` callable that :class:`LRSchedulersContainer` wraps with ``torch.optim.lr_scheduler.LambdaLR`` to drive each optimizer's param-group LRs. ``lr_scheduler_type`` matching is case-insensitive; valid values are ``"lambdalinear"`` (linear decay), ``"lambdacosine"`` (warmup + cosine decay), and ``"wsd"`` (warmup-stable-decay). Any other value raises ``NotImplementedError``. All remaining ``**lr_scheduler_kwargs`` are forwarded verbatim to the underlying scheduler constructor (e.g. ``warm_up_steps``, ``cycle_lengths``, ``total_steps``, ``decay_steps``, ``f_start``, ``f_max``, ``f_min``, ``verbosity_interval``). """ if lr_scheduler_type.lower() == "lambdalinear": lr_scheduler = LambdaLinearScheduler(**lr_scheduler_kwargs) elif lr_scheduler_type.lower() == "lambdacosine": lr_scheduler = LambdaWarmUpCosineScheduler(**lr_scheduler_kwargs) elif lr_scheduler_type.lower() == "wsd": lr_scheduler = WSDScheduler(**lr_scheduler_kwargs) else: raise NotImplementedError(f"LR Scheduler {lr_scheduler_type} not found.") return lr_scheduler class LRSchedulersContainer(Stateful): """Wraps one ``torch.optim.lr_scheduler.LambdaLR`` per inner optimizer. Mirrors the structure of :class:`OptimizersContainer`: one scheduler is created per element of an ``OptimizersContainer`` (i.e. one per device mesh), and ``step()`` / ``state_dict()`` / ``load_state_dict()`` fan out over them. All schedulers share the same underlying lambda (``lr_scheduler.schedule``), so they advance in lock-step but multiply each of their optimizer's per-group ``base_lr`` values independently. Note: Subclasses must keep the ``step``, ``state_dict``, and ``load_state_dict`` signatures compatible with ``torch.optim.lr_scheduler.LRScheduler`` so that downstream training and checkpoint code can treat the container as a drop-in scheduler. Checkpoint format: ``LambdaLR.state_dict`` contains both scalar fields (``last_epoch``, ``_step_count``, ``verbose``, ``_get_lr_called_within_step`` — all identical across schedulers because they step in lock-step) and per-param-group lists (``base_lrs``, ``_last_lr``, ``lr_lambdas`` — whose length and values depend on the underlying optimizer's ``param_groups``). When ``len(self.schedulers) == 1`` (the common single-mesh case), :meth:`state_dict` returns scheduler-0's dict verbatim. This matches the historical on-disk shape, so single-scheduler checkpoints written by earlier versions of this class continue to round-trip. When ``len(self.schedulers) > 1`` (heterogeneous mesh case), the per-scheduler list fields differ in length and content across schedulers. :meth:`state_dict` augments scheduler-0's dict with a ``"schedulers"`` entry — a ``list[dict]`` of per-scheduler ``LambdaLR.state_dict`` values — so each scheduler's ``base_lrs`` / ``_last_lr`` / ``lr_lambdas`` are persisted and restored without corruption. Args: optimizers: The optimizer container to drive. Each inner optimizer gets its own ``LambdaLR`` instance sharing the same schedule lambda. """ schedulers: list[LRScheduler] def __init__( self, optimizers: OptimizersContainer, lr_scheduler_type: str, **lr_scheduler_kwargs: Any, ) -> None: if len(optimizers) == 0: raise ValueError("Must have at least one optimizer to create LRScheduler") lr_scheduler = _lr_scheduler_cls(lr_scheduler_type, **lr_scheduler_kwargs) self.schedulers = [LambdaLR(optimizer, lr_scheduler.schedule) for optimizer in optimizers] log.info(f"Created {len(self.schedulers)} schedulers") def __iter__(self) -> Iterator[LRScheduler]: return iter(self.schedulers) def __len__(self) -> int: return len(self.schedulers) def step(self) -> None: for scheduler in self.schedulers: scheduler.step() def state_dict(self) -> dict[str, Any]: # Single-scheduler: return the legacy shape verbatim so existing # on-disk checkpoints continue to round-trip with strict_resume=True. # Multi-scheduler: also persist per-scheduler state so heterogeneous # param-group structures (different base_lrs / _last_lr lengths) are # captured without corrupting one scheduler with another's data. state_dict = self.schedulers[0].state_dict() if len(self.schedulers) > 1: state_dict["schedulers"] = [scheduler.state_dict() for scheduler in self.schedulers] return state_dict def load_state_dict(self, state_dict: dict[str, Any]) -> None: # Inverse of ``state_dict``: presence of ``"schedulers"`` implies the # multi-scheduler shape and *exactly* matches ``len(self.schedulers)``; # absence implies the single-scheduler shape. ``state_dict`` is always # produced by ``self.state_dict()`` on either the same or a compatible # container, so a mismatch is a real config drift, not a recoverable # path — raise with an actionable message. per_scheduler_states = state_dict.get("schedulers") if per_scheduler_states is None: if len(self.schedulers) != 1: raise ValueError( f"LRSchedulersContainer.load_state_dict: checkpoint holds a single " f"scheduler's state, but this container has {len(self.schedulers)}. " "Resume from a checkpoint saved with a matching number of schedulers, " "or add the scheduler to `checkpoint.keys_not_to_resume` to start " "scheduler state fresh." ) self.schedulers[0].load_state_dict(copy.deepcopy(state_dict)) return if len(per_scheduler_states) != len(self.schedulers): raise ValueError( f"LRSchedulersContainer.load_state_dict: checkpoint has state for " f"{len(per_scheduler_states)} schedulers, but this container has " f"{len(self.schedulers)}. Resume from a checkpoint with a matching number " "of schedulers, or add the scheduler to `checkpoint.keys_not_to_resume` to " "start scheduler state fresh." ) for scheduler, sub_state in zip(self.schedulers, per_scheduler_states): # Deepcopy so nested mutable values (lists) are not aliased across # schedulers after ``LambdaLR.__dict__.update``. scheduler.load_state_dict(copy.deepcopy(sub_state)) def get_last_lr(self) -> list[float | torch.Tensor]: # Concatenate the per-optimizer ``get_last_lr()`` lists. Each # underlying scheduler reports one entry per param group, and # different optimizers (one per device mesh) may carry distinct # ``base_lr`` values via ``lr_multipliers``, so the full snapshot # requires aggregating across schedulers. return [lr for scheduler in self.schedulers for lr in scheduler.get_last_lr()] def get_lr(self) -> list[float | torch.Tensor]: # Same aggregation as ``get_last_lr``; see comment there. return [lr for scheduler in self.schedulers for lr in scheduler.get_lr()] def build_lr_scheduler( optimizer: OptimizersContainer, lr_scheduler_type: str, **lr_scheduler_kwargs: Any, ) -> LRSchedulersContainer: """Create an :class:`LRSchedulersContainer` for the given optimizer container. Thin factory used by Hydra/LazyConfig scheduler registrations (see ``cosmos_framework/configs/base/defaults/optimizer.py``). At instantiation time, ``optimizer`` is supplied by the model's ``init_optimizer_scheduler`` hook via ``lazy_instantiate(scheduler_config, optimizer=optimizer)``. Note: Callers who need different scheduler behavior can subclass :class:`LRSchedulersContainer` and write their own factory; pointing the Hydra ``_target_`` at that factory swaps it in. Args: optimizer: The optimizer container the scheduler should drive. One :class:`torch.optim.lr_scheduler.LambdaLR` will be built per element of this container. lr_scheduler_type: Scheduler kind accepted by :func:`_lr_scheduler_cls` — ``"lambdalinear"``, ``"lambdacosine"``, or ``"wsd"`` (case-insensitive). **lr_scheduler_kwargs: Forwarded verbatim to the underlying lambda scheduler constructor (e.g. ``warm_up_steps``, ``cycle_lengths``, ``f_start``, ``f_max``, ``f_min``, ``verbosity_interval``). Returns: An :class:`LRSchedulersContainer` wrapping one ``LambdaLR`` per inner optimizer. """ return LRSchedulersContainer(optimizer, lr_scheduler_type, **lr_scheduler_kwargs)