update tests

Signed-off-by: Ananth Subramaniam <[email protected]>

Author: ananthsub

tests/unit_tests/conftest.py

@@ -19,8 +19,7 @@
 
 import pytest
 
-
-# from tests.unit_tests.download_unit_tests_dataset import get_oldest_release_and_assets
+from tests.unit_tests.download_unit_tests_dataset import get_oldest_release_and_assets
 
 
 logging.basicConfig(level=logging.INFO)
@@ -42,34 +41,34 @@ def cleanup_local_folder():
         rmtree("./nemo_experiments", ignore_errors=True)
 
 
-# @pytest.fixture(scope="session", autouse=True)
-# def ensure_test_data(tmp_path_factory):
-#     """Ensure test data is available in a temporary directory by downloading if necessary."""
-#     data_path = tmp_path_factory.mktemp("test_data")
+@pytest.fixture(scope="session", autouse=True)
+def ensure_test_data(tmp_path_factory):
+    """Ensure test data is available in a temporary directory by downloading if necessary."""
+    data_path = tmp_path_factory.mktemp("test_data")
 
-#     # Check if data directory exists and has content
-#     if not any(data_path.iterdir()):
-#         logger.info(f"Test data not found at {data_path}. Downloading...")
+    # Check if data directory exists and has content
+    if not any(data_path.iterdir()):
+        logger.info(f"Test data not found at {data_path}. Downloading...")
 
-#         try:
-#             # Download assets to data_path
-#             get_oldest_release_and_assets(assets_dir=str(data_path))
+        try:
+            # Download assets to data_path
+            get_oldest_release_and_assets(assets_dir=str(data_path))
 
-#             logger.info("Test data downloaded successfully.")
+            logger.info("Test data downloaded successfully.")
 
-#         except ImportError as e:
-#             logger.info(f"Failed to import download function: {e}")
-#         except ValueError as e:
-#             logger.error(e)
-#             pytest.exit(f"Failed to download test data: {e}", returncode=1)
-#             # Don't fail the tests, just warn
-#         except Exception as e:
-#             logger.info(f"Failed to download test data: {e}")
-#             # Don't fail the tests, just warn
-#     else:
-#         logger.info(f"Test data already available at {data_path}")
+        except ImportError as e:
+            logger.info(f"Failed to import download function: {e}")
+        except ValueError as e:
+            logger.error(e)
+            pytest.exit(f"Failed to download test data: {e}", returncode=1)
+            # Don't fail the tests, just warn
+        except Exception as e:
+            logger.info(f"Failed to download test data: {e}")
+            # Don't fail the tests, just warn
+    else:
+        logger.info(f"Test data already available at {data_path}")
 
-#     yield data_path
+    yield data_path
 
 
 @pytest.fixture(autouse=True)

tests/unit_tests/models/starcoder/test_starcoder2_provider.py

@@ -38,7 +38,7 @@ def test_starcoder2_model_provider_defaults(self):
         assert provider.hidden_size == 768
         assert provider.num_attention_heads == 12
 
-        # Check Starcoder2-specific defaults
+        # Check Starcoder2-specific defaults + transformer config post init
         assert provider.normalization == "LayerNorm"
         assert provider.activation_func == F.gelu
         assert provider.add_bias_linear is True
@@ -49,28 +49,14 @@ def test_starcoder2_model_provider_defaults(self):
         assert provider.attention_dropout == 0.0
         assert provider.init_method_std == 0.01
         assert provider.share_embeddings_and_output_weights is False
-        assert provider.kv_channels is None
-        assert provider.num_query_groups is None
+        assert provider.kv_channels == 64
+        assert provider.num_query_groups == 12
         assert provider.window_size is None
         assert provider.attention_softmax_in_fp32 is True
         assert provider.bias_activation_fusion is True
         assert provider.bias_dropout_fusion is True
         assert provider.layernorm_epsilon == 1e-5
 
-    def test_starcoder2_model_provider_inheritance(self):
-        """Test Starcoder2ModelProvider inherits from GPTModelProvider."""
-        from megatron.bridge.models.gpt_provider import GPTModelProvider
-
-        provider = Starcoder2ModelProvider(
-            num_layers=12,
-            hidden_size=768,
-            num_attention_heads=12,
-        )
-
-        assert isinstance(provider, GPTModelProvider)
-        assert hasattr(provider, "provide")
-        assert callable(provider.provide)
-
 
 class TestStarcoder2ModelProvider3B:
     """Test cases for Starcoder2ModelProvider3B class."""
@@ -98,21 +84,12 @@ def test_starcoder2_3b_defaults(self):
         assert provider.hidden_dropout == 0.0
         assert provider.attention_dropout == 0.0
         assert provider.share_embeddings_and_output_weights is False
-        assert provider.kv_channels is None
         assert provider.window_size is None
         assert provider.attention_softmax_in_fp32 is True
         assert provider.bias_activation_fusion is True
         assert provider.bias_dropout_fusion is True
         assert provider.layernorm_epsilon == 1e-5
 
