Make iter persistent and add kernel check (#2897)

Summary: Pull Request resolved: #2897 [fbgemm_gpu][GWD] Make `iter` persistent and add kernel check to fail when `prev_iter` > `iter`. Differential Revision: D58566070
pytorch · Jul 25, 2024 · 4dffb79 · 4dffb79
1 parent a369c9e
commit 4dffb79
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/fbgemm_gpu/codegen/genscript/jinja_environment.py b/fbgemm_gpu/codegen/genscript/jinja_environment.py
@@ -338,6 +338,7 @@ def compute_global_weight_decay(is_global_weight_decay_kernel: bool) -> str:
     if is_global_weight_decay_kernel:
         return """
         const auto prev_iter = prev_iter_dev[linear_index];
+        CUDA_KERNEL_ASSERT(prev_iter < iter);
         const auto global_weight_decay = prev_iter == 0 ? 1 : max(gwd_lower_bound, powf(weight_decay_base, iter - prev_iter - 1));
         if (threadIdx.x == 0) {
             prev_iter_dev[linear_index] = iter;

diff --git a/fbgemm_gpu/codegen/training/forward/embedding_forward_split_kernel_template.cu b/fbgemm_gpu/codegen/training/forward/embedding_forward_split_kernel_template.cu
@@ -188,6 +188,7 @@ using namespace fbgemm_gpu;
         {%- if is_gwd_kernel %}
         // if l > L or prev_iter == 0, global_weight_decay = 1
         const auto prev_it = prev_iter[idx];
+        CUDA_KERNEL_ASSERT(prev_it < iter);
         const auto global_weight_decay = (l > L || prev_it == 0) ? 1 : max(gwd_lower_bound, powf(weight_decay_base, iter - prev_it - 1));
         {%- endif %}