Lightning-AI / Lightning-AI/lightning-thunder
Logging, especially in `Executor.can_execute`, `Executor.can_fuse`, and splitter in `thunderfx`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 1.5k
- Forks
- 121
- PR merge metrics
- No merged PRs in 30d
Description
## 🚀 Feature
Logging. For the sake of better user/developer experience, it'd be necessary for thunder to have logging, e.g. in `thunderfx`'s splitter and `Executor.can_execute`/`Executor.can_fuse` so that we can tell where and how splits are made and why and how some ops are not taken care of by some Executors.
- https://github.com/Lightning-AI/lightning-thunder/blob/1c0ea1e3d7954ea8bcac997109c3e62c964141b9/thunder/dynamo/splitter.py#L28
- https://github.com/Lightning-AI/lightning-thunder/blob/1c0ea1e3d7954ea8bcac997109c3e62c964141b9/thunder/extend/__init__.py#L87-L120
### Motivation
When I was trying `thunder.dynamo.thunderfx` to torchao's [`to_mx`](https://github.com/pytorch/ao/blob/dd22777024e31eecd4ee2312a269e66e37c6e6f5/torchao/prototype/mx_formats/mx_tensor.py#L145) function that takes a `torch.Tensor` in fp32, fp16, or bf16 and returns mxdtype tensor and e8m0 scale tensor, even after having added missing ops, some ops don't seem to be taken care of by nvfuser executor. The trace in question is as follows:
```python
def computation(l_data_hp_):
# l_data_hp_: "cuda:0 bf16[4096, 4096]"
[add, isnan] = nvFusion0(l_data_hp_)
# data_hp = prims.reshape(l_data_hp_, (524288, 32)) # data_hp: "cuda:0 bf16[524288, 32]"
# t1 = prims.convert_element_type(data_hp, dtypes.float32) # t1: "cuda:0 f32[524288, 32]"
# t2 = prims.abs(t1) # t2: "cuda:0 f32[524288, 32]"
# t5 = prims.amax(t2, (1,)) # t5: "cuda:0 f32[524288]"
# max_abs = prims.convert_element_type(t5, dtypes.bfloat16) # max_abs: "cuda:0 bf16[524288]"
# eq = prims.eq(max_abs, 0.0) # eq: "cuda:0 b8[524288]"
# type_1 = prims.convert_element_type(eq, dtypes.bfloat16) # type_1: "cuda:0 bf16[524288]"
# t9 = prims.convert_element_type(type_1, dtypes.float32) # t9: "cuda:0 f32[524288]"
# t10 = prims.mul(1.1754943508222875e-38, t9) # t10: "cuda:0 f32[524288]"
# t14 = prims.add(t5, t10) # t14: "cuda:0 f32[524288]"
# add = prims.convert_element_type(t14, dtypes.bfloat16) # add: "cuda:0 bf16[524288]"
# isnan = prims.ne(max_abs, max_abs) # isnan: "cuda:0 b8[524288]"
# .4:12: max_abs_int32 = add.view(torch.int16); add = None
t73 = Tensor.view(add, torch.int16) # t73: "cuda:0 i16[524288]"
# t73 = ltorch.view(add, torch.int16) # t73: "cuda:0 i16[524288]"
# t73 = ltorch.to(add, None, None, device=None, dtype=torch.int16, copy=False, memory_format=None) # t73: "cuda:0 i16[524288]"
# t73 = prims.convert_element_type(add, dtypes.int16) # t73: "cuda:0 i16[524288]"
del add
# .4:13: rshift = max_abs_int32 >> 7; max_abs_int32 = None
t60 = torch.bitwise_right_shift(t73, 7) # t60: "cuda:0 i16[524288]"
# t60 = ltorch.bitwise_right_shift(t73, 7) # t60: "cuda:0 i16[524288]"
# t60 = prims.bitwise_right_shift(t73, 7) # t60: "cuda:0 i16[524288]"
del t73
# .4:14: and_ = rshift & 255; rshift = None
t61 = torch.bitwise_and(t60, 255) # t61: "cuda:0 i16[524288]"
# t61 = ltorch.bitwise_and(t60, 255) # t61: "cuda:0 i16[524288]"
# t61 = prims.bitwise_and(t60, 255) # t61: "cuda:0 i16[524288]"
del t60
# .4:15: extracted_pow2 = and_ - 127; and_ = None
t62 = torch.sub(t61, 127, alpha=1) # t62: "cuda:0 i16[524288]"
# t62 = ltorch.sub(t61, 127, alpha=1) # t62: "cuda:0 i16[524288]"
# t62 = prims.sub(t61, 127) # t62: "cuda:0 i16[524288]"
del t61
# .4:16: scale_e8m0_unbiased = extracted_pow2 - 8; extracted_pow2 = None
t63 = torch.sub(t62, 8, alpha=1) # t63: "cuda:0 i16[524288]"
# t63 = ltorch.sub(t62, 8, alpha=1) # t63: "cuda:0 i16[524288]"
# t63 = prims.sub(t62, 8) # t63: "cuda:0 i16[524288]"
del t62
# .4:17: scale_e8m0_unbiased_1 = torch.clamp(scale_e8m0_unbiased, min = -127, max = 128); scale_e8m0_unbiased = None
t71 = torch.clamp(t63, -127, 128) # t71: "cuda:0 i16[524288]"
# t71 = ltorch.clamp(t63, -127, 128) # t71: "cuda:0 i16[524288]"
# t64 = ltorch.ne(t63, t63) # t64: "cuda:0 b8[524288]"
# t64 = prims.ne(t63, t63) # t64: "cuda:0 b8[524288]"
# t65 = ltorch.gt(t63, -127) # t65: "cuda:0 b8[524288]"
# t65 = prims.gt(t63, -127) # t65: "cuda:0 b8[524288]"
# t66 = ltorch.where(t65, t63, -127) # t66: "cuda:0 i16[524288]"
# t66 = prims.where(t65, t63, -127) # t66: "cuda:0 i16[524288]"
# t67 = ltorch.where(t64, t63, t66) # t67: "cuda:0 i16[524288]"
# t67 = prims.where(t64, t63, t66) # t67: "cuda:0 i16[524288]"
# t68 = ltorch.ne(t67, t67) # t68: "cuda:0 b8[524288]"
# t68 = prims.ne(t67, t67) # t68: "cuda:0 b8[524288]"
# t69 = ltorch.lt(t67, 128) # t69: "cuda:0 b8[524288]"
# t69 = prims.lt(t67, 128) # t69: "cuda:0 b8[524288]"
# t70 = ltorch.where(t69, t67, 128) # t70: "cuda:0 i16[524288]"
# t70 = prims.where(t69, t67, 128) # t70: "cuda:0 i16[524288]"
# t71 = ltorch.where(t68, t67, t70) # t71: "cuda:0 i16[524288]"
# t71 = prims.where(t68, t67, t70) # t71: "cuda:0 i16[524288]"
del t63
# .4:18: scale_e8m0_biased = scale_e8m0_unbiased_1 + 127; scale_e8m0_unbiased_1 = None
t72 = torch.add(t71, 127, alpha=1) # t72: "cuda:0 i16[524288]"
# t72 = ltorch.add(t71, 127, alpha=1) # t72: "cuda:0 i16[524288]"
# t72 = prims.add(t71, 127) # t72: "cuda:0 i16[524288]"
del t71
# .4:19: scale_e8m0_biased_1 = scale_e8m0_biased.to(torch.uint8); scale_e8m0_biased = None
t74 = Tensor.to(t72, copy=False, dtype=torch.uint8) # t74: "cuda:0 ui8[524288]"
# t74 = ltorch.to(t72, None, None, device=None, dtype=torch.uint8, copy=False, memory_format=None) # t74: "cuda:0 ui8[524288]"
# t74 = prims.convert_element_type(t72, dtypes.uint8) # t74: "cuda:0 ui8[524288]"
del t72
# .4:21: scale_e8m0_biased_2 = torch.where(isnan, 255, scale_e8m0_biased_1); isnan = scale_e8m0_biased_1 = None
t75 = torch.where(isnan, 255, t74) # t75: "cuda:0 ui8[524288]"
# t75 = ltorch.where(isnan, 255, t74) # t75: "cuda:0 ui8[524288]"
# t75 = prims.where(isnan, 255, t74) # t75: "cuda:0 ui8[524288]"
del isnan, t74
# .4:22: to_1 = scale_e8m0_biased_2.to(torch.int32)
t76 = Tensor.to(t75, copy=False, dtype=torch.int32) # t76: "cuda:0 i32[524288]"
# t76 = ltorch.to(t75, None, None, device=None, dtype=torch.int32, copy=False, memory_format=None) # t76: "cuda:0 i32[524288]"
# t76 = prims.convert_element_type(t75, dtypes.int32) # t76: "cuda:0 i32[524288]"
# .4:30: scale_e8m0_biased_3 = scale_e8m0_biased_2.view(torch.float8_e8m0fnu); scale_e8m0_biased_2 = None
t77 = Tensor.view(t75, torch.float8_e8m0fnu) # t77: "cuda:0 f8_e8m0fnu[524288]"
# t77 = ltorch.view(t75, torch.float8_e8m0fnu) # t77: "cuda:0 f8_e8m0fnu[524288]"
# t77 = ltorch.to(t75, None, None, device=None, dtype=torch.float8_e8m0fnu, copy=False, memory_format=None) # t77: "cuda:0 f8_e8m0fnu[524288]"
# t77 = prims.convert_element_type(t75, dtypes.float8_e8m0fnu) # t77: "cuda:0 f8_e8m0fnu[524288]"
del t75
[data_lp_2] = nvFusion1(t76, l_data_hp_)
# lshift = prims.bitwise_left_shift(t76, 23) # lshift: "cuda:0 i32[524288]"
# scale_fp32 = prims.convert_element_type(lshift, dtypes.float32) # scale_fp32: "cuda:0 f32[524288]"
# t37 = prims.gt(scale_fp32, 1.1754943508222875e-38) # t37: "cuda:0 b8[524288]"
# t38 = prims.where(t37, scale_fp32, 1.1754943508222875e-38) # t38: "cuda:0 f32[524288]"
# t36 = prims.ne(scale_fp32, scale_fp32) # t36: "cuda:0 b8[524288]"
# scale_fp32_1 = prims.where(t36, scale_fp32, t38) # scale_fp32_1: "cuda:0 f32[524288]"
# unsqueeze = prims.broadcast_in_dim(scale_fp32_1, [524288, 1], [0]) # unsqueeze: "cuda:0 f32[524288, 1]"
# data_hp = prims.reshape(l_data_hp_, (524288, 32)) # data_hp: "cuda:0 bf16[524288, 32]"
# t41 = prims.broadcast_in_dim(unsqueeze, (524288, 32), (0, 1)) # t41: "cuda:0 f32[524288, 32]"
# t42 = prims.convert_element_type(data_hp, dtypes.float32) # t42: "cuda:0 f32[524288, 32]"
# data_lp = prims.div(t42, t41) # data_lp: "cuda:0 f32[524288, 32]"
# data_lp_1 = prims.convert_element_type(data_lp, dtypes.float8_e4m3fn) # data_lp_1: "cuda:0 f8_e4m3fn[524288, 32]"
# data_lp_2 = prims.reshape(data_lp_1, (4096, 4096)) # data_lp_2: "cuda:0 f8_e4m3fn[4096, 4096]"
del t76
return (t77, data_lp_2)
```
Obviously there are ops such as `add` and `sub` which I'd expect to be fused by nvfuser.
- https://github.com/Lightning-AI/lightning-thunder/blob/1c0ea1e3d7954ea8bcac997109c3e62c964141b9/thunder/executors/nvfuserex_impl.py#L1772-L1779
- https://github.com/Lightning-AI/lightning-thunder/blob/1c0ea1e3d7954ea8bcac997109c3e62c964141b9/thunder/executors/nvfuserex_impl.py#L1946-L1953
I'm struggling efficiently collecting the information which helps me understand why those `BoundSymbol`s are not fused.
So in conclusion, I'm expecting that logging in `can_execute` and `can_fuse` would be pretty beneficial.
Also, while `ThunderCompiler` has [`SplitReason`](https://github.com/Lightning-AI/lightning-thunder/blob/0c6f9a91f1ac955bd5c1087ae26d120d7ab184a3/thunder/dynamo/utils.py#L73), I think relevant log messages would be helpful as well.
### Pitch
Configurable, flexible logging like [`TORCH_LOGS`](https://docs.pytorch.org/docs/stable/logging.html) (and similar features of `TORCH_TRACE` and `tlparse` would also be great[^1]).
#### Configuration
For example, `THUNDER_LOGS="nvfuser"` would display log messages on console when ops get marked as not executable or not fusible. When `THUNDER_LOGS="+nvfuser"`, then debug messages are also displayed.
- `THUNDER_LOGS="executor_name"`: WARNING and higher messages of the specified executor.
- `THUNDER_LOGS="+executor_name"`: DEBUG and higher messages of the specified executor.
- `THUNDER_LOGS="executors"`: WARNING and higher messages of all executors.
- `THUNDER_LOGS=+executors"`: DEBUG and higher messages of all executors.
- `THUNDER_LOGS="splits": WARNINGS and higher messages of thunderfx's splitter: when and why splits are made.
[^1]: See https://github.com/Lightning-AI/lightning-thunder/pull/2182
#### Formats
For `logging.StreamHandler`[^2], human readable formats such as `[%(asctime)s] - %(levelname)s - %(name)s:%(lineno)d - %(message)s` would be reasonable. For `logging.FileHandler`, JSON compatible format would be a good fit, e.g. `{"asctime": "2024-06-12 23:48:01,234", "levelname": "INFO", "name": "thunder.executors.fusion", "lineno": 150, "message": "Fusion executor initialized."}`.
[^2]: https://docs.python.org/3/howto/logging-cookbook.html is a nice reference to get familiar with Python standard `logging` module.
### Alternatives
I might fill out later...
### Additional context
I might fill out later...
cc @borda @mruberry
Contributor guide
No contributing guide indexed for this repository
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with thunder/dynamo/splitter.py and thunder/extend/__init__.py, then inspect the nvfuser executor checks in thunder/executors/nvfuserex_impl.py and SplitReason in thunder/dynamo/utils.py. Run the thunderfx path described in the issue and trace where execution, fusion, and splitting decisions are made. Done means configurable logging explains those decisions for the requested executor and splitter settings.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- compilers, observability
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100