gopikrsmscs/torch-issues huggingface dataset:
[FSDP] using CPUOffload cannot make the code runing stop
Labels: triaged, module: fsdp

### 🐛 Describe the bug

I was tried to use offload fsdp to run the vicuna training, 

however, after I run this command. the code looks like keeping still for 3-4hours and prints nothing like and the last warning before it stay  is 

```
Please use torch.distributed.reduce_scatter_tensor instead.
  "torch.distributed._reduce_scatter_base is a private function and will "
```


### Versions



```
torchrun --nproc_per_node=8 --master_port=20002 fastchat/train/train_mem.py \
    --model_name_or_path /home/ma-user/work/jianqiao/alpca-lora/weights  \
    --data_path data/sharegpt_clean_gpt4_openchat_v2_maxlen1024.json \
    --bf16 False \
    --fp16 True \
    --output_dir output_vicuna \
    --num_train_epochs 3 \
    --per_device_train_batch_size 1 \
    --per_device_eval_batch_size 1 \
    --gradient_accumulation_steps 16 \
    --evaluation_strategy "no" \
    --save_strategy "steps" \
    --save_steps 1200 \
    --save_total_limit 10 \
    --learning_rate 2e-5 \
    --weight_decay 0. \
    --warmup_ratio 0.03 \
    --lr_scheduler_type "cosine" \
    --logging_steps 1 \
    --fsdp "full_shard offload auto_wrap" \
    --fsdp_transformer_layer_cls_to_wrap 'LlamaDecoderLayer' \
    --tf32 False \
    --model_max_length 1024 \
    --gradient_checkpointing True \
    --lazy_preprocess True \
    --data_key items
```

Could anybody help to solve this issue?

Thanks.

cc @zhaojuanmao @mrshenli @rohan-varma @awgu
Compile error PyTorch 2.0.1 / GCC 13.1.0
Labels: module: build, triaged

### 🐛 Describe the bug

I came across a strange error when compiling PyTorch (only C++ API, Python disabled) on A64FX using GCC 13.1.0.

When compiling the file `aten/src/ATen/native/mkldnn/Pooling.cpp` GCC turns on `-Werror` without turning on `-Wno-dangling-reference` so that
```
aten/src/ATen/core/IListRef_inl.h:171:17: error: possibly dangling reference to a temporary [-Werror=dangling-reference]
  171 |     const auto& ivalue = (*it).get();
```
This seems to be triggered by line 534 in `caffe2/CMakeLists.txt`:
```
# Required workaround for LLVM 9 includes.
if(NOT MSVC)
  set_source_files_properties(${TORCH_SRC_DIR}/csrc/jit/tensorexpr/llvm_jit.cpp PROPERTIES COMPILE_FLAGS -Wno-noexcept-type)
  # Force -Werror on several files
  set_source_files_properties(${CMAKE_CURRENT_LIST_DIR}/../aten/src/ATen/native/mkldnn/Pooling.cpp PROPERTIES COMPILE_FLAGS "-Werror")
endif()
```
Is this special treatment of the file `Pooling.cpp` really necessary? If so, one should maybe add `-Wno-dangling-reference`.

### Error logs

_No response_

### Minified repro

_No response_

### Versions

Collecting environment information...
PyTorch version: N/A
Is debug build: N/A
CUDA used to build PyTorch: N/A
ROCM used to build PyTorch: N/A

OS: Rocky Linux release 8.4 (Green Obsidian) (aarch64)
GCC version: (GCC) 8.5.0 20210514 (Red Hat 8.5.0-15)
Clang version: 12.0.1 (Red Hat 12.0.1-4.module+el8.5.0+715+58f51d49)
CMake version: version 3.20.2
Libc version: glibc-2.17

Python version: 3.6.8 (default, Nov  9 2021, 14:47:09)  [GCC 8.5.0 20210514 (Red Hat 8.5.0-3)] (64-bit runtime)
Python platform: Linux-4.18.0-305.25.1.el8_4.aarch64-aarch64-with-centos-8.4-Green_Obsidian
Is CUDA available: N/A
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: N/A

CPU:
Architecture:        aarch64
Byte Order:          Little Endian
CPU(s):              48
On-line CPU(s) list: 0-47
Thread(s) per core:  1
Core(s) per cluster: 12
Socket(s):           -
Cluster(s):          4
NUMA node(s):        4
Vendor ID:           FUJITSU
Model:               0
Model name:          A64FX
Stepping:            0x1
BogoMIPS:            200.00
NUMA node0 CPU(s):   0-11
NUMA node1 CPU(s):   12-23
NUMA node2 CPU(s):   24-35
NUMA node3 CPU(s):   36-47
Flags:               fp asimd evtstrm sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm fcma dcpop sve

Versions of relevant libraries:
[pip3] numpy==1.14.3
[conda] Could not collect

cc @malfet @seemethere @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
[pytorch] replace __FILE__ with __FILE_NAME__ for exceptions
Labels: fb-exported, Stale, topic: not user facing

Summary: Replace uses of the `__FILE__` macro with `__FILE_NAME__` for pytorch exceptions. This will prevent embedding build system paths in the build output, which will cause varying output depending on configuration.

