OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

Fix bug in FSDP wrapped module with zero argument (#147771)

Browse Source 
Fixes https://github.com/pytorch/pytorch/issues/147531

Pull Request resolved: https://github.com/pytorch/pytorch/pull/147771
Approved by: https://github.com/awgu
This commit is contained in:
mori360 2025-02-26 01:40:53 +00:00 committed by PyTorch MergeBot
parent 8de6fe8c0b
commit 12112fd198
2 changed files with 15 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
test/distributed/fsdp/test_wrap.py
View File

@ -395,6 +395,19 @@ class TestFSDPWrap(FSDPTest):
loss.backward() loss.backward()
optim.step() optim.step()
@skip_if_lt_x_gpu(1)
def test_zero_argument(self):
class ZeroArguModel(nn.Module):
def __init__(self) -> None:
super().__init__()
self.a = torch.tensor([1.0])
def forward(self):
return self.a
model = FSDP(ZeroArguModel())
self.assertEqual(model(), torch.tensor([1.0]))
class TestAutoWrap(TestCase): class TestAutoWrap(TestCase):
def setUp(self) -> None: def setUp(self) -> None:

4
torch/distributed/fsdp/_runtime_utils.py
View File

@ -596,8 +596,8 @@ def _root_pre_forward(
args_tuple, kwargs_tuple = _to_kwargs( args_tuple, kwargs_tuple = _to_kwargs(
args, kwargs, state.compute_device, False args, kwargs, state.compute_device, False
) )
args = args_tuple[0] args = args_tuple[0] if args_tuple else tuple()
kwargs = kwargs_tuple[0] kwargs = kwargs_tuple[0] if kwargs_tuple else {}
return _root_cast_forward_input(state, module, args, kwargs) return _root_cast_forward_input(state, module, args, kwargs)