[Inductor] fix device error for NopKernelSchedulerNode #141372
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.
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. |
| if not isinstance(node, NopKernelSchedulerNode) and ( | ||
| device := node.get_device() | ||
| ): | ||
| if device := node.get_device(): |
There was a problem hiding this comment.
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.
| 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.
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.
- 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.
Yea, i think we should just do this in GraphLowering class prior to invoking the op.. anyway we dont need to block this pr
| 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.
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
5 participants
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