chore: import upstream snapshot with attribution

nemo/core/optim/__init__.py

@@ -0,0 +1,38 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# flake8: noqa
+
+from nemo.core.optim.adafactor import Adafactor
+from nemo.core.optim.adan import Adan
+from nemo.core.optim.flash_optim import patch_flashoptim_uneven_shard_support
+from nemo.core.optim.lr_scheduler import (
+    CosineAnnealing,
+    InverseSquareRootAnnealing,
+    NoamAnnealing,
+    PolynomialDecayAnnealing,
+    PolynomialHoldDecayAnnealing,
+    SquareAnnealing,
+    SquareRootAnnealing,
+    T5InverseSquareRootAnnealing,
+    WarmupAnnealing,
+    WarmupHoldAnnealLinear,
+    WarmupHoldAnnealOneMinusSquareRoot,
+    WarmupHoldPolicy,
+    WarmupPolicy,
+    prepare_lr_scheduler,
+)
+from nemo.core.optim.novograd import Novograd
+from nemo.core.optim.optimizer_with_main_params import MainParamsOptimizerWrapper
+from nemo.core.optim.optimizers import get_optimizer, parse_optimizer_args, register_optimizer

nemo/core/optim/adafactor.py

@@ -0,0 +1,218 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Most of the code here has been copied from:
+# https://github.com/pytorch/fairseq/blob/main/fairseq/optim/adafactor.py
+
+import math
+
+import torch
+from torch.optim.optimizer import Optimizer
+
+__all__ = ['Adafactor']
+
+
+class Adafactor(Optimizer):
+    """Implements Adafactor algorithm.
+    This implementation is based on:
+    `Adafactor: Adaptive Learning Rates with Sublinear Memory Cost`
+    (see https://arxiv.org/abs/1804.04235)
+    Note that this optimizer internally adjusts the learning rate
+    depending on the *scale_parameter*, *relative_step* and
+    *warmup_init* options. To use a manual (external) learning rate
+    schedule you should set `scale_parameter=False` and
+    `relative_step=False`.
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): external learning rate (default: None)
+        eps (tuple[float, float]): regularization constans for square gradient
+            and parameter scale respectively (default: (1e-30, 1e-3))
+        clip_threshold (float): threshold of root mean square of
+            final gradient update (default: 1.0)
+        decay_rate (float): coefficient used to compute running averages of square
+            gradient (default: -0.8)
+        beta1 (float): coefficient used for computing running averages of gradient
+            (default: None)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        scale_parameter (bool): if True, learning rate is scaled by root mean square of
+            parameter (default: True)
+        relative_step (bool): if True, time-dependent learning rate is computed
+            instead of external learning rate (default: True)
+        warmup_init (bool): time-dependent learning rate computation depends on
+            whether warm-up initialization is being used (default: False)
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=None,
+        eps=(1e-30, 1e-3),
+        clip_threshold=1.0,
+        decay_rate=-0.8,
+        beta1=None,
+        weight_decay=0.0,
+        scale_parameter=True,
+        relative_step=True,
+        warmup_init=False,
+        min_step=1e-2,
+    ):
+        if lr is not None and relative_step:
+            raise ValueError("Cannot combine manual lr and relative_step options")
+        if warmup_init and not relative_step:
+            raise ValueError("warmup_init requires relative_step=True")
+        self.min_step = min_step
+
+        defaults = dict(
+            lr=lr,
+            eps=eps,
+            clip_threshold=clip_threshold,
+            decay_rate=decay_rate,
+            beta1=beta1,
+            weight_decay=weight_decay,
+            scale_parameter=scale_parameter,
+            relative_step=relative_step,
+            warmup_init=warmup_init,
+            min_step=min_step,
+        )
+        super(Adafactor, self).__init__(params, defaults)
+
+    @property
+    def supports_memory_efficient_fp16(self):
+        """Whether this optimizer supports memory-efficient fp16 training."""
+        return True
+
+    @property
+    def supports_flat_params(self):
+        """Whether this optimizer supports flattened parameters."""
+        return False
+
+    def _get_lr(self, param_group, param_state):
+        rel_step_sz = param_group["lr"]
+        if param_group["relative_step"]:
+            min_step = 1e-6 * param_state["step"] if param_group["warmup_init"] else self.min_step
+            rel_step_sz = min(min_step, 1.0 / math.sqrt(param_state["step"]))
+        param_scale = 1.0
+        if param_group["scale_parameter"]:
+            param_scale = max(param_group["eps"][1], param_state["RMS"])
+        return param_scale * rel_step_sz
+
+    def _get_options(self, param_group, param_shape):
+        factored = len(param_shape) >= 2
+        use_first_moment = param_group["beta1"] is not None
+        return factored, use_first_moment
+
+    def _rms(self, tensor):
+        return tensor.norm(2) / (tensor.numel() ** 0.5)
+
+    def _approx_sq_grad(self, exp_avg_sq_row, exp_avg_sq_col):
+        r_factor = (exp_avg_sq_row / exp_avg_sq_row.mean(dim=-1, keepdim=True)).rsqrt_().unsqueeze(-1)
+        c_factor = exp_avg_sq_col.unsqueeze(-2).rsqrt()
+        return torch.mul(r_factor, c_factor)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Args:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.dtype in {torch.float16, torch.bfloat16}:
+                    grad = grad.float()
+                if grad.is_sparse:
+                    raise RuntimeError("Adafactor does not support sparse gradients.")
+
+                state = self.state[p]
+                grad_shape = grad.shape
+
+                factored, use_first_moment = self._get_options(group, grad_shape)
+                # State Initialization
+                if len(state) == 0:
+                    state["step"] = 0
+
+                    if use_first_moment:
+                        # Exponential moving average of gradient values
+                        state["exp_avg"] = torch.zeros_like(grad)
+                    if factored:
+                        state["exp_avg_sq_row"] = torch.zeros(grad_shape[:-1]).to(grad)
+                        state["exp_avg_sq_col"] = torch.zeros(grad_shape[:-2] + grad_shape[-1:]).to(grad)
+                    else:
+                        state["exp_avg_sq"] = torch.zeros_like(grad)
+
+                    state["RMS"] = 0
+                else:
+                    if use_first_moment:
+                        state["exp_avg"] = state["exp_avg"].to(grad)
+                    if factored:
+                        state["exp_avg_sq_row"] = state["exp_avg_sq_row"].to(grad)
+                        state["exp_avg_sq_col"] = state["exp_avg_sq_col"].to(grad)
+                    else:
+                        state["exp_avg_sq"] = state["exp_avg_sq"].to(grad)
+
+                p_data_fp32 = p.data
+                if p.data.dtype in {torch.float16, torch.bfloat16}:
+                    p_data_fp32 = p_data_fp32.float()
+
+                state["step"] += 1
+                state["RMS"] = self._rms(p_data_fp32)
+                group["lr"] = self._get_lr(group, state)
+
+                beta2t = 1.0 - math.pow(state["step"], group["decay_rate"])
+                update = (grad**2) + group["eps"][0]
+                if factored:
+                    exp_avg_sq_row = state["exp_avg_sq_row"]
+                    exp_avg_sq_col = state["exp_avg_sq_col"]
+
+                    exp_avg_sq_row.mul_(beta2t).add_(update.mean(dim=-1), alpha=1.0 - beta2t)
+                    exp_avg_sq_col.mul_(beta2t).add_(update.mean(dim=-2), alpha=1.0 - beta2t)
+
+                    # Approximation of exponential moving average of square of gradient
+                    update = self._approx_sq_grad(exp_avg_sq_row, exp_avg_sq_col)
+                    update.mul_(grad)
+                else:
+                    exp_avg_sq = state["exp_avg_sq"]
+
+                    exp_avg_sq.mul_(beta2t).add_(update, alpha=1.0 - beta2t)
+                    update = exp_avg_sq.rsqrt().mul_(grad)
+
+                update.div_((self._rms(update) / group["clip_threshold"]).clamp_(min=1.0))
+                update.mul_(group["lr"])
+
+                if use_first_moment:
+                    exp_avg = state["exp_avg"]
+                    exp_avg.mul_(group["beta1"]).add_(update, alpha=1 - group["beta1"])
+                    update = exp_avg
+
+                if group["weight_decay"] != 0:
+                    p_data_fp32.add_(p_data_fp32, alpha=-group["weight_decay"] * group["lr"])
+
+                p_data_fp32.add_(-update)
+
+                if p.data.dtype in {torch.float16, torch.bfloat16}:
+                    p.data.copy_(p_data_fp32)
+
+        return loss

nemo/core/optim/adan.py

@@ -0,0 +1,453 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Copyright 2022 Garena Online Private Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from typing import List
+
+import torch
+from torch import Tensor
+from torch.optim.optimizer import Optimizer
+
+
+class MultiTensorApply(object):
+    available = False
+    warned = False
+
+    def __init__(self, chunk_size):
+        try:
+            MultiTensorApply.available = True
+            self.chunk_size = chunk_size
+        except ImportError as err:
+            MultiTensorApply.available = False
+            MultiTensorApply.import_err = err
+
+    def __call__(self, op, noop_flag_buffer, tensor_lists, *args):
+        return op(self.chunk_size, noop_flag_buffer, tensor_lists, *args)
+
+
+class Adan(Optimizer):
+    """
+    Implements a pytorch variant of Adan
+    Adan was proposed in
+    Adan: Adaptive Nesterov Momentum Algorithm for
+        Faster Optimizing Deep Models[J].arXiv preprint arXiv:2208.06677, 2022.
+    https://arxiv.org/abs/2208.06677
+    Arguments:
+        params (iterable): iterable of parameters to optimize or
+            dicts defining parameter groups.
+        lr (float, optional): learning rate. (default: 1e-3)
+        betas (Tuple[float, float, flot], optional): coefficients used for
+            first- and second-order moments. (default: (0.98, 0.92, 0.99))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability. (default: 1e-8)
+        weight_decay (float, optional): decoupled weight decay
+            (L2 penalty) (default: 0)
+        max_grad_norm (float, optional): value used to clip
+            global grad norm (default: 0.0 no clip)
+        no_prox (bool): how to perform the decoupled weight decay
+            (default: False)
+        foreach (bool): if True would use torch._foreach implementation.
+            It's faster but uses slightly more memory. (default: True)
+        fused (bool, optional): whether fused implementation is used.
+            (default: False)
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        betas=(0.98, 0.92, 0.99),
+        eps=1e-8,
+        weight_decay=0.0,
+        max_grad_norm=0.0,
+        no_prox=False,
+        foreach: bool = True,
+        fused: bool = False,
+    ):
+        if not 0.0 <= max_grad_norm:
+            raise ValueError('Invalid Max grad norm: {}'.format(max_grad_norm))
+        if not 0.0 <= lr:
+            raise ValueError('Invalid learning rate: {}'.format(lr))
+        if not 0.0 <= eps:
+            raise ValueError('Invalid epsilon value: {}'.format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError('Invalid beta parameter at index 0: {}'.format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError('Invalid beta parameter at index 1: {}'.format(betas[1]))
+        if not 0.0 <= betas[2] < 1.0:
+            raise ValueError('Invalid beta parameter at index 2: {}'.format(betas[2]))
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            max_grad_norm=max_grad_norm,
+            no_prox=no_prox,
+            foreach=foreach,
+            fused=fused,
+        )
+        super().__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(Adan, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('no_prox', False)
+
+    @torch.no_grad()
+    def restart_opt(self):
+        for group in self.param_groups:
+            group['step'] = 0
+            for p in group['params']:
+                if p.requires_grad:
+                    state = self.state[p]
+                    # State initialization
+
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p)
+                    # Exponential moving average of gradient difference
+                    state['exp_avg_diff'] = torch.zeros_like(p)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        """Performs a single optimization step."""
+
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        if self.defaults['max_grad_norm'] > 0:
+            device = self.param_groups[0]['params'][0].device
+            global_grad_norm = torch.zeros(1, device=device)
+
+            max_grad_norm = torch.tensor(self.defaults['max_grad_norm'], device=device)
+            for group in self.param_groups:
+
+                for p in group['params']:
+                    if p.grad is not None:
+                        grad = p.grad
+                        global_grad_norm.add_(grad.pow(2).sum())
+
+            global_grad_norm = torch.sqrt(global_grad_norm)
+
+            clip_global_grad_norm = torch.clamp(max_grad_norm / (global_grad_norm + group['eps']), max=1.0).item()
+        else:
+            clip_global_grad_norm = 1.0
+
+        for group in self.param_groups:
+            params_with_grad = []
+            grads = []
+            exp_avgs = []
+            exp_avg_sqs = []
+            exp_avg_diffs = []
+            neg_pre_grads = []
+
+            beta1, beta2, beta3 = group['betas']
+            # assume same step across group now to simplify things
+            # per parameter step can be easily support
+            # by making it tensor, or pass list into kernel
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            bias_correction1 = 1.0 - beta1 ** group['step']
+            bias_correction2 = 1.0 - beta2 ** group['step']
+            bias_correction3 = 1.0 - beta3 ** group['step']
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                params_with_grad.append(p)
+                grads.append(p.grad)
+
+                state = self.state[p]
+                if len(state) == 0:
+                    state['exp_avg'] = torch.zeros_like(p)
+                    state['exp_avg_sq'] = torch.zeros_like(p)
+                    state['exp_avg_diff'] = torch.zeros_like(p)
+
+                if 'neg_pre_grad' not in state or group['step'] == 1:
+                    state['neg_pre_grad'] = p.grad.clone().mul_(-clip_global_grad_norm)
+
+                exp_avgs.append(state['exp_avg'])
+                exp_avg_sqs.append(state['exp_avg_sq'])
+                exp_avg_diffs.append(state['exp_avg_diff'])
+                neg_pre_grads.append(state['neg_pre_grad'])
+
+            kwargs = dict(
+                params=params_with_grad,
+                grads=grads,
+                exp_avgs=exp_avgs,
+                exp_avg_sqs=exp_avg_sqs,
+                exp_avg_diffs=exp_avg_diffs,
+                neg_pre_grads=neg_pre_grads,
+                beta1=beta1,
+                beta2=beta2,
+                beta3=beta3,
+                bias_correction1=bias_correction1,
+                bias_correction2=bias_correction2,
+                bias_correction3_sqrt=math.sqrt(bias_correction3),
+                lr=group['lr'],
+                weight_decay=group['weight_decay'],
+                eps=group['eps'],
+                no_prox=group['no_prox'],
+                clip_global_grad_norm=clip_global_grad_norm,
+            )
+
+            if group['foreach']:
+                if group['fused']:
+                    if torch.cuda.is_available():
+                        _fused_adan_multi_tensor(**kwargs)
+                    else:
+                        raise ValueError('Fused Adan does not support CPU')
+                else:
+                    _multi_tensor_adan(**kwargs)
+            elif group['fused']:
+                if torch.cuda.is_available():
+                    _fused_adan_single_tensor(**kwargs)
+                else:
+                    raise ValueError('Fused Adan does not support CPU')
+            else:
+                _single_tensor_adan(**kwargs)
+
+        return loss
+
+
+def _single_tensor_adan(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    exp_avg_diffs: List[Tensor],
+    neg_pre_grads: List[Tensor],
+    *,
+    beta1: float,
+    beta2: float,
+    beta3: float,
+    bias_correction1: float,
+    bias_correction2: float,
+    bias_correction3_sqrt: float,
+    lr: float,
+    weight_decay: float,
+    eps: float,
+    no_prox: bool,
+    clip_global_grad_norm: Tensor,
+):
+    for i, param in enumerate(params):
+        grad = grads[i]
+        exp_avg = exp_avgs[i]
+        exp_avg_sq = exp_avg_sqs[i]
+        exp_avg_diff = exp_avg_diffs[i]
+        neg_grad_or_diff = neg_pre_grads[i]
+
+        grad.mul_(clip_global_grad_norm)
+
+        # for memory saving, we use `neg_grad_or_diff`
+        # to get some temp variable in a inplace way
+        neg_grad_or_diff.add_(grad)
+
+        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)  # m_t
+        exp_avg_diff.mul_(beta2).add_(neg_grad_or_diff, alpha=1 - beta2)  # diff_t
+
+        neg_grad_or_diff.mul_(beta2).add_(grad)
+        exp_avg_sq.mul_(beta3).addcmul_(neg_grad_or_diff, neg_grad_or_diff, value=1 - beta3)  # n_t
+
+        denom = ((exp_avg_sq).sqrt() / bias_correction3_sqrt).add_(eps)
+        step_size_diff = lr * beta2 / bias_correction2
+        step_size = lr / bias_correction1
+
+        if no_prox:
+            param.mul_(1 - lr * weight_decay)
+            param.addcdiv_(exp_avg, denom, value=-step_size)
+            param.addcdiv_(exp_avg_diff, denom, value=-step_size_diff)
+        else:
+            param.addcdiv_(exp_avg, denom, value=-step_size)
+            param.addcdiv_(exp_avg_diff, denom, value=-step_size_diff)
+            param.div_(1 + lr * weight_decay)
+
+        neg_grad_or_diff.zero_().add_(grad, alpha=-1.0)
+
+
+def _multi_tensor_adan(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    exp_avg_diffs: List[Tensor],
+    neg_pre_grads: List[Tensor],
+    *,
+    beta1: float,
+    beta2: float,
+    beta3: float,
+    bias_correction1: float,
+    bias_correction2: float,
+    bias_correction3_sqrt: float,
+    lr: float,
+    weight_decay: float,
+    eps: float,
+    no_prox: bool,
+    clip_global_grad_norm: Tensor,
+):
+    if len(params) == 0:
+        return
+
+    torch._foreach_mul_(grads, clip_global_grad_norm)
+
+    # for memory saving, we use `neg_pre_grads`
+    # to get some temp variable in a inplace way
+    torch._foreach_add_(neg_pre_grads, grads)
+
+    torch._foreach_mul_(exp_avgs, beta1)
+    torch._foreach_add_(exp_avgs, grads, alpha=1 - beta1)  # m_t
+
+    torch._foreach_mul_(exp_avg_diffs, beta2)
+    torch._foreach_add_(exp_avg_diffs, neg_pre_grads, alpha=1 - beta2)  # diff_t
+
+    torch._foreach_mul_(neg_pre_grads, beta2)
+    torch._foreach_add_(neg_pre_grads, grads)
+    torch._foreach_mul_(exp_avg_sqs, beta3)
+    torch._foreach_addcmul_(exp_avg_sqs, neg_pre_grads, neg_pre_grads, value=1 - beta3)  # n_t
+
+    denom = torch._foreach_sqrt(exp_avg_sqs)
+    torch._foreach_div_(denom, bias_correction3_sqrt)
+    torch._foreach_add_(denom, eps)
+
+    step_size_diff = lr * beta2 / bias_correction2
+    step_size = lr / bias_correction1
+
+    if no_prox:
+        torch._foreach_mul_(params, 1 - lr * weight_decay)
+        torch._foreach_addcdiv_(params, exp_avgs, denom, value=-step_size)
+        torch._foreach_addcdiv_(params, exp_avg_diffs, denom, value=-step_size_diff)
+    else:
+        torch._foreach_addcdiv_(params, exp_avgs, denom, value=-step_size)
+        torch._foreach_addcdiv_(params, exp_avg_diffs, denom, value=-step_size_diff)
+        torch._foreach_div_(params, 1 + lr * weight_decay)
+    torch._foreach_zero_(neg_pre_grads)
+    torch._foreach_add_(neg_pre_grads, grads, alpha=-1.0)
+
+
+def _fused_adan_multi_tensor(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    exp_avg_diffs: List[Tensor],
+    neg_pre_grads: List[Tensor],
+    *,
+    beta1: float,
+    beta2: float,
+    beta3: float,
+    bias_correction1: float,
+    bias_correction2: float,
+    bias_correction3_sqrt: float,
+    lr: float,
+    weight_decay: float,
+    eps: float,
+    no_prox: bool,
+    clip_global_grad_norm: Tensor,
+):
+    import fused_adan
+
+    multi_tensor_applier = MultiTensorApply(2048 * 32)
+    _dummy_overflow_buf = torch.cuda.IntTensor([0])
+    multi_tensor_applier(
+        fused_adan.adan_multi_tensor,
+        _dummy_overflow_buf,
+        [params, grads, exp_avgs, exp_avg_sqs, exp_avg_diffs, neg_pre_grads],
+        beta1,
+        beta2,
+        beta3,
+        bias_correction1,
+        bias_correction2,
+        bias_correction3_sqrt,
+        lr,
+        weight_decay,
+        eps,
+        no_prox,
+        clip_global_grad_norm,
+    )
+    torch._foreach_zero_(neg_pre_grads)
+    torch._foreach_add_(neg_pre_grads, grads, alpha=-1.0)
+
+
+def _fused_adan_single_tensor(
+    params: List[Tensor],
+    grads: List[Tensor],
+    exp_avgs: List[Tensor],
+    exp_avg_sqs: List[Tensor],
+    exp_avg_diffs: List[Tensor],
+    neg_pre_grads: List[Tensor],
+    *,
+    beta1: float,
+    beta2: float,
+    beta3: float,
+    bias_correction1: float,
+    bias_correction2: float,
+    bias_correction3_sqrt: float,
+    lr: float,
+    weight_decay: float,
+    eps: float,
+    no_prox: bool,
+    clip_global_grad_norm: Tensor,
+):
+    for i, param in enumerate(params):
+        p_data_fp32 = param.data.float()
+        out_p = param.data
+        grad = grads[i]
+        exp_avg = exp_avgs[i]
+        exp_avg_sq = exp_avg_sqs[i]
+        exp_avg_diff = exp_avg_diffs[i]
+        neg_grad = neg_pre_grads[i]
+        with torch.cuda.device(param.device):
+            import fused_adan
+
+            fused_adan.adan_single_tensor(
+                p_data_fp32,
+                out_p,
+                grad,
+                exp_avg,
+                exp_avg_sq,
+                exp_avg_diff,
+                neg_grad,
+                beta1,
+                beta2,
+                beta3,
+                bias_correction1,
+                bias_correction2,
+                bias_correction3_sqrt,
+                lr,
+                weight_decay,
+                eps,
+                no_prox,
+                clip_global_grad_norm,
+            )
+        neg_grad.zero_().add_(grad, alpha=-1.0)

