Merge pull request #60 from junjihashimoto/feature/enable-cross-entropy-backward

austinvhuang · web-flow · commit 8b0f487059ff · 2024-09-18T13:12:43.000-04:00
Fix crossentropy_softmax_backward kernel
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

Original file line number	Diff line number	Diff line change
`@@ -677,7 +677,7 @@ fn main(@builtin(global_invocation_id) global_id : vec3<u32>) {`
`677`	`677`	`for (var i : u32 = 0u; i < V; i++) {`
`678`	`678`	`let p : {{precision}} = probs[probs_bt + i];`
`679`	`679`	`let indicator : {{precision}} = select(0.0, 1.0, i == ix);`
`680`		`- atomicAdd(&dlogits[dlogits_bt + i], (p - indicator) * dloss);`
	`680`	`+ dlogits[dlogits_bt + i] += (p - indicator) * dloss;`
`681`	`681`	`}`
`682`	`682`	`}`
`683`	`683`	`}`