Skip to content

Add RMS Normalization Layer #2999

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Sep 7, 2024
Merged

Add RMS Normalization Layer #2999

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
c30e088
Add RMS Normalization Layer
Cydral Aug 25, 2024
d6517b2
Update dnn.cpp
Cydral Aug 26, 2024
cf462cd
Missing entry in visitors.h to take into account the new rms_norm_ layer
Cydral Aug 28, 2024
5a85514
Fix test function name
Cydral Aug 29, 2024
19a0cfe
Fix dangling pointer issue in CUDA implementation of rms_normalize_gr…
Cydral Aug 29, 2024
bd3d3f7
Merge branch 'davisking:master' into rms-norm-layer
Cydral Aug 29, 2024
6a12e0b
Fixing the dnn.cpp test program for the new rms_norm_ layer
Cydral Aug 29, 2024
03cdf93
Merge branch 'davisking:master' into rms-norm-layer
Cydral Sep 2, 2024
1083aab
General update of the rms_norm_ class
Cydral Sep 7, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 126 additions & 0 deletions dlib/cuda/cpu_dlib.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1431,6 +1431,132 @@ namespace dlib
}
}

// -----------------------------------------------------------------------------------

void rms_normalize(
const double eps,
resizable_tensor& dest,
resizable_tensor& scale,
const tensor& src,
const tensor& gamma
)
{
const long num = src.k() * src.nr() * src.nc();
DLIB_CASSERT(
src.k() == gamma.k() &&
src.nr() == gamma.nr() &&
src.nc() == gamma.nc() &&
eps > 0,
"\ngamma.k(): " << gamma.k() <<
"\ngamma.nr(): " << gamma.nr() <<
"\ngamma.nc(): " << gamma.nc() <<
"\nsrc.k(): " << src.k() <<
"\nsrc.nr(): " << src.nr() <<
"\nsrc.nc(): " << src.nc() <<
"\neps: " << eps
);

dest.copy_size(src);
scale.set_size(src.num_samples());

// Compute RMS
const auto p_scale = scale.host();
auto p_src = src.host();
for (long n = 0; n < src.num_samples(); ++n)
{
float sum_squares = 0;
for (long i = 0; i < num; ++i)
{
float val = p_src[n * num + i];
sum_squares += val * val;
}
p_scale[n] = sum_squares / num;
}
// Compute RMS inverse
for (long n = 0; n < src.num_samples(); ++n)
{
p_scale[n] = 1.0f / std::sqrt(p_scale[n] + eps);
}

p_src = src.host();
auto p_dest = dest.host();
auto p_gamma = gamma.host();
for (long n = 0; n < src.num_samples(); ++n)
{
for (long i = 0; i < num; ++i)
{
*p_dest = (*p_src) * p_scale[n] * p_gamma[i];
++p_src;
++p_dest;
}
}
}

void rms_normalize_gradient(
const double eps,
const tensor& gradient_input,
const tensor& scale,
const tensor& src,
const tensor& gamma,
tensor& src_grad,
tensor& gamma_grad
)
{
const long num = src.k() * src.nr() * src.nc();
DLIB_CASSERT(src.num_samples() == scale.size());
DLIB_CASSERT(src.k() == gamma.k());
DLIB_CASSERT(src.nr() == gamma.nr());
DLIB_CASSERT(src.nc() == gamma.nc());
DLIB_CASSERT(have_same_dimensions(gradient_input, src));
DLIB_CASSERT(have_same_dimensions(gradient_input, src_grad));
DLIB_CASSERT(have_same_dimensions(gamma_grad, gamma));
DLIB_CASSERT(eps > 0);

gamma_grad = 0;
auto p_grad = gradient_input.host();
auto p_src = src.host();
const auto p_gamma = gamma.host();
const auto p_gamma_grad = gamma_grad.host();
const auto p_scale = scale.host();

resizable_tensor dscale;
dscale.copy_size(scale);
dscale = 0;
const auto p_dscale = dscale.host();

for (long n = 0; n < src.num_samples(); ++n)
{
for (long i = 0; i < num; ++i)
{
const float x_hat = (*p_src) * p_scale[n];
p_gamma_grad[i] += (*p_grad) * x_hat;

const float dx = *p_grad * p_gamma[i];
p_dscale[n] += dx * (*p_src) * (-0.5) * p_scale[n] * p_scale[n] * p_scale[n];

++p_grad;
++p_src;
}
}

p_grad = gradient_input.host();
p_src = src.host();
auto p_src_grad = src_grad.host();
for (long n = 0; n < src.num_samples(); ++n)
{
for (long i = 0; i < num; ++i)
{
const float dx = *p_grad * p_gamma[i];

*p_src_grad += dx * p_scale[n] + p_dscale[n] * 2 * (*p_src) / num;

++p_grad;
++p_src;
++p_src_grad;
}
}
}

