diff --git a/orttraining/orttraining/python/training/optim/_ds_modifier.py b/orttraining/orttraining/python/training/optim/_ds_modifier.py index 6ae6ccee51..c41f2ae412 100644 --- a/orttraining/orttraining/python/training/optim/_ds_modifier.py +++ b/orttraining/orttraining/python/training/optim/_ds_modifier.py @@ -37,7 +37,7 @@ class DeepSpeedZeROModifier(FP16OptimizerModifier): return self.check_requirements( ["has_overflow_serial", "get_grad_norm_direct", "has_overflow_partitioned_grads_serial"], - require_apex=True, + require_apex=False, require_torch_non_finite_check=True, ) @@ -45,7 +45,7 @@ class DeepSpeedZeROModifier(FP16OptimizerModifier): warnings.warn("DeepSpeed fp16_optimizer functions are overrided with faster implementation.", UserWarning) def get_grad_norm_direct(target, gradients, params, norm_type=2): - import amp_C + from onnxruntime.training.ortmodule.torch_cpp_extensions import fused_ops def is_model_parallel_parameter(p): return hasattr(p, "model_parallel") and p.model_parallel @@ -93,7 +93,10 @@ class DeepSpeedZeROModifier(FP16OptimizerModifier): # Multi-tensor applier takes a function and a list of list # and performs the operation on that list all in one kernel. grad_norm, _ = multi_tensor_applier( - amp_C.multi_tensor_l2norm, dummy_overflow_buf, [grads_for_norm], False # no per-parameter norm + fused_ops.multi_tensor_l2norm, + dummy_overflow_buf, + [fused_ops.TorchTensorVector(grads_for_norm)], + False, # no per-parameter norm ) # Since we will be summing across data parallel groups, # we need the pow(norm-type). diff --git a/orttraining/orttraining/python/training/optim/_megatron_modifier.py b/orttraining/orttraining/python/training/optim/_megatron_modifier.py index 545c669615..b6c5823110 100644 --- a/orttraining/orttraining/python/training/optim/_megatron_modifier.py +++ b/orttraining/orttraining/python/training/optim/_megatron_modifier.py @@ -11,7 +11,7 @@ import types import warnings -from numpy import inf + from ._modifier import FP16OptimizerModifier, check_overflow, clip_grad_norm_fp32 @@ -23,7 +23,7 @@ class LegacyMegatronLMModifier(FP16OptimizerModifier): def can_be_modified(self): return self.check_requirements( - ["_check_overflow", "clip_master_grads"], require_apex=True, require_torch_non_finite_check=True + ["_check_overflow", "clip_master_grads"], require_apex=False, require_torch_non_finite_check=True ) def override_function(self): diff --git a/orttraining/orttraining/python/training/optim/_modifier.py b/orttraining/orttraining/python/training/optim/_modifier.py index b3ad73110d..952e90ee43 100644 --- a/orttraining/orttraining/python/training/optim/_modifier.py +++ b/orttraining/orttraining/python/training/optim/_modifier.py @@ -8,9 +8,11 @@ # - check_overflow_for_grads : https://github.com/NVIDIA/Megatron-LM/blob/5ac5571ba0265af4c491ee0af1508ca7589450c6/megatron/optimizer/optimizer.py#L341 # -------------------------------------------------------------------------- -import torch import warnings + +import torch from numpy import inf + from ._multi_tensor_apply import MultiTensorApply multi_tensor_applier = MultiTensorApply(2048 * 32) @@ -66,7 +68,7 @@ def check_overflow_for_grads(grad_data): def clip_grad_norm_fp32( parameters, max_norm, norm_type, get_horizontal_model_parallel_rank=None, get_horizontal_model_parallel_group=None ): - import amp_C + from onnxruntime.training.ortmodule.torch_cpp_extensions import fused_ops horizontal_model_parallel_grad_norm_aggregation = False if get_horizontal_model_parallel_rank and get_horizontal_model_parallel_group: @@ -120,7 +122,10 @@ def clip_grad_norm_fp32( # Multi-tensor applier takes a function and a list of list # and performs the operation on that list all in one kernel. grad_norm, _ = multi_tensor_applier( - amp_C.multi_tensor_l2norm, dummy_overflow_buf, [grads_for_norm], False # no per-parameter norm + fused_ops.multi_tensor_l2norm, + dummy_overflow_buf, + [fused_ops.TorchTensorVector(grads_for_norm)], + False, # no per-parameter norm ) if not horizontal_model_parallel_grad_norm_aggregation: @@ -145,6 +150,7 @@ def clip_grad_norm_fp32( # Filter parameters with gradients. grads = [p.grad for p in parameters if p.grad is not None] if clip_coef < 1.0: - multi_tensor_applier(amp_C.multi_tensor_scale, dummy_overflow_buf, [grads, grads], clip_coef) + grads_vec = fused_ops.TorchTensorVector(grads) + multi_tensor_applier(fused_ops.multi_tensor_scale, dummy_overflow_buf, [grads_vec, grads_vec], clip_coef) return total_norm diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/fused_ops_frontend.cpp b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/fused_ops_frontend.cpp index 6111214676..19ba6b17ab 100644 --- a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/fused_ops_frontend.cpp +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/fused_ops_frontend.cpp @@ -42,6 +42,12 @@ void multi_tensor_axpby_cuda(int chunk_size, float b, int arg_to_check); +// This function is adapted from NVIDIA/apex +// https://github.com/NVIDIA/apex/blob/0c7d8e3fa9a095a1641a2290877436d0314b69c6/csrc/amp_C_frontend.cpp#L30 +std::tuple multi_tensor_l2norm_cuda(int chunk_size, at::Tensor noop_flag, + std::vector> tensor_lists, + at::optional per_tensor_python); + class MemoryBuffer { public: MemoryBuffer(size_t numel, at::Tensor val) { @@ -218,4 +224,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("unscale_fp16_grads_into_fp32_grads", &unscale_fp16_grads_into_fp32_grads, "Unscale those fp16 gradients into fp32 gradient buffers."); + m.def("multi_tensor_scale", + &multi_tensor_scale_cuda, + "Fused overflow check + scale for a list of contiguous tensors"); + m.def("multi_tensor_l2norm", + &multi_tensor_l2norm_cuda, + "Computes L2 norm for a list of contiguous tensors"); } diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/multi_tensor_l2norm_kernel.cu b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/multi_tensor_l2norm_kernel.cu new file mode 100644 index 0000000000..362ea3311a --- /dev/null +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/multi_tensor_l2norm_kernel.cu @@ -0,0 +1,167 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Copyright NVIDIA/apex +// This file is adapted from NVIDIA/apex, commit 3ff1a10f72ec07067c4e44759442329804ac5162 + +#include +#include +#include +#include +#include +// Another possibility: +// #include + +#include + +#include "type_shim.h" +#include "multi_tensor_apply.cuh" + +#define BLOCK_SIZE 512 +#define ILP 4 + +template +__device__ __forceinline__ bool is_aligned(T* p) { + return ((uint64_t)p) % (ILP * sizeof(T)) == 0; +} + +template +__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset) { + typedef typename std::aligned_storage::type LT; + ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset]; +} + +template +struct L2NormFunctor { + __device__ __forceinline__ void operator()(int chunk_size, volatile int* noop_gmem, TensorListMetadata<1>& tl, + float* output, float* output_per_tensor, bool per_tensor, + int max_chunks_per_tensor) { + // I'd like this kernel to propagate infs/nans. + // if(*noop_gmem == 1) + // return; + + int tensor_loc = tl.block_to_tensor[blockIdx.x]; + int chunk_idx = tl.block_to_chunk[blockIdx.x]; + int n = tl.sizes[tensor_loc]; + + x_t* x = (x_t*)tl.addresses[0][tensor_loc]; + x += chunk_idx * chunk_size; + + n -= chunk_idx * chunk_size; + + __shared__ float s_vals[512]; + + float vals[ILP]; // = {0}; // this probably works too but I want to be sure... + x_t r_x[ILP]; + for (int i = 0; i < ILP; i++) { + vals[i] = 0.f; + r_x[i] = 0; + } + + // to make things simple, we put aligned case in a different code path + if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x)) { + for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) { + // load + load_store(r_x, x, 0, i_start); +#pragma unroll + for (int ii = 0; ii < ILP; ii++) { + float next = static_cast(r_x[ii]); + vals[ii] += next * next; + } + } + } else { + for (int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x * 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) { + float next = static_cast(x[i]); + vals[ii] += next * next; + } + } + } + } + + float val = 0.f; + for (int i = 0; i < ILP; i++) val += vals[i]; + + float final = reduce_block_into_lanes(s_vals, val); + + if (threadIdx.x == 0) { + if (!isfinite(final)) *noop_gmem = 1; // Blindly fire off a write. These will race but that's ok. + output[blockIdx.x] += final; + if (per_tensor) + output_per_tensor[(tl.start_tensor_this_launch + tensor_loc) * max_chunks_per_tensor + chunk_idx] = final; + } + } +}; + +__global__ void cleanup(float* output, float* output_per_tensor, float* ret, float* ret_per_tensor, bool per_tensor, + int max_chunks_per_tensor) { + __shared__ float vals[512]; + + if (blockIdx.x == 0) { + float val = 0; + if (threadIdx.x < 320) val = output[threadIdx.x]; + + float final = reduce_block_into_lanes(vals, val); + + if (threadIdx.x == 0) *ret = sqrt(final); + } + + if (per_tensor) { + float* output_this_tensor = output_per_tensor + blockIdx.x * max_chunks_per_tensor; + + float val = 0; + for (int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x) val += output_this_tensor[i]; + + float final = reduce_block_into_lanes(vals, val); + + if (threadIdx.x == 0) ret_per_tensor[blockIdx.x] = sqrt(final); + } +} + +std::tuple multi_tensor_l2norm_cuda(int chunk_size, at::Tensor noop_flag, + std::vector> tensor_lists, + at::optional per_tensor_python) { + bool per_tensor = per_tensor_python.has_value() ? per_tensor_python.value() : false; + + auto float_options = tensor_lists[0][0].options().dtype(at::kFloat); + auto output = at::zeros({320}, float_options); + + at::Tensor output_per_tensor; + at::Tensor ret_per_tensor; + + int ntensors = tensor_lists[0].size(); + int max_chunks_per_tensor = -1; + + if (per_tensor) { + for (int t = 0; t < ntensors; t++) { + int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1) / chunk_size; + if (max_chunks_this_tensor > max_chunks_per_tensor) max_chunks_per_tensor = max_chunks_this_tensor; + } + output_per_tensor = at::zeros({ntensors * max_chunks_per_tensor}, float_options); + ret_per_tensor = at::empty({ntensors}, float_options); + } else { + ret_per_tensor = at::empty({0}, float_options); + } + + DISPATCH_DOUBLE_FLOAT_AND_HALF( + tensor_lists[0][0].scalar_type(), 0, "multi_tensor_l2norm_cuda", + multi_tensor_apply<1>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists, L2NormFunctor(), + output.data_ptr(), per_tensor ? output_per_tensor.data_ptr() : nullptr, + per_tensor, max_chunks_per_tensor);) + + AT_CUDA_CHECK(cudaGetLastError()); + // AT_CUDA_CHECK(cudaDeviceSynchronize()); + + // This involves one more small kernel launches, but will be negligible end to end. + // I could get rid of these by hacking the functor + multi tensor harness with persistence + // logic, but keeping it simple for now + auto ret = at::empty({1}, output.options()); + const at::cuda::OptionalCUDAGuard device_guard(device_of(output)); + auto stream = at::cuda::getCurrentCUDAStream(); + cleanup<<>>( + output.data_ptr(), per_tensor ? output_per_tensor.data_ptr() : nullptr, ret.data_ptr(), + per_tensor ? ret_per_tensor.data_ptr() : nullptr, per_tensor, max_chunks_per_tensor); + + return std::tuple(ret, ret_per_tensor); +} diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/setup.py b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/setup.py index d0202f4637..86a9369cd5 100644 --- a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/setup.py +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/setup.py @@ -15,6 +15,7 @@ filenames = [ os.path.join(os.path.dirname(__file__), "multi_tensor_adam.cu"), os.path.join(os.path.dirname(__file__), "multi_tensor_scale_kernel.cu"), os.path.join(os.path.dirname(__file__), "multi_tensor_axpby_kernel.cu"), + os.path.join(os.path.dirname(__file__), "multi_tensor_l2norm_kernel.cu"), ] use_rocm = True if os.environ["ONNXRUNTIME_ROCM_VERSION"] else False diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/type_shim.h b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/type_shim.h index 07641e9597..0f68709a10 100644 --- a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/type_shim.h +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/fused_ops/type_shim.h @@ -24,3 +24,47 @@ } \ default: AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \ } + +template +__device__ __forceinline__ T reduce_block_into_lanes(T* x, T val, int lanes = 1, + bool share_result = false) // lanes is intended to be <= 32. +{ + int tid = threadIdx.x + threadIdx.y * blockDim.x; + int blockSize = blockDim.x * blockDim.y; // blockSize is intended to be a multiple of 32. + + if (blockSize >= 64) { + x[tid] = val; + __syncthreads(); + } + +#pragma unroll + for (int i = (blockSize >> 1); i >= 64; i >>= 1) { + if (tid < i) x[tid] = x[tid] + x[tid + i]; + __syncthreads(); + } + + T final; + + if (tid < 32) { + if (blockSize >= 64) + final = x[tid] + x[tid + 32]; + else + final = val; + // __SYNCWARP(); + +#pragma unroll +#if defined(CUDA_VERSION) && CUDA_VERSION >= 9000 + for (int i = 16; i >= lanes; i >>= 1) final = final + __shfl_down_sync(0xffffffff, final, i); +#else + for (int i = 16; i >= lanes; i >>= 1) final = final + __shfl_down(final, i); +#endif + } + + if (share_result) { + if (tid < lanes) x[tid] = final; // EpilogueOp + // Make sure the smem result is visible to all warps. + __syncthreads(); + } + + return final; +}