comm.py should maybe consider backend-specific support of different devices

[lucaslie]

Depending on the backend, distributed communication may only be supported on either CPU or GPU, see table here.

Right now, in comm.py communication is always done on the GPU, see here e.g.:

torchpack/torchpack/distributed/comm.py

Lines 32 to 34 in d3fda52

	# serialize
	if context.rank() == src:
	tensor = _serialize(obj).cuda()

I would suggest considering the backend-specific device support for both allgather() and broadcast() to ensure the functions are usable across multiple backends.

torch.distributed.broadcast_object_list and torch.distributed.all_gather_object might be a useful starting points for this.