address mapping clone

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>

Author: yechank-nvidia

tensorrt_llm/_torch/modules/gated_mlp.py

@@ -6,7 +6,6 @@
 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
@@ -42,19 +41,14 @@ 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 = 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,
-            )
+            mapping = config.mapping.clone()
+            mapping.world_size = overridden_tp_size * pp_size
+            mapping.tp_size = overridden_tp_size
+            mapping.pp_size = pp_size
         else:
             mapping = config.mapping
 

tensorrt_llm/_torch/modules/mlp.py

@@ -4,8 +4,6 @@
 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
@@ -32,19 +30,14 @@ 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 = 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,
-            )
+            mapping = config.mapping.clone()
+            mapping.world_size = overridden_tp_size * pp_size
+            mapping.tp_size = overridden_tp_size
+            mapping.pp_size = pp_size
         else:
             mapping = config.mapping
 

tensorrt_llm/mapping.py

@@ -12,6 +12,7 @@
 # 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
 
@@ -238,6 +239,9 @@ 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