| 151 | return tokens |
| 152 | |
| 153 | def sequence_gather(s, world_size, pad_tok_id): |
| 154 | local_size = torch.tensor(s.size(), device=s.device) |
| 155 | all_sizes = [torch.zeros_like(local_size) for _ in range(world_size)] |
| 156 | dist.all_gather(all_sizes, local_size) |
| 157 | max_length = max(size[1] for size in all_sizes) |
| 158 | length_diff = max_length.item() - local_size[1].item() |
| 159 | if length_diff: |
| 160 | pad_size = (*s.shape[:-1], length_diff) |
| 161 | padding = torch.ones(pad_size, device=s.device, dtype=s.dtype) * pad_tok_id |
| 162 | s = torch.concat((s, padding), dim = -1) |
| 163 | gathered_s = [torch.ones_like(s)*pad_tok_id for _ in range(world_size)] |
| 164 | dist.all_gather(gathered_s, s) |
| 165 | |
| 166 | return gathered_s |
| 167 | |
| 168 | @dataclass |
| 169 | class DataCollatorForSupervisedDataset(object): |