MPSNDArray Error: buffer is not large enough

[theRoodest]

🐛 Describe the bug

Installed latest PyTorch nightly build for MPS on M1, 1.13.0dev20220608, set environmental variable PYTORCH_ENABLE_MPS_FALLBACK=1 to resolve issue with aten::index.Tensor. Ran Stable-Baselines3 to train a PPO agent. Conducted one successful iteration then received an error:

import torch

from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env

# Parallel environments
env = make_vec_env("CartPole-v1", n_envs=4)
device = torch.device('mps')

model = PPO("MlpPolicy", env, device=device, verbose=1)
model.learn(total_timesteps=25000)
model.save("ppo_cartpole")

/AppleInternal/Library/BuildRoots/8d3bda53-8d9c-11ec-abd7-fa6a1964e34e/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShaders/MPSCore/Types/MPSNDArray.mm:782: failed assertion '[MPSNDArray, initWithBuffer:descriptor:] Error: buffer is not large enough. Must be 1024 bytes
'

Process finished with exit code 134 (interrupted by signal 6: SIGABRT)

Versions

Collecting environment information...
PyTorch version: 1.13.0.dev20220608
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.3.1 (arm64)
GCC version: Could not collect
Clang version: 13.1.6 (clang-1316.0.21.2.5)
CMake version: Could not collect
Libc version: N/A

Python version: 3.8.13 (default, Mar 28 2022, 06:13:39)  [Clang 12.0.0 ] (64-bit runtime)
Python platform: macOS-12.3.1-arm64-i386-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
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

Versions of relevant libraries:
[pip3] numpy==1.21.4
[pip3] pytorch-ignite==0.4.6
[pip3] torch==1.13.0.dev20220608
[pip3] torchvision==0.14.0a0+f9f721d
[conda] pytorch                   1.13.0.dev20220608         py3.8_0    pytorch-nightly
[conda] torch                     1.13.0.dev20220602          pypi_0    pypi
[conda] torchvision               0.14.0a0+f9f721d          pypi_0    pypi

cc @ezyang @gchanan @zou3519 @kulinseth @albanD

[malfet]

Likely a duplicate of #78916
And I can reproduce the crash easily (though TORCH_SHOW_CPP_STACKTRACES=1 does not work for it :( ), so here is the backtrace with lldb attached:

    frame #0: 0x000000019cd52d98 libsystem_kernel.dylib`__pthread_kill + 8
    frame #1: 0x000000019cd87ee0 libsystem_pthread.dylib`pthread_kill + 288
    frame #2: 0x000000019ccc2340 libsystem_c.dylib`abort + 168
    frame #3: 0x000000019ccc1754 libsystem_c.dylib`__assert_rtn + 272
  * frame #4: 0x00000001a57787a8 Metal`MTLReportFailure.cold.1 + 56
    frame #5: 0x00000001a57622bc Metal`MTLReportFailure + 480
    frame #6: 0x00000001a63b6984 MPSCore`___lldb_unnamed_symbol641$$MPSCore + 428
    frame #7: 0x00000002009bac90 MetalPerformanceShadersGraph`___lldb_unnamed_symbol2960$$MetalPerformanceShadersGraph + 536
    frame #8: 0x000000012bf040c0 libtorch_cpu.dylib`at::native::mps::_gatherViewTensor(at::Tensor const&, id<MTLBuffer>, at::native::mps::MPSCachedGraph*, at::Tensor&) + 176
    frame #9: 0x000000012bf0457c libtorch_cpu.dylib`at::native::mps::Placeholder::Placeholder(MPSGraphTensor*, at::Tensor const&, NSArray<NSNumber*>*) + 208
    frame #10: 0x000000012bf7d808 libtorch_cpu.dylib`at::native::structured_gather_out_mps::impl(at::Tensor const&, long long, at::Tensor const&, bool, at::Tensor const&) + 1468
    frame #11: 0x0000000129b6de60 libtorch_cpu.dylib`at::(anonymous namespace)::wrapper_gather(at::Tensor const&, long long, at::Tensor const&, bool) + 128
    frame #12: 0x000000012a98ff2c libtorch_cpu.dylib`c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor (c10::DispatchKeySet, at::Tensor const&, long long, at::Tensor const&, bool), &(torch::autograd::VariableType::(anonymous namespace)::gather(c10::DispatchKeySet, at::Tensor const&, long long, at::Tensor const&, bool))>, at::Tensor, c10::guts::typelist::typelist<c10::DispatchKeySet, at::Tensor const&, long long, at::Tensor const&, bool> >, at::Tensor (c10::DispatchKeySet, at::Tensor const&, long long, at::Tensor const&, bool)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, long long, at::Tensor const&, bool) + 1068
    frame #13: 0x0000000129222fd0 libtorch_cpu.dylib`at::_ops::gather::call(at::Tensor const&, long long, at::Tensor const&, bool) + 304
    frame #14: 0x00000001044875f0 libtorch_python.dylib`torch::autograd::THPVariable_gather(_object*, _object*, _object*) + 692

[malfet]

Here is the one line reproducer to the problem:

% python3 -c "import torch;x=(torch.rand(64, 1, device='mps')*1000).to(dtype=torch.int64); y=x.as_strided(size=(64,2), stride=(1, 0));z=y.as_strided(size=(64, 2), stride=(1, 0));z.to('cpu')"
/AppleInternal/Library/BuildRoots/b6051351-c030-11ec-96e9-3e7866fcf3a1/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShaders/MPSCore/Types/MPSNDArray.mm:782: failed assertion `[MPSNDArray, initWithBuffer:descriptor:] Error: buffer is not large enough. Must be 1024 bytes
'
zsh: abort      python3 -c 

[abhudev]

We aren't able to reproduce this issue, could you please update to latest nightly and try again?