Get the data parallel group the caller rank belongs to.
()
| 699 | |
| 700 | |
| 701 | def _get_data_parallel_group(): |
| 702 | """Get the data parallel group the caller rank belongs to.""" |
| 703 | assert dist.is_initialized(), 'dist is not initialized' |
| 704 | global mpu |
| 705 | if mesh_device is not None: |
| 706 | return mesh_device.get_group(mesh_dim="data_parallel") |
| 707 | if mpu is not None: |
| 708 | if hasattr(mpu, 'initialize_sequence_parallel'): |
| 709 | return None |
| 710 | else: |
| 711 | return mpu.get_data_parallel_group() |
| 712 | |
| 713 | # Return the clone of dist world group |
| 714 | return _clone_world_group() |
| 715 | |
| 716 | |
| 717 | def _get_data_parallel_group_ranks(): |
no test coverage detected