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 all 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
138 changes: 138 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 @@ -1447,6 +1447,144 @@ namespace dlib
}
}

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

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

const long ns = src.num_samples();
const long ks = src.k();
const long num = src.nr() * src.nc();

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

// Compute RMS values
scale = 0;
const float* p_src = src.host();
float* p_scale = scale.host();
for (long n = 0; n < ns; ++n)
{
for (long k = 0; k < ks; ++k)
{
for (long i = 0; i < num; ++i)
{
p_scale[n] += (*p_src) * (*p_src);
++p_src;
}
}
p_scale[n] = 1.0f / std::sqrt(p_scale[n] / (ks * num) + static_cast<float>(eps));
}
scale.host();

// Apply RMS normalization
p_src = src.host();
float* p_dest = dest.host();
const float* p_gamma = gamma.host();
for (long n = 0; n < ns; ++n)
{
for (long k = 0; k < ks; ++k)
{
for (long i = 0; i < num; ++i)
{
*p_dest = (*p_src) * p_scale[n] * p_gamma[k];
++p_src;
++p_dest;
}
}
}
}

void rms_normalize_gradient(
const tensor& gradient_input,
const tensor& scale,
const tensor& src,
const tensor& gamma,
tensor& src_grad,
tensor& gamma_grad,
resizable_tensor& dscale
)
{
DLIB_CASSERT(src.num_samples() == scale.size());
DLIB_CASSERT(have_same_dimensions(gamma, gamma_grad));
DLIB_CASSERT(gamma.k() == src.k());
DLIB_CASSERT(gamma.nr() == 1);
DLIB_CASSERT(gamma.nc() == 1);
DLIB_CASSERT(have_same_dimensions(gradient_input, src));
DLIB_CASSERT(have_same_dimensions(gradient_input, src_grad));

const long ns = src.num_samples();
const long ks = src.k();
const long num = src.nr() * src.nc();

gamma_grad = 0;
dscale.copy_size(scale);
dscale = 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();
auto p_dscale = dscale.host();

for (long n = 0; n < ns; ++n)
{
const float scale_pow = -0.5f * std::pow(p_scale[n], 3.0f);
for (long k = 0; k < ks; ++k)
{
for (long i = 0; i < num; ++i)
{
const float x_hat = *p_src * p_scale[n];
p_gamma_grad[k] += (*p_grad) * x_hat;

const float dx = *p_grad * p_gamma[k];
p_dscale[n] += dx * *p_src * scale_pow;

++p_grad;
++p_src;
}
}
}

p_grad = gradient_input.host();
p_src = src.host();
auto p_src_grad = src_grad.host();
const float invnum = 1.0f / (ks * num);
for (long n = 0; n < ns; ++n)
{
for (long k = 0; k < ks; ++k)
{
for (long i = 0; i < num; ++i)
{
const float dx = *p_grad * p_gamma[k];
*p_src_grad += dx * p_scale[n] + p_dscale[n] * 2 * *p_src * invnum;

++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 @@ -255,6 +255,26 @@ namespace dlib
resizable_tensor& dvars
);

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

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

void rms_normalize_gradient(
const tensor& gradient_input,
const tensor& scale,
const tensor& src,
const tensor& gamma,
tensor& src_grad,
tensor& gamma_grad,
resizable_tensor& dscale
);

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

void threshold (
Expand Down
160 changes: 160 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 @@ -2280,6 +2280,166 @@ namespace dlib
dmeans.device(), dvars.device(), eps, src.num_samples(), src.k(), num);
}

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

__global__ void _cuda_rms_normalize(
float* dest,
float* scale,
const float* src,
const float* gamma,
float eps,
size_t ns,
size_t ks,
size_t num
)
{
for (auto n : grid_stride_range_y(0, ns))
{
const auto ps = src + n * ks * num;
float sum_squares = 0.0f;
for (auto i : grid_stride_range(0, ks * num))
{
sum_squares += ps[i] * ps[i];
}
warp_reduce_atomic_add(scale[n], sum_squares / (ks * num));
}
__syncthreads();

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))
{
const auto ps = src + n * ks * num;
const auto pd = dest + n * ks * num;
for (auto i : grid_stride_range(0, ks * num))
{
pd[i] = ps[i] * scale[n] * gamma[i / num];
}
}
}

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

