From 14d08a1fa9a478a588c9e54d134224cbf7149252 Mon Sep 17 00:00:00 2001 From: Chien-Chin Huang Date: Tue, 7 Jul 2026 10:12:04 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- torchtitan/distributed/deepep/hybridep.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/torchtitan/distributed/deepep/hybridep.py b/torchtitan/distributed/deepep/hybridep.py index 0bd7ed702c..28c303cfbd 100644 --- a/torchtitan/distributed/deepep/hybridep.py +++ b/torchtitan/distributed/deepep/hybridep.py @@ -23,14 +23,15 @@ import torch import torch.distributed as dist -from torch._library.opaque_object import OpaqueBase, register_opaque_type +from torch._custom_class_base import CustomClassBase +from torch._library.opaque_object import register_custom_class from torchtitan.tools.logging import logger _buffer: Any = None # Global buffer instance -class DispatchHandle(OpaqueBase): +class DispatchHandle(CustomClassBase): """Opaque wrapper for HybridEP dispatch handle. Wraps the deep_ep dispatch handle as an opaque type so it can be returned @@ -58,7 +59,7 @@ def __fx_repr__(self): return "DispatchHandle()", {"DispatchHandle": DispatchHandle} -register_opaque_type(DispatchHandle, typ="reference") +register_custom_class(DispatchHandle, typ="reference") @dataclass