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

Fix reduce type in FSDP mixed precision #18818

Closed
wants to merge 10 commits into from
Closed

Fix reduce type in FSDP mixed precision #18818

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
5f3a29f
update
awaelchli Oct 18, 2023
a07007f
update
awaelchli Oct 18, 2023
cc8633d
update test
awaelchli Oct 18, 2023
f780c60
update
awaelchli Oct 18, 2023
51e3324
update tests
awaelchli Oct 18, 2023
42fc3e5
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Oct 18, 2023
229cff6
update
awaelchli Oct 18, 2023
3eed461
Merge branch 'master' into bugfix/reduce-type-mixed-precision
awaelchli Oct 18, 2023
fa18547
chlog
awaelchli Oct 18, 2023
9458575
Merge branch 'master' into bugfix/reduce-type-mixed-precision
awaelchli Oct 29, 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
Prev Previous commit
Next Next commit
update tests
  • Loading branch information
@awaelchli
awaelchli committed Oct 18, 2023
commit 51e3324757db719517621f39fef8a21bf1acfc70
8 changes: 4 additions & 4 deletions tests/tests_fabric/plugins/precision/test_fsdp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,19 +27,19 @@
("16-true", (torch.float16, torch.float16, torch.float16)),
("bf16-true", (torch.bfloat16, torch.bfloat16, torch.bfloat16)),
pytest.param(
"16-mixed", (torch.float32, torch.float16, torch.float16), marks=RunIf(min_torch="2.0"), id="16-mixed-ge2_0"
"16-mixed", (torch.float32, torch.float32, torch.float32), marks=RunIf(min_torch="2.0"), id="16-mixed-ge2_0"
),
pytest.param(
"16-mixed", (None, torch.float16, torch.float16), marks=RunIf(max_torch="2.0"), id="16-mixed-lt2_0"
"16-mixed", (None, None, None), marks=RunIf(max_torch="2.0"), id="16-mixed-lt2_0"
),
pytest.param(
"bf16-mixed",
(torch.float32, torch.bfloat16, torch.bfloat16),
(torch.float32, torch.float32, torch.float32),
marks=RunIf(min_torch="2.0"),
id="bf16-mixed-ge2_0",
),
pytest.param(
"bf16-mixed", (None, torch.bfloat16, torch.bfloat16), marks=RunIf(max_torch="2.0"), id="bf16-mixed-lt2_0"
"bf16-mixed", (None, None, None), marks=RunIf(max_torch="2.0"), id="bf16-mixed-lt2_0"
),
pytest.param(
"32-true", (torch.float32, torch.float32, torch.float32), marks=RunIf(min_torch="2.0"), id="32-true-ge2_0"
Expand Down
6 changes: 2 additions & 4 deletions tests/tests_fabric/strategies/test_fsdp_integration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,11 +46,9 @@ def step(self, model, batch):
precision = self._precision
assert isinstance(precision, FSDPPrecision)
if precision.precision == "16-mixed":
param_dtype = torch.float32
reduce_dtype = buffer_dtype = None
param_dtype = reduce_dtype = buffer_dtype = torch.float32
elif precision.precision == "bf16-mixed":
param_dtype = torch.float32
reduce_dtype = buffer_dtype = None
param_dtype = reduce_dtype = buffer_dtype = torch.float32
elif precision.precision == "16-true":
param_dtype = reduce_dtype = buffer_dtype = torch.float16
elif precision.precision == "bf16-true":
Expand Down
8 changes: 4 additions & 4 deletions tests/tests_pytorch/plugins/precision/test_fsdp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,19 +27,19 @@
("16-true", (torch.float16, torch.float16, torch.float16)),
("bf16-true", (torch.bfloat16, torch.bfloat16, torch.bfloat16)),
pytest.param(
"16-mixed", (torch.float32, torch.float16, torch.float16), marks=RunIf(min_torch="2.0"), id="16-mixed-ge2_0"
"16-mixed", (torch.float32, torch.float32, torch.float32), marks=RunIf(min_torch="2.0"), id="16-mixed-ge2_0"
),
pytest.param(
"16-mixed", (None, torch.float16, torch.float16), marks=RunIf(max_torch="2.0"), id="16-mixed-lt2_0"
"16-mixed", (None, None, None), marks=RunIf(max_torch="2.0"), id="16-mixed-lt2_0"
),
pytest.param(
"bf16-mixed",
(torch.float32, torch.bfloat16, torch.bfloat16),
(torch.float32, torch.float32, torch.float32),
marks=RunIf(min_torch="2.0"),
id="bf16-mixed-ge2_0",
),
pytest.param(
"bf16-mixed", (None, torch.bfloat16, torch.bfloat16), marks=RunIf(max_torch="2.0"), id="bf16-mixed-lt2_0"
"bf16-mixed", (None, None, None), marks=RunIf(max_torch="2.0"), id="bf16-mixed-lt2_0"
),
pytest.param(
"32-true", (torch.float32, torch.float32, torch.float32), marks=RunIf(min_torch="2.0"), id="32-true-ge2_0"
Expand Down
6 changes: 2 additions & 4 deletions tests/tests_pytorch/strategies/test_fsdp.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,11 +83,9 @@ def _assert_layer_fsdp_instance(self) -> None:
assert isinstance(self.trainer.strategy.precision_plugin, FSDPPrecisionPlugin)

if self.trainer.precision == "16-mixed":
param_dtype = None if not _TORCH_GREATER_EQUAL_2_0 else torch.float32
reduce_dtype = buffer_dtype = torch.float16
param_dtype = reduce_dtype = buffer_dtype = None if not _TORCH_GREATER_EQUAL_2_0 else torch.float32
elif self.trainer.precision == "bf16-mixed":
param_dtype = None if not _TORCH_GREATER_EQUAL_2_0 else torch.float32
reduce_dtype = buffer_dtype = torch.bfloat16
param_dtype = reduce_dtype = buffer_dtype = None if not _TORCH_GREATER_EQUAL_2_0 else torch.float32
elif self.trainer.precision == "16-true":
param_dtype = reduce_dtype = buffer_dtype = torch.float16
elif self.trainer.precision == "bf16-true":
Expand Down