[Utils] has_offloaded_params

[kylesayrs]

Purpose

• Add a has_offloaded_params utility function which returns True iff there is an AlignDevicesHook with offloading enabled
• Prepare for other utility functions to be implemented in future PRs

Changes

• Implement utility function in accelerate.modeling
• Replace repeat code use in accelerate.modeling and accelerate.accelerator with newly added function

Testing

test_has_offloaded_params.py

import torch

from accelerate.hooks import attach_align_device_hook
from accelerate.utils.modeling import has_offloaded_params
from accelerate.big_modeling import cpu_offload_with_hook

class ModelForTest(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = torch.nn.Linear(3, 4)
        self.batchnorm = torch.nn.BatchNorm1d(4)
        self.linear2 = torch.nn.Linear(4, 5)

    def forward(self, x):
        return self.linear2(self.batchnorm(self.linear1(x)))
    
model = ModelForTest()
assert not has_offloaded_params(model.linear1)
assert not has_offloaded_params(model.batchnorm)
assert not has_offloaded_params(model.batchnorm)

model = ModelForTest()
model, hook = cpu_offload_with_hook(model, execution_device="cuda:0")
assert not has_offloaded_params(model.linear1)
assert not has_offloaded_params(model.batchnorm)
assert not has_offloaded_params(model.batchnorm)

model = ModelForTest()
attach_align_device_hook(model, offload=False)
assert not has_offloaded_params(model.linear1)
assert not has_offloaded_params(model.batchnorm)
assert not has_offloaded_params(model.batchnorm)

model = ModelForTest()
attach_align_device_hook(model, offload=True)
assert has_offloaded_params(model.linear1)
assert has_offloaded_params(model.batchnorm)
assert has_offloaded_params(model.batchnorm)

Who can review?

@SunMarc
@LysandreJik
@mgoin