[Bug]: RuntimeErrors when using TorchInductor for stable diffusion

[mreso]

🐛 Describe the bug

Hi,

when running a stable diffusion example I am experiencing a high number of non deterministic runtime errors (~50% of the times when running my example with inductor). I observed two flavors.

First error is:

RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered

Its much more frequent than the second which is:

RuntimeError: CUDA: Error- illegal address
+
RuntimeError: false INTERNAL ASSERT FAILED at "../c10/cuda/CUDAGraphsC10Utils.h":73, please report a bug to PyTorch. Unknown CUDA graph CaptureStatus32650

@soumith reported here that he saw the same set of errors.

Any thoughts on how I can help to collect more data or could create more reliable repros?

My specific example looks as follows:

Repro:

pip install --upgrade diffusers transformers

Then run:

import torchdynamo

from diffusers import StableDiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4") #ACCESS TOKEN REQUIRED
pipe = pipe.to("cuda")

@torchdynamo.optimize("inductor")
def apply(x):
    return pipe(x)

prompt = "a photo of an astronaut riding a horse on mars"
image = apply(prompt).images[0]  

I collected the stack traces of the failed attempts here: https://gist.github.com/mreso/982ef0b4c915eea6170fd81f0de16e98

Error logs

Example for the first error is:

Traceback (most recent call last):
  File "/home/mreso/torchdynamo/test.py", line 40, in <module>
    image = apply(prompt).images[0]
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 157, in _fn
    return fn(*args, **kwargs)
  File "/home/mreso/torchdynamo/test.py", line 29, in apply
    return pipe(x)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 214, in __call__
    text_inputs = self.tokenizer(
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 269, in <graph break in __call__>
    uncond_embeddings = self.text_encoder(uncond_input.input_ids.to(self.device))[0]
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1357, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/transformers/models/clip/modeling_clip.py", line 705, in forward
    @add_start_docstrings_to_model_forward(CLIP_TEXT_INPUTS_DOCSTRING)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 157, in _fn
    return fn(*args, **kwargs)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 870, in forward
    return compiled_f(
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 861, in new_func
    return compiled_fn(args)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 299, in new_fn
    fw_outs = call_func_with_args(compiled_fw, args, disable_amp=disable_amp)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 255, in call_func_with_args
    out = normalize_as_list(f(args))
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/compile_fx.py", line 194, in run
    return model(new_inputs_to_cuda)
  File "/tmp/torchinductor_mreso/qt/cqt3tpyb2awwxhp3vhrmtamvb6b57ctovraioa5mbqzemitewns2.py", line 915, in call
    kernel2.run(arg197_1, arg0_1, arg196_1, arg1_1, arg2_1, arg3_1, buf7, 77, 768, grid=grid(77), stream=stream0)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/triton_ops/autotune.py", line 170, in run
    result = launcher(
  File "<string>", line 4, in launcher
RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered

For the second error its:

[2022-10-21 00:11:28,203] torch._dynamo.convert_frame: [ERROR] WON'T CONVERT forward /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/diffusers/models/unet_2d_condition.py line 232
due to:
Traceback (most recent call last):
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 401, in clone_input
    result.copy_(x.clone())
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 410, in clone_input
    y = torch.clone(x)
RuntimeError: CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.

from user code:
   File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/diffusers/models/unet_2d_condition.py", line 286, in forward
    t_emb = t_emb.to(dtype=self.dtype)

Set torch._dynamo.config.verbose=True for more information
==========
  0%|                                                                                                                                                                                                                                                                                                                                              | 0/51 [00:00<?, ?it/s]
Traceback (most recent call last):
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 65, in preserve_rng_state
    yield
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 579, in create_aot_dispatcher_function
    return aot_dispatch_base(flat_fn, fake_flat_tensor_args, aot_config)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 295, in aot_dispatch_base
    compiled_fw = aot_config.fw_compiler(fw_module, flat_args)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 85, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/compile_fx.py", line 362, in fw_compiler
    return compile_fx_inner(
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_dynamo/debug_utils.py", line 464, in debug_wrapper
    compiled_fn = compiler_fn(gm, example_inputs, **kwargs)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/debug.py", line 178, in inner
    return fn(*args, **kwargs)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/compile_fx.py", line 127, in compile_fx_inner
    compiled_fn = graph.compile_to_fn()
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/graph.py", line 339, in compile_to_fn
    return self.compile_to_module().call
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_dynamo/utils.py", line 85, in time_wrapper
    r = func(*args, **kwargs)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/graph.py", line 329, in compile_to_module
    mod = PyCodeCache.load(code)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/codecache.py", line 212, in load
    exec(code, mod.__dict__, mod.__dict__)
  File "/tmp/torchinductor_mreso/gb/cgb5kyxk5qcj67xbzxmiuwpxlrcjspq7dyscbrllcrgknkoahswm.py", line 94, in <module>
    async_compile.wait(globals())
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/codecache.py", line 352, in wait
    scope[key] = result.result()
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/codecache.py", line 260, in result
    kernel = self.kernel = _load_kernel(self.source_code)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/codecache.py", line 246, in _load_kernel
    kernel.precompile()
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/triton_ops/autotune.py", line 58, in precompile
    self.launchers = [
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/triton_ops/autotune.py", line 59, in <listcomp>
    self._precompile_config(c, warm_cache_only_with_cc)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/triton_ops/autotune.py", line 83, in _precompile_config
    binary = triton.compile(
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/triton/compiler.py", line 1263, in compile
    return CompiledKernel(name, so_cache_manager._make_path(so_name), fn_cache_manager.cache_dir, device)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/triton/compiler.py", line 1296, in __init__
    mod, func, n_regs, n_spills = _triton.code_gen.load_binary(metadata["name"], self.asm["cubin"], self.shared, device)
RuntimeError: CUDA: Error- illegal address

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/home/mreso/torchdynamo/test.py", line 24, in <module>
    image = apply(prompt).images[0]
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 157, in _fn
    return fn(*args, **kwargs)
  File "/home/mreso/torchdynamo/test.py", line 21, in apply
    return pipe(x)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 214, in __call__
    text_inputs = self.tokenizer(
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 326, in <graph break in __call__>
    noise_pred = self.unet(latent_model_input, t, encoder_hidden_states=text_embeddings).sample
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1357, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/diffusers/models/unet_2d_condition.py", line 281, in forward
    t_emb = self.time_proj(timesteps)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1357, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/diffusers/models/embeddings.py", line 91, in forward
    def forward(self, timesteps):
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 157, in _fn
    return fn(*args, **kwargs)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 870, in forward
    return compiled_f(
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 856, in new_func
    compiled_fn = create_aot_dispatcher_function(
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 579, in create_aot_dispatcher_function
    return aot_dispatch_base(flat_fn, fake_flat_tensor_args, aot_config)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/contextlib.py", line 137, in __exit__
    self.gen.throw(typ, value, traceback)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 69, in preserve_rng_state
    torch.cuda.set_rng_state(cuda_rng_state)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/cuda/random.py", line 64, in set_rng_state
    _lazy_call(cb)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/cuda/__init__.py", line 176, in _lazy_call
    callable()
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/cuda/random.py", line 62, in cb
    default_generator.set_state(new_state_copy)
RuntimeError: false INTERNAL ASSERT FAILED at "../c10/cuda/CUDAGraphsC10Utils.h":73, please report a bug to PyTorch. Unknown CUDA graph CaptureStatus32650

Did Dynamo succeed?

• Does dynamo.optimize("eager") succeed?

Did AOT succeed?

• Did dynamo.optimize("aot_eager") succeed?

Did Inductor succeed?

• Does dynamo.optimize("inductor") succeed?

Minified repro

Not applicable as the created minified examples do not produce errors.

$python /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/minifier_launcher.py
Traceback (most recent call last):
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/minifier_launcher.py", line 1028, in <module>
    minifier(
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/fx_minifier.py", line 96, in minifier
    raise RuntimeError("Input graph did not fail the tester")
RuntimeError: Input graph did not fail the tester

[mreso]

@ngimel After disabling cuda graphs and setting PYTORCH_NO_CUDA_MEMORY_CACHING=1 I don't see the runtime errors as often anymore. In 20+ runs I had this error only once so far:

terminate called after throwing an instance of 'c10::Error'
  what():  CUDA error: an illegal memory access was encountered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1.
Exception raised from c10_cuda_check_implementation at ../c10/cuda/CUDAException.cpp:31 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x3e (0x7f7c208decce in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libc10.so)
frame pytorch/torchdynamo#1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, std::string const&) + 0x5c (0x7f7c208a58b9 in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libc10.so)
frame pytorch/torchdynamo#2: c10::cuda::c10_cuda_check_implementation(std::string const&, std::string const&, int, bool) + 0xb4 (0x7f7c209869b4 in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libc10_cuda.so)
frame pytorch/torchdynamo#3: <unknown function> + 0x1ac97 (0x7f7c20956c97 in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libc10_cuda.so)
frame pytorch/torchdynamo#4: <unknown function> + 0x4f28f6 (0x7f7c4baf18f6 in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame pytorch/torchdynamo#5: <unknown function> + 0x40ea7 (0x7f7c208c2ea7 in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libc10.so)
frame pytorch/torchdynamo#6: c10::TensorImpl::~TensorImpl() + 0x1bb (0x7f7c208bc9bb in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libc10.so)
frame pytorch/torchdynamo#7: c10::TensorImpl::~TensorImpl() + 0x9 (0x7f7c208bca59 in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libc10.so)
frame pytorch/torchdynamo#8: <unknown function> + 0x7482e8 (0x7f7c4bd472e8 in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame pytorch/torchdynamo#9: THPVariable_subclass_dealloc(_object*) + 0x2a5 (0x7f7c4bd475d5 in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame pytorch/torchdynamo#10: python() [0x4f0b40]
frame pytorch/torchdynamo#11: python() [0x4f895d]
<omitting python frames>
frame pytorch/torchdynamo#13: python() [0x4e729a]
frame pytorch/torchdynamo#16: python() [0x4e729a]
frame pytorch/torchdynamo#19: python() [0x4e729a]
frame pytorch/torchdynamo#22: python() [0x4e729a]
frame pytorch/torchdynamo#25: python() [0x4e729a]
frame pytorch/torchdynamo#28: <unknown function> + 0x758b86 (0x7f7c4bd57b86 in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame pytorch/torchdynamo#29: python() [0x4e729a]
frame pytorch/torchdynamo#32: <unknown function> + 0x758b01 (0x7f7c4bd57b01 in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame pytorch/torchdynamo#33: python() [0x4e729a]
frame pytorch/torchdynamo#35: python() [0x505d51]
frame pytorch/torchdynamo#37: <unknown function> + 0x758b86 (0x7f7c4bd57b86 in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame pytorch/torchdynamo#38: python() [0x4e729a]
frame pytorch/torchdynamo#41: python() [0x5cc3f3]
frame pytorch/torchdynamo#44: <unknown function> + 0x758b01 (0x7f7c4bd57b01 in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame pytorch/torchdynamo#45: python() [0x4f8923]
frame pytorch/torchdynamo#47: <unknown function> + 0x758b01 (0x7f7c4bd57b01 in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame pytorch/torchdynamo#48: python() [0x4e729a]
frame pytorch/torchdynamo#51: <unknown function> + 0x758b86 (0x7f7c4bd57b86 in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame pytorch/torchdynamo#52: python() [0x4e729a]
frame pytorch/torchdynamo#55: python() [0x5cc3f3]
frame pytorch/torchdynamo#58: <unknown function> + 0x758b01 (0x7f7c4bd57b01 in /home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame pytorch/torchdynamo#59: python() [0x4f8923]
frame pytorch/torchdynamo#61: python() [0x4e729a]

Aborted

[ngimel]

Hm, I was hoping the error would become more frequent with no memory caching, can you also please enable CUDA_LAUNCH_BLOCKING=1 to get a better stack trace for when the error happens?

[mreso]

@ngimel This is the error with CUDA_LAUNCH_BLOCKING=1:

Traceback (most recent call last):
  File "/home/mreso/torchdynamo/test.py", line 42, in <module>
    image = apply(prompt).images[0]
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 157, in _fn
    return fn(*args, **kwargs)
  File "/home/mreso/torchdynamo/test.py", line 31, in apply
    return pipe(x)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 214, in __call__
    text_inputs = self.tokenizer(
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py", line 229, in <graph break in __call__>
    text_embeddings = self.text_encoder(text_input_ids.to(self.device))[0]
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1357, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/transformers/models/clip/modeling_clip.py", line 705, in forward
    @add_start_docstrings_to_model_forward(CLIP_TEXT_INPUTS_DOCSTRING)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py", line 157, in _fn
    return fn(*args, **kwargs)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 870, in forward
    return compiled_f(
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 861, in new_func
    return compiled_fn(args)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 299, in new_fn
    fw_outs = call_func_with_args(compiled_fw, args, disable_amp=disable_amp)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/functorch/_src/aot_autograd.py", line 255, in call_func_with_args
    out = normalize_as_list(f(args))
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/compile_fx.py", line 194, in run
    return model(new_inputs_to_cuda)
  File "/tmp/torchinductor_mreso/gx/cgxznf3od6zmaxtthle7gcqk3lcynngigfrcsnz6hn5s3haqbow3.py", line 1518, in call
    kernel16.run(buf302, buf301, buf303, 768, grid=grid(768), stream=stream0)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/triton_ops/autotune.py", line 170, in run
    result = launcher(
  File "<string>", line 4, in launcher
RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered

I've copied the code of the tmp file into a gist: https://gist.github.com/mreso/e2d38beb6ece8b3863cffe5e5e68931a

When running it separately it also fails in a non deterministic manner with the same error:

Traceback (most recent call last):
  File "/home/mreso/torchdynamo/cgxznf3od6zmaxtthle7gcqk3lcynngigfrcsnz6hn5s3haqbow3.py", line 1723, in <module>
    print_performance(lambda: call([arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1, arg9_1, arg10_1, arg11_1, arg12_1, arg13_1, arg14_1, arg15_1, arg16_1, arg17_1, arg18_1, arg19_1, arg20_1, arg21_1, arg22_1, arg23_1, arg24_1, arg25_1, arg26_1, arg27_1, arg28_1, arg29_1, arg30_1, arg31_1, arg32_1, arg33_1, arg34_1, arg35_1, arg36_1, arg37_1, arg38_1, arg39_1, arg40_1, arg41_1, arg42_1, arg43_1, arg44_1, arg45_1, arg46_1, arg47_1, arg48_1, arg49_1, arg50_1, arg51_1, arg52_1, arg53_1, arg54_1, arg55_1, arg56_1, arg57_1, arg58_1, arg59_1, arg60_1, arg61_1, arg62_1, arg63_1, arg64_1, arg65_1, arg66_1, arg67_1, arg68_1, arg69_1, arg70_1, arg71_1, arg72_1, arg73_1, arg74_1, arg75_1, arg76_1, arg77_1, arg78_1, arg79_1, arg80_1, arg81_1, arg82_1, arg83_1, arg84_1, arg85_1, arg86_1, arg87_1, arg88_1, arg89_1, arg90_1, arg91_1, arg92_1, arg93_1, arg94_1, arg95_1, arg96_1, arg97_1, arg98_1, arg99_1, arg100_1, arg101_1, arg102_1, arg103_1, arg104_1, arg105_1, arg106_1, arg107_1, arg108_1, arg109_1, arg110_1, arg111_1, arg112_1, arg113_1, arg114_1, arg115_1,
arg116_1, arg117_1, arg118_1, arg119_1, arg120_1, arg121_1, arg122_1, arg123_1, arg124_1, arg125_1, arg126_1, arg127_1, arg128_1, arg129_1, arg130_1, arg131_1, arg132_1, arg133_1, arg134_1, arg135_1, arg136_1, arg137_1, arg138_1, arg139_1, arg140_1, arg141_1, arg142_1, arg143_1, arg144_1, arg145_1, arg146_1, arg147_1, arg148_1, arg149_1, arg150_1, arg151_1, arg152_1, arg153_1, arg154_1, arg155_1, arg156_1, arg157_1, arg158_1, arg159_1, arg160_1, arg161_1, arg162_1, arg163_1, arg164_1, arg165_1, arg166_1, arg167_1, arg168_1, arg169_1, arg170_1, arg171_1, arg172_1, arg173_1, arg174_1, arg175_1, arg176_1, arg177_1, arg178_1, arg179_1, arg180_1, arg181_1, arg182_1, arg183_1, arg184_1, arg185_1, arg186_1, arg187_1, arg188_1, arg189_1, arg190_1, arg191_1, arg192_1, arg193_1, arg194_1, arg195_1, arg196_1, arg197_1]))
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/utils.py", line 120, in print_performance
    timings = [timed(fn, args, times) for _ in range(repeat)]
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/utils.py", line 120, in <listcomp>
    timings = [timed(fn, args, times) for _ in range(repeat)]
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/utils.py", line 111, in timed
    result = model(*example_inputs)
  File "/home/mreso/torchdynamo/cgxznf3od6zmaxtthle7gcqk3lcynngigfrcsnz6hn5s3haqbow3.py", line 1723, in <lambda>
    print_performance(lambda: call([arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1, arg9_1, arg10_1, arg11_1, arg12_1, arg13_1, arg14_1, arg15_1, arg16_1, arg17_1, arg18_1, arg19_1, arg20_1, arg21_1, arg22_1, arg23_1, arg24_1, arg25_1, arg26_1, arg27_1, arg28_1, arg29_1, arg30_1, arg31_1, arg32_1, arg33_1, arg34_1, arg35_1, arg36_1, arg37_1, arg38_1, arg39_1, arg40_1, arg41_1, arg42_1, arg43_1, arg44_1, arg45_1, arg46_1, arg47_1, arg48_1, arg49_1, arg50_1, arg51_1, arg52_1, arg53_1, arg54_1, arg55_1, arg56_1, arg57_1, arg58_1, arg59_1, arg60_1, arg61_1, arg62_1, arg63_1, arg64_1, arg65_1, arg66_1, arg67_1, arg68_1, arg69_1, arg70_1, arg71_1, arg72_1, arg73_1, arg74_1, arg75_1, arg76_1, arg77_1, arg78_1, arg79_1, arg80_1, arg81_1, arg82_1, arg83_1, arg84_1, arg85_1, arg86_1, arg87_1, arg88_1, arg89_1, arg90_1, arg91_1, arg92_1, arg93_1, arg94_1, arg95_1, arg96_1, arg97_1, arg98_1, arg99_1, arg100_1, arg101_1, arg102_1, arg103_1, arg104_1, arg105_1, arg106_1, arg107_1, arg108_1, arg109_1, arg110_1, arg111_1, arg112_1, arg113_1, arg114_1, arg115_1,
arg116_1, arg117_1, arg118_1, arg119_1, arg120_1, arg121_1, arg122_1, arg123_1, arg124_1, arg125_1, arg126_1, arg127_1, arg128_1, arg129_1, arg130_1, arg131_1, arg132_1, arg133_1, arg134_1, arg135_1, arg136_1, arg137_1, arg138_1, arg139_1, arg140_1, arg141_1, arg142_1, arg143_1, arg144_1, arg145_1, arg146_1, arg147_1, arg148_1, arg149_1, arg150_1, arg151_1, arg152_1, arg153_1, arg154_1, arg155_1, arg156_1, arg157_1, arg158_1, arg159_1, arg160_1, arg161_1, arg162_1, arg163_1, arg164_1, arg165_1, arg166_1, arg167_1, arg168_1, arg169_1, arg170_1, arg171_1, arg172_1, arg173_1, arg174_1, arg175_1, arg176_1, arg177_1, arg178_1, arg179_1, arg180_1, arg181_1, arg182_1, arg183_1, arg184_1, arg185_1, arg186_1, arg187_1, arg188_1, arg189_1, arg190_1, arg191_1, arg192_1, arg193_1, arg194_1, arg195_1, arg196_1, arg197_1]))
  File "/home/mreso/torchdynamo/cgxznf3od6zmaxtthle7gcqk3lcynngigfrcsnz6hn5s3haqbow3.py", line 1518, in call
    kernel16.run(buf302, buf301, buf303, 768, grid=grid(768), stream=stream0)
  File "/home/mreso/miniconda3/envs/torchdynamo2/lib/python3.9/site-packages/torch/_inductor/triton_ops/autotune.py", line 170, in run
    result = launcher(
  File "<string>", line 4, in launcher
RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered

[ngimel]

Ok it looks like the indices tensor that is sent to kernel16 is invalid, pytorch/pytorch#93538

[mreso]

I experimented a bit with kernel16 and looked at the inputs.

It has two input (buf302, buf301) and one output (303) buffer:

kernel16.run(buf302, buf301, buf303, 768, grid=grid(768), stream=stream0)

I compared the values of the input buffers for a successful and failing run and it seems like buf301 is different in the failing case.

Haven't yet looked closer into the other kernels, but can this behavior be easily explained @ngimel?

[ngimel]

I would expect buf302 (indices, of type int64) to be different in the failing case, can you save the failing buffers? Also arg197_1 from which buf302 is computed would be interesting.

[mreso]

@ngimel saved and uploaded the relevant buffers buffers.zip. The file contains fail and success case as pt files as well as the src. buf301 is changing as do the buffers it's computed from. This is the output of torch.all(torch.isclose()) on the relevant arguments and buffers between failing and success:

arg194_1: False
arg195_1: False
arg197_1: True
buf266: False
buf273: False
buf290: False
buf297: False
buf301: False
buf302: True
buf303: True

[ngimel]

argmax is computing completely wrong result, minimum repro:

import torch
import torch._inductor
import torch._inductor.config as config

#torch._dynamo.config.dynamic_shapes = True
torch._inductor.config.triton.cudagraphs = False
torch._inductor.config.debug=True
import torch._dynamo as torchdynamo
@torchdynamo.optimize("inductor")
def argmax(x):
    return x.argmax(0)

x=torch.ones(77, device="cuda", dtype = torch.int64)
print(x.argmax(0), argmax(x))

expected a number between 0 and 76 (eager outputs 0), result is 992

[ngimel]

Triton generated kernel seems correct, but with AOT compilation produces wrong results, filed triton-lang/triton#802