diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index b75111fee9..2e7f11b575 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -289,6 +289,9 @@ if (onnxruntime_ENABLE_TRAINING) file(GLOB onnxruntime_python_ortmodule_torch_cpp_ext_torch_gpu_allocator_srcs CONFIGURE_DEPENDS "${ORTTRAINING_SOURCE_DIR}/python/training/ortmodule/torch_cpp_extensions/cuda/torch_gpu_allocator/*" ) + file(GLOB onnxruntime_python_ortmodule_torch_cpp_ext_adam_optimizer_srcs CONFIGURE_DEPENDS + "${ORTTRAINING_SOURCE_DIR}/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/*" + ) else() file(GLOB onnxruntime_python_capi_training_srcs CONFIGURE_DEPENDS "${ONNXRUNTIME_ROOT}/python/training/*.py" @@ -525,6 +528,7 @@ if (onnxruntime_ENABLE_TRAINING) COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/training/ortmodule/torch_cpp_extensions/cpu/torch_interop_utils COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/training/ortmodule/torch_cpp_extensions/cuda/torch_gpu_allocator + COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer COMMAND ${CMAKE_COMMAND} -E copy ${onnxruntime_python_capi_training_srcs} $/onnxruntime/capi/training/ @@ -561,6 +565,9 @@ if (onnxruntime_ENABLE_TRAINING) COMMAND ${CMAKE_COMMAND} -E copy ${onnxruntime_python_ortmodule_torch_cpp_ext_torch_gpu_allocator_srcs} $/onnxruntime/training/ortmodule/torch_cpp_extensions/cuda/torch_gpu_allocator/ + COMMAND ${CMAKE_COMMAND} -E copy + ${onnxruntime_python_ortmodule_torch_cpp_ext_adam_optimizer_srcs} + $/onnxruntime/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/ ) endif() diff --git a/orttraining/orttraining/python/training/optim/fused_adam.py b/orttraining/orttraining/python/training/optim/fused_adam.py new file mode 100644 index 0000000000..c6ab65ed5c --- /dev/null +++ b/orttraining/orttraining/python/training/optim/fused_adam.py @@ -0,0 +1,187 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# fused_adam.py +# This file has been adapted from microsoft/DeepSpeed + +''' +Copyright 2020 The Microsoft DeepSpeed Team + +Copyright NVIDIA/apex +This file is adapted from fused adam in NVIDIA/apex, commit a109f85 +''' + +import torch +import importlib +from .multi_tensor_apply import MultiTensorApply +multi_tensor_applier = MultiTensorApply(2048 * 32) + + +class FusedAdam(torch.optim.Optimizer): + """Implements Adam algorithm. + + Currently GPU-only. + + This version of fused Adam implements 2 fusions. + + * Fusion of the Adam 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. + + Adam was 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! + adam_w_mode (boolean, optional): Apply L2 regularization or weight decay + True for decoupled weight decay(also known as AdamW) (default: True) + set_grad_none (bool, optional): whether set grad to None when zero_grad() + method is called. (default: True) + + .. _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, + adam_w_mode=True, + weight_decay=0., + amsgrad=False, + set_grad_none=True): + + if amsgrad: + raise RuntimeError('FusedAdam does not support the AMSGrad variant.') + defaults = dict(lr=lr, + bias_correction=bias_correction, + betas=betas, + eps=eps, + weight_decay=weight_decay) + super(FusedAdam, self).__init__(params, defaults) + self.adam_w_mode = 1 if adam_w_mode else 0 + self.set_grad_none = set_grad_none + + # Skip buffer + self._dummy_overflow_buf = torch.cuda.IntTensor([0]) + + from onnxruntime.training.ortmodule.torch_cpp_extensions import adam_optimizer + self.multi_tensor_adam = adam_optimizer.multi_tensor_adam + + 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(FusedAdam, self).zero_grad() + + def step(self, + closure=None, + grads=None, + output_params=None, + scale=None, + grad_norms=None): + """Performs a single optimization step. + + Arguments: + closure (callable, optional): A closure that reevaluates the model + and returns the loss. + + The remaining arguments are deprecated, and are only retained (for the moment) for error-checking purposes. + """ + if any(p is not None for p in [grads, output_params, scale, grad_norms]): + raise RuntimeError( + 'FusedAdam has been updated. Simply initialize it identically to torch.optim.Adam, and call step() with no arguments.' + ) + loss = None + if closure is not None: + loss = closure() + + for group in self.param_groups: + bias_correction = 1 if group['bias_correction'] else 0 + beta1, beta2 = 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 + + # create lists for multi-tensor apply + g_16, p_16, m_16, v_16 = [], [], [], [] + g_32, p_32, m_32, v_32 = [], [], [], [] + + for p in group['params']: + if p.grad is None: + continue + if p.grad.data.is_sparse: + raise RuntimeError( + 'FusedAdam does not support sparse gradients, please consider SparseAdam instead' + ) + + state = self.state[p] + # State initialization + if len(state) == 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_like(p.data) + + if p.dtype == torch.float16: + g_16.append(p.grad.data) + p_16.append(p.data) + m_16.append(state['exp_avg']) + v_16.append(state['exp_avg_sq']) + elif p.dtype == torch.float32: + g_32.append(p.grad.data) + p_32.append(p.data) + m_32.append(state['exp_avg']) + v_32.append(state['exp_avg_sq']) + else: + raise RuntimeError('FusedAdam only support fp16 and fp32.') + + if (len(g_16) > 0): + multi_tensor_applier(self.multi_tensor_adam, + self._dummy_overflow_buf, + [g_16, + p_16, + m_16, + v_16], + group['lr'], + beta1, + beta2, + group['eps'], + group['step'], + self.adam_w_mode, + bias_correction, + group['weight_decay']) + if (len(g_32) > 0): + multi_tensor_applier(self.multi_tensor_adam, + self._dummy_overflow_buf, + [g_32, + p_32, + m_32, + v_32], + group['lr'], + beta1, + beta2, + group['eps'], + group['step'], + self.adam_w_mode, + bias_correction, + group['weight_decay']) + + return loss diff --git a/orttraining/orttraining/python/training/optim/multi_tensor_apply.py b/orttraining/orttraining/python/training/optim/multi_tensor_apply.py new file mode 100644 index 0000000000..9a5451427f --- /dev/null +++ b/orttraining/orttraining/python/training/optim/multi_tensor_apply.py @@ -0,0 +1,19 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# multi_tensor_apply.py +# This file has been adapted from microsoft/DeepSpeed + +''' +Copyright 2020 The Microsoft DeepSpeed Team + +Copyright NVIDIA/apex +This file is adapted from NVIDIA/apex, commit a109f85 +''' + + +class MultiTensorApply(object): + def __init__(self, chunk_size): + self.chunk_size = chunk_size + + def __call__(self, op, noop_flag_buffer, tensor_lists, *args): + return op(self.chunk_size, noop_flag_buffer, tensor_lists, *args) diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/fused_adam_frontend.cpp b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/fused_adam_frontend.cpp new file mode 100644 index 0000000000..b2157e5a33 --- /dev/null +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/fused_adam_frontend.cpp @@ -0,0 +1,24 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// This file is adapted from microsoft/DeepSpeed +// fused_adam_frontend.cpp + +#include + +void multi_tensor_adam_cuda(int chunk_size, + at::Tensor noop_flag, + std::vector> tensor_lists, + const float lr, + const float beta1, + const float beta2, + const float epsilon, + const int step, + const int mode, + const int bias_correction, + const float weight_decay); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) +{ + m.def("multi_tensor_adam", + &multi_tensor_adam_cuda, + "Compute and apply gradient update to parameters for Adam optimizer"); +} diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/multi_tensor_adam.cu b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/multi_tensor_adam.cu new file mode 100644 index 0000000000..be096f3b75 --- /dev/null +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/multi_tensor_adam.cu @@ -0,0 +1,167 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// This file is adapted from microsoft/DeepSpeed +// multi_tensor_adam.cu + +/* Copyright 2020 The Microsoft DeepSpeed Team + Copyright NVIDIA/apex + This file is adapted from fused adam in NVIDIA/apex, commit a109f85 +*/ + +#include +#include +#include +#include +// Another possibility: +// #include + +#include + +#include "multi_tensor_apply.cuh" +#include "type_shim.h" + +#define BLOCK_SIZE 512 +#define ILP 4 + +typedef enum { + ADAM_MODE_0 = 0, // L2 regularization mode + ADAM_MODE_1 = 1 // Decoupled weight decay mode(AdamW) +} adamMode_t; + +using MATH_T = float; + +template +struct AdamFunctor { + __device__ __forceinline__ void operator()(int chunk_size, + volatile int* noop_gmem, + TensorListMetadata<4>& tl, + const float beta1, + const float beta2, + const float beta1_correction, + const float beta2_correction, + const float epsilon, + const float lr, + adamMode_t mode, + const float decay) + { + // I'd like this kernel to propagate infs/nans. + // if(*noop_gmem == 1) + // return; + + int tensor_loc = tl.block_to_tensor[blockIdx.x]; + + // potentially use to pass in list of scalar + // int tensor_num = tl.start_tensor_this_launch + tensor_loc; + + 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* m = (T*)tl.addresses[2][tensor_loc]; + m += chunk_idx * chunk_size; + + T* v = (T*)tl.addresses[3][tensor_loc]; + v += 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_m[ILP]; + MATH_T r_v[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_m[ii] = m[i]; + r_v[ii] = v[i]; + } else { + r_g[ii] = MATH_T(0); + r_p[ii] = MATH_T(0); + r_m[ii] = MATH_T(0); + r_v[ii] = MATH_T(0); + } + } +#pragma unroll + for (int ii = 0; ii < ILP; ii++) { + if (mode == ADAM_MODE_0) { // L2 + r_g[ii] = r_g[ii] + (decay * r_p[ii]); + r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; + r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii]; + MATH_T next_m_unbiased = r_m[ii] / beta1_correction; + MATH_T next_v_unbiased = r_v[ii] / beta2_correction; + MATH_T denom = sqrtf(next_v_unbiased) + epsilon; + MATH_T update = next_m_unbiased / denom; + r_p[ii] = r_p[ii] - (lr * update); + } else { // weight decay + r_m[ii] = beta1 * r_m[ii] + (1 - beta1) * r_g[ii]; + r_v[ii] = beta2 * r_v[ii] + (1 - beta2) * r_g[ii] * r_g[ii]; + MATH_T next_m_unbiased = r_m[ii] / beta1_correction; + MATH_T next_v_unbiased = r_v[ii] / beta2_correction; + MATH_T denom = sqrtf(next_v_unbiased) + epsilon; + MATH_T update = (next_m_unbiased / denom) + (decay * r_p[ii]); + r_p[ii] = r_p[ii] - (lr * update); + } + } +#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]; + m[i] = r_m[ii]; + v[i] = r_v[ii]; + } + } + } + } +}; + +void multi_tensor_adam_cuda(int chunk_size, + at::Tensor noop_flag, + std::vector> tensor_lists, + const float lr, + const float beta1, + const float beta2, + const float epsilon, + const int step, + const int mode, + const int bias_correction, + const float weight_decay) +{ + using namespace at; + + // Handle bias correction mode + float bias_correction1 = 1.0f, bias_correction2 = 1.0f; + if (bias_correction == 1) { + bias_correction1 = 1 - std::pow(beta1, step); + bias_correction2 = 1 - std::pow(beta2, step); + } + + // Assume single type across p,g,m1,m2 now + DISPATCH_DOUBLE_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), + 0, + "adam", + multi_tensor_apply<4>(BLOCK_SIZE, + chunk_size, + noop_flag, + tensor_lists, + AdamFunctor(), + beta1, + beta2, + bias_correction1, + bias_correction2, + epsilon, + lr, + (adamMode_t)mode, + weight_decay);) + + AT_CUDA_CHECK(cudaGetLastError()); +} diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/multi_tensor_apply.cuh b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/multi_tensor_apply.cuh new file mode 100644 index 0000000000..fabbda2e58 --- /dev/null +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/multi_tensor_apply.cuh @@ -0,0 +1,128 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// This file is adapted from microsoft/DeepSpeed +// multi_tensory_apply.cuh + +/* Copyright 2020 The Microsoft DeepSpeed Team + Copyright NVIDIA/apex + This file is adapted from fused adam in NVIDIA/apex, commit a109f85 +*/ + +#include +#include +#include +#include +#include + +#include + +// This header is the one-stop shop for all your multi-tensor apply needs. + +// TODO: Kernel arg size limit may be <4KB for some other cards (ie Jetson) +constexpr int depth_to_max_tensors[5] = {110, 64, 48, 36, 30}; +constexpr int depth_to_max_blocks[5] = {320, 320, 320, 320, 320}; + +template +struct TensorListMetadata { + void* addresses[n][depth_to_max_tensors[n - 1]]; + int sizes[depth_to_max_tensors[n - 1]]; + unsigned char block_to_tensor[depth_to_max_blocks[n - 1]]; + int block_to_chunk[depth_to_max_blocks[n - 1]]; // I fear this needs to be a full int. + int start_tensor_this_launch; +}; + +template +__global__ void multi_tensor_apply_kernel(int chunk_size, + volatile int* noop_flag, + T tl, + U callable, + ArgTypes... args) +{ + // Hand the chunk information to the user-supplied functor to process however it likes. + callable(chunk_size, noop_flag, tl, args...); +} + +template +void multi_tensor_apply(int block_size, + int chunk_size, + const at::Tensor& noop_flag, + const std::vector>& tensor_lists, + T callable, + ArgTypes... args) +{ + TORCH_CHECK(tensor_lists.size() == depth, "tensor_lists.size() != depth"); + int len0 = tensor_lists[0].size(); + TORCH_CHECK(len0 > 0, "tensor_lists[0].size() is not > 0"); + auto ref_device = tensor_lists[0][0].device(); + TORCH_CHECK(ref_device.type() == at::kCUDA, "expected input to be on cuda"); + for (int l = 0; l < tensor_lists.size(); l++) // No range-based for because I need indices + { + TORCH_CHECK(tensor_lists[l].size() == len0, "Size mismatch among tensor lists"); + for (int t = 0; t < tensor_lists[l].size(); t++) { + // TODO: Print which tensor fails. + bool contiguous_memory = tensor_lists[l][t].is_contiguous(); +#ifdef VERSION_GE_1_5 + contiguous_memory = (contiguous_memory || + tensor_lists[l][t].is_contiguous(at::MemoryFormat::ChannelsLast)); +#endif + TORCH_CHECK(contiguous_memory, "A tensor was not contiguous."); + TORCH_CHECK(tensor_lists[l][t].device() == ref_device, + "A tensor was not on the same device as the first tensor"); + TORCH_CHECK(tensor_lists[l][t].numel() == tensor_lists[0][t].numel(), "Size mismatch"); + } + } + + int ntensors = tensor_lists[0].size(); + + TensorListMetadata tl; + + const at::cuda::OptionalCUDAGuard device_guard(device_of(tensor_lists[0][0])); + auto stream = at::cuda::getCurrentCUDAStream(); + + tl.start_tensor_this_launch = 0; + int loc_block_info = 0; + int loc_tensor_info = 0; + for (int t = 0; t < ntensors; t++) { + tl.sizes[loc_tensor_info] = tensor_lists[0][t].numel(); + for (int d = 0; d < depth; d++) + tl.addresses[d][loc_tensor_info] = tensor_lists[d][t].data_ptr(); + loc_tensor_info++; + + int chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size; + + for (int chunk = 0; chunk < chunks_this_tensor; chunk++) { + // std::cout << chunks_this_tensor << std::endl; + tl.block_to_tensor[loc_block_info] = loc_tensor_info - 1; + tl.block_to_chunk[loc_block_info] = chunk; + loc_block_info++; + + bool tensors_full = (loc_tensor_info == depth_to_max_tensors[depth - 1] && + chunk == chunks_this_tensor - 1); + bool blocks_full = (loc_block_info == depth_to_max_blocks[depth - 1]); + bool last_chunk = (t == ntensors - 1 && chunk == chunks_this_tensor - 1); + if (tensors_full || blocks_full || last_chunk) { + // using accscalar_t = acc_type; + multi_tensor_apply_kernel<<>>( + chunk_size, noop_flag.data_ptr(), tl, callable, args...); + + AT_CUDA_CHECK(cudaGetLastError()); + + // Reset. The control flow possibilities here make my brain hurt. + loc_block_info = 0; + if (chunk == chunks_this_tensor - 1) { + // std::cout << "Hit case 1 " << cond1 << " " << cond2 << " " << cond3 << + // std::endl; + loc_tensor_info = 0; + tl.start_tensor_this_launch = t + 1; + } else { + // std::cout << "Hit case 2 " << cond1 << " " << cond2 << " " << cond3 << + // std::endl; + tl.sizes[0] = tl.sizes[loc_tensor_info - 1]; + for (int d = 0; d < depth; d++) + tl.addresses[d][0] = tl.addresses[d][loc_tensor_info - 1]; + loc_tensor_info = 1; + tl.start_tensor_this_launch = t; + } + } + } + } +} diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/setup.py b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/setup.py new file mode 100644 index 0000000000..39c13f51b2 --- /dev/null +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/setup.py @@ -0,0 +1,30 @@ +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. +# -------------------------------------------------------------------------- + +import fileinput +import os +import sys + +from setuptools import setup +from torch.utils import cpp_extension + +filenames = [os.path.join(os.path.dirname(__file__), 'fused_adam_frontend.cpp'), + os.path.join(os.path.dirname(__file__), 'multi_tensor_adam.cu')] + +use_rocm = True if os.environ['ONNXRUNTIME_ROCM_VERSION'] else False +extra_compile_args = { + 'cxx': ['-O3'] +} +if not use_rocm: + extra_compile_args.update({ + 'nvcc': ['-lineinfo', '-O3', '--use_fast_math'] + }) + +setup(name='adam_optimizer', + ext_modules=[cpp_extension.CUDAExtension(name='adam_optimizer', + sources=filenames, + extra_compile_args=extra_compile_args + )], + cmdclass={'build_ext': cpp_extension.BuildExtension}) diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/type_shim.h b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/type_shim.h new file mode 100644 index 0000000000..07641e9597 --- /dev/null +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/adam_optimizer/type_shim.h @@ -0,0 +1,26 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// This file is adapted from microsoft/DeepSpeed +// type_shim.h + +/* Taken from NVIDIA/apex commit 855808f3fc268e9715d613f3c2e56469d8c986d8 */ +#include + +#define DISPATCH_DOUBLE_FLOAT_AND_HALF(TYPE, LEVEL, NAME, ...) \ + switch (TYPE) { \ + case at::ScalarType::Double: { \ + using scalar_t_##LEVEL = double; \ + __VA_ARGS__; \ + break; \ + } \ + 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; \ + } \ + default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ + } diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py index 2ba544ae9f..c21a4a9811 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_api.py @@ -30,6 +30,10 @@ from onnxruntime.training.ortmodule import (ORTModule, LogLevel, _fallback, _graph_execution_manager) + +from onnxruntime.training.optim.fused_adam import FusedAdam +from transformers import AdamW + import _test_helpers # Import autocasting libs @@ -4015,3 +4019,44 @@ def test_override_pytorch_exporter_kwargs_using_ortmodule_extension(): assert prediction is not None prediction = prediction.sum() prediction.backward() + +def test_ortmodule_fused_adam_optimizer_correctness(): + + device = 'cuda' + N, D_in, H, D_out = 32, 128, 500, 10 + + pt_model = NeuralNetSinglePositionalArgument(D_in, H, D_out).to(device) + transformers_adamw_optimizer = AdamW(pt_model.parameters()) + + ort_model = ORTModule(copy.deepcopy(pt_model)) + ort_fused_adam_optimizer = FusedAdam(ort_model.parameters()) + + def run_step(model, x): + prediction = model(x) + loss = prediction.sum() + loss.backward() + return prediction, loss + + def run_optim_step(optimizer): + optimizer.step() + optimizer.zero_grad() + + ga_steps = 2 + pt_model.zero_grad() + ort_model.zero_grad() + + for step in range(10): + x = torch.randn(N, D_in, device=device) + + _, pt_loss = run_step(pt_model, x) + _, ort_loss = run_step(ort_model, x) + + _test_helpers.assert_values_are_close(pt_loss, ort_loss, rtol=1e-4) + _test_helpers.assert_gradients_match_and_reset_gradient(ort_model, pt_model) + + if (step+1) % ga_steps == 0: + run_optim_step(transformers_adamw_optimizer) + run_optim_step(ort_fused_adam_optimizer) + + for pt_param, ort_param in zip(pt_model.parameters(), ort_model.parameters()): + _test_helpers.assert_values_are_close(pt_param, ort_param) diff --git a/setup.py b/setup.py index 145064dde7..f888bdbbb7 100644 --- a/setup.py +++ b/setup.py @@ -331,10 +331,13 @@ if enable_training: 'onnxruntime.training.ortmodule.torch_cpp_extensions', 'onnxruntime.training.ortmodule.torch_cpp_extensions.cpu.aten_op_executor', 'onnxruntime.training.ortmodule.torch_cpp_extensions.cpu.torch_interop_utils', - 'onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.torch_gpu_allocator']) + 'onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.torch_gpu_allocator', + 'onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.adam_optimizer']) package_data['onnxruntime.training.ortmodule.torch_cpp_extensions.cpu.aten_op_executor'] = ['*.cc'] package_data['onnxruntime.training.ortmodule.torch_cpp_extensions.cpu.torch_interop_utils'] = ['*.cc'] package_data['onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.torch_gpu_allocator'] = ['*.cc'] + package_data['onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.adam_optimizer'] = \ + ['*.cpp', '*.cu', '*.cuh', '*.h'] requirements_file = "requirements-training.txt" # with training, we want to follow this naming convention: # stable: