Add expandeable segments flag (#1075)

mosaicml · Apr 1, 2024 · 349c2ff · 349c2ff
1 parent 7a8a156
commit 349c2ff
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/scripts/train/train.py b/scripts/train/train.py
@@ -130,11 +130,18 @@ def main(cfg: DictConfig) -> Trainer:
     # Create copy of config for logging
     logged_cfg: DictConfig = copy.deepcopy(cfg)
 
+    cuda_alloc_conf = []
     # Get max split size mb
     max_split_size_mb: Optional[int] = cfg.pop('max_split_size_mb', None)
     if max_split_size_mb is not None:
-        os.environ[
-            'PYTORCH_CUDA_ALLOC_CONF'] = f'max_split_size_mb:{max_split_size_mb}'
+        cuda_alloc_conf.append(f'max_split_size_mb:{max_split_size_mb}')
+
+    # Expandeable segments
+    if cfg.pop('expandeable_segments', False):
+        cuda_alloc_conf.append('expandeable_segments:True')
+
+    if len(cuda_alloc_conf) > 0:
+        os.environ['PYTORCH_CUDA_ALLOC_CONF'] = ','.join(cuda_alloc_conf)
 
     # Set CUDA lazy loading
     # This can save a bit of memory if not all modules are needed

diff --git a/scripts/train/yamls/finetune/dbrx-full-ft.yaml b/scripts/train/yamls/finetune/dbrx-full-ft.yaml
@@ -88,6 +88,7 @@ device_eval_batch_size: 1
 precision: amp_bf16
 autoresume: true
 dist_timeout: 3600
+expandeable_segments: true
 
 # FSDP
 fsdp_config:

diff --git a/scripts/train/yamls/finetune/dbrx-lora-ft.yaml b/scripts/train/yamls/finetune/dbrx-lora-ft.yaml
@@ -96,6 +96,7 @@ device_eval_batch_size: 1
 precision: amp_bf16
 autoresume: true
 dist_timeout: 3600
+expandeable_segments: true
 
 # FSDP
 fsdp_config: