Skip to content

Fix FSDP2 weight-access tests broken on main by #3091 #3106

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

Closed
kevalmorabia97 wants to merge 1 commit into from
Closed

Fix FSDP2 weight-access tests broken on main by #3091 #3106

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
13 changes: 9 additions & 4 deletions tests/unit/torch/quantization/test_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@

from modelopt.torch.quantization.utils import (
convert_quantization_axis_to_reduce_axis,
core_utils,
reduce_amax,
reduce_block_amax,
)
Expand Down Expand Up @@ -415,11 +416,13 @@ def _fsdp2_owns_its_params_in_every_state_job(rank, size):
finally:
layer.reshard()
# The set is kept for the unit: opening the window again costs the module's parameters only.
with mock.patch.object(fully_shard, "state", wraps=fully_shard.state) as state:
with mock.patch.object(
core_utils, "_fsdp2_param_groups", wraps=core_utils._fsdp2_param_groups
) as scans:
for _ in range(3):
with fsdp2_weight_access_and_writeback_context(model.dense, model):
pass
assert state.call_count == 0
assert scans.call_count == 0


def test_fsdp2_owns_its_parameters_in_every_shard_state(cpu_dist_workers):
Expand Down Expand Up @@ -447,11 +450,13 @@ def _window_leaves_what_fsdp2_does_not_own_job(rank, size):
assert weight.to_local().shape[0] == 2
# Every expert weight misses the kept set, and one scan of the unit answers them all.
with (
mock.patch.object(fully_shard, "state", wraps=fully_shard.state) as state,
mock.patch.object(
core_utils, "_fsdp2_param_groups", wraps=core_utils._fsdp2_param_groups
) as scans,
fsdp2_weight_access_and_writeback_context(layer, layer),
):
pass
assert state.call_count == 1
assert scans.call_count == 1
assert all(weight.to_local().shape[0] == 2 for weight in layer.experts.parameters())


Expand Down