def METHOD_NAME(self,
                                device,
                                dtype,
                                attention_mask: Optional[torch.ByteTensor] = None,
                                prefix_mask: Optional[torch.ByteTensor] = None,
                                sequence_id: Optional[torch.LongTensor] = None):
    (attn_bias, attention_mask) = self._attn_bias_orig(device,
                                                       dtype,
                                                       attention_mask=attention_mask,
                                                       prefix_mask=prefix_mask,
                                                       sequence_id=sequence_id)
    if dist.is_initialized():
        num_heads_per_rank = int(self.config.n_heads / dist.get_world_size())
        offset = dist.get_rank() * num_heads_per_rank
        attn_bias = attn_bias[:, offset:num_heads_per_rank + offset, :, :]
    return attn_bias, attention_mask