Repository: hpcaitech/ColossalAI
File: applications/ColossalChat/coati/distributed/utils.py (Python)

Code before:
    logits: torch.Tensor,
    sequences: torch.LongTensor,
    num_actions: int,
    shard_config,
    vocab_size: int = None,
) -> torch.Tensor:
    """Calculate action log probs.

    Args:
        logits (torch.Tensor): Output tensor of Actor.forward.logits.
        sequences (torch.LongTensor): Input sequences.
        num_actions (int): Number of actions.
        shard_config
        vocab_size


    Returns:
        torch.Tensor: Action log probs.
    """
    # labels: torch.Tensor,  # [B, S] or [B, S, Vocab_size]
    # logits: torch.Tensor,  # [B, S, Vocab_size]
    if shard_config.tensor_parallel_size > 1 and shard_config.parallel_output:
        log_probs = dist_log_prob(sequences, logits, shard_config, vocab_size, logits.dtype)
        log_probs = log_probs.squeeze(-1)
    else:
        log_probs = log_probs_from_logits(logits[:, :-1, :], sequences[:, 1:])
    return log_probs[:, -num_actions:]


def masked_mean(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1) -> torch.Tensor:
    """
    Compute the masked mean of a tensor along a specified dimension.

    Args:
        tensor (torch.Tensor): The input tensor.
        mask (torch.Tensor): The mask tensor with the same shape as the input tensor.
        dim (int, optional): The dimension along which to compute the mean. Default is 1.

    Returns:
        torch.Tensor: The masked mean tensor.

    """
    tensor = tensor * mask
    tensor = tensor.sum(dim=dim)
    mask_sum = mask.sum(dim=dim)
    mean = tensor / (mask_sum + 1e-8)
    return mean


Review comment (suggestion):
`dist_log_prob` and `log_probs_from_logits ` are same level of function but different version, which means that one is for non-tp and one is for tp. 

Thus, the input and output should keep as similar as possible.

For example, both functions have `log_probs = log_probs.squeeze(-1)`, but we have one inside the function but one outside the function, which is not a good practice. 

Code after:
    logits: torch.Tensor,
    sequences: torch.LongTensor,
    num_actions: int,
    shard_config,
    vocab_size: int = None,
) -> torch.Tensor:
    """Calculate action log probs.

    Args:
        logits (torch.Tensor): Output tensor of Actor.forward.logits.
        sequences (torch.LongTensor): Input sequences.
        num_actions (int): Number of actions.
        shard_config
        vocab_size


    Returns:
        torch.Tensor: Action log probs.
    """
    # labels: torch.Tensor,  # [B, S] or [B, S, Vocab_size]
    # logits: torch.Tensor,  # [B, S, Vocab_size]
    log_probs = dist_log_prob(sequences, logits, shard_config, vocab_size, logits.dtype)
    log_probs = log_probs.squeeze(-1)
    return log_probs[:, -num_actions:]


def masked_mean(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1) -> torch.Tensor:
    """
    Compute the masked mean of a tensor along a specified dimension.

    Args:
        tensor (torch.Tensor): The input tensor.
        mask (torch.Tensor): The mask tensor with the same shape as the input tensor.
        dim (int, optional): The dimension along which to compute the mean. Default is 1.

    Returns:
        torch.Tensor: The masked mean tensor.

    """
    tensor = tensor * mask
    tensor = tensor.sum(dim=dim)
    mask_sum = mask.sum(dim=dim)
    mean = tensor / (mask_sum + 1e-8)
    return mean