Prepare a batch for dynamic batching. Args: data (DataProto): The input data. max_token_len (int): The maximum token length for dynamic batching. Returns: Tuple[List[DataProto], List[List[int]]]: A tuple containing a list of DataProto objects and a list
(
data: DataProto,
max_token_len: int,
dp_group=None,
num_batches_divided_by=None,
same_micro_num_in_dp=True,
min_num_micro_batch=None,
use_dynamic_bsz_balance=True,
)
| 441 | |
| 442 | |
| 443 | def prepare_dynamic_batch( |
| 444 | data: DataProto, |
| 445 | max_token_len: int, |
| 446 | dp_group=None, |
| 447 | num_batches_divided_by=None, |
| 448 | same_micro_num_in_dp=True, |
| 449 | min_num_micro_batch=None, |
| 450 | use_dynamic_bsz_balance=True, |
| 451 | ) -> tuple[list[DataProto], list[list[int]]]: |
| 452 | """ |
| 453 | Prepare a batch for dynamic batching. |
| 454 | |
| 455 | Args: |
| 456 | data (DataProto): The input data. |
| 457 | max_token_len (int): The maximum token length for dynamic batching. |
| 458 | |
| 459 | Returns: |
| 460 | Tuple[List[DataProto], List[List[int]]]: A tuple containing a list of DataProto objects |
| 461 | and a list of index lists. |
| 462 | """ |
| 463 | batch, batch_idx_list = rearrange_micro_batches( |
| 464 | data.batch, |
| 465 | max_token_len=max_token_len, |
| 466 | dp_group=dp_group, |
| 467 | num_batches_divided_by=num_batches_divided_by, |
| 468 | same_micro_num_in_dp=same_micro_num_in_dp, |
| 469 | min_num_micro_batch=min_num_micro_batch, |
| 470 | use_dynamic_bsz_balance=use_dynamic_bsz_balance, |
| 471 | ) |
| 472 | micro_batches = [] |
| 473 | for i, batch_idx in enumerate(batch_idx_list): |
| 474 | tensors = dict(batch[i]) |
| 475 | non_tensors = {key: value[batch_idx] for key, value in data.non_tensor_batch.items()} |
| 476 | meta_info = copy.deepcopy(data.meta_info) |
| 477 | micro_batches.append(DataProto.from_dict(tensors, non_tensors, meta_info=meta_info)) |
| 478 | |
| 479 | return micro_batches, batch_idx_list |
| 480 | |
| 481 | |
| 482 | def restore_dynamic_batch(data: torch.Tensor, batch_idx_list: list[list[int]]) -> torch.Tensor: |