[v1] refactor registry plugin structure and params (#10641)

src/llamafactory/v1/config/training_args.py

@@ -76,7 +76,31 @@ class TrainingArguments:
     )
     dist_config: PluginConfig | None = field(
         default=None,
-        metadata={"help": "Distribution configuration for training."},
+        metadata={"help": "Distributed backend plugin configuration."},
+    )
+    dp_size: int | None = field(
+        default=None,
+        metadata={"help": "Data parallel size, default to world_size // cp_size."},
+    )
+    cp_size: int = field(
+        default=1,
+        metadata={"help": "Context parallel size."},
+    )
+    cp_mode: str = field(
+        default="ulysses",
+        metadata={"help": "Context parallel implementation."},
+    )
+    mp_replicate_size: int = field(
+        default=1,
+        metadata={"help": "Model parallel replicate size."},
+    )
+    mp_shard_size: int | None = field(
+        default=None,
+        metadata={"help": "Model parallel shard size, default to world_size // mp_replicate_size."},
+    )
+    dist_timeout: int = field(
+        default=18000,
+        metadata={"help": "Distributed process group initialization timeout in seconds."},
     )
     optim_config: PluginConfig | None = field(
         default=None,
@@ -149,6 +173,12 @@ class TrainingArguments:
         self.dist_config = get_plugin_config(self.dist_config)
         self.optim_config = get_plugin_config(self.optim_config)
         self.lr_scheduler_config = get_plugin_config(self.lr_scheduler_config)
+        try:
+            from ..plugins.model_plugins.deepspeed_utils import register_deepspeed_dist_config
+
+            register_deepspeed_dist_config(self.dist_config)
+        except ImportError:
+            pass
 
         # The optimizer learning rate has a single source of truth: ``learning_rate``.
         # Propagate it into ``optim_config["lr"]`` so optimizer plugins (e.g. Muon) pick it up