Add serialization logic to pytree types
Author: angelayi
Pull request: yes

# What does this PR do?

[`torch.export`](https://pytorch.org/docs/main/export.html) produces a graph representation of the program which contains a pytree `TreeSpec` specifying how to flatten the inputs from whatever structure they were in before to a flattened list of inputs to pass to the graph, and unflatten the outputs from the graph to the original structure that eager mode PyTorch produces. We would like to serialize the artifact from `torch.export` and later pass it to a runtime. This means that we need a good BC/FC surface.

The pytree `TreeSpec` contains a structure telling us how the original input/output looked like. For example, if we had an input containing a tuple of a class `InputClass`, the `TreeSpec` would look something like: `TreeSpec(tuple, None, [TreeSpec[InputClass, None, []], TreeSpec[InputClass, None, []]])`. We want to serialize this into a json string consisting of the serialized type, serialized context, and recursively serialized children. Since we want to convert to a json format, only json types are allowed (strings, ints...). Therefore in this PR you can see logic to serialize a python type to a fully qualified string name.

Note that this field only exists starting from PyTorch version 2.2.

Here's an example:
```
import torch
import io
from transformers import AlbertForMaskedLM
from transformers import file_utils

model_cls = AlbertForMaskedLM
model_config = model_cls.config_class()
model = model_cls(model_config)
model.to("cuda", torch.float32)

input_dict = {"input_ids": torch.randint(0, 30000, (1, 512), device="cuda", dtype=torch.int64, requires_grad=False)}

ep = torch.export.export(model, (), input_dict)

print(ep.module_call_graph[0].signature.in_spec)
# TreeSpec(tuple, None, [TreeSpec(tuple, None, []), TreeSpec(dict, ['input_ids'], [*])])
print(ep.module_call_graph[0].signature.out_spec)
# TreeSpec(MaskedLMOutput, (<class 'transformers.modeling_outputs.MaskedLMOutput'>, ['logits']), [*])

buffer = io.BytesIO()
torch.export.save(ep, buffer)
buffer.seek(0)
loaded_ep = torch.export.load(buffer)

print(loaded_ep.module_call_graph[0].signature.in_spec)
# TreeSpec(tuple, None, [TreeSpec(tuple, None, []), TreeSpec(dict, ['input_ids'], [*])])
print(loaded_ep.module_call_graph[0].signature.out_spec)
# TreeSpec(MaskedLMOutput, (<class 'transformers.modeling_outputs.MaskedLMOutput'>, ['logits']), [*])

input_dict = {"input_ids": torch.randint(0, 30000, (1, 512), device="cuda", dtype=torch.int64, requires_grad=False)}
assert(torch.allclose(model(**input_dict).logits, loaded_ep(**input_dict).logits))
```

## 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.
- [ ] 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.

@ydshieh who seems to have reviewed pytree-related stuff before :) 

Comments:
Hi @angelayi Thank you for opening a PR. I am not really familiar with this topic, so here are a few questions:

> To ensure that the serialized name is consistent and doesn't change often

Could you show us an example that we can see the serialized name being changed (without the change in this PR)?

And could you share a documentation (I guess it is pytorch doc) that  providing a serialized type name is a good/recommended practice name and why.

Finally, an example demonstrate that the change in this PR does its job as described.

Thank you in advance!



@ydshieh Sorry for the delay! The PR context changed a little, but I added more information to the PR summary, and added a test. Please let me know if you need any more clarifications!

Hi @angelayi 

Thanks for the update. I would still love to see a code snippet (that we can run easily) in action that address the following

> Could you show us an example that we can see the serialized name being changed (without the change in this PR)?

🙏 

Hi @ydshieh, I updated the PR so that comment is no longer really relevant!

bump @ydshieh 

Hi @angelayi .

Although I get the idea overall, and it LGTM, I still have a few questions and comments (like previously) and would love to see more detailed explanations.

- I still don't get why this change is `necessary`:
  - without this change, would we get some BC/FC issue, say a graph created with old torch version could not be used with newer torch version?
  - if this is the case, could you elaborate what the situation(s) is/are?
