example.py

def _register_comm_hook(
	self, 
	state: object, 
	hook: callable):

def fp16_compress_hook(
	process_group: object, 
	bucket: dist._GradBucket):

    compressed_tensor = 
		bucket.get_tensors()[0].to(torch.float16)

    fut = dist.all_reduce(
        compressed_tensor, 
	    group=process_group, 
        async_op=True
    ).get_future()

    def decompress(fut):
        return [fut.value()[0]
			.to(torch.float32)
			.div_(world_size)]

    return fut.then(decompress)