| 182 | |
| 183 | |
| 184 | class _Broadcast(Function): |
| 185 | def __init__(self, group=WORLD, device=None): |
| 186 | self.group = group |
| 187 | self.out_device = device |
| 188 | |
| 189 | def forward(self, data): |
| 190 | self.in_device = str(data.device) |
| 191 | return collective_comm( |
| 192 | data, CollectiveComm.Mode.BROADCAST, self.group, self.out_device, |
| 193 | ) |
| 194 | |
| 195 | def backward(self, grad): |
| 196 | # TODO backward with a part of grad |
| 197 | if grad is not None: |
| 198 | return reduce_sum(grad, self.group, self.in_device) |
| 199 | |
| 200 | |
| 201 | def broadcast( |