Run all_gather on data which is a dictionary of Tensors
(data)
| 157 | |
| 158 | |
| 159 | def all_gather_dict(data): |
| 160 | """ |
| 161 | Run all_gather on data which is a dictionary of Tensors |
| 162 | """ |
| 163 | assert isinstance(data, dict) |
| 164 | |
| 165 | gathered_dict = {} |
| 166 | for item_key in data: |
| 167 | if isinstance(data[item_key], torch.Tensor): |
| 168 | if is_distributed(): |
| 169 | data[item_key] = data[item_key].contiguous() |
| 170 | tensor_list = [torch.empty_like(data[item_key]) for _ in range(get_world_size())] |
| 171 | dist.all_gather(tensor_list, data[item_key]) |
| 172 | gathered_tensor = torch.cat(tensor_list, dim=0) |
| 173 | else: |
| 174 | gathered_tensor = data[item_key] |
| 175 | gathered_dict[item_key] = gathered_tensor |
| 176 | return gathered_dict |
| 177 |
nothing calls this directly
no test coverage detected