fix(torchrun): Omit empty arguments and correct nproc_per_node type (#661)

* fix(torchrun): Omit empty arguments and correct nproc_per_node type

The command generation logic is updated to dynamically
build the torchrun command, excluding arguments that
are empty or None. This prevents them from overriding
environment variables, ensuring that torchrun can
correctly inherit its configuration. An exception is
made for integer arguments where 0 is a valid value.

Additionally, the nproc_per_node argument type has been
changed from int to str to support special values
accepted by PyTorch, such as 'auto', 'gpu', and 'cpu'.

Reference: https://github.com/pytorch/pytorch/blob/main/torch/distributed/run.py#L77-L88

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

* only dynamically add torchrun args & change rdzv_id type to str

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

* fix smoke tests

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

* Enable both dtypes str, int for nproc_per_node, rdzv_id

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

* Use python3.11 style for pydatnic model

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

* add all torchrun args and validate them

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

* Remove non-required dependencies

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

* update datatypes only

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

* replace _ with - when passing torchrun args

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

* make nproc_per_node to only accept gpu or int

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

* add master_{addr, port} validate args

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

* check for not set or empty rdzv endpoint

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

* fix formatting error

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

* Update src/instructlab/training/config.py

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

* Update tests/smoke/test_train.py

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

* Update src/instructlab/training/main_ds.py

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

* fixes indentation

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

* formatting

* add standalone as the fallback when neither master_addr nor rdzv_endpoint are provided

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

* clarify rdzv-backend arg

---------

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

Author: szaher

src/instructlab/training/config.py

@@ -9,7 +9,7 @@
 from typing import List, Literal, Optional
 
 # Third Party
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, model_validator
 
 
 # public API
@@ -70,16 +70,33 @@ class DataProcessArgs(BaseModel):
 # public API
 class TorchrunArgs(BaseModel):
     """
-    Representation of the arguments being used by torchrun.
-    The full list of arguments can be found here:
-    https://pytorch.org/docs/stable/elastic/run.html#definitions
+    Arguments for torchrun (https://pytorch.org/docs/stable/elastic/run.html#definitions)
+
+    Precedence order: arg > env > defaults
+    Ensures that either `rdzv_endpoint` OR both `master_addr` and `master_port`
+    are provided, but not both.
     """
 
-    nproc_per_node: int
+    # Core distributed training arguments
+    nproc_per_node: Literal["gpu"] | int
     nnodes: int
     node_rank: int
-    rdzv_id: int
-    rdzv_endpoint: str
+    rdzv_id: str | int
+
+    # Rendezvous / master configuration
+    rdzv_endpoint: Optional[str] = None
+    master_addr: Optional[str] = None
+    master_port: Optional[int] = None
+
+    model_config = ConfigDict(extra="ignore")
+
+    @model_validator(mode="after")
+    def validate_endpoint_config(self):
+        if self.rdzv_endpoint and self.master_addr:
+            raise ValueError(
+                "Provide either `rdzv_endpoint` OR both `master_addr` and `master_port`, not both."
+            )
+        return self
 
 
 # public API

src/instructlab/training/main_ds.py

@@ -480,27 +480,54 @@ def run_training(torch_args: TorchrunArgs, train_args: TrainingArgs) -> None:
     if not os.path.exists(train_args.ckpt_output_dir):
         os.makedirs(train_args.ckpt_output_dir, exist_ok=True)
 
+    # build distributed training command
     command = [
         "torchrun",
+        f"--nproc-per-node={torch_args.nproc_per_node}",
         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}",
-        __file__,
-        f"--model_name_or_path={train_args.model_path}",
-        f"--data_path={train_args.data_output_dir}/data.jsonl",
-        f"--output_dir={train_args.ckpt_output_dir}",
-        f"--num_epochs={train_args.num_epochs}",
-        f"--effective_batch_size={train_args.effective_batch_size}",
-        f"--learning_rate={train_args.learning_rate}",
-        f"--num_warmup_steps={train_args.warmup_steps}",
-        f"--save_samples={train_args.save_samples}",
-        f"--log_level={train_args.log_level}",
-        f"--max_batch_len={train_args.max_batch_len}",
-        f"--seed={train_args.random_seed}",
+        f"--node-rank={torch_args.node_rank}",
+        f"--rdzv-id={torch_args.rdzv_id}",
     ]
 
+    # validation should have already caught the mutually exclusive case earlier, but here we check
+    # anyway just to be extra sure since Python is not type-safe and validation can bypassed (e.g. during testing)
+    if torch_args.master_addr and torch_args.rdzv_endpoint:
+        raise ValueError(
+            "`torch_args.master_addr` and `torch_args.rdzv_endpoint` cannot be passed at the same time; please pass only one"
+        )
+
+    if torch_args.master_addr:
+        command += [
+            f"--master-addr={torch_args.master_addr}",
+            "--rdzv-backend=static",
+        ]
+        command += (
+            [f"--master-port={torch_args.master_port}"]
+            if torch_args.master_port
+            else []
+        )
+    elif torch_args.rdzv_endpoint:
+        command += [f"--rdzv-endpoint={torch_args.rdzv_endpoint}"]
+    else:
+        command += ["--standalone"]
+
+    command.extend(
+        [
+            __file__,
+            f"--model_name_or_path={train_args.model_path}",
+            f"--data_path={train_args.data_output_dir}/data.jsonl",
+            f"--output_dir={train_args.ckpt_output_dir}",
+            f"--num_epochs={train_args.num_epochs}",
+            f"--effective_batch_size={train_args.effective_batch_size}",
+            f"--learning_rate={train_args.learning_rate}",
+            f"--num_warmup_steps={train_args.warmup_steps}",
+            f"--save_samples={train_args.save_samples}",
+            f"--log_level={train_args.log_level}",
+            f"--max_batch_len={train_args.max_batch_len}",
+            f"--seed={train_args.random_seed}",
+        ]
+    )
+
     if train_args.chat_tmpl_path is not None:
         command.append(f"--chat-tmpl-path={train_args.chat_tmpl_path}")