Merge pull request #13 from szaher/torchrun-args-from-env

feat(traininghub): Use torchrun environment variables for default configuration

Author: RobotSail

src/training_hub/algorithms/osft.py

@@ -1,10 +1,10 @@
 import os
-from typing import get_origin, get_args, Union
+from typing import Literal, get_origin, get_args, Union
 from dataclasses import fields
 
 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):
@@ -58,11 +58,13 @@ def train(
         data_output_dir: str | None = None,
 
         # Torchrun parameters for multi-node support
-        nproc_per_node: int | None = None,
+        nproc_per_node: Literal['auto', 'gpu'] | int | None = None,
         nnodes: int | None = None,
         node_rank: int | None = None,
-        rdzv_id: int | None = None,
+        rdzv_id: str | int | None = None,
         rdzv_endpoint: str | None = None,
+        master_addr: str | None = None,
+        master_port: int | None = None,
         **kwargs,
     ) -> any:
         """
@@ -121,11 +123,14 @@ def train(
                 Directory where outputs from data processing will be saved such as intermediate
                 files. When not provided, it defaults to `_internal_data_processing` under the
                 `ckpt_output_dir`.
-            nproc_per_node (int): Number of processes (GPUs) per node for distributed training.
+            nproc_per_node (Literal['auto', 'gpu'] | int): Number of processes (GPUs) per node for distributed training.
             nnodes (int): Total number of nodes for distributed training.
             node_rank (int): Rank of this node (0 to nnodes-1) for distributed training. 
-            rdzv_id (int): Unique job ID for rendezvous in distributed training.
+            rdzv_id (str | int): 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 (int): 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
@@ -222,11 +229,13 @@ def get_optional_params(self) -> dict[str, type]:
             'use_processed_dataset': bool,
             'unmask_messages': bool,
             'data_output_dir': str,
-            'nproc_per_node': int,
+            'nproc_per_node': Literal['auto', 'gpu'] | int,
             'nnodes': int,
             'node_rank': int,
-            'rdzv_id': int,
+            'rdzv_id': str | int,
             'rdzv_endpoint': str,
+            'master_addr': str,
+            'master_port': int,
         }
 
     def _validate_param_types(self, params: dict[str, any]):
@@ -333,6 +342,16 @@ def execute_training(self, algorithm_params: dict[str, any]) -> any:
         # Rename parameters before sending to backend
         algorithm_params = {renames.get(k, k): v for k, v in algorithm_params.items()}
 
+        # Separate parameters into their respective dataclass fields
+        torchrun_args_fields = {f.name for f in fields(TorchrunArgs)}
+        training_args_fields = {f.name for f in fields(TrainingArgs)}
+
+
+        # process this up here so we can exit early
+        torchrun_args_pre = {k: v for k, v in algorithm_params.items() if k in torchrun_args_fields and v is not None}
+        torchrun_args_pre = get_torchrun_params(torchrun_args_pre)
+        torch_args = TorchrunArgs(**torchrun_args_pre)
+
         # We separate this from `ckpt_output_dir` so that we can use `/dev/shm` for low-latency data
         # proceessing. But we do not want to make assumptions about the size of training data or the
         # amount of memory on the host. So by default we write to storage, but expose this as a separate
@@ -353,11 +372,6 @@ def execute_training(self, algorithm_params: dict[str, any]) -> any:
             unmask_messages=algorithm_params.get('unmask_messages', False),
         )
 
-
-        # Separate parameters into their respective dataclass fields
-        torchrun_args_fields = {f.name for f in fields(TorchrunArgs)}
-        training_args_fields = {f.name for f in fields(TrainingArgs)}
-
         # adjust arguments to align with the API definition 
         training_args_pre = {k: v for k, v in algorithm_params.items() if k in training_args_fields and v is not None}
         training_args_pre['data_path'] = training_ready_data_path  # replaces raw data path with processed
@@ -372,14 +386,9 @@ def execute_training(self, algorithm_params: dict[str, any]) -> any:
         # but default it to True
         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')
-
-
         # now we run training
         return run_training(
-            torch_args=TorchrunArgs(**torchrun_args_pre),
+            torch_args=torch_args,
             train_args=TrainingArgs(**training_args_pre),
         )
 
@@ -460,11 +469,13 @@ def osft(
     save_final_checkpoint: bool | None = None,
     num_epochs: int | None = None,
     # Torchrun parameters for multi-node support
-    nproc_per_node: int | None = None,
+    nproc_per_node: Literal['auto', 'gpu'] | int | None = None,
     nnodes: int | None = None,
     node_rank: int | None = None,
-    rdzv_id: int | None = None,
+    rdzv_id: str | int | None = None,
     rdzv_endpoint: str | None = None,
+    master_port: int | None = None,
+    master_addr: str | None = None,
     **kwargs
 ) -> any:
     from . import create_algorithm
@@ -496,5 +507,7 @@ def osft(
         node_rank=node_rank,
         rdzv_id=rdzv_id,
         rdzv_endpoint=rdzv_endpoint,
+        master_port=master_port,
+        master_addr=master_addr,
         **kwargs
     )

src/training_hub/algorithms/sft.py

@@ -2,6 +2,7 @@
 from instructlab.training import run_training, TorchrunArgs, TrainingArgs
 
 from . import Algorithm, Backend, AlgorithmRegistry
+from training_hub import utils
 
 
 class InstructLabTrainingSFTBackend(Backend):
@@ -10,7 +11,7 @@ class InstructLabTrainingSFTBackend(Backend):
     def execute_training(self, algorithm_params: Dict[str, Any]) -> Any:
         """Execute SFT training using instructlab-training."""
         # Separate torchrun parameters from training parameters
-        torchrun_keys = {'nproc_per_node', 'nnodes', 'node_rank', 'rdzv_id', 'rdzv_endpoint'}
+        torchrun_keys = {'nproc_per_node', 'nnodes', 'node_rank', 'rdzv_id', 'rdzv_endpoint', 'master_addr', 'master_port'}
 
         # Extract torchrun parameters
         torchrun_params = {k: v for k, v in algorithm_params.items() if k in torchrun_keys}
@@ -26,26 +27,9 @@ 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)
-        if torchrun_params:
-            torchrun_defaults = {
-                'nnodes': 1,
-                'node_rank': 0,
-                'rdzv_id': 0,
-                'rdzv_endpoint': ""
-            }
-            # 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(
-                nproc_per_node=1,
-                nnodes=1, 
-                node_rank=0,
-                rdzv_id=0,
-                rdzv_endpoint=""
-            )
-        
+        final_torchrun_params = utils.get_torchrun_params(torchrun_params)
+        torchrun_args = TorchrunArgs(**final_torchrun_params)
+
         # Execute training
         return run_training(
             torch_args=torchrun_args,
@@ -76,11 +60,13 @@ def train(self,
               accelerate_full_state_at_epoch: Optional[bool] = None,
               checkpoint_at_epoch: Optional[bool] = None,
               # Torchrun parameters for multi-node support
-              nproc_per_node: Optional[int] = None,
+              nproc_per_node: Optional[str | int] = None,
               nnodes: Optional[int] = None,
               node_rank: Optional[int] = None,
-              rdzv_id: Optional[int] = None,
+              rdzv_id: Optional[str | int] = None,
               rdzv_endpoint: Optional[str] = None,
+              master_addr: Optional[str] = None,
+              master_port: Optional[int] = None,
               **kwargs) -> Any:
         """Execute SFT training.
         