const long ns = src.num_samples();
const long ks = src.k();
const long num = src.nr() * src.nc();

dest.copy_size(src);
scale.set_size(ns);
scale = 0;

launch_kernel(_cuda_rms_normalize, max_jobs(ks * num, ns),
dest.device(), scale.device(), src.device(), gamma.device(), eps, ns, ks, num);
}

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

__global__ void _cuda_rms_normalize_gradient(
float* src_grad,
float* gamma_grad,
float* dscale,
const float* src,
const float* gradient_input,
const float* scale,
const float* gamma,
size_t ns,
size_t ks,
size_t num
)
{
for (auto nk : grid_stride_range_y(0, ns * ks))
{
const auto n = nk / ks;
const auto k = nk % ks;
const auto ps = src + (n * ks + k) * num;
const auto pgi = gradient_input + (n * ks + k) * num;
const float scale_pow = -0.5f * std::pow(scale[n], 3.0f);
float temp_gg = 0.0f;
float temp_ds = 0.0f;
for (auto i : grid_stride_range(0, num))
{
const float x_hat = ps[i] * scale[n];
const float dx = pgi[i] * gamma[i / num];
temp_gg += pgi[i] * x_hat;
temp_ds += dx * ps[i] * scale_pow;
}
warp_reduce_atomic_add(gamma_grad[k], temp_gg);
warp_reduce_atomic_add(dscale[n], temp_ds);
}
__syncthreads();

const float invnum = 1.0f / (ks * num);
for (auto n : grid_stride_range_y(0, ns))
{
const auto ps = src + n * ks * num;
const auto pgi = gradient_input + n * ks * num;
const auto psg = src_grad + n * ks * num;
for (auto i : grid_stride_range(0, ks * num))
{
const float dx = pgi[i] * gamma[i / num];
psg[i] += dx * scale[n] + dscale[n] * 2 * ps[i] * invnum;
}
}
}

void rms_normalize_gradient(
const tensor& gradient_input,
const tensor& scale,
const tensor& src,
const tensor& gamma,
tensor& src_grad,
tensor& gamma_grad,
resizable_tensor& dscale
)
{
DLIB_CASSERT(src.num_samples() == scale.size());
DLIB_CASSERT(have_same_dimensions(gamma, gamma_grad));
DLIB_CASSERT(gamma.k() == src.k());
DLIB_CASSERT(gamma.nr() == 1);
DLIB_CASSERT(gamma.nc() == 1);
DLIB_CASSERT(have_same_dimensions(gradient_input, src));
DLIB_CASSERT(have_same_dimensions(gradient_input, src_grad));

const long ns = src.num_samples();
const long ks = src.k();
const long num = src.nr() * src.nc();

gamma_grad = 0;
dscale.copy_size(scale);
dscale = 0;

// Lancement du kernel CUDA
launch_kernel(_cuda_rms_normalize_gradient, max_jobs(ks * num, ns),
src_grad.device(), gamma_grad.device(), dscale.device(),
src.device(), gradient_input.device(), scale.device(), gamma.device(),
ns, ks, 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 @@ -362,6 +362,26 @@ namespace dlib
resizable_tensor& dvars
);

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

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

void rms_normalize_gradient(
const tensor& gradient_input,
const tensor& scale,
const tensor& src,
const tensor& gamma,
tensor& src_grad,
tensor& gamma_grad,
resizable_tensor& dscale
);

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

void threshold (
Expand Down
Loading
Loading