| 22 | DEFAULT_PAD_TOKEN = "[PAD]" |
| 23 | |
| 24 | def sequence_gather(s, world_size, pad_tok_id): |
| 25 | local_size = torch.tensor(s.size(), device=s.device) |
| 26 | all_sizes = [torch.zeros_like(local_size) for _ in range(world_size)] |
| 27 | dist.all_gather(all_sizes, local_size) |
| 28 | max_length = max(size[1] for size in all_sizes) |
| 29 | length_diff = max_length.item() - local_size[1].item() |
| 30 | if length_diff: |
| 31 | pad_size = (*s.shape[:-1], length_diff) |
| 32 | padding = torch.ones(pad_size, device=s.device, dtype=s.dtype) * pad_tok_id |
| 33 | s = torch.concat((s, padding), dim = -1) |
| 34 | gathered_s = [torch.ones_like(s)*pad_tok_id for _ in range(world_size)] |
| 35 | dist.all_gather(gathered_s, s) |
| 36 | |
| 37 | return gathered_s |
| 38 | |
| 39 | def smart_tokenizer_and_embedding_resize( |
| 40 | special_tokens_dict: Dict, |