jax: On entry to GEMM_EX parameter number {9,12} had an illegal value

When I try to run my jax-based training script on the GPU, it prints the following messages several hundred times but proceeds without an exception.

** On entry to GEMM_EX  parameter number 12 had an illegal value
** On entry to GEMM_EX  parameter number 9 had an illegal value

The same script doesn’t print any errors when run with a CPU-only build. This error occurs with the pip build and a custom build of master.

To reproduce, run python jax_transformer.py alice.txt in https://github.com/joschu/jax-exp/

About this issue

Original URL
State: closed
Created 5 years ago
Comments: 16 (16 by maintainers)

Commits related to this issue

index_take in terms of gather, delete index_untake (c.f. #304) — committed to google/jax by mattjj 5 years ago
Update XLA release. Updates XLA to https://github.com/tensorflow/tensorflow/commit/e889ea1dd965c31c391106aa3518fc23d2689954, which fixes #304. — committed to hawkinsp/jax by hawkinsp 5 years ago

Most upvoted comments

The GEMM_EX problem is now tracked at https://github.com/tensorflow/tensorflow/issues/25761

hawkinsp on Feb 14, 2019

I think this happens during GEMM autotuning. It may be benign because eventually XLA gives up and uses a generic GEMM, which seems to work:

Snippet from a log (with TF_CPP_MIN_VLOG_LEVEL=2 set)

2019-02-01 14:41:54.525466: I external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gpu_executable.cc:142] Executing the thunk for %dot.3438 = f32[96,384]{0,1} dot(f32[8192,96]{1,0} %bitcast.123, f32[8192,384]{1,0} %fusion.250), lhs_contr
acting_dims={0}, rhs_contracting_dims={0} on stream 0
2019-02-01 14:41:54.525470: I external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc:314] Executing a GemmThunk
2019-02-01 14:41:54.525477: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=-1)
2019-02-01 14:41:54.525488: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=0)
2019-02-01 14:41:54.525498: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=1)
2019-02-01 14:41:54.525508: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=2)
2019-02-01 14:41:54.525520: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=3)
2019-02-01 14:41:54.525529: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=4)
 ** On entry to GEMM_EX  parameter number 9 had an illegal value
 ** On entry to GEMM_EX  parameter number 9 had an illegal value
 ** On entry to GEMM_EX  parameter number 9 had an illegal value
 ** On entry to GEMM_EX  parameter number 9 had an illegal value
 ** On entry to GEMM_EX  parameter number 9 had an illegal value
...
 ** On entry to GEMM_EX  par2019-02-01 14:41:54.525543: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpo
se, m=384, n=96, k=8192, alpha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=5)
2019-02-01 14:41:54.525553: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=6)
2019-02-01 14:41:54.525562: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=7)
2019-02-01 14:41:54.525572: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=8)
2019-02-01 14:41:54.525582: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=9)
2019-02-01 14:41:54.525592: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=10)
2019-02-01 14:41:54.525601: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=11)
2019-02-01 14:41:54.525611: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=12)
2019-02-01 14:41:54.525621: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=13)
2019-02-01 14:41:54.525633: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=14)
2019-02-01 14:41:54.525643: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=15)
2019-02-01 14:41:54.525653: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=16)
2019-02-01 14:41:54.525662: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=17)
2019-02-01 14:41:54.525672: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=99)
2019-02-01 14:41:54.525676: I external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_blas.cc:2226] DoBlasGemmWithAlgorithm returning false because algorithm 99 uses tensor ops, but the input data type is not fp16.
2019-02-01 14:41:54.525682: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=100)
2019-02-01 14:41:54.525686: I external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_blas.cc:2226] DoBlasGemmWithAlgorithm returning false because algorithm 100 uses tensor ops, but the input data type is not fp16.
2019-02-01 14:41:54.525691: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=101)
2019-02-01 14:41:54.525695: I external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_blas.cc:2226] DoBlasGemmWithAlgorithm returning false because algorithm 101 uses tensor ops, but the input data type is not fp16.
2019-02-01 14:41:54.525700: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=102)
2019-02-01 14:41:54.525704: I external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_blas.cc:2226] DoBlasGemmWithAlgorithm returning false because algorithm 102 uses tensor ops, but the input data type is not fp16.
2019-02-01 14:41:54.525710: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=103)
2019-02-01 14:41:54.525715: I external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_blas.cc:2226] DoBlasGemmWithAlgorithm returning false because algorithm 103 uses tensor ops, but the input data type is not fp16.
2019-02-01 14:41:54.525721: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:4209] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemmWithAlgorithm(transa=NoTranspose, transb=Transpose, m=384, n=96, k=8192, alp
ha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=384, computation_type=f32, algorithm=104)
2019-02-01 14:41:54.525724: I external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_blas.cc:2226] DoBlasGemmWithAlgorithm returning false because algorithm 104 uses tensor ops, but the input data type is not fp16.
2019-02-01 14:41:54.525729: I external/org_tensorflow/tensorflow/compiler/xla/util.cc:36] Internal: Unable to autotune cuBLAS gemm on stream 0x5610733f2430; none of the 25 algorithms ran successfully
2019-02-01 14:41:54.526256: I external/org_tensorflow/tensorflow/compiler/xla/util.cc:37] *** Begin stack trace ***
...
2019-02-01 14:41:54.526265: I external/org_tensorflow/tensorflow/compiler/xla/service/gpu/gemm_thunk.cc:444] Autotune on GemmThunk %dot.3438 = f32[96,384]{0,1} dot(f32[8192,96]{1,0} %bitcast.123, f32[8192,384]{1,0} %fusion.250), lhs_contracting_dims={0}, rhs_contracting_dims={0} unsuccessful.  Will use generic gemm.
2019-02-01 14:41:54.526272: I external/org_tensorflow/tensorflow/stream_executor/stream.cc:3888] [stream=0x5610733f2430,impl=0x561073379d50] Called Stream::ThenBlasGemm(transa=NoTranspose, transb=Transpose, m=96, n=384, k=8192, alpha=1, a=0x7f04dce00000, lda=96, b=0x7f04e4b78000, ldb=384, beta=0, c=0x7f051e224000, ldc=96)
2019-02-01 14:41:54.526279: I external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_blas.cc:1940] doing cuBLAS SGEMM: at=0 bt=1 m=96 n=384 k=8192 alpha=1.000000 a=0x7f04dce00000 lda=96 b=0x7f04e4b78000 ldb=384 beta=0.000000 c=0x7f051e224000 ldc=96

hawkinsp on Feb 1, 2019