From 5c7e99bed1dfe28ed4fe61a30c15a2b2d4527b8c Mon Sep 17 00:00:00 2001
From: Daniel King <43149077+dakinggg@users.noreply.github.com>
Date: Sun, 28 Jul 2024 18:39:24 -0700
Subject: [PATCH] Add spin_dataloaders flag (#1405)

---
 llmfoundry/command_utils/train.py | 1 +
 llmfoundry/utils/config_utils.py  | 1 +
 2 files changed, 2 insertions(+)

diff --git a/llmfoundry/command_utils/train.py b/llmfoundry/command_utils/train.py
index 77bb9dbcfe..c925e6e586 100644
--- a/llmfoundry/command_utils/train.py
+++ b/llmfoundry/command_utils/train.py
@@ -544,6 +544,7 @@ def train(cfg: DictConfig) -> Trainer:
         dist_timeout=train_cfg.dist_timeout,
         profiler=profiler,
         compile_config=compile_config,
+        spin_dataloaders=train_cfg.spin_dataloaders,
     )
 
     # Optionally just save an HF checkpoint
diff --git a/llmfoundry/utils/config_utils.py b/llmfoundry/utils/config_utils.py
index dcb97eb0de..84a3376718 100644
--- a/llmfoundry/utils/config_utils.py
+++ b/llmfoundry/utils/config_utils.py
@@ -167,6 +167,7 @@ class TrainConfig:
     # Dataloader
     device_train_microbatch_size: Union[str, int, float] = 'auto'
     global_train_batch_size: Optional[int] = None
+    spin_dataloaders: bool = True
 
     # Eval dataloader
     eval_subset_num_batches: int = -1