Merge pull request #8 from lcskrishna/ifu_05152020

[Upstream] IFU 05/15/2020

apex/contrib/csrc/optimizers/fused_adam_cuda.cpp

 #include <torch/extension.h>
 
 // CUDA forward declaration
+void fused_strided_check_finite(at::Tensor & overflow_flag, at::Tensor & p_copy, int stride, int clear_overflow_first);
+
 void fused_adam_cuda(at::Tensor & p, at::Tensor & p_copy, at::Tensor & m, at::Tensor & v, at::Tensor & g, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay);
+void fused_reversible_adam_cuda(at::Tensor & p, at::Tensor & p_copy, at::Tensor & m, at::Tensor & v, at::Tensor & g, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay);
+void fused_maybe_adam_undo_cuda(at::Tensor & overflow_flag, at::Tensor & p, at::Tensor & m, at::Tensor & v, at::Tensor & g, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay);
 
-void fused_adam_cuda_mt(int chunk_size, at::Tensor noop_flag, std::vector<std::vector<at::Tensor>> tensor_lists, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay);
+void fused_adam_cuda_mt(int chunk_size, at::Tensor overflow_flag, std::vector<std::vector<at::Tensor>> tensor_lists, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay);
 
+void maybe_cast_cuda(at::Tensor & overflow_flag, at::Tensor & p_in, at::Tensor & p_out);
+void maybe_cast_cuda_mt(int chunk_size, at::Tensor overflow_flag, std::vector<std::vector<at::Tensor>> tensor_lists);
 
 #define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
 #define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
 #define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
 
 // C++ interface
+void strided_check_finite(
+		at::Tensor& overflow_flag,
+		at::Tensor& p_copy,
+		int stride,
+		int clear_overflow_first
+	 ) {
+	CHECK_INPUT(p_copy);
+	fused_strided_check_finite(overflow_flag, p_copy, stride, clear_overflow_first);
+}
 void adam(at::Tensor & p, at::Tensor & p_copy, at::Tensor & m, at::Tensor & v, at::Tensor & g, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay) {
         CHECK_INPUT(p);
         if (p_copy.numel() > 0) CHECK_INPUT(p_copy);
@@ -25,8 +40,47 @@ void adam(at::Tensor & p, at::Tensor & p_copy, at::Tensor & m, at::Tensor & v, a
 
         fused_adam_cuda(p, p_copy, m, v, g, lr, beta1, beta2, eps, grad_scale, step, mode, bias_correction, decay);
 }
+void reversible_adam(at::Tensor & p, at::Tensor & p_copy, at::Tensor & m, at::Tensor & v, at::Tensor & g, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay) {
+        CHECK_INPUT(p);
+        if (p_copy.numel() > 0) CHECK_INPUT(p_copy);
+        CHECK_INPUT(m);
+        CHECK_INPUT(v);
+        CHECK_INPUT(g);
+        int64_t num_elem = p.numel();
+        AT_ASSERTM(m.numel() == num_elem, "number of elements in m and p tensors should be equal");
+        AT_ASSERTM(v.numel() == num_elem, "number of elements in v and p tensors should be equal");
+        AT_ASSERTM(g.numel() == num_elem, "number of elements in g and p tensors should be equal");
+        AT_ASSERTM(p_copy.numel() == num_elem || p_copy.numel() == 0, "number of elements in p_copy and p tensors should be equal, or p_copy should be empty");
+
+        fused_reversible_adam_cuda(p, p_copy, m, v, g, lr, beta1, beta2, eps, grad_scale, step, mode, bias_correction, decay);
+}
+void maybe_adam_undo(at::Tensor & overflow_flag, at::Tensor & p, at::Tensor & m, at::Tensor & v, at::Tensor & g, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay) {
+        CHECK_INPUT(p);
+        CHECK_INPUT(m);
+        CHECK_INPUT(v);
+        CHECK_INPUT(g);
+        int64_t num_elem = p.numel();
+        AT_ASSERTM(m.numel() == num_elem, "number of elements in m and p tensors should be equal");
+        AT_ASSERTM(v.numel() == num_elem, "number of elements in v and p tensors should be equal");
+        AT_ASSERTM(g.numel() == num_elem, "number of elements in g and p tensors should be equal");
+
+        fused_maybe_adam_undo_cuda(overflow_flag, p, m, v, g, lr, beta1, beta2, eps, grad_scale, step, mode, bias_correction, decay);
+}
+void maybe_cast(at::Tensor & overflow_flag, at::Tensor & p_in, at::Tensor & p_out) {
+	CHECK_INPUT(p_in);
+	CHECK_INPUT(p_out);
+	int64_t num_elem = p_in.numel();
+	AT_ASSERTM(p_out.numel() == num_elem, "number of elements in p_in and p_out should be equal");
+
+	maybe_cast_cuda(overflow_flag, p_in, p_out);
+}
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+        m.def("strided_check_finite", &strided_check_finite, "Strided finite check.");
         m.def("adam", &adam, "Adam optimized CUDA implementation.");
+        m.def("reversible_adam", &reversible_adam, "Reversible Adam optimized CUDA implementation.");
         m.def("adam_mt", &fused_adam_cuda_mt, "Multi tensor Adam optimized CUDA implementation.");
+        m.def("maybe_adam_undo", &maybe_adam_undo, "Undo function for Adam optimized CUDA implementation.");
+        m.def("maybe_cast", &maybe_cast, "Unpack byte tensor containing e5m2 floats.");
+        m.def("maybe_cast_mt", &maybe_cast_cuda_mt, "Unpack byte tensor containing e5m2 floats.");
 }

