Problem with JAX cuda install - convolutions?

[homerjed]

In my installation of jax

from jax.lib import xla_bridge
print(xla_bridge.get_backend().platform)

and

import jax
print(jax.devices())

indicate my GPU is being found.

I am not sure exactly what the problem is, but some of the jax/examples scripts that run without convolutions do not seem as fast as I would expect (maybe even at CPU speed) and I cannot run my own more complicated models (that run perfectly on colab).

I believe the GPU works after running the VAE example in the jax/examples directory. It seems instead that convolutions cause the problem.

My relevant versions are here

python 3.9.7
cuda 11.4.2
cudnn 8.2.4.15

and I install via the commands

pip install -U pip wheel
pip install https://storage.googleapis.com/jax-releases/cuda11/jaxlib-0.3.14+cuda11.cudnn82-cp39-none-manylinux2014_x86_64.whl
pip install -U "jax[cuda11_cudnn82]==0.3.14" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.htm

However, when I run some codes, I obtain the error here:

/software/opt/focal/x86_64/spack/2021.12/spack/opt/spack/linux-ubuntu20.04-x86_64_v2/gcc-11.2.0/cudnn-8.2.4.15-11.4-r5srvd2bjed7zlr75cesfus3nwsjprw6/lib64:/software/opt/focal/x86_64/spack/2021.12/spack/opt/spack/linux-ubuntu20.04-x86_64_v2/gcc-11.2.0/cuda-11.4.2-jefqkwdwi245u5nbdg5tw3ufrucvsnag/lib64:/opt/slurm/lib:/opt/slurm/lib:/opt/slurm/lib:
Jax version 0.3.14
Flax version 0.5.3
60000 training images
3750 training steps per epoch
Available jax devices: [GpuDevice(id=0, process_index=0)]
2022-08-04 16:23:20.156272: W tensorflow/core/kernels/data/cache_dataset_ops.cc:856] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset  will be discarded. This can happen if you have an input pipeline similar to `dataset.cache().take(k).repeat()`. You should use `dataset.take(k).cache().repeat()` instead.
2022-08-04 16:24:08.553578: W external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_conv_algorithm_picker.cc:727] None of the algorithms provided by cuDNN heuristics worked; trying fallback algorithms.  Conv: (f32[3,3,8,16]{2,1,0,3}, u8[0]{0}) custom-call(f32[16,4,4,8]{3,2,1,0}, f32[16,4,4,16]{3,2,1,0}), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, custom_call_target="__cudnn$convBackwardFilter", backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"
Traceback (most recent call last):
  File "/home/j/Jed.Homer/jax_test/cglow/conditional/train.py", line 472, in <module>
    model, params = train_glow(train_ds, val_ds=val_ds, **config_dict)
  File "/home/j/Jed.Homer/jax_test/cglow/conditional/train.py", line 307, in train_glow
    loss, state = train_step(state, batch)
  File "/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/jax/_src/api.py", line 522, in cache_miss
    out_flat = xla.xla_call(
  File "/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/jax/core.py", line 1836, in bind
    return call_bind(self, fun, *args, **params)
  File "/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/jax/core.py", line 1852, in call_bind
    outs = top_trace.process_call(primitive, fun_, tracers, params)
  File "/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/jax/core.py", line 683, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/jax/_src/dispatch.py", line 196, in _xla_call_impl
    compiled_fun = xla_callable(fun, device, backend, name, donated_invars,
  File "/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/jax/linear_util.py", line 286, in memoized_fun
    ans = call(fun, *args)
  File "/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/jax/_src/dispatch.py", line 245, in _xla_callable_uncached
    return lower_xla_callable(fun, device, backend, name, donated_invars, False,
  File "/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/jax/_src/dispatch.py", line 808, in compile
    self._executable = XlaCompiledComputation.from_xla_computation(
  File "/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/jax/_src/dispatch.py", line 911, in from_xla_computation
    compiled = compile_or_get_cached(backend, xla_computation, options)
  File "/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/jax/_src/dispatch.py", line 872, in compile_or_get_cached
    return backend_compile(backend, computation, compile_options)
  File "/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/jax/_src/profiler.py", line 312, in wrapper
    return func(*args, **kwargs)
  File "/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/jax/_src/dispatch.py", line 818, in backend_compile
    return backend.compile(built_c, compile_options=options)
jax._src.traceback_util.UnfilteredStackTrace: jaxlib.xla_extension.XlaRuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm for:
%cudnn-conv-bw-filter = (f32[3,3,8,16]{2,1,0,3}, u8[0]{0}) custom-call(f32[16,4,4,8]{3,2,1,0} %slice.26094, f32[16,4,4,16]{3,2,1,0} %add.2057), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, custom_call_target="__cudnn$convBackwardFilter", metadata={op_name="jit(train_step)/jit(main)/transpose(jvp(GLOW))/flow_scale_3//conv_prior/Conv_0/conv_general_dilated[window_strides=(1, 1) padding=((1, 1), (1, 1)) lhs_dilation=(1, 1) rhs_dilation=(1, 1) dimension_numbers=ConvDimensionNumbers(lhs_spec=(3, 0, 1, 2), rhs_spec=(3, 0, 1, 2), out_spec=(2, 3, 0, 1)) feature_group_count=1 batch_group_count=1 lhs_shape=(16, 4, 4, 8) rhs_shape=(16, 4, 4, 16) precision=None preferred_element_type=None]" source_file="/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/flax/linen/linear.py" source_line=425}, backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"

Original error: INTERNAL: All algorithms tried for %cudnn-conv-bw-filter = (f32[3,3,8,16]{2,1,0,3}, u8[0]{0}) custom-call(f32[16,4,4,8]{3,2,1,0} %slice.26094, f32[16,4,4,16]{3,2,1,0} %add.2057), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, custom_call_target="__cudnn$convBackwardFilter", metadata={op_name="jit(train_step)/jit(main)/transpose(jvp(GLOW))/flow_scale_3//conv_prior/Conv_0/conv_general_dilated[window_strides=(1, 1) padding=((1, 1), (1, 1)) lhs_dilation=(1, 1) rhs_dilation=(1, 1) dimension_numbers=ConvDimensionNumbers(lhs_spec=(3, 0, 1, 2), rhs_spec=(3, 0, 1, 2), out_spec=(2, 3, 0, 1)) feature_group_count=1 batch_group_count=1 lhs_shape=(16, 4, 4, 8) rhs_shape=(16, 4, 4, 16) precision=None preferred_element_type=None]" source_file="/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/flax/linen/linear.py" source_line=425}, backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}" failed. Falling back to default algorithm.  Per-algorithm errors:
  Profiling failure on cuDNN engine eng20{k2=6,k3=0}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng34{k5=2,k14=6,k0=14,k2=5}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng18{k11=0}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng25{k5=2,k14=6,k0=7,k2=5}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng17{}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng30{k23=0,k12=8,k13=2,k14=3,k15=1,k17=9,k18=1,k2=0}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng22{k14=6,k2=5,k5=2}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng1{k2=6,k3=0}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng23{k2=0,k18=1,k13=2,k14=3,k23=0}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng0{}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng0{}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng1{}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng20{}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'

To ignore this failure and try to use a fallback algorithm (which may have suboptimal performance), use XLA_FLAGS=--xla_gpu_strict_conv_algorithm_picker=false.  Please also file a bug for the root cause of failing autotuning.

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/j/Jed.Homer/jax_test/cglow/conditional/train.py", line 472, in <module>
    model, params = train_glow(train_ds, val_ds=val_ds, **config_dict)
  File "/home/j/Jed.Homer/jax_test/cglow/conditional/train.py", line 307, in train_glow
    loss, state = train_step(state, batch)
jaxlib.xla_extension.XlaRuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm for:
%cudnn-conv-bw-filter = (f32[3,3,8,16]{2,1,0,3}, u8[0]{0}) custom-call(f32[16,4,4,8]{3,2,1,0} %slice.26094, f32[16,4,4,16]{3,2,1,0} %add.2057), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, custom_call_target="__cudnn$convBackwardFilter", metadata={op_name="jit(train_step)/jit(main)/transpose(jvp(GLOW))/flow_scale_3//conv_prior/Conv_0/conv_general_dilated[window_strides=(1, 1) padding=((1, 1), (1, 1)) lhs_dilation=(1, 1) rhs_dilation=(1, 1) dimension_numbers=ConvDimensionNumbers(lhs_spec=(3, 0, 1, 2), rhs_spec=(3, 0, 1, 2), out_spec=(2, 3, 0, 1)) feature_group_count=1 batch_group_count=1 lhs_shape=(16, 4, 4, 8) rhs_shape=(16, 4, 4, 16) precision=None preferred_element_type=None]" source_file="/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/flax/linen/linear.py" source_line=425}, backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"

Original error: INTERNAL: All algorithms tried for %cudnn-conv-bw-filter = (f32[3,3,8,16]{2,1,0,3}, u8[0]{0}) custom-call(f32[16,4,4,8]{3,2,1,0} %slice.26094, f32[16,4,4,16]{3,2,1,0} %add.2057), window={size=3x3 pad=1_1x1_1}, dim_labels=b01f_01io->b01f, custom_call_target="__cudnn$convBackwardFilter", metadata={op_name="jit(train_step)/jit(main)/transpose(jvp(GLOW))/flow_scale_3//conv_prior/Conv_0/conv_general_dilated[window_strides=(1, 1) padding=((1, 1), (1, 1)) lhs_dilation=(1, 1) rhs_dilation=(1, 1) dimension_numbers=ConvDimensionNumbers(lhs_spec=(3, 0, 1, 2), rhs_spec=(3, 0, 1, 2), out_spec=(2, 3, 0, 1)) feature_group_count=1 batch_group_count=1 lhs_shape=(16, 4, 4, 8) rhs_shape=(16, 4, 4, 16) precision=None preferred_element_type=None]" source_file="/home/j/Jed.Homer/omf_env6/lib/python3.9/site-packages/flax/linen/linear.py" source_line=425}, backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}" failed. Falling back to default algorithm.  Per-algorithm errors:
  Profiling failure on cuDNN engine eng20{k2=6,k3=0}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng34{k5=2,k14=6,k0=14,k2=5}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng18{k11=0}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng25{k5=2,k14=6,k0=7,k2=5}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng17{}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng30{k23=0,k12=8,k13=2,k14=3,k15=1,k17=9,k18=1,k2=0}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng22{k14=6,k2=5,k5=2}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng1{k2=6,k3=0}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng23{k2=0,k18=1,k13=2,k14=3,k23=0}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng0{}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng0{}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng1{}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'
  Profiling failure on cuDNN engine eng20{}: UNKNOWN: CUDNN_STATUS_INTERNAL_ERROR
in external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_dnn.cc(4369): 'status'

To ignore this failure and try to use a fallback algorithm (which may have suboptimal performance), use XLA_FLAGS=--xla_gpu_strict_conv_algorithm_picker=false.  Please also file a bug for the root cause of failing autotuning.

[hawkinsp]

I don't know if it will help, but it's something to try: can you try lowering the amount of GPU memory JAX allocates (https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html) ?

Can you also share the output of nvidia-smi?

[homerjed]

OK I'm convinced, with the information I provided, plus this example script, that at least I am using the GPU when using jax.

The flax nn.Conv's use the jax.lax.conv_general_dilated so something must be working!

Not sure where my error arises though...

  @nn.compact
  def __call__(self, x):
    x = nn.Conv(features=32, kernel_size=(3, 3))(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = nn.Conv(features=64, kernel_size=(3, 3))(x)
    x = nn.relu(x)
    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
    x = x.reshape((x.shape[0], -1))  # flatten
    x = nn.Dense(features=256)(x)
    x = nn.relu(x)
    x = nn.Dense(features=10)(x)
    x = nn.log_softmax(x)
    return x

model = CNN()
batch = jnp.ones((128, 64, 64, 10))  # (N, H, W, C) format
variables = model.init(jax.random.PRNGKey(0), batch)
output = model.apply(variables, batch)

 key = jr.PRNGKey(0)                                                             
 for i in range(100):                                                            
     (key,) = jr.split(key, 1)                                                   
     batch = jr.normal(key, (128, 64, 64, 10))                                   
     print(model.apply(variables, batch).shape) 

[hawkinsp]

Can you try running with CUDA_LAUNCH_BLOCKING=1? Does the error change?

[homerjed]

Same error as before, thanks for the suggestion

[sudhakarsingh27]

Sure, the model does seem to use GPUs. Going by the error raised, seems like cuDNN doesn't have that very specific filter (apparently for the backward pass). I'm surprised that even fallback algorithm didn't work.
Can you share a min repro of the above error?

[homerjed]

Thanks for the message. Here are two minimal-ish examples that work on Colab.

This is a minimal example testing the jax.lax.conv which works for me:

import jax                                                                                                
import jax.numpy as jnp                                                         
from jax import lax                                                             
print(jax.devices())                                                            
                                                                                 
# 2D kernel - HWIO layout                                                       
kernel = jnp.zeros((3, 3, 3, 3), dtype=jnp.float32)                             
kernel += jnp.array([[1, 1, 0],                                                 
                    [1, 0,-1],                                                 
                    [0,-1,-1]])[:, :, jnp.newaxis, jnp.newaxis]                
                                                                                
# NHWC layout                                                                   
img = jnp.zeros((1, 200, 198, 3), dtype=jnp.float32)                            
for k in range(3):                                                              
  x = 30 + 60*k                                                                 
  y = 20 + 60*k                                                                 
  img = img.at[0, x:x+10, y:y+10, k].set(1.0)                                   

out = lax.conv(jnp.transpose(img,[0,3,1,2]),    # lhs = NCHW image tensor       
              jnp.transpose(kernel,[3,2,0,1]), # rhs = OIHW conv kernel tensor                
               (1, 1),  # window strides                                        
              'SAME') # padding mode                                           
print("out shape: ", out.shape)   

However, this other "minimal" example doesn't work for me. I assume that the stax modules depend on the same jax.lax.conv which is why this is interesting.

I have tried changing the usual environment variables listed on the jax page about the OOM errors.

import jax                                                                      
import jax.numpy as jnp                                                         
from jax import random                                                          
import numpy as onp                                                             
from jax.example_libraries import stax, optimizers                              
                                                                                 
rng = random.PRNGKey(1)                                                         
                                                                                 
conv_init, conv_apply = stax.serial(                                            
  stax.Conv(32,(3,3), padding="SAME"),                                        
  stax.Relu,                                                                  
  stax.Conv(16, (3,3), padding="SAME"),                                       
  stax.Relu,                                                                  
  stax.Flatten,                                                               
  stax.Dense(10),                                                             
  stax.Softmax)                                                               
                                                                                  
weights = conv_init(rng, (18, 28, 28, 1))                                       
weights = weights[1]                                                            
learning_rate = jnp.array(1e-4)                                                 
epochs = 25                                                                     
batch_size = 256                                                                
                                                                                 
opt_init, opt_update, opt_get_weights = optimizers.sgd(learning_rate)           
opt_state = opt_init(weights)                                                   
                                                                                
def loss_fn(weights, input_data, actual):                                       
  preds = conv_apply(weights, input_data)                                     
  one_hot_actual = jax.nn.one_hot(actual, num_classes=10)                     
  log_preds = jnp.log(preds)                                                  
  loss = -(one_hot_actual * log_preds).sum()                                  
  return loss                                                                 
                                                                                 
def train(epochs, opt_state, batch_size=32):                                                                                                                     
  for i in range(1, epochs+1):                                                
                                                                                 
    X_batch = random.normal(rng, shape=(batch_size, 28, 28, 1))             
    Y_batch = random.randint(rng, minval=0, maxval=10, shape=(batch_size, 1))
                                                                                                 
    loss, gradients = jax.value_and_grad(loss_fn)(
        opt_get_weights(opt_state),                             
        X_batch, Y_batch)                                       
                                                                             
    opt_state = opt_update(i, gradients, opt_state)                         
                                                                                 
    print("Loss : {:.3f}".format(loss))                 
                                                                                
    return opt_state                                                            
                                                                                 
final_opt_state = train(epochs, opt_state, batch_size=batch_size) 

Sorry for the long message. Thanks for your help!