Skip to content

GitLab

"...Chat/git@developer.sourcefind.cn:OpenDAS/colossalai.git" did not exist on "e4fc57c3de6204ac66df75aa1752db1ec284f31f"
Commit c1bed0d9 authored by superhao1995's avatar superhao1995 Committed by binmakeswell
Browse files

[NFC] polish colossalai/kernel/cuda_native/csrc/multi_tensor_lamb.cu code stype (#628)

parent 0a96338b
Hide whitespace changes
Inline Side-by-side
Showing with 304 additions and 377 deletions
colossalai/kernel/cuda_native/csrc/multi_tensor_lamb.cu
View file @ c1bed0d9
// modified from https://github.com/NVIDIA/apex/blob/master/csrc/multi_tensor_lamb.cu // modified from
// https://github.com/NVIDIA/apex/blob/master/csrc/multi_tensor_lamb.cu
#include <ATen/ATen.h> #include <ATen/ATen.h>
#include <ATen/AccumulateType.h> #include <ATen/AccumulateType.h>
#include <ATen/cuda/CUDAContext.h> #include <ATen/cuda/CUDAContext.h>
...@@ -8,420 +9,346 @@ ...@@ -8,420 +9,346 @@
#include <assert.h> #include <assert.h>
#include "type_shim.h"
#include "multi_tensor_apply.cuh" #include "multi_tensor_apply.cuh"
#include "type_shim.h"
#define BLOCK_SIZE 512 #define BLOCK_SIZE 512
#define ILP 4 #define ILP 4
template <typename T> template <typename T> __device__ __forceinline__ bool is_aligned(T *p) {
__device__ __forceinline__ bool is_aligned(T *p) return ((uint64_t)p) % (ILP * sizeof(T)) == 0;
{
return ((uint64_t)p) % (ILP * sizeof(T)) == 0;
} }
template <typename T> template <typename T>
__device__ __forceinline__ void load_store(T *dst, T *src, int dst_offset, int src_offset) __device__ __forceinline__ void load_store(T *dst, T *src, int dst_offset,
{ int src_offset) {
typedef typename std::aligned_storage<ILP * sizeof(T), ILP * alignof(T)>::type LT; typedef
((LT *)dst)[dst_offset] = ((LT *)src)[src_offset]; typename std::aligned_storage<ILP * sizeof(T), ILP * alignof(T)>::type LT;
((LT *)dst)[dst_offset] = ((LT *)src)[src_offset];
} }
typedef enum typedef enum {
{ MOMENT_MODE_0 = 0, // L2 regularization mode
MOMENT_MODE_0 = 0, // L2 regularization mode MOMENT_MODE_1 = 1 // Decoupled weight decay mode
MOMENT_MODE_1 = 1 // Decoupled weight decay mode
} adamMode_t; } adamMode_t;
std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_cuda( std::tuple<at::Tensor, at::Tensor>
int chunk_size, multi_tensor_l2norm_cuda(int chunk_size, at::Tensor noop_flag,
at::Tensor noop_flag, std::vector<std::vector<at::Tensor>> tensor_lists,
std::vector<std::vector<at::Tensor>> tensor_lists, at::optional<bool> per_tensor_python);
at::optional<bool> per_tensor_python);
using MATH_T = float; using MATH_T = float;
template <typename T> template <typename T> struct LAMBStage1Functor {
struct LAMBStage1Functor __device__ __forceinline__ void
{ operator()(int chunk_size, volatile int *noop_gmem, TensorListMetadata<4> &tl,
__device__ __forceinline__ void operator()( const float beta1, const float beta2, const float beta3,
int chunk_size, const float beta1_correction, const float beta2_correction,
volatile int *noop_gmem, const float epsilon, adamMode_t mode, const float decay,
TensorListMetadata<4> &tl, const float *global_grad_norm, const float max_global_grad_norm) {
const float beta1, // I'd like this kernel to propagate infs/nans.
const float beta2, // if(*noop_gmem == 1)
const float beta3, // return;
const float beta1_correction,
const float beta2_correction, int tensor_loc = tl.block_to_tensor[blockIdx.x];
const float epsilon, int chunk_idx = tl.block_to_chunk[blockIdx.x];
adamMode_t mode, int n = tl.sizes[tensor_loc];
const float decay,
const float *global_grad_norm, float clipped_global_grad_norm =
const float max_global_grad_norm) (*global_grad_norm) > max_global_grad_norm
{ ? (*global_grad_norm) / max_global_grad_norm
// I'd like this kernel to propagate infs/nans. : 1.0f;
// if(*noop_gmem == 1)
// return; T *g = (T *)tl.addresses[0][tensor_loc];
g += chunk_idx * chunk_size;
int tensor_loc = tl.block_to_tensor[blockIdx.x];
int chunk_idx = tl.block_to_chunk[blockIdx.x]; T *p = (T *)tl.addresses[1][tensor_loc];
int n = tl.sizes[tensor_loc]; p += chunk_idx * chunk_size;
float clipped_global_grad_norm = (*global_grad_norm) > max_global_grad_norm ? (*global_grad_norm) / max_global_grad_norm : 1.0f; T *m = (T *)tl.addresses[2][tensor_loc];
m += chunk_idx * chunk_size;
T *g = (T *)tl.addresses[0][tensor_loc];
g += chunk_idx * chunk_size; T *v = (T *)tl.addresses[3][tensor_loc];
v += chunk_idx * chunk_size;
T *p = (T *)tl.addresses[1][tensor_loc];
p += chunk_idx * chunk_size; n -= chunk_idx * chunk_size;
T *m = (T *)tl.addresses[2][tensor_loc]; MATH_T r_g[ILP];
m += chunk_idx * chunk_size; MATH_T r_p[ILP];
MATH_T r_m[ILP];
T *v = (T *)tl.addresses[3][tensor_loc]; MATH_T r_v[ILP];
v += chunk_idx * chunk_size; // to make things simple, we put aligned case in a different code path
if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(g) &&
n -= chunk_idx * chunk_size; is_aligned(p) && is_aligned(m) && is_aligned(v)) {
T l_g[ILP];
T l_p[ILP];
T l_m[ILP];
T l_v[ILP];
for (int i_start = threadIdx.x;
i_start * ILP < n && i_start * ILP < chunk_size;
i_start += blockDim.x) {
// load
load_store(l_g, g, 0, i_start);
if (decay != 0)
load_store(l_p, p, 0, i_start);
load_store(l_m, m, 0, i_start);
load_store(l_v, v, 0, i_start);
// unpack
#pragma unroll
for (int ii = 0; ii < ILP; ii++) {
r_g[ii] = l_g[ii];
if (decay == 0) {
r_p[ii] = MATH_T(0);
} else {
r_p[ii] = l_p[ii];
}
r_m[ii] = l_m[ii];
r_v[ii] = l_v[ii];
}
#pragma unroll
for (int ii = 0; ii < ILP; ii++) {
if (mode == MOMENT_MODE_0) {
MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;
// L2 on scaled grad
scaled_grad = scaled_grad + decay * r_p[ii];
r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
r_v[ii] = r_v[ii] * beta2 + (1 - beta2) * scaled_grad * scaled_grad;
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;
r_p[ii] = next_m_unbiased / denom;
} else {
MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;
r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
r_v[ii] = r_v[ii] * beta2 + (1 - beta2) * scaled_grad * scaled_grad;
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;
r_p[ii] = (next_m_unbiased / denom) + (decay * r_p[ii]);
}
}
#pragma unroll
for (int ii = 0; ii < ILP; ii++) {
l_p[ii] = r_p[ii];
l_m[ii] = r_m[ii];
l_v[ii] = r_v[ii];
}
// store
load_store(g, l_p, i_start, 0);
load_store(m, l_m, i_start, 0);
load_store(v, l_v, i_start, 0);
}
} else {
// 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_g[ILP];
MATH_T r_p[ILP]; MATH_T r_p[ILP];
MATH_T r_m[ILP]; MATH_T r_m[ILP];
MATH_T r_v[ILP]; MATH_T r_v[ILP];
// to make things simple, we put aligned case in a different code path
if (n % ILP == 0 &&
chunk_size % ILP == 0 &&
is_aligned(g) &&
is_aligned(p) &&
is_aligned(m) &&
is_aligned(v))
{
T l_g[ILP];
T l_p[ILP];
T l_m[ILP];
T l_v[ILP];
for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x)
{
// load
load_store(l_g, g, 0, i_start);
if (decay != 0)
load_store(l_p, p, 0, i_start);
load_store(l_m, m, 0, i_start);
load_store(l_v, v, 0, i_start);
// unpack
#pragma unroll
for (int ii = 0; ii < ILP; ii++)
{
r_g[ii] = l_g[ii];
if (decay == 0)
{
r_p[ii] = MATH_T(0);
}
else
{
r_p[ii] = l_p[ii];
}
r_m[ii] = l_m[ii];
r_v[ii] = l_v[ii];
}
#pragma unroll
for (int ii = 0; ii < ILP; ii++)
{
if (mode == MOMENT_MODE_0)
{
MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;
// L2 on scaled grad
scaled_grad = scaled_grad + decay * r_p[ii];
r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
r_v[ii] = r_v[ii] * beta2 + (1 - beta2) * scaled_grad * scaled_grad;
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;
r_p[ii] = next_m_unbiased / denom;
}
else
{
MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;
r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
r_v[ii] = r_v[ii] * beta2 + (1 - beta2) * scaled_grad * scaled_grad;
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;
r_p[ii] = (next_m_unbiased / denom) + (decay * r_p[ii]);
}
}
#pragma unroll #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;
l_p[ii] = r_p[ii]; if (i < n && i < chunk_size) {
l_m[ii] = r_m[ii]; r_g[ii] = g[i];
l_v[ii] = r_v[ii]; // special ?optimization? for lamb stage 1
} if (decay == 0) {
// store r_p[ii] = MATH_T(0);
load_store(g, l_p, i_start, 0); } else {
load_store(m, l_m, i_start, 0); r_p[ii] = p[i];
load_store(v, l_v, i_start, 0);
} }
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);
}
} }
else
{
// 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 #pragma unroll
for (int ii = 0; ii < ILP; ii++) for (int ii = 0; ii < ILP; ii++) {
{ if (mode == MOMENT_MODE_0) {
int i = i_start + threadIdx.x + ii * blockDim.x; MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;
if (i < n && i < chunk_size) // L2 on scaled grad
{ scaled_grad = scaled_grad + decay * r_p[ii];
r_g[ii] = g[i]; r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
// special ?optimization? for lamb stage 1 r_v[ii] = r_v[ii] * beta2 + (1 - beta2) * scaled_grad * scaled_grad;
if (decay == 0) MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
{ MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
r_p[ii] = MATH_T(0); MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
} r_p[ii] = next_m_unbiased / denom;
else } else {
{ MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;
r_p[ii] = p[i]; r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
} r_v[ii] = r_v[ii] * beta2 + (1 - beta2) * scaled_grad * scaled_grad;
r_m[ii] = m[i]; MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
r_v[ii] = v[i]; MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
} MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
else r_p[ii] = (next_m_unbiased / denom) + (decay * r_p[ii]);
{ }
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 == MOMENT_MODE_0)
{
MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;
// L2 on scaled grad
scaled_grad = scaled_grad + decay * r_p[ii];
r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
r_v[ii] = r_v[ii] * beta2 + (1 - beta2) * scaled_grad * scaled_grad;
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;
r_p[ii] = next_m_unbiased / denom;
}
else
{
MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;
r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
r_v[ii] = r_v[ii] * beta2 + (1 - beta2) * scaled_grad * scaled_grad;
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;
r_p[ii] = (next_m_unbiased / denom) + (decay * r_p[ii]);
}
}
#pragma unroll #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;
int i = i_start + threadIdx.x + ii * blockDim.x; if (i < n && i < chunk_size) {
if (i < n && i < chunk_size) g[i] = r_p[ii];
{ m[i] = r_m[ii];
g[i] = r_p[ii]; v[i] = r_v[ii];
m[i] = r_m[ii]; }
v[i] = r_v[ii];
}
}
}
} }
}
} }
}
}; };
// Step 2 reads in 'update' value and per-tensor param_norm and update_norm. // Step 2 reads in 'update' value and per-tensor param_norm and update_norm.
// It computes new parameter value. // It computes new parameter value.
template <typename T> template <typename T> struct LAMBStage2Functor {
struct LAMBStage2Functor __device__ __forceinline__ void
{ operator()(int chunk_size, volatile int *noop_gmem, TensorListMetadata<2> &tl,
__device__ __forceinline__ void operator()( const float *per_tensor_param_norm,
int chunk_size, const float *per_tensor_update_norm, const float learning_rate,
volatile int *noop_gmem, const float decay, bool use_nvlamb) {
TensorListMetadata<2> &tl, // I'd like this kernel to propagate infs/nans.
const float *per_tensor_param_norm, // if(*noop_gmem == 1)
const float *per_tensor_update_norm, // return;
const float learning_rate,
const float decay, int tensor_loc = tl.block_to_tensor[blockIdx.x];
bool use_nvlamb) int tensor_num = tl.start_tensor_this_launch + tensor_loc;
{ int chunk_idx = tl.block_to_chunk[blockIdx.x];
// I'd like this kernel to propagate infs/nans. int n = tl.sizes[tensor_loc];
// if(*noop_gmem == 1)
// return; MATH_T ratio = learning_rate;
// nvlamb: apply adaptive learning rate to all parameters
int tensor_loc = tl.block_to_tensor[blockIdx.x]; // otherwise, only apply to those with non-zero weight decay
int tensor_num = tl.start_tensor_this_launch + tensor_loc; if (use_nvlamb || (decay != 0.0)) {
int chunk_idx = tl.block_to_chunk[blockIdx.x]; float param_norm = per_tensor_param_norm[tensor_num];
int n = tl.sizes[tensor_loc]; float update_norm = per_tensor_update_norm[tensor_num];
ratio = (update_norm != 0.0f && param_norm != 0.0f)
MATH_T ratio = learning_rate; ? learning_rate * (param_norm / update_norm)
// nvlamb: apply adaptive learning rate to all parameters : learning_rate;
// otherwise, only apply to those with non-zero weight decay }
if (use_nvlamb || (decay != 0.0))
{
float param_norm = per_tensor_param_norm[tensor_num];
float update_norm = per_tensor_update_norm[tensor_num];
ratio = (update_norm != 0.0f && param_norm != 0.0f) ? learning_rate * (param_norm / update_norm) : learning_rate;
}
T *update = (T *)tl.addresses[0][tensor_loc]; T *update = (T *)tl.addresses[0][tensor_loc];
update += chunk_idx * chunk_size; update += chunk_idx * chunk_size;
T *p = (T *)tl.addresses[1][tensor_loc]; T *p = (T *)tl.addresses[1][tensor_loc];
p += chunk_idx * chunk_size; p += chunk_idx * chunk_size;
n -= chunk_idx * chunk_size; n -= chunk_idx * chunk_size;
// to make things simple, we put aligned case in a different code path // to make things simple, we put aligned case in a different code path
if (n % ILP == 0 && if (n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(p) &&
chunk_size % ILP == 0 && is_aligned(update)) {
is_aligned(p) && T r_p[ILP];
is_aligned(update)) T r_update[ILP];
{ for (int i_start = threadIdx.x;
T r_p[ILP]; i_start * ILP < n && i_start * ILP < chunk_size;
T r_update[ILP]; i_start += blockDim.x) {
for (int i_start = threadIdx.x; i_start * ILP < n && i_start * ILP < chunk_size; i_start += blockDim.x) // load
{ load_store(r_p, p, 0, i_start);
// load load_store(r_update, update, 0, i_start);
load_store(r_p, p, 0, i_start);
load_store(r_update, update, 0, i_start);
#pragma unroll #pragma unroll
for (int ii = 0; ii < ILP; ii++) for (int ii = 0; ii < ILP; ii++) {
{ r_p[ii] = static_cast<MATH_T>(r_p[ii]) -
r_p[ii] = static_cast<MATH_T>(r_p[ii]) - (ratio * static_cast<MATH_T>(r_update[ii])); (ratio * static_cast<MATH_T>(r_update[ii]));
}
load_store(p, r_p, i_start, 0);
}
} }
else load_store(p, r_p, i_start, 0);
{ }
for (int i_start = 0; } else {
i_start < n && i_start < chunk_size; for (int i_start = 0; i_start < n && i_start < chunk_size;
i_start += blockDim.x * ILP) i_start += blockDim.x * ILP) {
{ MATH_T r_p[ILP];
MATH_T r_p[ILP]; MATH_T r_update[ILP];
MATH_T r_update[ILP];
#pragma unroll #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;
int i = i_start + threadIdx.x + ii * blockDim.x; if (i < n && i < chunk_size) {
if (i < n && i < chunk_size) r_p[ii] = p[i];
{ r_update[ii] = update[i];
r_p[ii] = p[i]; }
r_update[ii] = update[i]; }
}
}
#pragma unroll #pragma unroll
for (int ii = 0; ii < ILP; ii++) for (int ii = 0; ii < ILP; ii++) {
{ r_p[ii] = r_p[ii] - (ratio * r_update[ii]);
r_p[ii] = r_p[ii] - (ratio * r_update[ii]); }
}
#pragma unroll #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;
int i = i_start + threadIdx.x + ii * blockDim.x; if (i < n && i < chunk_size) {
if (i < n && i < chunk_size) p[i] = r_p[ii];
{ }
p[i] = r_p[ii];
}
}
}
} }
}
} }
}
}; };
void multi_tensor_lamb_cuda( void multi_tensor_lamb_cuda(int chunk_size, at::Tensor noop_flag,
int chunk_size, std::vector<std::vector<at::Tensor>> tensor_lists,
at::Tensor noop_flag, const float lr, const float beta1,
std::vector<std::vector<at::Tensor>> tensor_lists, const float beta2, const float epsilon,
const float lr, const int step, const int bias_correction,
const float beta1, const float weight_decay, const int grad_averaging,
const float beta2, const int mode, at::Tensor global_grad_norm,
const float epsilon, const float max_grad_norm,
const int step, at::optional<bool> use_nvlamb_python) {
const int bias_correction, using namespace at;
const float weight_decay, // Master weight and 32bit momentum(potentially changing) is not handled by
const int grad_averaging, // this So we assume every tensor are all in the same type
const int mode,
at::Tensor global_grad_norm, bool use_nvlamb =
const float max_grad_norm, use_nvlamb_python.has_value() ? use_nvlamb_python.value() : false;
at::optional<bool> use_nvlamb_python)
{ // Handle bias correction mode
using namespace at; float bias_correction1 = 1.0f, bias_correction2 = 1.0f;
// Master weight and 32bit momentum(potentially changing) is not handled by this if (bias_correction == 1) {
// So we assume every tensor are all in the same type bias_correction1 = 1 - std::pow(beta1, step);
bias_correction2 = 1 - std::pow(beta2, step);
bool use_nvlamb = use_nvlamb_python.has_value() ? use_nvlamb_python.value() : false; }
// Handle bias correction mode // Handle grad averaging mode
float bias_correction1 = 1.0f, bias_correction2 = 1.0f; float beta3 = 1.0f;
if (bias_correction == 1) if (grad_averaging == 1)
{ beta3 = 1 - beta1;
bias_correction1 = 1 - std::pow(beta1, step);
bias_correction2 = 1 - std::pow(beta2, step); std::vector<std::vector<at::Tensor>> grad_list(tensor_lists.begin(),
} tensor_lists.begin() + 1);
std::vector<std::vector<at::Tensor>> param_list(tensor_lists.begin() + 1,
// Handle grad averaging mode tensor_lists.begin() + 2);
float beta3 = 1.0f;
if (grad_averaging == 1) // Compute per tensor param norm
beta3 = 1 - beta1; auto param_norm_tuple =
multi_tensor_l2norm_cuda(chunk_size, noop_flag, param_list, true);
std::vector<std::vector<at::Tensor>> grad_list(tensor_lists.begin(), tensor_lists.begin() + 1);
std::vector<std::vector<at::Tensor>> param_list(tensor_lists.begin() + 1, tensor_lists.begin() + 2); // We now in-place modify grad to store update before compute its norm
// Generally this is not a issue since people modify grad in step() method all
// Compute per tensor param norm // the time We can also grab list of empty tensor to avoid this, but I'd like
auto param_norm_tuple = multi_tensor_l2norm_cuda(chunk_size, noop_flag, param_list, true); // to save space/cpu code
DISPATCH_FLOAT_AND_HALF(
// We now in-place modify grad to store update before compute its norm tensor_lists[0][0].scalar_type(), 0, "lamb_stage_1",
// Generally this is not a issue since people modify grad in step() method all the time multi_tensor_apply<4>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
// We can also grab list of empty tensor to avoid this, but I'd like to save space/cpu code LAMBStage1Functor<scalar_t_0>(), beta1, beta2,
DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, "lamb_stage_1", beta3, // 1-beta1 or 1 depends on averaging mode
multi_tensor_apply<4>( bias_correction1, bias_correction2, epsilon,
BLOCK_SIZE, (adamMode_t)mode, weight_decay,
chunk_size, global_grad_norm.DATA_PTR<float>(), max_grad_norm);)
noop_flag,
tensor_lists, // Compute update norms
LAMBStage1Functor<scalar_t_0>(), auto update_norm_tuple =
beta1, multi_tensor_l2norm_cuda(chunk_size, noop_flag, grad_list, true);
beta2,
beta3, // 1-beta1 or 1 depends on averaging mode std::vector<std::vector<at::Tensor>> grad_param_list(
bias_correction1, tensor_lists.begin(), tensor_lists.begin() + 2);
bias_correction2,
epsilon, DISPATCH_FLOAT_AND_HALF(
(adamMode_t)mode, tensor_lists[0][0].scalar_type(), 0, "lamb_stage_2",
weight_decay, multi_tensor_apply<2>(BLOCK_SIZE, chunk_size, noop_flag, grad_param_list,
global_grad_norm.DATA_PTR<float>(), LAMBStage2Functor<scalar_t_0>(),
max_grad_norm);) std::get<1>(param_norm_tuple).DATA_PTR<float>(),
std::get<1>(update_norm_tuple).DATA_PTR<float>(),
// Compute update norms lr, weight_decay, use_nvlamb);)
auto update_norm_tuple = multi_tensor_l2norm_cuda(chunk_size, noop_flag, grad_list, true);
AT_CUDA_CHECK(cudaGetLastError());
std::vector<std::vector<at::Tensor>> grad_param_list(tensor_lists.begin(), tensor_lists.begin() + 2);
DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, "lamb_stage_2",
multi_tensor_apply<2>(
BLOCK_SIZE,
chunk_size,
noop_flag,
grad_param_list,
LAMBStage2Functor<scalar_t_0>(),
std::get<1>(param_norm_tuple).DATA_PTR<float>(),
std::get<1>(update_norm_tuple).DATA_PTR<float>(),
lr,
weight_decay,
use_nvlamb);)
AT_CUDA_CHECK(cudaGetLastError());
} }
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment