torchtnt.utils.test_utils.spawn_multi_process

torchtnt.utils.test_utils.spawn_multi_process(world_size: int, backend: str, test_method: Callable[[], TReturn], *args: Any) → Dict[int, TReturn]

Spawn single node, multi-rank function. Uses localhost and free port to communicate.

Parameters:

• world_size – number of processes
• backend – backend to use. for example, “nccl”, “gloo”, etc
• test_method – callable to spawn. first 3 arguments are rank, world_size and mp output dict
• args – additional args for func

Returns:

A dictionary of rank -> func return value