coremltools icon indicating copy to clipboard operation
coremltools copied to clipboard

Potentially wrong number of argument in cumsum op

Open msluszniak opened this issue 1 year ago • 6 comments

I discovered that the expected number of arguments to cumsum op in coremltools might be wrong while exporting PyTorch module to *.pte file (executorch). See https://github.com/apple/coremltools/blob/8b7048ebc034a8b40e48c39df3f33ba9f60feda9/coremltools/converters/mil/frontend/torch/ops.py#L931 Number of arguments is set to 3 while only two are used. Indeed, in my export, I got an error that cumsum op expects 3 arguments while only 2 were provided. Changing this value to 2 fixed suppressed the error.

To Reproduce

import torch
from torch.export import export
from executorch.exir import to_edge
from torch import nn
 
import coremltools as ct
 
class TestCumsum(nn.Module):
    """
    Test class
    """

    def __init__(self) -> None:
        super().__init__()

    def forward(self) -> torch.Tensor:
        """Test forward"""
        h, w = 100, 100
        grid = torch.ones([h, w], device="cpu", dtype=torch.float32)
        
        _y_embed = grid.cumsum(dim=0)
        x_embed = grid.cumsum(dim=1)
        
        return x_embed
 
 
if __name__ == "__main__":
 
    example_args = tuple([])
    aten_dialect = export(TestCumsum(), example_args)
    edge_dialect = to_edge(aten_dialect).exported_program()

    mlmodel = ct.convert(edge_dialect)

System environment (please complete the following information):

  • coremltools version: 8.0b1
  • OS (e.g. MacOS version or Linux type): MacOS Version 14.5
  • Any other relevant version information (e.g. PyTorch or TensorFlow version): PyTorch 2.5.0.dev20240716

Additional context

Discovered this while exporting Pytorch model to executorch with CoreML backend.

msluszniak avatar Aug 09 '24 14:08 msluszniak

Right now the torch version coremltools officially supports is 2.3.0 (https://github.com/apple/coremltools/blob/main/coremltools/_deps/init.py#L156), while your enviroment got 2.5.0.

jakesabathia2 avatar Aug 15 '24 04:08 jakesabathia2

@jakesabathia2 From what I can tell this issue persists over all versions of pytorch from at least 2.0, as the signature of the torch.cumsum has not changed since. Can You suggest why this would work with older version of PyTorch?

mkopcins avatar Oct 11 '24 09:10 mkopcins

The third argument is optional. https://pytorch.org/docs/stable/generated/torch.cumsum.html

It is treated as optional in torch.onnx: https://github.com/pytorch/pytorch/blob/e30c55ee527b40d67555464b9e402b4b7ce03737/torch/onnx/symbolic_opset11.py#L412-L424

and you guys actually not use the third argument here - you ignore it https://github.com/apple/coremltools/blob/8b7048ebc034a8b40e48c39df3f33ba9f60feda9/coremltools/converters/mil/frontend/torch/ops.py#L929-L936

wojtke avatar Oct 11 '24 09:10 wojtke

As an update I upload the code that generates the issue. Below are all version of libraries:

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

OS: macOS 15.0.1 (arm64)
GCC version: Could not collect
Clang version: 16.0.0 (clang-1600.0.26.3)
CMake version: version 3.30.4
Libc version: N/A

Python version: 3.10.0 (default, Mar 3 2022, 03:54:28) [Clang 12.0.0 ] (64-bit runtime)
Python platform: macOS-15.0.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 M3 Pro

Versions of relevant libraries:
[pip3] executorch==0.5.0a0+cb3a546
[pip3] executorchcoreml==0.0.1
[pip3] numpy==1.21.3
[pip3] torch==2.6.0.dev20241007
[pip3] torchaudio==2.5.0.dev20241007
[pip3] torchsr==1.0.4
[pip3] torchvision==0.20.0.dev20241007
[conda] executorch 0.5.0a0+cb3a546 pypi_0 pypi
[conda] executorchcoreml 0.0.1 pypi_0 pypi
[conda] numpy 1.21.3 pypi_0 pypi
[conda] torch 2.6.0.dev20241007 pypi_0 pypi
[conda] torchaudio 2.5.0.dev20241007 pypi_0 pypi
[conda] torchsr 1.0.4 pypi_0 pypi
[conda] torchvision 0.20.0.dev20241007 pypi_0 pypi

and the actual code

import torch
from torch.export import export
from executorch.exir import to_edge
from torch import nn
 
import coremltools as ct
 
class TestCumsum(nn.Module):
    """
    Test class
    """

    def __init__(self) -> None:
        super().__init__()

    def forward(self) -> torch.Tensor:
        """Test forward"""
        h, w = 100, 100
        grid = torch.ones([h, w], device="cpu", dtype=torch.float32)
        
        _y_embed = grid.cumsum(dim=0)
        x_embed = grid.cumsum(dim=1)
        
        return x_embed
 
 
if __name__ == "__main__":
 
    example_args = tuple([])
    aten_dialect = export(TestCumsum(), example_args)
    edge_dialect = to_edge(aten_dialect).exported_program()

    mlmodel = ct.convert(edge_dialect)

msluszniak avatar Oct 16 '24 13:10 msluszniak

Locally confirmed fix with the (slightly modified) reproduce with torch 2.4

import numpy as np

import torch
from torch.export import export
from executorch.exir import to_edge
from torch import nn

import coremltools as ct


class TestCumsum(nn.Module):
    def forward(self, grid) -> torch.Tensor:
        _y_embed = grid.cumsum(dim=0)
        x_embed = grid.cumsum(dim=1)
        
        return _y_embed, x_embed
 
 
if __name__ == "__main__":
    torch_model = TestCumsum()
    torch_model.eval()

    h, w = 100, 100
    grid = torch.ones([h, w], device="cpu", dtype=torch.float32)

    aten_dialect = export(torch_model, (grid,))
    edge_dialect = to_edge(aten_dialect).exported_program()

    coreml_model = ct.convert(edge_dialect)


    outputs_torch = torch_model(grid)
    y_embed_torch = outputs_torch[0].detach().numpy()
    x_embed_torch = outputs_torch[1].detach().numpy()
    
    outputs_coreml = coreml_model.predict({"grid": grid.detach().numpy()})
    y_embed_coreml = outputs_coreml["aten_cumsum_default"]
    x_embed_coreml = outputs_coreml["aten_cumsum_default_1"]

    np.testing.assert_allclose(y_embed_coreml, y_embed_torch)
    np.testing.assert_allclose(x_embed_coreml, x_embed_torch)

Will include this fix in next release

YifanShenSZ avatar Oct 17 '24 00:10 YifanShenSZ

Great! Thank you much for your help 😁

msluszniak avatar Oct 17 '24 09:10 msluszniak

Fixed in coremltools 8.1 release with updated cumsum support 🎊 Please try it out

YifanShenSZ avatar Nov 20 '24 20:11 YifanShenSZ