// -----------------------------------------------------------------------------------

void threshold (
Expand Down
20 changes: 20 additions & 0 deletions dlib/cuda/cpu_dlib.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -253,6 +253,26 @@ namespace dlib
tensor& beta_grad
);

// -----------------------------------------------------------------------------------

void rms_normalize(
const double eps,
resizable_tensor& dest,
resizable_tensor& scale,
const tensor& src,
const tensor& gamma
);

void rms_normalize_gradient(
const double eps,
const tensor& gradient_input,
const tensor& scale,
const tensor& src,
const tensor& gamma,
tensor& src_grad,
tensor& gamma_grad
);

// -----------------------------------------------------------------------------------

void threshold (
Expand Down
126 changes: 126 additions & 0 deletions dlib/cuda/cuda_dlib.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2242,6 +2242,132 @@ namespace dlib
dmeans.device(), dvars.device(), eps, src.num_samples(), num);
}

// ----------------------------------------------------------------------------------------

__global__ void _cuda_rms_normalize(float* out, const float* s, float* scale, const float* g, float eps, size_t ns, size_t num)
{
// Compute sum of squares
for (auto n : grid_stride_range_y(0, ns))
{
auto p = s + n * num;
float sum_squares = 0;
for (auto i : grid_stride_range(0, num))
{
sum_squares += p[i] * p[i];
}
warp_reduce_atomic_add(scale[n], sum_squares / num);
}
__syncthreads();

// Compute RMS inverse
for (auto n : grid_stride_range_y(0, ns))
{
for (auto i : grid_stride_range(0, 1))
{
scale[n] = 1.0f / std::sqrt(scale[n] + eps);
}
}
__syncthreads();

for (auto n : grid_stride_range_y(0, ns))
{
for (auto i : grid_stride_range(0, num))
{
const float val = s[n * num + i] * scale[n];
out[n * num + i] = val * g[i];
}
}
}

__global__ void _cuda_rms_normalize_gradient(float* out, float* gg, const float* s, const float* gi, const float* scale, const float* g, float* dscale, float eps, size_t ns, size_t num)
{
for (auto n : grid_stride_range_y(0, ns))
{
float temp_dscale = 0;
for (auto i : grid_stride_range(0, num))
{
auto idx = n * num + i;
const float x_hat = s[idx] * scale[n];
gg[i] += gi[idx] * x_hat;

const float dx = gi[idx] * g[i];
temp_dscale += dx * s[idx] * -0.5 * scale[n] * scale[n] * scale[n];
}
warp_reduce_atomic_add(dscale[n], temp_dscale);
}
__syncthreads();

for (auto n : grid_stride_range_y(0, ns))
{
for (auto i : grid_stride_range(0, num))
{
auto idx = n * num + i;
const float dx = gi[idx] * g[i];
out[idx] += dx * scale[n] + dscale[n] * 2 * s[idx] / num;
}
}
}

