move cudnn backend

Author: acisseJZhong

torchtitan/models/common/attention.py

@@ -72,11 +72,6 @@ class LocalMapModule(Module):
     (e.g., CP's ``sdpa_input_fn``) fires, and wraps outputs back to TP
     DTensors **after** all ``forward_hook``s complete.
 
-    Under ``torch.compile``, the ``__call__`` override is skipped (dynamo
-    bypasses custom ``__call__`` and inlines ``forward`` directly), so
-    ``forward`` calls ``_forward_local`` without ``local_map``. This is
-    correct because dynamo decomposes DTensor ops at trace time.
-
     Placements and device mesh are inferred from the input DTensors.
     """
 

torchtitan/models/qwen3/parallelize.py

@@ -10,7 +10,6 @@
 import torch
 import torch._inductor.config
 import torch.nn as nn
-from torch.backends.cuda import SDPBackend
 from torch.distributed.device_mesh import DeviceMesh
 from torch.distributed.tensor import Replicate, Shard
 from torch.distributed.tensor.parallel import (
@@ -135,19 +134,6 @@ def parallelize_qwen3(
             attn_backend,
         )
 
-    if parallel_dims.tp_enabled and parallel_dims.cp_enabled:
-        # Workaround: cuDNN SDPA backward has a stride mismatch bug with CP.
-        # Exclude cuDNN until PyTorch fix lands. See https://github.com/pytorch/pytorch/issues/176915.
-        if attn_backend == "sdpa":
-            # pyrefly: ignore [missing-attribute, not-callable]
-            for block in model.layers.values():
-                block.attention.inner_attention.sdpa_backends = (
-                    [  # pyrefly: ignore [missing-attribute]
-                        SDPBackend.FLASH_ATTENTION,
-                        SDPBackend.MATH,
-                    ]
-                )
-
     if ac_config.mode != "none":
         apply_ac(
             model,