()
| 48 | |
| 49 | |
| 50 | def test_dynamic_batch(): |
| 51 | input_ids = torch.randint(low=0, high=10, size=(20, 100)) |
| 52 | |
| 53 | attention_mask = create_random_mask( |
| 54 | input_ids=input_ids, max_ratio_of_left_padding=0.1, max_ratio_of_valid_token=0.9, min_ratio_of_valid_token=0.5 |
| 55 | ) |
| 56 | data = {"input_ids": input_ids, "attention_mask": attention_mask} |
| 57 | dataproto = DataProto.from_single_dict(data) |
| 58 | micro_batches, micro_bsz_idx_lst = prepare_dynamic_batch(dataproto, max_token_len=300) |
| 59 | input_ids = torch.cat([micro_batch.batch["input_ids"] for micro_batch in micro_batches], dim=0) |
| 60 | input_ids = restore_dynamic_batch(input_ids, micro_bsz_idx_lst) |
| 61 | torch.testing.assert_close(input_ids, dataproto.batch["input_ids"]) |
| 62 | |
| 63 | |
| 64 | def _worker(rank, world_size, init_method, max_token_len, use_same_dp, min_mb): |
nothing calls this directly
no test coverage detected