Test Plan:
Build and inspect assert strings before change:
```
$ buck2 build --show-output //fbobjc/Libraries/FBPyTorchCore:torch_core_msg_opsApple
$ strings buck-out/v2/gen/fbsource/c969ccb4ab4ae099/fbobjc/Libraries/FBPyTorchCore/__torch_core_msg_opsApple__/libtorch_core_msg_opsApple.pic.a | grep FAIL | grep buck | head -n 3
*current_device == options.device() INTERNAL ASSERT FAILED at "buck-out/v2/gen/fbsource/c969ccb4ab4ae099/fbobjc/Libraries/FBPyTorchCore/__torch_core_msg_ops_aten__/out/RegisterCompositeExplicitAutogradNonFunctional.cpp":99, please report a bug to PyTorch.
*current_device == options.device() INTERNAL ASSERT FAILED at "buck-out/v2/gen/fbsource/c969ccb4ab4ae099/fbobjc/Libraries/FBPyTorchCore/__torch_core_msg_ops_aten__/out/RegisterCompositeExplicitAutogradNonFunctional.cpp":118, please report a bug to PyTorch.
*current_device == options.device() INTERNAL ASSERT FAILED at "buck-out/v2/gen/fbsource/c969ccb4ab4ae099/fbobjc/Libraries/FBPyTorchCore/__torch_core_msg_ops_aten__/out/RegisterCompositeExplicitAutogradNonFunctional.cpp":151, please report a bug to PyTorch.
```
After change:
```
$ buck2 build --show-output //fbobjc/Libraries/FBPyTorchCore:torch_core_msg_opsApple
$ strings buck-out/v2/gen/fbsource/c969ccb4ab4ae099/fbobjc/Libraries/FBPyTorchCore/__torch_core_msg_opsApple__/libtorch_core_msg_opsApple.pic.a | grep FAIL | head -n 10
isBool() INTERNAL ASSERT FAILED at "ivalue.h":666, please report a bug to PyTorch.
new_refcount != 1 INTERNAL ASSERT FAILED at "intrusive_ptr.h":268, please report a bug to PyTorch.
false INTERNAL ASSERT FAILED at "ivalue.h":1239, please report a bug to PyTorch.
owning_ptr == NullType::singleton() || owning_ptr->refcount_.load() == 0 || owning_ptr->weakcount_.load() INTERNAL ASSERT FAILED at "intrusive_ptr.h":474, please report a bug to PyTorch.
target_->refcount_ == 0 && target_->weakcount_ == 0 INTERNAL ASSERT FAILED at "intrusive_ptr.h":315, please report a bug to PyTorch.
refcount_.load() == 0 || refcount_.load() >= 2147483647 INTERNAL ASSERT FAILED at "intrusive_ptr.h":125, please report a bug to PyTorch.
weakcount_.load() == 1 || weakcount_.load() == 0 || weakcount_.load() == 2147483647 - 1 || weakcount_.load() == 2147483647 INTERNAL ASSERT FAILED at "intrusive_ptr.h":131, please report a bug to PyTorch.
isDouble() INTERNAL ASSERT FAILED at "ivalue.h":542, please report a bug to PyTorch.
isTensorList() INTERNAL ASSERT FAILED at "ivalue_inl.h":2019, please report a bug to PyTorch.
isTensorList() INTERNAL ASSERT FAILED at "ivalue_inl.h":2023, please report a bug to PyTorch.

$ strings buck-out/v2/gen/fbsource/c969ccb4ab4ae099/fbobjc/Libraries/FBPyTorchCore/__torch_core_msg_opsApple__/libtorch_core_msg_opsApple.pic.a | grep FAIL | grep buck
```

Reviewed By: milend

Differential Revision: D47637795
Dynamo test pipeline failed on MaxPool2d test when changed to use f-string
Labels: module: ci, module: tests, triaged, module: dynamo

### 🐛 Describe the bug

Tests

