pass through

mosaicml · Oct 16, 2024 · 78fa9a9 · 78fa9a9
1 parent 13e0a6d
commit 78fa9a9
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
@@ -592,6 +592,7 @@ def train(cfg: DictConfig) -> Trainer:
         profiler=profiler,
         compile_config=compile_config,
         spin_dataloaders=train_cfg.spin_dataloaders,
+        accumulate_train_batch_on_tokens=train_cfg.accumulate_train_batch_on_tokens,
     )
 
     _sort_callbacks(trainer)

diff --git a/llmfoundry/utils/config_utils.py b/llmfoundry/utils/config_utils.py
@@ -121,6 +121,7 @@ class TrainConfig:
     dist_timeout: Union[int, float] = 600.0
     fsdp_config: Optional[dict[str, Any]] = None
     tp_config: Optional[dict[str, Any]] = None
+    accumulate_train_batch_on_tokens: bool = False
 
     # Evaluation parameters
     eval_interval: Union[int, str] = 1