attempt

Author: acisseJZhong

torchtitan/models/common/moe/moe.py

@@ -4,8 +4,9 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Callable, Literal
+from typing import Literal
 
 import torch
 import torch.nn.functional as F

torchtitan/models/common/moe/moe_deepep.py

@@ -6,8 +6,8 @@
 
 """MoE with DeepEP backend for efficient expert-parallel communication."""
 
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Callable
 
 import torch
 from torch.distributed.tensor import DTensor, Partial