Merge branch 'dev' into optimize_hybrid_ep

Author: Autumn1998

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]