apex/contrib/optimizers/distributed_fused_adam_v3.py

+import math
+import torch
+import importlib
+import amp_C
+from apex.multi_tensor_apply import multi_tensor_applier
+
+class DistributedFusedAdamV3(torch.optim.Optimizer):
+
+    """Implements Adam algorithm. Currently GPU-only.  Requires Apex to be installed via
+    ``python setup.py install --cuda_ext --cpp_ext``.
+
+    It has been proposed in `Adam: A Method for Stochastic Optimization`_.
+
+    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`_
+            (default: False) NOT SUPPORTED in FusedAdam!
+        eps_inside_sqrt (boolean, optional): in the 'update parameters' step,
+            adds eps to the bias-corrected second moment estimate before
+            evaluating square root instead of adding it to the square root of
+            second moment estimate as in the original paper. (default: False)
+        use_mt (boolean, optional): use multi tensor apply for lower launch
+            latency. (default: False)
+        overlap_reductions(boolean, optional): whether to overlap reductions
+            with bprop (default: True)
+        num_prestats (integer, optional): number of fp64 stats that will be
+            reduced during first fp16 gradient reduction block. 
+
+    .. _Adam\: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(self, params,
+                 lr=1e-3, bias_correction = True,
+                 betas=(0.9, 0.999), eps=1e-8, eps_inside_sqrt = False,
+                 weight_decay=0., max_grad_norm=0., amsgrad=False, use_mt=False,
+                 amp_scale_adjustment=1.0, overlap_reductions=True, full_pipeline=True,
+                 compute_L2_grad_norm=False, distributed_weight_update=0,
+                 dwu_group_size=0, dwu_num_blocks=4, dwu_num_rs_pg=1, dwu_num_ar_pg=4,
+                 dwu_num_ag_pg=0, revert_method=1, flat_mt=False,
+                 dwu_num_chunks=4, predivide=True, e5m2_allgather=False,
+                 do_not_flatten_model=False):
+        global fused_adam_cuda
+        fused_adam_cuda = importlib.import_module("fused_adam_cuda")
+
+        self._amp_scale_adjustment = amp_scale_adjustment
+
+        if use_mt:
+            raise RuntimeError('DistributedFusedAdam does not support use_mt.')
+        if amsgrad:
+            raise RuntimeError('DistributedFusedAdam does not support the AMSGrad variant.')
+
+        defaults = dict(lr=lr, bias_correction=bias_correction,
+                        betas=betas, eps=eps, weight_decay=weight_decay,
+                        max_grad_norm=max_grad_norm)
+        super(DistributedFusedAdamV3, self).__init__(params, defaults)
+        self.eps_mode = 0 if  eps_inside_sqrt else 1
+
+        self._overflow_buf = torch.cuda.IntTensor([0])
+
+        assert (len(self.param_groups) == 1), "More than one parameter group is not supported."
+
+        # Way to revert a step
+        # 3 -> undo kernel + double buffer (debug, print norm of difference)
+        # 2 -> double buffer fp32 parameters
+        # 1 -> undo kernel
+        self._revert_method = revert_method
+        if self._revert_method > 1:
+            print("revert_method -> double buffer fp32 parameters, will consume more memory")
+
+        self._last_step = False
+        self._overlap_reductions = overlap_reductions
+        self._global_scale = None
+        self._num_blocks = dwu_num_blocks
+        self._predivide = predivide
+        self._e5m2_allgather = e5m2_allgather
+        self._do_not_flatten_model = do_not_flatten_model
+        self._full_pipeline = full_pipeline
+        self._L2_grad_norm = None
+        self._group_size = torch.cuda.device_count() if dwu_group_size <= 0 else dwu_group_size
+        self._world_size = torch.distributed.get_world_size()
+        self._num_groups = self._world_size // self._group_size
+        self._rank_in_group = torch.distributed.get_rank() % self._group_size
+
+        p_offset = 0
+        p_i = 0
+        self._param_state = None
+        self._model_params = []
+        self._grads_info = []
+        self._grad_accs = []
+        for group in self.param_groups:
+            self._param_group = group
+            prev = None
+            for p in group['params']:
+                torch.distributed.broadcast(p,0)
+                if not p.requires_grad:
+                    continue
+                self._model_params.append(p)
+                state = self.state[p]
+                if len(state) == 0:
+                    state['step'] = 0
+                if self._param_state is None:
+                    self._param_state = state
+                p_grads_size = p.numel()
+                def wrapper(param, param_i, param_grads_size, param_offset):
+                    param_tmp = param.expand_as(param)
+                    grad_acc = param_tmp.grad_fn.next_functions[0][0]
+                    def allreduce_hook(*unused):
+                        self._do_overlapped_reduction(param_i, param_grads_size, param_offset, param)
+                    grad_acc.register_hook(allreduce_hook)
+                    self._grad_accs.append(grad_acc)
+                self._grads_info.append({"param_grads_size":p_grads_size, "param_offset":p_offset})
+                wrapper(p, p_i, p_grads_size, p_offset)
+                p_offset += p_grads_size
+                # Only enforce 128b alignment (64 * fp16) for non-consecutive parameters
+                # RNN is one example of consecutive parameters:
+                # (weight_ih, weight_hh, bias_ih, bias_hh)
+                if prev is not None and (prev.data_ptr() + prev.numel() * prev.element_size() != p.data_ptr()):
+                    p_offset = ((p_offset + 63) // 64) * 64
+                prev = p
+                p_i += 1
+        self._grads_generated = [False]*len(self._grads_info)
+        self._flat_mt = flat_mt
+        self._grads = []
+        self._current_block = self._num_blocks
+
+        self._net_total_param_size = p_offset
+        self._total_param_size = p_offset
+        dwu_min_page_size = 256 * self._num_blocks * self._group_size
+        self._total_param_size = ((self._total_param_size + dwu_min_page_size - 1) // dwu_min_page_size) * dwu_min_page_size
+        self._block_size = self._total_param_size // self._num_blocks
+        self._shard_size = self._total_param_size // self._group_size
+        print("self._net_total_param_size=%d, self._total_param_size=%d, dwu_min_page_size=%d, self._block_size=%d, self._shard_size=%d" % (self._net_total_param_size, self._total_param_size,dwu_min_page_size,self._block_size,self._shard_size))
+
+        self._low_param_i = [0]*self._num_blocks
+        for block_id in range(self._num_blocks-1,-1,-1):
+            p_i = len(self._grads_info)-1
+            while p_i > 0 and self._grads_info[p_i]["param_offset"] > block_id*self._block_size:
+                p_i -= 1
+            self._low_param_i[block_id] = p_i
+        print(self._low_param_i)
+
+        self._flat_grads = torch.zeros([self._total_param_size], dtype=torch.float16, device='cuda')
+        self._flat_params = torch.zeros_like(self._flat_grads)
+
+        def _flat_split(flat):
+            def __flat_blockify(flat):
+                return [flat[block_id*self._block_size:(block_id+1)*self._block_size] for block_id in range(self._num_blocks)]
+            def __flat_shardify(flat):
+                return [flat[shard_id*self._shard_size:(shard_id+1)*self._shard_size] for shard_id in range(self._group_size)]
+            return __flat_blockify(flat), __flat_shardify(flat)
+        self._flat_grads_blocks, self._flat_grads_shards = _flat_split(self._flat_grads)
+        self._flat_params_blocks, self._flat_params_shards = _flat_split(self._flat_params)
+
+        # master params
+        self._fp32_p = torch.zeros([self._shard_size], dtype=torch.float32, device='cuda')
+        self._fp32_m = torch.zeros([self._shard_size], dtype=torch.float32, device='cuda')
+        self._fp32_v = torch.zeros([self._shard_size], dtype=torch.float32, device='cuda')
+
+        # copy model params to flat_params and set_ model params to flat_params.
+        self._individual_flat_grads = []
+        with torch.no_grad():
+            for p, grads_info in zip(self._model_params, self._grads_info):
+                start = grads_info["param_offset"]
+                end = start + grads_info["param_grads_size"]
+                flat_p = self._flat_params[start:end].view_as(p)
+                flat_p.copy_(p)
+                p.set_(flat_p)
+                flat_grad = self._flat_grads[start:end]
+                self._individual_flat_grads.append(flat_grad)
+        self._fp32_p.copy_(self._flat_params_shards[self._rank_in_group].float())
+
+        self._dwu_st = torch.cuda.Stream()
+        self._l2_grad_norm_st = torch.cuda.Stream()
+        for group_i in range(self._num_groups):
+            ranks = [group_i*self._group_size+local_rank for local_rank in range(self._group_size)]
+            pg = torch.distributed.new_group(ranks=ranks)
+            if torch.distributed.get_rank() in ranks:
+                self._ag_pg = pg
+                torch.distributed.all_reduce(self._overflow_buf, group=self._ag_pg)
+
+        import inspect
+        assert ('no_copy' in inspect.getfullargspec(torch.distributed.reduce_scatter).args), "This version of c10d does not support no_copy option"
+
+    @property
+    def has_overflow(self):
+        return True if not self.L2_grad_norm is None and not math.isfinite(self.L2_grad_norm) else False
+
+    def set_last_step(self, last_step):
+        self._last_step = last_step
+        
+    def _get_flush_block(self):
+        flush_block = []
+        if self._current_block > 0 and self._grads_generated[self._low_param_i[self._current_block-1]]:
+            num_grads = len(self._grads_generated)
+            contiguous_idx = num_grads
+            while contiguous_idx > 0 and self._grads_generated[contiguous_idx-1]:
+                contiguous_idx -= 1
+
+            if contiguous_idx < num_grads and self._grads_info[contiguous_idx]["param_offset"] <= (self._current_block-1)*self._block_size:
+                self._current_block -= 1
+                start = self._current_block * self._block_size
+                end = (self._current_block+1) * self._block_size
+                flush_block = [start, end]
+
+        return flush_block
+
+    def __launch_step_kernel(self, p, p_copy, m, v, g):
+        combined_scale = self._global_scale
+        if self._param_group['max_grad_norm'] > 0 and math.isfinite(self.L2_grad_norm):
+            combined_scale = self._param_group['max_grad_norm'] / (self.L2_grad_norm / self._global_scale + 1e-6)
+            combined_scale = self._global_scale / min(1, combined_scale)
+        bias_correction = 1 if self._param_group['bias_correction'] else 0
+        beta1, beta2 = self._param_group['betas']
+        fused_adam_cuda.reversible_adam(
+                p, p_copy, m, v, g,
+                self._param_group['lr'],
+                beta1,
+                beta2,
+                self._param_group['eps'],
+                combined_scale,
+                self._param_state['step']+1,
+                self.eps_mode,
+                bias_correction,
+                self._param_group['weight_decay'])
+
+    def _flatten_grad_mt(self, scale):
+        if self._flat_mt and len(self._grads) > 0:
+            self._overflow_buf.zero_()
+            multi_tensor_applier(
+                    amp_C.multi_tensor_scale,
+                    self._overflow_buf,
+                    list(zip(*self._grads)),
+                    scale)
+            self._grads = []
+
+    def _do_overlapped_reduction(self, param_i, param_grads_size, param_offset, param):
+        # handle overlapped reductions
+        if self._flat_mt:
+            self._grads.append( (param.grad, self._individual_flat_grads[param_i]) )
+        else:
+            torch.div(param.grad, self._world_size if self._predivide else 1.0, out=self._individual_flat_grads[param_i])
+        self._grads_generated[param_i]=True
+        if not self._last_step and self._overlap_reductions:
+            flush_block = self._get_flush_block()
+            while flush_block:
+                block_id = flush_block[0] // self._block_size
+                self._dwu_st.wait_stream(torch.cuda.current_stream())
+                with torch.cuda.stream(self._dwu_st):
+                    self._flatten_grad_mt(1.0/self._world_size if self._predivide else 1.0)
+                    torch.distributed.all_reduce(self._flat_grads_blocks[block_id])
+                if block_id == 0:
+                    self._l2_grad_norm_st.wait_stream(self._dwu_st)
+                    with torch.cuda.stream(self._l2_grad_norm_st):
+                        self._L2_grad_norm = self._flat_grads.norm(dtype=torch.float32, p=2).item()
+                flush_block = self._get_flush_block()
+
+    def set_global_scale(self, global_scale):
+        """Set global scale.
+        """
+        self._global_scale = global_scale
+
+    @property
+    def global_scale(self):
+        return self._global_scale
+
+    @property
+    def L2_grad_norm(self):
+        torch.cuda.current_stream().wait_stream(self._l2_grad_norm_st)
+        return self._L2_grad_norm
+
+    def complete_reductions(self):
+        """Complete reductions if full pipeline is not selected or overlap is not allowed.
+        """
+
+        if self._last_step:
+            # zero out gradients that have not been completed yet
+            for param_i, flat_grad in enumerate(self._individual_flat_grads):
+                if not self._grads_generated[param_i]:
+                    flat_grad.zero_()
+                    self._grads_generated[param_i] = True
+
+        if self._last_step or not self._overlap_reductions:
+            # nothing done so far, run full pipeline after reductions
+            self._dwu_st.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(self._dwu_st):
+                self._flatten_grad_mt(1.0/self._world_size if self._predivide else 1.0)
+                torch.distributed.all_reduce(self._flat_grads)
+            self._l2_grad_norm_st.wait_stream(self._dwu_st)
+            with torch.cuda.stream(self._l2_grad_norm_st):
+                self._L2_grad_norm = self._flat_grads.norm(dtype=torch.float32, p=2).item()
+
+        self._current_block = self._num_blocks
+        self._grads_generated = [False]*len(self._grads_info)
+
+    def step(self, closure=None, skip_overflow_check=False):
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        with torch.cuda.stream(self._dwu_st):
+            self.__launch_step_kernel(
+                self._fp32_p,
+                self._flat_params_shards[self._rank_in_group],
+                self._fp32_m,
+                self._fp32_v,
+                self._flat_grads_shards[self._rank_in_group])
+            torch.distributed.all_gather(self._flat_params_shards, self._flat_params_shards[self._rank_in_group], group=self._ag_pg, no_copy=True)
+            for p in self._model_params: self.state[p]['step'] += 1
+
+        torch.cuda.current_stream().wait_stream(self._dwu_st)
+
+        return loss
+
+

apex/contrib/optimizers/fused_adam.py

@@ -29,7 +29,7 @@ class FusedAdam(torch.optim.Optimizer):
         use_mt (boolean, optional): use multi tensor apply for lower launch
             latency. (default: False)
 
-    .. _Adam\: A Method for Stochastic Optimization:
+    .. _Adam - A Method for Stochastic Optimization:
         https://arxiv.org/abs/1412.6980
     .. _On the Convergence of Adam and Beyond:
         https://openreview.net/forum?id=ryQu7f-RZ

apex/contrib/optimizers/fused_lamb.py

@@ -54,7 +54,7 @@ class FusedLAMB(torch.optim.Optimizer):
         max_grad_norm (float, optional): value used to clip global grad norm
             (default: 1.0)
 
-    .. _Large Batch Optimization for Deep Learning\: Training BERT in 76 minutes:
+    .. _Large Batch Optimization for Deep Learning - Training BERT in 76 minutes:
         https://arxiv.org/abs/1904.00962
     .. _On the Convergence of Adam and Beyond:
         https://openreview.net/forum?id=ryQu7f-RZ

apex/optimizers/__init__.py

@@ -2,3 +2,4 @@ from .fused_sgd import FusedSGD
 from .fused_adam import FusedAdam
 from .fused_novograd import FusedNovoGrad
 from .fused_lamb import FusedLAMB
+from .fused_adagrad import FusedAdagrad
\ No newline at end of file

apex/optimizers/fused_adagrad.py

+import torch
+from apex.multi_tensor_apply import multi_tensor_applier
+
+
+class FusedAdagrad(torch.optim.Optimizer):
+    """Implements Adagrad algorithm.
+
+    Currently GPU-only.  Requires Apex to be installed via
+    ``pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./``.
+
+    This version of fused Adagrad implements 2 fusions.
+      * Fusion of the Adagrad update's elementwise operations
+      * A multi-tensor apply launch that batches the elementwise updates applied to all the model's parameters into one or a few kernel launches.
+
+    :class:`apex.optimizers.FusedAdagrad`'s usage is identical to any ordinary Pytorch optimizer::
+        opt = apex.optimizers.FusedAdagrad(model.parameters(), lr = ....)
+        ...
+        opt.step()
+
+    :class:`apex.optimizers.FusedAdagrad` may be used with or without Amp.  If you wish to use :class:`FusedAdagrad` with Amp,
+    you may choose any ``opt_level``::
+        opt = apex.optimizers.FusedAdagrad(model.parameters(), lr = ....)
+        model, opt = amp.initialize(model, opt, opt_level="O0" or "O1 or "O2")
+        ...
+        opt.step()
+    In general, ``opt_level="O1"`` is recommended.
+
+    It has been proposed in `Adaptive Subgradient Methods for Online Learning
+    and Stochastic Optimization`_.
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-2)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-10)
+        adagrad_w_mode (boolean, optional): Apply L2 regularization or weight decay
+            True for decoupled weight decay (also known as AdamW) (default: False)
+
+    .. _Adaptive Subgradient Methods for Online Learning and Stochastic
+        Optimization: http://jmlr.org/papers/v12/duchi11a.html
+    """
+    def __init__(self, params, lr=1e-2, eps=1e-10,
+                 weight_decay=0., set_grad_none=True, adagrad_w_mode=False):
+
+        defaults = dict(lr=lr, eps=eps, weight_decay=weight_decay)
+        super(FusedAdagrad, self).__init__(params, defaults)
+        self.adagrad_w_mode = 1 if adagrad_w_mode else 0
+        self.set_grad_none = set_grad_none
+
+        if multi_tensor_applier.available:
+            import amp_C
+            # Skip buffer
+            self._dummy_overflow_buf = torch.cuda.IntTensor([0])
+            self.multi_tensor_adagrad = amp_C.multi_tensor_adagrad
+        else:
+            raise RuntimeError('apex.optimizers.FusedAdagrad requires cuda extensions')
+
+    def zero_grad(self):
+        if self.set_grad_none:
+            for group in self.param_groups:
+                for p in group['params']:
+                    p.grad = None
+        else:
+            super(FusedAdagrad, self).zero_grad()
+
+    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:
+            # create lists for multi-tensor apply
+            g_16, p_16, h_16 = [], [], []
+            g_32, p_32, h_32 = [], [], []
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                if p.grad.data.is_sparse:
+                    raise RuntimeError('FusedAdagrad does not support sparse gradients')
+
+                state = self.state[p]
+                # State initialization
+                if len(state) == 0:
+                    # Exponential moving average of gradient values
+                    state['sum'] = torch.zeros_like(p.data)
+                if p.dtype == torch.float16:
+                    g_16.append(p.grad.data)
+                    p_16.append(p.data)
+                    h_16.append(state['sum'])
+                elif p.dtype == torch.float32:
+                    g_32.append(p.grad.data)
+                    p_32.append(p.data)
+                    h_32.append(state['sum'])
+                else:
+                    raise RuntimeError('FusedAdagrad only support fp16 and fp32.')
+
+            if(len(g_16) > 0):
+                multi_tensor_applier(self.multi_tensor_adagrad,
+                                     self._dummy_overflow_buf,
+                                     [g_16, p_16, h_16],
+                                     group['lr'],
+                                     group['eps'],
+                                     self.adagrad_w_mode,
+                                     group['weight_decay'])
+            if(len(g_32) > 0):
+                multi_tensor_applier(self.multi_tensor_adagrad,
+                                     self._dummy_overflow_buf,
+                                     [g_32, p_32, h_32],
+                                     group['lr'],
+                                     group['eps'],
+                                     self.adagrad_w_mode,
+                                     group['weight_decay'])
+
+        return loss
\ No newline at end of file

apex/optimizers/fused_adam.py

@@ -53,7 +53,7 @@ class FusedAdam(torch.optim.Optimizer):
         set_grad_none (bool, optional): whether set grad to None when zero_grad()
             method is called. (default: True)
 
-    .. _Adam\: A Method for Stochastic Optimization:
+    .. _Adam - A Method for Stochastic Optimization:
         https://arxiv.org/abs/1412.6980
     .. _On the Convergence of Adam and Beyond:
         https://openreview.net/forum?id=ryQu7f-RZ

apex/optimizers/fused_lamb.py

@@ -52,7 +52,7 @@ class FusedLAMB(torch.optim.Optimizer):
         max_grad_norm (float, optional): value used to clip global grad norm
             (default: 1.0)
 
-    .. _Large Batch Optimization for Deep Learning\: Training BERT in 76 minutes:
+    .. _Large Batch Optimization for Deep Learning - Training BERT in 76 minutes:
         https://arxiv.org/abs/1904.00962
     .. _On the Convergence of Adam and Beyond:
         https://openreview.net/forum?id=ryQu7f-RZ

apex/optimizers/fused_novograd.py

@@ -58,7 +58,7 @@ class FusedNovoGrad(torch.optim.Optimizer):
         set_grad_none (bool, optional): whether set grad to None when zero_grad()
             method is called. (default: True)
 
-    .. _Jasper\: An End-to-End Convolutional Neural Acoustic Model:
+    .. _Jasper - An End-to-End Convolutional Neural Acoustic Model:
         https://arxiv.org/abs/1904.03288
     .. _On the Convergence of Adam and Beyond:
         https://openreview.net/forum?id=ryQu7f-RZ

apex/reparameterization/__init__.py

@@ -2,7 +2,7 @@ from .weight_norm import WeightNorm
 from .reparameterization import Reparameterization
 
 def apply_weight_norm(module, name='', dim=0, hook_child=True):
-    """
+    r"""
     Applies weight normalization to a parameter in the given module.
     If no parameter is provided, applies weight normalization to all
     parameters in model (except 1-d vectors and scalars).

apex/reparameterization/weight_norm.py

@@ -20,7 +20,7 @@ def _norm(p, dim):
 HALF_TYPES = (torch.cuda.HalfTensor, torch.HalfTensor)
 
 class WeightNorm(Reparameterization):
-    """
+    r"""
     Weight normalization is a reparameterization that decouples the magnitude
     of a weight tensor from its direction. This replaces the parameter specified
     by `name` (e.g. "weight") with two parameters: one specifying the magnitude

csrc/amp_C_frontend.cpp

@@ -66,6 +66,17 @@ void multi_tensor_adam_cuda(
   const int bias_correction,
   const float weight_decay);
 
+
+void multi_tensor_adagrad_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  const float lr,
+  const float epsilon,
+  const int mode,
+  const float weight_decay);
+
+
 void multi_tensor_novograd_cuda(
   int chunk_size,
   at::Tensor noop_flag,
@@ -112,6 +123,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         "Completes application of gradient to parameters for LAMB optimizer");
   m.def("multi_tensor_adam", &multi_tensor_adam_cuda,
         "Compute and apply gradient update to parameters for Adam optimizer");
+  m.def("multi_tensor_adagrad", &multi_tensor_adagrad_cuda,
+        "Compute and apply gradient update to parameters for Adam optimizer");
   m.def("multi_tensor_novograd", &multi_tensor_novograd_cuda,
         "Compute and apply gradient update to parameters for Adam optimizer");
   m.def("multi_tensor_lamb", &multi_tensor_lamb_cuda,

csrc/multi_tensor_adagrad.cu

+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+
+#include "multi_tensor_apply.cuh"
+#include "type_shim.h"
+
+#define BLOCK_SIZE 1024
+#define ILP 4
+
+typedef enum {
+  ADAGRAD_MODE_0 = 0, // L2 regularization mode.
+  ADAGRAD_MODE_1 = 1, // AdamW-style weight decay.
+
+} adagradMode_t;
+
+using MATH_T = float;
+
+template <typename T> struct AdagradFunctor {
+  __device__ __forceinline__ void
+  operator()(int chunk_size, volatile int *noop_gmem, TensorListMetadata<3> &tl,
+             const float epsilon, const float lr, adagradMode_t mode,
+             const float weight_decay) {
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    T *g = (T *)tl.addresses[0][tensor_loc];
+    g += chunk_idx * chunk_size;
+
+    T *p = (T *)tl.addresses[1][tensor_loc];
+    p += chunk_idx * chunk_size;
+
+    T *h = (T *)tl.addresses[2][tensor_loc];
+    h += chunk_idx * chunk_size;
+
+    n -= chunk_idx * chunk_size;
+
+    // see note in multi_tensor_scale_kernel.cu
+    for (int i_start = 0; i_start < n && i_start < chunk_size;
+         i_start += blockDim.x * ILP) {
+      MATH_T r_g[ILP];
+      MATH_T r_p[ILP];
+      MATH_T r_h[ILP];
+#pragma unroll
+      for (int ii = 0; ii < ILP; ii++) {
+        int i = i_start + threadIdx.x + ii * blockDim.x;
+        if (i < n && i < chunk_size) {
+          r_g[ii] = g[i];
+          r_p[ii] = p[i];
+          r_h[ii] = h[i];
+        } else {
+          r_g[ii] = MATH_T(0);
+          r_p[ii] = MATH_T(0);
+          r_h[ii] = MATH_T(0);
+        }
+      }
+#pragma unroll
+      for (int ii = 0; ii < ILP; ii++) {
+        if (mode == ADAGRAD_MODE_0) { // L2
+          r_g[ii] = r_g[ii] + weight_decay * r_p[ii];
+          r_h[ii] = r_h[ii] + r_g[ii] * r_g[ii];
+          r_p[ii] = r_p[ii] - lr * (r_g[ii] / (sqrtf(r_h[ii]) + epsilon));
+        } else { // AdamW-style
+          r_h[ii] = r_h[ii] + r_g[ii] * r_g[ii];
+          r_p[ii] = r_p[ii] - lr * (r_g[ii] / (sqrtf(r_h[ii]) + epsilon) + weight_decay * r_p[ii]);
+        }
+      }
+#pragma unroll
+      for (int ii = 0; ii < ILP; ii++) {
+        int i = i_start + threadIdx.x + ii * blockDim.x;
+        if (i < n && i < chunk_size) {
+          p[i] = r_p[ii];
+          h[i] = r_h[ii];
+        }
+      }
+    }
+  }
+};
+
+void multi_tensor_adagrad_cuda(
+    int chunk_size, at::Tensor noop_flag,
+    std::vector<std::vector<at::Tensor>> tensor_lists, const float lr,
+    const float epsilon, const int mode, const float weight_decay) {
+  using namespace at;
+
+  // Assume single type across p,g,h now
+  DISPATCH_DOUBLE_FLOAT_AND_HALF(
+      tensor_lists[0][0].scalar_type(), 0, "adagrad",
+      multi_tensor_apply<3>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
+                            AdagradFunctor<scalar_t_0>(), epsilon, lr,
+                            (adagradMode_t)mode, weight_decay);)
+
+  AT_CUDA_CHECK(cudaGetLastError());
+}

csrc/type_shim.h

@@ -34,6 +34,32 @@
   }
 
 
+#define DISPATCH_FLOAT_HALF_AND_BYTE(TYPE, LEVEL, NAME, ...) \
+  switch(TYPE) \
+  { \
+    case at::ScalarType::Float: \
+    { \
+      using scalar_t_##LEVEL = float; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    case at::ScalarType::Half: \
+    { \
+      using scalar_t_##LEVEL = at::Half; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    case at::ScalarType::Byte: \
+    { \
+      using scalar_t_##LEVEL = uint8_t; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    default: \
+      AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'");  \
+  }
+
+
 #define DISPATCH_DOUBLE_FLOAT_AND_HALF(TYPE, LEVEL, NAME, ...) \
   switch(TYPE) \
   { \

requirements_dev.txt

+-r requirements.txt
+flake8>=3.7.9
+Sphinx>=3.0.3
\ No newline at end of file

setup.py

@@ -132,6 +132,7 @@ if "--cuda_ext" in sys.argv:
                                        'csrc/multi_tensor_lamb_stage_1.cu',
                                        'csrc/multi_tensor_lamb_stage_2.cu',
                                        'csrc/multi_tensor_adam.cu',
+                                       'csrc/multi_tensor_adagrad.cu',
                                        'csrc/multi_tensor_novograd.cu',
                                        'csrc/multi_tensor_lamb.cu'],
                               extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,

tests/L0/run_optimizers/test_adagrad.py

+import unittest
+
+import apex
+import torch
+
+
+class TestFusedAdagrad(unittest.TestCase):
+    def setUp(self, max_abs_diff=1e-6, max_rel_diff=1, iters=7):
+        self.max_abs_diff = max_abs_diff
+        self.max_rel_diff = max_rel_diff
+        self.iters = iters
+        torch.cuda.manual_seed(9876)
+
+    def tearDown(self):
+        pass
+
+    def gen_param_optim(self, tensors, adagrad_option):
+        ref_param = []
+        tst_param = []
+        for tensor in tensors:
+            ref_param.append(torch.nn.Parameter(tensor.clone()))
+            tst_param.append(torch.nn.Parameter(tensor.clone()))
+
+        ref_optim = torch.optim.Adagrad(ref_param, **adagrad_option)
+        tst_optim = apex.optimizers.FusedAdagrad(tst_param, **adagrad_option)
+
+        return (ref_param, tst_param, ref_optim, tst_optim)
+
+    def gen_grad(self, ref_param, tst_param):
+        for p_ref, p_tst in zip(ref_param, tst_param):
+            p_ref.grad = torch.rand_like(p_ref)
+            p_tst.grad = p_ref.grad
+
+    def gen_mixed_grad(self, ref_param, tst_param, scale=1.0):
+        half_grads = []
+        for p_ref, _ in zip(ref_param, tst_param):
+            half_grads.append(torch.rand_like(p_ref).half())
+            p_ref.grad = half_grads[-1].float() / scale
+        return half_grads
+
+    def get_max_diff(self, ref_param, tst_param):
+        max_abs_diff = max_rel_diff = 0
+        for p_ref, p_tst in zip(ref_param, tst_param):
+            max_abs_diff_p = (p_ref - p_tst).abs().max().item()
+            max_rel_diff_p = ((p_ref - p_tst) / p_ref).abs().max().item()
+
+            if max_abs_diff_p > max_abs_diff:
+                max_abs_diff = max_abs_diff_p
+            if max_rel_diff_p > max_rel_diff:
+                max_rel_diff = max_rel_diff_p
+
+        return max_abs_diff, max_rel_diff
+
+    def gen_single_type_test(self, param_type=torch.float):
+        nelem = 278011
+        adagrad_option = {"lr": 5e-4, "eps": 1e-08, "weight_decay": 1.0e-5}
+
+        tensor = torch.rand(nelem, dtype=param_type, device="cuda")
+        ref_param, tst_param, ref_optim, tst_optim = self.gen_param_optim(
+            [tensor], adagrad_option
+        )
+
+        for _ in range(self.iters):
+            self.gen_grad(ref_param, tst_param)
+            ref_optim.step()
+            tst_optim.step()
+            max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
+
+            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
+            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
+
+    def test_float(self):
+        self.gen_single_type_test(param_type=torch.float)
+
+    @unittest.skip("PyTorch optimizer is not numerically correct for fp16")
+    def test_half(self):
+        self.gen_single_type_test(param_type=torch.float16)
+
+    def test_multi_params(self):
+        sizes = [[4096, 1024], [4096], [4096, 2048], [32320, 1024], [1]]
+        adagrad_option = {"lr": 5e-4, "eps": 1e-08, "weight_decay": 0}
+
+        tensors = []
+        for size in sizes:
+            tensors.append(torch.rand(size, dtype=torch.float, device="cuda"))
+        ref_param, tst_param, ref_optim, tst_optim = self.gen_param_optim(
+            tensors, adagrad_option
+        )
+
+        for _ in range(self.iters):
+            self.gen_grad(ref_param, tst_param)
+            ref_optim.step()
+            tst_optim.step()
+            max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
+            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
+            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
+
+    def test_adagrad_option(self):
+        nelem = 1
+        adagrad_option = {"lr": 0.01, "eps": 3e-06, "weight_decay": 0}
+
+        tensor = torch.rand(nelem, dtype=torch.float, device="cuda")
+        ref_param, tst_param, ref_optim, tst_optim = self.gen_param_optim(
+            [tensor], adagrad_option
+        )
+
+        for _ in range(self.iters):
+            self.gen_grad(ref_param, tst_param)
+            ref_optim.step()
+            tst_optim.step()
+            max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
+
+            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
+            self.assertLessEqual(max_rel_diff, self.max_rel_diff)