Merge pull request #60 from junjihashimoto/feature/enable-cross-entro…

…py-backward Fix crossentropy_softmax_backward kernel
AnswerDotAI · Sep 18, 2024 · 8b0f487 · 8b0f487
2 parents be0971a + 353fd1f
commit 8b0f487
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/experimental/kernels/kernels.h b/experimental/kernels/kernels.h
@@ -677,7 +677,7 @@ fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {
  for (var i : u32 = 0u; i < V; i++) {
  let p : {{precision}} = probs[probs_bt + i];
  let indicator : {{precision}} = select(0.0, 1.0, i == ix);
- atomicAdd(&dlogits[dlogits_bt + i], (p - indicator) * dloss);
+ dlogits[dlogits_bt + i] += (p - indicator) * dloss;
  }
  }
 }

diff --git a/experimental/kernels/unittest_llmc/unittest_kernels.h b/experimental/kernels/unittest_llmc/unittest_kernels.h
@@ -20,7 +20,7 @@ extern "C" {
 #define USE_GPU_FOR_RESIDUAL_BACKWARD 1
 #define USE_GPU_FOR_SOFTMAX_FORWARD 1
 #define USE_GPU_FOR_CROSSENTROPY_FORWARD 1
-// #define USE_GPU_FOR_CROSSENTROPY_SOFTMAX_BACKWARD 1
+#define USE_GPU_FOR_CROSSENTROPY_SOFTMAX_BACKWARD 1
 
 
 #ifdef USE_GPU_FOR_ENCODER_FORWARD