mirror of
https://github.com/hanzoai/gym-v1.git
synced 2026-10-30 18:35:21 +00:00
146 lines
4.2 KiB
Python
146 lines
4.2 KiB
Python
from torch.optim.lr_scheduler import LambdaLR
|
|
|
|
import math
|
|
import torch
|
|
import torch.nn.utils as nn_utils
|
|
|
|
from typing import Dict, Any
|
|
|
|
from .communicate import all_reduce
|
|
|
|
from exogym.aux.utils import LogModule
|
|
|
|
from abc import ABC, abstractmethod
|
|
|
|
from .optim import OptimSpec, ensure_optim_spec
|
|
|
|
|
|
class Strategy(ABC, LogModule):
|
|
def __init__(
|
|
self,
|
|
lr_scheduler: str = None,
|
|
lr_scheduler_kwargs: Dict[str, Any] = None,
|
|
**kwargs: Dict[str, Any],
|
|
):
|
|
|
|
self.lr_scheduler = lr_scheduler
|
|
self.lr_scheduler_kwargs = lr_scheduler_kwargs
|
|
|
|
self.kwargs = kwargs
|
|
|
|
for k, v in kwargs.items():
|
|
setattr(self, k, v)
|
|
|
|
# Initialize scheduler as None; will be set after self.optim is defined in subclasses.
|
|
self.scheduler = None
|
|
|
|
# List of callbacks to record learning rate changes.
|
|
self.lr_callbacks = []
|
|
|
|
self.max_steps = 1 # Needs to be initialized for first call of lr_lambda.
|
|
|
|
def _init_node(self, model, rank, num_nodes):
|
|
self.model = model
|
|
self.rank = rank
|
|
self.num_nodes = num_nodes
|
|
|
|
self.local_step = 0
|
|
|
|
if hasattr(self, 'optim_spec'):
|
|
self.optim = self.optim_spec.build(model)
|
|
|
|
@abstractmethod
|
|
def step(self):
|
|
self.nbytes = 0
|
|
|
|
if self.scheduler is not None:
|
|
self.scheduler.step()
|
|
|
|
if self.rank == 0:
|
|
for callback in self.lr_callbacks:
|
|
callback(self.scheduler.get_last_lr()[0])
|
|
|
|
self.local_step += 1
|
|
|
|
def zero_grad(self):
|
|
self.optim.zero_grad()
|
|
|
|
def _setup_scheduler(self):
|
|
def lr_lambda(current_step):
|
|
warmup_steps = self.lr_scheduler_kwargs.get("warmup_steps", 1)
|
|
# If max steps not set,
|
|
if "max_steps" in self.lr_scheduler_kwargs:
|
|
max_steps = min(self.lr_scheduler_kwargs["max_steps"], self.max_steps)
|
|
else:
|
|
max_steps = self.max_steps
|
|
cosine_anneal = self.lr_scheduler_kwargs.get("cosine_anneal", False)
|
|
|
|
if current_step < warmup_steps:
|
|
return float(current_step) / float(max(warmup_steps, 1))
|
|
elif cosine_anneal:
|
|
min_lr_factor = 0.1
|
|
progress = (current_step - warmup_steps) / float(
|
|
max(1, max_steps - warmup_steps)
|
|
)
|
|
cosine_term = 0.5 * (1.0 + math.cos(math.pi * progress))
|
|
return (1 - min_lr_factor) * cosine_term + min_lr_factor
|
|
else:
|
|
return 1.0
|
|
|
|
if self.lr_scheduler == "lambda_cosine":
|
|
self.scheduler = LambdaLR(self.optim, lr_lambda)
|
|
elif self.lr_scheduler is not None:
|
|
lr_sched_kwargs = (
|
|
self.lr_scheduler_kwargs if self.lr_scheduler_kwargs is not None else {}
|
|
)
|
|
self.scheduler = self.lr_scheduler(self.optim, **lr_sched_kwargs)
|
|
else:
|
|
self.scheduler = None
|
|
|
|
def __config__(self):
|
|
remove_keys = [
|
|
"iteration",
|
|
"local_step",
|
|
"lr_callbacks",
|
|
"model",
|
|
"optim",
|
|
"scheduler",
|
|
]
|
|
|
|
config = super().__config__(remove_keys)
|
|
|
|
config["strategy"] = self.__class__.__name__
|
|
|
|
return config
|
|
|
|
|
|
class SimpleReduceStrategy(Strategy):
|
|
def __init__(self, optim_spec=None, max_norm=None, **kwargs):
|
|
super().__init__(**kwargs)
|
|
|
|
self.optim_spec = ensure_optim_spec(optim_spec) or OptimSpec(torch.optim.AdamW)
|
|
|
|
self.max_norm = max_norm
|
|
|
|
def _init_node(self, model, rank, num_nodes):
|
|
super()._init_node(model, rank, num_nodes)
|
|
|
|
self.optim = self.optim_spec.build(model)
|
|
self._setup_scheduler()
|
|
|
|
def step(self):
|
|
if self.num_nodes > 1 or True:
|
|
for param in self.model.parameters():
|
|
if param.grad is not None:
|
|
all_reduce(param.grad)
|
|
param.grad.div_(self.num_nodes)
|
|
|
|
if self.max_norm:
|
|
nn_utils.clip_grad_norm_(
|
|
self.model.parameters(), max_norm=self.max_norm
|
|
)
|
|
|
|
self.optim.step()
|
|
|
|
super().step()
|