Change varname

Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

Author: ksivaman

tests/pytorch/test_sanity.py

@@ -608,7 +608,7 @@ def test_sanity_grouped_linear(
     num_tokens = bs * config.max_seqlen_q * (num_gemms - 1)
 
     if single_param:
-        os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"] = "1"
+        os.environ["NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS"] = "1"
 
     if fp8_recipe is not None:
         if not is_fp8_supported(config):
@@ -650,7 +650,7 @@ def test_sanity_grouped_linear(
     assert out.shape == (num_tokens, ffn_hidden_size)
 
     if single_param:
-        del os.environ["NVTE_GROUPED_LINEAR_SINGLE_PARAM"]
+        del os.environ["NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS"]
 
 
 @pytest.mark.parametrize("dtype", param_types)

transformer_engine/pytorch/module/grouped_linear.py

@@ -795,7 +795,7 @@ def make_grouped_weights(self, defer_init=False) -> None:
     def reset_parameters(self, defer_init=False):
         super().reset_parameters(defer_init=defer_init)
         # Grouped tensor weights is an opt-in feature.
-        if bool(int(os.getenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0"))):
+        if bool(int(os.getenv("NVTE_ALLOC_CONTIGUOUS_GROUPED_LINEAR_WEIGHTS", "0"))):
             self.make_grouped_weights(defer_init=defer_init)
 
     def set_tensor_parallel_attributes(self, defer_init=False) -> None: