Hi, I am running into an issue when using FSDP2’s fully_shard API with models that share submodules. Specifically, when I try to shard two models that share a submodule, I encounter an error related to DTensors already existing in the shared module
Example code
import os
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.distributed._composable.fsdp import fully_shard
class Block(nn.Module):
def __init__(self) -> None:
super().__init__()
self.nested_lin1 = nn.Linear(16, 16)
self.nested_lin2 = nn.Linear(16, 16)
def forward(self, input):
lin_out = self.nested_lin1(input)
lin_out = self.nested_lin2(lin_out)
return lin_out
class MyModel(nn.Module):
def __init__(self, block: Optional[nn.Module] = None) -> None:
super().__init__()
self.lin1 = nn.Linear(4, 16)
self.block = block or Block()
self.lin3 = nn.Linear(16, 7)
def forward(self, input):
output = self.lin1(input)
output = self.block(output)
output = self.lin3(output)
return output
def main():
rank = int(os.environ["LOCAL_RANK"])
world_size = int(os.environ["WORLD_SIZE"])
dist.init_process_group("nccl", rank=rank, world_size=world_size)
torch.cuda.set_device(rank)
torch.set_default_device("cuda:" + str(rank))
root_1 = MyModel()
root_2 = MyModel(block=root_1.block) # Shared submodule
root_1 = fully_shard(root_1)
root_2 = fully_shard(root_2) # Error occurs here
# Example forward pass
for i in range(num_iters):
if i % 2:
output = root_2(input)
else:
output = root_1(input)
...
if __name__ == "__main__":
main()
Error message
When applying fully_shard to root_2, I get the following error:
AssertionError: FSDP requires named DeviceMesh dims for ND parallelism
Question
- Is sharing submodules between models unsupported in FSDP2’s
fully_shard? - Is there a recommended way to handle shared modules in this scenario?
- Should I avoid sharing modules entirely, or is there a workaround or best practice for this use case?
Any clarification or suggestions would be greatly appreciated. Thank you!