diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py index 70dea10bf36..7b85939e1e7 100755 --- a/megatron/core/transformer/multi_token_prediction.py +++ b/megatron/core/transformer/multi_token_prediction.py @@ -130,7 +130,7 @@ def tie_output_layer_state_dict( ) -def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=None): +def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=None, fill_value=0): """Roll the tensor input along the sequence dimension with Context Parallelism (CP) support. This function extends the original roll_tensor to support Context Parallelism, which allows @@ -153,17 +153,21 @@ def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=Non falls back to standard rolling behavior. packed_seq_params (PackedSeqParams): Parameters for packed sequence processing. If provided, respects sequence boundaries. + fill_value: Value to fill at sequence boundaries where no source token exists. + Defaults to 0; padding masks should pass ``True``. Returns: tuple: (rolled_tensor, sum_of_rolled_tensor) """ # Handle packed sequences cases if packed_seq_params is not None: - return _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group) + return _roll_tensor_packed_seq( + tensor, shifts, dims, packed_seq_params, cp_group, fill_value=fill_value + ) # Standard rolling behavior when CP is not enabled (cp_group is None or size=1) if cp_group is None or cp_group.size() == 1: rolled_tensor = torch.roll(tensor, shifts=shifts, dims=dims) - rolled_tensor.select(dims, shifts).fill_(0) + rolled_tensor.select(dims, shifts).fill_(fill_value) return rolled_tensor, rolled_tensor.sum() # CP-enabled rolling: Split tensor into chunks and handle boundary communication @@ -200,8 +204,7 @@ def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=Non req_recv_second_part = torch.distributed.irecv(tensor=tensor_recv_list[1], src=prev_rank) ops.append(req_recv_second_part) else: - # Inserted elements are set to be 0.0. - tensor_recv_list[1] = 0 + tensor_recv_list[1] = fill_value if local_rank != len(global_ranks) - 1: req_recv_first_part = torch.distributed.irecv(tensor=tensor_recv_list[0], src=next_rank) ops.append(req_recv_first_part) @@ -228,7 +231,7 @@ def roll_tensor(tensor, shifts=-1, dims=-1, cp_group=None, packed_seq_params=Non return rolled_tensor, rolled_tensor.sum() -def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=None): +def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=None, fill_value=0): """Roll tensor with packed sequence support. This function handles rolling for packed sequences by respecting sequence boundaries """ @@ -251,21 +254,38 @@ def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=No ) assert cu_seqlens is not None, "Packed sequence parameters must provide cu_seqlens_q." - rolled_tensor = tensor.clone() - cp_size = cp_group.size() if cp_group is not None else 1 if cp_size == 1: + rolled_tensor = tensor.clone() # CP disabled: roll each packed sequence independently within its boundaries for i in range(len(cu_seqlens) - 1): start_idx = cu_seqlens[i] end_idx = cu_seqlens[i + 1] seq_slice = tensor[..., start_idx:end_idx] rolled_seq = torch.roll(seq_slice, shifts=shifts, dims=dims) - # Zero out the last position(s) that would cross sequence boundaries - rolled_seq[..., shifts:] = 0 + rolled_seq[..., shifts:] = fill_value rolled_tensor[..., start_idx:end_idx] = rolled_seq return rolled_tensor, rolled_tensor.sum() + cp_partition_mode = getattr(packed_seq_params, 'cp_partition_mode', 'zigzag') + if cp_partition_mode == 'zigzag': + rolled_tensor = _roll_tensor_packed_seq_zigzag_cp( + tensor, shifts, dims, cu_seqlens, cp_group, fill_value=fill_value + ) + return rolled_tensor, rolled_tensor.sum() + if cp_partition_mode == 'contiguous': + rolled_tensor = _roll_tensor_packed_seq_contiguous_cp( + tensor, dims, cu_seqlens, cp_group, fill_value=fill_value + ) + return rolled_tensor, rolled_tensor.sum() + raise ValueError(f"Unsupported packed sequence CP partition mode: {cp_partition_mode}") + + +def _roll_tensor_packed_seq_zigzag_cp(tensor, shifts, dims, cu_seqlens, cp_group, fill_value=0): + """Roll a zigzag-CP THD shard without crossing packed sequence boundaries.""" + cp_size = cp_group.size() + rolled_tensor = tensor.clone() + # CP enabled: each rank owns two chunks per sequence (front and mirrored tail). local_rank = torch.distributed.get_rank(group=cp_group) global_ranks = torch.distributed.get_process_group_ranks(group=cp_group) @@ -314,7 +334,7 @@ def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=No ops.append(torch.distributed.isend(tensor=tensor_send_list[0], dst=prev_rank)) ops.append(torch.distributed.irecv(tensor=tensor_recv_list[1], src=prev_rank)) else: - tensor_recv_list[1].zero_() + tensor_recv_list[1].fill_(fill_value) if local_rank != cp_size - 1: ops.append(torch.distributed.irecv(tensor=tensor_recv_list[0], src=next_rank)) @@ -338,7 +358,55 @@ def _roll_tensor_packed_seq(tensor, shifts, dims, packed_seq_params, cp_group=No # update the rolled tensor rolled_tensor[..., local_start_idx:local_end_idx] = seq_result - return rolled_tensor, rolled_tensor.sum() + return rolled_tensor + + +def _roll_tensor_packed_seq_contiguous_cp(tensor, dims, cu_seqlens, cp_group, fill_value=0): + """Roll a contiguous-CP THD shard without crossing packed sequence boundaries.""" + local_seq_len = tensor.size(dims) + rolled_tensor = torch.roll(tensor, shifts=-1, dims=dims) + if local_seq_len == 0: + return rolled_tensor + + cp_size = cp_group.size() + local_rank = torch.distributed.get_rank(group=cp_group) + global_ranks = torch.distributed.get_process_group_ranks(group=cp_group) + + cu = cu_seqlens.to(device=tensor.device, dtype=torch.long) + if cu.numel() > 1: + nonduplicate_boundaries = torch.ones(cu.numel(), device=cu.device, dtype=torch.bool) + nonduplicate_boundaries[1:] = cu[1:] != cu[:-1] + cu = cu[nonduplicate_boundaries] + if cu.numel() <= 1: + rolled_tensor.fill_(fill_value) + return rolled_tensor + + global_start = local_rank * local_seq_len + global_positions = global_start + torch.arange(local_seq_len, device=tensor.device) + seq_idx = torch.bucketize(global_positions, cu[1:], right=True).clamp(max=cu.numel() - 2) + seq_ends = cu[1:][seq_idx] + valid_next = (global_positions < cu[-1]) & (global_positions + 1 < seq_ends) + + invalid_next = ~valid_next + rolled_tensor[..., invalid_next] = fill_value + + recv_next_first = torch.empty_like(tensor.select(dims, 0)) + ops = [] + if local_rank < cp_size - 1: + next_rank = global_ranks[local_rank + 1] + ops.append(torch.distributed.irecv(tensor=recv_next_first, src=next_rank)) + if local_rank > 0: + prev_rank = global_ranks[local_rank - 1] + send_first = tensor.select(dims, 0).contiguous() + ops.append(torch.distributed.isend(tensor=send_first, dst=prev_rank)) + for op in ops: + op.wait() + + if local_rank < cp_size - 1: + last = rolled_tensor.select(dims, -1) + last.copy_(torch.where(valid_next[-1], recv_next_first, last)) + + return rolled_tensor class MTPLossLoggingHelper: diff --git a/tests/unit_tests/transformer/test_multi_token_prediction.py b/tests/unit_tests/transformer/test_multi_token_prediction.py index 1d201e611e7..37ec75c5994 100644 --- a/tests/unit_tests/transformer/test_multi_token_prediction.py +++ b/tests/unit_tests/transformer/test_multi_token_prediction.py @@ -611,6 +611,80 @@ def test_roll_tensor_with_packed_sequences(self, cp): Utils.destroy_model_parallel() + def test_roll_tensor_with_packed_sequences_contiguous_cp(self): + """Contiguous THD CP rolls across rank boundaries without crossing sequence boundaries.""" + cp = 2 + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=cp) + cp_group = get_context_parallel_group() + cp_rank = torch.distributed.get_rank(group=cp_group) + + # Full padded layout: + # seq1: [1,2,3,4,5,6,7,0] + # seq2: [11,12,13,14,15,16,17,18,19,20,21,0] + # Contiguous CP rank 0 owns global rows [0, 10), rank 1 owns [10, 20). + if cp_rank == 0: + tensor = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 0, 11, 12]], dtype=torch.float32).cuda() + expected = torch.tensor([[2, 3, 4, 5, 6, 7, 0, 0, 12, 13]], dtype=torch.float32).cuda() + padding_mask = torch.tensor( + [[False, False, False, False, False, False, False, True, False, False]] + ).cuda() + expected_padding_mask = torch.tensor( + [[False, False, False, False, False, False, True, True, False, False]] + ).cuda() + else: + tensor = torch.tensor( + [[13, 14, 15, 16, 17, 18, 19, 20, 21, 0]], dtype=torch.float32 + ).cuda() + expected = torch.tensor( + [[14, 15, 16, 17, 18, 19, 20, 21, 0, 0]], dtype=torch.float32 + ).cuda() + padding_mask = torch.tensor( + [[False, False, False, False, False, False, False, False, False, True]] + ).cuda() + expected_padding_mask = torch.tensor( + [[False, False, False, False, False, False, False, False, True, True]] + ).cuda() + + cu_seqlens = torch.tensor([0, 7, 18], dtype=torch.int32).cuda() + cu_seqlens_padded = torch.tensor([0, 8, 20], dtype=torch.int32).cuda() + packed_seq_params = PackedSeqParams( + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + cu_seqlens_q_padded=cu_seqlens_padded, + cu_seqlens_kv_padded=cu_seqlens_padded, + max_seqlen_q=11, + max_seqlen_kv=11, + qkv_format='thd', + ) + # The partition-mode field is introduced by the dynamic-CP stack. Keep this + # focused fix independently reviewable on top of #4495 by exercising the + # same duck-typed contract used in production. + packed_seq_params.cp_partition_mode = 'contiguous' + + rolled, sum_val = roll_tensor( + tensor, shifts=-1, dims=-1, cp_group=cp_group, packed_seq_params=packed_seq_params + ) + rolled_padding_mask, _ = roll_tensor( + padding_mask, + shifts=-1, + dims=-1, + cp_group=cp_group, + packed_seq_params=packed_seq_params, + fill_value=True, + ) + + assert torch.equal(rolled, expected), ( + f"CP Rank {cp_rank}: Expected\n{expected}\nbut got\n{rolled}\nDiff:\n" + f"{rolled - expected}" + ) + assert torch.equal(rolled_padding_mask, expected_padding_mask), ( + f"CP Rank {cp_rank}: Expected padding mask\n{expected_padding_mask}\nbut got\n" + f"{rolled_padding_mask}" + ) + assert sum_val.numel() == 1, "Sum should be a scalar" + + Utils.destroy_model_parallel() + @pytest.mark.parametrize("cp", [1, 2]) def test_roll_tensor_with_packed_sequences_odd_seqlen(self, cp): """Test roll_tensor with ODD packed seqlens.