Fix: don't enter branch if mtp_num_layers == 0 (#2581)

Co-authored-by: Xin Yao <xiny@nvidia.com>

Author: rj42

megatron/core/models/gpt/gpt_model.py

@@ -562,7 +562,8 @@ def _postprocess(
         if not self.post_process:
             return hidden_states
 
-        if self.config.mtp_num_layers is not None:
+        # Skip when mtp_num_layers is None or 0
+        if self.config.mtp_num_layers:
             mtp_labels = labels.clone()
             hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0)
             hidden_states = hidden_states_list[0]