def test_partial_flattening(self):
module = self._get_transformer()
num_params = sum(p.numel() for p in module.parameters())
params_to_flatten = list(module.encoder.layers[1].parameters()) + list(
module.decoder.layers[0].parameters())
num_params_to_flatten = sum(p.numel() for p in params_to_flatten)
module = FlattenParamsWrapper(module, param_list=params_to_flatten)
assert module.flat_param.numel() == num_params_to_flatten
assert sum(p.numel() for p in module.parameters()) == num_params
# flattened parameters are removed
assert len(list(module.encoder.layers[1].parameters())) == 0
assert len(list(module.decoder.layers[0].parameters())) == 0
# non-flattened parameters remain
assert len(list(module.encoder.layers[0].parameters())) > 0
assert len(list(module.decoder.layers[1].parameters())) > 0
# test that changing the module dtype works properly
orig_dtype = params_to_flatten[0].dtype
new_dtype = torch.float32 if orig_dtype == torch.float16 else torch.float16
assert module.flat_param.dtype == orig_dtype
assert all(p.dtype == orig_dtype
for p in module.encoder.layers[0].parameters())
module = module.to(dtype=new_dtype)
assert module.flat_param.dtype == new_dtype
assert all(p.dtype == new_dtype
for p in module.encoder.layers[0].parameters())
def _test_num_params(self, module):
ref_num_params = sum(p.numel() for p in module.parameters())
flat_module = FlattenParamsWrapper(module)
flat_num_params = sum(p.numel() for p in flat_module.parameters())
assert ref_num_params == flat_num_params
assert flat_num_params == flat_module.flat_param.numel()
def test_empty_module(self):
module = self._get_empty_module()
in_data = torch.rand(1)
ref_out = module(in_data)
module = FlattenParamsWrapper(module)
assert len(list(module.parameters())) == 0
assert len(module.state_dict()) == 0
fpw_out = module(in_data)
torch.testing.assert_allclose(ref_out, fpw_out)
def test_two_flattening_group(self):
module = self._get_transformer()
num_params = sum(p.numel() for p in module.parameters())
params_to_flatten1 = list(
module.encoder.layers[1].parameters()) + list(
module.decoder.layers[0].parameters())
params_to_flatten2 = list(
module.encoder.layers[0].parameters()) + list(
module.decoder.layers[1].parameters())
num_params_to_flatten1 = sum(p.numel() for p in params_to_flatten1)
num_params_to_flatten2 = sum(p.numel() for p in params_to_flatten2)
module = FlattenParamsWrapper(
module, param_list=[params_to_flatten1, params_to_flatten2])
assert module.flat_params[0].numel() == num_params_to_flatten1
assert module.flat_params[1].numel() == num_params_to_flatten2
assert sum(p.numel() for p in module.parameters()) == num_params