Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make Casanovo model compatible with torch.compile #347

Open
Lilferrit opened this issue Jun 25, 2024 · 0 comments
Open

Make Casanovo model compatible with torch.compile #347

Lilferrit opened this issue Jun 25, 2024 · 0 comments
Labels
enhancement New feature or request

Comments

@Lilferrit
Copy link
Contributor

Lilferrit commented Jun 25, 2024

When attempting to do a sequencing run using a compiled Casanovo model I encountered this runtime error:

...
torch._dynamo.exc.TorchRuntimeError: Failed running call_module L__self___encoder_transformer_encoder(*(FakeTensor(..., device='cuda:0', size=(128, 121, 512)),), **{'src_key_padding_mask': FakeTensor(..., device='cuda:0', size=(128, 121), dtype=torch.bool)}):
strided nested tensors are not supported by meta conversion

from user code:
   File "C:\Users\gavin\OneDrive\Documents\casanovo\casanovo\denovo\model.py", line 844, in predict_step
    self.forward(batch[0], batch[1]),
  File "C:\Users\gavin\OneDrive\Documents\casanovo\casanovo\denovo\model.py", line 210, in forward
    return self.beam_search_decode(
  File "C:\Users\gavin\OneDrive\Documents\casanovo\casanovo\denovo\model.py", line 240, in beam_search_decode
    memories, mem_masks = self.encoder(spectra)
  File "D:\anaconda3\envs\casanovo_env\lib\site-packages\torch\nn\modules\module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "D:\anaconda3\envs\casanovo_env\lib\site-packages\depthcharge\components\transformers.py", line 105, in forward
    return self.transformer_encoder(peaks, src_key_padding_mask=mask), mask

From the best I can tell this specific error is an issue with PyTorch's built-in TransformerEncoder module, which is known to be incompatible with torch.compile: https://discuss.pytorch.org/t/runtime-error-when-running-inference-on-a-compiled-nn-transformerencoder/198010

In order to run with a compiled Casanovo model I added self.model = torch.compile(self.model) to the end of ModelRunner.initialize_model.

From my understanding making the Casanovo model compatible with torch.compile would be useful for addressing #328.

@Lilferrit Lilferrit added the enhancement New feature or request label Jun 25, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
enhancement New feature or request
Projects
None yet
Development

No branches or pull requests

1 participant