Change:
Hello I have a training loop running without any errors or warnings on GPUs using Pytorch 2.0.1. When I try to modify my learning loop to run on TPUs I get the error Pytorch XLA developers solved the problem. But after the problem is solved I get the warning below. /home/mfatih/env38/lib/python3.8/site- UserWarning: aten::reshape: an autograd kernel was not registered to the Autograd key(s) but we are trying to backprop through it. This may lead to silently incorrect behavior. This behavior is deprecated and will be removed in a future version of PyTorch. If your operator is differentiable, please ensure you have registered an autograd kernel to the correct Autograd key (e.g. DispatchKey::Autograd, DispatchKey::CompositeImplicitAutograd). If your operator is not differentiable, or to squash this warning and use the previous behavior, please register torch::CppFunction::makeFallthrough() to DispatchKey::Autograd. (Triggered internally at ) Variable.executionengine.runbackward( # Calls into the C++ engine to run the backward pass Pytorch XLA developers told me that the warning is related to Pytorch, not Pytorch-XLA. Is this warning an indication of the wrong backward computations? You can regenerate the error using the public . Please change the working directory on in the config file. Run the for debugging. After a successful short train and val cycle we get the warning. The warning occurs in operations placed in the functions in during backward pass. Here is my environment PyTorch version: 2.2.0.dev20231213+cpu Is debug build: False CUDA used to build PyTorch: None ROCM used to build PyTorch: N/A OS: Ubuntu 20.04.4 LTS (x8664) GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0 Clang version: Could not collect CMake version: version 3.26.0 Libc version: glibc-2.31 Python version: 3.8.10 (default, Nov 22 2023, 10:22:35) [GCC 9.4.0] (64-bit runtime) Python platform: Linux-5.13.0-1027-gcp-x8664-with-glibc2.29 Is CUDA available: False CUDA runtime version: No CUDA CUDAMODULELOADING set to: N/A GPU models and configuration: No CUDA Nvidia driver version: No CUDA cuDNN version: No CUDA HIP runtime version: N/A MIOpen runtime version: N/A Is XNNPACK available: True CPU: Architecture: x8664 CPU op-mode(s): 32-bit, 64-bit Byte Order: Little Endian Address sizes: 46 bits physical, 48 bits virtual CPU(s): 96 On-line CPU(s) list: 0-95 Thread(s) per core: 2 Core(s) per socket: 24 Socket(s): 2 NUMA node(s): 2 Vendor ID: GenuineIntel CPU family: 6 Model: 85 Model name: Intel(R) Xeon(R) CPU @ 2.00GHz Stepping: 3 CPU MHz: 2000.186 BogoMIPS: 4000.37 Hypervisor vendor: KVM Virtualization type: full L1d cache: 1.5 MiB L1i cache: 1.5 MiB L2 cache: 48 MiB L3 cache: 77 MiB NUMA node0 CPU(s): 0-23,48-71 NUMA node1 CPU(s): 24-47,72-95 Vulnerability Itlb multihit: Not affected Vulnerability L1tf: Mitigation; PTE Inversion Vulnerability Mds: Mitigation; Clear CPU buffers; SMT Host state unknown Vulnerability Meltdown: Mitigation; PTI Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp Vulnerability Spectre v1: Mitigation; usercopy/swapgs barriers and _user pointer sanitization Vulnerability Spectre v2: Mitigation; Retpolines, IBPB conditional, IBRSFW, STIBP conditional, RSB filling Vulnerability Srbds: Not affected Vulnerability Tsx async abort: Mitigation; Clear CPU buffers; SMT Host state unknown Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constanttsc repgood nopl xtopology nonstoptsc cpuid tscknownfreq pni pclmulqdq ssse3 fma cx16 pcid sse41 sse42 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahflm abm 3dnowprefetch invpcidsingle pti ssbd ibrs ibpb stibp fsgsbase tscadjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves arat mdclear archcapabilities Versions of relevant libraries: [pip3] numpy==1.24.1 [pip3] torch==2.2.0.dev20231213+cpu [pip3] torch-xla==2.2.0+git5577dd7 [pip3] torchaudio==2.2.0.dev20231213+cpu [pip3] torchmetrics==1.2.1 [pip3] torchsummary==1.5.1 [pip3] torchvision==0.18.0.dev20231213+cpu [conda] Could not collect cc
From the log it seems like it suggese XLA need to manually register the backward for reshape which is a bit weird. This is the first time I see this error, if you can shed some light on this issue that would be great!
Hey! this warning recently as a set of more testing to ensure custom ops are properly registered. This is expected to trigger if you register a backend kernel (XLA key) but no autograd imple (AutogradXLA) for a given op. Since reshape() cannot have an autograd formula (it is sometimes a view and sometimes not), then you should register your kernel as a CompositeImplicitAutograd kernel (this will properly register onto AutogradXLA and remove this warning)
Thanks ! This make sense! I have a follow up question, we registered reshape and reshape symint in and This file will then get codegen into and . What's the correct way of marking to be ?
Ho I'm not familiar with how your codegen works... But in principles, I guess you want to do the same as NestedTensor here which uses for the dispatch key in do you have more context on this by any chance?
it's probably enough for you to just register that kernel to both the and dispatch keys (this is similar to what CompositeImplicitAutograd does, but it does it for all backends). I'm a bit confused though - Why do you need a custom reshape impl? And how are you going to get training support for it? (do you have a custom implementation of your reshape's backward)
Haha this is actually by you in for functionization. Do you think we still need this trick for functionization?
Yes - I think we should be able to remove it. Do you mind trying that? That should just run the existing kernel defined . And that should be identical to the current behavior, since the "registration" for reshape in XLA just calls back into ATen ()
I merged , let's wait and see if tmr's nightly will resolve this issue.
Thank you. I will test it and give feedback. Till now, I have run my experiments with the warning. Do I need to rerun the experiments?
OK, with the recent nighly versions of Pytorch and Pytorch-XLA I do not observe the warning anymore while running my experiments. Thank you. Should I rerun the experiments that were run with past Pytorch versions throwing the warning? Does this update only remove the warning or does it also change the calculations that I ran with the previous Pytorch versions? Can I rely on my previous experiments?
Sorry we have to revert the pr since it creates regression on nightly.
Hello I am using nightly without any problem. What should I do? Wait for new update?
if you don't see any regression in performance, you can likely just keep using today's nightly. I am not even sure if that warning has any real impact on the model yet.
Thank you for the answer Reverting will cause problem for me. But reverting will only regenerate the warning. In my experiments, till now, I have not observed any difference with or without the warning. But I am not 100% sure. I'm looking forward to any more information you can give me.
Seeing the same error. Any solutions?
Yea I need to find someone to look into this issue. It seems like removing the reshape lowering will introduce a regression in performance.
I am facing the same issue, I don't even have a "reshape" call in my code anywhere.
reopened , will try to fix this before the 2.3 release branch cut.
will lead the investigation on the regression that blocks merging Regression only happens in dynamo so my guess is it has something to do with functionization pass..
Running through a small example with on HEAD and on , the snippet HLOo generated for in both cases are: Now will try to experiment with functionalizaton explicitly on/off to see if there are some unexpected decompositions.
I am also experiencing the same problem. I have also seen significant reduction in the learning capabilities of my model. My model has more than 80% accuracy when I trained it in another device without xla. However, its accuracy has been reduced to less than 10% compared to when I trained on a GCP TPU VM with pytorch-xla. It is also apparent that loss values do not change between consecutive epochs during training. Here is a portion of the output including the error message: {'loss': 0.0288, 'gradnorm': 0., 'learningrate': 4.-05, 'epoch': 6.0} {'evalloss': 0., 'evalruntime': 5.4579, 'evalsamplespersecond': 615.621, 'evalstepspersecond': 19.238, 'epoch': 6.0} 20% 2502/ [28:36&lt;1:08:07, 2.45it/s/home/egecitepred/miniconda3/envs/ege/lib/python3.8/site- UserWarning: aten::reshape: an autograd kernel was not registered to the Autograd key(s) but we are trying to backprop through it. This may lead to silently incorrect behavior. This behavior is deprecated and will be removed in a future version of PyTorch. If your operator is differentiable, please ensure you have registered an autograd kernel to the correct Autograd key (e.g. DispatchKey::Autograd, DispatchKey::CompositeImplicitAutograd). If your operator is not differentiable, or to squash this warning and use the previous behavior, please register torch::CppFunction::makeFallthrough() to DispatchKey::Autograd. (Triggered internally at ) Variable.executionengine.runbackward( # Calls into the C++ engine to run the backward pass {'loss': 0.0289, 'gradnorm': 0., 'learningrate': 3.-05, 'epoch': 7.0} {'evalloss': 0., 'evalruntime': 5.4422, 'evalsamplespersecond': 617.398, 'evalstepspersecond': 19.294, 'epoch': 7.0} 23% 2919/ [31:44&lt;1:05:13, 2.45it/s/home/egecitepred/miniconda3/envs/ege/lib/python3.8/site- UserWarning: aten::reshape: an autograd kernel was not registered to the Autograd key(s) but we are trying to backprop through it. This may lead to silently incorrect behavior. This behavior is deprecated and will be removed in a future version of PyTorch. If your operator is differentiable, please ensure you have registered an autograd kernel to the correct Autograd key (e.g. DispatchKey::Autograd, DispatchKey::CompositeImplicitAutograd). If your operator is not differentiable, or to squash this warning and use the previous behavior, please register torch::CppFunction::makeFallthrough() to DispatchKey::Autograd. (Triggered internally at ) Variable.executionengine.runbackward( # Calls into the C++ engine to run the backward pass {'loss': 0.0288, 'gradnorm': 0., 'learningrate': 3.-05, 'epoch': 8.0} {'evalloss': 0., 'evalruntime': 5.4432, 'evalsamplespersecond': 617.286, 'evalstepspersecond': 19.29, 'epoch': 8.0} I am also receiving this error message at every epoch. Is it possible to register CompositeImplicitAutograd key to autograd kernel like the error message says? If someone could provide me instructions on registering this key, I could try to see if this fixes the error or at least the reduction of the performance. I don't have enough experience in this area yet. So, I couldn't find a way of registering the autograd key.
Hey thanks for reporting the issue. Just to confirm, the warning messages in the original issue should not have an impact in the performance. We tried to remove this warning message and merged , which caused a performance regression (which is now reverted in head). Which nightlies did you observe the performance regression in your case?
Running a quick experiment on resnet18, the performance is more or so similar with HEAD and reshape op removed. With HEAD, timing some runs of the resnet18 dynamo unit test: And timing some runs of the resnet18 dynamo unit test with removed: Before we continue further with the investigation, and I discussed that if simply registering this kernel to both the AutogradXLA and XLA dispatch keys would remove the warning message, we might just do that.. though I'm not too familiar this. do you know how can we register the kernel?
Hello This message seems to slightly reduced the performance of my model unlike others. I am not absolutely sure about this but I made some additional observations that lead me to believe that. I have tried the same model on a GPU without XLA. It has somewhat higher performance compared to TPU VM with PytorchXLA. However, my main problem is the reduction of learning capabilities after training. This error message reduced my result metrics after training from ~80% to ~20% under the same conditions. Also, my output for each epoch shows the same evalloss value. My other experiments on a GPU had the same eval loss value on epoch 1. It seems the model (with TPU XLA) is learning only for the first epoch and fails to properly backprop. Thus, it keeps repeating the first epoch again and again. Although I am not sure about this, it is the only explanation I could come up. I am not sure which nightly version I am using. However, I know that I downloaded this library on February 29th using this command: "pip install torch~=2.2.0 torch_xla[tpu]~=2.2.0 -f "
Hey I wonder if that's a different issue -- could you open a separate issue in PyTorch/XLA? Since Brian is OOO, do you have any idea on the kernel registration mentioned in
From the warning, you are hitting the AutogradXLA key since the fallback at that key is triggering this warning. So the problem is that your differentiable reshape implementation should be registered at that key. If you say above that "ResigisterAutogradXLA.h/cpp." gets generated then I would check if reshare is there. If it is, I would double check that this file is properly included in your binary and properly registers the kernels. You can also check to see all the kernels registered for that op and from where
Hey thanks for the pointers! I tried to register the op to with With this PR, I could see it get registered: However, now I'm getting the error: I assume something is wrong with our new lowering, which is just a redispatch to functionalization -- However, this seems to be what we're doing for some of similar ops -- Maybe I'm just confused here lol, do you see anything obvious that is wrong? Thanks for the help, again!

Discussion:
// TODO: gradOutput shape check // Resize and initialize result tensor. THCTensor_(resizeAs)(state, gradInput, input); <del> THCTensor_(newContiguous)(state, gradInput); </del> THCTensor_(zero)(state, gradInput); int batchSize;