Fix assertion to run pipeline engine with a compiled module (#5197)

This PR fixes assertion in the pipeline engine to compute a compile
module with pipeline parallelism.

deepspeed/runtime/pipe/engine.py

@@ -66,7 +66,9 @@ class PipelineEngine(DeepSpeedEngine):
 
     def __init__(self, has_bool_tensors=False, *super_args, **super_kwargs):
         super().__init__(*super_args, **super_kwargs)
-        assert isinstance(self.module, PipelineModule), "model must base PipelineModule"
+        assert isinstance(self.module, PipelineModule) \
+            or (hasattr(self.module, 'wrapped') and isinstance(self.module.wrapped, PipelineModule)), \
+            "model must base PipelineModule"
 
         assert self.zero_optimization_stage(
         ) < ZeroStageEnum.gradients, "ZeRO-2 and ZeRO-3 are incompatible with pipeline parallelism"