delay imports for replace policies and fix missing req
diff --git a/deepspeed/module_inject/replace_module.py b/deepspeed/module_inject/replace_module.py
index c58130f23185..31f16cce8109 100755
--- a/deepspeed/module_inject/replace_module.py
+++ b/deepspeed/module_inject/replace_module.py
@@ -387,6 +387,9 @@ def replace_module(model, orig_class, replace_fn, _replace_policy):
         policy.update({orig_class: (replace_fn, _replace_policy)})
     else:
         for plcy in replace_policies:
+            # instantiate a throw-away policy in order to populate the _orig_layer_class
+            _ = plcy(None)
+            assert plcy._orig_layer_class != None
             policy.update({plcy._orig_layer_class: (replace_fn, plcy)})
 
     replaced_module, _ = _replace_module(model, policy)
diff --git a/deepspeed/module_inject/replace_policy.py b/deepspeed/module_inject/replace_policy.py
index 2bbec0f9ff62..325bb37efd74 100755
--- a/deepspeed/module_inject/replace_policy.py
+++ b/deepspeed/module_inject/replace_policy.py
@@ -41,16 +41,18 @@ def layerNorm(self):
 
 
 class HFBertLayerPolicy(DSPolicy):
-    try:
-        import transformers
-        _orig_layer_class = transformers.models.bert.modeling_bert.BertLayer
-    except:
-        _orig_layer_class = None
+    _orig_layer_class = None
 
     def __init__(self, client_module, inference=False, preln=False):
         super().__init__(inference)
         self.client_module = client_module
         self.preln = preln
+        if HFBertLayerPolicy._orig_layer_class is None:
+            try:
+                import transformers
+                HFBertLayerPolicy._orig_layer_class = transformers.models.bert.modeling_bert.BertLayer
+            except:
+                HFBertLayerPolicy._orig_layer_class = None
 
     def get_hidden_heads(self):
         return self.client_module.attention.self.query.weight.data.shape[1], \
@@ -98,15 +100,16 @@ def layerNorm(self):
 
 
 class HFGPTNEOLayerPolicy(DSPolicy):
-    try:
-        import transformers
-        _orig_layer_class = transformers.models.gpt_neo.modeling_gpt_neo.GPTNeoBlock
-    except:
-        _orig_layer_class = None
+    _orig_layer_class = None
 
     def __init__(self, client_module, inference=True):
         super().__init__(inference, scale_attention=False)
         self.client_module = client_module
+        try:
+            import transformers
+            HFGPTNEOLayerPolicy._orig_layer_class = transformers.models.gpt_neo.modeling_gpt_neo.GPTNeoBlock
+        except:
+            HFGPTNEOLayerPolicy._orig_layer_class = None
 
     def get_hidden_heads(self):
         return self.client_module.attn.attention.q_proj.weight.data.shape[1], \
@@ -141,11 +144,7 @@ def layerNorm(self):
 
 
 class MegatronLayerPolicy(DSPolicy):
-    try:
-        import megatron
-        _orig_layer_class = megatron.model.ParallelTransformerLayer
-    except:
-        _orig_layer_class = None
+    _orig_layer_class = None
 
     def __init__(self, client_module, version=0, inference=True):
         super().__init__(inference)
@@ -153,6 +152,13 @@ def __init__(self, client_module, version=0, inference=True):
         # we use megatron version to differentiate between the old and new
         # megatron-lm source code
         self.version = version
+        if MegatronLayerPolicy._orig_layer_class is None:
+            try:
+                import megatron
+                from megatron.model.transformer import ParallelTransformerLayer
+                MegatronLayerPolicy._orig_layer_class = ParallelTransformerLayer
+            except ImportError:
+                MegatronLayerPolicy._orig_layer_class = None
 
     def get_hidden_heads(self):
         return self.client_module.attention.query_key_value.weight.data.shape[1], \
@@ -187,16 +193,17 @@ def layerNorm(self):
 
 
 class HFGPT2LayerPolicy(DSPolicy):
-    try:
-        import transformers
-        _orig_layer_class = transformers.models.gpt2.modeling_gpt2.GPT2Block
-    except:
-        _orig_layer_class = None
+    _orig_layer_class = None
 
     def __init__(self, client_module, inference=True):
         # HuggingFace GPT2 uses convolutional layer instead of linear layer
         super().__init__(inference, linear_layer=False)
         self.client_module = client_module
+        try:
+            import transformers
+            HFGPT2LayerPolicy._orig_layer_class = transformers.models.gpt2.modeling_gpt2.GPT2Block
+        except ImportError:
+            HFGPT2LayerPolicy._orig_layer_class = None
 
     def get_hidden_heads(self):
         return self.client_module.attn.embed_dim, \
diff --git a/requirements/requirements-dev.txt b/requirements/requirements-dev.txt
index caedac15d709..1cf960afcfab 100644
--- a/requirements/requirements-dev.txt
+++ b/requirements/requirements-dev.txt
@@ -6,3 +6,5 @@ clang-format
 sphinx
 recommonmark
 sphinx-rtd-theme
+megatron-lm==1.1.5
+importlib-metadata>=4

- [x] Add requirement for dev on importlib-metadata > 4, we were seeing pytest issues with importlib-metadata 2.0
- [x] Delay import of static class variables, this was causing our unit test framework to crash since importing megatron seems to initialize cuda which causes us python/pytorch multi-process issues :(
- [x] fix incorrect import path for ParallelTransformerLayer