Repository: huggingface/transformers
File: tests/deepspeed/test_deepspeed.py (Python)

Code before:
                        self.assertEqual(param.shape[0], 8, f"Should have 8 experts, got {param.shape[0]}")

    def test_init_zero3_siglip(self):
        """
        Test that SigLIP model initialization works correctly with DeepSpeed ZeRO-3.
        This tests the fix for lecun_normal_ and default_flax_embed_init initialization
        functions that now check for _is_hf_initialized flag to prevent re-initialization
        in ZeRO-3 environment.
        """
        ds_config = {
            "train_batch_size": 1,
            "zero_optimization": {
                "stage": 3,
            },
        }

        dschf = HfDeepSpeedConfig(ds_config)

        self.assertTrue(dschf.is_zero3())
        self.assertTrue(is_deepspeed_zero3_enabled())

        with LoggingLevel(logging.INFO):
            with mockenv_context(**self.dist_env_1_gpu):
                logger = logging.get_logger("transformers.modeling_utils")
                with CaptureLogger(logger) as cl:
                    model = AutoModel.from_pretrained("google/siglip-base-patch16-224")

        self.assertIn("Detected DeepSpeed ZeRO-3", cl.out)

        self.assertIsNotNone(model)


class TrainerIntegrationDeepSpeedWithCustomConfig(TestCasePlus):
    def setUp(self):
        super().setUp()

        args = TrainingArguments(".")
        self.n_epochs = args.num_train_epochs
        self.batch_size = args.train_batch_size

        master_port = get_master_port(real_launcher=False)
        self.dist_env_1_gpu = {
            "MASTER_ADDR": "localhost",
            "MASTER_PORT": master_port,
            "RANK": "0",
            "LOCAL_RANK": "0",
            "WORLD_SIZE": "1",
        }

        self.ds_config_file = {
            "zero2": f"{self.test_file_dir_str}/ds_config_zero2.json",

Review comment (suggestion):
The test is looking good I just have one nit here: We do not need to have a full model, we can create a dummy model and save to a temporary directory and load from that - no hub calls needed and should run much faster.

You can also look into test modeling siglip2 file for some dummy values. It's just there to have something small 

Code after:
                        self.assertEqual(param.shape[0], 8, f"Should have 8 experts, got {param.shape[0]}")

    def test_init_zero3_variance_scaling(self):
        """
        Tests whether variance scaling initializations (`lecun_normal_`, `default_flax_embed_init_`) work correctly
        with DeepSpeed ZeRO-3, e.g. as in SigLIP models. It indirectly checks for the `_is_hf_initialized` flag to
        prevent re-initialization in ZeRO-3 environments. See #43574
        """
        import tempfile

        from transformers import (
            SiglipConfig,
            SiglipModel,
            SiglipTextConfig,
            SiglipVisionConfig,
        )

        text_cfg = SiglipTextConfig(
            vocab_size=64,
            hidden_size=16,
            num_hidden_layers=1,
            num_attention_heads=2,
            intermediate_size=32,
            max_position_embeddings=16,
        )

        vision_cfg = SiglipVisionConfig(
            image_size=4,
            patch_size=2,
            num_channels=3,
            hidden_size=16,
            num_hidden_layers=1,
            num_attention_heads=2,
            intermediate_size=32,
        )

        cfg = SiglipConfig(text_config=text_cfg.to_dict(), vision_config=vision_cfg.to_dict())

        with tempfile.TemporaryDirectory() as tmpdirname:
            model = SiglipModel(cfg).eval()
            model.save_pretrained(tmpdirname)

            ds_config = {
                "train_batch_size": 1,
                "zero_optimization": {
                    "stage": 3,
                },
            }

            dschf = HfDeepSpeedConfig(ds_config)