| import torch.distributed as dist |
| from torch.distributed.device_mesh import init_device_mesh |
|
|
| _SP_MESH = None |
| _DP_MESH = None |
| _TP_MESH = None |
| _SP_GROUP = None |
| _DP_GROUP = None |
| _TP_GROUP = None |
| _SP_WORLD_SIZE = None |
| _DP_WORLD_SIZE = None |
| _TP_WORLD_SIZE = None |
|
|
|
|
| def setup_sp(sp_size): |
| world_size = dist.get_world_size() |
| assert world_size % sp_size == 0 |
| dp_size = world_size // sp_size |
| device_mesh = init_device_mesh( |
| 'cuda', (dp_size, sp_size), mesh_dim_names=('dp', 'sp')) |
|
|
| global _SP_MESH, _DP_MESH |
| _SP_MESH = device_mesh['sp'] |
| _DP_MESH = device_mesh['dp'] |
|
|
| global _SP_GROUP, _DP_GROUP |
| _SP_GROUP = device_mesh.get_group('sp') |
| _DP_GROUP = device_mesh.get_group('dp') |
|
|
|
|
| def setup_tp(tp_size): |
| world_size = dist.get_world_size() |
| assert world_size % tp_size == 0 |
| dp_size = world_size // tp_size |
| device_mesh = init_device_mesh( |
| 'cuda', (dp_size, tp_size), mesh_dim_names=('dp', 'tp')) |
|
|
| global _TP_MESH, _DP_MESH |
| _TP_MESH = device_mesh['tp'] |
| _DP_MESH = device_mesh['dp'] |
|
|
| global _TP_GROUP, _DP_GROUP |
| _TP_GROUP = device_mesh.get_group('tp') |
| _DP_GROUP = device_mesh.get_group('dp') |
|
|
|
|
| def setup_dp(): |
| world_size = dist.get_world_size() |
| device_mesh = init_device_mesh( |
| 'cuda', (world_size, ), mesh_dim_names=('dp', )) |
|
|
| global _DP_MESH |
| _DP_MESH = device_mesh['dp'] |
|
|
| global _DP_GROUP |
| _DP_GROUP = device_mesh.get_group('dp') |
|
|
|
|
| def setup_parallel(sp_size=1, tp_size=1): |
| assert not (sp_size > 1 and tp_size > 1), \ |
| ('DeepSpeed Sequence Parallel can not be used with ' |
| 'Megatron-LM Tensor Parallel') |
|
|
| if sp_size > 1: |
| setup_sp(sp_size) |
| elif tp_size > 1: |
| setup_tp(tp_size) |
| else: |
| setup_dp() |
|
|
|
|
| def get_dp_mesh(): |
| return _DP_MESH |
|
|
|
|
| def get_dp_group(): |
| return _DP_GROUP |
|
|
|
|
| def get_dp_world_size(): |
| global _DP_WORLD_SIZE |
| if _DP_WORLD_SIZE is not None: |
| return _DP_WORLD_SIZE |
| if not dist.is_initialized() or (_DP_GROUP is None): |
| _DP_WORLD_SIZE = 1 |
| else: |
| _DP_WORLD_SIZE = dist.get_world_size(_DP_GROUP) |
| return _DP_WORLD_SIZE |
|
|
|
|
| def get_sp_mesh(): |
| return _SP_MESH |
|
|
|
|
| def get_sp_group(): |
| return _SP_GROUP |
|
|
|
|
| def get_sp_world_size(): |
| global _SP_WORLD_SIZE |
| if _SP_WORLD_SIZE is not None: |
| return _SP_WORLD_SIZE |
| if not dist.is_initialized() or (_SP_GROUP is None): |
| _SP_WORLD_SIZE = 1 |
| else: |
| _SP_WORLD_SIZE = dist.get_world_size(_SP_GROUP) |
| return _SP_WORLD_SIZE |
|
|
|
|
| def get_tp_mesh(): |
| return _TP_MESH |
|
|
|
|
| def get_tp_group(): |
| return _TP_GROUP |
|
|
|
|
| def get_tp_world_size(): |
| global _TP_WORLD_SIZE |
| if _TP_WORLD_SIZE is not None: |
| return _TP_WORLD_SIZE |
| if not dist.is_initialized() or (_TP_GROUP is None): |
| _TP_WORLD_SIZE = 1 |
| else: |
| _TP_WORLD_SIZE = dist.get_world_size(_TP_GROUP) |
| return _TP_WORLD_SIZE |
|
|
