初始化项目，由ModelHub XC社区提供模型
Model: ayh015/myLightningOPD Source: Original Platform
This commit is contained in:
ModelHub XC
210
slime/backends/megatron_utils/cp_utils.py
Normal file
210
slime/backends/megatron_utils/cp_utils.py
Normal file
@@ -0,0 +1,210 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
import torch.nn.functional as F
|
||||
from megatron.core import mpu
|
||||
|
||||
|
||||
def get_logits_and_tokens_offset_with_cp(
|
||||
total_length: int,
|
||||
response_length: int,
|
||||
):
|
||||
"""
|
||||
All offsets start from the begining of the prompt.
|
||||
"""
|
||||
cp_rank = mpu.get_context_parallel_rank()
|
||||
cp_size = mpu.get_context_parallel_world_size()
|
||||
assert cp_size > 1
|
||||
|
||||
prompt_length = total_length - response_length
|
||||
chunk_size = (total_length + 2 * cp_size - 1) // (2 * cp_size)
|
||||
|
||||
# the offset of 2 chunks
|
||||
chunk_0 = (cp_rank * chunk_size, (cp_rank + 1) * chunk_size)
|
||||
chunk_1 = ((2 * cp_size - cp_rank - 1) * chunk_size, (2 * cp_size - cp_rank) * chunk_size)
|
||||
|
||||
# the offset of 2 logits, note that the logits need a "-1".
|
||||
logits_0 = (max(chunk_0[0], prompt_length - 1), min(chunk_0[1], total_length - 1))
|
||||
logits_1 = (max(chunk_1[0], prompt_length - 1), min(chunk_1[1], total_length - 1))
|
||||
|
||||
# when the sequence is empty, make an empty slice to continue the gradient flow.
|
||||
if logits_0[0] < logits_0[1]:
|
||||
token_0 = (logits_0[0] + 1, logits_0[1] + 1)
|
||||
else:
|
||||
logits_0 = (0, 0)
|
||||
token_0 = (0, 0)
|
||||
|
||||
if logits_1[0] < logits_1[1]:
|
||||
token_1 = (logits_1[0] + 1, logits_1[1] + 1)
|
||||
else:
|
||||
logits_1 = (0, 0)
|
||||
token_1 = (0, 0)
|
||||
|
||||
return chunk_size, (chunk_0, chunk_1), (logits_0, logits_1), (token_0, token_1)
|
||||
|
||||
|
||||
def get_sum_of_sample_mean(
|
||||
total_lengths: list[int],
|
||||
response_lengths: list[int],
|
||||
loss_masks: list[torch.Tensor],
|
||||
calculate_per_token_loss: bool = False,
|
||||
) -> Callable[[torch.Tensor], torch.Tensor]:
|
||||
"""
|
||||
Calculate correct sample mean for CP
|
||||
"""
|
||||
cp_size = mpu.get_context_parallel_world_size()
|
||||
if cp_size == 1:
|
||||
|
||||
def sum_of_sample_mean(x: torch.Tensor) -> torch.Tensor:
|
||||
return sum(
|
||||
[
|
||||
(x_i * loss_mask_i).sum() / torch.clamp_min(loss_mask_i.sum(), 1)
|
||||
for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=False)
|
||||
]
|
||||
)
|
||||
|
||||
def sum_of_token(x: torch.Tensor) -> torch.Tensor:
|
||||
return sum(
|
||||
[
|
||||
(x_i * loss_mask_i).sum()
|
||||
for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=False)
|
||||
]
|
||||
)
|
||||
|
||||
else:
|
||||
cp_chunk_lengths = []
|
||||
chunked_loss_masks = []
|
||||
for i, (total_length, response_length, loss_mask) in enumerate(
|
||||
zip(total_lengths, response_lengths, loss_masks, strict=False)
|
||||
):
|
||||
prompt_length = total_length - response_length
|
||||
_, _, _, tokens_offset = get_logits_and_tokens_offset_with_cp(total_length, response_length)
|
||||
loss_mask_0 = loss_mask[tokens_offset[0][0] - prompt_length : tokens_offset[0][1] - prompt_length]
|
||||
loss_mask_1 = loss_mask[tokens_offset[1][0] - prompt_length : tokens_offset[1][1] - prompt_length]
|
||||
chunked_loss_masks.append(torch.cat([loss_mask_0, loss_mask_1], dim=0))
|
||||
cp_chunk_lengths.append(chunked_loss_masks[i].size(0))
|
||||
|
||||
def sum_of_sample_mean(x: torch.Tensor) -> torch.Tensor:
|
||||
return sum(
|
||||
[
|
||||
(x_i * chunked_loss_mask).sum() / torch.clamp_min(loss_mask.sum(), 1)
|
||||
for x_i, chunked_loss_mask, loss_mask in zip(
|
||||
x.split(cp_chunk_lengths, dim=0), chunked_loss_masks, loss_masks, strict=False
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
def sum_of_token(x: torch.Tensor) -> torch.Tensor:
|
||||
return sum(
|
||||
[
|
||||
(x_i * chunked_loss_mask).sum()
|
||||
for x_i, chunked_loss_mask in zip(
|
||||
x.split(cp_chunk_lengths, dim=0), chunked_loss_masks, strict=False
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
return sum_of_sample_mean if not calculate_per_token_loss else sum_of_token
|
||||
|
||||
|
||||
def all_gather_with_cp(tensor: torch.Tensor, total_length: int, response_length: int) -> torch.Tensor:
|
||||
"""
|
||||
Gather tensors across all ranks in the context parallel group.
|
||||
The first dimension of the output tensor will be the `response_length`.
|
||||
"""
|
||||
cp_group = mpu.get_context_parallel_group()
|
||||
cp_size = mpu.get_context_parallel_world_size()
|
||||
|
||||
if cp_size == 1:
|
||||
return tensor
|
||||
|
||||
_, _, logits_offset, _ = get_logits_and_tokens_offset_with_cp(total_length, response_length)
|
||||
|
||||
prompt_length = total_length - response_length
|
||||
|
||||
chunk_0 = tensor[: logits_offset[0][1] - logits_offset[0][0]]
|
||||
chunk_1 = tensor[logits_offset[0][1] - logits_offset[0][0] :]
|
||||
assert chunk_1.shape[0] == logits_offset[1][1] - logits_offset[1][0]
|
||||
|
||||
def zero(len: int) -> torch.Tensor:
|
||||
return torch.zeros(
|
||||
[len] + list(tensor.shape[1:]),
|
||||
dtype=tensor.dtype,
|
||||
device=tensor.device,
|
||||
requires_grad=True,
|
||||
)
|
||||
|
||||
# logprob should be within the range of [prompt_length - 1, total_length - 1]
|
||||
if chunk_0.shape[0] == 0 and chunk_1.shape[0] == 0:
|
||||
# all empty
|
||||
full_tensor = zero(response_length)
|
||||
elif chunk_0.shape[0] != 0 and chunk_1.shape[0] == 0:
|
||||
# only first chunk
|
||||
left = zero(logits_offset[0][0] - (prompt_length - 1))
|
||||
right = zero(total_length - 1 - logits_offset[0][1])
|
||||
full_tensor = torch.cat([left, chunk_0, right], dim=0)
|
||||
elif chunk_0.shape[0] == 0 and chunk_1.shape[0] != 0:
|
||||
# only second chunk
|
||||
left = zero(logits_offset[1][0] - (prompt_length - 1))
|
||||
right = zero(total_length - 1 - logits_offset[1][1])
|
||||
full_tensor = torch.cat([left, chunk_1, right], dim=0)
|
||||
else:
|
||||
left = zero(logits_offset[0][0] - (prompt_length - 1))
|
||||
mid = zero(logits_offset[1][0] - logits_offset[0][1])
|
||||
right = zero(total_length - 1 - logits_offset[1][1])
|
||||
full_tensor = torch.cat([left, chunk_0, mid, chunk_1, right], dim=0)
|
||||
|
||||
assert full_tensor.shape[0] == response_length, f"Expected {response_length}, got {full_tensor.shape}"
|
||||
full_tensor = dist.nn.all_reduce(full_tensor, group=cp_group)
|
||||
return full_tensor
|
||||
|
||||
|
||||
def slice_with_cp(tokens: torch.Tensor, pad_value: tuple[int, float, Callable]) -> torch.Tensor:
|
||||
cp_rank = mpu.get_context_parallel_rank()
|
||||
cp_size = mpu.get_context_parallel_world_size()
|
||||
|
||||
if cp_size == 1:
|
||||
return tokens
|
||||
|
||||
# pad
|
||||
chunk_size = (len(tokens) + 2 * cp_size - 1) // (2 * cp_size)
|
||||
pad = 2 * cp_size * chunk_size - len(tokens)
|
||||
if isinstance(pad_value, Callable):
|
||||
pad_func = pad_value
|
||||
tokens = pad_func(tokens, pad)
|
||||
else:
|
||||
# pad on the first dimension
|
||||
pad_tuple = (0, 0) * (tokens.dim() - 1) + (0, pad)
|
||||
tokens = F.pad(tokens, pad_tuple, value=pad_value)
|
||||
# get 2 chunk for thd cp
|
||||
start_1, end_1 = chunk_size * cp_rank, chunk_size * (cp_rank + 1)
|
||||
start_2, end_2 = chunk_size * (2 * cp_size - cp_rank - 1), chunk_size * (2 * cp_size - cp_rank)
|
||||
return torch.cat([tokens[start_1:end_1], tokens[start_2:end_2]])
|
||||
|
||||
|
||||
def slice_log_prob_with_cp(
|
||||
log_prob: list[float] | torch.Tensor,
|
||||
total_length: int,
|
||||
response_length: int,
|
||||
) -> list[float] | torch.Tensor:
|
||||
assert len(log_prob) == response_length
|
||||
|
||||
cp_size = mpu.get_context_parallel_world_size()
|
||||
|
||||
if cp_size == 1:
|
||||
return log_prob
|
||||
|
||||
prompt_length = total_length - response_length
|
||||
_, _, logits_offset, _ = get_logits_and_tokens_offset_with_cp(total_length, response_length)
|
||||
|
||||
chunk_1 = log_prob[logits_offset[0][0] - (prompt_length - 1) : logits_offset[0][1] - (prompt_length - 1)]
|
||||
chunk_2 = log_prob[logits_offset[1][0] - (prompt_length - 1) : logits_offset[1][1] - (prompt_length - 1)]
|
||||
|
||||
if isinstance(log_prob, list):
|
||||
return chunk_1 + chunk_2
|
||||
else:
|
||||
return torch.cat([chunk_1, chunk_2], dim=0)
|
||||
Reference in New Issue
Block a user