From 32c990848a37274f0d4f0925e46d6230efe5cb03 Mon Sep 17 00:00:00 2001
From: Mihir Patel <mihir.v.patel7@gmail.com>
Date: Tue, 20 Aug 2024 22:53:41 -0400
Subject: [PATCH] add triansubset numbatches

---
 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 8fac739544..6e5ca9726a 100644
--- a/llmfoundry/command_utils/train.py
+++ b/llmfoundry/command_utils/train.py
@@ -512,6 +512,7 @@ def train(cfg: DictConfig) -> Trainer:
         seed=seed,
         model=model,
         train_dataloader=train_loader,
+        train_subset_num_batches=train_cfg.train_subset_num_batches,
         eval_dataloader=evaluators,
         optimizers=optimizer,
         schedulers=scheduler,
diff --git a/llmfoundry/utils/config_utils.py b/llmfoundry/utils/config_utils.py
index 1d6303a2cf..f9ba311325 100644
--- a/llmfoundry/utils/config_utils.py
+++ b/llmfoundry/utils/config_utils.py
@@ -162,6 +162,7 @@ class TrainConfig:
     only_composer_checkpoint: bool = False
 
     # Dataloader
+    train_subset_num_batches: int = -1
     device_train_microbatch_size: Union[str, int, float] = 'auto'
     global_train_batch_size: Optional[int] = None
     spin_dataloaders: bool = True