make TransformForColumnWiseParallel frozen

Signed-off-by: Masaki Kozuki <mkozuki@nvidia.com>
Lightning-AI · May 17, 2024 · c1e398d · c1e398d
1 parent 5b2f2b9
commit c1e398d
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/thunder/distributed/tensor_parallel.py b/thunder/distributed/tensor_parallel.py
@@ -54,7 +54,7 @@ def __call__(self, bsym: BoundSymbol) -> VISIT_TYPE:
         return VISIT_TYPE.REPLACE
 
 
-@dataclass
+@dataclass(frozen=True)
 class TransformForColumnWiseParallel:
     rank: int
     world_size: int
@@ -69,8 +69,6 @@ def __post_init__(self):
         if getattr(self.compile_data, "use_fsdp", False) or getattr(self.compile_data.fn, "use_fsdp", False):
             raise NotImplementedError("Currently thunder does not support the combination of fsdp and tensor parallel")
 
-        self.swap_map: dict[VariableInterface, ProxyInterface] = {}
-
     def __call__(
         self,
         prologue_trace: TraceCtx,