-    def test_starcoder2_3b_inheritance(self):
-        """Test Starcoder2ModelProvider3B inherits from Starcoder2ModelProvider."""
-        provider = Starcoder2ModelProvider3B()
-
-        assert isinstance(provider, Starcoder2ModelProvider)
-        assert hasattr(provider, "provide")
-        assert callable(provider.provide)
-
 
 class TestStarcoder2ModelProvider7B:
     """Test cases for Starcoder2ModelProvider7B class."""
@@ -140,21 +117,13 @@ def test_starcoder2_7b_defaults(self):
         assert provider.hidden_dropout == 0.0
         assert provider.attention_dropout == 0.0
         assert provider.share_embeddings_and_output_weights is False
-        assert provider.kv_channels is None
+        assert provider.kv_channels is 128
         assert provider.window_size is None
         assert provider.attention_softmax_in_fp32 is True
         assert provider.bias_activation_fusion is True
         assert provider.bias_dropout_fusion is True
         assert provider.layernorm_epsilon == 1e-5
 
-    def test_starcoder2_7b_inheritance(self):
-        """Test Starcoder2ModelProvider7B inherits from Starcoder2ModelProvider."""
-        provider = Starcoder2ModelProvider7B()
-
-        assert isinstance(provider, Starcoder2ModelProvider)
-        assert hasattr(provider, "provide")
-        assert callable(provider.provide)
-
 
 class TestStarcoder2ModelProvider15B:
     """Test cases for Starcoder2ModelProvider15B class."""
@@ -182,21 +151,13 @@ def test_starcoder2_15b_defaults(self):
         assert provider.hidden_dropout == 0.0
         assert provider.attention_dropout == 0.0
         assert provider.share_embeddings_and_output_weights is False
-        assert provider.kv_channels is None
+        assert provider.kv_channels == 128
         assert provider.window_size is None
         assert provider.attention_softmax_in_fp32 is True
         assert provider.bias_activation_fusion is True
         assert provider.bias_dropout_fusion is True
         assert provider.layernorm_epsilon == 1e-5
 
-    def test_starcoder2_15b_inheritance(self):
-        """Test Starcoder2ModelProvider15B inherits from Starcoder2ModelProvider."""
-        provider = Starcoder2ModelProvider15B()
-
-        assert isinstance(provider, Starcoder2ModelProvider)
-        assert hasattr(provider, "provide")
-        assert callable(provider.provide)
-
 
 class TestStarcoder2ProviderInheritance:
     """Test inheritance relationships between Starcoder2 providers."""

tests/unit_tests/models/starcoder/test_starcoder_provider.py

@@ -47,26 +47,12 @@ def test_starcoder_model_provider_defaults(self):
         assert provider.init_method_std == 0.01
         assert provider.layernorm_epsilon == 1e-5
         assert provider.share_embeddings_and_output_weights is False
-        assert provider.kv_channels is None
+        assert provider.kv_channels == 64
         assert provider.num_query_groups == 1
         assert provider.attention_softmax_in_fp32 is True
         assert provider.bias_activation_fusion is True
         assert provider.bias_dropout_fusion is True
 
-    def test_starcoder_model_provider_inheritance(self):
-        """Test StarcoderModelProvider inherits from GPTModelProvider."""
-        from megatron.bridge.models.gpt_provider import GPTModelProvider
-
-        provider = StarcoderModelProvider(
-            num_layers=12,
-            hidden_size=768,
-            num_attention_heads=12,
-        )
-
-        assert isinstance(provider, GPTModelProvider)
-        assert hasattr(provider, "provide")
-        assert callable(provider.provide)
-
 
 class TestStarcoderConfig15B:
     """Test cases for StarcoderConfig15B class."""
@@ -92,20 +78,12 @@ def test_starcoder_config_15b_defaults(self):
         assert provider.attention_dropout == 0.2
         assert provider.layernorm_epsilon == 1e-5
         assert provider.share_embeddings_and_output_weights is False
-        assert provider.kv_channels is None
+        assert provider.kv_channels == 128
         assert provider.num_query_groups == 1
         assert provider.attention_softmax_in_fp32 is True
         assert provider.bias_activation_fusion is True
         assert provider.bias_dropout_fusion is True
 
-    def test_starcoder_config_15b_inheritance(self):
-        """Test StarcoderConfig15B inherits from StarcoderModelProvider."""
-        provider = StarcoderConfig15B()
-
-        assert isinstance(provider, StarcoderModelProvider)
-        assert hasattr(provider, "provide")
-        assert callable(provider.provide)
-
 
 class TestStarcoderProviderInheritance:
     """Test inheritance relationships between Starcoder providers."""
@@ -120,9 +98,6 @@ def test_starcoder_models_inherit_from_gpt(self):
 
     def test_provide_method_inherited(self):
         """Test that provide method works correctly in inherited classes."""
-        # Test with StarcoderConfig15B
         provider = StarcoderConfig15B()
-
-        # The provide method should be inherited from GPTModelProvider
         assert hasattr(provider, "provide")
         assert callable(provider.provide)