mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-06-30 17:22:21 +00:00
[NFC] polish colossalai/kernel/cuda_native/csrc/multi_tensor_sgd_kernel.cu code style (#978)
This commit is contained in:
binmakeswell
parent
18542b47fc
commit
f28c021376
@ -1,14 +1,15 @@
|
||||
// modified from https://github.com/NVIDIA/apex/blob/master/csrc/multi_tensor_sgd_kernel.cu
|
||||
// modified from
|
||||
// https://github.com/NVIDIA/apex/blob/master/csrc/multi_tensor_sgd_kernel.cu
|
||||
#include <ATen/ATen.h>
|
||||
#include <ATen/AccumulateType.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <ATen/cuda/Exceptions.h>
|
||||
#include "multi_tensor_apply.cuh"
|
||||
#include "compat.h"
|
||||
|
||||
#include <assert.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include "compat.h"
|
||||
#include "multi_tensor_apply.cuh"
|
||||
|
||||
#define BLOCK_SIZE 512
|
||||
#define ILP 4
|
||||
|
||||
@ -28,24 +29,13 @@
|
||||
* wd_after_momentum : apply weight decay _after_ momentum instead of before
|
||||
**/
|
||||
template <int N, typename T_grad, typename T_weight>
|
||||
struct SGDFunctor
|
||||
{
|
||||
struct SGDFunctor {
|
||||
__device__ __forceinline__ void operator()(
|
||||
int chunk_size,
|
||||
volatile int *noop_gmem,
|
||||
TensorListMetadata<N> &tl,
|
||||
float wd,
|
||||
float momentum,
|
||||
float dampening,
|
||||
float lr,
|
||||
bool nesterov,
|
||||
bool first_run,
|
||||
bool wd_after_momentum,
|
||||
float scale)
|
||||
{
|
||||
int chunk_size, volatile int *noop_gmem, TensorListMetadata<N> &tl,
|
||||
float wd, float momentum, float dampening, float lr, bool nesterov,
|
||||
bool first_run, bool wd_after_momentum, float scale) {
|
||||
// Early exit if we don't need to do anything
|
||||
if (*noop_gmem)
|
||||
return;
|
||||
if (*noop_gmem) return;
|
||||
|
||||
int tensor_loc = tl.block_to_tensor[blockIdx.x];
|
||||
int chunk_idx = tl.block_to_chunk[blockIdx.x];
|
||||
@ -61,8 +51,7 @@ struct SGDFunctor
|
||||
mom_in += chunk_idx * chunk_size;
|
||||
|
||||
at::Half *model_weights_out = nullptr;
|
||||
if (N == 4)
|
||||
{
|
||||
if (N == 4) {
|
||||
model_weights_out = (at::Half *)tl.addresses[3][tensor_loc];
|
||||
model_weights_out += chunk_idx * chunk_size;
|
||||
}
|
||||
@ -73,19 +62,15 @@ struct SGDFunctor
|
||||
float incoming_grads[ILP];
|
||||
float incoming_weights[ILP];
|
||||
float incoming_moms[ILP];
|
||||
for (int i_start = 0;
|
||||
i_start < n && i_start < chunk_size;
|
||||
i_start += blockDim.x * ILP)
|
||||
{
|
||||
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++)
|
||||
{
|
||||
for (int ii = 0; ii < ILP; ii++) {
|
||||
incoming_grads[ii] = 0;
|
||||
incoming_weights[ii] = 0;
|
||||
incoming_moms[ii] = 0;
|
||||
int i = i_start + threadIdx.x + ii * blockDim.x;
|
||||
if (i < n && i < chunk_size)
|
||||
{
|
||||
if (i < n && i < chunk_size) {
|
||||
incoming_grads[ii] = static_cast<float>(grad_in[i]) * scale;
|
||||
incoming_weights[ii] = static_cast<float>(weight_in[i]);
|
||||
incoming_moms[ii] = static_cast<float>(mom_in[i]);
|
||||
@ -98,19 +83,17 @@ struct SGDFunctor
|
||||
// Put another way, the STGs are dependent on the LDGs, but not on each other.
|
||||
// There is still compute ILP benefit from unrolling the loop though.
|
||||
#pragma unroll
|
||||
for (int ii = 0; ii < ILP; ii++)
|
||||
{
|
||||
for (int ii = 0; ii < ILP; ii++) {
|
||||
int i = i_start + threadIdx.x + ii * blockDim.x;
|
||||
if (i < n && i < chunk_size)
|
||||
{
|
||||
if (i < n && i < chunk_size) {
|
||||
// apply weight decay before momentum if necessary
|
||||
if (wd != 0.f && !wd_after_momentum)
|
||||
incoming_grads[ii] += wd * incoming_weights[ii];
|
||||
|
||||
if (momentum != 0.f)
|
||||
{
|
||||
if (momentum != 0.f) {
|
||||
if (!first_run)
|
||||
incoming_moms[ii] = incoming_moms[ii] * momentum + (1.f - dampening) * incoming_grads[ii];
|
||||
incoming_moms[ii] = incoming_moms[ii] * momentum +
|
||||
(1.f - dampening) * incoming_grads[ii];
|
||||
else // initialize momentums to current incoming grads
|
||||
incoming_moms[ii] = incoming_grads[ii];
|
||||
|
||||
@ -132,27 +115,18 @@ struct SGDFunctor
|
||||
model_weights_out[i] = static_cast<at::Half>(weight_in[i]);
|
||||
|
||||
// also write out the new momentum
|
||||
if (momentum != 0.f)
|
||||
mom_in[i] = incoming_moms[ii];
|
||||
if (momentum != 0.f) mom_in[i] = incoming_moms[ii];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
void multi_tensor_sgd_cuda(
|
||||
int chunk_size,
|
||||
at::Tensor noop_flag,
|
||||
void multi_tensor_sgd_cuda(int chunk_size, at::Tensor noop_flag,
|
||||
std::vector<std::vector<at::Tensor>> tensor_lists,
|
||||
float wd,
|
||||
float momentum,
|
||||
float dampening,
|
||||
float lr,
|
||||
bool nesterov,
|
||||
bool first_run,
|
||||
bool wd_after_momentum,
|
||||
float scale)
|
||||
{
|
||||
float wd, float momentum, float dampening, float lr,
|
||||
bool nesterov, bool first_run,
|
||||
bool wd_after_momentum, float scale) {
|
||||
auto num_tensors = tensor_lists.size();
|
||||
auto grad_type = tensor_lists[0][0].scalar_type();
|
||||
auto weight_type = tensor_lists[1][0].scalar_type();
|
||||
@ -162,7 +136,8 @@ void multi_tensor_sgd_cuda(
|
||||
TORCH_CHECK(tensor_lists[3][i].scalar_type() == at::ScalarType::Half,
|
||||
"Additional output tensors should always be fp16.");
|
||||
|
||||
TORCH_CHECK(noop_flag.device() == tensor_lists[0][0].device(), "expected noop flag to be on the same device as tensors");
|
||||
TORCH_CHECK(noop_flag.device() == tensor_lists[0][0].device(),
|
||||
"expected noop flag to be on the same device as tensors");
|
||||
|
||||
// We have 3 possibilities to handle here, in terms of
|
||||
// grad_type, param_type, momentum_type, requires_fp16_copy
|
||||
@ -176,22 +151,10 @@ void multi_tensor_sgd_cuda(
|
||||
|
||||
// Case 1. fp16, fp16, fp16, No
|
||||
if (grad_type == at::ScalarType::Half &&
|
||||
weight_type == at::ScalarType::Half &&
|
||||
num_tensors == 3)
|
||||
{
|
||||
multi_tensor_apply<3>(
|
||||
BLOCK_SIZE,
|
||||
chunk_size,
|
||||
noop_flag,
|
||||
tensor_lists,
|
||||
SGDFunctor<3, at::Half, at::Half>(),
|
||||
wd,
|
||||
momentum,
|
||||
dampening,
|
||||
lr,
|
||||
nesterov,
|
||||
first_run,
|
||||
wd_after_momentum,
|
||||
weight_type == at::ScalarType::Half && num_tensors == 3) {
|
||||
multi_tensor_apply<3>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
|
||||
SGDFunctor<3, at::Half, at::Half>(), wd, momentum,
|
||||
dampening, lr, nesterov, first_run, wd_after_momentum,
|
||||
scale);
|
||||
}
|
||||
// Case 2. fp16, fp32, fp32, No
|
||||
@ -214,68 +177,33 @@ void multi_tensor_sgd_cuda(
|
||||
// }
|
||||
// Case 2. fp32, fp32, fp32, No
|
||||
else if (grad_type == at::ScalarType::Float &&
|
||||
weight_type == at::ScalarType::Float &&
|
||||
num_tensors == 3)
|
||||
{
|
||||
multi_tensor_apply<3>(
|
||||
BLOCK_SIZE,
|
||||
chunk_size,
|
||||
noop_flag,
|
||||
tensor_lists,
|
||||
SGDFunctor<3, float, float>(),
|
||||
wd,
|
||||
momentum,
|
||||
dampening,
|
||||
lr,
|
||||
nesterov,
|
||||
first_run,
|
||||
wd_after_momentum,
|
||||
weight_type == at::ScalarType::Float && num_tensors == 3) {
|
||||
multi_tensor_apply<3>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
|
||||
SGDFunctor<3, float, float>(), wd, momentum,
|
||||
dampening, lr, nesterov, first_run, wd_after_momentum,
|
||||
scale);
|
||||
}
|
||||
// Case 3. fp16, fp32, fp32, Yes
|
||||
else if (grad_type == at::ScalarType::Half &&
|
||||
weight_type == at::ScalarType::Float &&
|
||||
num_tensors == 4)
|
||||
{
|
||||
multi_tensor_apply<4>(
|
||||
BLOCK_SIZE,
|
||||
chunk_size,
|
||||
noop_flag,
|
||||
tensor_lists,
|
||||
SGDFunctor<4, at::Half, float>(),
|
||||
wd,
|
||||
momentum,
|
||||
dampening,
|
||||
lr,
|
||||
nesterov,
|
||||
first_run,
|
||||
wd_after_momentum,
|
||||
weight_type == at::ScalarType::Float && num_tensors == 4) {
|
||||
multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
|
||||
SGDFunctor<4, at::Half, float>(), wd, momentum,
|
||||
dampening, lr, nesterov, first_run, wd_after_momentum,
|
||||
scale);
|
||||
}
|
||||
// Case 4. fp32, fp32, fp32, Yes
|
||||
else if (grad_type == at::ScalarType::Float &&
|
||||
weight_type == at::ScalarType::Float &&
|
||||
num_tensors == 4)
|
||||
{
|
||||
multi_tensor_apply<4>(
|
||||
BLOCK_SIZE,
|
||||
chunk_size,
|
||||
noop_flag,
|
||||
tensor_lists,
|
||||
SGDFunctor<4, float, float>(),
|
||||
wd,
|
||||
momentum,
|
||||
dampening,
|
||||
lr,
|
||||
nesterov,
|
||||
first_run,
|
||||
wd_after_momentum,
|
||||
weight_type == at::ScalarType::Float && num_tensors == 4) {
|
||||
multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
|
||||
SGDFunctor<4, float, float>(), wd, momentum,
|
||||
dampening, lr, nesterov, first_run, wd_after_momentum,
|
||||
scale);
|
||||
}
|
||||
else
|
||||
{
|
||||
AT_ERROR("multi_tensor_sgd only supports some combinations of gradient & weight types. Given: ",
|
||||
"gradient: ", grad_type, ", weight: ", weight_type, ", num_lists: ", num_tensors);
|
||||
} else {
|
||||
AT_ERROR(
|
||||
"multi_tensor_sgd only supports some combinations of gradient & weight "
|
||||
"types. Given: ",
|
||||
"gradient: ", grad_type, ", weight: ", weight_type,
|
||||
", num_lists: ", num_tensors);
|
||||
}
|
||||
|
||||
AT_CUDA_CHECK(cudaGetLastError());
|
||||
|
||||
Loading…
Reference in New Issue
Block a user