| CARVIEW |
Navigation Menu
-
Notifications
You must be signed in to change notification settings - Fork 24.7k
[Inductor] fix device error for NopKernelSchedulerNode #141372
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
Conversation
BoyuanFeng
added
ciflow/trunk
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/141372
Note: Links to docs will display an error until the docs builds have been completed. âś… You can merge normally! (2 Unrelated Failures)As of commit f412f85 with merge base 51b7528 ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
BROKEN TRUNK - The following job failed but was present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
| isinstance(buffer, ir.ComputedBuffer) | ||
| and buffer.is_zero_elements() | ||
| and device == torch.device("cpu") | ||
| ) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Prior implementation also skips empty CUDA tensors, leading to wrong V.graph.device_ops.
|
@BoyuanFeng has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I am gunna let @eellison give the stamp on this one
| if not isinstance(node, NopKernelSchedulerNode) and ( | ||
| device := node.get_device() | ||
| ): | ||
| if device := node.get_device(): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
empty_strided_cuda generates a NopKernelSchedulerNode. Prior to this PR, we do not create device guard for NopKernelSchedulerNode. This leads to an error if empty_strided_cuda should create a tensor on cuda:1.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks good but a question on decode_device. where is the unnormalized device coming in ?
| @@ -3027,7 +3027,7 @@ def _new_constant( | |||
| dtype = decode_dtype(dtype) or x.get_dtype() | |||
| device = device or x.get_device() | |||
| size = [sympy.Integer(s) for s in size] | |||
| return _full(fill_value, device, dtype, size) | |||
| return _full(fill_value, decode_device(device), dtype, size) | |||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why do we have all the decode_device calls ? I would sort of expect this to be normalized already. or maybe we can normalize them in a more central place. it's going to be really easy to forget to do this
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
- unnormalized device comes from dynamo. E.g., for this user code:
def f(x):
return torch.ops.aten.empty_strided([2,3], [3, 1], device="cuda")We will get the dynamo graph:
def forward(self):
empty_strided: "f32[2, 3][3, 1]cuda:0" = torch.ops.aten.empty_strided([2, 3], [3, 1], device = 'cuda')
return (empty_strided,)
Note that dynamo just copies device="cuda" from the user code. This unnormalized device is directly passed to aotautograd and finally to inductor.
- I talked with @yanboliang on whether we should modify dynamo or inductor. Historically device is normalized in
lowering.pywithdecode_device. Now there are 9 callsites ofdecode_devicein lowering.py on main branch. I think a better way might be removing all thesedecode_deviceand normalize during GraphLowering. I can have a separate PR for that.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yea, i think we should just do this in GraphLowering class prior to invoking the op.. anyway we dont need to block this pr
| @@ -3027,7 +3027,7 @@ def _new_constant( | |||
| dtype = decode_dtype(dtype) or x.get_dtype() | |||
| device = device or x.get_device() | |||
| size = [sympy.Integer(s) for s in size] | |||
| return _full(fill_value, device, dtype, size) | |||
| return _full(fill_value, decode_device(device), dtype, size) | |||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Yea, i think we should just do this in GraphLowering class prior to invoking the op.. anyway we dont need to block this pr
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Merge failedReason: 1 jobs have failed, first few of them are: Meta Internal-Only Changes Check Details for Dev Infra teamRaised by workflow job |
|
@BoyuanFeng has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator. |
|
@pytorchbot merge -f "skip unrelated errors" |
Merge startedYour change will be merged immediately since you used the force (-f) flag, bypassing any CI checks (ETA: 1-5 minutes). Please use Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
This PR adds device guard support for NopKernelSchedulerNode which may create a tensor. Prior to this PR, we do not codegen device guard for NopKernelSchedulerNode, leading to errors.
Prior to the PR:
```python
def call(args):
arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1, arg9_1, arg10_1 = args
args.clear()
assert_size_stride(arg0_1, (1, 1, 2048, 128), (262144, 262144, 128, 1))
assert_size_stride(arg1_1, (1, 1, 2048, 128), (262144, 262144, 128, 1))
assert_size_stride(arg2_1, (1, 1, 2048, 128), (262144, 262144, 128, 1))
assert_size_stride(arg3_1, (1, 1, 16), (16, 16, 1))
assert_size_stride(arg4_1, (1, 1, 16, 16), (256, 256, 16, 1))
assert_size_stride(arg5_1, (1, 1, 16), (16, 16, 1))
assert_size_stride(arg6_1, (1, 1, 16, 16), (256, 256, 16, 1))
assert_size_stride(arg7_1, (1, 1, 16), (16, 16, 1))
assert_size_stride(arg8_1, (1, 1, 16, 16), (256, 256, 16, 1))
assert_size_stride(arg9_1, (1, 1, 16), (16, 16, 1))
assert_size_stride(arg10_1, (1, 1, 16, 16), (256, 256, 16, 1))
buf0 = empty_strided_cuda((1, 1, 2048), (2048, 2048, 1), torch.float32) # TODO: ERROR here. Should be cuda:1
with torch.cuda._DeviceGuard(1):
torch.cuda.set_device(1)
buf1 = empty_strided_cuda((1, 1, 2048, 128), (262144, 262144, 128, 1), torch.bfloat16)
# Topologically Sorted Source Nodes: [flex_attention], Original ATen: []
stream1 = get_raw_stream(1)
breakpoint()
triton_tem_fused_0.run(arg0_1, arg1_1, arg2_1, buf0, arg3_1, arg4_1, arg5_1, arg6_1, buf1, grid=torch._inductor.kernel.flex_attention.flex_attention_grid(1, 1, 2048, 128, meta0), stream=stream1)
del arg0_1
del arg1_1
del arg2_1
del arg3_1
del arg4_1
del arg5_1
del arg6_1
del buf0
return (buf1, )
```
After the PR:
```python
def call(args):
arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1, arg9_1, arg10_1 = args
args.clear()
assert_size_stride(arg0_1, (1, 1, 2048, 128), (262144, 262144, 128, 1))
assert_size_stride(arg1_1, (1, 1, 2048, 128), (262144, 262144, 128, 1))
assert_size_stride(arg2_1, (1, 1, 2048, 128), (262144, 262144, 128, 1))
assert_size_stride(arg3_1, (1, 1, 16), (16, 16, 1))
assert_size_stride(arg4_1, (1, 1, 16, 16), (256, 256, 16, 1))
assert_size_stride(arg5_1, (1, 1, 16), (16, 16, 1))
assert_size_stride(arg6_1, (1, 1, 16, 16), (256, 256, 16, 1))
assert_size_stride(arg7_1, (1, 1, 16), (16, 16, 1))
assert_size_stride(arg8_1, (1, 1, 16, 16), (256, 256, 16, 1))
assert_size_stride(arg9_1, (1, 1, 16), (16, 16, 1))
assert_size_stride(arg10_1, (1, 1, 16, 16), (256, 256, 16, 1))
with torch.cuda._DeviceGuard(1):
torch.cuda.set_device(1)
buf0 = empty_strided_cuda((1, 1, 2048), (2048, 2048, 1), torch.float32) # New: move into device guard
buf1 = empty_strided_cuda((1, 1, 2048, 128), (262144, 262144, 128, 1), torch.bfloat16)
# Topologically Sorted Source Nodes: [flex_attention], Original ATen: []
stream1 = get_raw_stream(1)
triton_tem_fused_0.run(arg0_1, arg1_1, arg2_1, buf0, arg3_1, arg4_1, arg5_1, arg6_1, buf1, grid=torch._inductor.kernel.flex_attention.flex_attention_grid(1, 1, 2048, 128, meta0), stream=stream1)
del arg0_1
del arg1_1
del arg2_1
del arg3_1
del arg4_1
del arg5_1
del arg6_1
del buf0
return (buf1, )
```
Fixes #141010
Pull Request resolved: #141372
Approved by: https://github.com/eellison
This PR adds device guard support for NopKernelSchedulerNode which may create a tensor. Prior to this PR, we do not codegen device guard for NopKernelSchedulerNode, leading to errors.
Prior to the PR:
After the PR:
Fixes #141010
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @aakhundov