From d26f2dcbdd0b7d5c13996ed89b96f48b6987f064 Mon Sep 17 00:00:00 2001 From: Rahul Parundekar Date: Sun, 5 May 2024 12:08:31 -0700 Subject: [PATCH] load distributed config --- src/aihero/research/finetuning/train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aihero/research/finetuning/train.py b/src/aihero/research/finetuning/train.py index 755311f..6c3fc87 100644 --- a/src/aihero/research/finetuning/train.py +++ b/src/aihero/research/finetuning/train.py @@ -136,7 +136,7 @@ def fetch_dataset(self) -> DatasetDict: raise Exception("Data Abort") else: print(f"LOCAL RANK {self.local_rank}: Data ready") - + print(f"LOCAL RANK {self.local_rank}: loading data") try: splits = {} bos_token = self.tokenizer.bos_token @@ -273,7 +273,7 @@ def fetch_dataset(self) -> DatasetDict: except: # pylint: disable=bare-except # noqa: E722 print("Unable to create test dataset") - print("Loaded dataset") + print(f"LOCAL RANK {self.local_rank}: data loaded") with open(f"{DATASET_DIR}/data_ready.txt", "w") as f: f.write("Data Ready") return DatasetDict(splits)