Issue 1:
Steps to reproduce the behavior:

Call torch.ceil on a sparse tensor

Below is a code sample:

import torch

i = torch.cuda.LongTensor([[0, 1, 1],

[2, 0, 2],

[2, 0, 2]])

v = torch.cuda.FloatTensor([3.2, -44.5, -5.1])

s = torch.sparse_coo_tensor(i, v, torch.Size([2, 4, 4]), device=torch.device('cuda'))

print(torch.ceil(s))

Actual behavior

Traceback (most recent call last):

File "sparse_floor.py", line 9, in

print(torch.ceil(s))

RuntimeError: Could not run 'aten::ceil.out' with arguments from the 'SparseCUDA' backend. 'aten::ceil.out' is only available for these backends: [CPU, CUDA, Named, Autograd, Profiler, Tracer].

Expected behavior

ceil of the sparse tensor should get calculated and output as a sparse tensor

Environment

Versions of relevant libraries:

[pip3] numpy==1.19.1

[pip3] torch==1.8.0a0

[pip3] torchtext==0.8.0a0+7e267d2

[pip3] torchvision==0.8.0a0+be8192e

[conda] blas 1.0 mkl

[conda] cudatoolkit 10.2.89 hfd86e86_1

[conda] magma-cuda102 2.5.2 1 pytorch

[conda] mkl 2020.1 217

[conda] mkl-include 2020.1 217

[conda] mkl-service 2.3.0 py38he904b0f_0

[conda] mkl_fft 1.1.0 py38h23d657b_0

[conda] mkl_random 1.1.1 py38h0573a6f_0

[conda] numpy 1.19.1 py38hbc911f0_0

[conda] numpy-base 1.19.1 py38hfa32c7d_0

[conda] torch 1.8.0a0 dev_0

[conda] torchtext 0.8.0a0+7e267d2 pypi_0 pypi

[conda] torchvision 0.8.0a0+be8192e pypi_0 pypi

PyTorch Version: 1.8.0a0

OS: Linux

How you installed PyTorch (conda, pip, source): source

Build command you used (if compiling from source): python setup.py develop

Python version: 3.8

CUDA/cuDNN version: 10.2

cc @aocsa @nikitaved @pearu @mruberry

Issue 2:
-- with @Tierex, @mrshenli, @kiukchung, @pritamdamania, @lw, @agolynski , @osalpekar, @zhaojuanmao, @rohan-varma

🚀 Feature

Concept and a new set of rpc APIs in torch.distributed.rpc aiming to help users construct distributed applications more easily, apis maily comprise of:

Higher level torch rpc APIs, including group creation, scoped service registering and easier value sharing among processes.

Succinct init methods to for apps using both rpc and process groups with a heterogeneous topology (roles),

Ensure all the above works well when launching these apps using TorchElastic, or with a simple launch script

such as torch.distributed.launch.

Motivation

This RFC is a refined and extended version of #41425

Homogeneous/Heterogeneous Torch RPC Apps

This section has been described in #41425 with details, and therefore it will be omitted here.

Service registering

Currently torch.distributed.rpc has only defined three basic point to point style APIs: rpc_sync, rpc_async, and rpc_remote. They require users to manually specify a specific action: func to perform, and a specific destination: to where this action will be performed on:

torch.distributed.rpc.rpc_sync(to, func, args=None, kwargs=None)
torch.distributed.rpc.rpc_async(to, func, args=None, kwargs=None)
torch.distributed.rpc.remote(to, func, args=None, kwargs=None)

However, in most real application scenarios, a single process does not provide sufficient computing power to provide a complete service, therefore users are usually dealing with a complex of different processes serving different roles, such as a parameter server architecture diagram below:

If users want to invoke push() or the pull() method, it would be most sensible to provide a mechanism to register these two methods along with the method provider (i.e: job queue, model distributors) in the scope where they would like to use them, and let processes query for the registered method.

Compared to the mechanism mentioned in #41425

group-rpc calls (e.g. call the same function on all trainers)

which can also provide a similar function as "Service registering", "Service registering" enables users to call different processes with different roles, as well as hiding the actual service provider and make detail implementations transparent to users. This is especially important if users wants to maintain modularity or write complex applications.

Resource sharing

Although RPC is an "action" based paradigm, compared to the "resource" based paradigm used by http & restful apis, sometimes it could be useful to borrow some ideas from there. For example, to implement a simple synchronization mechanism using a shared value in a group of processes:

def process1():
rpc.init_rpc(...)
some_signal = Switch(False) # a switch class
rpc.pair("signal", some_signal) # expose this variable in the implicit global rpc group
sleep(10) # do some initialization
some_signal.set(True)
# synchronized

