First get decorator working

[muellerzr]

Fix process decorators and make them more flexible

What does this add?

This PR completely refactors/reworks the decorator methodology introduced in #488 into something that fully completes the proposed API

Who is it for?

Closes #923

Why is it needed?

The decorators that were introduced did not work as intended, and only worked to have the Accelerator object itself actually use it. To use it outside the Accelerator was impossible. As a result, this PR completely reworks how the existing decorators were designed by making them utilize the state instead, and ensures that a special internal use case is added when using the decorators back on the Accelerator class itself, with a raised error if a user tries to do the same.

What parts of the API does this impact?

User-facing:

A user can now actually perform @accelerator.on_main_process as a decorator, along with the rest of the existing ones.
The AcceleratorState also now stores the use_distributed value originally in the Accelerator object

Internal structure:

The existing implementation in Accelerator was gutted, and reworked into a usable fashion.

Specifically, this logic has now been added to the decorators:

    def on_local_process(self, function: Callable[..., Any] = None, local_process_index: int = None):
        ...
        if function is None:
            if "Accelerator." in self.__qualname__:
                function = self
            else:
                raise ValueError(
                    "The `on_main_process` decorator must be called with a function on an instantiated `Accelerator` object."
                )

This check is to ensure if self is a function from the Accelerator itself, or the behavior that occurs when performing:

# Inside the Accelerator object definition
    @on_main_process
    def log(self, ...)
        ...

Instead of requiring complex workarounds. As you can see in the above, a proper error will also be raised if a user breaks this "magic" with a clear description of what is happening.

Inside the on_process* functions there also exists this snippet:

    def on_process(self, function: Callable[..., Any] = None, process_index: int = None):
        # Initial construction of the decorator.
        if (self is not None) and (process_index is not None) and (function is None):
            return partial(self.on_process, process_index=process_index)

This is to check if an overloaded decorator was used, potentially again on the Accelerator class, before moving on to the proper constructor.

Basic Usage Example(s):

from accelerate import Accelerator
accelerator = Accelerator()
@accelerator.on_local_main_process
def print_local_main():
    print(f"Printing from the local main process {state.local_process_index}")
print_local_main()

[sgugger]

Thanks for splitting the work in two. LGTM!

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.