add lazy load option

mosaicml · Oct 14, 2023 · c1d00e9 · c1d00e9
1 parent 4e7c628
commit c1d00e9
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/scripts/train/train.py b/scripts/train/train.py
@@ -209,6 +209,11 @@ def main(cfg: DictConfig) -> Trainer:
     if max_split_size_mb is not None:
         os.environ[
             'PYTORCH_CUDA_ALLOC_CONF'] = f'max_split_size_mb:{max_split_size_mb}'
+
+    # Set CUDA lazy loading
+    cuda_load_lazy: bool = cfg.pop('cuda_load_lazy', True)
+    if cuda_load_lazy:
+        os.environ['CUDA_MODULE_LOADING'] = 'LAZY'
 
     # Set seed first
     seed: int = pop_config(cfg, 'seed', must_exist=True)