[linux-bionic-py3.11-clang9 / test (dynamo, 1, 2, linux.2xlarge)](https://hud.pytorch.org/pr/pytorch/pytorch/105436#15182560189) ([gh](https://github.com/pytorch/pytorch/actions/runs/5604098622/jobs/10251876391))
[linux-bionic-py3.8-clang9 / test (dynamo, 1, 2, linux.2xlarge)](https://hud.pytorch.org/pr/pytorch/pytorch/105436#15182563911) ([gh](https://github.com/pytorch/pytorch/actions/runs/5604098622/jobs/10251879196))

fails when the string format call in `cls_name = 'MaxPool{}d'.format(num_dim)` is changed to f-string in `_test_maxpool_indices`. For some reason the name becomes `MaxPool1d` even when it should be `MaxPool2d`

I could not reproduce in either a linux environment or on an m1 mac. It only happens in CI.

https://github.com/pytorch/pytorch/pull/105436#issuecomment-1642786366

### Versions

```
PyTorch version: 2.1.0a0+git2d16f88
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.4.1 (arm64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: version 3.26.4
Libc version: N/A

Python version: 3.11.4 (main, Jun 20 2023, 17:23:00) [Clang 14.0.3 (clang-1403.0.22.14.1)] (64-bit runtime)
Python platform: macOS-13.4.1-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING 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:
Apple M1 Max

Versions of relevant libraries:
[pip3] flake8==6.0.0
[pip3] numpy==1.25.1
[pip3] pytorch==2.1.0a0+git2d16f88
[conda] Could not collect
```

cc @seemethere @malfet @pytorch/pytorch-dev-infra @mruberry @voznesenskym @penguinwu @anijain2305 @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @ipiszy @chenyang78 @aakhundov
Slightly improve AOTAutograd logging with ViewAndMutationMeta
Labels: Merged, Reverted, Stale, ciflow/trunk, ciflow/inductor, release notes: AO frontend, module: export

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #105702

Signed-off-by: Edward Z. Yang <ezyang@meta.com>
Don't use weak ref finalization for freeing resources when code objects die
Labels: triaged, oncall: pt2, module: dynamo

### 🐛 Describe the bug

Today, we use weakref on codegen'ed code objects to trigger deallocation of compiled functions when the cache is cleared (using CleanupManager). I find this to be a bit fragile in the presence of pytest/exceptions, because if an exception is raised inside our compilation stack, the exception will hold references to the frame, and thus the code object, and then if this exception is kept around, this prevents the compiled function from ever being deallocated. It seems better to treat the cache entries as owning their respective code objects etc, and directly trigger deallocation when they are cleared.

### Versions

main

cc @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov
Support symmetry in einsum
Labels: triaged, module: linear algebra

### 🚀 The feature, motivation and pitch

PyTorch is an incredibly powerful tool, particularly due to its capability to parallelize and leverage GPU resources for operations like `einsum`. One area where I'd love to see development is in the support for tensor symmetries. 

For example, let's consider a tensor contraction like `A[a,b,c,d,e] * B[d,e,f,g,h] = C[a,b,c,f,g,h]`. If we have symmetries such as `A[a,b,c,d,e] = A[b,a,c,d,e]` or `B[d,e,f,g,h]=B[d,e,f,h,g]`, it would be great if we could use these symmetries to boost computational efficiency significantly. 

### Alternatives


The `ctf` package does offer some functionality along these lines, but it seems limited to symmetric/antisymmetric operations for adjacent indices. Here's the GitHub repository for `ctf`:

https://github.com/cyclops-community/ctf

For more specifics, the documentation can be found here:

https://solomon2.web.engr.illinois.edu/ctf_python/ctf.html#module-ctf.core

That said, from my observations, the performance isn't quite as efficient as one might hope. You can check out the issue I'm referring to here:

https://github.com/cyclops-community/ctf/issues/136


### Additional context

_No response_

cc @jianyuh @nikitaved @pearu @mruberry @walterddr @IvanYashchuk @xwang233 @Lezcano
[WIP] low mem max_pool2d_with_indices
Labels: Stale, module: inductor, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #105966
* __->__ #105687



cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @ngimel
Turn indexing with a scalar tensor into an copy into a view and avoid a D2H synchronization.
Labels: module: bc-breaking, triaged, module: numpy, module: advanced indexing, topic: bc breaking

### 🚀 The feature, motivation and pitch

Today, this triggers a cuda synchronization:
```
import torch

torch.set_default_device('cuda')

def f(x, y):
    return x[y]

inps = (torch.randn(5), torch.tensor(0))
torch.cuda.set_sync_debug_mode(2)
f(*inps)
```

The reason why is that when the tensor is a 0-dim value, instead of launching a gather kernel, we move the tensor to the hsot and do a slice instead (https://github.com/pytorch/pytorch/pull/105518/files#diff-2574bfb0ffa78d685fb7bd2ebc0c64b1a5f87dd55ec74ae67b41b31adc566020L466).

We could just fix this, but unfortunately, this does change the semantics. In particular, now, this operation would create a copy instead of a view, which could cause issues for downstream in-place operations.

I think these are bad semantics, for 3 reasons:

1. Cuda synchronizations are very bad in general. They're slow, prevent the use of many different features (streams, cudagraphs, don't play well with collectives, etc.), and should strongly be avoided. This, however, is a very implicit coercion we're doing. It's not obvious at all that if the tensor is 3-dim/2-dim/1-dim it doesn't do a sync, but if the tensor is 0-dim it does do a sync. In addition, it makes this much harder to trace and compile/not composite compliant in general.

2. Moreover, it's *not* consistent!!

Why should `x[torch.tensor(0)]` return a view but `x[torch.tensor([0])` return a copy? Why should the first one do a synchronization and the second one doesn't?

To drive this point home further, we also diverge from Numpy semantics here.

```
import numpy as np
x = np.ones(5)
y = np.array(1)
z = x[y]
z += 1
print(x)
>>> array([1., 1., 1., 1., 1.])
```

3. It's actually *slower* than just doing the index operator on GPUs! Benchmarking `x[torch.tensor(0)]` vs. `x[torch.tensor([0])`, we see that the first takes `35 us` per iteration while the second one takes `8 us`.

PS: I've also done a brief survey of use cases with this pattern I could find (https://github.com/search?q=%2F%28%5Cw%2B%29%5C%5Btorch.tensor%5C%28%2F+language%3APython&type=code), and I couldn't find many use cases of this code path at all.


cc: @ezyang @zou3519 @ngimel 

cc @ezyang @gchanan @mruberry @rgommers
Add z3-solver as dependency to dynamo tests
Labels: fb-exported, Stale, topic: not user facing, module: dynamo

Test Plan: sandcastle

Reviewed By: malfet, huydhn

Differential Revision: D47438456



cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng @anijain2305 @ipiszy
[MPS] Add mps support for max unpool2d 
Labels: triaged, open source, release notes: mps, ciflow/mps

Fixes one of the missing ops listed in #77764

Adds support for max_unpool2d forward & backward on the mps backend.

Since I don't think this op is natively supported in MPS, I've added an MSL kernel that mirrors the max_unpool2d cuda kernel.
FSDP with gradient checkpointing lead to redundant allgathers during backward
Labels: triaged, module: fsdp

### 🐛 Describe the bug

While training huggingface Llama 13B with FSDP with full shard and gradient checkpointing enabled on a single node, I observed that the backward pass has two allgathers per layer. Compared to non checkpointed training, this additional allgather also affects the overlap between the reduce scatter and gradient computation.  I think ideally only one allgather is needed to gather the weights per FSDP module. 

Trace: (browm is reduce scatter, blue is allgather)
![allgather](https://github.com/pytorch/pytorch/assets/21034047/96d6527f-94cb-4c94-9256-42127eb663c0)


```python3
        model = FSDP(model,
                    sharding_strategy=ShardingStrategy.FULL_SHARD,
                    sync_module_states=True,
                    mixed_precision=mixed_precision,
                    auto_wrap_policy=functools.partial(transformer_auto_wrap_policy, transformer_layer_cls={LlamaDecoderLayer}),
                    limit_all_gathers=True,
                    device_id=dev,
                    param_init_fn=param_init_fn,
                    )
```



### Versions

I'm using pytorch nightly: `2.1.0.dev20230709+cu121`

cc @zhaojuanmao @mrshenli @rohan-varma @awgu
torch.nn.TransformerDecoderLayer lacks parameter validation check
Labels: oncall: transformer/mha

### 🐛 Describe the bug

### Description:
torch.nn.TransformerDecoderLayer lacks parameter validation check, and when invalid values are given, they are used in subsequent computations, leading to errors such as division by zero.

### Examples:
input:
```
p = torch.nn.TransformerDecoderLayer(d_model=10, nhead=0)
```
error_message:
```
Traceback (most recent call last):
  File "D:\PythonProjects\venv\lib\site-packages\IPython\core\interactiveshell.py", line 3508, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "<ipython-input-32-6f43c243aac4>", line 1, in <module>
    p = torch.nn.TransformerDecoderLayer(d_model=10, nhead=0)
  File "D:\PythonProjects\venv\lib\site-packages\torch\nn\modules\transformer.py", line 653, in __init__
    self.self_attn = MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=batch_first,
  File "D:\PythonProjects\venv\lib\site-packages\torch\nn\modules\activation.py", line 968, in __init__
    self.head_dim = embed_dim // num_heads
ZeroDivisionError: integer division or modulo by zero

```


### Versions

PS D:\PythonProjects\venv\Lib\site-packages\torch\utils> python collect_env.py
Collecting environment information...
PyTorch version: 2.0.1+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 11 家庭中文版
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.10.11 (tags/v3.10.11:7d4cc5a, Apr  5 2023, 00:38:17) [MSC v.1929 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.22621-SP0
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4070 Laptop GPU
Nvidia driver version: 532.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=2200
DeviceID=CPU0
Family=207
L2CacheSize=16384
L2CacheSpeed=
Manufacturer=GenuineIntel
MaxClockSpeed=2200
Name=13th Gen Intel(R) Core(TM) i9-13900HX
ProcessorType=3
Revision=

Versions of relevant libraries:
[pip3] numpy==1.24.3
[pip3] torch==2.0.1+cu117
[pip3] torchaudio==2.0.2+cu117
[pip3] torchvision==0.15.2+cu117
[conda] Could not collect

cc @jbschlosser @bhosmer @cpuhrsch @erichan1 @drisspg
[ONNX] fix `test_fx_op_consistency.py` test failure when running on torch built with cuda
Labels: module: onnx, triaged

Step to repro

`pytest test/onnx/test_fx_op_consistency.py -k test_output_match_full_like_cpu_float32`

raises

`AssertionError: The values for attribute 'device' do not match: cuda:0 != cpu.`

cc @justinchuby
Enable Mypy checking for scheduler.py
Labels: topic: not user facing, module: inductor, ciflow/inductor

ATT, add type annotations and type assertions to pass Mypy checks.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ngimel @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov
Add sdpa op prototype
Labels: Stale, release notes: fx, module: inductor, module: dynamo, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #105596
* #105518



cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @ngimel @anijain2305
Change default autograd fallback mode to "Warn"
Labels: Stale, ciflow/trunk, topic: not user facing

Stack from [ghstack](https://github.com/ezyang/ghstack):
* __->__ #105592
* #105845
* #105660

We changed the default autograd fallback mode to "Nothing" in #105505
because it was breaking internal tests.

We have fixed the problems in #105587, so this PR changes the fallback
back to "Warn".

Test Plan:
- internal tests
[Inductor] Add support for NEON ISA in the Inductor C++ backend
Labels: module: cpu, triaged, open source, module: inductor, ciflow/inductor

Fixes #104729

As suggested in the [blog](https://dev-discuss.pytorch.org/t/torchinductor-update-5-cpu-backend-backend-performance-update-and-deep-dive-on-key-optimizations/1117#:~:text=It%20can%20be,sub%2Dclasses.), I subclassed the `VecISA` class and implemented a NEON version of the `vec_reduce_all()` function, to go along with the existing AVX2 and AVX512 versions. Any operation that calls `vec_reduce_all()` will also take the NEON path and benefit from its vectorization.

The `vec_reduce_all()` is invoked by Softmax and other operations like norms. Using the fast path results in 30% time savings for Softmax as compared to the previously taken slow path.


  | Slow path | Fast path (NEON intrinsics)
-- | -- | --
Softmax (100 passes, 1024 dimension) | 623.706ms | 452.011ms


cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @voznesenskym @penguinwu @EikanWang @Guobing-Chen @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler @ngimel @malfet
torch.sparse.sampled_addmm doesn't compute gradients for 3D tensors
Labels: module: sparse, triaged

### 🐛 Describe the bug

Hi, I hope you don't mind me raising so many issues 😄

`torch.sparse.sampled_addmm` works fine in the forward phase for both 2D and 3D tensors. However, in the backward pass it fails for 3D tensors and throws a cryptic error message.

```
import torch

B, N, D, p = 4, 100, 30, 0.01
M1 = torch.randn(B, N, D).cuda().requires_grad_(True)
M2 = torch.randn(B, D, N).cuda()
mask = torch.bernoulli(p * torch.ones((N, N))).to_sparse_csr().cuda()

out = torch.sparse.sampled_addmm(mask, M1[0], M2[0])
out.to_dense()[0, 0].backward()  # works

out = torch.sparse.sampled_addmm(mask, M1, M2)
out.to_dense()[0, 0, 0].backward()  # doesn't work
```

I receive the following error:
```
RuntimeError: crow_indices is supposed to be a vector, but got 2 dimensional tensor.
```

### Versions

PyTorch version: 2.0.1+cu117

cc @alexsamardzic @nikitaved @pearu @cpuhrsch @amjames @bhosmer
UserWarning: There is a performance drop because we have not yet implemented the batching rule for aten::index_add_. Please file us an issue on GitHub so that we can prioritize its implementation.
Labels: triaged, actionable, module: vmap, module: functorch

### 🚀 The feature, motivation and pitch

I'm working on a system that requires inplace updates to a state inside of a vmap (think the performer model with multiple heads and sparse updates/matmul). This used to not work, but as of 2.0.1 it does but I get this warning:

```
UserWarning: There is a performance drop because we have not yet implemented the batching rule for aten::index_add_. Please file us an issue on GitHub so that we can prioritize its implementation.
```

This occurred on an M2 macbook, I'll test later whether this also occurs on a cuda device. It would be nice to have indexed inplace updates fully supported in a future update.

### Alternatives

_No response_

### Additional context

_No response_

cc @zou3519 @Chillee @samdow @kshitij12345 @janeyx99
[POC] DynamicTensor
Labels: release notes: fx

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #105535

DynamicTensor is a regular CPU/CUDA tensor, but where some dimensions are treated symbolically. If you have an (s0, 2) dynamic tensor, it is intended to vary dynamically over s0. This size propagates over operations, and always behaves uniformly, even if s0 = 0, 1 (so, for example, you cannot broadcast (s0, 2) with (5, 2), even when s0 == 1).

More discussion at https://docs.google.com/document/d/1Dge173HVbXnTysnvp8716mi_0BlEpJZbmhXbI_-WdqI/edit#heading=h.ey3z8v3k3z06

The implementation strategy is to subclass FakeTensor into DynamicTensor, but then embed a real backing tensor which also gets operated on. The implementation does some naughty things, putting the POC out there to get some comments.

TODO: When a DynamicTensor has no more symbolic dimensions, decay it into a regular tensor.

Signed-off-by: Edward Z. Yang <ezyang@meta.com>
tts_angular: fail_to_run, torch._dynamo.exc.Unsupported: call_method NNModuleVariable() flatten_parameters [] {}
Labels: triaged, oncall: pt2

Repro:
```
python benchmarks/dynamo/torchbench.py --accuracy --inference --bfloat16 --export-aot-inductor --disable-cudagraphs --device cuda --only tts_angular
```

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
Efficient BMM for sparse-dense tensors
Labels: module: sparse, triaged, topic: new features

### 🚀 The feature, motivation and pitch

Hi,
I want to perform a sparse-dense BMM and compute gradients for the sparse matrix. Is there an operation in torch which does it efficiently? According to [this table](https://pytorch.org/docs/stable/sparse.html#supported-operations), `torch.bmm` computes only the dense gradients. Also, it's limited to the COO layout.

### Alternatives

_No response_

### Additional context

_No response_

cc @alexsamardzic @nikitaved @pearu @cpuhrsch @amjames @bhosmer
DISABLED test_conv_with_as_strided_dynamic_shapes_cuda (__main__.DynamicShapesCudaTests)
Labels: module: rocm, triaged, module: flaky-tests, skipped

Platforms: rocm

This test was disabled because it is failing in CI. See [recent examples](https://hud.pytorch.org/failure/test_conv_with_as_strided_dynamic_shapes_cuda) and the most recent trunk [workflow logs](https://github.com/pytorch/pytorch/runs/15120975791).

Over the past 72 hours, it has flakily failed in 6 workflow(s).

**Debugging instructions (after clicking on the recent samples link):**
To find relevant log snippets:
1. Click on the workflow logs linked above
2. Grep for `test_conv_with_as_strided_dynamic_shapes_cuda`


Test file path: `inductor/test_torchinductor_dynamic_shapes.py`

ResponseTimeoutError: Response timeout for 5000ms, GET https://raw.githubusercontent.com/pytorch/pytorch/main/test/inductor/test_torchinductor_dynamic_shapes.py -1 (connected: true, keepalive socket: false, socketHandledRequests: 1, socketHandledResponses: 0)
headers: {}

cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang
torch.onnx.export error
Labels: module: onnx, triaged

### 🐛 Describe the bug

while running the code

```python
input_s = torch.zeros((1, 500), dtype=torch.long).to('cuda')
input_names = ["x"] # specify the name of the input tensor
output_names = ["out"] # specify the name of the output tensor
dynamic_axes = {"x":{0: "batch_size"}, "out":{0: "batch_size"}}
torch.onnx.export(model, input_s, 'model_test.onnx', input_names=input_names, output_names=output_names)
```

The following problem occurs, what is the situation, is this a common warning, or an error? do i need to edit again

### Versions

 UserWarning: The exported ONNX model failed ONNX shape inference.The model will not be executable by the ONNX Runtime.If this is unintended and you believe there is a bug,please report an issue at https://github.com/pytorch/pytorch/issues.Error reported by strict ONNX shape inference: [ShapeInferenceError] Shape inference error(s): (op_type:MaxPool, node name: /MaxPool): [ShapeInferenceError] Attribute strides has incorrect size
(op_type:MaxPool, node name: /MaxPool_1): [ShapeInferenceError] Attribute strides has incorrect size
(op_type:MaxPool, node name: /MaxPool_2): [ShapeInferenceError] Attribute strides has incorrect size
 (Triggered internally at ../torch/csrc/jit/serialization/export.cpp:1407.)
  _C._check_onnx_proto(proto)
============= Diagnostic Run torch.onnx.export version 2.0.1+cu117 =============
verbose: False, log level: Level.ERROR
======================= 0 NONE 0 NOTE 0 WARNING 0 ERROR ========================
[ONNX] Exporting the operator 'aten::exponential' to opset version 13 is not supported
Labels: module: onnx, triaged

### 🐛 Describe the bug

'aten::exponential' seems not supported by onnx. One way I know is to create a custom function and register it, but it doesn't work.

### Versions

--2023-07-19 14:32:17--  https://raw.githubusercontent.com/pytorch/pytorch/main/torch/utils/collect_env.py
Resolving raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.111.133, 185.199.109.133, 185.199.108.133, ...
Connecting to raw.githubusercontent.com (raw.githubusercontent.com)|185.199.111.133|:443... connected.
HTTP request sent, awaiting response... 200 OK
Length: 21653 (21K) [text/plain]
Saving to: ‘collect_env.py’

collect_env.py                      100%[===================================================================>]  21.15K  --.-KB/s    in 0.001s  

2023-07-19 14:32:18 (22.8 MB/s) - ‘collect_env.py’ saved [21653/21653]
Avoid synchronization when using scalar tensor as index
Labels: Stale, module: dynamo

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #105596
* __->__ #105518



cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng @anijain2305 @ipiszy
[ONNX] FX produce valid node names in models
Labels: module: onnx, triaged

Currently we see names like `LayerNorm(L__self___embed_layer_norm)_14` for functions and `weight.output.1` for tensors. We need to make them valid C variable names to follow the ONNX standard: https://onnx.ai/onnx/repo-docs/IR.html#names-within-a-graph

> All names MUST adhere to [C90 identifier syntax rules](https://en.cppreference.com/w/c/language/identifier).
> Names of nodes, inputs, outputs, initializers, and attributes are organized into several namespaces. Within a namespace, each name MUST be unique for each given graph.

cc @BowenBao
[FSDP] Revisit mixed-precision casting logic
Labels: triaged, module: fsdp

**Overview**
We may want to revisit FSDP's behavior for casting forward input tensors and buffers. Our mixed precision API has options `param_dtype` ,`reduce_dtype`, `buffer_dtype`, `cast_forward_inputs `, and `cast_root_forward_inputs`:
https://github.com/pytorch/pytorch/blob/abc1cadddba00b7412240b56314d9592d7ad29c7/torch/distributed/fsdp/api.py#L211-L216
In addition, we allow using full precision in eval, configured by an environment variable:
https://github.com/pytorch/pytorch/blob/abc1cadddba00b7412240b56314d9592d7ad29c7/torch/distributed/fsdp/flat_param.py#L482-L484

After https://github.com/pytorch/pytorch/pull/104408, every FSDP state has either 0 or 1 handles (see https://github.com/pytorch/pytorch/pull/104488). Note that `any(...)` in Python returns `True` if the `...` is empty. This means that currently:
- For the case where `cast_root_forward_inputs=True` and the root does not have a handle, the `(module.training or not state._use_full_prec_in_eval)` is there to represent the `(not self._fully_sharded_module.training and self._use_full_prec_in_eval)` check, which cannot happen without a handle.
    - [ ] We should be able to simplify this to `should_cast_forward_inputs = (module.training or not state._use_full_prec_in_eval) and state.mixed_precision.cast_root_forward_inputs` (removing the check on `not handle._force_full_precision`). However, this creates some asymmetry with the condition for `cast_forward_inputs=True` (`_pre_forward()`).
- For the case where `cast_forward_inputs=True`, the `len(state._handles) > 0` check avoids the need for something similar. If the FSDP instance does not manage any parameters (and hence has no handle), it never casts its forward inputs.
    - [ ] Unlike `cast_root_forward_inputs=True`, if the module does not manage any parameters, then this will not cast. We should document this behavior explicitly.
- [ ] If the root does not have a handle, then it always casts buffers to full precision. This may not be the desired behavior, and we may want to add a `state._force_full_precision` clause as part of this check.

**Code Pointers**
`cast_root_forward_inputs=True` (`_root_pre_forward()`):
https://github.com/pytorch/pytorch/blob/91ab32e4b1f9e601cd42b7e9887b93a444c99dfb/torch/distributed/fsdp/_runtime_utils.py#L655-L658

`cast_forward_inputs=True` (`_pre_forwarrd()`):
https://github.com/pytorch/pytorch/blob/91ab32e4b1f9e601cd42b7e9887b93a444c99dfb/torch/distributed/fsdp/_runtime_utils.py#L461-L463

(`_root_pre_forward()`):
https://github.com/pytorch/pytorch/blob/91ab32e4b1f9e601cd42b7e9887b93a444c99dfb/torch/distributed/fsdp/_runtime_utils.py#L587-L589

https://github.com/pytorch/pytorch/blob/abc1cadddba00b7412240b56314d9592d7ad29c7/torch/distributed/fsdp/flat_param.py#L2469-L2478



---



cc @zhaojuanmao @mrshenli @rohan-varma
torch.save throws an error when the path uses mixed separators on Windows
Labels: module: windows, triaged

### 🐛 Describe the bug

Using some combinations of \ and / as the path separator throws an exception. Windows prefers \ as the path separator, but also accepts /. torch.save throws an exception for paths that should be valid windows paths.

Reproducible using the following example
``` python
import os
import torch

data = [{'1': 1}]

torch.save(data, "H:/a\\a.ckpt")
```

replacing the path can lead to a few different results:

`torch.save(data, "H:\\a\\a.ckpt")` -> works
`torch.save(data, "H:/a/a.ckpt")` -> works
`torch.save(data, "H:/a\\a.ckpt")` -> Throws "Parent directory H: does not exist."
`torch.save(data, "H:\\a/a.ckpt")` -> works

But now it gets really strange:
`torch.save(data, "H:/a.ckpt")` -> Throws "Parent directory H: does not exist.". The file is created, but remains empty
`torch.save(data, "H:\\a.ckpt")` -> Throws "Parent directory H: does not exist.". The file is created, but remains empty

Stack trace
```
...
  File "H:\stable-diffusion\one-trainer\scripts\debug.py", line 17, in main
    torch.save(data, "H:\\a.ckpt")
  File "H:\stable-diffusion\one-trainer\venv\lib\site-packages\torch\serialization.py", line 440, in save
    with _open_zipfile_writer(f) as opened_zipfile:
  File "H:\stable-diffusion\one-trainer\venv\lib\site-packages\torch\serialization.py", line 315, in _open_zipfile_writer
    return container(name_or_buffer)
  File "H:\stable-diffusion\one-trainer\venv\lib\site-packages\torch\serialization.py", line 288, in __init__
    super().__init__(torch._C.PyTorchFileWriter(str(name)))
RuntimeError: Parent directory H: does not exist.
```

### Versions

PyTorch version: 2.0.1+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 10 Home
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.10.8 (tags/v3.10.8:aaaf517, Oct 11 2022, 16:50:30) [MSC v.1933 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.19045-SP0
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA RTX A5000
Nvidia driver version: 535.98
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture=9
CurrentClockSpeed=3600
DeviceID=CPU0
Family=205
L2CacheSize=1536
L2CacheSpeed=
Manufacturer=GenuineIntel
MaxClockSpeed=3600
Name=Intel(R) Core(TM) i5-8600K CPU @ 3.60GHz
ProcessorType=3
Revision=

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.23.5
[pip3] pytorch-lightning==2.0.3
[pip3] torch==2.0.1+cu118
[pip3] torchmetrics==1.0.1
[pip3] torchvision==0.15.2+cu118
[conda] Could not collect

cc @peterjc123 @mszhanyi @skyline75489 @nbcsm @vladimir-aubrecht @iremyux @Blackhex @cristianPanaite
[benchmark] Rename the count field FunctionCount
Labels: oncall: profiler

The `count` field is in `FunctionCount` defined in `torch/utils/benchmark/utils/valgrind_wrapper/timer_interface.py` is redefining the built in field of tuple. Consider renaming.

>              Hmm, do you mind renaming it to something else? How about `counter`? or `invocation_count` (though strictly speaking out of scope of this PR)

_Originally posted by @malfet in https://github.com/pytorch/pytorch/pull/105424#discussion_r1266832538_
            

cc @robieta @chaekit @aaronenyeshi @nbcsm @guotuofeng @guyang3532 @gaoteng-git @tiffzhaofb @dzhulgakov @davidberard98
[proposal] Bit ops: e.g. setbit/getbit/togglebit/byteswap + introduce well-standardized unsigned dtypes (uint16, uint32, uint64)
Labels: feature, triaged, needs research, module: python frontend

### 🐛 Describe the bug

Some functions are reversible/invertible (e.g. https://pytorch.org/docs/stable/generated/torch.nn.SiLU.html) if we record one more bit of information: whether the input is larger than some cut-off value.

Of course one can save a bitmask for backward. Another funny/hacky way to do the same would be record this information in the LSB mantissa bit. Currently there exists a single such op: torch.copysign_ (which is also not very flexible as it can't accept BoolTensors now)

According to https://stackoverflow.com/a/47990/445810, forcing a certain bit can be done by (but with a lot of allocations):

```python
import torch
torch.manual_seed(1)

a = torch.rand(4, 3, dtype = torch.float32)
sign = a.ge(0.5).to(torch.int32)
a_, a_repr = a.clone(), a.view(torch.int32)
which_bit = 0
a_repr ^= (-sign ^ a_repr) & (1 << which_bit)

# for lsb : a_repr ^= ((-sign ^ a_repr) & 1)

print(a, a_, a == a_)
```

It would be good to maybe support more of these bitops as native ops? Ideally, inductor would generate efficient impls of these ops and fuse them with the rest of the computation, but clear API would be great.

Also, it might be good to support torch.uint32 for guaranteed bit ops? I think for int32 some bitops result are not well defined in C++, so at least for bit manipulations being able to clearly express uint32 might be useful. Existing issue: https://github.com/pytorch/pytorch/issues/58734

Related: https://github.com/pytorch/pytorch/issues/32867 on supporting BitTensor natively (and especially as outcome for boolean ops like torch.ge)

### Versions

N/A

cc @albanD
[ONNX] Support Fake Tensor Mode on new Dynamo based ONNX exporter
Labels: module: onnx, triaged, enhancement, release notes: onnx

### 🐛 Describe the bug

This task is an umbrella for all tasks related to exporting a model to ONNX using the new PyTorch Dynamo API with Fake Tensor support

Idea for alternative API design: https://github.com/pytorch/pytorch/issues/104144

### Versions

PyTorch main branch
```[tasklist]
### Tasks
- [ ] Revisit serialization of models with Fake Tensor support
- [ ] https://github.com/pytorch/pytorch/issues/105467
- [ ] Support large model export without special hardware (e.e.g A100)
- [x] Create public API for ONNX export with Fake Tensor support
- [ ] Address mix of fake and real tensor as reported by https://github.com/pytorch/pytorch/issues/105077
- [ ] https://github.com/pytorch/pytorch/issues/105490
- [ ] https://github.com/pytorch/pytorch/issues/105751
- [ ] Support Fake Mode with dynamic shapes natively
- [ ] https://github.com/pytorch/pytorch/issues/106412
```
Specify version
Labels: module: docs, triaged

### 📚 The doc issue

In many cases, I know that some functionality was only introduced in a recent PyTorch version. E.g. the use of `torch.device` as a context manager, or the function `torch.set_default_device`. I would expect that the documentation mentions in what version this was introduced, but this information is lacking.

### Suggest a potential alternative/fix

For every function, specify since what version it is available.

Also, similarly, if new arguments are added to a function, you could specify since what version the argument is available.

cc @svekars @carljparker
Adding documentation an diagram on code base
Labels: module: cpu, triaged, module: mkldnn, open source, module: amp (automated mixed precision), NNC, ciflow/trunk, release notes: quantization, topic: not user facing, ciflow/mps, module: inductor, module: dynamo, ciflow/inductor, module: export

Fixes #104962


cc @jgong5 @mingfeima @XiaobingSuper @sanchitintel @ashokei @jingxu10 @gujinghui @PenghuiCheng @jianyuh @min-jean-cho @yanbing-j @Guobing-Chen @Xia-Weiwen @mcarilli @ptrblck @leslie-fang-intel @EikanWang @voznesenskym @penguinwu @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @ngimel @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @anijain2305
Top level Glossary for users (not contributers)
Labels: module: docs, triaged

### 📚 The doc issue

Related to [https://github.com/pytorch/pytorch/issues/104623](https://github.com/pytorch/pytorch/issues/104623) i propose a new top level glassy that defines concepts that a PyTorch developer is interested in - the rest of the documentation could then link to the standard definitions of words. It should start with definitions of :
-PyTorch
- the associated various repos
- the various packages  
For example as a newby. I would have liked clear definitions of :
- frontend/ backend
- device
- tables of capabilities for different backends 
- what is pickleable an what is not.

I am happy to start this off with a starter set for people to review , then add to.

I know we have GLOSSARY.md - but this seems to be contributor focused, I am not sure if users would need to know about the dispatcher, but may be I am wrong. 



### Suggest a potential alternative/fix

_No response_

cc @svekars @carljparker
Will nn.unfold support non-4D-tensor input in future version?
Labels: module: nn, triaged, enhancement, actionable

### 🚀 The feature, motivation and pitch

For pretty many versions there is the warning "Currently, only 4-D input tensors (batched image-like tensors) are supported". Such program can be frequently used but nonnative implementation is slow. Perhaps you may consider to add the support for "arbitrary spatial dimensions" as the docs said in the ongoing version? Thanks!

### Alternatives

_No response_

### Additional context

_No response_

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki
DISABLED test_cross_entropy_large_tensor_reduction_none_cuda (__main__.TestNNDeviceTypeCUDA)
Labels: module: nn, module: rocm, triaged, module: flaky-tests, skipped

Platforms: rocm

This test was disabled because it is failing in CI. See [recent examples](https://hud.pytorch.org/failure/test_cross_entropy_large_tensor_reduction_none_cuda) and the most recent trunk [workflow logs](https://github.com/pytorch/pytorch/runs/15112704347).

Over the past 72 hours, it has flakily failed in 6 workflow(s).

**Debugging instructions (after clicking on the recent samples link):**
To find relevant log snippets:
1. Click on the workflow logs linked above
2. Grep for `test_cross_entropy_large_tensor_reduction_none_cuda`


Test file path: `test_nn.py`

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang
Silent Error of torch.fx.symbolic_trace when forward hooks are registered
Labels: triaged, module: fx

### 🐛 Describe the bug

When there are some forward hooks in the module, the `torch.fx.symbolic_trace` silently ignores the computation.

```python
from torch.nn.utils import spectral_norm
from torch import fx
from torch import nn
m = spectral_norm(nn.Linear(20, 40))
m.weight.data.zero_()
m.weight.data += 500

import torch
input = torch.ones(32, 20)
fx_model = torch.fx.symbolic_trace(m)

output1 = m(input)
output2 = fx_model(input)
print((output1 - output2).abs().max().item()) # 9999.2939453125
```

This is because the `torch.fx.symbolic_trace` ignores the forward hooks, as discussed in https://github.com/pytorch/vision/issues/5193 .

However, we should at least report some errors/warning in such a case.

The fix is simple: just raise errors/warnings when hooks are deteced.

### Versions

It affects `torch.fx` in all versions.

cc @ezyang @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv
Need support and testing for Adam optimizer for MPS
Labels: high priority, module: optimizer, triaged, enhancement, module: mps

### 🚀 The feature, motivation and pitch

Environment of Mac M2
```
Python3.10
torch              2.1.0.dev20230717
torchaudio         2.1.0.dev20230717
torchvision        0.15.2a0
```

I want to use the Adam optimizer to train my model. And got an error:
```
NotImplementedError: The operator 'aten::lerp.Scalar_out' is not currently implemented for the MPS device. If you want this op to be added in priority during the prototype phase of this feature, please comment on https://github.com/pytorch/pytorch/issues/77764. As a temporary fix, you can set the environment variable `PYTORCH_ENABLE_MPS_FALLBACK=1` to use the CPU as a fallback for this op. WARNING: this will be slower than running natively on MPS
```
And when I set the `PYTORCH_ENABLE_MPS_FALLBACK=1`, the training speed is quite slower.

I'm testing the [tiny VIT model](https://github.com/UdbhavPrasad072300/Transformer-Implementations/blob/main/transformer_package/models/transformer.py#L406) with minist dataset.

The details is following:
M2 chip takes about **2.4 minutes** on CPU with Adam for one epoch.
M2 chip takes about **2.0 minutes** on GPU with Adam for one epoch( with `PYTORCH_ENABLE_MPS_FALLBACK=1`).
M2 chip takes about **30 seconds** on GPU with SGD for one epoch.


### Alternatives

_No response_

### Additional context

_No response_

cc @vincentqb @jbschlosser @albanD @janeyx99 @crcrpar @kulinseth @malfet @DenisVieriu97 @razarmehr @abhudev @ezyang @gchanan @zou3519