diff --git a/orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu b/orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu index ddd5855e67..0bec959e14 100644 --- a/orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu +++ b/orttraining/orttraining/training_ops/cuda/math/softmax_grad_impl.cu @@ -28,20 +28,119 @@ namespace onnxruntime { namespace cuda { + template + __global__ void softmax_warp_backward(output_t* gradInput, const input_t* grad, const input_t* output, + int element_count, int batch_count) { + // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and warp_size of method + // warp_softmax_backward_kernel. + constexpr int next_power_of_two = 1 << log2_elements; + constexpr int WARP_SIZE = (next_power_of_two < GPU_WARP_SIZE) ? next_power_of_two : GPU_WARP_SIZE; + constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE; + #ifdef USE_ROCM + constexpr int WARP_BATCH = 1; + #else + constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1; + #endif + int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH; + + // batch_count might not be a multiple of WARP_BATCH. Check how + // many batches have to computed within this WARP. + int local_batches = batch_count - first_batch; + if (local_batches > WARP_BATCH) + local_batches = WARP_BATCH; + + // there might be multiple batches per warp. compute the index within the batch + int local_idx = threadIdx.x % WARP_SIZE; + + // the first element to process by the current thread + int thread_offset = first_batch * element_count + local_idx; + grad += thread_offset; + output += thread_offset; + gradInput += thread_offset; + + // The nested loops over WARP_BATCH and then WARP_ITERATIONS can be simplified to one loop, + // but I think doing so would obfuscate the logic of the algorithm, thus I chose to keep + // the nested loops. + // This should have no impact on performance because the loops are unrolled anyway. + + // load data from global memory + acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS]; + acc_t output_reg[WARP_BATCH][WARP_ITERATIONS]; + acc_t grad_output_reg[WARP_BATCH][WARP_ITERATIONS]; + for (int i = 0; i < WARP_BATCH; ++i) { + int batch_element_count = (i >= local_batches) ? 0 : element_count; + for (int it = 0; it < WARP_ITERATIONS; ++it) { + int element_index = local_idx + it * WARP_SIZE; + if (element_index < batch_element_count) { + grad_reg[i][it] = grad[i * element_count + it * WARP_SIZE]; + output_reg[i][it] = output[i * element_count + it * WARP_SIZE]; + grad_output_reg[i][it] = grad_reg[i][it] * output_reg[i][it]; + } else { + grad_reg[i][it] = acc_t(0); + output_reg[i][it] = acc_t(0); + grad_output_reg[i][it] = acc_t(0); + } + } + } + + acc_t sum[WARP_BATCH]; + if (!is_log_softmax) { + #pragma unroll + for (int i = 0; i < WARP_BATCH; ++i) { + sum[i] = grad_output_reg[i][0]; + #pragma unroll + for (int it = 1; it < WARP_ITERATIONS; ++it) { + sum[i] += grad_output_reg[i][it]; + } + } + warp_reduce(sum); + } + else { + #pragma unroll + for (int i = 0; i < WARP_BATCH; ++i) { + sum[i] = grad_reg[i][0]; + #pragma unroll + for (int it = 1; it < WARP_ITERATIONS; ++it) { + sum[i] += grad_reg[i][it]; + } + } + warp_reduce(sum); + } + + // store result + #pragma unroll + for (int i = 0; i < WARP_BATCH; ++i) { + if (i >= local_batches) + break; + #pragma unroll + for (int it = 0; it < WARP_ITERATIONS; ++it) { + int element_index = local_idx + it * WARP_SIZE; + if (element_index < element_count) { + // compute gradients + if (is_log_softmax) { + gradInput[i * element_count + it * WARP_SIZE] = (grad_reg[i][it] - std::exp(output_reg[i][it]) * sum[i]); + } else { + gradInput[i * element_count + it * WARP_SIZE] = (grad_reg[i][it] - sum[i] ) * output_reg[i][it]; + } + } + } + } + } + +// The function "softmax_warp_backward" saves intermediate results in float32 using registers to prevent recomputing, which can be beneficial for small shapes. +// However, for larger shapes, the usage of a large register resource can lead to low CUDA warp occupancy and poor performance. +// In contrast, "softmax_warp_backward_register_efficicent" recomputes intermediate results instead of saving them and also saves the inputs in float16 format to further reduce register usage. +// TODO: If the dimension to do softmax is greater than 2048, saving the input into shared memory can further reduce register usage. template -__global__ void softmax_warp_backward(output_t* gradInput, const input_t* grad, const input_t* output, +__global__ void softmax_warp_backward_register_efficicent(output_t* gradInput, const input_t* grad, const input_t* output, int element_count, int batch_count) { // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and warp_size of method // warp_softmax_backward_kernel. constexpr int next_power_of_two = 1 << log2_elements; constexpr int WARP_SIZE = (next_power_of_two < GPU_WARP_SIZE) ? next_power_of_two : GPU_WARP_SIZE; constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE; -#ifdef USE_ROCM constexpr int WARP_BATCH = 1; -#else - constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1; -#endif int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH; @@ -66,9 +165,8 @@ __global__ void softmax_warp_backward(output_t* gradInput, const input_t* grad, // This should have no impact on performance because the loops are unrolled anyway. // load data from global memory - acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS]; - acc_t output_reg[WARP_BATCH][WARP_ITERATIONS]; - acc_t grad_output_reg[WARP_BATCH][WARP_ITERATIONS]; + input_t grad_reg[WARP_BATCH][WARP_ITERATIONS]; + input_t output_reg[WARP_BATCH][WARP_ITERATIONS]; for (int i = 0; i < WARP_BATCH; ++i) { int batch_element_count = (i >= local_batches) ? 0 : element_count; for (int it = 0; it < WARP_ITERATIONS; ++it) { @@ -76,11 +174,9 @@ __global__ void softmax_warp_backward(output_t* gradInput, const input_t* grad, if (element_index < batch_element_count) { grad_reg[i][it] = grad[i * element_count + it * WARP_SIZE]; output_reg[i][it] = output[i * element_count + it * WARP_SIZE]; - grad_output_reg[i][it] = grad_reg[i][it] * output_reg[i][it]; } else { - grad_reg[i][it] = acc_t(0); - output_reg[i][it] = acc_t(0); - grad_output_reg[i][it] = acc_t(0); + grad_reg[i][it] = input_t(0); + output_reg[i][it] = input_t(0); } } } @@ -89,10 +185,10 @@ __global__ void softmax_warp_backward(output_t* gradInput, const input_t* grad, if (!is_log_softmax) { #pragma unroll for (int i = 0; i < WARP_BATCH; ++i) { - sum[i] = grad_output_reg[i][0]; + sum[i] = (acc_t)(grad_reg[i][0]) * (acc_t)(output_reg[i][0]); #pragma unroll for (int it = 1; it < WARP_ITERATIONS; ++it) { - sum[i] += grad_output_reg[i][it]; + sum[i] += (acc_t)(grad_reg[i][it]) * (acc_t)(output_reg[i][it]); } } warp_reduce(sum); @@ -100,10 +196,10 @@ __global__ void softmax_warp_backward(output_t* gradInput, const input_t* grad, else { #pragma unroll for (int i = 0; i < WARP_BATCH; ++i) { - sum[i] = grad_reg[i][0]; + sum[i] = (acc_t)grad_reg[i][0]; #pragma unroll for (int it = 1; it < WARP_ITERATIONS; ++it) { - sum[i] += grad_reg[i][it]; + sum[i] += (acc_t)grad_reg[i][it]; } } warp_reduce(sum); @@ -120,9 +216,9 @@ __global__ void softmax_warp_backward(output_t* gradInput, const input_t* grad, if (element_index < element_count) { // compute gradients if (is_log_softmax) { - gradInput[i * element_count + it * WARP_SIZE] = (grad_reg[i][it] - std::exp(output_reg[i][it]) * sum[i]); + gradInput[i * element_count + it * WARP_SIZE] = ((acc_t)(grad_reg[i][it]) - std::exp((acc_t)output_reg[i][it]) * sum[i]); } else { - gradInput[i * element_count + it * WARP_SIZE] = (grad_reg[i][it] - sum[i] ) * output_reg[i][it]; + gradInput[i * element_count + it * WARP_SIZE] = ((acc_t)grad_reg[i][it] - sum[i] ) * (acc_t)output_reg[i][it]; } } } @@ -133,7 +229,11 @@ template Status SoftmaxGradImpl(cudaStream_t stream, cudnnHandle_t cudnn_handle, T* input_grad, const T* output_grad, const T* softmax_output, int element_count, int batch_count, bool is_log_softmax) { if (element_count == 0) return Status::OK(); +#ifdef USE_ROCM if (element_count <= 1024 && element_count * sizeof(T) <= 4096) { +#else + if (element_count <= 2048 && element_count * sizeof(T) <= 4096) { +#endif typedef AccumulationType_t AccT; int log2_elements = log2_ceil(element_count); const int next_power_of_two = 1 << log2_elements; @@ -155,17 +255,26 @@ Status SoftmaxGradImpl(cudaStream_t stream, cudnnHandle_t cudnn_handle, T* input int blocks = (batch_count + batches_per_block - 1) / batches_per_block; dim3 threads(warp_size, warps_per_block, 1); // Launch code would be more elegant if C++ supported FOR CONSTEXPR + constexpr int start_to_use_register_efficient_func = 11; switch (log2_elements) { +#define LAUNCH_KERNEL(log2_elements_value, kernel_name) \ + if (is_log_softmax) { \ + kernel_name \ + <<>>(input_grad, output_grad, softmax_output, element_count, batch_count); \ + } else { \ + kernel_name \ + <<>>(input_grad, output_grad, softmax_output, element_count, batch_count); \ + } + #define CASE_LOG2_ELEMENTS(log2_elements_value) \ case log2_elements_value: { \ - if (is_log_softmax) { \ - softmax_warp_backward \ - <<>>(input_grad, output_grad, softmax_output, element_count, batch_count); \ + if (log2_elements_value < start_to_use_register_efficient_func) { \ + LAUNCH_KERNEL(log2_elements_value, softmax_warp_backward); \ } else { \ - softmax_warp_backward \ - <<>>(input_grad, output_grad, softmax_output, element_count, batch_count); \ + LAUNCH_KERNEL(log2_elements_value, softmax_warp_backward_register_efficicent); \ } \ } break + CASE_LOG2_ELEMENTS(0); // 1 CASE_LOG2_ELEMENTS(1); // 2 CASE_LOG2_ELEMENTS(2); // 4 @@ -177,6 +286,8 @@ Status SoftmaxGradImpl(cudaStream_t stream, cudnnHandle_t cudnn_handle, T* input CASE_LOG2_ELEMENTS(8); // 256 CASE_LOG2_ELEMENTS(9); // 512 CASE_LOG2_ELEMENTS(10); // 1024 + CASE_LOG2_ELEMENTS(11); // 2048, start to use softmax_warp_backward_register_efficicent, decided by value of start_to_use_register_efficient_func +#undef LAUNCH_KERNEL #undef CASE_LOG2_ELEMENTS } return Status::OK();