address comments from Yuxian

Signed-off-by: Balaram Buddharaju <[email protected]>

Author: brb-nv

tensorrt_llm/_torch/device_mesh.py

@@ -120,13 +120,17 @@ def build_mesh(self):
 
         # Dimensions go from slowest-varying (outermost) to fastest-varying (innermost).
         # Layout: pp is outermost, then tp, then cp is innermost (consecutive).
-        dims = ["pp", "tp"]
-        shape = [self.pp_size, self.tp_size]
+        dims = ["pp"]
+        shape = [self.pp_size]
 
         if self.moe_ep_size > 1:
             dims += ["moe_tp", "moe_ep"]
             shape += [self.moe_tp_size, self.moe_ep_size]
         else:
+            dims += ["tp"]
+            shape += [self.tp_size]
+
+        if self.cp_size > 1:
             dims += ["cp"]
             shape += [self.cp_size]
 

tensorrt_llm/models/modeling_utils.py

@@ -741,10 +741,10 @@ def from_checkpoint(
         rank = config.mapping.rank
         if config.mapping.cp_size > 1:
             # cp_tp_pp rank -> tp_pp rank: because different cp ranks share the same ckpt.
-            tp_size = config.mapping.tp_size
             cp_size = config.mapping.cp_size
-            rank = (rank % (tp_size * cp_size)) // cp_size + rank // (
-                tp_size * cp_size) * tp_size
+            # rank = pp_rank × tp_size × cp_size + tp_rank × cp_size + cp_rank.
+            # rank // cp_size is equivalent to pp_rank × tp_size + tp_rank.
+            rank = rank // cp_size
         weights_path = os.path.join(ckpt_dir, f'rank{rank}.safetensors')
 
         assert os.path.isfile(weights_path)