diff --git a/tests/pytorch/test_hybrid_quantization.py b/tests/pytorch/test_hybrid_quantization.py index f17a475831..73d69a9231 100644 --- a/tests/pytorch/test_hybrid_quantization.py +++ b/tests/pytorch/test_hybrid_quantization.py @@ -4825,6 +4825,13 @@ def unexpected_validation(*_args, **_kwargs): model(tensor, m_splits) assert model._custom_quantizer_cache == {} + def unexpected_grad_mode_query(): + pytest.fail("built-in recipe validation must not query grad mode") + + with monkeypatch.context() as context: + context.setattr(torch, "is_grad_enabled", unexpected_grad_mode_query) + model._validate_custom_recipe_quantizers(False, recipe.DelayedScaling()) + def unexpected_custom_validation(*_args, **_kwargs): pytest.fail("built-in recipes must not validate custom quantizers per forward") @@ -4869,12 +4876,29 @@ def tracked_validate(quantizers, *, operand_name="operand"): first_call_count = len(validation_calls) first_generation = model._custom_quantizer_cache["scaling_fwd"] assert first_call_count > 0 + assert "scaling_bwd" not in model._custom_quantizer_cache with torch.no_grad(), autocast(enabled=True, recipe=original_recipe): model(tensor, m_splits) assert len(validation_calls) == first_call_count assert model._custom_quantizer_cache["scaling_fwd"] is first_generation + # Backward quantizers are validated the first time quantizers are selected + # with gradients enabled. Do not run a full forward here: this validation + # test intentionally uses a columnwise-only configuration that is not + # supported by the split-quantization kernel. + with torch.enable_grad(), autocast(enabled=True, recipe=original_recipe): + model._get_quantizers() + assert len(validation_calls) == first_call_count + 1 + assert model._custom_quantizer_cache["scaling_bwd"] is model.quantizers["scaling_bwd"] + + def unexpected_grad_mode_query(): + pytest.fail("cached validation must not query grad mode") + + with monkeypatch.context() as context: + context.setattr(torch, "is_grad_enabled", unexpected_grad_mode_query) + model._validate_custom_recipe_quantizers(False, original_recipe) + rebuilt_recipe = recipe.CustomRecipe(qfactory=make_qfactory("rowwise_dequantized")) with torch.no_grad(), autocast(enabled=True, recipe=rebuilt_recipe): model(tensor, m_splits) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 612a430966..80d6f9596b 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -1810,6 +1810,8 @@ def _validate_custom_recipe_quantizers(self, fwd: bool, recipe: Recipe) -> None: return if self._custom_quantizer_cache.get(meta_key) is generation: return + if not fwd and not torch.is_grad_enabled(): + return if fwd: stride = self._num_fp8_tensors_per_gemm["fwd"] @@ -2455,8 +2457,7 @@ def _get_quantizers(self): # a failed generation remains installed when the caller catches the error. recipe = FP8GlobalStateManager.get_fp8_recipe() self._validate_custom_recipe_quantizers(True, recipe) - if torch.is_grad_enabled(): - self._validate_custom_recipe_quantizers(False, recipe) + self._validate_custom_recipe_quantizers(False, recipe) weight_quantizers = self._get_weight_quantizers() input_quantizers, output_quantizers = (