@@ -103,6 +89,8 @@ def train(self,
             node_rank: Rank of this node (0 to nnodes-1)
             rdzv_id: Unique job ID for rendezvous
             rdzv_endpoint: Master node endpoint for multi-node training
+            master_addr: Master node address for distributed training
+            master_port: Master node port for distributed training
             **kwargs: Additional parameters passed to the backend
             
         Returns:
@@ -128,6 +116,8 @@ def train(self,
             'node_rank': node_rank,
             'rdzv_id': rdzv_id,
             'rdzv_endpoint': rdzv_endpoint,
+            'master_addr': master_addr,
+            'master_port': master_port,
         }
 
         # Only add non-None parameters (let TrainingArgs handle defaults)
@@ -161,11 +151,13 @@ def get_optional_params(self) -> Dict[str, Type]:
             'warmup_steps': int,
             'accelerate_full_state_at_epoch': bool,
             'checkpoint_at_epoch': bool,
-            'nproc_per_node': int,
+            'nproc_per_node': str | int,
             'nnodes': int,
             'node_rank': int,
-            'rdzv_id': int,
+            'rdzv_id': str | int,
             'rdzv_endpoint': str,
+            'master_addr': str,
+            'master_port': int,
         }
 
 
@@ -190,11 +182,13 @@ def sft(model_path: str,
         accelerate_full_state_at_epoch: Optional[bool] = None,
         checkpoint_at_epoch: Optional[bool] = None,
         # Torchrun parameters for multi-node support
-        nproc_per_node: Optional[int] = None,
+        nproc_per_node: Optional[str | int] = None,
         nnodes: Optional[int] = None,
         node_rank: Optional[int] = None,
-        rdzv_id: Optional[int] = None,
+        rdzv_id: Optional[str | int] = None,
         rdzv_endpoint: Optional[str] = None,
+        master_addr: Optional[str] = None,
+        master_port: Optional[int] = None,
         **kwargs) -> Any:
     """Convenience function to run SFT training.
     
@@ -218,6 +212,9 @@ def sft(model_path: str,
         node_rank: Rank of this node (0 to nnodes-1) for distributed training
         rdzv_id: Unique job ID for rendezvous in distributed training
         rdzv_endpoint: Master node endpoint for multi-node training
+        master_addr: Master node address for distributed training
+        master_port: Master node port for distributed training
+
         **kwargs: Additional parameters passed to the backend
     
     Returns:
@@ -245,6 +242,8 @@ def sft(model_path: str,
         node_rank=node_rank,
         rdzv_id=rdzv_id,
         rdzv_endpoint=rdzv_endpoint,
+        master_addr=master_addr,
+        master_port=master_port,
         **kwargs
     )