use str | int for nproc_per_node and rdzv_id

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

Author: szaher

src/training_hub/algorithms/osft.py

@@ -58,7 +58,7 @@ def train(
         data_output_dir: str | None = None,
 
         # Torchrun parameters for multi-node support
-        nproc_per_node: str | None = None,
+        nproc_per_node: str | int | None = None,
         nnodes: int | None = None,
         node_rank: int | None = None,
         rdzv_id: str | None = None,
@@ -222,7 +222,7 @@ def get_optional_params(self) -> dict[str, type]:
             'use_processed_dataset': bool,
             'unmask_messages': bool,
             'data_output_dir': str,
-            'nproc_per_node': str,
+            'nproc_per_node': str | int,
             'nnodes': int,
             'node_rank': int,
             'rdzv_id': str,
@@ -460,7 +460,7 @@ def osft(
     save_final_checkpoint: bool | None = None,
     num_epochs: int | None = None,
     # Torchrun parameters for multi-node support
-    nproc_per_node: str | None = None,
+    nproc_per_node: str | int | None = None,
     nnodes: int | None = None,
     node_rank: int | None = None,
     rdzv_id: str | None = None,

src/training_hub/algorithms/sft.py

@@ -73,10 +73,10 @@ 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,
               **kwargs) -> Any:
         """Execute SFT training.
@@ -158,10 +158,10 @@ 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,
         }
 
@@ -187,10 +187,10 @@ 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,
         **kwargs) -> Any:
     """Convenience function to run SFT training.