HybridEP DispatchHandle not working with torch.compile

### Bug description Hi I'm opening an issue here because I plan to open a PR in pytorch to fix the underlying bug. Basically when using HybridEP with gpt-oss, the DispatchHandle returned from the underlying impl causes the following bug in inductor: ``` Traceback (most recent call last): File "torchtitan/.venv/lib/python3.12/site-packages/torch/distributed/elastic/multiprocessing/errors/__init__.py", line 367, in wrapper return f(*args, **kwargs) ^^^^^^^^^^^^^^^^^^ File "torchtitan/torchtitan/trainer.py", line 846, in train self.train_step(data_iterator) File "torchtitan/torchtitan/trainer.py", line 748, in train_step loss = self.forward_backward_step( ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "torchtitan/torchtitan/observability/structured_logger/structured_logging.py", line 441, in sync_wrapper return func(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "torchtitan/torchtitan/trainer.py", line 692, in forward_backward_step pred = model_parts[0](inputs, **extra_inputs, **extra_kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "torchtitan/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1778, in _wrapped_call_impl return self._call_impl(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "torchtitan/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1884, in _call_impl return inner() ^^^^^^^ File "torchtitan/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1832, in inner result = forward_call(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "torchtitan/torchtitan/models/common/decoder.py", line 139, in forward h = layer(h, self.freqs_cis, attention_masks, positions) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "torchtitan/.venv/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl return self._compiled_call_impl(*args, **kwargs) # type: ignore[misc] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "torchtitan/.venv/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 1158, in compile_wrapper raise e.remove_dynamo_frames() from None # see TORCHDYNAMO_VERBOSE=1 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "torchtitan/.venv/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 1078, in _compile_fx_inner raise InductorError(e, currentframe()).with_traceback( File "torchtitan/.venv/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 1058, in _compile_fx_inner mb_compiled_graph = fx_codegen_and_compile( ^^^^^^^^^^^^^^^^^^^^^^^ File "torchtitan/.venv/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 1845, in fx_codegen_and_compile return scheme.codegen_and_compile(gm, example_inputs, inputs_to_check, graph_kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "torchtitan/.venv/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 1378, in codegen_and_compile _recursive_post_grad_passes(gm, is_inference=is_inference) File "torchtitan/.venv/lib/python3.12/site-packages/torch/_inductor/compile_fx.py", line 581, in _recursive_post_grad_passes post_grad_passes(gm, is_inference) File "torchtitan/.venv/lib/python3.12/site-packages/torch/_inductor/fx_passes/post_grad.py", line 411, in post_grad_passes GraphTransformObserver(gm, "reinplace_inplaceable_ops").apply_graph_pass( File "torchtitan/.venv/lib/python3.12/site-packages/torch/fx/passes/graph_transform_observer.py", line 103, in apply_graph_pass return pass_fn(self.gm.graph) ^^^^^^^^^^^^^^^^^^^^^^ File "torchtitan/.venv/lib/python3.12/site-packages/torch/_inductor/fx_passes/reinplace.py", line 949, in reinplace_inplaceable_ops fake_tensor_updater.incremental_update() File "torchtitan/.venv/lib/python3.12/site-packages/torch/_inductor/fx_utils.py", line 230, in incremental_update if "val" in node.meta and is_fake_tensor_same( ^^^^^^^^^^^^^^^^^^^^ File "torchtitan/.venv/lib/python3.12/site-packages/torch/_inductor/fx_utils.py", line 114, in is_fake_tensor_same assert isinstance(new, (torch.SymInt, torch.SymBool, torch.SymFloat)), ( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ torch._inductor.exc.InductorError: AssertionError: Unknown type <class 'torchtitan.distributed.deepep.hybridep.DispatchHandle'> in graph(): ``` The issue being that FakeTensorUpdater in fx_utils tries to get a sense if the inputs & outputs are the same for the dispatch impl, but it doesn't know how to handle DispatchHandle, which inherits from pytorch's OpaqueBase. I have a monkey patch fix that I am about to open a PR into pytorch for! ### Versions pytorch-nightly torchao-nightly torchtitan-main config: ```python def gpt_oss_20b() -> Trainer.Config: return Trainer.Config( loss=CrossEntropyLoss.Config(), hf_assets_path="./assets/hf/gpt-oss-20b", model_spec=model_registry( flavor="20b", moe_comm_backend="hybridep", converters=[ MXFP8LinearConverter.Config( model_compile_enabled=True, fqns=["qkv_linear", "wo"], ), MXFP8GroupedExpertsConverter.Config( model_compile_enabled=True ), ]), dataloader=HuggingFaceTextDataLoader.Config(dataset="c4", pin_memory=True, num_workers=1, prefetch_factor=2), optimizer=OptimizersContainer.Config(lr=8e-4), lr_scheduler=LRSchedulersContainer.Config( warmup_steps=2000, decay_ratio=0.8, decay_type="cosine", min_lr_factor=0.1, ), training=TrainingConfig( local_batch_size=6, seq_len=8192, steps=10000, ), parallelism=ParallelismConfig( expert_parallel_degree=8, # fsdp_reshard_after_forward="never", ), checkpoint=CheckpointManager.Config(interval=500), activation_checkpoint=ActivationCheckpointConfig(mode="selective"), compile=CompileConfig(enable=True, components=["model", "loss"]), ) ```