pytorch - 💡(How to fix) Fix Incompatibility between FSDP and nn.utils.parametrize [9 comments, 5 participants]

import os
import torch
import torch.nn as nn


class SubModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.my_param = nn.Parameter(torch.randn(1024, 1024))

    def forward(self, x):
        return x @ self.my_param


class TransformerBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.sub = SubModule()
        self.linear = nn.Linear(1024, 1024)

    def forward(self, x):
        return self.linear(self.sub(x))


class Model(nn.Module):
    def __init__(self, num_layers=4):
        super().__init__()
        self.layers = nn.ModuleList([TransformerBlock() for _ in range(num_layers)])

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        return x


class IdentityParametrization(nn.Module):
    def forward(self, x):
        return x


def patch_forward(original_forward):
    # Monkey-patch SubModule.forward to call register_parametrization during forward

    def patched_forward(self, x):
        nn.utils.parametrize.register_parametrization(self, "my_param", IdentityParametrization())
        result = x @ self.my_param
        nn.utils.parametrize.remove_parametrizations(self, "my_param", leave_parametrized=False)
        return result

    SubModule.forward = patched_forward


def main():
    use_fsdp = "WORLD_SIZE" in os.environ
    if use_fsdp:
        import torch.distributed as dist

        dist.init_process_group("nccl")
        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
        from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
        from functools import partial

        rank = dist.get_rank()
        torch.cuda.set_device(rank)
        device = f"cuda:{rank}"
    else:
        device = "cuda:0"

    model = Model().to(device)

    if use_fsdp:
        auto_wrap_policy = partial(
            transformer_auto_wrap_policy,
            transformer_layer_cls={TransformerBlock},
        )
        model = FSDP(model, auto_wrap_policy=auto_wrap_policy, use_orig_params=True)

    x = torch.randn(2, 1024, device=device)
    model(x)

    original_forward = SubModule.forward
    patch_forward(original_forward)

    # With FSDP this fails with: AssertionError
    #   File .../torch/nn/utils/parametrize.py", line 387, in _inject_property
    #     assert not hasattr(module, tensor_name)
    model(x)
    print("passes")
    if use_fsdp:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()

---

[rank1]:     result = x @ self.my_param
[rank1]:              ~~^~~~~~~~~~~~~~~
[rank1]: RuntimeError: Output 0 of ViewBackward0 is a view and its base or another view of its base has been modified inplace. This view is the output of a function that returns multiple views. Such functions do not allow the output views to be modified inplace. You should replace the inplace operation by an out-of-place one.