Increase the number of dataloader workers

Signed-off-by: Fabrice Normandin <normandf@mila.quebec>
mila-iqia · Nov 8, 2024 · 159de06 · 159de06
1 parent 9b282dc
commit 159de06
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/project/algorithms/llm_finetuning_example.py b/project/algorithms/llm_finetuning_example.py
@@ -406,6 +406,7 @@ def train_dataloader(self):
             self.train_dataset,
             shuffle=True,
             collate_fn=default_data_collator,
+            num_workers=self.dataset_config.preprocessing_num_workers,
             batch_size=self.dataset_config.per_device_train_batch_size,
         )
 
@@ -414,6 +415,7 @@ def val_dataloader(self):
         return DataLoader(
             self.valid_dataset,
             collate_fn=default_data_collator,
+            num_workers=self.dataset_config.preprocessing_num_workers,
             batch_size=self.dataset_config.per_device_eval_batch_size,
         )