- And since this (the code content been added in this PR) is not something (very) public that a lot of users will use (it is more internal), I personally would appreciate a lot if an example that uses `torch.export` (I assume this is more public) that you mentioned in the PR description to demonstrate the situation (instead of just using `pytree.treespec_dumps` - it's fine to use this in the testing, but would help a lot to put an example in the PR description using `torch.export`).


Looking forward to your comments !

@ydshieh 

> I still don't get why this change is necessary:

Without this change, we cannot serialize exported programs for models with outputs of type `ModelOutput`, aka most HuggingFace models. We would like for users to be able to save their exported programs and load it in another time.

> I personally would appreciate a lot if an example that uses torch.export

I added an example to the PR description which exports AlbertMaskedLM, saves, and loads it. Under the hood in torch.export.save, we use pytree.treespec_dumps to save the treespec, and in torch.export.load, we use pytree.treespec_loads to load the treespec. I also linked an issue where some of our benchmarks are failing due to not having pytree serialization support.

Thank you a lot @angelayi I will check the example ❤️ .


review this today

Hi! I reviewed the change along with the example you provided, I have one major question and another minor one.

Minor one: So there is no way to make `torch.export.save` work with torch < 2.2 for HF models?

### Major issue:

A variable `SERIALIZED_CLASS_TO_PYTHON_CLASS` is created and used. It works well with your script as both saving and loading are done in the same python process (which uses a buffer).

However, despite I don't have much experience on this topic, I think a common scenario/use case is: A user saves the exported program to a file. And later on, they will load the file, but **in another python process**. 

So the 2 scripts (saving/loading) will give `KeyError: 'transformers.modeling_outputs.MaskedLMOutput'` (see the full log at the end).

**Question**: Do you have any idea how to address this common use case?

-------------------------------------------------------------------------------------------------------------------------

#### saving
```python
import torch
import io
from transformers import AlbertForMaskedLM
from transformers import file_utils

model_cls = AlbertForMaskedLM
model_config = model_cls.config_class()
model = model_cls(model_config)
model.to("cuda", torch.float32)

input_dict = {"input_ids": torch.randint(0, 30000, (1, 512), device="cuda", dtype=torch.int64, requires_grad=False)}
ep = torch.export.export(model, (), input_dict)

fn = 'my_exported'
torch.export.save(ep, fn)
```

#### loading
```python
import torch
import io
from transformers import AlbertForMaskedLM
from transformers import file_utils

model_cls = AlbertForMaskedLM
model_config = model_cls.config_class()
model = model_cls(model_config)
model.to("cuda", torch.float32)

input_dict = {"input_ids": torch.randint(0, 30000, (1, 512), device="cuda", dtype=torch.int64, requires_grad=False)}

fn = 'my_exported'
loaded_ep = torch.export.load(fn)


input_dict = {"input_ids": torch.randint(0, 30000, (1, 512), device="cuda", dtype=torch.int64, requires_grad=False)}
assert(torch.allclose(model(**input_dict).logits, loaded_ep(**input_dict).logits))
```

#### error log

```bash
Traceback (most recent call last):
  File "temp3.py", line 14, in <module>
    loaded_ep = torch.export.load(buffer)
  File "/usr/local/lib/python3.8/dist-packages/torch/export/__init__.py", line 582, in load
    return load(
  File "/usr/local/lib/python3.8/dist-packages/torch/_export/__init__.py", line 1079, in load
    ep = deserialize(artifact)
  File "/usr/local/lib/python3.8/dist-packages/torch/_export/serde/serialize.py", line 1714, in deserialize
    ExportedProgramDeserializer(expected_opset_version)
  File "/usr/local/lib/python3.8/dist-packages/torch/_export/serde/serialize.py", line 1565, in deserialize
    GraphModuleDeserializer()
  File "/usr/local/lib/python3.8/dist-packages/torch/_export/serde/serialize.py", line 1275, in deserialize
    module_call_graph = self.deserialize_module_call_graph(serialized_graph_module.module_call_graph)
  File "/usr/local/lib/python3.8/dist-packages/torch/_export/serde/serialize.py", line 1513, in deserialize_module_call_graph
    return [
  File "/usr/local/lib/python3.8/dist-packages/torch/_export/serde/serialize.py", line 1516, in <listcomp>
    signature=self.deserialize_module_call_signature(entry.signature) if entry.signature else None,
  File "/usr/local/lib/python3.8/dist-packages/torch/_export/serde/serialize.py", line 1509, in deserialize_module_call_signature
    out_spec=treespec_loads(module_call_signature.out_spec),
  File "/usr/local/lib/python3.8/dist-packages/torch/utils/_pytree.py", line 991, in treespec_loads
    return _SUPPORTED_PROTOCOLS[protocol].json_to_treespec(json_schema)
  File "/usr/local/lib/python3.8/dist-packages/torch/utils/_pytree.py", line 953, in _json_to_treespec
    context = serialize_node_def.from_dumpable_context(json_schema["context"])
  File "/transformers/src/transformers/utils/generic.py", line 465, in _model_output_from_dumpable_context
    python_class = SERIALIZED_CLASS_TO_PYTHON_CLASS[serialized_class]
KeyError: 'transformers.modeling_outputs.MaskedLMOutput'
```
 

@ydshieh 

Thanks for the detailed review!

> So there is no way to make torch.export.save work with torch < 2.2 for HF models?

No sorry, torch.export.save won't work with torch < 2.2 for HF models. But also, torch.export.save was introduced in 2.1.

> A user saves the exported program to a file. And later on, they will load the file, but in another python process.

Thanks for catching this! I made a minor change to how the flattening/unflattening is implemented, and this issue should no longer occur. Please let me know what you think!

Hi @angelayi 

Thanks a lot! The save/load works with the latest change. But the last line

> assert(torch.allclose(model(**input_dict).logits, loaded_ep(**input_dict).logits))

in `loading` script fails as the diff is `> 1.0`

(I am using this block at the end)
```python
diff = model(**input_dict).logits - loaded_ep(**input_dict).logits

diff = torch.amax(torch.abs(diff))
print(diff)
```

However, if we put the saving/loading logic in the same script, the diff is `0.0`.

I am feeling something very strange happens.

Let me know if you are able to reproduce with the scripts.

The bash output looks like

```bash
root@e6558c895956:/transformers# python3 temp1.py
tensor(0., device='cuda:0', grad_fn=<AmaxBackward0>)

root@e6558c895956:/transformers# python3 temp2.py
tensor(1.3523, device='cuda:0', grad_fn=<AmaxBackward0>)
```

sorry, it's my bad!  Let me check again.

(I was using 2 models with different initialized weights)

thanks so much for the detailed reviews!!

@amyeroberts do you when we could get this merged into the repo? 

@amyeroberts I should've addressed the issues you mentioned, except I'm failing the formatting CI job. I tried running `make quality --fix` but `--fix` does not seem to be a valid option 😅 Do you know how I could fix this? 

@angelayi To apply the quality changes you can run `make fixup` and push the changes applies

@amyeroberts ah it seemed to work now! thanks!

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_27871). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Thanks so much for all the help and reviews!