OLMo-Bitnet-1B / optim.py

update inference code

2e055bf 2 months ago

No virus

34.2 kB

	import logging
	from abc import ABCMeta, abstractmethod
	from dataclasses import dataclass, replace
	from math import cos, pi, sqrt
	from typing import Any, Dict, List, Optional, Tuple

	import torch
	import torch.distributed as dist
	import torch.nn as nn
	from torch.distributed.fsdp import FullyShardedDataParallel
	from torch.optim.optimizer import Optimizer as OptimizerBase

	from .model import LayerNormBase, BitLinear158
	from .config import OptimizerType, SchedulerConfig, SchedulerType, TrainConfig
	from .torch_util import get_default_device, is_distributed

	__all__ = [
	"Optimizer",
	"LionW",
	"AdamW",
	"Scheduler",
	"CosWithWarmup",
	"LinearWithWarmup",
	"InvSqrtWithWarmup",
	"MaxScheduler",
	"ConstantScheduler",
	"BoltOnWarmupScheduler",
	"build_optimizer",
	"build_scheduler",
	]


	log = logging.getLogger(__name__)


	class Optimizer(OptimizerBase):
	def _clean_param_name(self, name: str) -> str:
	return name.replace("_fsdp_wrapped_module.", "")

	@torch.no_grad()
	def clip_grads_and_collect_metrics(
	self, global_step: int, collect_param_metrics: bool = True
	) -> Dict[str, torch.Tensor]:
	"""
	Clips gradients for every group that has the field `max_grad_norm`.
	At the same time collect metrics for each parameter and its gradient.
	"""
	device = get_default_device()

	# NOTE (epwalsh): during distributed training we're making an assumption that the order of
	# the param groups and the params within each group are the same across all ranks.
	# This is justified since we initialize the parameter groups in every rank by iterating over
	# `module.parameters()` or `module.named_modules()` / `module.named_parameters()`, each of which
	# provides a consistent order.
	# For each parameter (with a gradient) we'll collect:
	# - min, max, avg, norm of the param itself
	# - min, max, avg, norm of the param's gradient
	# - min, max, avg, norm of any additional per-parameter optimizer state metrics returned from
	# `self.get_state_for_param()`.
	# Afterwards we'll reduce these all over all ranks.
	per_param_min_metrics: List[torch.Tensor] = []
	per_param_max_metrics: List[torch.Tensor] = []
	per_param_sum_metrics: List[torch.Tensor] = []
	per_param_norm_metrics: List[torch.Tensor] = []
	per_param_numel_metrics: List[torch.Tensor] = []

	per_param_min_metric_names: List[str] = []
	per_param_max_metric_names: List[str] = []
	per_param_avg_metric_names: List[str] = []
	per_param_norm_metric_names: List[str] = []

	# Collect metrics locally.
	for group in self.param_groups:
	if is_distributed():
	# TODO (epwalsh): handle non-sharded params. We don't have any right now but we would
	# with ReLoRa, for example.
	assert group.get("sharded", True) is True

	for name, p in zip(group["param_names"], group["params"]):
	name = self._clean_param_name(name)
	# Always need to collect the norm of gradients for clipping, even if we're not collecting
	# other metrics.
	tensors: List[Optional[torch.Tensor]] = [p.grad]
	prefixes: List[str] = [f"grad/{name}"]
	if collect_param_metrics:
	state = self.get_state_for_param(p)
	sorted_state_keys = sorted([k for k in state.keys()])
	tensors.extend([p] + [state[key] for key in sorted_state_keys])
	prefixes.extend([f"param/{name}"] + [f"{key}/{name}" for key in sorted_state_keys])
	assert len(tensors) == len(prefixes)

	# Get min, max, avg, and norm for all `tensors` associated with the parameter.
	for x, prefix in zip(tensors, prefixes):
	# grad or state tensors could be none for params that have their shards completely on
	# other ranks.
	if x is not None and x.numel() > 0:
	if collect_param_metrics:
	x_abs = x.abs()
	per_param_min_metrics.append(x_abs.min().unsqueeze(0).to(dtype=torch.float32))
	per_param_max_metrics.append(x_abs.max().unsqueeze(0).to(dtype=torch.float32))
	per_param_sum_metrics.append(x.sum().unsqueeze(0).to(dtype=torch.float32))
	per_param_numel_metrics.append(
	torch.tensor([x.numel()], device=device, dtype=torch.float32)
	)
	per_param_norm_metrics.append(
	torch.linalg.vector_norm(x, 2.0, dtype=torch.float32).unsqueeze(0)
	)
	else:
	if collect_param_metrics:
	per_param_min_metrics.append(
	torch.tensor([float("inf")], device=device, dtype=torch.float32)
	)
	per_param_max_metrics.append(torch.tensor([0.0], device=device, dtype=torch.float32))
	per_param_sum_metrics.append(torch.tensor([0.0], device=device, dtype=torch.float32))
	per_param_numel_metrics.append(torch.tensor([0.0], device=device, dtype=torch.float32))
	per_param_norm_metrics.append(torch.tensor([0.0], device=device, dtype=torch.float32))
	if collect_param_metrics:
	per_param_min_metric_names.append(f"{prefix}.min")
	per_param_max_metric_names.append(f"{prefix}.max")
	per_param_avg_metric_names.append(f"{prefix}.avg")
	per_param_norm_metric_names.append(f"{prefix}.norm")

	assert (
	len(per_param_min_metrics)
	== len(per_param_min_metric_names)
	== len(per_param_max_metrics)
	== len(per_param_max_metric_names)
	== len(per_param_sum_metrics)
	== len(per_param_numel_metrics)
	== len(per_param_avg_metric_names)
	)
	assert len(per_param_norm_metrics) == len(per_param_norm_metric_names)

	def is_grad_norm_metric(metric_name: str) -> bool:
	return metric_name.startswith("grad/") and metric_name.endswith(".norm")

	# Now reduce metrics over all ranks.
	total_grad_norm: torch.Tensor
	per_param_avg_metrics: List[torch.Tensor] = []
	if is_distributed(): # TODO (epwalsh): skip for non-sharded params
	# Reduce metrics across all ranks. Note that we can use a `reduce` for most cases
	# instead of an `all_reduce`, but we need `all_reduce` for norms so that all ranks
	# get the right value for gradient norms so they can clip correctly.
	# Reduce mins.
	if per_param_min_metrics:
	all_mins = torch.cat(per_param_min_metrics).to(device)
	dist.reduce(all_mins, 0, op=dist.ReduceOp.MIN)
	per_param_min_metrics = all_mins.split(1)
	# Reduce maxs.
	if per_param_max_metrics:
	all_maxs = torch.cat(per_param_max_metrics).to(device)
	dist.reduce(all_maxs, 0, op=dist.ReduceOp.MAX)
	per_param_max_metrics = all_maxs.split(1)
	# Reduce sums or just norms.
	all_norms = torch.cat(per_param_norm_metrics).to(device) ** 2.0
	if per_param_sum_metrics and per_param_numel_metrics:
	all_sums = torch.cat(per_param_sum_metrics).to(device)
	all_numels = torch.cat(per_param_numel_metrics).to(device)
	all_sums_norms_numels = torch.cat(
	[all_sums.unsqueeze(0), all_norms.unsqueeze(0), all_numels.unsqueeze(0)], dim=0
	)
	dist.all_reduce(all_sums_norms_numels, op=dist.ReduceOp.SUM)
	all_sums, all_norms, all_numels = all_sums_norms_numels.split(1)
	# Get averages.
	# NOTE: could get infs for non-rank0 processes but that's okay.
	per_param_avg_metrics = (all_sums / all_numels).squeeze(0).split(1)
	else:
	dist.all_reduce(all_norms, op=dist.ReduceOp.SUM)
	grad_norm_metric_mask = torch.tensor(
	[float(is_grad_norm_metric(n)) for n in per_param_norm_metric_names], device=all_norms.device
	)
	total_grad_norm = (all_norms * grad_norm_metric_mask).sum() ** 0.5
	per_param_norm_metrics = (all_norms ** (0.5)).squeeze(0).split(1)
	else:
	total_grad_norm = (
	torch.cat(
	[
	m
	for m, n in zip(per_param_norm_metrics, per_param_norm_metric_names)
	if is_grad_norm_metric(n)
	]
	)
	** 2.0
	).sum() ** 0.5
	per_param_avg_metrics = [x / n for x, n in zip(per_param_sum_metrics, per_param_numel_metrics)]

	assert len(per_param_avg_metrics) == len(per_param_avg_metric_names)

	# Collect all metrics into a single dict.
	all_metrics: Dict[str, torch.Tensor] = {}
	for metric_name, metric in zip(per_param_min_metric_names, per_param_min_metrics):
	all_metrics[metric_name] = metric.squeeze(0)
	for metric_name, metric in zip(per_param_max_metric_names, per_param_max_metrics):
	all_metrics[metric_name] = metric.squeeze(0)
	for metric_name, metric in zip(per_param_avg_metric_names, per_param_avg_metrics):
	all_metrics[metric_name] = metric.squeeze(0)
	for metric_name, metric in zip(per_param_norm_metric_names, per_param_norm_metrics):
	all_metrics[metric_name] = metric.squeeze(0)
	all_metrics["total_grad_norm"] = total_grad_norm

	# Clip gradients.
	num_grads_clipped = 0
	num_eligible_grads = 0
	for group in self.param_groups:
	if (max_norm_ratio := group.get("max_grad_norm_ratio")) is not None:
	num_clipped = self._do_adaptive_clipping(
	group, max_norm_ratio, global_step, all_metrics, collect_param_metrics=collect_param_metrics
	)
	elif (max_norm := group.get("max_grad_norm")) is not None:
	num_clipped = self._do_global_fixed_clipping(
	group, max_norm, all_metrics, collect_param_metrics=collect_param_metrics
	)
	else:
	# No clipping needed.
	continue
	num_eligible_grads += len(group["params"])
	if num_clipped is not None:
	num_grads_clipped += num_clipped

	if collect_param_metrics:
	if num_eligible_grads > 0:
	clipping_rate = torch.tensor(num_grads_clipped / num_eligible_grads, device="cpu")
	else:
	clipping_rate = torch.tensor(0.0, device="cpu")
	all_metrics["clipping_rate"] = clipping_rate
	return all_metrics
	else:
	return {}

	@torch.no_grad()
	def _do_adaptive_clipping(
	self,
	group: Dict[str, Any],
	max_norm_ratio: float,
	global_step: int,
	all_metrics: Dict[str, torch.Tensor],
	collect_param_metrics: bool = True,
	) -> Optional[int]:
	"""
	Do adaptive gradient clipping on a param group.

	If ``collect_param_metrics`` is ``True`` this will return the total number of gradients clipped.
	"""
	device = get_default_device()
	num_grads_clipped = 0
	# We'll use the bigger of beta1 and beta2 to update the exponential average of the norm of
	# the gradient (a scalar), not to be confused with the exponential average of the gradient.
	# TODO (epwalsh): handle optimizers that don't have betas.
	beta1, beta2 = group["betas"]
	beta = max(beta1, beta2)
	for name, p in zip(group["param_names"], group["params"]):
	name = self._clean_param_name(name)
	grad_norm = all_metrics.get(f"grad/{name}.norm")
	if grad_norm is None:
	continue

	# Get or initialize the exponential average of grad norm.
	# TODO: The way we have it right now, every rank tracks the `grad_norm_exp_avg` of every parameter,
	# even parameters for which the corresponding local shard is empty. This has the potential to
	# cause some issues with the optimizer, as we ran into with https://github.com/allenai/LLM/pull/372.
	# So we should consider changing how we do this at some point so that we don't add any state
	# to parameters for which the local shard is empty. That would probably add extra distributed
	# communication, at least on steps where we have to log (i.e. when `collect_param_metrics=True`).
	state = self.state[p]
	grad_norm_exp_avg = state.get("grad_norm_exp_avg")
	if grad_norm_exp_avg is None:
	grad_norm_exp_avg = grad_norm.clone().to(device)
	# We don't want to add anything to `state` until `state` has been initialized, otherwise
	# this will crash some optimizers which rely on checking `len(state)`. The downside here
	# is that we won't start tracking `grad_norm_exp_avg` until the 2nd training step.
	if global_step > 1:
	state["grad_norm_exp_avg"] = grad_norm_exp_avg

	max_allowed_norm = max_norm_ratio * grad_norm_exp_avg
	clip_coef = max_allowed_norm / (grad_norm + 1e-6)

	# Clip the gradients and update the exponential average.
	# Note that multiplying by the clamped coefficient is meaningless when it is
	# equal to 1, but it avoids the host-device sync that would result from `if clip_coef_clamped < 1`.
	clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
	if p.grad is not None:
	# p.grad could be none for some ranks when using FSDP.
	p.grad.detach().mul_(clip_coef_clamped.to(p.grad.device, p.grad.dtype))

	# Update the exponential average of the norm of the gradient with the clipped norm of the gradient.
	grad_norm_exp_avg.lerp_((grad_norm * clip_coef_clamped).to(grad_norm_exp_avg.device), 1 - beta)
	# Alternative: update with the unclipped norm of the gradient.
	# grad_norm_exp_avg.lerp_(grad_norm.to(grad_norm_exp_avg.device), 1 - beta)

	if collect_param_metrics:
	# Can't avoid host-device sync here.
	if clip_coef_clamped < 1.0:
	num_grads_clipped += 1
	all_metrics[f"grad_norm_exp_avg/{name}"] = grad_norm_exp_avg
	return num_grads_clipped if collect_param_metrics else None

	@torch.no_grad()
	def _do_global_fixed_clipping(
	self,
	group: Dict[str, Any],
	max_norm: float,
	all_metrics: Dict[str, torch.Tensor],
	collect_param_metrics: bool = True,
	) -> Optional[int]:
	"""
	Do global fixed gradient clipping on a param group.

	If ``collect_param_metrics`` is ``True`` this will return the total number of gradients clipped.
	"""
	device = get_default_device()
	total_grad_norm = all_metrics["total_grad_norm"]
	clip_coef = max_norm / (total_grad_norm.to(device) + 1e-6)
	clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
	num_grads_clipped: Optional[int] = None
	if collect_param_metrics:
	# Can't avoid host-device sync here.
	if clip_coef_clamped < 1.0:
	num_grads_clipped = len(group["params"])
	for p in group["params"]:
	# Clip the gradients.
	# Note that multiplying by the clamped coefficient is meaningless when it is
	# equal to 1, but it avoids the host-device sync that would result from `if clip_coef_clamped < 1`.
	if p.grad is not None:
	# p.grad could be none for some ranks when using FSDP.
	p.grad.detach().mul_(clip_coef_clamped.to(p.grad.device, p.grad.dtype))
	return num_grads_clipped

	def get_post_step_metrics(self, module: nn.Module) -> Dict[str, torch.Tensor]:
	del module
	return {}

	def get_state_for_param(self, param: nn.Parameter) -> Dict[str, Optional[torch.Tensor]]:
	del param
	return {}


	class LionW(Optimizer):
	"""
	Adapted from https://github.com/google/automl/blob/master/lion/lion_pytorch.py
	"""

	def __init__(
	self,
	params,
	lr: float = 1e-4,
	betas: Tuple[float, float] = (0.9, 0.99),
	weight_decay: float = 0.0,
	):
	assert lr > 0.0
	assert all([0.0 <= beta <= 1.0 for beta in betas])
	defaults = dict(lr=lr, betas=betas, weight_decay=weight_decay)
	super().__init__(params, defaults)
	for group in self.param_groups:
	group["initial_lr"] = group["lr"]
	self._update_total_dot_prod: Optional[torch.Tensor] = None
	self._update_total_norm: Optional[torch.Tensor] = None
	self._signed_update_total_norm: Optional[torch.Tensor] = None

	def get_post_step_metrics(self, module: nn.Module) -> Dict[str, torch.Tensor]:
	update_total_dot_prod = self._update_total_dot_prod
	update_total_norm = self._update_total_norm
	signed_update_total_norm = self._signed_update_total_norm
	if update_total_dot_prod is None or update_total_norm is None or signed_update_total_norm is None:
	return {}

	if is_distributed() and isinstance(module, FullyShardedDataParallel):
	# Reduce total dot prod and norms across all ranks.
	update_total_norm = update_total_norm**2.0
	signed_update_total_norm = signed_update_total_norm**2.0
	# Reduce all together to avoid multiple communication calls.
	all_together = torch.stack([update_total_dot_prod, update_total_norm, signed_update_total_norm])
	# Only need the final result on rank0, since that's where we log from.
	dist.reduce(all_together, 0)
	update_total_dot_prod, update_total_norm, signed_update_total_norm = all_together
	update_total_norm = update_total_norm**0.5
	signed_update_total_norm = signed_update_total_norm**0.5

	update_cos_sim = update_total_dot_prod / torch.max(
	update_total_norm * signed_update_total_norm, torch.tensor(1e-8, device=get_default_device())
	)
	return {"update_cos_sim": update_cos_sim}

	@torch.no_grad()
	def step(self, closure=None) -> None:
	if closure is not None:
	with torch.enable_grad():
	closure()

	update_total_dot_prod = torch.tensor(0.0, dtype=torch.float32)
	update_norms = []
	signed_update_norms = []

	for group in self.param_groups:
	for p in group["params"]:
	if p.grad is None:
	continue

	# Perform step weight decay
	p.data.mul_(1 - group["lr"] * group["weight_decay"])

	grad = p.grad
	state = self.state[p]

	# State initialization
	if len(state) == 0:
	# Exponential moving average of gradient values
	state["exp_avg"] = torch.zeros_like(p)

	exp_avg = state["exp_avg"]
	beta1, beta2 = group["betas"]

	# Weight update
	update = exp_avg * beta1 + grad * (1 - beta1)
	signed_update = torch.sign(update)
	p.add_(signed_update, alpha=-group["lr"])

	# Decay the momentum running average coefficient
	exp_avg.mul_(beta2).add_(grad, alpha=1 - beta2)

	# Track dot product and norms of update vs signed update in order to calculate
	# their cosine similarity.
	update_total_dot_prod = update_total_dot_prod.to(update.device)
	update_total_dot_prod += torch.tensordot(update, signed_update, dims=len(update.shape))
	update_norms.append(torch.linalg.vector_norm(update, 2.0, dtype=torch.float32))
	signed_update_norms.append(torch.linalg.vector_norm(signed_update, 2.0, dtype=torch.float32))

	# Compute cosine similarity between update and signed update.
	self._update_total_dot_prod = update_total_dot_prod.to(get_default_device())
	self._update_total_norm = torch.linalg.vector_norm(
	torch.stack(update_norms),
	2.0,
	dtype=torch.float32,
	).to(get_default_device())
	self._signed_update_total_norm = torch.linalg.vector_norm(
	torch.stack(signed_update_norms),
	2.0,
	dtype=torch.float32,
	).to(get_default_device())


	class AdamW(torch.optim.AdamW, Optimizer):
	def get_state_for_param(self, param: nn.Parameter) -> Dict[str, Optional[torch.Tensor]]:
	return {key: self.state[param].get(key) for key in ("exp_avg", "exp_avg_sq")} # type: ignore


	@dataclass
	class Scheduler(metaclass=ABCMeta):
	# NOTE: these fields are not given default values because otherwise dataclasses complains
	# about how the scheduler subclasses are defined.
	grad_clip_warmup_steps: Optional[int]
	grad_clip_warmup_factor: Optional[float]

	@abstractmethod
	def get_lr(self, initial_lr: float, step: int, max_steps: int) -> float:
	raise NotImplementedError

	def _get_max_grad_norm_coeff(
	self, initial_value: Optional[float], step: int, max_steps: int
	) -> Optional[float]:
	del max_steps # might need this in the future, but for now I just wanted to match the API of `get_lr()`.
	if initial_value is None:
	return None
	elif (
	self.grad_clip_warmup_steps is None
	or self.grad_clip_warmup_factor is None
	or step > self.grad_clip_warmup_steps
	):
	return initial_value
	else:
	return self.grad_clip_warmup_factor * initial_value

	def get_max_grad_norm(
	self, initial_max_grad_norm: Optional[float], step: int, max_steps: int
	) -> Optional[float]:
	return self._get_max_grad_norm_coeff(initial_max_grad_norm, step, max_steps)

	def get_max_grad_norm_ratio(
	self, initial_max_grad_norm_ratio: Optional[float], step: int, max_steps: int
	) -> Optional[float]:
	return self._get_max_grad_norm_coeff(initial_max_grad_norm_ratio, step, max_steps)

	def _linear_warmup(self, initial_lr: float, step: int, warmup_steps: int = 2000) -> float:
	return initial_lr * (0.1 + 0.9 * min(step, warmup_steps) / warmup_steps)


	@dataclass
	class CosWithWarmup(Scheduler):
	warmup_steps: int
	alpha_f: float = 0.1
	t_max: Optional[int] = None

	def get_lr(self, initial_lr: float, step: int, max_steps: int) -> float:
	max_steps = max_steps if self.t_max is None else self.t_max
	eta_min = initial_lr * self.alpha_f
	if step < self.warmup_steps:
	return self._linear_warmup(initial_lr, step, self.warmup_steps)
	elif step >= max_steps:
	return eta_min
	else:
	step = step - self.warmup_steps
	max_steps = max_steps - self.warmup_steps
	return eta_min + (initial_lr - eta_min) * (1 + cos(pi * step / max_steps)) / 2


	@dataclass
	class LinearWithWarmup(Scheduler):
	warmup_steps: int
	alpha_f: float = 0.1
	t_max: Optional[int] = None

	def get_lr(self, initial_lr: float, step: int, max_steps: int) -> float:
	max_steps = max_steps if self.t_max is None else self.t_max
	eta_min = initial_lr * self.alpha_f
	if step < self.warmup_steps:
	return self._linear_warmup(initial_lr, step, self.warmup_steps)
	elif step >= max_steps:
	return eta_min
	else:
	step = step - self.warmup_steps
	max_steps = max_steps - self.warmup_steps
	return initial_lr - (initial_lr - eta_min) * (step / max_steps)


	@dataclass
	class InvSqrtWithWarmup(Scheduler):
	warmup_steps: int

	def get_lr(self, initial_lr: float, step: int, max_steps: int) -> float:
	if step < self.warmup_steps:
	return self._linear_warmup(initial_lr, step, self.warmup_steps)
	del max_steps
	return initial_lr * sqrt(self.warmup_steps / max(self.warmup_steps, step))


	@dataclass
	class MaxScheduler(Scheduler):
	sched1: Scheduler
	sched2: Scheduler

	def get_lr(self, initial_lr: float, step: int, max_steps: int) -> float:
	return max(
	self.sched1.get_lr(initial_lr, step, max_steps), self.sched2.get_lr(initial_lr, step, max_steps)
	)


	@dataclass
	class BoltOnWarmupScheduler(Scheduler):
	inner: Scheduler
	warmup_start: int
	warmup_end: int

	@classmethod
	def wrap(cls, scheduler: Scheduler, warmup_start: int, warmup_end: int) -> "BoltOnWarmupScheduler":
	return cls(
	grad_clip_warmup_steps=None,
	grad_clip_warmup_factor=None,
	inner=scheduler,
	warmup_start=warmup_start,
	warmup_end=warmup_end,
	)

	def get_lr(self, initial_lr: float, step: int, max_steps: int) -> float:
	if step < self.warmup_start:
	return 0.0
	if step < self.warmup_end:
	lr_at_intercept = self.inner.get_lr(initial_lr, self.warmup_end, max_steps)
	return lr_at_intercept * (step - self.warmup_start) / (self.warmup_end - self.warmup_start)
	else:
	return self.inner.get_lr(initial_lr, step, max_steps)

	def _get_max_grad_norm_coeff(
	self, initial_value: Optional[float], step: int, max_steps: int
	) -> Optional[float]:
	return self.inner._get_max_grad_norm_coeff(initial_value, step, max_steps)


	@dataclass
	class ConstantScheduler(Scheduler):
	def get_lr(self, initial_lr: float, step: int, max_steps: int) -> float:
	del step, max_steps
	return initial_lr


	PARAM_GROUP_FIELDS = ("sharded", "max_grad_norm", "max_grad_norm_ratio", "param_names")


	def get_param_groups(cfg: TrainConfig, model: nn.Module) -> List[Dict[str, Any]]:
	"""
	Separate parameters into weight decay and non weight decay groups.
	"""
	param_groups: List[Dict[str, Any]]
	param_group_defaults = {
	"sharded": isinstance(model, FullyShardedDataParallel),
	"max_grad_norm": cfg.max_grad_norm,
	"max_grad_norm_ratio": cfg.max_grad_norm_ratio,
	}

	# Separate out parameters that we don't want to apply weight decay to, like norms and biases.
	decay = set()
	no_decay = set()
	all_params = {}
	for mn, m in model.named_modules():
	for pn, p in m.named_parameters():
	# NOTE: because named_modules and named_parameters are recursive
	# we will see the same tensors p many many times, but doing it this way
	# allows us to know which parent module any tensor p belongs to...
	if not p.requires_grad:
	continue

	fpn = f"{mn}.{pn}" if mn else pn
	all_params[fpn] = p

	if pn.endswith("bias"):
	if cfg.optimizer.decay_norm_and_bias:
	decay.add(fpn)
	else:
	no_decay.add(fpn)
	elif pn.endswith("weight") and (isinstance(m, nn.Linear) or isinstance(m, BitLinear158)):
	decay.add(fpn)
	elif pn.endswith("weight") and isinstance(m, (LayerNormBase, nn.LayerNorm)):
	if cfg.optimizer.decay_norm_and_bias:
	decay.add(fpn)
	else:
	no_decay.add(fpn)
	elif pn.endswith("weight") and isinstance(m, nn.Embedding):
	if cfg.optimizer.decay_embeddings:
	decay.add(fpn)
	else:
	no_decay.add(fpn)

	# Validate that we've considered every parameter
	inter_params = decay & no_decay
	union_params = decay \| no_decay
	assert len(inter_params) == 0, f"parameters {inter_params} made it into both decay/no_decay sets!"
	assert (
	len(all_params.keys() - union_params) == 0
	), f"parameters {all_params.keys() - union_params} were not separated into either decay/no_decay set!"

	# Create the pytorch optimizer groups.
	decay_sorted = sorted(list(decay))
	no_decay_sorted = sorted(list(no_decay))
	param_groups = []
	if len(decay_sorted) > 0:
	param_groups.append(
	{
	"params": [all_params[pn] for pn in decay_sorted],
	"param_names": decay_sorted,
	**param_group_defaults,
	}
	)
	if len(no_decay_sorted) > 0:
	param_groups.append(
	{
	"params": [all_params[pn] for pn in no_decay_sorted],
	"param_names": no_decay_sorted,
	"weight_decay": 0.0,
	**param_group_defaults,
	}
	)

	# Validate fields.
	for group in param_groups:
	for key in PARAM_GROUP_FIELDS:
	assert key in group

	return param_groups


	def fix_optim_state_dict(optimizer: Optimizer, state_dict: Dict[str, Any]) -> Dict[str, Any]:
	"""
	Make sure old optim state dicts are compatible with new versions.
	"""
	if len(state_dict["param_groups"]) == 1 and len(optimizer.param_groups) == 2:
	assert optimizer.param_groups[1]["weight_decay"] == 0.0

	# Decay
	decay_param_group = {k: v for k, v in state_dict["param_groups"][0].items() if k != "params"}
	decay_param_group["params"] = optimizer.state_dict()["param_groups"][0]["params"]

	# No decay.
	no_decay_param_group = {k: v for k, v in state_dict["param_groups"][0].items() if k != "params"}
	no_decay_param_group["weight_decay"] = 0.0
	no_decay_param_group["params"] = optimizer.state_dict()["param_groups"][1]["params"]

	state_dict["param_groups"] = [decay_param_group, no_decay_param_group]

	assert len(optimizer.param_groups) == len(state_dict["param_groups"])

	# Make sure:
	# - All required fields are included in the state dict,
	# - And that the values of those fields doesn't change from what's currently set in the optimizer,
	# since we might have changed those fields on purpose after a restart.
	for group, sd_group in zip(optimizer.param_groups, state_dict["param_groups"]):
	for key in PARAM_GROUP_FIELDS:
	sd_group[key] = group[key]

	return state_dict


	def build_optimizer(cfg: TrainConfig, model: nn.Module) -> Optimizer:
	param_groups = get_param_groups(cfg, model)
	log.info(f"Constructing optimizer with {len(param_groups)} param groups")
	if cfg.optimizer.name == OptimizerType.lionw:
	return LionW(
	param_groups,
	lr=cfg.optimizer.learning_rate,
	betas=cfg.optimizer.betas,
	weight_decay=cfg.optimizer.weight_decay,
	)
	elif cfg.optimizer.name == OptimizerType.adamw:
	return AdamW(
	param_groups,
	lr=cfg.optimizer.learning_rate,
	betas=cfg.optimizer.betas,
	weight_decay=cfg.optimizer.weight_decay,
	eps=1e-5,
	)
	else:
	raise NotImplementedError


	def build_scheduler(cfg: TrainConfig, sched_cfg: Optional[SchedulerConfig] = None) -> Scheduler:
	sched_cfg = sched_cfg if sched_cfg is not None else cfg.scheduler
	if sched_cfg.name == SchedulerType.cosine_with_warmup:
	return CosWithWarmup(
	grad_clip_warmup_steps=None
	if sched_cfg.grad_clip_warmup_steps is None
	else int(sched_cfg.grad_clip_warmup_steps),
	grad_clip_warmup_factor=sched_cfg.grad_clip_warmup_factor,
	warmup_steps=int(sched_cfg.t_warmup),
	alpha_f=sched_cfg.alpha_f,
	t_max=None if sched_cfg.t_max is None else int(sched_cfg.t_max),
	)
	elif sched_cfg.name == SchedulerType.linear_with_warmup:
	return LinearWithWarmup(
	grad_clip_warmup_steps=None
	if sched_cfg.grad_clip_warmup_steps is None
	else int(sched_cfg.grad_clip_warmup_steps),
	grad_clip_warmup_factor=sched_cfg.grad_clip_warmup_factor,
	warmup_steps=int(sched_cfg.t_warmup),
	alpha_f=sched_cfg.alpha_f,
	t_max=None if sched_cfg.t_max is None else int(sched_cfg.t_max),
	)
	elif sched_cfg.name == SchedulerType.inverse_sqrt_with_warmup:
	return InvSqrtWithWarmup(
	grad_clip_warmup_steps=None
	if sched_cfg.grad_clip_warmup_steps is None
	else int(sched_cfg.grad_clip_warmup_steps),
	grad_clip_warmup_factor=sched_cfg.grad_clip_warmup_factor,
	warmup_steps=int(sched_cfg.t_warmup),
	)
	elif sched_cfg.name == SchedulerType.max_scheduler:
	return MaxScheduler(
	grad_clip_warmup_steps=None
	if sched_cfg.grad_clip_warmup_steps is None
	else int(sched_cfg.grad_clip_warmup_steps),
	grad_clip_warmup_factor=sched_cfg.grad_clip_warmup_factor,
	sched1=build_scheduler(cfg, replace(sched_cfg, name=SchedulerType.cosine_with_warmup)),
	sched2=build_scheduler(cfg, replace(sched_cfg, name=SchedulerType.inverse_sqrt_with_warmup)),
	)
	elif sched_cfg.name == SchedulerType.constant:
	return ConstantScheduler(
	grad_clip_warmup_steps=None
	if sched_cfg.grad_clip_warmup_steps is None
	else int(sched_cfg.grad_clip_warmup_steps),
	grad_clip_warmup_factor=sched_cfg.grad_clip_warmup_factor,
	)
	else:
	raise NotImplementedError