code: def get_alibi_mask(self, tensor, seq_length_with_past):
      mask = self.get_alibi_mask_orig(tensor, seq_length_with_past)
      if not self.training and dist.is_initialized():
          num_heads_per_rank = get_shard_size(
              self.n_head, dist.get_world_size()
          )
          offset = sum(
              get_shard_size_list(self.n_head, dist.get_world_size())[
                  0 : dist.get_rank()
              ]
          )
          mask = mask[
              offset : num_heads_per_rank + offset,
              :seq_length_with_past,
              :seq_length_with_past,
          ]
  
      return mask
  
file_path: DeepSpeed/deepspeed/module_inject/auto_tp_model_utils.py
imports: from deepspeed import comm
  import torch
  from typing import Optional
  from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list
label_desc: Use image node_id 2 for calling a global function with example usage: get_alibi_mask(tensor, seq_length_with_past) and returns: mask