diff --git a/examples/multimodal_dev/forward_step.py b/examples/multimodal_dev/forward_step.py index 2729d4cf720..ce01cea6c30 100644 --- a/examples/multimodal_dev/forward_step.py +++ b/examples/multimodal_dev/forward_step.py @@ -227,9 +227,23 @@ def pack_or_pad_batch( cu_seqlens = list(accumulate(seqlens_list, initial=0)) cu_seqlens_padded = list(accumulate(seqlens_padded_list, initial=0)) + # padding_mask: True at collate-padded positions within each packed + # sample. Real tokens occupy [cu_seqlens_padded[i], +seqlens_list[i]); + # the tail up to cu_seqlens_padded[i+1] is padding. Consumed by MoE + # routing in megatron.core to exclude padded tokens from aux loss, + # z-loss, and expert-bias accumulation. + total_tokens_padded = cu_seqlens_padded[-1] + padding_mask_thd = torch.zeros(total_tokens_padded, dtype=torch.bool) + for i, real_seqlen in enumerate(seqlens_list): + pad_start = cu_seqlens_padded[i] + real_seqlen + pad_end = cu_seqlens_padded[i + 1] + if pad_end > pad_start: + padding_mask_thd[pad_start:pad_end] = True + packed_batch["input_ids"] = torch.concat(input_ids_list, dim=0).unsqueeze(0) packed_batch["labels"] = torch.concat(labels_list, dim=0).unsqueeze(0) packed_batch["loss_mask"] = torch.concat(loss_mask_list, dim=0).unsqueeze(0) + packed_batch["padding_mask"] = padding_mask_thd.unsqueeze(0) packed_batch["pixel_values"] = torch.concat(pixel_values_list) packed_batch["image_grid_thw"] = torch.concat(image_grid_thw_list) # cu_seqlens / cu_seqlens_padded need to reach non-source TP ranks @@ -273,6 +287,10 @@ def pack_or_pad_batch( if divisible_by > 1: target_seqlens = math.ceil(target_seqlens / divisible_by) * divisible_by + # Capture real lengths before in-place padding so we can build a + # padding_mask for MoE routing (True at collate-padded positions). + real_seqlens = [s["input_ids"].shape[0] for s in batch] + for sample in batch: sample["input_ids"] = F.pad( sample["input_ids"], (0, target_seqlens - sample["input_ids"].shape[0]), value=0 @@ -291,6 +309,8 @@ def pack_or_pad_batch( padded_batch["loss_mask"] = torch.concat( [x["loss_mask"].unsqueeze(0) for x in batch], dim=0 ) + positions = torch.arange(target_seqlens).unsqueeze(0) + padded_batch["padding_mask"] = positions >= torch.tensor(real_seqlens).unsqueeze(1) padded_batch["pixel_values"] = torch.concat([x["pixel_values"] for x in batch]) padded_batch["image_grid_thw"] = torch.concat([x["image_grid_thw"] for x in batch]) @@ -392,6 +412,7 @@ def forward_step(data_iterator, model): attention_mask=batch.get("attention_mask", None), labels=batch.get("labels", None), loss_mask=batch.get("loss_mask", None), + padding_mask=batch.get("padding_mask", None), pixel_values=pixel_values, image_grid_thw=batch.get("image_grid_thw", None), packed_seq_params=batch.get("packed_seq_params", None), diff --git a/examples/multimodal_dev/models/base.py b/examples/multimodal_dev/models/base.py index 00c3b10a740..a85f36c468f 100644 --- a/examples/multimodal_dev/models/base.py +++ b/examples/multimodal_dev/models/base.py @@ -202,6 +202,7 @@ def _cp_split_for_forward( attention_mask, position_ids, packed_seq_params, + padding_mask=None, ): """Apply CP split to model-forward inputs. @@ -215,7 +216,10 @@ def _cp_split_for_forward( """ cp_size = parallel_state.get_context_parallel_world_size() if cp_size <= 1: - return (decoder_input, input_ids, labels, loss_mask, attention_mask, position_ids) + return ( + decoder_input, input_ids, labels, loss_mask, + attention_mask, position_ids, padding_mask, + ) cp_rank = parallel_state.get_context_parallel_rank() if packed_seq_params is not None: @@ -233,6 +237,8 @@ def _cp_split_for_forward( labels = labels.index_select(1, idx) if loss_mask is not None: loss_mask = loss_mask.index_select(1, idx) + if padding_mask is not None: + padding_mask = padding_mask.index_select(1, idx) else: def _split(t, seq_dim): @@ -247,8 +253,12 @@ def _split(t, seq_dim): labels = _split(labels, 1) loss_mask = _split(loss_mask, 1) attention_mask = _split(attention_mask, 1) + padding_mask = _split(padding_mask, 1) - return (decoder_input, input_ids, labels, loss_mask, attention_mask, position_ids) + return ( + decoder_input, input_ids, labels, loss_mask, + attention_mask, position_ids, padding_mask, + ) @staticmethod def cp_split_loss_mask(loss_mask, packed_seq_params): @@ -301,6 +311,7 @@ def forward( attention_mask: Tensor = None, labels: Tensor = None, loss_mask: Tensor = None, + padding_mask: Tensor = None, pixel_values: Tensor = None, image_grid_thw: Tensor = None, decoder_input: Tensor = None, @@ -316,6 +327,11 @@ def forward( attention_mask: ``[B, S]`` attention mask (None in THD). labels: ``[B, S]`` target token IDs (``[1, T]`` in THD). loss_mask: ``[B, S]`` mask for loss (``[1, T]`` in THD). + padding_mask: ``[B, S]`` bool mask, True at collate-padded + positions (``[1, T]`` in THD). Forwarded to the language + decoder so MoE routing excludes padded tokens from aux + loss / z-loss / expert-bias accumulation. Distinct from + ``loss_mask``: only true padding, not SFT prompt tokens. pixel_values: Preprocessed image pixels. image_grid_thw: ``[num_images, 3]`` grid dimensions. decoder_input: Pre-computed decoder input (skip embed). @@ -345,16 +361,18 @@ def forward( else: decoder_input = text_embeddings - (decoder_input, input_ids, labels, loss_mask, attention_mask, position_ids) = ( - self._cp_split_for_forward( - decoder_input=decoder_input, - input_ids=input_ids, - labels=labels, - loss_mask=loss_mask, - attention_mask=attention_mask, - position_ids=position_ids, - packed_seq_params=packed_seq_params, - ) + ( + decoder_input, input_ids, labels, loss_mask, + attention_mask, position_ids, padding_mask, + ) = self._cp_split_for_forward( + decoder_input=decoder_input, + input_ids=input_ids, + labels=labels, + loss_mask=loss_mask, + attention_mask=attention_mask, + position_ids=position_ids, + packed_seq_params=packed_seq_params, + padding_mask=padding_mask, ) with self._thd_mrope_no_cp_override(packed_seq_params): @@ -365,5 +383,6 @@ def forward( decoder_input=decoder_input, labels=labels, loss_mask=loss_mask, + padding_mask=padding_mask, packed_seq_params=packed_seq_params, )