🍻 Add static typing

[julien-blanchon]

What does this PR do?

Add strong typing to Accelerator methods.

It's super annoying to loose your typing when using accelerate wrapper.

Here is some typing that are preserved with this PR:

accelerator.split_between_processes is typed

with accelerator.split_between_processes(["A", "B", "C"]) as inputs:
    reveal_type(inputs)  # E: list[str]

accelerator.print is typed, all the print *args are passed

accelerator.print("Hello world!", file=...) # all the `print` args are passed

accelerator.prepare is typed, output types exactly match input type. Same for accelerator.prepare_model, accelerator.prepare_optimizer and accelerator.prepare_scheduler, accelerator.unwrap_model

model2, optimizer2, scheduler2, = accelerator.prepare(model, optimizer, scheduler)

model2 = accelerator.prepare_model(model)
optimizer2 = accelerator.prepare_optimizer(optimizer)
scheduler2 = accelerator.prepare_scheduler(scheduler)

accelerator.gather, accelerator.pad_across_processes and accelerator.reduce are also typed

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, Pull Request section?
• [ ] Was this discussed/approved via a Github issue or the forum? 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, and here are tips on formatting docstrings.
• [ ] 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.

[muellerzr]

Thanks for your PR. At this time we have zero intention of turning accelerate into a strongly-typed library in a dynamically-typed language, very similar to transformers. It does not make sense, given the fact that much of this library is dynamic based on the backend used, plugin used, etc.

Similar comment in transformers can be found here: https://github.com/huggingface/transformers/issues/18464#issuecomment-1206386420

[BenjaminBossan]

I have no strong (pun not intended) opinion on this matter. As is, the existing type annotations are more type hints for convenience but not really checked for correctness, which makes relying on them a bit of a gamble.

If we wanted to really add full static typing (not strong typing, Python is already strongly typed), this would help with what Julien mentioned but also add more clutter and maintenance burden. Some of the more dynamic features will be difficult to express, especially as we want to support older Python versions (as an example, TypeVarTuple was only added in Python 3.11).

Although I haven't used them myself, might I suggest to you, Julien, to try adding a pyi stub file to your project with the added type hints. This way, you should still be able to benefit from the work you've done on this PR but we don't have to change the accelerate code base. Maybe even upload the stub(s) to GitHub for others to contribute.

[julien-blanchon]

Personally, I don't really like working with such libraries that get mixed up with all your types. It's a huge pain in terms of DX (developer experience) and code robustness. The idea of having a pyi stub separate from the code base is great, I didn't know about it. It could be a great addition to speed up, but in a long term vision, but it likelly not a priority now.

I'm converting the PR to a Draft so we can discuss it here and it can serve as a starting point in the future :)

PS: I seem to have got some type wrong in my initial PR: like the prepare_* that actually change the object class type.

[BenjaminBossan]

Personally, I don't really like working with such libraries that get mixed up with all your types.

If existing type annotations are incorrect, that's certainly not on purpose. We would generally accept PRs that fix incorrect type annotations. But I hope that you can see that types like TypeVar("T_split_between_processes", bound=Union[Sequence, dict, torch.Tensor]) are not easy to understand and increase the maintenance/contribution burden. Also, as mentioned earlier, we have to be careful about supporting older Python versions, which lack some of the more ergonomic typing features.

The idea of having a pyi stub separate from the code base is great, I didn't know about it. It could be a great addition to speed up, but in a long term vision, but it likelly not a priority now.

AFAIK, the type annotations in the stub file don't need to be complete. You could create this file and only add the annotations you created for this PR if that's enough for you to improve DX for the time being.

[github-actions[bot]]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.