feat(python, core): support mutable bucket tensors #271
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
NOBLES5E
reviewed
Oct 12, 2021
NOBLES5E
reviewed
Oct 12, 2021
NOBLES5E
reviewed
Oct 12, 2021
NOBLES5E
reviewed
Oct 12, 2021
wangraying
changed the title
NOBLES5E
reviewed
Oct 20, 2021
| Create a zero gradient tensor for the current parameter if not exist. | ||
|
|
||
| Returns: | ||
| The original tensor. |
wangraying
changed the title
|
close #287 |
NOBLES5E
reviewed
Oct 27, 2021
| @@ -10,29 +10,54 @@ | |||
| @gorilla.patches(torch.Tensor, filter=lambda name, obj: "bagua" in name) | |||
| class BaguaTensor: | |||
| """ | |||
| This class patch `torch.Tensor <https://pytorch.org/docs/stable/tensors.html?highlight=tensor#torch.Tensor>`_ with additional methods. | |||
| This class patch `torch.Tensor <https://pytorch.org/docs/stable/tensors.html?highlight=tensor#torch.Tensor>`_ | |||
| with additional methods. | |||
There was a problem hiding this comment.
Bagua Tensor features a proxy structure, where the actual tensor used by backend is accessed via a "Proxy Tensor".
The proxy tensor is registered in Bagua, whenever the Bagua backend needs a tensor (for example use it for communication), it calls
the `getter_closure` on the proxy tensor to get the tensor that is actually worked on. We call this tensor "Effective Tensor".
Their relation can be seen in the following diagram:
┌───────────────┐
│ Bagua Backend │
└──────▲────────┘
│
access
│
┌────────────────┼────────────────┐
│Bagua Tensor │ │
│ ┌───────┴────────┐ │
│ │ Proxy Tensor │ │
│ └───┬──────▲─────┘ │
│ │ │ │
│ setter_closure getter_closure │
│ │ │ │
│ ┌──────▼──────┴───────┐ │
│ │ Effective Tensor │ │
│ └─────────────────────┘ │
│ │
└─────────────────────────────────┘
For example, in the gradient allreduce algorithm, the effective tensor that
needs to be exchanged between machines is the gradient. In this case, we will
register the model parameters as proxy tensor, and register `getter_closure` to
be `lambda proxy_tensor: proxy_tensor.grad`. In this way, even if the gradient
tensor is recreated or changed during runtime, Bagua can still use the correct
tensor for communication, since the `proxy_tensor` serves as the root for
access and is never replaced.
The `setter_closure` is used to replace the effective tensor during runtime. It
is intended to be used to replace the effective tensor with customized
workflow.
getter_closure takes the registered the tensor as input and returns a PyTorch tensor.
setter_closure takes XXXX and returns XXX.
![github-actions[bot]](https://github.com/avatars/in/15368?s=60&v=4){kind=small}
NOBLES5E
reviewed
Oct 28, 2021
bagua/torch_api/tensor.py
Outdated
| """ | ||
| Sets the underlying storage using an existing `torch.Storage <https://pytorch.org/docs/stable/storage.html?highlight=storage>`_. | ||
| Sets the underlying storage for the tensor returned by :meth:`bagua_getter_closure` with an existing |
There was a problem hiding this comment.
use the terminology we defined in class documentation
There was a problem hiding this comment.
similar for other methods' doc
NOBLES5E
pushed a commit
that referenced
this pull request
Oct 28, 2021
BREAKING CHANGE: `BaguaTensor::bagua_ensure_grad` returns the tensor itself now
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
BREAKING CHANGE:
BaguaTensor::bagua_ensure_gradreturns the tensor itself nowBaguaTensor::bagua_set_storageis renamed toBaguaTensor::bagua_set_registered_storage