def process2():
rpc.init_rpc(...)
some_signal = Switch(False) # a switch class
rpc.pair("signal2", some_signal) # expose this variable in the implicit global rpc group
sleep(5) # do some initialization
some_signal.set(True)
while not rpc.get_paired("signal"):
sleep(1)
# synchronized

def process2():
rpc.init_rpc(...)
while not rpc.get_paired("signal") or not rpc.get_paired("signal2"):
sleep(1)
# synchronized

compared to

signal1 = False
signal2 = False

def get_signal1():
return signal1

def get_signal2():
return signal2

def process1():
rpc.init_rpc(...)
sleep(10) # do some initialization
some_signal = True

# synchronized

def process2():
rpc.init_rpc(...)
sleep(5) # do some initialization
while not rpc.rpc_sync("proc1", get_signal1):
sleep(1)
# synchronized

def process2():
rpc.init_rpc(...)
while (not rpc.rpc_sync("proc1", get_signal1) or
not rpc.rpc_sync("proc1", get_signal2)):
sleep(1)
# synchronized

is simpler, because in the second example, users will have to roll out a new function for each signal they would like to look up, or

they have to implement a custom lookup function taking a signal name and return a signal value, like below:

signals = {}
def lookup(signal_name: str):
return signals.get(signal_name, None)

Therefore, it is reasonable to provide a simple api like rpc.pair and rpc.get_paired to expose the value to the respective scope, so that processes can easily look up other values.

Groups

If we want to either share a resource, or access a service, by RPC, it is crutial to limit their effect scope, so we can define services / pair values in a controllable, safe range, without overwriting other paired values and creating hidden conflicts. RPC has a implicit global group, as stated in #41425 :

rpc not really needing a group (other than the implicit global one).

Again, we must be clear that rpc_sync, rpc_async, and rpc_remote are all operating on a global scope, and only having one scope is a severe disadvantage in OOP designs.

There is one difference in the group conception between this RFC and #41425, that is group may consists of different roles, for example, in the parameter server architecture above, users may implement each sub-module and group them by their kind, or consider all submodules as members of the parameter server RPC group:

Pitch

Role based auto initialization

Ideally we'd like the application code to look like the following:

def master_main():
# do something
pass

def trainer_main():
# do something
pass

if __name__ == "__main__":
rpc.init_role(backend=backend, backend_options)
if rpc.get_worker_info().role == "master":
master_main()
elif rpc.get_worker_info().role == "trainer":
trainer_main()
rpc.shutdown()

So users do not specify world size and process rank by hand, instead they should pass how many roles to launch to torchelastic launcher or a simpler launcher scipt similar to torch.distributed.launch if they don't want to configure torchelastic with etcd:

>>> python -m torchelastic.distributed.launch
--nnodes=$NUM_NODES
--nproc_per_node=$NUM_TRAINERS
--role=trainer:4 # launch 4 trainer processes
--role=master:1 # launch 1 master process
--rdzv_id=$JOB_ID
--rdzv_backend=etcd
--rdzv_endpoint=$ETCD_HOST:$ETCD_PORT
YOUR_TRAINING_SCRIPT.py (--arg1 ... train script args...)

And torchelastic or the launching script will set RANK, WORLD_SIZE, NODE_RANK, NODE_DIMS, LOCAL_RANK, ROLE, ROLE_NUM, ROLE_NUMS, ROLE_RANK, NAME, MASTER_ADDR, MASTER_PORT in the env vars, and rpc.init_role will pick these env vars up.

After the application is initialized the following hold true:

Every process (regardless of role) can invoke rpc APIs on one another.

Every process has a global rank (e.g. one that was passed as RANK env var) that can be obtained via dist.get_rank()

Each node is assigned a rank. (this is equivalent to the elastic agent’s rank - a.k.a GROUP_RANK in torchelastic)

The global rank can be derived from the following formula:

Where node_dims[i] returns the number of local processes on node rank i.

But cannot be derived from ROLE_NUMS since there is no explicit garantee on the order of roles.

the launching utility (torchelastic or a simple script launcher) will set RANK, WORLD_SIZE, NODE_RANK(GROUP_RANK), NODE_DIMS, LOCAL_RANK, ROLE, ROLE_NUM, ROLE_NUMS, ROLE_RANK, NAME, MASTER_ADDR, MASTER_PORT in the env vars.

ROLE_NUM is the number of processes in current role, and ROLE_NUMS is the number of processes of each role.

NODE_DIMS is a comma seperated string of >0 integers: 2,3,4,5

ROLE_NUMS is a comma sepeated string of role names and >0 integers: trainer,4,master,1

There is no trailling comma in them.

Role names must not contain commas.

Every process has a unique name in form name = {role}:{role_id}, {role_id} is a number in [0, num_workers_in_role),

