From e24439ca3e0bfd371263fd6e6c54231cf0557eb7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Piotr=20=C5=BBelasko?= Date: Mon, 29 Apr 2024 15:27:36 -0400 Subject: [PATCH 1/3] Option to auto-set expandable_segments in PyTorch CUDA allocator MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Piotr Żelasko --- .../common/data/lhotse/dataloader.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/nemo/collections/common/data/lhotse/dataloader.py b/nemo/collections/common/data/lhotse/dataloader.py index 5bb3bf2988ea..53469f9e2fc0 100644 --- a/nemo/collections/common/data/lhotse/dataloader.py +++ b/nemo/collections/common/data/lhotse/dataloader.py @@ -74,6 +74,7 @@ class LhotseDataLoadingConfig: drop_last: bool = False shard_seed: int | str = "trng" max_open_streams: int | None = None + cuda_expandable_segments: bool = True # 2.1 Multimodal sampling override options use_multimodal_sampling: bool = False @@ -147,6 +148,8 @@ def get_lhotse_dataloader_from_config( """ logging.info("We will be using a Lhotse DataLoader.") + maybe_set_cuda_expandable_segments(enabled=config.cuda_expandable_segments) + config = make_structured_with_schema_warnings(config) # First, resolve the random seed in case a string value was provided. @@ -443,3 +446,17 @@ def _flatten_alt_text(cut) -> list: text_instance.custom = {"text": data.pop("text"), "lang": data.pop("lang"), **data} ans.append(text_instance) return ans + + +def maybe_set_cuda_expandable_segments(enabled: bool): + """ + Configures PyTorch memory allocator to expand existing allocated segments + instead of re-allocating them when tensor shape grows. + This can help speed up the training when sequence length and/or batch size change often, + and makes GPU more robust towards OOM. + + See here for more details: + https://pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf + """ + if enabled and torch.cuda.is_available(): + torch.cuda.memory._set_allocator_settings("expandable_segments:True") From 1c05fe029534408760f249398bbdea4378d8b7ba Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Piotr=20=C5=BBelasko?= Date: Mon, 29 Apr 2024 15:30:28 -0400 Subject: [PATCH 2/3] warning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Piotr Żelasko --- nemo/collections/common/data/lhotse/dataloader.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/nemo/collections/common/data/lhotse/dataloader.py b/nemo/collections/common/data/lhotse/dataloader.py index 53469f9e2fc0..f3d4efbd2d75 100644 --- a/nemo/collections/common/data/lhotse/dataloader.py +++ b/nemo/collections/common/data/lhotse/dataloader.py @@ -11,7 +11,7 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - +import os import warnings from dataclasses import dataclass from functools import partial @@ -459,4 +459,12 @@ def maybe_set_cuda_expandable_segments(enabled: bool): https://pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf """ if enabled and torch.cuda.is_available(): + if ( + (value := os.environ.get("PYTORCH_CUDA_ALLOC_CONF")) is not None + and len(value) > 0 + and "expandable_segments:True" not in value + ): + warnings.warn( + "You have set PYTORCH_CUDA_ALLOC_CONF without expandable_segments:True option. We're setting that option anyway. To disable it, set cuda_expandable_segments=False in NeMo dataloader configuration." + ) torch.cuda.memory._set_allocator_settings("expandable_segments:True") From 58cfa7ea8f111d659637f7837693d59b0b8f2bec Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Piotr=20=C5=BBelasko?= Date: Tue, 30 Apr 2024 12:55:31 -0400 Subject: [PATCH 3/3] set opts after parsing config MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Piotr Żelasko --- nemo/collections/common/data/lhotse/dataloader.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nemo/collections/common/data/lhotse/dataloader.py b/nemo/collections/common/data/lhotse/dataloader.py index f3d4efbd2d75..40ea41147b25 100644 --- a/nemo/collections/common/data/lhotse/dataloader.py +++ b/nemo/collections/common/data/lhotse/dataloader.py @@ -148,10 +148,10 @@ def get_lhotse_dataloader_from_config( """ logging.info("We will be using a Lhotse DataLoader.") - maybe_set_cuda_expandable_segments(enabled=config.cuda_expandable_segments) - config = make_structured_with_schema_warnings(config) + maybe_set_cuda_expandable_segments(enabled=config.cuda_expandable_segments) + # First, resolve the random seed in case a string value was provided. seed = resolve_seed(config.seed) fix_random_seed(seed)