Initialize DeviceMesh from torchrun env vars. Returns a context dict with is_distributed, world_size, is_main, device_mesh.
()
| 12 | |
| 13 | |
| 14 | def init_distributed_from_env() -> Dict[str, object]: |
| 15 | """Initialize DeviceMesh from torchrun env vars. |
| 16 | Returns a context dict with is_distributed, world_size, is_main, device_mesh. |
| 17 | """ |
| 18 | world_size = int(os.environ.get('WORLD_SIZE', '1')) |
| 19 | is_distributed = world_size > 1 and torch.cuda.is_available() |
| 20 | |
| 21 | if is_distributed: |
| 22 | device_mesh = init_device_mesh("cuda", (world_size,), mesh_dim_names=('fsdp',)) |
| 23 | else: |
| 24 | device_mesh = None |
| 25 | |
| 26 | rank = device_mesh.get_rank() if is_distributed else 0 |
| 27 | is_main = (rank == 0) |
| 28 | |
| 29 | return { |
| 30 | 'is_distributed': is_distributed, |
| 31 | 'world_size': world_size, |
| 32 | 'is_main': is_main, |
| 33 | 'device_mesh': device_mesh, |
| 34 | } |
| 35 | |
| 36 | |
| 37 | def prepare_model_for_distributed(model: torch.nn.Module, config: DistributedConfig, model_type: ModelType, device_mesh: DeviceMesh) -> torch.nn.Module: |