-
Notifications
You must be signed in to change notification settings - Fork 22k
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
Wrap output with FakeTensor if input FakeTensor is not preserved #128206
Conversation
…ng type conversion
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/128206
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New FailureAs of commit 6492290 with merge base 70724bd ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
@ezyang @laithsakka I am also checking how a test case can be added for the case showing here. |
I made a deep dive to what is happening as to why fakeness is not preserved during type conversion for this case:
The key part is the mul_ operator call on tensor (int type) and argument 1.0 (float type): Here is what happening for this line when compiling:
|
This change is probably not correct but I need to study the bug report in more detail |
Thanks @ezyang. The key area causing the problem under layer is the following piece of code of elementwise_type_promotion_wrapper in torch/_prims_common/wrappers.py:
_maybe_convert_to_dtype will do the type conversion and will finally calls _to_copy in decompositions. The fakeness is not preserved under this context (in_kernel_invocation_manager). and the promoted argument is tensor while the original argument is a FakeTensor. During the process of fn(**bound.arguments), the tensor argument may be used. If a device check is done within any of the code path, there is a problem. For CPU device, it has not problem for the above test case. For HPU device, there is a device check in is_noncontiguous_supported function of torch/_refs. So the problem is resulted in combination of two things:
For the bug test case running on CPU, only condition 1 is happening and shows no problem. When running on HPU, both condition 1 and condition 2 are happening and thus cause a problem. Just as mentioned in the previous "deep dive" comment, in fake_tensor.py, there is wrap_meta_outputs_with_default_device_logic function called over the returned meta tensor at a later stage. After that, a FakeTensor is wrapped again with the real device and will have no problem here after. |
@laithsakka did we ever end up talking about what we thought the "correct" way to fix this problem was? I read through all of the old issues/PRs but all it says is that this was "pending discussion" at #104689 (comment) and we never followed up. |
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.
Pending discussion
The example in your linked issue fails to run in eager:
Gives:
Do you have another example of this which is non-erroring ? If so, please re-request review. |
@eellison
Please note that the problematic behavior mentioned in the issue shows only on compile mode (related a lot to FakeTensor behaviors related to compile). Eager mode doesn't go with the problematic path. |
@eellison |
I'm still not sure what the error is:
|
@eellison Thank you for checking this. There is no error shown up running on CPU (so I use the term behavior). When I run with HPU device, the data of result tensor is incorrect which is caused by the combination of the behavior (same as on CPU) and another check done specific by HPU depending on it. Since there is no error on CPU, the behavior on CPU is not directly observable without analysis into the internal code path. See my comment for the code path analysis #128206 (comment) and #128206 (comment). Here is the summary of the behavior that is not expected:
This is caused by exactly the following line of code ( of course tensor data types matters) during the compilation: Because b is bfloat16 tensor, 1.0 is float32 type, there is implicit type conversion happening dispatching aten.mul.Tensor in FakeTensorMode. This triggers the behavior that is not expected. To showing up this for analysis, one simple thing to do is to print some tensor info in the _to_copy defined in torch/_decomp/decompositions.py as following (two prints before and after :torch._prims.convert_element_type)
You will see the FakeTensor becomes meta Tensor here. My comment #128206 (comment) has details as to from where it comes here in FakeTensorMode dispatch function. The behavior doesn't show when you do something with explicitly FakeTensorMode context (like you do in the above test code, for all this type of usage you will see FakeTensor to FakeTensor before and after type conversion). It happens in a very specific program pattern I shown above (compile mode, a function with implicit type conversion for its args) Hope this would help start for understanding the problem. |
@eellison Any chance to check this again? |
We are invoking Those outputs of the meta kernel then get wrapped back to the appropriate device with |
@eellison Thank your for checking this and response. Yes, you are right that _to_copy is called under in_kernel_invocation_manager and the meta device tensor is wrapped back by wrap_meta_outputs_with_default_device_logic. So the current problem behavior (state) I got is just after _to_copy and before wrap_meta_outputs_with_default_device_logic. Which is the following code piece in torch/_prims_common/wrappers.py:
The _maybe_convert_to_dtype calls _to_copy to type conversion of arguments and fn(**bound.arguments) may use the arguments lacking of the real device. wrap_meta_outputs_with_default_device_logic happens afterwards and so it doesn't help here. Any thoughts how to improve going forward? |
Is it just pointwise ops that are |
Checked with the original author, the is_noncontiguous_supported handling is not strictly to pointwise ops.
I am not sure whether there are equivalent places in FakeTensor to add these device checking and contiguous wrappings. I saw quite a few places of handling this in torch/_refs/init.py. @ezyang Would you please help a check this part? I am not sure what @eellison suggested (move the contiguous wrapping to FakeTensor) is feasible or not. There are quite a few layers of wrappers and intermediate logics between FakeTensor dispatch function and noncontiguous handling in _refs. |
So, certainly if we reimplement all the pointwise rules directly in fake tensor, we can ensure we never decay to meta (since we're always operating on fake tensors) and then you would be guaranteed to have accurate device info. It does sound a bit duplicatey. But it definitely will work. |
@ezyang Thank you checking on this. |
@jerrychenhf you can check if the operator has the |
Fixes #128202
For cases that the input is a FakeTensor with "meta" device, the output may not preserve its fakeness and come out as Tensor (see _convert_element_type_meta). In this case, the original device type is lost in the output.
This PR pick up the work from #104689 and apply the check of the input and output fakeness originally showing in #119868.