r"""Reduce tensors with sum operation on each value across the specified group. Note: ``inp`` tensor must have identical shape in all processes across the group. Args: inp (Tensor): tensor to be reduced. Keyword args: group (Group or sequence of ints): the proc
(
inp: Tensor, group: Optional[Group] = WORLD, device: Optional[str] = None,
)
| 368 | |
| 369 | |
| 370 | def all_reduce_sum( |
| 371 | inp: Tensor, group: Optional[Group] = WORLD, device: Optional[str] = None, |
| 372 | ) -> Tensor: |
| 373 | r"""Reduce tensors with sum operation on each value across the specified group. |
| 374 | |
| 375 | Note: |
| 376 | ``inp`` tensor must have identical shape in all processes across the group. |
| 377 | |
| 378 | Args: |
| 379 | inp (Tensor): tensor to be reduced. |
| 380 | |
| 381 | Keyword args: |
| 382 | group (Group or sequence of ints): the process group to work on. Default: ``WORLD``. |
| 383 | ``WORLD`` group selects all processes available. |
| 384 | list of process rank as parameter will create a new group to work on. |
| 385 | device (:attr:`.Tensor.device`): the specific device to execute this operator. Default: ``None`` |
| 386 | ``None`` will select the device of ``inp`` to execute. |
| 387 | Specially, ``GPU`` device can assign a different stream to execute |
| 388 | by adding a number right after a colon following the device name while |
| 389 | ``:0`` denotes default stream of GPU, otherwise will use default stream. |
| 390 | |
| 391 | Returns: |
| 392 | A tensor with sum operation on each value across the group. |
| 393 | |
| 394 | The shape of the output tensor must be the same as ``inp``, and the output |
| 395 | tensor is going to be bitwise identical in all processes across the group. |
| 396 | |
| 397 | |
| 398 | Examples: |
| 399 | |
| 400 | >>> # We execute all_reduce_sum on rank 0 and rank 1 |
| 401 | >>> input = F.arange(2) + 1 + 2 * rank # doctest: +SKIP |
| 402 | >>> input # doctest: +SKIP |
| 403 | Tensor([1. 2.], device=xpux:0) # Rank 0 |
| 404 | Tensor([3. 4.], device=xpux:0) # Rank 1 |
| 405 | >>> F.distributed.all_reduce_sum(input, group=[0, 1]) # doctest: +SKIP |
| 406 | Tensor([4. 6.], device=xpux:0) # Rank 0 |
| 407 | Tensor([4. 6.], device=xpux:0) # Rank 1 |
| 408 | |
| 409 | >>> # We execute all_reduce_sum with on gpu0 with cuda stream 1 |
| 410 | >>> megengine.set_default_device("gpu0") # doctest: +SKIP |
| 411 | >>> input = F.arange(2) + 1 + 2 * rank # doctest: +SKIP |
| 412 | >>> input # doctest: +SKIP |
| 413 | Tensor([1. 2.], device=gpu0:0) # Rank 0 |
| 414 | Tensor([3. 4.], device=gpu0:0) # Rank 1 |
| 415 | >>> F.distributed.all_reduce_sum(input, device="gpu0:1") # doctest: +SKIP |
| 416 | Tensor([4. 6.], device=gpu0:0) # Rank 0 |
| 417 | Tensor([4. 6.], device=gpu0:0) # Rank 1 |
| 418 | |
| 419 | """ |
| 420 | mode = CollectiveComm.Mode.ALL_REDUCE_SUM |
| 421 | return collective_comm(inp, mode, group, device) |
| 422 | |
| 423 | |
| 424 | def all_reduce_max( |