add persistent workers (#6205) (#6211)

Signed-off-by: ericharper <[email protected]> Co-authored-by: Eric Harper <[email protected]>
NVIDIA · Mar 16, 2023 · 296bce1 · 296bce1
1 parent c3a9d63
commit 296bce1
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py
@@ -784,7 +784,11 @@ def build_pretraining_data_loader(
             raise ValueError('cfg.data.dataloader_type not found. Must be "single" or "cyclic"')
 
         return torch.utils.data.DataLoader(
-            dataset, batch_sampler=batch_sampler, num_workers=self.cfg.data.num_workers, pin_memory=True,
+            dataset,
+            batch_sampler=batch_sampler,
+            num_workers=self.cfg.data.num_workers,
+            pin_memory=True,
+            persistent_workers=True,
         )
 
     def setup(self, stage=None):

diff --git a/nemo/collections/nlp/modules/common/megatron/clip_grads.py b/nemo/collections/nlp/modules/common/megatron/clip_grads.py
@@ -17,7 +17,7 @@
 import itertools
 
 import torch
-from torch._six import inf
+from torch import inf
 
 from nemo.collections.nlp.modules.common.megatron.module import param_is_not_shared