rollback to mapping class

Signed-off-by: yechank <[email protected]>

Author: yechank-nvidia

tensorrt_llm/_torch/modules/gated_mlp.py

@@ -6,6 +6,7 @@
 from torch import nn
 
 from tensorrt_llm.logger import logger
+from tensorrt_llm.mapping import Mapping
 
 from ..distributed import AllReduceParams
 from ..model_config import ModelConfig
@@ -41,14 +42,19 @@ def __init__(
         self.activation = activation
 
         config = config or ModelConfig()
+        self.mapping = config.mapping
         if overridden_tp_size is not None:
             assert config.mapping.tp_size % overridden_tp_size == 0
+            tp_size = overridden_tp_size
             # "Misuse" pp_size here to perform all-reduce within smaller groups
             pp_size = config.mapping.pp_size * config.mapping.tp_size // overridden_tp_size
-            mapping = config.mapping.clone()
-            mapping.world_size = overridden_tp_size * pp_size
-            mapping.tp_size = overridden_tp_size
-            mapping.pp_size = pp_size
+            mapping = Mapping(
+                world_size=tp_size * pp_size,
+                rank=self.mapping.rank,
+                gpus_per_node=self.mapping.gpus_per_node,
+                tp_size=tp_size,
+                pp_size=pp_size,
+            )
         else:
             mapping = config.mapping
 

tensorrt_llm/_torch/modules/mlp.py

@@ -4,6 +4,8 @@
 import torch
 from torch import nn
 
+from tensorrt_llm.mapping import Mapping
+
 from ..model_config import ModelConfig
 from ..peft.lora.layer import LoraLayer, LoraModuleType
 from .linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig
@@ -30,14 +32,19 @@ def __init__(self,
         self.activation = activation
 
         config = config or ModelConfig()
+        self.mapping = config.mapping
         if overridden_tp_size is not None:
             assert config.mapping.tp_size % overridden_tp_size == 0
+            tp_size = overridden_tp_size
             # "Misuse" pp_size here to perform all-reduce within smaller groups
             pp_size = config.mapping.pp_size * config.mapping.tp_size // overridden_tp_size
-            mapping = config.mapping.clone()
-            mapping.world_size = overridden_tp_size * pp_size
-            mapping.tp_size = overridden_tp_size
-            mapping.pp_size = pp_size
+            mapping = Mapping(
+                world_size=tp_size * pp_size,
+                rank=self.mapping.rank,
+                gpus_per_node=self.mapping.gpus_per_node,
+                tp_size=tp_size,
+                pp_size=pp_size,
+            )
         else:
             mapping = config.mapping
 

tensorrt_llm/mapping.py

@@ -12,7 +12,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import copy
 from enum import IntEnum
 from typing import List
 
@@ -239,9 +238,6 @@ def local_rank(self):
     def dp_size(self):
         return self.tp_size if self.enable_attention_dp else 1
 
-    def clone(self):
-        return copy.deepcopy(self)
-
     def has_cp_ulysses(self):
         return self.cp_size > 1 and self.cp_config.get(
             "cp_type") == CpType.ULYSSES