Collective ops in optimum/neuron/accelerate
#1042
Conversation
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
tengomucho
left a comment
tengomucho
left a comment
There was a problem hiding this comment.
can you provide more information about why this change is necessary and what does it affect?
I do not see the training workflow running, and I do not really know what this changes for training.
| """ | ||
| Broadcasts arbitrary objects across XLA-distributed processes. | ||
| Returns the object from the source rank on all ranks. | ||
| If `groups` is specified, broadcast is done separately in each group, and the `src` rank is relative to each group. |
There was a problem hiding this comment.
Consider mentioning that object needs to be pickle compatible.
It's a bunch of features required to transfer non-tensor values during training. It is needed for GRPO. I just broken down the GRPO PR (#1020 ) into smaller PRs. |
What does this PR do?
Provide collective ops on arbitrary python objects.