From 007b71707115b31f756bb1aa4f1f51177f9a12e0 Mon Sep 17 00:00:00 2001 From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> Date: Wed, 28 Oct 2026 15:02:17 -0700 Subject: [PATCH] Count FSDP2 unit scans, not fully_shard.state calls, in the weight-access tests #3091 made _get_fsdp2_mesh call fully_shard.state, so the window now calls it once per entry and the #3088 tests' call counts were off. Count _fsdp2_param_groups, the scan. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com> --- tests/unit/torch/quantization/test_utils.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/unit/torch/quantization/test_utils.py b/tests/unit/torch/quantization/test_utils.py index 9b06862e56c..bd661affd6a 100644 --- a/tests/unit/torch/quantization/test_utils.py +++ b/tests/unit/torch/quantization/test_utils.py @@ -24,6 +24,7 @@ from modelopt.torch.quantization.utils import ( convert_quantization_axis_to_reduce_axis, + core_utils, reduce_amax, reduce_block_amax, ) @@ -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): @@ -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())