nemo/core/optim/distributed_adam.py

@@ -0,0 +1,853 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import collections
+import contextlib
+import itertools
+from typing import Callable, Dict, Iterable, Optional, Tuple, Union
+
+import torch
+from apex.contrib.optimizers.distributed_fused_adam import (
+    DistributedFusedAdam,
+    _disable_pre_forward_hook,
+    _multi_tensor_copy,
+)
+
+try:
+    import apex.contrib.nccl_allocator as nccl_allocator
+except ImportError:
+    nccl_allocator = None
+
+from megatron.core import parallel_state
+from megatron.core.dist_checkpointing.dict_utils import dict_list_map_inplace
+from megatron.core.dist_checkpointing.mapping import ShardedTensor
+from megatron.core.dist_checkpointing.optimizer import get_param_id_to_sharded_param_map, optim_state_to_sharding_state
+
+from nemo.utils import logging, str_to_dtype
+from nemo.utils.te_utils import is_float8tensor, is_mxfp8tensor, te_version
+
+if te_version() >= (2, 0):
+
+    # TE quantization logic using quantizer API
+    # Supported TE versions: 2.0+
+
+    from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor
+
+    def _quantize_param_fragment_impl(
+        input_: torch.Tensor,
+        *,
+        out: torch.Tensor,
+        param: torch.nn.Parameter,
+    ) -> None:
+        quantizer = param._quantizer
+        out = Float8Tensor(
+            shape=input_.size(),
+            dtype=param.dtype,
+            requires_grad=False,
+            data=out,
+            fp8_scale_inv=param._scale_inv,
+            fp8_dtype=param._fp8_dtype,
+            quantizer=quantizer,
+        )
+        quantizer.update_quantized(input_, out)
+
+    def _get_fp8_scale_and_amax_impl(tensor: Float8Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        quantizer = tensor._quantizer
+        return quantizer.scale, quantizer.amax
+
+elif te_version() >= (1, 0):
+
+    # TE quantization logic with fp8_meta dicts
+    # Supported TE versions: 1.0 - 1.14
+
+    from transformer_engine.pytorch.cpp_extensions import cast_to_fp8
+
+    def _quantize_param_fragment_impl(
+        input_: torch.Tensor,
+        *,
+        out: torch.Tensor,
+        param: torch.nn.Parameter,
+    ) -> None:
+        cast_to_fp8(
+            src.view(1, -1),
+            param._fp8_meta["scaling_fwd"],
+            param._fp8_meta_index,
+            param._fp8_dtype,
+            out=dst.view(1, -1),
+        )
+
+    def _get_fp8_scale_and_amax_impl(tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        fp8_meta = tensor._fp8_meta["scaling_fwd"]
+        fp8_meta_index = tensor._fp8_meta_index
+        return fp8_meta.scale[fp8_meta_index], fp8_meta.amax_history[0][fp8_meta_index]
+
+else:
+
+    # Fallback impl if TE version is invalid
+    def _quantize_param_fragment_impl(*args, **kwargs) -> None:
+        raise RuntimeError("Invalid Transformer Engine version for FP8 distributed optimizer")
+
+    def _get_fp8_scale_and_amax_impl(*args, **kwargs):
+        raise RuntimeError("Invalid Transformer Engine version for FP8 distributed optimizer")
+
+
+def quantize_param_fragment(
+    input_: torch.Tensor,
+    *,
+    out: torch.Tensor,
+    param: torch.nn.Parameter,
+) -> None:
+    """Cast values in parameter fragment to FP8
+
+    Arguments:
+      input_ (torch.Tensor): Values to quantize.
+      out (torch.Tensor): Raw UINT8 buffer to fill with FP8 values.
+          Dimensions should match input_.
+      param (torch.nn.Parameter): Parameter containing this parameter
+          fragment. Must be a Float8Tensor.
+
+    """
+    _quantize_param_fragment_impl(input_, out=out, param=param)
+
+
+def get_fp8_scale_and_amax(tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Get FP8 scale and amax from Float8Tensor"""
+    return _get_fp8_scale_and_amax_impl(tensor)
+
+
+_distributed_pgs = {}
+
+
+def create_distributed_pgs(*, distributed_size: int) -> Dict:
+    """Create process groups for distributing within multiple devices.
+
+    User can reuse this function to reorder communicators for SHArP.
+
+    Arguments:
+        distributed_size (int): the number of devices to distribute optimizer
+            state over.
+
+    """
+    global _distributed_pgs
+    assert torch.distributed.is_initialized()
+    if _distributed_pgs:
+        return _distributed_pgs
+
+    world_size = torch.distributed.get_world_size()
+    rank = torch.distributed.get_rank()
+    devices = distributed_size
+    nodes = world_size // devices
+
+    if nodes * devices != world_size:
+        logging.warning("Expected all nodes have the same amout of devices, disable distribute_within_nodes.")
+        return {}
+
+    node_id = rank // devices
+    device_id = rank % devices
+
+    distributed_pgs = []
+    for i in range(nodes):
+        ranks = [i * devices + j for j in range(devices)]
+        pg = torch.distributed.new_group(ranks=ranks)
+        distributed_pgs.append(pg)
+
+    redundant_pgs = []
+    for i in range(devices):
+        ranks = [i + j * devices for j in range(nodes)]
+        pg = torch.distributed.new_group(ranks=ranks)
+        redundant_pgs.append(pg)
+
+    # To re-order SHArP communicator right after distributed init,
+    # we have to expose redundant_process_group to user.
+    # User has too invoke allreduce through redundant_process_group
+    # before all other communicators to lock SHArP tree.
+    _distributed_pgs = {
+        'world_size': world_size,
+        'rank': rank,
+        'devices': devices,
+        'nodes': nodes,
+        'node_id': node_id,
+        'device_id': device_id,
+        'distributed_process_group': distributed_pgs[node_id],
+        'redundant_process_group': redundant_pgs[device_id],
+    }
+    return _distributed_pgs
+
+
+def create_distribute_within_nodes_pgs():
+    """Create process groups for distributing within nodes.
+
+    User can reuse this function to reorder communicators for SHArP.
+    This funcion is kept for backward compatibility.
+    """
+    return create_distributed_pgs(distributed_size=torch.cuda.device_count())
+
+
+class MegatronDistributedFusedAdam(DistributedFusedAdam):
+    """Adam optimizer with ZeRO algorithm
+
+    Child class of Apex DistributedFusedAdam, with optimizations for
+    NeMo-Megatron.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts
+            defining parameter groups.
+        disable_distributed_parameters (bool, optional): use standard
+            data-parallel communication instead of ZeRO.
+            (default: False)
+        distribute_within_nodes (bool, optional): distribute states
+            within the same node, e.g. DGX. This can improve performance
+            but requires larger memory than distributing within all
+            ranks, especially for pure data parallel models.
+            (default: False).
+        distributed_size (int, optional): the number of devices to
+            distribute optimizer state over.
+        lock_timeout (float, optional): timeout for callback mutex in
+            seconds.
+        **kwargs: keyword arguments to pass to Apex
+            DistributedFusedAdam.
+
+    """
+
+    def __init__(
+        self,
+        params: Union[Iterable[torch.nn.Parameter], Iterable[dict]],
+        disable_distributed_parameters: bool = False,
+        distribute_within_nodes: bool = False,
+        distributed_size: Optional[int] = None,
+        lock_timeout: Optional[float] = None,
+        **kwargs,
+    ):
+
+        # Update distributed_size settings
+        if distribute_within_nodes:
+            if distributed_size is not None and distributed_size != torch.cuda.device_count():
+                raise ValueError("Inconsistent distributed_size value")
+            distributed_size = torch.cuda.device_count()
+
+        # Initialize process groups
+        if 'process_group' not in kwargs and parallel_state.is_initialized():
+            kwargs['process_group'] = parallel_state.get_data_parallel_group(with_context_parallel=True)
+        if disable_distributed_parameters:
+            world_size = torch.distributed.get_world_size()
+            rank = torch.distributed.get_rank()
+            self_groups = [torch.distributed.new_group(ranks=[i]) for i in range(world_size)]
+            kwargs['distributed_process_group'] = self_groups[rank]
+            kwargs['redundant_process_group'] = kwargs['process_group']
+        elif distributed_size is not None:
+            dist_pg_infos = create_distributed_pgs(distributed_size=distributed_size)
+            if dist_pg_infos:
+                kwargs['distributed_process_group'] = dist_pg_infos['distributed_process_group']
+                kwargs['redundant_process_group'] = dist_pg_infos['redundant_process_group']
+                global _distributed_pgs
+                _distributed_pgs = {}
+
+        # Make sure dtypes are in right type
+        for keyword in ('dtype', 'grad_sync_dtype', 'param_sync_dtype'):
+            if keyword in kwargs:
+                kwargs[keyword] = str_to_dtype(kwargs[keyword])
+
+        # Make sure params are in consistent format (list of param group dicts)
+        param_groups = list(params)
+        assert param_groups
+        if not isinstance(param_groups[0], dict):
+            param_groups = [{'params': param_groups}]
+
+        # Construct distributed optimizer
+        super().__init__(param_groups, **kwargs)
+
+        # Create mutex with timeout
+        self._lock_with_timeout = None
+        if lock_timeout is not None:
+
+            @contextlib.contextmanager
+            def lock_with_timeout():
+                result = self._lock.acquire(timeout=lock_timeout)
+                try:
+                    yield result
+                finally:
+                    if result:
+                        # Acquired lock before timeout
+                        self._lock.release()
+                    else:
+                        # Failed to acquire lock before timeout
+                        print(f'MegatronDistributedFusedAdam: Failed to acquire lock within {lock_timeout} seconds.')
+
+            self._lock_with_timeout = lock_with_timeout
+
+        # Check for MXFP8 parameters
+        if any(is_mxfp8tensor(param) for param in self.parameters()):
+            raise ValueError("Distributed optimizer currently does not support MXFP8 parameters")
+
+    def _broadcast_params(self) -> None:
+        # Assume params have already been synchronized
+        pass
+
+    def _make_post_backward_hook(self, param: torch.nn.Parameter, param_group_id: int, param_id: int) -> Callable:
+        def hook(*unused):
+            if getattr(param, '_pre_forward_hook_is_enabled', False):
+                raise RuntimeError(
+                    'A parameter called its post-backward hook '
+                    'before its pre-forward hook. '
+                    'Please manually interact with the parameter '
+                    'before the forward pass (e.g. by calling data_ptr) '
+                    'or run DistributedFusedAdam with overlap_param_sync=False.'
+                )
+            lock = self._lock
+            if self._lock_with_timeout is not None:
+                lock = self._lock_with_timeout()
+            with lock:
+                need_to_initialize = 'fragments' not in self.state[param]
+                if need_to_initialize:
+                    self._init_param_state(param, param_group_id, param_id)
+                if self.greedy_grad_copy and not getattr(param, '_disable_greedy_grad_copy', False):
+                    self._grad_copy(param)
+                    if self.overlap_grad_sync and not getattr(param, '_disable_overlap_grad_sync', False):
+                        self._try_start_bucket_grad_sync(
+                            params=[param],
+                            ignore_last_bucket=need_to_initialize,
+                        )
+
+        return hook
+
+    def init_params(
+        self,
+        params: Optional[Iterable[torch.nn.Parameter]] = None,
+        param_sync_dtype: Optional[torch.dtype] = None,
+        **kwargs,
+    ) -> None:
+        """Initialize optimizer state for parameters
+
+        Initializes FP8 and non-FP8 params separately.
+
+        """
+
+        # Default cases
+        if params is None:
+            params = self.parameters()
+        elif isinstance(params, torch.Tensor):
+            params = [params]
+
+        # Ignore parameters that have already been initialized
+        params = [param for param in params if "fragments" not in self.state[param]]
+        if not params:
+            return
+
+        # Initialize FP8 and non-FP8 tensors separately
+        if any(is_float8tensor(param) for param in params):
+            super().init_params(
+                filter(is_float8tensor, params),
+                param_sync_dtype=torch.uint8,
+                **kwargs,
+            )
+        super().init_params(
+            params,
+            param_sync_dtype=param_sync_dtype,
+            **kwargs,
+        )
+
+    def init_params_bucket(
+        self,
+        params: Iterable[torch.nn.Parameter],
+        grad_sync_dtype: Optional[torch.dtype] = None,
+        param_sync_dtype: Optional[torch.dtype] = None,
+        **kwargs,
+    ) -> None:
+        """Initialize optimizer state for parameters in one effective bucket"""
+
+        # Ignore parameters that have already been initialized
+        if isinstance(params, torch.Tensor):
+            params = [params]
+        params = [param for param in params if "fragments" not in self.state[param]]
+        if not params:
+            return
+
+        # Initialize parameters with FP32 grads
+        fp32_params = []
+        remaining_params = []
+        for param in params:
+            if getattr(param, '_with_fp32_optimizer', False):
+                fp32_params.append(param)
+            else:
+                remaining_params.append(param)
+        params = remaining_params
+        start_bucket_id = len(self.state["buckets"])
+        super().init_params_bucket(
+            fp32_params,
+            grad_sync_dtype=torch.float32,
+            param_sync_dtype=param_sync_dtype,
+            **kwargs,
+        )
+        end_bucket_id = len(self.state["buckets"])
+        fp32_buckets = self.state["buckets"][start_bucket_id:end_bucket_id]
+
+        # Initialize FP8 parameters
+        fp8_params = []
+        remaining_params = []
+        for param in params:
+            if is_float8tensor(param):
+                fp8_params.append(param)
+            else:
+                remaining_params.append(param)
+        params = remaining_params
+        start_bucket_id = len(self.state["buckets"])
+        super().init_params_bucket(
+            fp8_params,
+            grad_sync_dtype=grad_sync_dtype,
+            param_sync_dtype=torch.uint8,
+            **kwargs,
+        )
+        end_bucket_id = len(self.state["buckets"])
+        fp8_buckets = self.state["buckets"][start_bucket_id:end_bucket_id]
+
+        # Initialize remaining parameters as usual
+        normal_buckets = []
+        start_bucket_id = len(self.state["buckets"])
+        super().init_params_bucket(
+            params,
+            grad_sync_dtype=grad_sync_dtype,
+            param_sync_dtype=param_sync_dtype,
+            **kwargs,
+        )
+        end_bucket_id = len(self.state["buckets"])
+        normal_buckets = self.state["buckets"][start_bucket_id:end_bucket_id]
+
+        def add_param_to_bucket(
+            param: torch.nn.Parameter,
+            bucket: self.StateBucket,
+        ) -> None:
+            """Add trivial param fragment to bucket"""
+            param_fragments = self.state[param]["fragments"]
+            param_group_id = param_fragments[0].param_group_id
+            param_id = param_fragments[0].param_id
+            bucket_id = bucket.fragments[0].bucket_id
+            param_size = param.numel()
+            bucket_size = bucket.bucket_size
+            fragment = self.ParameterFragment(
+                param_group_id=param_group_id,
+                param_id=param_id,
+                bucket_id=bucket_id,
+                param_range=(param_size, param_size),
+                bucket_range=(bucket_size, bucket_size),
+                in_local_shard=False,
+                shard_range=None,
+                shard_bucket_range=None,
+                shard_param_range=None,
+            )
+            param_fragments.append(fragment)
+            bucket.fragments.append(fragment)
+
+        # Make sure all added buckets depend on provided params
+        for bucket in fp32_buckets:
+            for param in itertools.chain(fp8_params, params):
+                add_param_to_bucket(param, bucket)
+        for bucket in fp8_buckets:
+            for param in itertools.chain(fp32_params, params):
+                add_param_to_bucket(param, bucket)
+        for bucket in normal_buckets:
+            for param in itertools.chain(fp32_params, fp8_params):
+                add_param_to_bucket(param, bucket)
+
+    def _init_param_state(
+        self,
+        param: torch.nn.Parameter,
+        param_group_id: int,
+        param_id: int,
+        param_sync_dtype: Optional[torch.dtype] = None,
+        **kwargs,
+    ) -> None:
+        """Initialize optimizer state for a parameter
+
+        Initializing the master weights requires slicing a flattened
+        view of the param. FP8 tensors do not handle these operations
+        gracefully, so we hack around it by explicitly casting to
+        FP32.
+
+        """
+
+        # Initialize non-FP8 params as usual
+        if not is_float8tensor(param):
+            super()._init_param_state(
+                param,
+                param_group_id,
+                param_id,
+                param_sync_dtype=param_sync_dtype,
+                **kwargs,
+            )
+
+        # Return immediately if already initialized
+        if "fragments" in self.state[param]:
+            return
+
+        # Initialize with FP32 copy of param
+        fp32_param = param.float()
+        super()._init_param_state(
+            fp32_param,
+            param_group_id,
+            param_id,
+            param_sync_dtype=torch.uint8,
+            **kwargs,
+        )
+        self.state[param].update(self.state[fp32_param])
+        del self.state[fp32_param]
+
+    @torch.no_grad()
+    def init_param_buffer(self) -> None:
+        """Allocate contiguous buffers for param buckets
+
+        For FP8 params, the FP8 data buffer is made a view into a
+        contiguous buffer.
+
+        """
+
+        # Make sure all params are initialized
+        self.contiguous_param_buffer = True
+        self.init_params()
+
+        # Construct param buffers
+        buffer_sizes = collections.defaultdict(lambda: 0)
+        for bucket in self.state["buckets"]:
+            dtypes = bucket.dtypes()
+            buffer_sizes[dtypes] = max(bucket.contiguous_buffer_offset + bucket.bucket_size, buffer_sizes[dtypes])
+        for dtypes, buffer_size in buffer_sizes.items():
+            _, _, param_sync_dtype = dtypes
+            if getattr(self, "nccl_ub", False):
+                if not nccl_allocator:
+                    raise RuntimeError("NCCL allocator importing failed but nccl ub is still requested")
+                with nccl_allocator.nccl_mem():
+                    self._param_buffers[dtypes] = torch.zeros(
+                        [buffer_size], dtype=param_sync_dtype, device=self.device
+                    )
+            else:
+                self._param_buffers[dtypes] = torch.zeros([buffer_size], dtype=param_sync_dtype, device=self.device)
+        # Figure out corresponding positions in params and param buffer
+        params = list(self.parameters())
+        param_flat_views = []
+        param_buffer_views = []
+        for i, param in enumerate(params):
+            fragment = self.state[param]["fragments"][0]
+            bucket_id = fragment.bucket_id
+            bucket = self.state["buckets"][bucket_id]
+            param_size = param.numel()
+            bucket_start, _ = fragment.bucket_range
+            buffer_offset = bucket.contiguous_buffer_offset
+            buffer_start = buffer_offset + bucket_start
+            buffer_end = buffer_start + param_size
+            param_buffer = self._param_buffers[bucket.dtypes()]
+            param_buffer_view = param_buffer[buffer_start:buffer_end].detach()
+            if param_buffer_view.device != param.device:
+                raise RuntimeError(
+                    "Attempted to change a parameter with device={param.device} "
+                    f"into a buffer view with device={param_buffer_view.device}"
+                )
+            if is_float8tensor(param):
+                param_flat_views.append(param._data.detach().view(-1))
+            else:
+                if param_buffer_view.dtype != param.dtype:
+                    raise RuntimeError(
+                        f"Attempted to change a parameter with dtype={param.dtype} "
+                        f"into a buffer view with dtype={param_buffer_view.dtype}"
+                    )
+                if param.is_contiguous(memory_format=torch.channels_last):
+                    param = param.permute(0, 2, 3, 1)
+                param_flat_views.append(param.detach().view(-1))
+            param_buffer_views.append(param_buffer_view)
+
+        # Copy values into param buffer
+        _multi_tensor_copy(
+            param_flat_views,
+            param_buffer_views,
+            dummy_overflow_buf=self._dummy_overflow_buf,
+        )
+
+        # Make all params a view into the param buffer
+        for param, buffer_view in zip(params, param_buffer_views):
+            if is_float8tensor(param):
+                param._data = buffer_view.view(param.size())
+            else:
+                # Preserve memory format for param here, i.e. NHWC tensors
+                # `param.data.set_()` failed to change storage.
+                # `param.set_()` invalidates bprop hook.
+                param.data = torch.as_strided(
+                    buffer_view,
+                    param.size(),
+                    param.stride(),
+                    storage_offset=buffer_view.storage_offset(),
+                )
+
+    def try_grad_sync(self, params: Iterable[torch.nn.Parameter]) -> None:
+        """Attempt to launch gradient synchronization"""
+
+        def is_grad_copy_enabled(param: torch.nn.Parameter) -> bool:
+            return not getattr(param, '_disable_greedy_grad_copy', False) and not getattr(
+                param, '_disable_overlap_grad_sync', False
+            )
+
+        params = list(filter(is_grad_copy_enabled, params))
+        for p in params:
+            self._grad_copy(p)
+        self._try_start_bucket_grad_sync(params=params)
+
+    def zero_grad(self, *args, **kwargs) -> None:
+        """Clear parameter gradients"""
+        super().zero_grad(*args, **kwargs)
+
+        # Reset main grads
+        if self.contiguous_grad_buffer:
+            for param in self.parameters():
+                with _disable_pre_forward_hook(param):
+                    param.main_grad = self.grad_buffer_view(param)
+
+    def grad_norm(
+        self,
+        parameters: Optional[Iterable[torch.nn.Parameter]] = None,
+        norm_type: float = 2.0,
+        force: bool = False,
+    ) -> torch.Tensor:
+        """L2 norm of parameter gradients"""
+        assert norm_type == 2
+
+        if parameters is not None:
+            # Make sure we can access iterable multiple times
+            parameters = list(parameters)
+
+        # Compute grad norm
+        if force or self._grad_norm is None:
+
+            # Compute norm of local gradients for distributed optimizer
+            grad_norm_sq = self._local_grad_norm(parameters=parameters, norm_type=norm_type)
+            if self.redundant_size > 1:
+                grad_norm_sq /= self.redundant_size
+
+            # Sum over all procs to get grad norm
+            torch.distributed.all_reduce(
+                grad_norm_sq,
+                op=torch.distributed.ReduceOp.SUM,
+            )
+            self._grad_norm = grad_norm_sq.sqrt()
+
+        # Use cached grad norm
+        return super().grad_norm()
+
+    @torch.no_grad()
+    def _param_copy_fragments(self, fragments: Iterable[DistributedFusedAdam.ParameterFragment]) -> None:
+        """Update parameter fragments with values from parameter buckets
+
+        For FP8 params, values are copied directly into the FP8 data
+        buffer.
+
+        """
+
+        # Figure out corresponding positions in param buckets and params
+        buffers_in = []
+        buffers_out = []
+        fragments = list(fragments)
+        for fragment in fragments:
+
+            # Check if fragment needs to be updated
+            bucket_id = fragment.bucket_id
+            bucket_start, bucket_end = fragment.bucket_range
+            param_start, param_end = fragment.param_range
+            if param_end <= param_start or bucket_id not in self._params_buckets:
+                continue
+
+            # Corresponding positions in bucket and param
+            param_bucket = self._params_buckets[bucket_id]
+            param = self.parameter(fragment)
+            buffer_in = param_bucket.params_bucket[bucket_start:bucket_end]
+            if is_float8tensor(param):
+                # Copy into FP8 params's data buffer
+                assert (
+                    param_bucket.params_bucket.dtype == torch.uint8
+                ), "Expected FP8 params to perform param sync in UINT8"
+                buffer_out = param._data.view(-1)[param_start:param_end]
+                buffers_in.append(buffer_in)
+                buffers_out.append(buffer_out)
+            elif torch.is_floating_point(buffer_in) and torch.is_floating_point(param):
+                # Conv with NHWC layout, i.e. shape (N, C, H, W) and stride
+                # (HWC, 1, WC, C), can't `.view(-1)`. Here to turn it to
+                # tensor with shape (N, H, W, C) and stride (HWC, WC, C, 1).
+                # Note: https://github.com/NVIDIA/apex/pull/1794
+                if param.is_contiguous(memory_format=torch.channels_last):
+                    param = param.permute(0, 2, 3, 1)
+
+                # Cast between floating-point dtypes
+                buffer_out = param.detach().view(-1)[param_start:param_end]
+                buffers_in.append(buffer_in)
+                buffers_out.append(buffer_out)
+            else:
+                # Copy most significant bytes for non-floating-point
+                # dtypes
+                # Note: Assume dtypes are little-endian
+                buffer_out = param.detach().view(-1)[param_start:param_end]
+                in_bytes = buffer_in.unsqueeze(-1).view(torch.uint8)
+                out_bytes = buffer_out.unsqueeze(-1).view(torch.uint8)
+                copy_size = min(in_bytes.size(-1), out_bytes.size(-1))
+                buffers_in.append(in_bytes[..., -copy_size:])
+                buffers_out.append(out_bytes[..., -copy_size:])
+                if copy_size < out_bytes.size(-1):
+                    out_bytes[..., :-copy_size].zero_()
+
+        # Copy data from parameter buckets to parameters
+        _multi_tensor_copy(
+            buffers_in,
+            buffers_out,
+            dummy_overflow_buf=self._dummy_overflow_buf,
+        )
+
+        # Update transpose caches
+        params = set(self.parameter(fragment) for fragment in fragments)
+        for param in params:
+            if is_float8tensor(param):
+                param._reset_caches()
+
+    @torch.no_grad()
+    def _check_params_shard_dtypes(self, params_buckets: Dict[int, DistributedFusedAdam.ParameterBucket]) -> None:
+        """Make sure local shards of parameters are in expected datatypes
+
+        For FP8 params, FP32 values are cast into FP8 using per-param
+        scaling factors and per-param amaxes are computed and reduced.
+
+        """
+
+        # Just call base class function if there are no FP8 tensors
+        num_fp8_params = sum(1 for param in self.parameters() if is_float8tensor(param))
+        if num_fp8_params == 0:
+            super()._check_params_shard_dtypes(params_buckets)
+            return
+
+        # Cast local data to FP8
+        fp8_params_shards = dict()
+        for bucket_id, param_bucket in params_buckets.items():
+            state_bucket = self.state["buckets"][bucket_id]
+            if state_bucket.param_sync_dtype != torch.uint8:
+                continue
+
+            # Initialize FP8 buffer for param sync
+            params_shard = param_bucket.params_shard
+            if self.contiguous_param_buffer:
+                shard_size = state_bucket.shard_size
+                buffer_offset = state_bucket.contiguous_buffer_offset
+                buffer_start = buffer_offset + self.distributed_rank * shard_size
+                buffer_end = buffer_start + shard_size
+                param_buffer = self._param_buffers[state_bucket.dtypes()]
+                fp8_params_shard = param_buffer[buffer_start:buffer_end]
+            else:
+                fp8_params_shard = torch.empty_like(params_shard, dtype=torch.uint8)
+            param_bucket.params_shard = fp8_params_shard
+
+            # Cast param fragments to FP8
+            for fragment in self.state["buckets"][bucket_id].fragments:
+                param = self.parameter(fragment)
+                if not is_float8tensor(param):
+                    continue
+                if not fragment.in_local_shard:
+                    continue
+                shard_start, shard_end = fragment.shard_range
+                if shard_end <= shard_start:
+                    continue
+                shard_range = slice(shard_start, shard_end)
+                quantize_param_fragment(
+                    params_shard[shard_range],
+                    out=fp8_params_shard[shard_range],
+                    param=param,
+                )
+
+        # Update FP8 scaling factors when all buckets have processed
+        if getattr(self, "_check_params_shard_dtypes_progress", None) is None:
+            self._check_params_shard_dtypes_progress = []
+        self._check_params_shard_dtypes_progress.extend(params_buckets.keys())
+        if len(self._check_params_shard_dtypes_progress) == len(self.state["buckets"]):
+            assert len(set(self._check_params_shard_dtypes_progress)) == len(self.state["buckets"])
+
+            # FP8 scaling factors
+            amaxes = []
+            scales = []
+            scale_invs = []
+            i = -1
+            for param in self.parameters():
+                if not is_float8tensor(param):
+                    continue
+                i += 1
+                scale, amax = get_fp8_scale_and_amax(param)
+                amaxes.append(amax.view(1))
+                scales.append(scale.view(1))
+                scale_invs.append(param._scale_inv.view(1))
+
+            # Update cached scale-inverses
+            packed_scales = torch.empty(num_fp8_params, dtype=torch.float32, device=self.device)
+            packed_scale_views = [packed_scales[i].view(1) for i in range(num_fp8_params)]
+            _multi_tensor_copy(
+                scales,
+                packed_scale_views,
+                dummy_overflow_buf=self._dummy_overflow_buf,
+            )
+            torch.reciprocal(packed_scales, out=packed_scales)
+            _multi_tensor_copy(
+                packed_scale_views,
+                scale_invs,
+                dummy_overflow_buf=self._dummy_overflow_buf,
+            )
+
+            # Reduce amaxes
+            # Note: Assume each param has a separate amax
+            packed_amaxes = torch.empty(num_fp8_params, dtype=torch.float32, device=self.device)
+            packed_amax_views = [packed_amaxes[i].view(1) for i in range(num_fp8_params)]
+            _multi_tensor_copy(
+                amaxes,
+                packed_amax_views,
+                dummy_overflow_buf=self._dummy_overflow_buf,
+            )
+            torch.distributed.all_reduce(
+                packed_amaxes,
+                op=torch.distributed.ReduceOp.MAX,
+                group=self.distributed_process_group,
+            )
+            _multi_tensor_copy(
+                packed_amax_views,
+                amaxes,
+                dummy_overflow_buf=self._dummy_overflow_buf,
+            )
+
+            # Reset
+            self._check_params_shard_dtypes_progress = None
+
+        # Handle any remaining dtype conversions
+        super()._check_params_shard_dtypes(params_buckets)
+
+    def sharded_state_dict(self, model_sharded_state_dict, optimizer_state_dict=None):
+        """Create sharded state dict"""
+        if optimizer_state_dict is None:
+            optimizer_state_dict = self.state_dict()
+
+        id_to_sharded_param_map = get_param_id_to_sharded_param_map(
+            model_sharded_state_dict=model_sharded_state_dict,
+            optim_params_iter=self.parameters(),
+        )
+        # Convert state
+        step = optimizer_state_dict['state'].pop('step')
+        state_dict_format = optimizer_state_dict.pop('format', None)
+        optim_state_to_sharding_state(optimizer_state_dict, id_to_sharded_param_map)
+        optimizer_state_dict['state']['step'] = step
+        if state_dict_format is not None:
+            optimizer_state_dict['format'] = state_dict_format
+
+        def rename_fp32_params(x):
+            if isinstance(x, ShardedTensor) and x.key.startswith('optimizer.state.param'):
+                x.key = x.key.replace('optimizer.state.param', 'optimizer.state.fp32_param')
+            return x
+
+        dict_list_map_inplace(rename_fp32_params, optimizer_state_dict)
+
+        return optimizer_state_dict

nemo/core/optim/flash_optim.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import torch
+
+from nemo.utils import logging
+
+
+def patch_flashoptim_uneven_shard_support(optimizer) -> None:
+    """Patch flashoptim to handle FSDP2 unevenly-sharded parameters in DCP.
+
+    FlashOptim <= 0.1.3 raises ``ValueError`` when saving optimizer state for
+    parameters whose shard dimension is not evenly divisible by the FSDP mesh
+    size. The root cause is that ``DTensor.from_local()`` is called without an
+    explicit ``shape``, so it infers ``global = local * mesh_size`` which is
+    wrong for padded (uneven) shards.
+
+    This patch replaces ``_wrap_state_as_dtensor`` on the optimizer class so
+    that ``shape=param.shape`` and ``stride=param.stride()`` are always passed,
+    which is correct for both even and uneven shards.
+    """
+
+    klass = type(optimizer)
+    if not hasattr(klass, "_wrap_state_as_dtensor"):
+        return
+    if getattr(klass, "_nemo_patched_uneven_shard", False):
+        return
+
+    @staticmethod
+    def _fixed_wrap_state_as_dtensor(state: dict[str, Any], param: torch.Tensor) -> None:  # noqa: UP006
+        if not hasattr(param, "device_mesh"):
+            return
+
+        from torch.distributed.tensor import DTensor
+
+        mesh = param.device_mesh
+        placements = param.placements
+
+        for key, val in state.items():
+            if isinstance(val, torch.Tensor) and not isinstance(val, DTensor) and val.dim() > 0:
+                state[key] = DTensor.from_local(
+                    val,
+                    mesh,
+                    placements,
+                    shape=param.shape,
+                    stride=param.stride(),
+                )
+
+    klass._wrap_state_as_dtensor = _fixed_wrap_state_as_dtensor
+    klass._nemo_patched_uneven_shard = True
+    logging.info("Patched flashoptim %s to support unevenly-sharded FSDP2 parameters in DCP.", klass.__name__)

nemo/core/optim/novograd.py

@@ -0,0 +1,150 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch.optim.optimizer import Optimizer
+
+__all__ = ['Novograd']
+
+
+def _check_valid_opt_params(lr, eps, betas):
+    if lr < 0:
+        raise ValueError(f"Invalid learning rate: {lr}")
+    if eps < 0:
+        raise ValueError(f"Invalid epsilon value: {eps}")
+    if not (0.0 <= betas[0] < 1.0 and 0.0 <= betas[1] < 1.0):
+        raise ValueError(f"Betas have to be between 0 and 1: {betas}")
+
+
+class Novograd(Optimizer):
+    """Implements Novograd algorithm.
+    It has been proposed  in "Stochastic Gradient Methods with Layer-wise
+    Adaptive Moments for Training of Deep Networks"
+    (https://arxiv.org/abs/1905.11286)
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper "On the Convergence of Adam and Beyond"
+    """
+
+    def __init__(
+        self,
+        params,
+        lr=1e-3,
+        betas=(0.95, 0.98),
+        eps=1e-8,
+        weight_decay=0,
+        grad_averaging=False,
+        amsgrad=False,
+        luc=False,
+        luc_trust=1e-3,
+        luc_eps=1e-8,
+    ):
+        _check_valid_opt_params(lr, eps, betas)
+        defaults = dict(
+            lr=lr,
+            betas=betas,
+            eps=eps,
+            weight_decay=weight_decay,
+            grad_averaging=grad_averaging,
+            amsgrad=amsgrad,
+        )
+        self.luc = luc
+        self.luc_trust = luc_trust
+        self.luc_eps = luc_eps
+        super(Novograd, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(Novograd, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault("amsgrad", False)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError("Sparse gradients are not supported.")
+                amsgrad = group["amsgrad"]
+                state = self.state[p]
+
+                # State initialization
+                if not state:
+                    state["step"] = 0
+                    # Exponential moving average of gradient values
+                    state["exp_avg"] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state["exp_avg_sq"] = torch.zeros([]).to(state["exp_avg"].device)
+                    if amsgrad:
+                        # Maintains max of all exp moving avg of squared grad
+                        state["max_exp_avg_sq"] = torch.zeros([]).to(state["exp_avg"].device)
+
+                exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+                if amsgrad:
+                    max_exp_avg_sq = state["max_exp_avg_sq"]
+                beta1, beta2 = group["betas"]
+
+                state["step"] += 1
+
+                norm = grad.norm().pow(2)
+
+                if exp_avg_sq == 0:
+                    exp_avg_sq.copy_(norm)
+                else:
+                    exp_avg_sq.mul_(beta2).add_(norm, alpha=1.0 - beta2)
+
+                if amsgrad:
+                    # Maintains max of all 2nd moment running avg till now
+                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    # Use the max for normalizing running avg. of gradient
+                    denom = max_exp_avg_sq.sqrt().add_(group["eps"])
+                else:
+                    denom = exp_avg_sq.sqrt().add_(group["eps"])
+
+                grad.div_(denom)
+                if group["weight_decay"] != 0:
+                    grad.add_(p.data, alpha=group["weight_decay"])
+                if group["grad_averaging"]:
+                    grad.mul_(1 - beta1)
+                exp_avg.mul_(beta1).add_(grad)
+
+                if self.luc:
+                    # Clip update so that updates are less than eta*weights
+                    data_norm = torch.norm(p.data)
+                    grad_norm = torch.norm(exp_avg.data)
+                    luc_factor = self.luc_trust * data_norm / (grad_norm + self.luc_eps)
+                    luc_factor = min(luc_factor, group["lr"])
+                    p.data.add_(exp_avg, alpha=-luc_factor)
+                else:
+                    p.data.add_(exp_avg, alpha=-group["lr"])
+
+        return loss

nemo/core/optim/optimizer_with_main_params.py

@@ -0,0 +1,566 @@
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from contextlib import contextmanager
+
+import torch
+
+from nemo.utils import logging
+
+try:
+    import amp_C
+    from apex.multi_tensor_apply import multi_tensor_applier
+
+    HAVE_APEX = True
+
+except (ImportError, ModuleNotFoundError):
+
+    HAVE_APEX = False
+
+try:
+    from megatron.core.parallel_state import (
+        get_data_parallel_group,
+        get_data_parallel_world_size,
+        get_expert_data_parallel_group,
+    )
+    from megatron.core.tensor_parallel import copy_tensor_model_parallel_attributes
+
+    HAVE_MEGATRON_CORE = True
+
+except (ImportError, ModuleNotFoundError):
+
+    HAVE_MEGATRON_CORE = False
+
+
+def _zero_grad_group_helper(group, set_to_none):
+    """Zero out the gradient for a group of parameters.
+    Note: copied from torch.optim.optimizer."""
+    for param in group:
+        if param.grad is not None:
+            if set_to_none:
+                param.grad = None
+            else:
+                if param.grad.grad_fn is not None:
+                    param.grad.detach_()
+                else:
+                    param.grad.requires_grad_(False)
+                param.grad.zero_()
+
+
+def _multi_tensor_copy_this_to_that(this, that, overflow_buf):
+    """Use multi-tensor-applier to copy values from one list to another.
+    We don't have a blfoat16 implementation so for now if the overflow_buf
+    is not provided, we default back to simple loop copy to be compatible
+    with bfloat16."""
+    if overflow_buf:
+        # Scaling with factor `1.0` is equivalent to copy.
+        multi_tensor_applier(amp_C.multi_tensor_scale, overflow_buf, [this, that], 1.0)
+    else:
+        # FIXME: use multi-tensor applier for bf16
+        for this_, that_ in zip(this, that):
+            that_.copy_(this_)
+
+
+def _get_grad_data_group(is_expert_group):
+    if is_expert_group:
+        data_group = get_expert_data_parallel_group()
+    else:
+        data_group = get_data_parallel_group(with_context_parallel=True)
+    return data_group
+
+
+class GradBucket(object):
+    """
+    Persistent buffer for main gradients that remains allocated between training iterations
+    """
+
+    def __init__(self, numel, chunk_size_mb, data_group):
+        if not HAVE_APEX:
+            raise ImportError(
+                "Apex was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+            )
+
+        if not HAVE_MEGATRON_CORE:
+            raise ImportError(
+                "megatron-core was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+            )
+
+        self.numel = numel
+        self.data = torch.zeros(self.numel, dtype=torch.float, device=torch.cuda.current_device(), requires_grad=False)
+
+        self._data_group = data_group
+        self.chunk_size_mb = chunk_size_mb
+        if self.chunk_size_mb > 0:
+            chunk_size_bytes = chunk_size_mb * 1024 * 1024
+            self.chunk_size_numel = chunk_size_bytes // 4
+            self.num_chunks = self.numel // self.chunk_size_numel
+            self.numel_per_chunk = [self.chunk_size_numel] * self.num_chunks
+            if self.numel % self.chunk_size_numel != 0:
+                self.num_chunks += 1
+                self.numel_per_chunk.append(self.numel % self.chunk_size_numel)
+
+            self.start_index_per_chunk = torch.cumsum(torch.tensor([0] + self.numel_per_chunk[:-1]), dim=0)
+            self.current_chunk = 0
+            self.computed_numel_per_chunk = [0] * self.num_chunks
+
+    def zero(self):
+        """Reset the buffer to zero."""
+        self.data.zero_()
+
+    def allreduce_buffer(self):
+        """Synchronous buffer data allreduce"""
+        self.data.div_(get_data_parallel_world_size())
+        torch.distributed.all_reduce(self.data, group=self._data_group)
+
+    def get(self, shape, start_index):
+        """Return a tensor with the input `shape` as a view into the
+        1-D data starting at `start_index`."""
+        end_index = start_index + shape.numel()
+        assert end_index <= self.numel, 'requested tensor is out of the buffer range.'
+        buffer_tensor = self.data[start_index:end_index]
+        buffer_tensor = buffer_tensor.view(shape)
+
+        grad_chunk_info = None
+        if self.chunk_size_mb > 0:
+            grad_chunk_info = {}
+            chunk = start_index // self.chunk_size_numel
+            chunk_start_index = self.start_index_per_chunk[chunk]
+            chunk_end_index = chunk_start_index + self.numel_per_chunk[chunk]
+            grad_chunk_info[chunk] = min(chunk_end_index, end_index) - start_index
+            while chunk_end_index < end_index:
+                chunk += 1
+                chunk_start_index = self.start_index_per_chunk[chunk]
+                chunk_end_index = chunk_start_index + self.numel_per_chunk[chunk]
+                grad_chunk_info[chunk] = min(chunk_end_index, end_index) - chunk_start_index
+
+        return buffer_tensor, grad_chunk_info
+
+    def update_chunk_info(self, grad_chunk_info):
+        for chunk in grad_chunk_info.keys():
+            self.computed_numel_per_chunk[chunk] += grad_chunk_info[chunk]
+
+    def get_allreduce_tensor(self):
+        if self.computed_numel_per_chunk[self.current_chunk] == self.numel_per_chunk[self.current_chunk]:
+            chunk_start_index = self.start_index_per_chunk[self.current_chunk]
+            chunk_end_index = chunk_start_index + self.numel_per_chunk[self.current_chunk]
+            allreduce_tensor = self.data[chunk_start_index:chunk_end_index]
+
+            self.computed_numel_per_chunk[self.current_chunk] = 0
+            self.current_chunk += 1
+            if self.current_chunk == self.num_chunks:
+                self.current_chunk = 0
+
+            return allreduce_tensor
+
+        return None
+
+
+class MainParamsOptimizerWrapper(torch.optim.Optimizer):
+    """
+    Float16 optimizer wrapper for half precision (fp16 and bf16) data types.
+    This optimizer wrapper holds main parameters and gradients in fp32 to support
+    stable convergence.
+
+    Arguments:
+        optimizer: base optimizer such as Adam or SGD.
+        fp32_grad_accum: to enable the use of fp32 in gradient accumulation and allreduce.
+        contiguous_grad_bucket: to enable allocating the master gradients in the
+            contiguous memory space to reduce memory fragmentation.
+        async_grad_allreduce: enable asynchronous gradient allreduce that is executed
+            along with the training step backprop.
+    """
+
+    def __init__(
+        self,
+        optimizer,
+        fp32_grad_accum=False,
+        contiguous_grad_bucket=False,
+        async_grad_allreduce=False,
+        grad_div_ar_fusion=True,
+        grad_allreduce_chunk_size_mb=0,
+    ):
+        if not HAVE_APEX:
+            raise ImportError(
+                "Apex was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+            )
+
+        if not HAVE_MEGATRON_CORE:
+            raise ImportError(
+                "megatron-core was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+            )
+
+        self.optimizer = optimizer
+        assert self.optimizer, 'no optimizer is provided.'
+        if contiguous_grad_bucket:
+            assert fp32_grad_accum, 'contiguous gradient buffer assumes using fp32 grad.'
+        if async_grad_allreduce:
+            assert fp32_grad_accum, (
+                'async allreduce applies to master gradients only, '
+                'which is supposed to be accumulated after grad op.'
+            )
+            assert contiguous_grad_bucket, (
+                'currently async_grad_allreduce is supported only ' 'with contiguous_grad_bucket.'
+            )
+
+        self._fp32_grad_accum = fp32_grad_accum
+        self._contiguous_grad_bucket = contiguous_grad_bucket
+
+        # used with tensor parallel only (no pipeline parallelism)
+        # be careful, weight update cannot start until all async grad AR works are done
+        self._async_grad_allreduce = (
+            async_grad_allreduce and get_data_parallel_world_size(with_context_parallel=True) > 1
+        )
+
+        if self._async_grad_allreduce:
+            # use @no_sync to disable backward grad sync during gradient accumulation
+            self._require_backward_grad_sync = True
+            self._grad_div_ar_fusion = grad_div_ar_fusion
+            self._grad_allreduce_chunk_size_mb = grad_allreduce_chunk_size_mb
+        else:
+            self._require_backward_grad_sync = False
+            self._grad_div_ar_fusion = False
+            self._grad_allreduce_chunk_size_mb = 0
+
+        # Dummy tensor needed for apex multi-apply tensor.
+        self._dummy_overflow_buf = None
+
+        # Create persistent buffers for main gradients in contiguous memory space
+        # - Chunked element-wise and allreduce ops without creating a temporary buffer for merged operation
+        # - Low memory fragmentation
+        self._main_grad_buffers = None
+        if self._contiguous_grad_bucket:
+            self._main_grad_buffers = {}
+            # get the size of buffers
+            num_elements = {}
+            for i, param_group in enumerate(self.optimizer.param_groups):
+                num_elements[i] = sum(
+                    map(lambda x: x.data.nelement(), filter(lambda p: p.requires_grad, param_group['params']))
+                )
+
+                # Allocate gradient memory buffers for each data type
+                if num_elements[i] > 0:
+                    self._main_grad_buffers[i] = GradBucket(
+                        num_elements[i],
+                        self._grad_allreduce_chunk_size_mb,
+                        _get_grad_data_group(param_group.get('is_expert', False)),
+                    )
+
+        # Three groups of parameters:
+        self.float16_groups = []  # original float16 parameters
+        self.fp32_from_float16_groups = []  # fp32 copy of float16 parameters
+        self.fp32_from_fp32_groups = []  # original fp32 parameters
+
+        # gradient function hooks
+        if self._fp32_grad_accum:
+            self.grad_accs = []
+
+        # For all the groups in the original optimizer:
+        for i, param_group in enumerate(self.optimizer.param_groups):
+            float16_params_this_group = []
+            fp32_params_this_group = []
+            fp32_from_float16_params_this_group = []
+            # For all the parameters in this group:
+            is_expert_group = param_group.get('is_expert', False)
+            for j, param in enumerate(param_group['params']):
+                main_param = None
+                if param.requires_grad:
+                    # float16 params:
+                    if param.type() in ['torch.cuda.HalfTensor', 'torch.cuda.BFloat16Tensor']:
+                        float16_params_this_group.append(param)
+
+                        # Allocate the main parameter
+                        main_param = param.detach().clone().float()
+
+                        # Copy tensor model parallel attributes.
+                        copy_tensor_model_parallel_attributes(main_param, param)
+                        if hasattr(param, 'shared'):
+                            main_param.shared = param.shared
+                        if hasattr(param, 'allreduce'):
+                            main_param.allreduce = param.allreduce
+
+                        # Assign the grad buffer offset to main parameters
+                        if self._contiguous_grad_bucket:
+                            num_elements[i] -= param.data.nelement()
+                            main_param.grad, grad_chunk_info = self._main_grad_buffers[i].get(
+                                param.data.shape, num_elements[i]
+                            )
+                            # Add a pointer to main_grad in model param for first-last stage embedding param reduction
+                            param.main_grad = main_param.grad
+
+                        # Replace the optimizer params with the new fp32 copy.
+                        param_group['params'][j] = main_param
+                        fp32_from_float16_params_this_group.append(main_param)
+                        # Reset existing state dict key to the new main param.
+                        if param in self.optimizer.state:
+                            self.optimizer.state[main_param] = self.optimizer.state.pop(param)
+                    # fp32 params.
+                    elif param.type() == 'torch.cuda.FloatTensor':
+                        fp32_params_this_group.append(param)
+                        param_group['params'][j] = param
+
+                    else:
+                        raise TypeError(
+                            'Wrapped parameters must be one of '
+                            'torch.cuda.FloatTensor,  '
+                            'torch.cuda.HalfTensor, or '
+                            'torch.cuda.BFloat16Tensor. '
+                            'Received {}'.format(param.type())
+                        )
+
+                # Add gradient accumulation hook for fp32 grad accumulation
+                if main_param is not None and self._fp32_grad_accum and param.requires_grad:
+                    # Expand so we get access to grad_fn
+                    param_tmp = param.expand_as(param)
+                    # Get the gradient accumulator function.
+                    grad_acc = param_tmp.grad_fn.next_functions[0][0]
+                    grad_acc.register_hook(
+                        self._make_param_hook(param, main_param, i, grad_chunk_info, is_expert_group)
+                    )
+                    self.grad_accs.append(grad_acc)
+
+            self.float16_groups.append(float16_params_this_group)
+            self.fp32_from_float16_groups.append(fp32_from_float16_params_this_group)
+            self.fp32_from_fp32_groups.append(fp32_params_this_group)
+
+        # Leverage state_dict() and load_state_dict() to
+        # recast preexisting per-param state tensors
+        self.optimizer.load_state_dict(self.optimizer.state_dict())
+
+    def _make_param_hook(self, param, main_param, i, grad_chunk_info, is_expert_group):
+        """Create the grad accumulation and all-reduce hook for backprop."""
+
+        # Hook used for back-prop.
+        def param_hook(*unused):
+            # Accumulates gradients on main gradients
+            if param.grad is not None:
+                if main_param.grad is None:
+                    main_param.grad = param.grad.float()
+                else:
+                    main_param.grad.add_(param.grad.data)
+                # Deallocate grad memory.
+                param.grad = None
+
+            def allreduce_grads(use_fused_div, tensor, data_group, grad_mult):
+                if use_fused_div:
+                    torch.distributed.all_reduce(
+                        tensor,
+                        group=data_group,
+                        async_op=True,
+                        op=torch.distributed._make_nccl_premul_sum(1 / grad_mult),
+                    )
+                else:
+                    tensor.div_(grad_mult)
+                    torch.distributed.all_reduce(
+                        tensor,
+                        group=data_group,
+                        async_op=True,
+                    )
+
+            # Asynchronous gradients allreduce accross data_parallel ranks
+            grad_mult = get_data_parallel_world_size()
+            if self._require_backward_grad_sync:
+                data_group = _get_grad_data_group(is_expert_group)
+                if self._grad_allreduce_chunk_size_mb > 0:
+                    self._main_grad_buffers[i].update_chunk_info(grad_chunk_info)
+                    while True:
+                        allreduce_tensor = self._main_grad_buffers[i].get_allreduce_tensor()
+                        if allreduce_tensor is None:
+                            break
+                        allreduce_grads(self._grad_div_ar_fusion, allreduce_tensor, data_group, grad_mult)
+                else:
+                    allreduce_grads(self._grad_div_ar_fusion, main_param.grad, data_group, grad_mult)
+
+        return param_hook
+
+    def zero_grad(self, set_to_none=True):
+        """We only need to zero the model related parameters, i.e.,
+        float16_groups & fp32_from_fp32_groups. We additionally zero
+        fp32_from_float16_groups as a memory optimization to reduce
+        fragmentation; in the case of set_to_none==True, the space
+        used by this field can be safely deallocated at this point."""
+        for group in self.float16_groups:
+            _zero_grad_group_helper(group, set_to_none)
+        if self._contiguous_grad_bucket:
+            for i in self._main_grad_buffers:
+                self._main_grad_buffers[i].zero()
+        else:
+            for group in self.fp32_from_float16_groups:
+                _zero_grad_group_helper(group, set_to_none)
+        for group in self.fp32_from_fp32_groups:
+            _zero_grad_group_helper(group, set_to_none)
+
+    def copy_model_grads_to_main_grads(self):
+        # This only needs to be done for the float16 group.
+        for model_group, main_group in zip(self.float16_groups, self.fp32_from_float16_groups):
+            for model_param, main_param in zip(model_group, main_group):
+                if model_param.grad is not None:
+                    main_param.grad = model_param.grad.float()
+
+                # Safe to deallocate model's grad after copying.
+                # (If using contiguous buffers, main_grad's memory should
+                # persist and therefore should not be deallocated.)
+                model_param.grad = None
+
+    def _get_model_and_main_params_data_float16(self):
+        model_data = []
+        main_data = []
+        half_dtype = None
+        for model_group, main_group in zip(self.float16_groups, self.fp32_from_float16_groups):
+            for model_param, main_param in zip(model_group, main_group):
+                if half_dtype is None:
+                    half_dtype = model_param.data.dtype
+                model_data.append(model_param.data)
+                main_data.append(main_param.data)
+        return model_data, main_data, half_dtype
+
+    def _set_overflow_buffer(self, half_dtype):
+        if half_dtype == torch.float16:
+            if self._dummy_overflow_buf is None:
+                self._dummy_overflow_buf = torch.cuda.IntTensor([0])
+            else:
+                self._dummy_overflow_buf.fill_(0)
+
+    def _copy_main_params_to_model_params(self):
+        # Only needed for the float16 params.
+        model_data, main_data, half_dtype = self._get_model_and_main_params_data_float16()
+        self._set_overflow_buffer(half_dtype)
+        _multi_tensor_copy_this_to_that(this=main_data, that=model_data, overflow_buf=self._dummy_overflow_buf)
+
+    def _copy_model_params_to_main_params(self):
+        # Only needed for the float16 params.
+        model_data, main_data, half_dtype = self._get_model_and_main_params_data_float16()
+        self._set_overflow_buffer(half_dtype)
+        _multi_tensor_copy_this_to_that(this=model_data, that=main_data, overflow_buf=self._dummy_overflow_buf)
+
+    def reload_model_params(self):
+        self._copy_model_params_to_main_params()
+
+    @torch.no_grad()
+    def step(self, **kwargs):
+        # while async grad allreduce is enabled, bprop will keep moving forward without waiting for
+        # the finish of async grad AR works. Hence, to guarantee the correctness of grads reduction,
+        # we cannot start weight update until all async grad AR works are done.
+        if self._async_grad_allreduce:
+            torch.cuda.synchronize()
+
+        closure = kwargs.pop('closure', None)
+        # Allows applications to specify closure as None without erroring due to duplicate specification
+        assert closure is None, f"closure should be None but was passed {closure}"
+
+        # Step the optimizer.
+        self.optimizer.step(closure=None, **kwargs)
+
+        # Update params from main params.
+        with torch.no_grad():
+            self._copy_main_params_to_model_params()
+
+        # Successful update.
+        return True
+
+    def state_dict(self):
+        state_dict = {}
+        state_dict['optimizer'] = self.optimizer.state_dict()
+        state_dict['fp32_from_fp16_params'] = self.fp32_from_float16_groups
+        return state_dict
+
+    def load_state_dict(self, state_dict):
+        # Optimizer.
+        optimizer_key = 'optimizer'
+        if optimizer_key not in state_dict:
+            optimizer_key = 'optimizer_state_dict'
+            logging.info('***WARNING*** loading optimizer from ' 'an old checkpoint ...')
+        if 'state' not in state_dict[optimizer_key]:
+            state_dict[optimizer_key]['state'] = {}
+        self.optimizer.load_state_dict(state_dict[optimizer_key])
+
+        # Copy data for the main params.
+        fp32_from_float16_params_key = 'fp32_from_fp16_params'
+        if fp32_from_float16_params_key not in state_dict:
+            fp32_from_float16_params_key = 'fp32_from_fp16'
+        if fp32_from_float16_params_key not in state_dict:
+            state_dict[fp32_from_float16_params_key] = []
+        for current_group, saved_group in zip(self.fp32_from_float16_groups, state_dict[fp32_from_float16_params_key]):
+            for current_param, saved_param in zip(current_group, saved_group):
+                current_param.data.copy_(saved_param.data)
+
+    def allreduce_main_grads(self):
+        for i in self._main_grad_buffers:
+            self._main_grad_buffers[i].allreduce_buffer()
+
+    @contextmanager
+    def no_sync(self):
+        """A context manager to disable gradient synchronizations across
+        data-parallel ranks."""
+        old_require_backward_grad_sync = self._require_backward_grad_sync
+        self._require_backward_grad_sync = False
+        try:
+            yield
+        finally:
+            self._require_backward_grad_sync = old_require_backward_grad_sync
+
+    @property
+    def async_master_grads_allreudce(self):
+        return self._async_grad_allreduce
+
+    @property
+    def fp32_grad_accumulation(self):
+        return self._fp32_grad_accum
+
+    def get_parameters_with_grad(self):
+        params = []
+        for param_group in self.optimizer.param_groups:
+            for param in param_group['params']:
+                if param.grad is not None:  # (@adithyare) added to enable pp>1 training for adapters
+                    params.append(param)
+        return params
+
+    # Promote state so it can be retrieved or set via
+    # "optimizer_instance.state"
+    def _get_state(self):
+        if hasattr(self, 'optimizer'):
+            return self.optimizer.state
+        else:
+            return []
+
+    def _set_state(self, value):
+        self.optimizer.state = value
+
+    state = property(_get_state, _set_state)
+
+    # Promote param_groups so it can be retrieved or set via
+    # "optimizer_instance.param_groups"
+    # (for example, to adjust the learning rate)
+    def _get_param_groups(self):
+        if hasattr(self, 'optimizer'):
+            return self.optimizer.param_groups
+        else:
+            return []
+
+    def _set_param_groups(self, value):
+        self.optimizer.param_groups = value
+
+    param_groups = property(_get_param_groups, _set_param_groups)
+
+    # Promote defaults so it can be retrieved or set via
+    # "optimizer_instance.defaults
+    def _get_defaults(self):
+        if hasattr(self, 'optimizer'):
+            return self.optimizer.defaults
+        else:
+            return []
+
+    def _set_defaults(self, value):
+        self.optimizer.defaults = value
+
+    defaults = property(_get_defaults, _set_defaults)

nemo/core/optim/optimizers.py

@@ -0,0 +1,226 @@
+# Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import copy
+from functools import partial
+from typing import Any, Dict, Optional, Union
+
+import torch
+import torch.optim as optim
+from omegaconf import DictConfig, OmegaConf
+from torch.optim import adadelta, adagrad, adamax, rmsprop, rprop
+from torch.optim.optimizer import Optimizer
+
+from nemo.core.classes.common import safe_instantiate
+from nemo.core.config.optimizers import OptimizerParams, get_optimizer_config, register_optimizer_params
+from nemo.core.optim.adafactor import Adafactor
+from nemo.core.optim.adan import Adan
+from nemo.core.optim.novograd import Novograd
+
+from nemo.utils.model_utils import maybe_update_config_version
+
+AVAILABLE_OPTIMIZERS = {
+    'sgd': optim.SGD,
+    'adam': optim.Adam,
+    'adamw': optim.AdamW,
+    'adadelta': adadelta.Adadelta,
+    'adamax': adamax.Adamax,
+    'adagrad': adagrad.Adagrad,
+    'rmsprop': rmsprop.RMSprop,
+    'rprop': rprop.Rprop,
+    'novograd': Novograd,
+    'adafactor': Adafactor,
+    'adan': Adan,
+}
+
+try:
+    from apex.optimizers import FusedAdam, FusedLAMB
+
+    HAVE_APEX = True
+
+    AVAILABLE_OPTIMIZERS['lamb'] = FusedLAMB
+    AVAILABLE_OPTIMIZERS['fused_adam'] = FusedAdam
+except ModuleNotFoundError:
+    HAVE_APEX = False
+
+HAVE_APEX_DISTRIBUTED_ADAM = False
+if HAVE_APEX:
+    try:
+        # Try importing wrapper for Apex distributed Adam optimizer
+        from nemo.core.optim.distributed_adam import MegatronDistributedFusedAdam
+
+        HAVE_APEX_DISTRIBUTED_ADAM = True
+
+        AVAILABLE_OPTIMIZERS['distributed_fused_adam'] = MegatronDistributedFusedAdam
+    except (ImportError, ModuleNotFoundError):
+        HAVE_APEX_DISTRIBUTED_ADAM = False
+
+    try:
+        # Try importing wrapper for Apex FusedAdam optimizer
+        from nemo.core.optim.megatron_fused_adam import MegatronFusedAdam
+
+        AVAILABLE_OPTIMIZERS['megatron_fused_adam'] = MegatronFusedAdam
+    except (ImportError, ModuleNotFoundError):
+        pass
+
+__all__ = ['get_optimizer', 'register_optimizer', 'parse_optimizer_args']
+
+
+def parse_optimizer_args(
+    optimizer_name: str, optimizer_kwargs: Union[DictConfig, Dict[str, Any]]
+) -> Union[Dict[str, Any], DictConfig]:
+    """
+    Parses a list of strings, of the format "key=value" or "key2=val1,val2,..."
+    into a dictionary of type {key=value, key2=[val1, val2], ...}
+
+    This dictionary is then used to instantiate the chosen Optimizer.
+
+    Args:
+        optimizer_name: string name of the optimizer, used for auto resolution of params
+        optimizer_kwargs: Either a list of strings in a specified format,
+            or a dictionary. If a dictionary is provided, it is assumed the dictionary
+            is the final parsed value, and simply returned.
+            If a list of strings is provided, each item in the list is parsed into a
+            new dictionary.
+
+    Returns:
+        A dictionary
+    """
+    kwargs = {}
+
+    if optimizer_kwargs is None:
+        return kwargs
+
+    optimizer_kwargs = copy.deepcopy(optimizer_kwargs)
+    optimizer_kwargs = maybe_update_config_version(optimizer_kwargs)
+
+    if isinstance(optimizer_kwargs, DictConfig):
+        optimizer_kwargs = OmegaConf.to_container(optimizer_kwargs, resolve=True)
+
+    # If it is a dictionary, perform stepwise resolution
+    if hasattr(optimizer_kwargs, 'keys'):
+        # Attempt class path resolution
+        if '_target_' in optimizer_kwargs:  # captures (target, _target_)
+            optimizer_kwargs_config = OmegaConf.create(optimizer_kwargs)
+            optimizer_instance = safe_instantiate(optimizer_kwargs_config)  # type: DictConfig
+            optimizer_instance = vars(optimizer_instance)
+            return optimizer_instance
+
+        # If class path was not provided, perhaps `name` is provided for resolution
+        if 'name' in optimizer_kwargs:
+            # If `auto` is passed as name for resolution of optimizer name,
+            # then lookup optimizer name and resolve its parameter config
+            if optimizer_kwargs['name'] == 'auto':
+                optimizer_params_name = "{}_params".format(optimizer_name)
+                optimizer_kwargs.pop('name')
+            else:
+                optimizer_params_name = optimizer_kwargs.pop('name')
+
+            # Override arguments provided in the config yaml file
+            if 'params' in optimizer_kwargs:
+                # If optimizer kwarg overrides are wrapped in yaml `params`
+                optimizer_params_override = optimizer_kwargs.get('params')
+            else:
+                # If the kwargs themselves are a DictConfig
+                optimizer_params_override = optimizer_kwargs
+
+            if isinstance(optimizer_params_override, DictConfig):
+                optimizer_params_override = OmegaConf.to_container(optimizer_params_override, resolve=True)
+
+            optimizer_params_cls = get_optimizer_config(optimizer_params_name, **optimizer_params_override)
+
+            # If we are provided just a Config object, simply return the dictionary of that object
+            if optimizer_params_name is None:
+                optimizer_params = vars(optimizer_params_cls)
+                return optimizer_params
+
+            else:
+                # If we are provided a partial class instantiation of a Config,
+                # Instantiate it and retrieve its vars as a dictionary
+                optimizer_params = optimizer_params_cls()  # instantiate the parameters object
+                optimizer_params = vars(optimizer_params)
+                return optimizer_params
+
+        # simply return the dictionary that was provided
+        return optimizer_kwargs
+
+    return kwargs
+
+
+def register_optimizer(name: str, optimizer: Optimizer, optimizer_params: OptimizerParams):
+    """
+    Checks if the optimizer name exists in the registry, and if it doesnt, adds it.
+
+    This allows custom optimizers to be added and called by name during instantiation.
+
+    Args:
+        name: Name of the optimizer. Will be used as key to retrieve the optimizer.
+        optimizer: Optimizer class
+        optimizer_params: The parameters as a dataclass of the optimizer
+    """
+    if name in AVAILABLE_OPTIMIZERS:
+        raise ValueError(f"Cannot override pre-existing optimizers. Conflicting optimizer name = {name}")
+
+    AVAILABLE_OPTIMIZERS[name] = optimizer
+
+    optim_name = "{}_params".format(optimizer.__name__)
+    register_optimizer_params(name=optim_name, optimizer_params=optimizer_params)
+
+
+def get_optimizer(name: str, **kwargs: Optional[Dict[str, Any]]) -> Optimizer:
+    """
+    Convenience method to obtain an Optimizer class and partially instantiate it with optimizer kwargs.
+
+    Args:
+        name: Name of the Optimizer in the registry.
+        kwargs: Optional kwargs of the optimizer used during instantiation.
+
+    Returns:
+        a partially instantiated Optimizer
+    """
+    if name not in AVAILABLE_OPTIMIZERS:
+        raise ValueError(
+            f"Cannot resolve optimizer '{name}'. Available optimizers are : " f"{AVAILABLE_OPTIMIZERS.keys()}"
+        )
+    if name == 'fused_adam':
+        if not torch.cuda.is_available():
+            raise ValueError('CUDA must be available to use fused_adam.')
+
+    optimizer = AVAILABLE_OPTIMIZERS[name]
+    optimizer = partial(optimizer, **kwargs)
+    return optimizer
+
+
+def init_optimizer_states(optimizer: Optimizer):
+    """
+    Initialize optimizer states for Adam-based optimizers.
+
+    This function initializes the exponential moving averages (exp_avg and exp_avg_sq)
+    for Adam, AdamW, and FusedAdam optimizers if they haven't been initialized yet.
+
+    Args:
+        optimizer: The optimizer instance to initialize states for
+    """
+    adam_nondist_optims = (optim.Adam, optim.AdamW)
+    if HAVE_APEX:
+        adam_nondist_optims += (FusedAdam,)
+    if isinstance(optimizer, adam_nondist_optims):
+        for group in optimizer.param_groups:
+            for p in group['params']:
+                state = optimizer.state[p]
+                if len(state) == 0:
+                    state['exp_avg'] = torch.zeros_like(p.data, memory_format=torch.preserve_format)
+                    state['exp_avg_sq'] = torch.zeros_like(p.data, memory_format=torch.preserve_format)
+                    if group.get('amsgrad'):
+                        state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format)

nemo/core/optim/radam.py

@@ -0,0 +1,129 @@
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""RAdam
+Original source taken from https://github.com/LiyuanLucasLiu/RAdam
+
+Copyright 2019 Liyuan Liu
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+"""
+import math
+
+import torch
+from torch.optim.optimizer import Optimizer
+
+
+class RAdam(Optimizer):
+    """RAdam optimizer"""
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, weight_decay=0):
+        """
+        Init
+
+        :param params: parameters to optimize
+        :param lr: learning rate
+        :param betas: beta
+        :param eps: numerical precision
+        :param weight_decay: weight decay weight
+        """
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+        self.buffer = [[None, None, None] for _ in range(10)]
+        super().__init__(params, defaults)
+
+    def step(self, closure=None):
+
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data.float()
+                if grad.is_sparse:
+                    raise RuntimeError('RAdam does not support sparse gradients')
+
+                p_data_fp32 = p.data.float()
+
+                state = self.state[p]
+
+                if len(state) == 0:
+                    state['step'] = 0
+                    state['exp_avg'] = torch.zeros_like(p_data_fp32)
+                    state['exp_avg_sq'] = torch.zeros_like(p_data_fp32)
+                else:
+                    state['exp_avg'] = state['exp_avg'].type_as(p_data_fp32)
+                    state['exp_avg_sq'] = state['exp_avg_sq'].type_as(p_data_fp32)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=(1.0 - beta2))
+                exp_avg.mul_(beta1).add_(grad, alpha=(1.0 - beta1))
+
+                state['step'] += 1
+                buffered = self.buffer[int(state['step'] % 10)]
+                if state['step'] == buffered[0]:
+                    N_sma, step_size = buffered[1], buffered[2]
+                else:
+                    buffered[0] = state['step']
+                    beta2_t = beta2 ** state['step']
+                    N_sma_max = 2 / (1 - beta2) - 1
+                    N_sma = N_sma_max - 2 * state['step'] * beta2_t / (1 - beta2_t)
+                    buffered[1] = N_sma
+
+                    # more conservative since it's an approximated value
+                    if N_sma >= 5:
+                        step_size = (
+                            group['lr']
+                            * math.sqrt(
+                                (1 - beta2_t)
+                                * (N_sma - 4)
+                                / (N_sma_max - 4)
+                                * (N_sma - 2)
+                                / N_sma
+                                * N_sma_max
+                                / (N_sma_max - 2)
+                            )
+                            / (1 - beta1 ** state['step'])
+                        )
+                    else:
+                        step_size = group['lr'] / (1 - beta1 ** state['step'])
+                    buffered[2] = step_size
+
+                if group['weight_decay'] != 0:
+                    p_data_fp32.add_(-group['weight_decay'] * group['lr'], p_data_fp32)
+
+                # more conservative since it's an approximated value
+                if N_sma >= 5:
+                    denom = exp_avg_sq.sqrt().add_(group['eps'])
+                    p_data_fp32.addcdiv_(-step_size, exp_avg, denom)
+                else:
+                    p_data_fp32.add_(-step_size, exp_avg)
+
+                p.data.copy_(p_data_fp32)
+
+        return loss