Function _custom_mesh_worker

tests/models/testing_utils/parallelism.py:171–227

Worker function for context parallel testing with a user-provided custom DeviceMesh.

(
    rank,
    world_size,
    master_port,
    model_class,
    init_dict,
    cp_dict,
    mesh_shape,
    mesh_dim_names,
    inputs_dict,
    return_dict,
)

def _custom_mesh_worker(
    rank,
    world_size,
    master_port,
    model_class,
    init_dict,
    cp_dict,
    mesh_shape,
    mesh_dim_names,
    inputs_dict,
    return_dict,
):
    """Worker function for context parallel testing with a user-provided custom DeviceMesh."""
    try:
        os.environ["MASTER_ADDR"] = "localhost"
        os.environ["MASTER_PORT"] = str(master_port)
        os.environ["RANK"] = str(rank)
        os.environ["WORLD_SIZE"] = str(world_size)

        # Get device configuration
        device_config = DEVICE_CONFIG.get(torch_device, DEVICE_CONFIG["cuda"])
        backend = device_config["backend"]
        device_module = device_config["module"]

        dist.init_process_group(backend=backend, rank=rank, world_size=world_size)

        # Set device for this process
        device_module.set_device(rank)
        device = torch.device(f"{torch_device}:{rank}")

        model = model_class(**init_dict)
        model.to(device)
        model.eval()

        inputs_on_device = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in inputs_dict.items()}

        # DeviceMesh must be created after init_process_group, inside each worker process.
        mesh = torch.distributed.device_mesh.init_device_mesh(
            torch_device, mesh_shape=mesh_shape, mesh_dim_names=mesh_dim_names
        )
        cp_config = ContextParallelConfig(**cp_dict, mesh=mesh)
        model.enable_parallelism(config=cp_config)

        with torch.no_grad():
            output = model(**inputs_on_device, return_dict=False)[0]

        if rank == 0:
            return_dict["status"] = "success"
            return_dict["output_shape"] = list(output.shape)

    except Exception as e:
        if rank == 0:
            return_dict["status"] = "error"
            return_dict["error"] = str(e)
    finally:
        if dist.is_initialized():
            dist.destroy_process_group()