Support dist.all_gather related collective ops

[zpcore]

Add dynamo/nondynamo support for torch.distributed.all_reduce and torch.distributed.all_gather_into_tensor.

Motivation: We want to deprecate the collective ops in xla_model.py and be consist with the torch.distributed.

Issue: dist.all_reduct doesn't work with dynamo openxla backend at this time.

[zpcore]

Hi @JackCaoG , I commented out assert met.metric_data("ExecuteTime")[0] == 1 in the test_traceable_collectives.py since the value changed to 3 for all_gather ops due to upstream changes. Shall we enable the test first so we can run the test?

[zpcore]

I have no idea how this works:

xla/torch_xla/distributed/xla_backend.py

Lines 71 to 72 in 37312c1

	xm.all_reduce(reduce_type, tensors, groups=self._mesh, pin_layout=False)
	return _ret_work(tensors)

xm.all_reduce should return a tensor instead of modifying inputs argument.

I will probably give up supporting dist.all_reduce in this PR.

Update: turns out for output_tensor=dist.all_reduce(intput_tensor...)both intput_tensor and output_tensor got updated. We have to use the input_tensor as the final results for the nondynamo path.

[rpsilva-aws]

Motivation: We want to deprecate the collective ops in xla_model.py and be consist with the torch.distributed.

@zpcore This makes sense, and we would like to help close that gap. Do you have a more descriptive motivation and/or opened issue to migrate the remaining collective ops? I tried finding, but checking in first before I create one. cc: @miladm @tengyifei