Keep getting the warning: external/org_tensorflow/tensorflow/compiler/xla/service/gpu/ir_emitter_triton.cc:758] Shared memory size limit exceeded

[JyChang012]

I got the following warning when running a jitted JAX function on Nvidia v100 GPU.

external/org_tensorflow/tensorflow/compiler/xla/service/gpu/ir_emitter_triton.cc:758] Shared memory size limit exceeded.     

The warning was only raised for the first several calls to the function, so I suspect it was during compilation.
Does anyone know the cause?

[BugMaker-Boyan]

I also got the warning, is there anyone can help?

[samblouir]

Same here, haven’t had this in Jax in the past two years or so. Are you on the latest version?

[lockwo]

I have also been encountering this, breaking systems that since it registered as an error. Is there any fix for this or way to suppress it? I tried

import os
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

But still see it.

[sokrypton]

Try add this at the top of the script

import os
os.environ["XLA_FLAGS"] = "--xla_gpu_enable_triton_gemm=false"

thanks @pschuh

[mattjj]

Thanks, @sokrypton ! Context is the discussion on #15290.

I hadn't seen this thread until just now.

XLA:GPU recently started doing some codegen using Triton, and it looks like this is a new error resulting from that codegen path.

Can folks provide repros we can share with the XLA:GPU team?

[JyChang012]

Thanks @mattjj @sokrypton ! I noticed this warning when comparing training speed of a JAX ASR model I wrote with the same model written in PyTorch on V100 GPU. JAX was much slower. Hope this could also fix the performance issue.

[hawkinsp]

@JyChang012 If anything, Triton should make run time faster. It may make compile time slower. If you'd like us to look at the runtime speed of something please share a benchmark.

[sokrypton]

@hawkinsp the issue is when we want to run the model for inference.
A simple model that used to take seconds to compile now takes over a minute. The prediction takes milliseconds.

Maybe another option would be to provide means to compile model for inference vs. training? Assuming triition provides a speedup for training?

[sergachev]

The warning "Shared memory size limit exceeded." is harmless and was suppressed in XLA 5th of April.
We are working on making compilation with Triton GEMM (autotuning) fast.