clear mtp code

Author: Jintao-Huang

src/mcore_bridge/model/gpt_model.py

@@ -391,6 +391,8 @@ def _postprocess(
         output_weight = None
         if self.share_embeddings_and_output_weights:
             output_weight = self.shared_embedding_or_output_weight()
+        if self.config.is_multimodal and self.config.context_parallel_size > 1:
+            input_ids = split_cp_inputs(input_ids, getattr(packed_seq_params, 'cu_seqlens_q', None), 1)
 
         if self.mtp_process:
             hidden_states = self.mtp(
@@ -407,55 +409,56 @@ def _postprocess(
                 embedding=self.embedding,
                 **(extra_block_kwargs or {}),
             )
+            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]
-
-            if labels is not None:
-                mtp_labels = labels.clone()
-                if loss_mask is None:
-                    # if loss_mask is not provided, use all ones as loss_mask
-                    if packed_seq_params is None:
-                        loss_mask = torch.ones_like(mtp_labels)
-                    else:
-                        loss_mask = mtp_labels.new_ones((1, packed_seq_params.cu_seqlens_q[-1]))
-                cu_seqlens = packed_seq_params.cu_seqlens_q if packed_seq_params is not None else None
-                for mtp_layer_number in range(self.config.mtp_num_layers):
-                    # output
-                    mtp_logits, _ = self.output_layer(
-                        hidden_states_list[mtp_layer_number + 1],
-                        weight=output_weight,
-                        runtime_gather_output=runtime_gather_output,
+            if loss_mask is None:
+                # if loss_mask is not provided, use all ones as loss_mask
+                loss_mask = torch.ones_like(mtp_labels)
+            for mtp_layer_number in range(self.config.mtp_num_layers):
+                # output
+                mtp_logits, _ = self.output_layer(
+                    hidden_states_list[mtp_layer_number + 1],
+                    weight=output_weight,
+                    runtime_gather_output=runtime_gather_output,
+                )
+                # Calc loss for the current Multi-Token Prediction (MTP) layers.
+                mtp_labels, _ = roll_tensor(
+                    mtp_labels,
+                    shifts=-1,
+                    dims=-1,
+                    cp_group=self.cp_group,
+                    packed_seq_params=packed_seq_params,
+                )
+                loss_mask, num_tokens = roll_tensor(
+                    loss_mask,
+                    shifts=-1,
+                    dims=-1,
+                    cp_group=self.cp_group,
+                    packed_seq_params=packed_seq_params,
+                )
+                mtp_loss = self.compute_language_model_loss(mtp_labels, mtp_logits)
+                mtp_loss = loss_mask * mtp_loss
+                if self.training:
+                    # TODO(shifangx): remove the use of parallel_state here
+                    # after moving loss logging to loss_func in pretrain_gpt.py
+                    MTPLossLoggingHelper.save_loss_to_tracker(
+                        torch.sum(mtp_loss) / num_tokens,
+                        mtp_layer_number,
+                        self.config.mtp_num_layers,
+                        avg_group=parallel_state.get_data_parallel_group(
+                            with_context_parallel=True
+                        ),
+                    )
+                mtp_loss_scale = self.config.mtp_loss_scaling_factor / self.config.mtp_num_layers
+                if self.config.calculate_per_token_loss:
+                    hidden_states = MTPLossAutoScaler.apply(
+                        hidden_states, mtp_loss_scale * mtp_loss
+                    )
+                else:
+                    hidden_states = MTPLossAutoScaler.apply(
+                        hidden_states, mtp_loss_scale * mtp_loss / num_tokens
                     )
-                    # Calc loss for the current Multi-Token Prediction (MTP) layers.
-                    mtp_labels, _ = roll_tensor(mtp_labels, shifts=-1, dims=-1, cp_group=self.cp_group)
-                    if cu_seqlens is None:
-                        loss_mask, _ = roll_tensor(loss_mask, shifts=-1, dims=-1, cp_group=self.cp_group)
-                        loss_mask_ = loss_mask
-                    else:
-                        loss_mask[:, cu_seqlens[:-1]] = 0
-                        loss_mask, _ = roll_tensor(loss_mask, shifts=-1, dims=-1)
-                        if self.config.context_parallel_size > 1:
-                            loss_mask_ = split_cp_inputs(loss_mask, cu_seqlens, dim=1)
-                        else:
-                            loss_mask_ = loss_mask.clone()
-                    mtp_loss = self.compute_language_model_loss(mtp_labels, mtp_logits)
-                    loss_mask_ = loss_mask_ & (mtp_labels != -100)
-                    mtp_loss = loss_mask_ * mtp_loss
-                    num_tokens = loss_mask_.sum()
-                    if self.training:
-                        mtp_loss_for_log = (
-                            torch.sum(mtp_loss) / num_tokens if num_tokens > 0 else mtp_loss.new_tensor(0.0))
-                        MTPLossLoggingHelper.save_loss_to_tracker(
-                            mtp_loss_for_log,
-                            mtp_layer_number,
-                            self.config.mtp_num_layers,
-                            avg_group=parallel_state.get_data_parallel_group(with_context_parallel=True),
-                        )
-                    mtp_loss_scale = self.config.mtp_loss_scaling_factor / self.config.mtp_num_layers
-                    if self.config.calculate_per_token_loss:
-                        hidden_states = MTPLossAutoScaler.apply(hidden_states, mtp_loss_scale * mtp_loss)
-                    else:
-                        hidden_states = MTPLossAutoScaler.apply(hidden_states, mtp_loss_scale * mtp_loss / num_tokens)
         sequence_parallel_override = False
         if in_inference_mode and inference_context.materialize_only_last_token_logits:
             if inference_context.is_static_batching():