Add train subset num batches (#1472)

mosaicml · Aug 21, 2024 · 3648cef · 3648cef
1 parent b32e48c
commit 3648cef
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git 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
@@ -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