CUDA: Getting zero gradients
- Dominant language
- LLVM
- Stars
- 1.7k
- Forks
- 188
- Avg merge
- 2d 4h
- Merged PRs (30d)
- 22
Description
I am trying Enzyme for CUDA, but I am getting zero gradient. Specifically, for `y = f(x)`, I get `∇x = 0` while `∇y` is left as its original value. It seems like the gradient is not propagated at all. This only happens on CUDA, and it works fine on CPU.
To reproduce, you can use the following code, which is a modified version of the code in the guide.
```cuda
#include
#define n (5856)
#define m (65536)
void __device__ foo_impl(double *faces_ptr, double *soft_colors_ptr) {
#define face(fn, a, b) (faces_ptr[(fn)*3 * 3 + (a)*3 + (b)])
#define soft_colors(k, pn) (soft_colors_ptr[(k)*m + (pn)])
for (int pn = 0; pn < m; pn++) {
double soft_color[3] = {0, 0, 0};
for (int fn = 0; fn < n; fn++) {
double w_clip[3] = {1, 1, 1};
// CORRECT
/*double zp = w_clip[0] * face(fn, 0, 2) +
w_clip[1] * face(fn, 1, 2) + w_clip[2] * face(fn, 2,
2);*/
// WRONG
double zp =
1. / (w_clip[0] / face(fn, 0, 2) + w_clip[1] / face(fn, 1, 2) +
w_clip[2] / face(fn, 2, 2));
for (int k = 0; k < 3; k++) {
soft_color[k] += zp;
}
}
for (int k = 0; k < 3; k++) {
soft_colors(k, pn) = soft_color[k];
}
}
#undef face
#undef soft_colors
}
typedef void (*f_ptr)(double *, double *);
extern void __device__ __enzyme_autodiff(f_ptr, int, double *, double *, int,
double *, double *);
void __global__ foo(double *x_in, double *x_out) { foo_impl(x_in, x_out); }
int __device__ enzyme_dup;
int __device__ enzyme_out;
int __device__ enzyme_const;
void __global__ foo_grad(double *x, double *d_x, double *y, double *d_y) {
__enzyme_autodiff(foo_impl, enzyme_dup, x, d_x, enzyme_dup, y, d_y);
}
int main() {
double *x, *d_x, *y, *d_y; // device pointers
cudaMalloc(&x, n * 3 * 3 * sizeof(*x));
cudaMalloc(&d_x, n * 3 * 3 * sizeof(*d_x));
cudaMalloc(&y, 3 * m * sizeof(*y));
cudaMalloc(&d_y, 3 * m * sizeof(*d_y));
double *host_x = new double[n * 3 * 3];
double *host_d_x = new double[n * 3 * 3];
double *host_y = new double[3 * m];
double *host_d_y = new double[3 * m];
for (int i = 0; i < n; i++) {
for (int j = 0; j < 3; j++) {
for (int k = 0; k < 3; k++) {
host_x[i * 3 * 3 + j * 3 + k] = 1;
host_d_x[i * 3 * 3 + j * 3 + k] = 0;
}
}
}
for (int i = 0; i < 3; i++) {
for (int j = 0; j < m; j++) {
host_y[i * m + j] = 0;
host_d_y[i * m + j] = 1;
}
}
cudaMemcpy(x, host_x, n * 3 * 3 * sizeof(*x), cudaMemcpyHostToDevice);
cudaMemcpy(d_x, host_d_x, n * 3 * 3 * sizeof(*d_x), cudaMemcpyHostToDevice);
cudaMemcpy(y, host_y, 3 * m * sizeof(*y), cudaMemcpyHostToDevice);
cudaMemcpy(d_y, host_d_y, 3 * m * sizeof(*d_y), cudaMemcpyHostToDevice);
// foo<<<1,1>>>(x, y); fwd-pass only
foo_grad<<<1, 1>>>(x, d_x, y, d_y); // fwd and bkwd pass
cudaDeviceSynchronize(); // synchroniz
cudaMemcpy(host_x, x, n * 3 * 3 * sizeof(*x), cudaMemcpyDeviceToHost);
cudaMemcpy(host_d_x, d_x, n * 3 * 3 * sizeof(*d_x), cudaMemcpyDeviceToHost);
cudaMemcpy(host_y, y, 3 * m * sizeof(*y), cudaMemcpyDeviceToHost);
cudaMemcpy(host_d_y, d_y, 3 * m * sizeof(*d_y), cudaMemcpyDeviceToHost);
for (int i = 0; i < n; i++) {
for (int j = 0; j < 3; j++) {
for (int k = 0; k < 3; k++) {
printf("x(%d, %d, %d) = %f\n", i, j, k,
host_x[i * 3 * 3 + j * 3 + k]);
}
}
}
for (int i = 0; i < 3; i++) {
for (int j = 0; j < m; j++) {
printf("y(%d, %d) = %f\n", i, j, host_y[i * m + j]);
}
}
for (int i = 0; i < n; i++) {
for (int j = 0; j < 3; j++) {
for (int k = 0; k < 3; k++) {
printf("d_x(%d, %d, %d) = %f\n", i, j, k,
host_d_x[i * 3 * 3 + j * 3 + k]);
}
}
}
for (int i = 0; i < 3; i++) {
for (int j = 0; j < m; j++) {
printf("d_y(%d, %d) = %f\n", i, j, host_d_y[i * m + j]);
}
}
}
```
Compile and run it with the following commands:
```sh
/path/to/clang test.cu -Xclang -load -Xclang /path/to/Enzyme/enzyme/build/Enzyme/ClangEnzyme-12.so -O2 -fno-vectorize -fno-unroll-loops -fPIC --cuda-gpu-arch=sm_70 -lcudart -L/path/to/cuda-10.1.243/lib64 -std=c++17 -lstdc++
./a.out
```
You will get `host_d_x` to be 0, and `host_d_y` left as 1.
I have checked there are no CUDA errors by running it with `cuda-memcheck`.
You may also try commenting out the `WRONG` line and uncommenting the `CORRECT` line. This will cause the gradient to follow again and `host_d_x` will no longer be 0. However, this doesn’t mean that the WRONG line is the root cause of the problem. There are many changes to the code that can make the gradients follow again, including change `n` or `m` to a smaller value. I think expression folding or other optimizations from LLVM is hiding the real cause of the problem. Another guess of the cause could be overflow of 32-bit integers somewhere in the differentiated program.
Tested with Enzyme 0.0.51, Clang 12.0.1, CUDA 10.1.243.
Contributor guide
Assessment
This issue has not been assessed yet.