feat: support broadcasting picklable objects in idist.broadcast

[zongyang078]

Related to #1757, prerequisite for #3789. Follows @vfdev-5's suggestion in #3789 (comment).

Description:

idist.broadcast only handles torch.Tensor, numbers, str and pathlib.Path, and raises TypeError otherwise, while idist.all_gather already falls back to an object collective. This gives broadcast the same capability via torch.distributed.broadcast_object_list.

• _do_broadcast_object_list on ComputationModel: dist.broadcast_object_list for native, hvd.broadcast_object for horovod, NotImplementedError for xla (as with _do_all_gather_object)
• ComputationModel.broadcast dispatches to it for any input that is not a tensor/number/str/Path
• Any in the idist.broadcast typehint; broadcast and all_gather docstrings now document picklable-object support

_setup_placeholder now also returns the input type tag agreed on by every rank. An object input has no meaningful placeholder, so a non-source rank cannot decide from its own value which collective the source is about to call — it has to branch on the negotiated tag instead. Details in the code comments.

Behaviour change: broadcast no longer raises TypeError("Unhandled input type"). The test asserting that now asserts a successful object broadcast, except on xla where it asserts NotImplementedError.

Check list:

• New tests are added (if a new feature is added)
• New doc strings: description and/or example code are in RST format
• Documentation is updated (if required)

[zongyang078]

@vfdev-5 one design question. broadcast_object_list pickles tensors together with their device, so a CUDA tensor sent from rank 0 arrives on the other ranks still pointing at rank 0's device — all_gather's object path behaves the same today.

I only documented it with a .. warning:: on both methods, since normalizing would mean idist rewrites the caller's payload. Happy to add it if you prefer — it matters for #3789, where otherwise every caller broadcasting a container of tensors repeats the fixup.

[vfdev-5]

@zongyang078 thanks for raising this point. I had to read the pytorch docstring until the end yesterday :(
They explicitly say that:

Calling broadcast_object_list() with GPU tensors is not well supported and inefficient as it incurs GPU -> CPU transfer since tensors would be pickled. Please consider using broadcast() instead.

(https://docs.pytorch.org/docs/2.14/distributed.html#torch.distributed.broadcast_object_list)

I'm now hesitating to use it as it may not output what user is expecting...
Maybe, at some point we could simply code ourselves broadcast_object_list using broadcast and pytree flatten/unflatten... Let me think a bit about it.

[zongyang078]

Agreed — the CPU round trip is worse than just the device mismatch.

I looked into the pytree route: treespec_dumps returns a str (confirmed back to 2.2), so the structure can go through the existing broadcast string path and the leaves through the existing tensor/number path. That means no new backend method at all, tensors stay on each rank's own device since every rank builds its own placeholder, and it works on xla too, unlike broadcast_object_list.

I'm not sure how you'd weigh the trade-offs, though. It costs one collective per leaf plus one for the spec instead of a single one. And it narrows the feature: a leaf that is not a tensor/number/str — a custom class instance, say — is a pytree leaf, so keeping "any picklable object" would need a pickle fallback for those leaves only, with tensors still taking the native path. torch.utils._pytree is also private, and ignite has no torch-private imports today.

If it would help, I can push it as a second commit here so the two approaches can be compared side by side.

[vfdev-5]

@zongyang078 thanks for digging into that!
I agree that the route "pytree flatten -> collective op -> pytree unflatten" takes N (number of leaves) of collective ops which is perf-wise is not ideal. IMO, the best we can do following this route is to group compatible leaves (all numbers, all strings, all same shape/dtype tensors) and perform less collective ops, but well...

I was then checking your original PR with specification for list/mappings and you use all_gather object list route, but as we may doubt this torch.dist.all_gather_object is also doing GPU->CPU route (due to pickle): https://docs.pytorch.org/docs/2.14/distributed.html#torch.distributed.all_gather_object

I would like to run some benchmarks: A vs B where (if you would like to help with that)
A: input is a list of tensors with different shapes: [T1, T2, T3, T3] on the rank 0 and we do torch dist.broadcast src=0 to 2 or 4 or 8 ranks (depending on your available hardware). Output should be: on every rank we have [T1, T2, T3, T4] with all tensors on rank device: T_i.device == cuda:rank (on this point I'm not very sure as each rank sees its device as cuda:0 but in reality it was set in the beginning torch.cuda.set_device(self._local_rank))
B: same input (list of tensors with different shapes: [T1, T2, T3, T3] on the rank 0) and we do torch dist.broadcast_object_list src = 0 (if needed on each rank we move the tensors to theirs appropriate ranks; we should first verify whether moved tensors are not on their rank device).