void rms_normalize(
const double eps,
resizable_tensor& dest,
resizable_tensor& scale,
const tensor& src,
const tensor& gamma
)
{
const long num = src.k() * src.nr() * src.nc();
DLIB_CASSERT(
src.k() == gamma.k() &&
src.nr() == gamma.nr() &&
src.nc() == gamma.nc() &&
eps > 0,
"\ngamma.k(): " << gamma.k() <<
"\ngamma.nr(): " << gamma.nr() <<
"\ngamma.nc(): " << gamma.nc() <<
"\nsrc.k(): " << src.k() <<
"\nsrc.nr(): " << src.nr() <<
"\nsrc.nc(): " << src.nc() <<
"\neps: " << eps
);

dest.copy_size(src);
scale.set_size(src.num_samples());
scale = 0;
launch_kernel(_cuda_rms_normalize, max_jobs(num, src.num_samples()), dest.device(), src.device(),
scale.device(), gamma.device(), eps, src.num_samples(), num);
}

void rms_normalize_gradient(
const double eps,
const tensor& gradient_input,
const tensor& scale,
const tensor& src,
const tensor& gamma,
tensor& src_grad,
tensor& gamma_grad
)
{
const long num = src.k() * src.nr() * src.nc();
DLIB_CASSERT(src.num_samples() == scale.size());
DLIB_CASSERT(src.k() == gamma.k());
DLIB_CASSERT(src.nr() == gamma.nr());
DLIB_CASSERT(src.nc() == gamma.nc());
DLIB_CASSERT(have_same_dimensions(gradient_input, src));
DLIB_CASSERT(have_same_dimensions(gradient_input, src_grad));
DLIB_CASSERT(have_same_dimensions(gamma_grad, gamma));
DLIB_CASSERT(eps > 0);

gamma_grad = 0;
resizable_tensor dscale;
dscale.copy_size(scale);
dscale = 0;
launch_kernel(_cuda_rms_normalize_gradient, max_jobs(num, src.num_samples()),
src_grad.device(), gamma_grad.device(), src.device(),
gradient_input.device(), scale.device(), gamma.device(),
dscale.device(), eps, src.num_samples(), num);
}

// ----------------------------------------------------------------------------------------

__global__ void _cuda_copy_tensor_add_to (float* dest, size_t size, const float* src, size_t dest_stride, size_t src_stride, size_t block_size)
Expand Down
20 changes: 20 additions & 0 deletions dlib/cuda/cuda_dlib.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -360,6 +360,26 @@ namespace dlib
tensor& beta_grad
);

// -----------------------------------------------------------------------------------

void rms_normalize(
const double eps,
resizable_tensor& dest,
resizable_tensor& scale,
const tensor& src,
const tensor& gamma
);

void rms_normalize_gradient(
const double eps,
const tensor& gradient_input,
const tensor& scale,
const tensor& src,
const tensor& gamma,
tensor& src_grad,
tensor& gamma_grad
);

// -----------------------------------------------------------------------------------

void threshold (
Expand Down
34 changes: 34 additions & 0 deletions dlib/cuda/tensor_tools.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -694,6 +694,40 @@ namespace dlib { namespace tt
#endif
}

// ----------------------------------------------------------------------------------------

void rms_normalize(
const double eps,
resizable_tensor& dest,
resizable_tensor& scale,
const tensor& src,
const tensor& gamma
)
{
#ifdef DLIB_USE_CUDA
cuda::rms_normalize(eps, dest, scale, src, gamma);
#else
cpu::rms_normalize(eps, dest, scale, src, gamma);
#endif
}

void rms_normalize_gradient(
const double eps,
const tensor& gradient_input,
const tensor& scale,
const tensor& src,
const tensor& gamma,
tensor& src_grad,
tensor& gamma_grad
)
{
#ifdef DLIB_USE_CUDA
cuda::rms_normalize_gradient(eps, gradient_input, scale, src, gamma, src_grad, gamma_grad);
#else
cpu::rms_normalize_gradient(eps, gradient_input, scale, src, gamma, src_grad, gamma_grad);
#endif
}

// ----------------------------------------------------------------------------------------

void threshold (
Expand Down
Loading
Loading