Fix os.environ.pop

Signed-off-by: Guyue Huang <[email protected]>
NVIDIA · Jan 30, 2025 · 530719a · 530719a
1 parent 83d35d5
commit 530719a
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/nemo/lightning/pytorch/callbacks/megatron_comm_overlap.py b/nemo/lightning/pytorch/callbacks/megatron_comm_overlap.py
@@ -231,7 +231,8 @@ def _set_num_cuda_device_max_connections(self):
                 """
                 os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = "32"
             else:
-                os.environ.pop('CUDA_DEVICE_MAX_CONNECTIONS')
+                if 'CUDA_DEVICE_MAX_CONNECTIONS' in os.environ:
+                    os.environ.pop('CUDA_DEVICE_MAX_CONNECTIONS')
         else:
             if tp_size > 1 or cp_size > 1:
                 """
@@ -244,7 +245,8 @@ def _set_num_cuda_device_max_connections(self):
                 """
                 os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = "1"
             else:
-                os.environ.pop('CUDA_DEVICE_MAX_CONNECTIONS')
+                if 'CUDA_DEVICE_MAX_CONNECTIONS' in os.environ:
+                    os.environ.pop('CUDA_DEVICE_MAX_CONNECTIONS')
 
     def setup(self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: str) -> None:
         assert isinstance(trainer.strategy, MegatronStrategy), "MegatronCommOverlapCallback requires MegatronStrategy"