[PyTorch] Use torch's register_custom_class API for opaque quantizers when available - #3495
Conversation
github-actions
Bot
added
the
community-contribution
Greptile SummaryThis PR updates Transformer Engine’s PyTorch Dynamo integration to use the current opaque-object API without invoking deprecated wrappers.
Confidence Score: 5/5The PR appears safe to merge because the renamed opaque-object APIs are available throughout the repository’s supported PyTorch range and no outstanding correctness issue remains. The current code consistently uses Important Files Changed
Flowchart%%{init: {'theme': 'neutral'}}%%
flowchart LR
Import[Transformer Engine import] --> Register[register_custom_class]
Register --> Quantizers[Opaque quantizer classes]
Register --> Bundle[OpaqueValueBundle]
Quantizers --> Compile[torch.compile specialization]
Bundle --> Compile
Reviews (4): Last reviewed commit: "[PyTorch] Use torch's register_custom_cl..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| if hasattr(oo, "register_custom_class") and hasattr(oo, "is_opaque_constant_type"): | ||
| return oo.register_custom_class, oo.is_opaque_constant_type, "constant" | ||
| return oo.register_opaque_type, oo.is_opaque_value_type, "value" |
There was a problem hiding this comment.
This helper introduces separate paths for the new and legacy PyTorch APIs, but the updated compile test only runs the path exposed by the installed PyTorch version and does not check for deprecation warnings. Add isolated tests that force both API generations and verify the matching function and typ value; otherwise the warning regression or legacy fallback could break without CI detecting it.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
f1e2312 to
dee7226
Compare
| from ..utils import record_compile_disabled | ||
|
|
||
|
|
||
| def opaque_object_api() -> Tuple[Any, Any, str]: |
There was a problem hiding this comment.
You can remove it and just import new API directly. It already is behind try/except. One would be affected only when one tries to use torch.compile with pytorch with older API, but this would not work due to other issues in Pytorch. Also, first version with torch.compile support will be 2.20, so we can safely remove this logic and test and just change imports.
There was a problem hiding this comment.
Done, thanks. Dropped the helper and its tests and switched the imports directly; the existing try/except handles older torch as before.
PyTorch renamed the opaque-object entry points: register_opaque_type -> register_custom_class, is_opaque_value_type -> is_opaque_constant_type, typ="value" -> typ="constant". The old names remain as thin wrappers that log a deprecation warning on every call, so each TE import now emits 21 warning lines (5 quantizer classes registered, 11 type checks) -- once per process, including every torch.compile worker. A 16-rank training job logs them from ~320 processes. Switch quantizer_opaque.py, custom_op.py and the torch.compile test to the new names. Both import sites already sit behind try/except that records the torch.compile custom-op path as disabled, so a PyTorch that predates the rename (torch < 2.14) keeps importing and running eagerly with the compile path off -- the same behaviour as a PyTorch without the opaque-object API at all. Signed-off-by: Michal Marcinkiewicz <michalm@nvidia.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
cc06308 to
296e419
Compare
|
/te-ci pytorch |
2 participants
Description
Switch quantizer_opaque.py, custom_op.py and the torch.compile test to the new names. Both import sites already sit behind try/except that records the torch.compile custom-op path as disabled, so a PyTorch that predates the rename (torch < 2.14) keeps importing and running eagerly with the compile path off, the same as a PyTorch without the opaque-object API at all.
Fixes # (issue)
Type of change
Checklist: