[PyTorch] Optimize GroupedLinear validation (#3405)

* Optimize GroupedLinear validation

Signed-off-by: Evgeny <etsykunov@nvidia.com>

* CI fix

Signed-off-by: Evgeny <etsykunov@nvidia.com>

---------

Signed-off-by: Evgeny <etsykunov@nvidia.com>
Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>

Author: negvet

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)

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 = (