verl/utils/seqlen_balancing.py

• karmarkar_karp · function · L25-L130 — def karmarkar_karp(seqlen_list: List[int], k_partitions: int, equal_size: bool): # see: https://en.wikipedia.org/wiki/Largest_differencing_method
• Set · class · L27-L47 — class Set
• init · method · L29-L31 — def init(self) -> None
• add · method · L33-L35 — def add(self, idx: int, val: int)
• merge · method · L37-L40 — def merge(self, other)
• lt · method · L42-L47 — def lt(self, other)
• State · class · L49-L101 — class State
• init · method · L51-L58 — def init(self, items: List[Tuple[int, int]], k: int) -> None
• spread · method · L60-L61 — def spread(self)
• get_partitions · method · L63-L70 — def get_partitions(self)
• merge · method · L72-L75 — def merge(self, other)
• spread · method · L78-L79 — def spread(self) -> int
• lt · method · L81-L87 — def lt(self, other): # least heap, let the state with largest spread to be popped first, # if the spread is the same, let the state who has the largest set # to be popped first.
• repr · method · L89-L101 — def repr(self) -> str
• greedy_partition · function · L133-L149 — def greedy_partition(seqlen_list: List[int], k_partitions: int, equal_size: bool)
• get_seqlen_balanced_partitions · function · L152-L183 — def get_seqlen_balanced_partitions(seqlen_list: List[int], k_partitions: int, equal_size: bool)
• _check_and_sort_partitions · function · L170-L180 — def _check_and_sort_partitions(partitions)
• log_seqlen_unbalance · function · L186-L217 — def log_seqlen_unbalance(seqlen_list: List[int], partitions: List[List[int]], prefix): # add some metrics of seqlen sum on dp ranks
• ceildiv · function · L220-L221 — def ceildiv(a, b)
• rearrange_micro_batches · function · L224-L256 — def rearrange_micro_batches(batch: TensorDict, max_token_len, dp_group=None)
• get_reverse_idx · function · L259-L265 — def get_reverse_idx(idx_map)