Add Apex import Guard for MM collection

nemo/collections/multimodal/models/vision_language_foundation/megatron_nsfw_clip_models.py

@@ -19,7 +19,14 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from apex.transformer.pipeline_parallel.utils import get_num_microbatches
+
+try:
+    from apex.transformer.pipeline_parallel.utils import get_num_microbatches
+
+    HAVE_APEX = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_APEX = False
+
 from megatron.core import parallel_state
 from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
 from omegaconf.dictconfig import DictConfig

nemo/collections/multimodal/modules/imagen/diffusionmodules/layers.py

@@ -43,8 +43,13 @@
 import torch as th
 import torch.nn as nn
 import torch.nn.functional as F
-from apex.contrib.group_norm import GroupNorm
 
+try:
+    from apex.contrib.group_norm import GroupNorm
+
+    HAVE_APEX = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_APEX = False
 
 def conv_nd(dims, *args, **kwargs):
     """

nemo/collections/multimodal/modules/stable_diffusion/attention.py

@@ -16,7 +16,14 @@
 
 import torch
 import torch.nn.functional as F
-from apex.contrib.group_norm import GroupNorm
+
+try:
+    from apex.contrib.group_norm import GroupNorm
+
+    HAVE_APEX = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_APEX = False
+
 from einops import rearrange, repeat
 from torch import einsum, nn
 from torch._dynamo import disable
@@ -142,6 +149,11 @@ def forward(self, x):
 
 class SpatialSelfAttention(nn.Module):
     def __init__(self, in_channels):
+        if not HAVE_APEX:
+            raise ImportError(
+                "Apex was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+            )
+
         super().__init__()
         self.in_channels = in_channels
 
@@ -432,6 +444,11 @@ def __init__(
         use_flash_attention=False,
         lora_network_alpha=None,
     ):
+        if not HAVE_APEX:
+            raise ImportError(
+                "Apex was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+            )
+
         super().__init__()
         logging.info(
             f"constructing {self.__class__.__name__} of depth {depth} w/ {in_channels} channels and {n_heads} heads"

nemo/collections/multimodal/modules/stable_diffusion/diffusionmodules/model.py

@@ -17,7 +17,14 @@
 import numpy as np
 import torch
 import torch.nn as nn
-from apex.contrib.group_norm import GroupNorm
+
+try:
+    from apex.contrib.group_norm import GroupNorm
+
+    HAVE_APEX = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_APEX = False
+
 from einops import rearrange
 
 from nemo.collections.multimodal.modules.stable_diffusion.attention import LinearAttention

nemo/collections/multimodal/modules/stable_diffusion/diffusionmodules/openaimodel.py

@@ -22,7 +22,13 @@
 import torch as th
 import torch.nn as nn
 import torch.nn.functional as F
-from apex.contrib.group_norm import GroupNorm
+
+try:
+    from apex.contrib.group_norm import GroupNorm
+
+    HAVE_APEX = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_APEX = False
 
 from nemo.collections.multimodal.modules.stable_diffusion.attention import SpatialTransformer
 from nemo.collections.multimodal.modules.stable_diffusion.diffusionmodules.util import (
@@ -554,6 +560,11 @@ def __init__(
         lora_network_alpha=None,
         timesteps=1000,
     ):
+        if not HAVE_APEX:
+            raise ImportError(
+                "Apex was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+            )
+
         super().__init__()
         from omegaconf.listconfig import ListConfig
 
@@ -1229,6 +1240,11 @@ def __init__(
         *args,
         **kwargs,
     ):
+        if not HAVE_APEX:
+            raise ImportError(
+                "Apex was not found. Please see the NeMo README for installation instructions: https://github.com/NVIDIA/NeMo#megatron-gpt."
+            )
+
         super().__init__()
         if num_heads_upsample == -1:
             num_heads_upsample = num_heads

nemo/collections/multimodal/modules/stable_diffusion/diffusionmodules/util.py

@@ -29,7 +29,14 @@
 import numpy as np
 import torch
 import torch.nn as nn
-from apex.contrib.group_norm import GroupNorm
+
+try:
+    from apex.contrib.group_norm import GroupNorm
+
+    HAVE_APEX = True
+except (ImportError, ModuleNotFoundError):
+    HAVE_APEX = False
+
 from einops import repeat
 from torch._dynamo import disable
 from torch.cuda.amp import custom_bwd, custom_fwd