r"""Gather tensors across the specified group and concat them at first dimension. Args: inp: Input tensor. group: The process group to work on. The default group is WORLD which means all processes available. You can use a list of process ranks to create n
(
inp: Tensor, group: Optional[Group] = WORLD, device: Optional[str] = None, axis=0,
)
| 252 | |
| 253 | |
| 254 | def all_gather( |
| 255 | inp: Tensor, group: Optional[Group] = WORLD, device: Optional[str] = None, axis=0, |
| 256 | ) -> Tensor: |
| 257 | r"""Gather tensors across the specified group and concat them at first dimension. |
| 258 | |
| 259 | Args: |
| 260 | inp: Input tensor. |
| 261 | group: The process group to work on. |
| 262 | The default group is WORLD which means all processes available. |
| 263 | You can use a list of process ranks to create new group to work on it, e.g. [1, 3, 5]. |
| 264 | device: The specific device to execute this operator. |
| 265 | None default device means the device of inp will be used. |
| 266 | Specify "gpu0:1" to execute this operator on diffrent cuda stream, |
| 267 | 1 is stream id, and default stream id is 0. |
| 268 | axis: The concat axis for collective_comm result |
| 269 | The default axis is 0 |
| 270 | |
| 271 | Returns: |
| 272 | Result tensor. |
| 273 | |
| 274 | Examples: |
| 275 | |
| 276 | .. code-block:: |
| 277 | |
| 278 | input = Tensor([rank]) |
| 279 | # Rank 0 # input: Tensor([0]) |
| 280 | # Rank 1 # input: Tensor([1]) |
| 281 | output = all_gather(input) |
| 282 | # Rank 0 # output: Tensor([0 1]) |
| 283 | # Rank 1 # output: Tensor([0 1]) |
| 284 | |
| 285 | input = Tensor([rank]) |
| 286 | group = Group([1, 0]) |
| 287 | output = all_gather(input, group) |
| 288 | # Rank 0 # output: Tensor([1 0]) |
| 289 | # Rank 1 # output: Tensor([1 0]) |
| 290 | """ |
| 291 | mode = CollectiveComm.Mode.ALL_GATHER |
| 292 | out = collective_comm(inp, mode, group, device) |
| 293 | if axis == 0: |
| 294 | return out |
| 295 | else: |
| 296 | group_size = group.size if group is not None else 1 |
| 297 | transformed_shape = list(inp._tuple_shape) |
| 298 | transformed_shape[axis] *= group_size |
| 299 | n, *shp = out._tuple_shape |
| 300 | index = ( |
| 301 | [_ for _ in range(1, axis)] |
| 302 | + [axis, 0] |
| 303 | + [_ for _ in range(axis + 1, out.ndim + 1)] |
| 304 | ) |
| 305 | return ( |
| 306 | out.reshape(group_size, n // group_size, *shp) |
| 307 | .transpose(index) |
| 308 | .reshape(transformed_shape) |
| 309 | ) |
| 310 | |
| 311 |