code: def test_unbalanced_mha_sharding(
      head_size: int, n_heads: int, n_shards: int
  ):
      """
      Test MHA sharding when the distribution of heads will not be equal across all ranks.
      """
      param = fill_with_head_ids(head_size, n_heads)
  
      max_heads = 0
      min_heads = n_heads
      total_heads = 0
      seen_heads = set()
  
      for shard_rank in range(n_shards):
          shard = shard_qkv_param(
              param, shard_rank, n_shards, head_size, n_heads, n_heads
          )
          n_local_q_heads, n_local_kv_heads = get_local_heads(
              shard_rank, n_shards, n_heads, n_heads
          )
          validate_inferred_shape(
              shard, head_size, n_local_q_heads, n_local_kv_heads
          )
  
          n_heads_in_shard = shard.shape[0] // head_size // 3
  
          max_heads = max(max_heads, n_heads_in_shard)
          min_heads = min(min_heads, n_heads_in_shard)
          total_heads += n_heads_in_shard
  
          heads = shard.chunk(n_heads_in_shard * 3, dim=0)
  
          for local_head_id in range(n_heads_in_shard):
              head_qkv = torch.cat(
                  [
                      heads[local_head_id],
                      heads[local_head_id + n_heads_in_shard],
                      heads[local_head_id + 2 * n_heads_in_shard],
                  ],
                  dim=0,
              )
              assert head_qkv.shape == (
                  3 * head_size,
                  head_size * n_heads,
              )
  
              global_head_id = torch.unique_consecutive(head_qkv)
              assert len(global_head_id) == 1
  
              seen_heads.add(global_head_id.item())
  
      assert max_heads - min_heads <= 1
      assert total_heads == n_heads
      assert len(seen_heads) == n_heads
  
file_path: deepspeed/tests/unit/inference/v2/model_implementations/sharding/test_qkv_sharding.py
imports: from typing import Optional
  import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.model_implementations.sharding import *
label_desc: Use image node_id 4 for calling a global function with example usage: test_unbalanced_mha_sharding(head_size, n_heads, n_shards) without return types