Register ModelOutput as supported torch pytree nodes
Author: XuehaiPan
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

This registers `ModelOutput` as supported `torch` pytree nodes. Currently all subclasses of `ModelOutput` are already registered by `ModelOutput.__init_subclass__()`. This PR additionally registers `ModelOutput` itself as supported `torch` pytree nodes.

See also:

- #25357
- #25358
- https://github.com/pytorch/pytorch/pull/109684#discussion_r1347384106

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [X] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [X] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [X] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@sgugger

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @pacman100

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc and @younesbelkada

Documentation: @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
This looks important but I don't have the bandwidth to dive into it, could you take a look if you have bandwidth @ydshieh ?

Hi @XuehaiPan,

Could you elaborate a bit more on why we need to register `ModelOutput` itself: why #25358 isn't enough?



> Could you elaborate a bit more on why we need to register `ModelOutput` itself: why #25358 isn't enough?

@ydshieh

1. `ModelOutput` is a container-like type, a subclass of `OrderedDict`. This fits the definition of a pytree node rather than a leaf.
2. The PyTorch upstream provides a context manager to register `ModelOutput` itself and its subclasses as pytree nodes. When we exit the context manager in `__exit__` method, it will unregister the types that registered in the `__enter__` method.
https://github.com/pytorch/pytorch/blob/53a9ac534c777047abc2b6b293f238865592291e/torch/onnx/_internal/fx/dynamo_graph_extractor.py#L31-L103

    We have already registered the `ModelOutput`'s subclasses as pytree node type. So in the `__enter__` method, it will additionally register `ModelOutput` as pytree node type, and unregister it in the `__exit__` method. All subclasses of `ModelOutput` in the pytree node type registry will remain untouched. See my comment in https://github.com/pytorch/pytorch/pull/109684#discussion_r1347384106.

    I think deleting or updating the pytree node registry is dangerous. It could lead to potential bugs when the developer passes the treespec object out of the scope of the context manager.

    ```python
    with PyTreeExtensionContext():
        # in the context
        # ModelOutput is a container-like node type, it will live in the treespec object
        leaves, treespec = pytree.flatten(tree)

    # out of the context
    # ModelOutput is not in the pytree node type registry and it will be considered as leaf
    # but in the treespec object, ModelOutput a node type
    pytree.tree_unflatten(leaves, treespec)  # -> KeyError: cannot found ModelOutput in pytree.SUPPORTED_NODES
    ```

     I could either do not register the `ModelOutput` type in the context manager in the PyTorch upstream, or register it as pytree node in the first place in transformers. I think it is more practical to register both `ModelOutput` and its subclasses as node type (e.g., for case `isinstance(obj, ModelOutput)`). 

Hi @XuehaiPan : I am not familiar with this, so would love to (and need) a bit more information 😅 

-  In the PyTorch upstream, `ModelOutput ` is already registered as as pytree nodes`?
    - and if so, why we need do this at both in `pytorch` and `transformers`
    - if they are not doing exactly the same job, what are the difference between `register in torch` and `register in transformers`

- `... in the __enter__ method, it will additionally register ModelOutput as pytree node type`:
    - so #25358 (dealing with subclasses) will register the parent class `ModelOutput` as pytree node type ...? (because of your work in  the torch upstream?)
    - and because of this, we need this PR, so `__exit__` won't unregister it?

> `... in the __enter__ method, it will additionally register ModelOutput as pytree node type`:
> 
> * so [Register ModelOutput subclasses as supported torch.utils._pytree nodes #25358](https://github.com/huggingface/transformers/pull/25358) (dealing with subclasses) will register the parent class `ModelOutput` as pytree node type ...? ~(because of your work in  the torch upstream?)~
> * and because of this, we need this PR, so `__exit__` won't unregister it?

@ydshieh This is correct. Let me elaborate this.

In the PyTorch upstream:

https://github.com/pytorch/pytorch/blob/359336e3e9a0f67974e53805b5207fbbbc149490/torch/onnx/_internal/fx/dynamo_graph_extractor.py#L93-L98

```python
        # All 'ModelOutput' subclasses are defined under module 'modeling_outputs'.
        named_model_output_classes = inspect.getmembers(
            modeling_outputs,
            lambda x: inspect.isclass(x)
            and issubclass(x, modeling_outputs.ModelOutput),
        )
```

The model output classes are determined by `issubclass(cls, ModelOutput)`. It will treat the `ModelOutput` itself as model output class.

```python
>>> issubclass(ModelOutput, ModelOutput)
True
```

```python
In [1]: from transformers import modeling_outputs

In [2]: import inspect

In [3]: named_model_output_classes = inspect.getmembers(
   ...:     modeling_outputs,
   ...:     lambda x: inspect.isclass(x)
   ...:     and issubclass(x, modeling_outputs.ModelOutput),
   ...: )

