Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[FSDP] Added DDP parity test for CPU training #114372

Closed
wants to merge 4 commits into from
Closed

[FSDP] Added DDP parity test for CPU training #114372

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
ba13215
[FSDP] Added DDP parity test for CPU training
awgu Nov 22, 2023
f96f66d
Update on "[FSDP] Added DDP parity test for CPU training"
awgu Nov 22, 2023
f1c9b4c
Update on "[FSDP] Added DDP parity test for CPU training"
awgu Nov 22, 2023
99e090b
Update on "[FSDP] Added DDP parity test for CPU training"
awgu Nov 22, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 16 additions & 3 deletions test/distributed/fsdp/test_fsdp_misc.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
)
from torch.distributed.optim import _apply_optimizer_in_backward
from torch.nn import TransformerDecoderLayer, TransformerEncoderLayer
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
from torch.testing._internal.common_fsdp import (
_assert_module_states,
Expand Down Expand Up @@ -505,7 +506,6 @@ def test_fsdp_optimizer_overlap(self):
@skip_if_lt_x_gpu(2)
def test_fsdp_cpu_training(self):
"""Tests FSDP training on CPU."""
torch.manual_seed(0)
gloo_pg = dist.new_group(backend="gloo")
for ss in [
ShardingStrategy.NO_SHARD,
Expand All @@ -514,15 +514,28 @@ def test_fsdp_cpu_training(self):
ShardingStrategy.HYBRID_SHARD,
ShardingStrategy._HYBRID_SHARD_ZERO2,
]:
torch.manual_seed(42)
model = MyModel()
fsdp = FSDP(
ref_model = DDP(deepcopy(model), process_group=gloo_pg)
model = FSDP(
model,
auto_wrap_policy=always_wrap_policy,
process_group=gloo_pg,
device_id=torch.device("cpu"),
)
ref_optim = torch.optim.Adam(ref_model.parameters(), lr=1e-2)
optim = torch.optim.Adam(model.parameters(), lr=1e-2)
torch.manual_seed(42 + self.rank)
inp = torch.randn(2, 2)
fsdp(inp, inp).sum().backward()
for _ in range(10):
losses = []
for _model, _optim in ((ref_model, ref_optim), (model, optim)):
loss = _model(inp, inp).sum()
losses.append(loss)
loss.backward()
_optim.step()
_optim.zero_grad()
self.assertEqual(losses[0], losses[1])

@skip_if_lt_x_gpu(2)
def test_fsdp_cpu_init_stays_on_cpu(self):
Expand Down
Loading