LSTM Output Transposed w/MPS on 1.13 nightly build

[PHRABAL]

🐛 Describe the bug

The 1.13 nightly build, when sending an LSTM model to device="mps" reverses the expected order of batch and seq in the output.

Please see this discussion for code examples and further details:

https://discuss.pytorch.org/t/lstm-output-transposed/154820/2

Versions

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

OS: macOS 12.4 (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.9.13 | packaged by conda-forge | (main, May 27 2022, 17:01:00) [Clang 13.0.1 ] (64-bit runtime)
Python platform: macOS-12.4-arm64-arm-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.22.4
[pip3] torch==1.13.0.dev20220620
[conda] numpy 1.22.4 pypi_0 pypi
[conda] pytorch 1.13.0.dev20220620 py3.9_0 pytorch-nightly

cc @kulinseth @albanD

[gscalise]

Minimal case confirming this issue. Verified on PyTorch version 1.13.0.dev20220627, Python version '3.8.13 | packaged by conda-forge | (default, Mar 25 2022, 06:05:16) \n[Clang 12.0.1 ]'

import torch
from torch.nn import LSTM

assert torch.backends.mps.is_available()

tensor_shapes = {}

for device_type in ("mps", "cpu"):
    device = torch.device(device_type)
    lstm = LSTM(64, 128, 1, batch_first=True)
    lstm.to(device)
    input_tensor = torch.randn((2, 4, 64)).to(device)
    output, _ = lstm(input_tensor)
    tensor_shapes[device_type] = output.shape
    
print(tensor_shapes) # prints {'mps': torch.Size([4, 2, 128]), 'cpu': torch.Size([2, 4, 128])}
                     #                            ^^^^^                           ^^^^^

Interestingly, this doesn't happen when batch_first is False:

import torch
from torch.nn import LSTM

assert torch.backends.mps.is_available()

tensor_shapes = {}

for device_type in ("mps", "cpu"):
    device = torch.device(device_type)
    lstm = LSTM(64, 128, 1)
    lstm.to(device)
    input_tensor = torch.randn((2, 4, 64)).to(device)
    output, _ = lstm(input_tensor)
    tensor_shapes[device_type] = output.shape
    
print(tensor_shapes) # prints {'mps': torch.Size([2, 4, 128]), 'cpu': torch.Size([2, 4, 128])}

[gscalise]

Also, this doesn't happen with GRU:

import torch
from torch.nn import GRU

assert torch.backends.mps.is_available()

tensor_shapes = {}

for device_type in ("mps", "cpu"):
    device = torch.device(device_type)
    gru = GRU(64, 128, 1, batch_first=True)
    gru.to(device)
    input_tensor = torch.randn((2,4,64)).to(device)
    output, _ = gru(input_tensor)
    tensor_shapes[device_type] = output.shape
    
print(tensor_shapes) # prints {'mps': torch.Size([2, 4, 128]), 'cpu': torch.Size([2, 4, 128])}

[qqaatw]

Should be fixed by the linked PR.

[PHRABAL]

Hi,

The committed changes did not fix the issue.

While the model's output order is now correct for batch_first=True, when you send it thru backprop it results in an error.

This simple example will demonstrate the problem. It runs fine with batch_first=False, but when changed to batch_first=True, it results in the error shown below.

import torch
import torch.nn as nn


device = torch.device("mps")
print(device)

batch_size = 64
seq_len = 10
input_size = 128
hidden_size = 32

class Model(nn.Module):

    def __init__(self):
        super(Model, self).__init__()
    
        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size, num_layers=1, batch_first=True)
        self.fc = nn.Linear(hidden_size, 1)

    def forward(self, batch):
        x, _ = self.lstm(batch)
        x = self.fc(x)
        return x              

model = Model().to(device)
loss_func = nn.BCEWithLogitsLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=.0001)

x = torch.randn(batch_size, seq_len, input_size).to(device)
y = torch.randn(batch_size, seq_len, 1).to(device)

print(x.shape)
print(y.shape)

for epoch in range(1, 10):
    out = model(x)
    loss = loss_func(out, y)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print(loss.data.cpu())

print(out.shape)

The error is:

loc("total derivative last state"("(mpsFileLoc): /AppleInternal/Library/BuildRoots/b6051351-c030-11ec-96e9-3e7866fcf3a1/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShadersGraph/mpsgraph/MetalPerformanceShadersGraph/Core/Files/MPSGraphUtilities.mm":219:0)): error: input types 'tensor<1x10x32xf32>' and 'tensor<1x64x32xf32>' are not broadcast compatible LLVM ERROR: Failed to infer result type(s). Abort trap: 6

It looks like it is still not flipping the batch (64)and sequence dimensions (10) when training.

[qqaatw]

@PHRABAL Hi, the backward of LSTM on mps has issues before the fix:

pytorch/test/test_mps.py

Line 5712 in f3008be

@unittest.skipIf(True, "Backward of lstm returns wrong result")

So currently only the forward pass is supported and tested.

[domschl]

LSTM backward path is still broken. (torch 1.12.1 and nightly of 2022-09-11)

As @PHRABAL noted, there is most probably a single permute() statement missing that permutes temporal and lstm hidden dimensions. This most probably also fixes the numerical issues.

[AngledLuffa]

This is still broken, and it's hugely problematic for my use case, since our library has bidirectional LSTMs everywhere

import torch

torch.manual_seed(1234)
lstm = torch.nn.LSTM(5, 5, num_layers=2, bidirectional=True, batch_first=True, dropout=0.5)
inp = torch.randn(1, 4, 5)
print(lstm(inp)[0])
print(lstm.to("mps")(inp.to("mps"))[0])

This is using a nightly build from 2022-12-06