code: def all_gather_dp_groups(
      partitioned_param_groups,
      dp_process_group,
      start_alignment_factor,
      allgather_bucket_size,
  ):
      for group_id, partitioned_params in enumerate(
          partitioned_param_groups
      ):
          # Sequential AllGather Best of both worlds
          partition_id = dist.get_rank(group=dp_process_group[group_id])
          dp_world_size = dist.get_world_size(
              group=dp_process_group[group_id]
          )
  
          num_shards = max(
              1,
              partitioned_params[partition_id].numel()
              * dp_world_size
              // allgather_bucket_size,
          )
  
          shard_size = (
              partitioned_params[partition_id].numel() // num_shards
          )
  
          # Enforce nccl/rccl alignment of start location of each shard
          shard_size = shard_size - (
              shard_size % start_alignment_factor
          )
  
          num_elements = shard_size
  
          assert (
              shard_size * num_shards
              <= partitioned_params[partition_id].numel()
          )
  
          for shard_id in range(num_shards):
              if shard_id == (num_shards - 1):
                  num_elements = (
                      partitioned_params[partition_id].numel()
                      - shard_id * shard_size
                  )
  
              shard_list = []
              for dp_id in range(dp_world_size):
                  curr_shard = (
                      partitioned_params[dp_id]
                      .narrow(0, shard_id * shard_size, num_elements)
                      .detach()
                  )
                  shard_list.append(curr_shard)
  
              dist.all_gather(
                  shard_list,
                  shard_list[partition_id],
                  dp_process_group[group_id],
              )
  
imports: from collections.abc import Iterable
  from deepspeed.moe.utils import is_moe_param
  import os
  import psutil
  import gc
  from math import sqrt
  from math import floor
  from bisect import bisect_left
  import torch
  from torch import inf
  from deepspeed import comm
  from deepspeed.utils import groups, logger
  from deepspeed.runtime.constants import PIPE_REPLICATED
  from numpy import prod
label_desc: Use image node_id 28 for calling a global function with example usage: all_gather_dp_groups(partitioned_param_groups, dp_process_group, start_alignment_factor, allgather_bucket_size) without return types