add num_layers & pp == 0 assertion

Signed-off-by: dimapihtar <[email protected]>

nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py

@@ -322,6 +322,12 @@ def __init__(self, cfg: DictConfig, trainer: Trainer):
                     'Expert parallelism is currently not supporting Apex distributed optimizer, use Mcore distributed optimizer instead'
                 )
 
+        if self.cfg.get('num_layers', 12) % self.cfg.get('pipeline_model_parallel_size', 1) != 0:
+            raise ValueError(
+                f"num_layers ({self.cfg.get('num_layers', 12)}) should be divisible by "
+                f"pipeline_model_parallel_size ({self.cfg.get('pipeline_model_parallel_size', 1)})"
+            )
+
         self.transformer_engine = cfg.get('transformer_engine', False)
         if self.megatron_amp_O2 and not self.transformer_engine:
             logging.warning('megatron_amp_O2 is enabled but transformer-engine is not.')