Lightning-AI / Lightning-AI/lightning-thunder

Strides of 2D column major Tensor seem to be unexpectedly changed

Open
#1,572 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
1.5k
Forks
121
PR merge metrics
No merged PRs in 30d

Description

## 🐛 Bug

Related: #1415

A 2D column major tensor seems to get its strides changed before the transform for `_scaled_mm` so that it becomes row major.

If `f` in the following snippet returns `y`, the strides of `y` are kept as is.

### To Reproduce

Steps to reproduce the behavior:

1. lightning-thunder @ d4c28e359d28bc34227b7cab361665ee5442ea7c

#### Code sample

```python
import torch
import thunder

def f(x, y, scale_x, scale_y):
return torch._scaled_mm(x, y, scale_a=scale_x, scale_b=scale_y, out_dtype=torch.float32)

def main():
device = torch.device("cuda")

x = torch.randn(32, 64, device=device).to(dtype=torch.float8_e4m3fn)
y = torch.randn(64, 96, device=device).to(dtype=torch.float8_e4m3fn).t().contiguous().t()
print(f"$$$ {x.stride() = }, {y.stride() = }")
scale_x = torch.tensor(1.0, device=device)
scale_y = torch.tensor(1.0, device=device)

expected = f(x, y, scale_x, scale_y)
jitted = thunder.jit(f)
actual = jitted(x, y, scale_x, scale_y)
print(thunder.last_traces(jitted)[-1])
torch.testing.assert_close(actual, expected)

if __name__ == "__main__":
main()
```

The output of this script is as follows and it indicates that the transform's `if not column_major` branch of https://github.com/Lightning-AI/lightning-thunder/blob/crpa/subclass-torchao_float8tensor/thunder/executors/torchex.py#L1410 kicks in.

```
$$$ x.stride() = (64, 1), y.stride() = (1, 64)
# Constructed by Unwrap the actual return value
import torch
from torch import Tensor
from thunder.executors.torchex import no_autocast

@torch.no_grad()
@no_autocast
def computation(x, y, scale_x, scale_y):
# x: "cuda:0 f8_e4m3fn[32, 64]"
# y: "cuda:0 f8_e4m3fn[64, 96]"
# scale_x: "cuda:0 f32[]"
# scale_y: "cuda:0 f32[]"
t5 = torch.transpose(y, 0, 1) # t5: "cuda:0 f8_e4m3fn[96, 64]"
# t5 = ltorch.transpose(y, 0, 1) # t5: "cuda:0 f8_e4m3fn[96, 64]"
# t5 = prims.transpose(y, (1, 0)) # t5: "cuda:0 f8_e4m3fn[96, 64]"

# /opt/pytorch/lightning-thunder/nvfuser_scaled_mm.py:7: return torch._scaled_mm(x, y, scale_a=scale_x, scale_b=scale_y, out_dtype=torch.float32)
t6 = Tensor.contiguous(t5, memory_format=_torch_memory_format_0) # t6: "cuda:0 f8_e4m3fn[96, 64]"
# t6 = ltorch.contiguous(t5, memory_format=_torch_memory_format_0) # t6: "cuda:0 f8_e4m3fn[96, 64]"
# t6 = prims.stride_order(t5, (1, 0)) # t6: "cuda:0 f8_e4m3fn[96, 64]"
del t5
t7 = torch.transpose(t6, 0, 1) # t7: "cuda:0 f8_e4m3fn[64, 96]"
# t7 = ltorch.transpose(t6, 0, 1) # t7: "cuda:0 f8_e4m3fn[64, 96]"
# t7 = prims.transpose(t6, (1, 0)) # t7: "cuda:0 f8_e4m3fn[64, 96]"
del t6

# /opt/pytorch/lightning-thunder/nvfuser_scaled_mm.py:7: return torch._scaled_mm(x, y, scale_a=scale_x, scale_b=scale_y, out_dtype=torch.float32)
t0 = torch._scaled_mm(x, t7, scale_x, scale_y, None, None, torch.float32, False) # t0: "cuda:0 f32[32, 96]"
del t7
return (t0,)
```

### Expected behavior

From my perspective, it doesn't feel intuitive that the input's strides are changed.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Run the provided CUDA reproduction with lightning-thunder at commit d4c28e359d28bc34227b7cab361665ee5442ea7c and compare the input strides with the generated trace. Read thunder/executors/torchex.py around line 1410 and the related issue #1415; done means the 2D column-major input is handled without the unexpected stride change before _scaled_mm.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.