code: def test_mha_unbalanced_sharding(n_heads: int, n_shards: int):
      """
      Unbalanced head sharding for MHA.
  
      Args:
          n_heads (int): The number QKV heads.
          n_shards (int): The number of shards to test for.
      """
      param = fill_with_head_ids(HEAD_SIZE, n_heads)
  
      max_heads = 0
      min_heads = n_heads
      seen_heads = set()
      total_heads = 0
  
      for shard_rank in range(n_shards):
          sharded_param = shard_attn_out_param(
              param, shard_rank, n_shards, HEAD_SIZE
          )
          n_heads_local_q, _ = get_local_heads(
              shard_rank, n_shards, n_heads
          )
  
          assert sharded_param.shape[-1] == HEAD_SIZE * n_heads_local_q
  
          n_local_heads = sharded_param.shape[1] // HEAD_SIZE
          total_heads += n_local_heads
          max_heads = max(max_heads, n_local_heads)
          min_heads = min(min_heads, n_local_heads)
  
          for i in range(n_local_heads):
              head_ids = torch.unique_consecutive(
                  sharded_param[:, i * HEAD_SIZE : (i + 1) * HEAD_SIZE]
              )
              assert len(head_ids) == 1
              seen_heads.add(head_ids.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_attn_out_sharding.py
imports: import pytest
  import torch
  from deepspeed.accelerator import get_accelerator
  from deepspeed.inference.v2.model_implementations.sharding import *
label_desc: Use image node_id 3 for calling a global function with example usage: test_mha_unbalanced_sharding(n_heads, n_shards) without return types