Remove accelerate.prepare hooks and FP32 conversion

[DavidePaglieri]

When I use accelerator.prepare(model) the model gets some hooks used in accelerate. After I'm done with the training and want to save the model, those hooks are still somewhere in the .state_dict() of the model and I can't figure out where exactly and how to remove them.

This is a problem especially when I later want to load the model without accelerate and run it in FP16 with other models, because despite the model being loaded in FP16, I have the following accelerate decorator still somewhere moving the output in FP32:

class ConvertOutputsToFp32:
    """
    Decorator to apply to a function outputing tensors (like a model forward pass) that ensures the outputs in FP16
    precision will be convert back to FP32.

    Use a class instead of a decorator because otherwise, the prepared model can no longer be pickled (issue #273).

    Args:
        model_forward (`Callable`):
            The function which outputs we want to treat.

    Returns:
        The same function as `model_forward` but with converted outputs.
    """

    def __init__(self, model_forward):
        self.model_forward = model_forward
        update_wrapper(self, model_forward)

    def __call__(self, *args, **kwargs):
        return convert_to_fp32(self.model_forward(*args, **kwargs))

Is there a way I can remove all such hooks from the model trained with accelerate? I have tried remove_hook_from_module but it doesn't work.

Thank you!

[sgugger]

There is no util to remove those for now (it's not the Hooks, which we use for big model inference, but more wrappers around the forward) but I see why you'd want a util like this! I think the unwrap_model method should remove them.

@muellerzr Would you have time to work on that?

[muellerzr]

@sgugger yep! Will assign myself

[DavidePaglieri]

Hi @sgugger thank you for your answer! I actually use:

model = accelerator.unwrap_model(model).to(dtype=dtype)

before saving it, still this ConvertOutputToFp32 remains somewhere in the state_dict of the model

[sgugger]

Yes, like I said there is nothing to remove it yet. The goal would be for this snippet of code to work once Zach is done :-)

[muellerzr]

Hi @DavidePaglieri, should be fixed with #860!