fix sharding <100 limitation (#8146)

Author: sneaxiy

paddlenlp/trainer/training_args.py

@@ -1576,20 +1576,20 @@ def pipeline_parallel_rank(self):
         else:
             return 0
 
+    def _format_name(self, prefix, rank, degree):
+        size = max(2, len(str(degree)))
+        return f"{prefix}{rank:0>{size}d}"
+
     @property
     def optimizer_name_suffix(self):
         if self.use_hybrid_parallel:
             name = []
             if self.tensor_parallel_degree > 1:
-                assert self.tensor_parallel_degree < 100, "tensor parallel degree should be less than 100."
-                name.append(f"tp{self.tensor_parallel_rank:0>2d}")
+                name.append(self._format_name("tp", self.tensor_parallel_rank, self.tensor_parallel_degree))
             if self.pipeline_parallel_degree > 1:
-                assert self.pipeline_parallel_degree < 100, "pipeline parallel degree should be less than 100."
-                name.append(f"pp{self.pipeline_parallel_rank:0>2d}")
+                name.append(self._format_name("pp", self.pipeline_parallel_rank, self.pipeline_parallel_degree))
             if self.sharding_parallel_degree > 1:
-                assert self.sharding_parallel_degree < 100, "sharding parallel degree should be less than 100."
-                name.append(f"shard{self.sharding_parallel_rank:0>2d}")
-
+                name.append(self._format_name("shard", self.sharding_parallel_rank, self.sharding_parallel_degree))
             return "_".join(name)
         else:
             return None
@@ -1599,11 +1599,9 @@ def weight_name_suffix(self):
         if self.use_hybrid_parallel:
             name = []
             if self.tensor_parallel_degree > 1:
-                assert self.tensor_parallel_rank < 100, "tensor parallel rank should be less than 100."
-                name.append(f"tp{self.tensor_parallel_rank:0>2d}")
+                name.append(self._format_name("tp", self.tensor_parallel_rank, self.tensor_parallel_degree))
             if self.pipeline_parallel_degree > 1:
-                assert self.pipeline_parallel_degree < 100, "tensor parallel rank should be less than 100."
-                name.append(f"pp{self.pipeline_parallel_rank:0>2d}")
+                name.append(self._format_name("pp", self.pipeline_parallel_rank, self.pipeline_parallel_degree))
             return "_".join(name)
 
         else:
@@ -1613,20 +1611,17 @@ def sharded_name_suffix(self, shard_id=None, pp_id=None):
         if self.use_hybrid_parallel:
             name = []
             if self.tensor_parallel_degree > 1:
-                assert self.tensor_parallel_rank < 100, "tensor parallel rank should be less than 100."
-                name.append(f"tp{self.tensor_parallel_rank:0>2d}")
+                name.append(self._format_name("tp", self.tensor_parallel_rank, self.tensor_parallel_degree))
             if self.pipeline_parallel_degree > 1:
                 if pp_id is None:
                     pp_id = self.pipeline_parallel_rank
                 assert isinstance(pp_id, int)
-                assert pp_id < 100, "pp_id should be less than 100."
-                name.append(f"pp{pp_id:0>2d}")
+                name.append(self._format_name("pp", pp_id, self.pipeline_parallel_degree))
             if self.sharding_parallel_degree > 1:
                 if shard_id is None:
                     shard_id = self.sharding_parallel_rank
                 assert isinstance(shard_id, int)
-                assert shard_id < 100, "shard_id should be less than 100."
-                name.append(f"shard{shard_id:0>2d}")
+                name.append(self._format_name("shard", shard_id, self.sharding_parallel_degree))
             return "_".join(name)
         else:
             return None