load distributed config

launch.py

@@ -1,20 +1,14 @@
 """run script for fine-tuning a model."""
-import yaml  # type: ignore
 from aihero.research.config.schema import BatchInferenceJob, TrainingJob
 from aihero.research.finetuning.infer import BatchInferenceJobRunner
 from aihero.research.finetuning.train import TrainingJobRunner
 from fire import Fire
 
 
-def train(training_config_file: str = "/mnt/config/training/config.yaml", distributed_config_file: str = "") -> None:
+def train(training_config_file: str = "/mnt/config/training/config.yaml") -> None:
     """Run Training."""
     training_config = TrainingJob.load(training_config_file)
-    if distributed_config_file:
-        with open(distributed_config_file) as f:
-            distributed_training_config = yaml.safe_load(f)
-    else:
-        distributed_training_config = ""
-    TrainingJobRunner(training_config, distributed_config=distributed_training_config).run()
+    TrainingJobRunner(training_config).run()
 
 
 def infer(batch_inference_config_file: str = "/mnt/config/batch_inference/config.yaml") -> None:

src/aihero/research/finetuning/train.py

@@ -1,6 +1,6 @@
 """Launch the training job inside a container."""
 import os
-from typing import Any, Optional, Tuple
+from typing import Any, Tuple
 
 import torch
 from datasets import Dataset, DatasetDict
@@ -25,14 +25,9 @@
 class TrainingJobRunner:
     """Class to run a training job."""
 
-    def __init__(self, training_job: TrainingJob, distributed_config: Optional[dict[str, Any]] = None):
+    def __init__(self, training_job: TrainingJob):
         """Initialize the training job runner."""
         self.training_job = training_job
-        if distributed_config:
-            self.distributed_config = distributed_config
-            print("Training LOCAL RANK: {} ...".format(os.getenv("LOCAL_RANK", "Unknown")))
-            print("Training RANK: {} ...".format(os.getenv("RANK", "Unknown")))
-            print("Training LOCAL WORLD SIZE: {} ...".format(os.getenv("LOCAL_WORLD_SIZE", "Unknown")))
 
         print("Loading model")
         self.model, self.tokenizer = self.load_model()