code: def test_gqa_uneven_sharding(
      head_size: int, n_heads_q: int, n_heads_kv: int, n_shards: int
  ):
      """
      Test GQA sharding when there are more shards than KV heads.
      """
      param = fill_with_head_ids(head_size, n_heads_q, n_heads_kv)
  
      n_kv_heads_in_shard = 1
      n_shards_per_kv_head = n_shards // n_heads_kv
  
      max_heads = 0
      min_heads = n_heads_q
      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_q,
              n_heads_kv,
          )
          n_local_q_heads, n_local_kv_heads = get_local_heads(
              shard_rank, n_shards, n_heads_q, n_heads_kv
          )
          validate_inferred_shape(
              shard, head_size, n_local_q_heads, n_local_kv_heads
          )
  
          local_n_heads_q = (
              shard.shape[0] - 2 * n_kv_heads_in_shard * head_size
          ) // head_size
  
          max_heads = max(max_heads, local_n_heads_q)
          min_heads = min(min_heads, local_n_heads_q)
          total_heads += local_n_heads_q
  
          q = shard[: local_n_heads_q * head_size]
          kv = shard[local_n_heads_q * head_size :]
  
          for local_head_id in range(local_n_heads_q):
              q_head_id = torch.unique_consecutive(
                  q[
                      local_head_id
                      * head_size : (local_head_id + 1)
                      * head_size
                  ]
              )
              assert len(q_head_id) == 1
  
              seen_heads.add(q_head_id.item())
  
          kv_id_calc = shard_rank // n_shards_per_kv_head
          kv_id = torch.unique_consecutive(kv)
          assert len(kv_id) == 1
          assert kv_id.item() == kv_id_calc
  
      assert max_heads - min_heads <= 1
      assert total_heads == n_heads_q
      assert len(seen_heads) == n_heads_q
  
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 6 for calling a global function with example usage: test_gqa_uneven_sharding(head_size, n_heads_q, n_heads_kv, n_shards) without return types