In [4]: dict(named_model_output_classes)
Out[4]: 
{'BackboneOutput': transformers.modeling_outputs.BackboneOutput,
 ...,
 'MoEModelOutputWithPastAndCrossAttentions': transformers.modeling_outputs.MoEModelOutputWithPastAndCrossAttentions,
 'ModelOutput': transformers.utils.generic.ModelOutput,
 'MultipleChoiceModelOutput': transformers.modeling_outputs.MultipleChoiceModelOutput,
 ...,
 'XVectorOutput': transformers.modeling_outputs.XVectorOutput}
```

In the PyTreeExtensionContext manager, it will ignore types that are already registered in the `pytree.SUPPORTED_NODES`.

https://github.com/pytorch/pytorch/blob/359336e3e9a0f67974e53805b5207fbbbc149490/torch/onnx/_internal/fx/dynamo_graph_extractor.py#L67-L72

```python
        if class_type in pytree.SUPPORTED_NODES or class_type in self._extensions:
            # PyTree node already registered.
            # E.g., `huggingface/transformer` registers `ModelOutput` as PyTree node after
            # https://github.com/huggingface/transformers/pull/25358.
            return
        self._extensions[class_type] = (flatten_func, unflatten_func)
```

Without this PR, the upstream code behaves like:

```python
# This command will register all subclasses of `ModelOutput` except `ModelOutput`
# itself as pytree node.
import transformers

# <<< at here
#     - all subclasses of `ModelOutput` are pytree node
#     - `ModelOutput` is not pytree node

# This command will add `ModelOutput` into `context._extension` because it is
# not registered as pytree node.
# All subclasses of `ModelOutput` will be ignored because they are already
# registered as pytree node by transformers.
context = PyTreeExtensionContext()

# <<< at here
#     - all subclasses of `ModelOutput` are pytree node
#     - `ModelOutput` is not pytree node
#     - context._extension = {ModelOutput: (..., ...)}

# This command will register elements in `context._extension` (which is only
# `ModelOutput`) as pytree node.
context.__enter__()

# <<< at here
#     - all subclasses of `ModelOutput` are pytree node
#     - `ModelOutput` is pytree node
#     - context._extension = {ModelOutput: (..., ...)}

# do something in the context
...

# This command will unregister elements in `context._extension` (which is only
# `ModelOutput`) as pytree node.
context.__exit__()

# <<< at here
#     - all subclasses of `ModelOutput` are pytree node
#     - `ModelOutput` is not pytree node
```

with this PR:

```python
# This command will register all subclasses of `ModelOutput` and `ModelOutput`
# itself as pytree node.
import transformers

# <<< at here
#     - all subclasses of `ModelOutput` are pytree node
#     - `ModelOutput` itself is also pytree node

# All subclasses of `ModelOutput` and `ModelOutput` itself will be ignored
# because they are already registered as pytree node by transformers.
# This is a no-op and `context._extension` is empty.
context = PyTreeExtensionContext()

# <<< at here
#     - all subclasses of `ModelOutput` are pytree node
#     - `ModelOutput` itself is also pytree node
#     - context._extension = {}  # empty

# This command will register elements in `context._extension` (which is empty)
# as pytree node. This is a no-op.
context.__enter__()

# <<< at here
#     - all subclasses of `ModelOutput` are pytree node
#     - `ModelOutput` is pytree node
#     - context._extension = {}  # empty

# do something in the context
...

# This command will unregister elements in `context._extension` (which is empty)
# as pytree node. This is a no-op.
context.__exit__()