The WorkerInfo object returned by torch.distributed.rpc.get_worker_info() should contain 4 properties: name, id, role, role_id. id is equal to the global rank, users can look up role and role_id of current process by passing None to get_worker_info(). If users are calling rpc.init_rpc instead of rpc.init_role to initialize, role and role_id should be None.

Clarification or different ranks:

! RANK depicts the global rank of the current process among all launched processes, it is an int in [0, num_total_processes)

! NODE_RANK depicts the rank of container / physical server, it is an int in [0, num_container_or_server)

! LOCAL_RANK depicts the rank of the current process among processes on this container/server, it is an int in [0, NODE_DIMS[NODE_RANK])

! ROLE_RANK depicts the rank of current process among processes in this role, it is an int in [0, ROLE_NUMS[ROLE])

Role-based rpc APIs

Since we are working with the concept of roles this is really a syntactic sugar that allows users to perform group rpc calls on all workers belonging to a role rather than having to for-loop around them. It can trivially be implemented by first getting all the worker names for the role (or derive it by using role_world_size = ROLE_NUMS[ROLE] = ROLE_NUM)

ret_futures = {}
role = os.environ['ROLE']
role_num = int(os.environ['ROLE_NUM'])
for idx in range(0, role_num):
name=f"trainer:{idx}"
ret_futures[name] = rpc.rpc_async(to=name, func=foobar, args=(...)))

# versus
ret_futures = rpc.rpc_async_role(role="trainer", selector=lambda _: True, func=foobar, args=(...))

# selector is a function which takes in a role rank and returns True for selecting it, False for deselecting it.
# could be None, and selecting all processes with this role

Note The original design also included an rpc.wait_all() API for completeness. This has already been implemented in the form of torch.futures.collect_all (https://github.com/pytorch/pytorch/blob/master/torch/futures/__init__.py#L88)

Create groups

In this section:

MPI style APIs will be refered to as "collective communication primitives" or "coll comm" for short.

"sub" means a group whose members is a subset of a larger group of same kind.

Collective group

In current distributed apis, we can create a collective group for coll comm using:

torch.distributed.new_group(ranks=None, timeout=datetime.timedelta(0, 1800), backend=None)

this function takes in a list of raw ranks, is blocking, and requires all mentioned processes in ranks to enter. However, many users find coll comm extremely useful in many scenarios, therefore, we decide to import all of its functions into the scope of rpc,

so that instead of using a list of raw ranks, users may now create a collective group with:

a list of names, such as ["worker:0", "worker:1", "manager:0"]

a role name, with a index selector: role="worker", selector=lambda id: id in (0, 1)

api has the following signature:

rpc.new_collective_group(names, role: str = None, selector: Callable = None, timeout=datetime.timedelta(0, 1800), backend=None) -> CollectiveGroup

CollectiveGroup has the same coll comm apis as torch.distributed, except the group keyword argument is the created group itself, and dst and src are process names instead of global ranks:

def send(self, tensor, dst, tag=0)
def recv(self, tensor, src=None, tag=0)
def isend(self, tensor, dst, tag=0)
def irecv(self, tensor, src, tag=0)
...

RPC group

RPC group is a name scope, each process may join multiple RPC groups, and each RPC group may contain multiple processes, RPC group is a place where users may:

share variables only among members by pairing their value with a key to the group

expose services so that only its member processes may access.

The default RPC group is the rpc world, that is the rpc module itself:

torch.distributed.rpc.rpc_sync(to, func, args=None, kwargs=None)
torch.distributed.rpc.rpc_async(to, func, args=None, kwargs=None)
torch.distributed.rpc.remote(to, func, args=None, kwargs=None)

are all operating on this implicit global group.

In order to create a sub rpc group:

rpc.new_rpc_group(self, names, role: str = None, selector: Callable = None, group_name: str) -> RpcGroup

compared to the collective group, group_name is necessary here because it will identify the scope of shared variables and exposed services.

RpcGroup class, as well as the implicit global group rpc, provides the following APIs:

This part is open to discussion

# pair a value to this group, the same key cannot be used twice, otherwise an exception must be thrown
pair(key: str, Value: Any) -> None

# unpair a value from this group, the key must exist, otherwise an exception must be thrown
unpair(key: str) -> None

# query whether a key has been used to pair
is_paired(key: str) -> bool

# get a paired value, if it exist, otherwise an exception must be thrown
get_paired(key: str) -> Any

# expose a service, the service function will be executed when called
# could be a global scope function, a class method, a local function returned by a decorator, or lambda,
# or torch script, etc.
register(key: str, func: Callable) -> None

# deregister an exposed service, the key must exist, otherwise an exception must be thrown
deregister(key: str) -> None

# query whether a service has been registered:
is_registered(key: str) -> bool

# call a service
registered_sync(key: str, args=None, kwargs=None)
registered_async(key: str, args=None, kwargs=None)
registered_remote(key: str, args=None, kwargs=None)

# destroy the group, global rpc does not have this function, since it has "shutdown"
# also deregister the group along with values, services from the global discovery manager
destroy()

# info accessors
# returns a list of member names
get_members() -> List[str]

# judge whether the target process is a member of this group, by default it is the current process
is_member(name: str=None) -> bool

# group name
get_group_name() -> str

# number of group members
size() -> int

the pair/register process happens with two stages:

register to local map

push the result to a global discovery manager

the unpair/deregister process happens with two stages:

degister from local map

push the result to a global discovery manager

the get_paired/registered call process happens with two stages:

look up the local discovery cache (every process has it)

if cache does not contain the address: (group, process name, key, query_type), query the global discovery manager

otherwise go to the target address

if the target group has been destroyed / target value deregistered, return look up failed exception

otherwise go on as normal, return the result or execute the service

In order to discover a shared value / an exposed service, a register thread / process must be started once rpc.init_rpc or rpc.init_auto has been called, since torchelastic tears down and restart all processes when at least one processes becomes unhealthy, it is reasonable to regard the period in between as stable.

Therefore, we can create a register service on a fixed process such as the process with rank=0, or use a more sophisticated service management infrastructure such as zookeeper, we should allow users to choose the discovery backend by specifying it in the ProcessGroupRpcBackendOptions passed to rpc.init_rpc and rpc.init_auto.

To prevent the explosion of discovery queries, a local query cache should be placed on each process and store looked up values, outdated address values (caused by unpair or deregister) will be cleared when get_paired or registered_sync etc. returns a lookup-failed signal.

RPC Info Accessors

Some times users may want to dive deeper and see the actual mapping relation. rpc.get_role_map() will return a dictionary mapping role name to global process rank.

This part is open to discussion

Launching utilities

Since torchelastic currently is not included in the default pytorch installation, a new launching utility similar to torch.distributed.launch should be included, its name could be torch.distributed.launch_roles, or it could be merged into torch.distributed.launch. this utility should be able to takein a list of roles and spawn subprocesses, allocating roles to global ranks like:

worker:0 -> Process 0, Node 0
worker:1 -> Process 1, Node 0
worker:2 -> Process 2, Node 0
worker:3 -> Process 0, Node 1
manager:0 -> Process 1, Node 1
manager:1 -> Process 2, Node 1

when invoked with:

Node 0:
python -m torch.distributed.launch --nproc_per_node=3 --role=worker:4 --role=manager:2
--nnodes=2 --node_rank=0 --master_addr="192.168.1.1"
--master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
and all other arguments of your training script)

