(use_hybrid: bool = False, tp_size: int = None)
| 58 | |
| 59 | |
| 60 | def get_device_mesh(use_hybrid: bool = False, tp_size: int = None): |
| 61 | world_size = get_world_size() |
| 62 | n_nodes = get_nodes() |
| 63 | replicate_size = n_nodes if use_hybrid else 1 |
| 64 | |
| 65 | assert (world_size % ((tp_size or 1) * replicate_size)) == 0 |
| 66 | shard_size = world_size // (replicate_size * (tp_size or 1)) |
| 67 | if use_hybrid: |
| 68 | if tp_size is None: |
| 69 | mesh_shape = (replicate_size, shard_size) |
| 70 | mesh_dim_names = ('replicate', 'dp') |
| 71 | else: |
| 72 | mesh_shape = (replicate_size, shard_size, tp_size) |
| 73 | mesh_dim_names = ('replicate', 'dp', 'tp') |
| 74 | else: |
| 75 | if tp_size is None: |
| 76 | mesh_shape = (shard_size, ) |
| 77 | mesh_dim_names = ('dp', ) |
| 78 | else: |
| 79 | mesh_shape = (shard_size, tp_size) |
| 80 | mesh_dim_names = ('dp', 'tp') |
| 81 | device_mesh = init_device_mesh( |
| 82 | 'cuda', mesh_shape, mesh_dim_names=mesh_dim_names) |
| 83 | return device_mesh |
no test coverage detected