From c1d00e960d943137842a202e32144a753e19cd04 Mon Sep 17 00:00:00 2001 From: Daniel King Date: Sat, 14 Oct 2023 14:07:53 -0700 Subject: [PATCH] add lazy load option --- scripts/train/train.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/scripts/train/train.py b/scripts/train/train.py index 7358d58d2e..64edaa1ba0 100644 --- 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)