diff --git a/megatron/core/models/gpt/fine_grained_callables.py b/megatron/core/models/gpt/fine_grained_callables.py index fa2a2ec4934..1ea17d60d56 100644 --- a/megatron/core/models/gpt/fine_grained_callables.py +++ b/megatron/core/models/gpt/fine_grained_callables.py @@ -669,14 +669,17 @@ def submodule_mtp_attn_forward(node, hidden_states): node.chunk_state.mtp_hidden_states = list(torch.chunk(hidden_states, 1 + offset, dim=0)) hidden_states = node.chunk_state.mtp_hidden_states[offset] - input_ids, position_ids, decoder_input, hidden_states = layer._get_embeddings( + input_ids, position_ids, padding_mask, decoder_input, hidden_states = layer._get_embeddings( input_ids=node.chunk_state.input_ids, position_ids=node.chunk_state.position_ids, embedding=node.chunk_state.model.embedding, hidden_states=hidden_states, + packed_seq_params=node.chunk_state.packed_seq_params, + padding_mask=node.chunk_state.padding_mask, ) node.chunk_state.input_ids = input_ids node.chunk_state.position_ids = position_ids + node.chunk_state.padding_mask = padding_mask # MTP Layer Preprocess # norm, linear projection and transformer diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index dedada837b7..3f3a2f7f675 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -561,6 +561,7 @@ def forward( loss_mask=loss_mask, decoder_input=decoder_input, attention_mask=attention_mask, + padding_mask=padding_mask, inference_params=inference_params, packed_seq_params=packed_seq_params, sequence_len_offset=sequence_len_offset, @@ -583,6 +584,7 @@ def _postprocess( loss_mask=None, decoder_input=None, attention_mask=None, + padding_mask=None, inference_params=None, packed_seq_params=None, sequence_len_offset=None, @@ -626,6 +628,7 @@ def _postprocess( rotary_pos_sin=rotary_pos_sin, packed_seq_params=packed_seq_params, sequence_len_offset=sequence_len_offset, + padding_mask=padding_mask, embedding=self.embedding, **(extra_block_kwargs or {}), ) diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py index 57315ec48d9..47e88dc9dd7 100644 --- a/megatron/core/ssm/gated_delta_net.py +++ b/megatron/core/ssm/gated_delta_net.py @@ -297,9 +297,39 @@ def forward( # TODO: support inference raise NotImplementedError("GDN does not support inference for now.") - if packed_seq_params is not None: - # TODO: support packed sequence - raise NotImplementedError("GDN does not support packed sequence for now.") + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + assert batch == 1, "Packed sequence expects batch dimension to be 1" + assert ( + not self.config.deterministic_mode + ), "Packed sequence does not support deterministic mode." + + # Resolve cu_seqlens with alignment padding handling. + cu_seqlens_q = self._resolve_cu_seqlens( + packed_seq_params.cu_seqlens_q_padded, + packed_seq_params.cu_seqlens_q, + seq_len, + "cu_seqlens_q", + cp_size=self.cp_size, + ) + cu_seqlens_kv = self._resolve_cu_seqlens( + packed_seq_params.cu_seqlens_kv_padded, + packed_seq_params.cu_seqlens_kv, + seq_len, + "cu_seqlens_kv", + cp_size=self.cp_size, + ) + assert torch.equal(cu_seqlens_q, cu_seqlens_kv), ( + "Currently only support cu_seqlens_q equals to cu_seqlens_kv, " + f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}" + ) + num_packed_seqs = cu_seqlens_q.shape[0] - 1 + assert num_packed_seqs > 0, ( + "Number of packed sequences must be greater than 0, " + f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}" + ) + else: + cu_seqlens_q = None + cu_seqlens_kv = None # Input projection nvtx_range_push(suffix="in_proj") @@ -307,20 +337,41 @@ def forward( nvtx_range_pop(suffix="in_proj") # CP All to All: CP to HP - qkvzba = tensor_a2a_cp2hp( - qkvzba, - seq_dim=0, - head_dim=-1, - cp_group=self.pg_collection.cp, - split_sections=[ - self.qk_dim_local_tp, - self.qk_dim_local_tp, - self.v_dim_local_tp, - self.v_dim_local_tp, - self.num_value_heads // self.tp_size, - self.num_value_heads // self.tp_size, - ], - ) + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + unpacked_qkvzba = _unpack_sequence(qkvzba, cu_seqlens_q // self.cp_size, dim=0) + outputs = [] + for qkvzba_i in unpacked_qkvzba: + qkvzba_i = tensor_a2a_cp2hp( + qkvzba_i, + seq_dim=0, + head_dim=-1, + cp_group=self.pg_collection.cp, + split_sections=[ + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + self.v_dim_local_tp, + self.num_value_heads // self.tp_size, + self.num_value_heads // self.tp_size, + ], + ) + outputs.append(qkvzba_i) + qkvzba = torch.cat(outputs, dim=0) + else: + qkvzba = tensor_a2a_cp2hp( + qkvzba, + seq_dim=0, + head_dim=-1, + cp_group=self.pg_collection.cp, + split_sections=[ + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + self.v_dim_local_tp, + self.num_value_heads // self.tp_size, + self.num_value_heads // self.tp_size, + ], + ) # Transpose: s b x --> b s x # From sbhd to bshd format @@ -387,6 +438,7 @@ def forward( activation=self.activation, initial_state=None, output_final_state=False, + cu_seqlens=cu_seqlens_q, ) nvtx_range_pop(suffix="conv1d") @@ -416,6 +468,7 @@ def forward( initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=False, + cu_seqlens=cu_seqlens_q, ) nvtx_range_pop(suffix="gated_delta_rule") @@ -430,9 +483,19 @@ def forward( norm_out = norm_out.transpose(0, 1).contiguous() # CP all to all: HP to CP - norm_out = tensor_a2a_hp2cp( - norm_out, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp - ) + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + unpacked_norm_out = _unpack_sequence(norm_out, cu_seqlens_q, dim=0) + outputs = [] + for norm_out_i in unpacked_norm_out: + norm_out_i = tensor_a2a_hp2cp( + norm_out_i, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp + ) + outputs.append(norm_out_i) + norm_out = torch.cat(outputs, dim=0) + else: + norm_out = tensor_a2a_hp2cp( + norm_out, seq_dim=0, head_dim=-1, cp_group=self.pg_collection.cp + ) # Output projection nvtx_range_push(suffix="out_proj") @@ -504,6 +567,32 @@ def _compute_g_and_beta(self, A_log_local_cp, dt_bias_local_cp, alpha, beta): beta = beta.sigmoid() return g, beta + def _resolve_cu_seqlens( + self, cu_seqlens_padded, cu_seqlens_actual, total_seq_len, name, cp_size: int = 1 + ): + """Resolve cu_seqlens for packed sequence all-to-all, handling alignment padding.""" + if cu_seqlens_padded is not None: + cu_seqlens = cu_seqlens_padded + else: + cu_seqlens = cu_seqlens_actual + + total_cu = cu_seqlens[-1].cpu().item() + if total_cu != total_seq_len: + raise ValueError( + f"GDN: {name}[-1]={total_cu} does not match " + f"total_sequence_length={total_seq_len}. " + f"({cu_seqlens_padded=}, {cu_seqlens_actual=})." + ) + + seq_lengths = cu_seqlens[1:] - cu_seqlens[:-1] + if (seq_lengths % cp_size != 0).any(): + raise ValueError( + f"All per-sequence lengths in cu_seqlens must be divisible by cp_size={cp_size}, " + f"but got lengths: {seq_lengths.tolist()}" + ) + + return cu_seqlens + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_group=None): """Provide a sharded state dictionary for distributed checkpointing.""" # Guard for cases metadata is not provided @@ -602,6 +691,18 @@ def _backward_out_proj(self): self.out_proj.backward_dw() +def _unpack_sequence(x, cu_seqlens, dim=1): + unpacked_x = [] + cu_seqlens_list = cu_seqlens.tolist() + num_seqs = len(cu_seqlens_list) - 1 + for i in range(num_seqs): + idx_start = cu_seqlens_list[i] + idx_end = cu_seqlens_list[i + 1] + chunked_index = [slice(None)] * dim + [slice(idx_start, idx_end)] + unpacked_x.append(x[tuple(chunked_index)]) + return unpacked_x + + #################### # Sharded state dict utilities #################### @@ -853,6 +954,7 @@ def torch_chunk_gated_delta_rule( initial_state=None, output_final_state=False, use_qk_l2norm_in_kernel=False, + cu_seqlens=None, ): # pylint: disable=line-too-long ''' @@ -862,6 +964,10 @@ def torch_chunk_gated_delta_rule( Reference: https://github.com/huggingface/transformers/blob/144c8ce2809a2e21914017652700e1ecb450501e/src/transformers/models/qwen3_next/modeling_qwen3_next.py#L470-L547 ''' + assert ( + cu_seqlens is None + ), "cu_seqlens is not supported for torch_chunk_gated_delta_rule for now." + initial_dtype = query.dtype if use_qk_l2norm_in_kernel: query = l2norm(query, dim=-1, eps=1e-6) diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index 2e0461e365c..d6cd437721f 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -870,6 +870,7 @@ def _get_embeddings( embedding: Callable, hidden_states: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None, + padding_mask: Optional[torch.Tensor] = None, ): """ Preprocesses input data for the Multi-Token Prediction (MTP) layers. @@ -901,12 +902,20 @@ def _get_embeddings( cp_group=self.cp_group, packed_seq_params=packed_seq_params, ) + if padding_mask is not None: + padding_mask, _ = roll_tensor( + padding_mask, + shifts=-1, + dims=-1, + cp_group=self.cp_group, + packed_seq_params=packed_seq_params, + ) # embedding decoder_input = embedding(input_ids=input_ids, position_ids=position_ids) hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) - return input_ids, position_ids, decoder_input, hidden_states + return input_ids, position_ids, padding_mask, decoder_input, hidden_states def _concat_embeddings(self, hidden_states: torch.Tensor, decoder_input: torch.Tensor): """ @@ -940,6 +949,7 @@ def _proj_and_transformer_layer( hidden_states: Tensor, decoder_input: Tensor, attention_mask: Optional[torch.Tensor] = None, + padding_mask: Optional[torch.Tensor] = None, context: Optional[torch.Tensor] = None, context_mask: Optional[torch.Tensor] = None, rotary_pos_emb: Optional[torch.Tensor] = None, @@ -980,6 +990,7 @@ def _proj_and_transformer_layer( hidden_states = self.mtp_model_layer( hidden_states=hidden_states, attention_mask=attention_mask, + padding_mask=padding_mask, rotary_pos_emb=rotary_pos_emb, inference_context=inference_params, packed_seq_params=packed_seq_params, @@ -998,6 +1009,7 @@ def _proj_and_transformer_layer( inference_params=inference_params, packed_seq_params=packed_seq_params, sequence_len_offset=sequence_len_offset, + padding_mask=padding_mask, ) hidden_states = self._postprocess(hidden_states) @@ -1107,6 +1119,7 @@ def forward( position_ids: Tensor, hidden_states: Tensor, attention_mask: Tensor, + padding_mask: Optional[Tensor] = None, context: Optional[Tensor] = None, context_mask: Optional[Tensor] = None, rotary_pos_emb: Optional[Tensor] = None, @@ -1141,9 +1154,10 @@ def forward( [s, b, h], and optionally the updated context tensor if cross-attention is used. """ assert context is None, "multi token prediction + cross attention is not yet supported." - input_ids, position_ids, decoder_input, hidden_states = self._get_embeddings( + input_ids, position_ids, padding_mask, decoder_input, hidden_states = self._get_embeddings( input_ids=input_ids, position_ids=position_ids, + padding_mask=padding_mask, embedding=embedding, hidden_states=hidden_states, packed_seq_params=packed_seq_params, @@ -1155,6 +1169,7 @@ def forward( hidden_states=hidden_states, decoder_input=decoder_input, attention_mask=attention_mask, + padding_mask=padding_mask, context=context, context_mask=context_mask, rotary_pos_emb=rotary_pos_emb, @@ -1170,6 +1185,7 @@ def forward( hidden_states=hidden_states, decoder_input=decoder_input, attention_mask=attention_mask, + padding_mask=padding_mask, context=context, context_mask=context_mask, rotary_pos_emb=rotary_pos_emb, @@ -1181,7 +1197,7 @@ def forward( sequence_len_offset=sequence_len_offset, ) - return hidden_states, input_ids, position_ids + return hidden_states, input_ids, position_ids, padding_mask def sharded_state_dict( self, prefix: str = '', sharded_offsets: tuple = (), metadata: Optional[dict] = None @@ -1428,6 +1444,7 @@ def forward( position_ids: Tensor, hidden_states: Tensor, attention_mask: Tensor, + padding_mask: Optional[Tensor] = None, context: Optional[Tensor] = None, context_mask: Optional[Tensor] = None, rotary_pos_emb: Optional[Tensor] = None, @@ -1458,11 +1475,12 @@ def forward( hidden_states = hidden_states_list[offset] for iteration in range(self.config.mtp_num_layers): layer_idx = 0 if self.mtp_use_repeated_layer else iteration - (hidden_states, input_ids, position_ids) = self.layers[layer_idx]( + (hidden_states, input_ids, position_ids, padding_mask) = self.layers[layer_idx]( input_ids=input_ids, position_ids=position_ids, hidden_states=hidden_states, attention_mask=attention_mask, + padding_mask=padding_mask, inference_params=inference_params, rotary_pos_emb=rotary_pos_emb, rotary_pos_cos=rotary_pos_cos, diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/test_gated_delta_net.py index ec4d7a86ecf..9af06d07305 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/test_gated_delta_net.py @@ -31,6 +31,10 @@ ) from tests.unit_tests.test_utilities import Utils from tests.unit_tests.transformer.test_attention import _test_parallel_attention_correctness +from tests.unit_tests.transformer.test_multi_latent_attention import ( + make_test_packed_seq_params, + make_test_packed_seq_params_with_padding, +) try: import fla @@ -201,7 +205,153 @@ def test_jit_compiled_helpers(self): assert g.shape == alpha.shape assert beta_sig.shape == beta.shape + def test_gpu_forward_thd_correctness(self): + if self.sp_size > 1: + pytest.skip("Sequence parallel is not supported for this test case.") + + atol, rtol = 3e-4, 3e-4 + + # Input shape + sequence_length = 32 + micro_batch_size = 4 + cu_seqlens = [0, 32, 64, 96, 128] + # sbhd input shape: [sequence length, batch size, hidden size] + sub_sequence_length = sequence_length // self.cp_size + hidden_states_sbhd = torch.rand( + (sub_sequence_length, micro_batch_size, self.gdn.config.hidden_size) + ) + attention_mask_sbhd = None + hidden_states_sbhd = hidden_states_sbhd.cuda().bfloat16() + # thd input shape: [sequence length * batch size, 1, hidden size] + hidden_states_thd = hidden_states_sbhd.transpose(0, 1).contiguous() + hidden_states_thd = hidden_states_thd.view(-1, 1, self.gdn.config.hidden_size) + attention_mask_thd = None + packed_seq_params = make_test_packed_seq_params(cu_seqlens=cu_seqlens) + + # THD format + output_thd, _ = self.gdn( + hidden_states_thd, attention_mask_thd, packed_seq_params=packed_seq_params + ) + # SBHD format + output_sbhd, _ = self.gdn(hidden_states_sbhd, attention_mask_sbhd) + output_sbhd_T = output_sbhd.transpose(0, 1).contiguous().view(*output_thd.shape) + + rank = torch.distributed.get_rank() + assert output_thd.shape[0] == sub_sequence_length * micro_batch_size + assert output_thd.shape[1] == 1 + assert output_thd.shape[2] == self.gdn.config.hidden_size + torch.testing.assert_close( + output_sbhd_T, + output_thd, + atol=atol, + rtol=rtol, + msg=lambda msg: f"Output mismatch ({rank=}): {msg}", + ) + + def test_gpu_forward_thd_padding_correctness(self): + if self.sp_size > 1: + pytest.skip("Sequence parallel is not supported for this test case.") + + atol, rtol = 3e-4, 3e-4 + sequence_length = 32 + micro_batch_size = 4 + + # sbhd input shape: [sequence length, batch size, hidden size] + sub_sequence_length = sequence_length // self.cp_size + hidden_states_sbhd = torch.rand( + (sub_sequence_length, micro_batch_size, self.gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + output_sbhd, _ = self.gdn(hidden_states_sbhd, None) + + # thd input shape: [sequence length * batch size, 1, hidden size] + hidden_states_thd = hidden_states_sbhd.transpose(0, 1).contiguous() + hidden_states_thd = hidden_states_thd.view(-1, 1, self.gdn.config.hidden_size) + output_bshd = output_sbhd.transpose(0, 1).contiguous() + + rank = torch.distributed.get_rank() + + # A) padded branch: prefer *_padded when available. + padded_params = make_test_packed_seq_params_with_padding( + cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 128] + ) + output_thd_padded, _ = self.gdn(hidden_states_thd, None, packed_seq_params=padded_params) + output_thd2bshd = output_thd_padded.view(*output_bshd.shape) + torch.testing.assert_close( + output_bshd[:, :30, :], + output_thd2bshd[:, :30, :], + atol=atol, + rtol=rtol, + msg=lambda msg: f"THD padded output mismatch ({rank=}): {msg}", + ) + + # B) no-padded branch: use actual cu_seqlens when it matches total_sequence_length. + no_padding_params = make_test_packed_seq_params(cu_seqlens=[0, 32, 64, 96, 128]) + output_thd_no_padding, _ = self.gdn( + hidden_states_thd, None, packed_seq_params=no_padding_params + ) + assert output_thd_no_padding.shape == output_thd_padded.shape + + # C) padded mismatch branch: if *_padded[-1] mismatches total_sequence_length, should raise. + padded_mismatch_params = make_test_packed_seq_params_with_padding( + cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 126] + ) + with pytest.raises(ValueError, match="does not match"): + self.gdn(hidden_states_thd, None, packed_seq_params=padded_mismatch_params) + + # D) actual mismatch branch without *_padded: should raise. + actual_mismatch_params = make_test_packed_seq_params(cu_seqlens=[0, 32, 64, 96, 129]) + with pytest.raises(ValueError, match="does not match"): + self.gdn(hidden_states_thd, None, packed_seq_params=actual_mismatch_params) + + +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.internal +class TestGDNCuSeqlensResolve: + + @pytest.fixture + def mock_gdn(self): + class MockGDN: + _resolve_cu_seqlens = GatedDeltaNet._resolve_cu_seqlens + + return MockGDN() + + def test_padded_preferred_when_available(self, mock_gdn): + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + padded = torch.tensor([0, 504, 1008], dtype=torch.int32) + result = mock_gdn._resolve_cu_seqlens(padded, actual, 1008, "cu_seqlens_q", cp_size=2) + assert torch.equal(result, padded) + + def test_actual_used_when_no_padding(self, mock_gdn): + actual = torch.tensor([0, 504, 1008], dtype=torch.int32) + result = mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=2) + assert torch.equal(result, actual) + + def test_raises_when_padding_mismatch(self, mock_gdn): + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + with pytest.raises(ValueError, match="does not match"): + mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=2) + + def test_raises_when_padded_mismatches_total(self, mock_gdn): + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + padded = torch.tensor([0, 504, 1004], dtype=torch.int32) + with pytest.raises(ValueError, match="does not match"): + mock_gdn._resolve_cu_seqlens(padded, actual, 1008, "cu_seqlens_q", cp_size=2) + + def test_raises_when_not_divisible_by_cp_size(self, mock_gdn): + actual = torch.tensor([0, 505, 1008], dtype=torch.int32) + with pytest.raises(ValueError, match="must be divisible by cp_size"): + mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=2) + + def test_cp1_still_validates_total(self, mock_gdn): + mock_gdn.cp_size = 1 + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + with pytest.raises(ValueError, match="does not match"): + mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=1) + +@pytest.mark.parametrize("sequence_packing", [False, True]) @pytest.mark.parametrize( ("tp", "sp", "cp"), [ @@ -213,7 +363,7 @@ def test_jit_compiled_helpers(self): ], ) @pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") -def test_parallel_gated_delta_net_correctness(tmp_path_dist_ckpt, tp, sp, cp): +def test_parallel_gated_delta_net_correctness(tmp_path_dist_ckpt, sequence_packing, tp, sp, cp): transformer_config = TransformerConfig( hidden_size=128, linear_conv_kernel_dim=2, @@ -254,4 +404,5 @@ def test_parallel_gated_delta_net_correctness(tmp_path_dist_ckpt, tp, sp, cp): seed=123, sequence_length=256, micro_batch_size=4, + sequence_packing=sequence_packing, ) diff --git a/tests/unit_tests/transformer/test_attention.py b/tests/unit_tests/transformer/test_attention.py index ab5a33aa61b..9951fc2c82b 100644 --- a/tests/unit_tests/transformer/test_attention.py +++ b/tests/unit_tests/transformer/test_attention.py @@ -34,6 +34,7 @@ init_checkpointing_mock_args, ) from tests.unit_tests.test_utilities import Utils +from tests.unit_tests.transformer.test_multi_latent_attention import make_test_packed_seq_params try: from transformer_engine.pytorch.attention.rope import apply_fused_qkv_rotary_pos_emb @@ -479,6 +480,7 @@ def _test_parallel_attention_correctness( seed=123, sequence_length=256, micro_batch_size=4, + sequence_packing=False, ): # Model initialization function def initialize_gpt_model( @@ -572,17 +574,24 @@ def initialize_gpt_model( def get_tensor_on_this_rank(tensor): if cp > 1: tensor = get_tensor_on_this_cp_rank(tensor, 0, cp_group) + if sequence_packing: + tensor = tensor.transpose(0, 1).contiguous().view(-1, 1, *tensor.shape[2:]) if tp > 1 and sp: - sp_seg = sequence_length // tp // cp + sp_seg = tensor.shape[0] // tp tensor = tensor[tp_rank * sp_seg : (tp_rank + 1) * sp_seg] return tensor # Calculate parallel model output + if sequence_packing: + cu_seqlens = [i * sequence_length for i in range(micro_batch_size + 1)] + packed_seq_params = make_test_packed_seq_params(cu_seqlens=cu_seqlens) + else: + packed_seq_params = None input_hidden_states = get_tensor_on_this_rank(input_hidden_states) input_hidden_states = input_hidden_states.detach().requires_grad_(True) parallel_attention = gpt_model[0].decoder.layers[0].self_attention output_hidden_states_parallel, bias_hidden_states_parallel = parallel_attention( - input_hidden_states, attention_mask=None + input_hidden_states, attention_mask=None, packed_seq_params=packed_seq_params ) output_hidden_states_parallel.sum().backward() input_grad_parallel = input_hidden_states.grad.detach() @@ -647,6 +656,7 @@ def get_tensor_on_this_rank(tensor): Utils.destroy_model_parallel() +@pytest.mark.parametrize("sequence_packing", [False, True]) @pytest.mark.parametrize("apply_rope_fusion", [False, True]) @pytest.mark.parametrize( ("tp", "sp", "cp"), @@ -661,7 +671,7 @@ def get_tensor_on_this_rank(tensor): @pytest.mark.parametrize("qk_layernorm", [False, True]) @pytest.mark.parametrize("output_gate", [False, True]) def test_parallel_attention_correctness( - tmp_path_dist_ckpt, apply_rope_fusion, tp, sp, cp, qk_layernorm, output_gate + tmp_path_dist_ckpt, sequence_packing, apply_rope_fusion, tp, sp, cp, qk_layernorm, output_gate ): transformer_config = TransformerConfig( num_layers=1, @@ -690,6 +700,7 @@ def test_parallel_attention_correctness( cp=cp, seed=123, sequence_length=256, + sequence_packing=sequence_packing, ) diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py index d4d7edfe44b..07e7d87c594 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -2,6 +2,7 @@ import os import sys +import types import pytest import torch @@ -130,6 +131,101 @@ def test_constructor_ues_te(self, tp, cp): elif tp == 4: assert num_weights == 15216 * config.mtp_num_layers + def test_get_embeddings_rolls_padding_mask(self): + """Test that _get_embeddings rolls padding_mask alongside input ids.""" + torch.manual_seed(_SEED) + config, mtp_block_spec = self._create_config_and_mtp_block_spec(tp=1, cp=1) + mtp = MultiTokenPredictionBlock(config=config, spec=mtp_block_spec) + mtp_layer = mtp.layers[0] + + seq_len = 6 + batch_size = 2 + input_ids = torch.tensor([[1, 2, 3, 4, 0, 0], [5, 6, 7, 0, 0, 0]], dtype=torch.int64) + position_ids = torch.arange(seq_len, dtype=torch.int64).repeat(batch_size, 1) + padding_mask = torch.tensor( + [[True, True, True, True, False, False], [True, True, True, False, False, False]] + ) + hidden_states = torch.randn(seq_len, batch_size, config.hidden_size) + + def fake_embedding(input_ids, position_ids): + return torch.zeros(seq_len, batch_size, config.hidden_size, dtype=hidden_states.dtype) + + rolled_input_ids, rolled_position_ids, rolled_padding_mask, _, _ = ( + mtp_layer._get_embeddings( + input_ids=input_ids, + position_ids=position_ids, + padding_mask=padding_mask, + embedding=fake_embedding, + hidden_states=hidden_states, + packed_seq_params=None, + ) + ) + + expected_input_ids, _ = roll_tensor(input_ids, shifts=-1, dims=-1) + expected_position_ids, _ = roll_tensor(position_ids, shifts=-1, dims=-1) + expected_padding_mask, _ = roll_tensor(padding_mask, shifts=-1, dims=-1) + + assert torch.equal(rolled_input_ids, expected_input_ids) + assert torch.equal(rolled_position_ids, expected_position_ids) + assert torch.equal(rolled_padding_mask, expected_padding_mask) + + def test_forward_propagates_rolled_padding_mask(self, monkeypatch): + """Test forward passes rolled padding_mask to transformer path.""" + torch.manual_seed(_SEED) + config, mtp_block_spec = self._create_config_and_mtp_block_spec(tp=1, cp=1) + mtp = MultiTokenPredictionBlock(config=config, spec=mtp_block_spec) + mtp_layer = mtp.layers[0] + + seq_len = 4 + batch_size = 2 + input_ids = torch.tensor([[1, 2, 3, 0], [4, 5, 0, 0]], dtype=torch.int64) + position_ids = torch.arange(seq_len, dtype=torch.int64).repeat(batch_size, 1) + padding_mask = torch.tensor([[True, True, True, False], [True, True, False, False]]) + hidden_states = torch.randn(seq_len, batch_size, config.hidden_size) + attention_mask = torch.ones((batch_size, 1, seq_len, seq_len), dtype=torch.bool) + seen = {} + + def fake_embedding(input_ids, position_ids): + return torch.zeros(seq_len, batch_size, config.hidden_size, dtype=hidden_states.dtype) + + def fake_proj_and_transformer_layer( + self, + hidden_states, + decoder_input, + attention_mask=None, + padding_mask=None, + context=None, + context_mask=None, + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + attention_bias=None, + inference_params=None, + packed_seq_params=None, + sequence_len_offset=None, + ): + seen["padding_mask"] = padding_mask + return hidden_states + + monkeypatch.setattr( + mtp_layer, + "_proj_and_transformer_layer", + types.MethodType(fake_proj_and_transformer_layer, mtp_layer), + ) + + _, _, _, returned_padding_mask = mtp_layer.forward( + input_ids=input_ids, + position_ids=position_ids, + hidden_states=hidden_states, + attention_mask=attention_mask, + padding_mask=padding_mask, + embedding=fake_embedding, + ) + + expected_padding_mask, _ = roll_tensor(padding_mask, shifts=-1, dims=-1) + assert torch.equal(seen["padding_mask"], expected_padding_mask) + assert torch.equal(returned_padding_mask, expected_padding_mask) + class TestMultiTokenPrediction: def setup_method(self, method):