validate torchrunargs for backends

Signed-off-by: Saad Zaher <[email protected]>

Author: szaher

src/training_hub/algorithms/osft.py

@@ -4,7 +4,7 @@
 
 import datasets
 from training_hub.algorithms import Algorithm, Backend, AlgorithmRegistry
-from training_hub.utils import format_type_name
+from training_hub.utils import format_type_name, get_torchrun_params
 
 
 class OSFTAlgorithm(Algorithm):
@@ -63,6 +63,8 @@ def train(
         node_rank: int | None = None,
         rdzv_id: str | None = None,
         rdzv_endpoint: str | None = None,
+        master_addr: str | None = None,
+        master_port: str | None = None,
         **kwargs,
     ) -> any:
         """
@@ -126,6 +128,9 @@ def train(
             node_rank (int): Rank of this node (0 to nnodes-1) for distributed training. 
             rdzv_id (str): Unique job ID for rendezvous in distributed training.
             rdzv_endpoint (str): Master node endpoint for multi-node training.
+            master_addr (str): Master node address for distributed training (only used with
+                static rdzv_backend).
+            master_port (str): Master node port for distributed training.
             **kwargs: Additional parameters passed to the backend.
 
         Returns:
@@ -176,6 +181,8 @@ def train(
             'node_rank': node_rank,
             'rdzv_id': rdzv_id,
             'rdzv_endpoint': rdzv_endpoint,
+            'master_addr': master_addr,
+            'master_port': master_port,
         }
 
         # now do validation now that we've set everything up
@@ -227,6 +234,8 @@ def get_optional_params(self) -> dict[str, type]:
             'node_rank': int,
             'rdzv_id': str,
             'rdzv_endpoint': str,
+            'master_addr': str,
+            'master_port': int,
         }
 
     def _validate_param_types(self, params: dict[str, any]):
@@ -373,8 +382,7 @@ def execute_training(self, algorithm_params: dict[str, any]) -> any:
         training_args_pre['osft'] = training_args_pre.get('osft', True)
 
         torchrun_args_pre = {k: v for k, v in algorithm_params.items() if k in torchrun_args_fields and v is not None}
-        # TODO: update this default in mini-trainer
-        torchrun_args_pre['rdzv_endpoint'] = torchrun_args_pre.get('rdzv_endpoint', 'localhost:1738')
+        torchrun_args_pre = get_torchrun_params(torchrun_params=torchrun_args_pre)
 
 
         # now we run training

src/training_hub/algorithms/sft.py

@@ -3,6 +3,7 @@
 from instructlab.training import run_training, TorchrunArgs, TrainingArgs
 
 from . import Algorithm, Backend, AlgorithmRegistry
+from training_hub import utils
 
 
 class InstructLabTrainingSFTBackend(Backend):
@@ -27,21 +28,13 @@ def execute_training(self, algorithm_params: Dict[str, Any]) -> Any:
         training_args = TrainingArgs(**training_params)
 
         # Set up torchrun arguments with single-node defaults (except nproc_per_node)
-        torchrun_defaults = {
-            'nproc_per_node': os.getenv("LOCAL_WORLD_SIZE", os.getenv("PET_NPROC_PER_NODE", "1")),
-            'nnodes': int(os.getenv("WORLD_SIZE", os.getenv("PET_NNODES", "1"))),
-            'node_rank': int(os.getenv("PET_NODE_RANK", os.getenv("RANK", "0"))),
-            'rdzv_id': 0,
-            'rdzv_endpoint': ""
-        }
+        final_torchrun_params = utils.get_torchrun_params(training_args.dict())
 
         if torchrun_params:
-            # Merge provided params with defaults
-            final_torchrun_params = {**torchrun_defaults, **torchrun_params}
             torchrun_args = TorchrunArgs(**final_torchrun_params)
         else:
             # Use single-node defaults including nproc_per_node
-            torchrun_args = TorchrunArgs(**torchrun_defaults)
+            torchrun_args = TorchrunArgs(**final_torchrun_params)
 
         # Execute training
         return run_training(

src/training_hub/utils.py

@@ -1,3 +1,6 @@
+import os
+from curses.ascii import isdigit
+from importlib.metadata import pass_none
 from typing import get_origin, get_args
 
 def format_type_name(tp):
@@ -26,3 +29,73 @@ def format_type_name(tp):
         return type_str[8:-2]
 
     return type_str
+
+
+def get_torchrun_params(args: dict):
+    """
+    Parse and load PyTorch variables from dict with fallback to environment variables.
+
+    Args:
+        args (dict): Dictionary containing PyTorch configuration parameters
+
+    Returns:
+        dict: Dictionary with PyTorch parameters loaded from args or environment
+    """
+    pytorch_vars = ['nproc_per_node', 'nnodes', 'node_rank', 'rdzv_id', 'rdzv_endpoint', 'master_addr', 'master_port']
+    torchrun_args = {}
+
+    def validate_nproc_per_node(value):
+        """Validate and convert nproc_per_node value."""
+        if isinstance(value, str):
+            if value.lower() == 'auto':
+                return 'gpu'
+            elif value.lower() == 'gpu':
+                return 'gpu'
+            else:
+                try:
+                    return int(value)
+                except ValueError:
+                    raise ValueError(f"nproc_per_node must be 'auto', 'gpu', or an integer, got: {value}")
+        elif isinstance(value, int):
+            return value
+        else:
+            raise ValueError(f"nproc_per_node must be 'auto', 'gpu', or an integer, got: {value}")
+
+    def get_env_var_name(var_name):
+        """Get environment variable name based on PyTorch convention."""
+        return var_name.upper() if var_name in ['master_addr', 'master_port'] else f"PET_{var_name.upper()}"
+
+    for var_name in pytorch_vars:
+        # Try args dict first
+        if var_name in args and args[var_name] is not None and args[var_name] != "":
+            value = args[var_name]
+            if var_name == 'nproc_per_node':
+                torchrun_args[var_name] = validate_nproc_per_node(value)
+            elif var_name in ['nnodes', 'node_rank', 'rdzv_id', 'master_port']:
+                torchrun_args[var_name] = int(value) if isinstance(value, (str, int)) else value
+            else:
+                torchrun_args[var_name] = value
+        else:
+            # Fallback to environment variable
+            env_value = os.getenv(get_env_var_name(var_name))
+            if env_value is not None:
+                if var_name == 'nproc_per_node':
+                    torchrun_args[var_name] = validate_nproc_per_node(env_value)
+                elif var_name in ['nnodes', 'node_rank', 'rdzv_id', 'master_port']:
+                    try:
+                        torchrun_args[var_name] = int(env_value)
+                    except ValueError:
+                        torchrun_args[var_name] = env_value
+                else:
+                    torchrun_args[var_name] = env_value
+            else:
+                # Set defaults
+                defaults = {'nnodes': 1, 'rdzv_id': 0}
+                torchrun_args[var_name] = defaults.get(var_name, "")
+
+    # Validate mutually exclusive parameters
+    if (torchrun_args.get('rdzv_endpoint', '') != "" and
+        (torchrun_args.get('master_addr', '') != "" or torchrun_args.get('master_port', '') != "")):
+        raise ValueError("Cannot specify both rdzv_endpoint and master_addr/master_port. These are mutually exclusive parameters.")
+
+    return torchrun_args