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

Code before:
            # Without the fix, expert_params_to_check would be empty (all MISSING)
            self.assertGreater(
                len(expert_params_to_check),
                0,
                "No expert weights found - weight conversion failed! "
                "Expected fused gate_up_proj and down_proj but found none.",
            )

            with deepspeed.zero.GatheredParameters([param for _, param in expert_params_to_check], modifier_rank=0):
                for name, param in expert_params_to_check:
                    if "mlp.experts.gate_up_proj" in name:
                        self.assertEqual(len(param.shape), 3, f"gate_up_proj should be 3D, got {param.shape}")
                        self.assertEqual(param.shape[0], 8, f"Should have 8 experts, got {param.shape[0]}")
                    elif (
                        "mlp.experts.down_proj" in name
                        and name[name.rfind(".mlp.experts.down_proj") :] == ".mlp.experts.down_proj"
                    ):
                        self.assertEqual(len(param.shape), 3, f"down_proj should be 3D, got {param.shape}")
                        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):

Review comment (suggestion):
```suggestion
        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.
```

Code after:
            # Without the fix, expert_params_to_check would be empty (all MISSING)
            self.assertGreater(
                len(expert_params_to_check),
                0,
                "No expert weights found - weight conversion failed! "
                "Expected fused gate_up_proj and down_proj but found none.",
            )

            with deepspeed.zero.GatheredParameters([param for _, param in expert_params_to_check], modifier_rank=0):
                for name, param in expert_params_to_check:
                    if "mlp.experts.gate_up_proj" in name:
                        self.assertEqual(len(param.shape), 3, f"gate_up_proj should be 3D, got {param.shape}")
                        self.assertEqual(param.shape[0], 8, f"Should have 8 experts, got {param.shape[0]}")
                    elif (
                        "mlp.experts.down_proj" in name
                        and name[name.rfind(".mlp.experts.down_proj") :] == ".mlp.experts.down_proj"
                    ):
                        self.assertEqual(len(param.shape), 3, f"down_proj should be 3D, got {param.shape}")
                        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,