code: def partition_balanced(weights, pipeline_parallel_size, num_chunks):
      num_total = pipeline_parallel_size * num_chunks
      num_items = len(weights)
      if num_items <= num_total:
          return partition_uniform(
              num_items, pipeline_parallel_size, num_chunks
          )
  
      intervals = _binary_search(weights, num_total)
  
      current = 0
      parts = [[] for _ in range(pipeline_parallel_size)]
      for inter in intervals:
          parts[current].append(inter)
          current = (current + 1) % pipeline_parallel_size
  
      return parts
  
file_path: ColossalAI/colossalai/legacy/pipeline/utils.py
imports: import heapq
  import inspect
  from collections import OrderedDict
  from typing import List
  import torch
  from colossalai.legacy.nn.layer.utils import CheckpointModule
  from colossalai.logging import get_dist_logger
label_desc: Use image node_id 6 for calling a global function with example usage: partition_balanced(weights, pipeline_parallel_size, num_chunks) and returns: parts, partition_uniform