Node 1:
python -m torch.distributed.launch --nproc_per_node=3 --role=worker:4 --role=manager:2
--nnodes=2 --node_rank=1 --master_addr="192.168.1.1"
--master_port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3
and all other arguments of your training script)

and nproc_per_node * nnodes >= num of all roles, otherwise an exception must be thrown.

Failure Handling Behaviors (naive launching tool)

No failure could be dealt with, one process fails, all processes fail.

Failure Handling Behaviors (torchelastic)

This section describes the types of failures and who/how those failures are handled, when users are launching the application with torchelastic.

We focus on worker process(es) failure since a node failure can be viewed as multiple worker failures. In general TorchElastic views failures as two scaling events: scale-down + scale-up. In TorchElastic the world size is fixed between "rounds" of quorum where "rounds" is defined as the state of the world between two rendezvous versions. Between rounds ALL existing workers are shut down and restarted, hence TorchElastic follows a "all-or-nothing" model when it comes to dealing with faults.

There are three types of torch applications that we care about:

RPC only (no process groups)

RPC + DDP

RPC + DDP + Pipelining

In all three cases the following behavior on failures is guaranteed by TorchElastic:

Worker process(es) fail

elastic agent detects the failure

elastic agent tears down all other workers on the same host

elastic agent enters a re-rendezvous

other elastic agents are notified of the re-rendezvous event

other elastic agents tear down all their respective local workers

all elastic agents re-rendezvous (e.g. next round) and get assigned a rank

each elastic agent computes the worker ranks based on the agent ranks

all elastic agents restart the worker processes

Note: This implies that work between checkpoints are lost.

Note on node failure: TorchElastic is NOT a scheduler hence it cannot deal with replacement of NODES. We rely on the scheduler to replace failed nodes (or containers). Assuming the scheduler replaces the failed node the following behavior is observed on node failures:

Node(s) fail and are replaced by the scheduler

Elastic agents are started on the replaced nodes and attempt to join the rendezvous

The surviving agents are notified of such event and they all tear down their local workers and re-rendezvous

A new version of rendezvous (e.g. round) is created and the worker processes are started again.

cc @pietern @mrshenli @kiukchung @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @rohan-varma @xush6528 @jjlilley @osalpekar @jiayisuse