linkedin / linkedin/Liger-Kernel

LigerGEGLUMLP error with torch.compile

Open
#96 7 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
6.6k
Forks
603
Avg merge
1d 20h
Merged PRs (30d)
47

Description

### 🐛 Describe the bug

When using ligergeglumlp with torch complie i get the following error.

```
UserWarning: Traceback (most recent call last):
Encountered an exception in identify_mutated_tensors, assuming every input is mutated:
File "~/.local/lib/python3.10/site-packages/torch/_higher_order_ops/triton_kernel_wrap.py", line 595, in identify_mutated_tensors
ttir_module, ordered_tensor_names = generate_ttir(kernel, kwargs)
Encountered an exception in identify_mutated_tensors, assuming every input is mutated:
File "~/.local/lib/python3.10/site-packages/torch/_higher_order_ops/triton_kernel_wrap.py", line 115, in generate_ttir
raise Exception("Incorrect number of arguments passed to kernel")
Encountered an exception in identify_mutated_tensors, assuming every input is mutated:
Exception: Incorrect number of arguments passed to kernel

BackendCompilerFailed: backend='inductor' raised:
CompilationError: at 21:18:
col_offsets = tl.arange(0, BLOCK_SIZE)
mask = col_offsets < n_cols
a_row = tl.load(a + col_offsets, mask=mask, other=0).to(tl.float32)
b_row = tl.load(b + col_offsets, mask=mask, other=0)

# tanh approximation form of GELU is computed with:
# 0.5 * a * (1 + tanh(sqrt(2 / pi) * (a + 0.044715 * a^3)))
sqrt_2_over_pi = 0.7978845608028654 # sqrt(2 / pi)
a_cubed = a_row * a_row * a_row
tanh_arg = sqrt_2_over_pi * (a_row + 0.044715 * a_cubed)
tanh_result = tanh(tanh_arg)
^
NameError('tanh is not defined')

```
I can run it with the torch.compile line commented out with no problem.

### Reproduce

```
from liger_kernel.transformers.geglu import LigerGEGLUMLP
from dataclasses import dataclass
import torch
torch.set_float32_matmul_precision('high')

@dataclass
class Config:
hidden_size: int = 768
intermediate_size: int = 768 * 4
hidden_act: str = 'gelu_pytorch_tanh'

cfg = Config()

gegelu = LigerGEGLUMLP(cfg).cuda()

gegelu = torch.compile(gegelu)

x = torch.randn(1,1024,768,device='cuda')

```

### Versions

Python Version: 3.10.12
CUDA Version: 12.1
PyTorch Version: 2.3.0+cu121
Triton Version: 2.3.0

Contributor guide

Open the contributing guide

Research direction

Start by running the provided LigerGEGLUMLP reproducer with PyTorch 2.3.0 and torch.compile, then inspect the LigerGEGLUMLP entry point and the generated kernel around the reported tanh failure. Done means the reproduced model compiles and runs without the reported argument or NameError failures, while the uncompiled path continues to work.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.