frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

[export] Add draft-export to error msg (#151065)

Browse Source 
Given an exception in torch.export, I want to try/catch it to add the message "hey try out draft-export!". Currently I only add this message for errors that draft-export is known to fix, like DataDependentErrors, ConstraintViolationErrors, and no fake impl.

Originally the error message looks like:
```
  File "/data/users/angelayi/pytorch/torch/_library/custom_ops.py", line 626, in fake_impl
    raise RuntimeError(
RuntimeError: There was no fake impl registered for <CustomOpDef(mylib::foo2)>. This is necessary for torch.compile/export/fx tracing to work. Please use `foo2_impl.register_fake` to add an fake impl.
```

Now, the error msg now looks something like:
```
  File "/data/users/angelayi/pytorch/torch/_library/custom_ops.py", line 626, in fake_impl
    raise RuntimeError(
RuntimeError: There was no fake impl registered for <CustomOpDef(mylib::foo2)>. This is necessary for torch.compile/export/fx tracing to work. Please use `foo2_impl.register_fake` to add an fake impl.

The error above occurred when calling torch.export.export. If you would like to view some more information about this error, and get a list of all other errors that may occur in your export call, you can rerun your program with the `DRAFT_EXPORT=1` envvar, or replace your `export()` call with `draft_export()`.
```

In python versions >= 3.11, we can use `exception.add_note` to add to the error message. However with previous versions I did a hack to modify `e.args`.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/151065
Approved by: https://github.com/pianpwk
ghstack dependencies: #151051
This commit is contained in:
angelayi
2025-04-15 11:18:47 -07:00
committed by PyTorch MergeBot
parent 84e633e09d
commit 41c97a72a1
1 changed files with 36 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
torch/export/__init__.py
View File

@ -272,6 +272,8 @@ def export(
"Maybe try converting your ScriptModule to an ExportedProgram "
"using `TS2EPConverter(mod, args, kwargs).convert()` instead."
)
try:
return _export(
mod,
args,
@ -281,6 +283,31 @@ def export(
preserve_module_call_signature=preserve_module_call_signature,
pre_dispatch=True,
)
except Exception as e:
draft_export_msg = (
"The error above occurred when calling torch.export.export. If you would "
"like to view some more information about this error, and get a list "
"of all other errors that may occur in your export call, you can "
"replace your `export()` call with `draft_export()`."
)
# For errors that we know can be caught by draft-export, add the message
# to ask users to try out draft-export
if isinstance(
e,
(
torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode,
torch._subclasses.fake_tensor.UnsupportedOperatorException,
torch._dynamo.exc.UserError,
torch.fx.experimental.symbolic_shapes.ConstraintViolationError,
),
):
new_msg = str(e) + "\n\n" + draft_export_msg
e.args = (new_msg,)
elif isinstance(e, RuntimeError) and "no fake impl registered" in str(e):
new_msg = str(e) + "\n\n" + draft_export_msg
e.args = (new_msg,)
raise e
DEFAULT_PICKLE_PROTOCOL = 2