# <<< at here
#     - all subclasses of `ModelOutput` are pytree node
#     - `ModelOutput` itself is also pytree node
```

@XuehaiPan Thank you a lot. I will take a look next week! This is really new to me 💪 !

@ydshieh Hi, any update for this?

Hi @XuehaiPan, thank you a lot of your super clean and detailed comment, it helps me a lot to get started. So yes, **`ModelOutput` itself will behave differently depending on if it is inside the context manager or outside it.**

> I think deleting or updating the pytree node registry is dangerous

I have 2 (general) questions:

  - I understand what you mentioned. But does this mean pytorch unstream has to register all standard container-like classes (list, dict, etc.) as pytree nodes?
    - (I guess so, but want to hear from you 😄 ) 
  - `ModelOutput` is not designed to be used directly (well, to be fair, within `transformers` code base).
    - All usage should inheritate `ModelOutput` and specify the attributes an output have), see, for example, https://github.com/huggingface/transformers/blob/929134bf65ac986c12c423c30b0db8a239f3b195/src/transformers/modeling_outputs.py#L25-L47
    - (it's kind abstract class, despite we can still use it directly, but that usage is outside `transformers` scope.)
    - Moreover, I would say it is super rare users will use `ModelOutput` directly  (in a manually way) to store their model output (without creating a subclass first). 
      - (and even more rare, this usage along with `PyTreeExtensionContext` context manager) 

Overall, **I am fine with the idea of registering `ModelOutput` as torch pytree nodes**, but would like the above 2 questions addressed. Then I can take a review on the changes themselves and request a core maintainer's review.

Thank you again @XuehaiPan !
 



@ydshieh

> * I understand what you mentioned. But does this mean pytorch unstream has to register all standard container-like classes (list, dict, etc.) as pytree nodes?
>   
>   * (I guess so, but want to hear from you 😄 )

Currently, the PyTorch upstream registers:

- `tuple`
- `list`
- `dict`
- `OrderedDict`
- classes created by `collections.namedtuple`
- `torch.return_types.*` (PyStructSequence, e.g., `torch.return_types.sort` for function `torch.sort()`)
- `torch.fx.immutable_collections.immutable_list`
- `torch.fx.immutable_collections.immutable_dict`

as pytree nodes.

We are going to add more container-like classes in the stdlib (`collections.defaultdict` and `collections.deque`) as type pytree nodes in the follow-up PRs in the PyTorch upstream.

For classes from third-party packages (non-Python stdlib, and non-PyTorch internals), users or developers need to manually register the classes as pytree types themselves. For example, we manually register the subclasses of ModelOutput in the transformers package.

------

> * `ModelOutput` is not designed to be used directly (well, to be fair, within `transformers` code base).
>   
>   * All usage should inheritate `ModelOutput` and specify the attributes an output have)
>
> Moreover, I would say it is super rare users will use `ModelOutput` directly (in a manually way) to store their model output (without creating a subclass first).
> * (and even more rare, this usage along with `PyTreeExtensionContext` context manager)

I agree that is a rare case. So we can resolve this issue in two ways:

- register ModelOutput as pytree node in transformers (this is what this PR does).
- not registering the ModelOutput as pytree node in the `PyTreeExtensionContext` context manager in the PyTorch upstream.

Both solutions are fine for me.

Let's go for `register ModelOutput as pytree node in transformers (this is what this PR does).` to align it with `OrderedDict` (it's parent). I will take a review on the changes 🤗 today or next Monday!

@XuehaiPan 

Could you provide a complete, self-contained (i.e. with all imports and variables defined) code snippet of the following (the one you provided before):

```python
with PyTreeExtensionContext():
    # in the context
    # ModelOutput is a container-like node type, it will live in the treespec object
    leaves, treespec = pytree.flatten(tree)

# out of the context
# ModelOutput is not in the pytree node type registry and it will be considered as leaf
# but in the treespec object, ModelOutput a node type
pytree.tree_unflatten(leaves, treespec)  # -> KeyError: cannot found ModelOutput in pytree.SUPPORTED_NODES
```

Here is a self-contained snippet with `torch==2.1.0` and `transformers==4.34.1`.

```python
import torch
import transformers
import torch.utils._pytree as pytree
from torch.onnx._internal.fx.dynamo_graph_extractor import _PyTreeExtensionContext as PyTreeExtensionContext
from transformers.modeling_utils import ModelOutput

print('===== before entering context =====')
leaves, treespec = pytree.tree_flatten(ModelOutput(a=1, b=2))  # ModelOutput is a leaf
print('leaves:', leaves)
print('treespec:', treespec)
print('reconstructed:', pytree.tree_unflatten(leaves, treespec))

with PyTreeExtensionContext():
    print('===== entered context =====')
    leaves, treespec = pytree.tree_flatten(ModelOutput(a=1, b=2))  # ModelOutput is a node
    print('leaves:', leaves)
    print('treespec:', treespec)
    print('reconstructed:', pytree.tree_unflatten(leaves, treespec))

print('===== exited context =====')
print('leaves:', leaves)
print('treespec:', treespec)
print('reconstructed:', pytree.tree_unflatten(leaves, treespec))
```

Output:

```console
$ python3 test.py             
===== before entering context =====
leaves: [ModelOutput([('a', 1), ('b', 2)])]
treespec: *
reconstructed: ModelOutput([('a', 1), ('b', 2)])
===== entered context =====
leaves: [1, 2]
treespec: TreeSpec(ModelOutput, (<class 'transformers.utils.generic.ModelOutput'>, ['a', 'b']), [*,
  *])
reconstructed: ModelOutput([('a', 1), ('b', 2)])
===== exited context =====
leaves: [1, 2]
treespec: TreeSpec(ModelOutput, (<class 'transformers.utils.generic.ModelOutput'>, ['a', 'b']), [*,
  *])
Traceback (most recent call last):
  File ".../test.py", line 23, in <module>
    print('reconstructed:', pytree.tree_unflatten(leaves, treespec))
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File ".../venv/lib/python3.11/site-packages/torch/utils/_pytree.py", line 268, in tree_unflatten
    unflatten_fn = SUPPORTED_NODES[spec.type].unflatten_fn
                   ~~~~~~~~~~~~~~~^^^^^^^^^^^
KeyError: <class 'transformers.utils.generic.ModelOutput'>
```

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_26618). All of your documentation changes will be reflected on that endpoint.