fix(torchrun): use correct datatypes for torchrun args (Red-Hat-AI-Innovation-Team#44)

* fix(torchrun): use correct datatypes for torchrun args

Torchrun supports nproc_per_node and rdzv_id as str.
TorchrunArgs only supports int, which is permissible
by pytorch.

This change will enable TorchrunArgs to support both str, int.
Also, remove unset or empty parameters before passing it to
torchrun args.

Signed-off-by: Saad Zaher <szaher@redhat.com>

* Use python3.11 style for pydatnic model

Signed-off-by: Saad Zaher <szaher@redhat.com>

* replace - with _ for cli args

Signed-off-by: Saad Zaher <szaher@redhat.com>

* make nproc_per_node to only accept gpu or int. Remove Defaults

Signed-off-by: Saad Zaher <szaher@redhat.com>

* add master_{addr, port} validate args

Signed-off-by: Saad Zaher <szaher@redhat.com>

* deep check if variables are set and not empty

Co-authored-by: Oleg Silkin <97077423+RobotSail@users.noreply.github.com>

* Update src/mini_trainer/training_types.py

Co-authored-by: Oleg Silkin <97077423+RobotSail@users.noreply.github.com>

* Update src/mini_trainer/api_train.py

Co-authored-by: Oleg Silkin <97077423+RobotSail@users.noreply.github.com>

* does not automatically set --master-port

* Update api_train.py

* use standalone when neither rdzv_endpoint nor master_addr are provided

* Update training_types.py

* update tests

---------

Signed-off-by: Saad Zaher <szaher@redhat.com>
Co-authored-by: Oleg Silkin <97077423+RobotSail@users.noreply.github.com>

Author: szaher

src/mini_trainer/api_train.py

@@ -85,14 +85,34 @@ def run_training(torch_args: TorchrunArgs, train_args: TrainingArgs) -> None:
 
     # Build torchrun command
     train_script = Path(__file__).parent / "train.py"
-    
+
     command = [
         "torchrun",
         f"--nnodes={torch_args.nnodes}",
-        f"--node_rank={torch_args.node_rank}",
-        f"--nproc_per_node={torch_args.nproc_per_node}",
-        f"--rdzv_id={torch_args.rdzv_id}",
-        f"--rdzv_endpoint={torch_args.rdzv_endpoint}",
+        f"--node-rank={torch_args.node_rank}",
+        f"--nproc-per-node={torch_args.nproc_per_node}",
+        f"--rdzv-id={torch_args.rdzv_id}",
+    ]
+
+    if torch_args.master_addr and torch_args.rdzv_endpoint:
+        raise ValueError("Provide either `rdzv_endpoint` OR `master_addr`, not both.")
+
+    if torch_args.master_addr:
+        # master-addr + master-port are only compatible with the static backend
+        # so here we pass it explicitly
+        command += [
+            f"--master-addr={torch_args.master_addr}",
+            "--rdzv-backend=static"
+        ]
+        if torch_args.master_port:
+            command += [f"--master-port={torch_args.master_port}"]
+
+    elif torch_args.rdzv_endpoint:
+        command += [f"--rdzv-endpoint={torch_args.rdzv_endpoint}"]
+    else:
+        command += ["--standalone"]
+
+    command.extend([
         str(train_script),
         f"--model-name-or-path={train_args.model_name_or_path}",
         f"--data-path={train_args.data_path}",
@@ -109,7 +129,7 @@ def run_training(torch_args: TorchrunArgs, train_args: TrainingArgs) -> None:
         f"--max-steps={train_args.max_steps}",
         f"--max-tokens={train_args.max_tokens}",
         f"--train-dtype={train_args.train_dtype}",
-    ]
+    ])
 
 
     # wandb-related arguments

src/mini_trainer/training_types.py

@@ -7,7 +7,7 @@
 
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, Literal
 
 
 class TrainingMode(str, Enum):
@@ -22,11 +22,24 @@ class TrainingMode(str, Enum):
 class TorchrunArgs:
     """Arguments for torchrun distributed training configuration."""
     nnodes: int = 1
-    nproc_per_node: int = 1
+    nproc_per_node: Literal["gpu"] | int = 1
     node_rank: int = 0
-    rdzv_id: int = 123
-    rdzv_endpoint: str = "127.0.0.1:1738"
+    rdzv_id: str | int = 123
 
+    # Optional rendezvous / master fields
+    rdzv_endpoint: Optional[str] = None
+    master_addr: Optional[str] = None
+    master_port: Optional[int] = None
+
+    def __post_init__(self):
+        # in order to support systems which are still relying on `master_addr`
+        # to construct the rendezvous address, torchrun must not be given a non-empty value
+        # for rdzv_endpoint:
+        # https://github.com/pytorch/pytorch/blob/ecb53078faf86ca1b33277df33b82985675bb011/torch/distributed/run.py#L799
+        if self.rdzv_endpoint and self.master_addr: 
+            raise ValueError(
+                "Provide either `rdzv_endpoint` OR both `master_addr` and `master_port`, not both."
+            )
 
 @dataclass
 class TrainingArgs:

tests/test_api_train.py

@@ -28,7 +28,7 @@ def test_torchrun_args_defaults(self):
         assert args.nproc_per_node == 1
         assert args.node_rank == 0
         assert args.rdzv_id == 123
-        assert args.rdzv_endpoint == "127.0.0.1:1738"
+        assert args.rdzv_endpoint == None
 
         # Test with custom nproc_per_node only
         args = TorchrunArgs(nproc_per_node=8)
@@ -405,10 +405,10 @@ def test_run_training_command_construction(self, mock_popen_class):
             # Verify command structure
             assert command[0] == "torchrun"
             assert "--nnodes=2" in command
-            assert "--node_rank=1" in command
-            assert "--nproc_per_node=4" in command
-            assert "--rdzv_id=999" in command
-            assert "--rdzv_endpoint=master:1234" in command
+            assert "--node-rank=1" in command
+            assert "--nproc-per-node=4" in command
+            assert "--rdzv-id=999" in command
+            assert "--rdzv-endpoint=master:1234" in command
 
             # Verify training arguments
             assert "--model-name-or-path=my-model" in command