Repository: huggingface/transformers
Pull request: fix(models): Fix dtype mismatch in SwitchTransformers and TimmWrapperModel

### What does this PR do?

The following dtype mismatch use cases were identified and fixed in this PR:
  
→ **Switch Transformers:** [7938e91fa](https://github.com/harshaljanjani/transformers/commit/7938e91faabb051f3a001cd39c173d4697c2d81c) refactored all MoE models for vLLM compatibility; in that refactor, the `_cast_classifier()` method was removed from `SwitchTransformersTop1Router` but no dtype cast was added. Casting hidden_states to `classifier.weight.dtype` before the linear call fixes that!
→ **TimmWrapper:** [6217adc6c8](https://github.com/huggingface/transformers/commit/6217adc6c8f0be7b5374e6a46129ad2214e4c6ed) changed the default dtype behavior to "auto"; in that commit, `pixel_values.to(self.device, self.dtype)` was regressed to [pixel_values.to(self.device)](https://github.com/huggingface/transformers/blob/main/src/transformers/models/timm_wrapper/modeling_timm_wrapper.py#L246) dropping the dtype cast. I'm not too sure why it was dropped; but restoring it seems logical to fix the use case.
→ For more details on reproducing the bug and the output screenshots, please visit the linked issue!

cc: @Rocketknight1

Fixes #45072

**CI run test coverage of this behavior (as suggested by @ydshieh) :):**

[SwitchTransformers](https://github.com/huggingface/transformers/actions/runs/23629599078/job/68826304897):
→ `test_modeling_switch_transformers.py::SwitchTransformersModelTest::test_generate_with_past_key_values`
→ `test_modeling_switch_transformers.py::SwitchTransformersModelTest::test_model_fp16_forward`
→ `test_modeling_switch_transformers.py::SwitchTransformerModelIntegrationTests::test_small_logits`
[TimmWrapper](https://github.com/huggingface/transformers/actions/runs/23629599078/job/68826305100):
→ `TimmWrapperModelTest` does not have explicit bfloat16 forward pass tests; added one in this PR for complete coverage.

**Repro output after the fixes (feel free to cross-check):**

<img width="500" height="300" alt="1" src="https://github.com/user-attachments/assets/6c97e0ca-e3ce-4fe7-b501-31e1b144d4ad" />
<img width="380" height="175" alt="1-1" src="https://github.com/user-attachments/assets/5cad72ed-9deb-43d4-b54d-5a3ad8ad793d" />

### Code Agent Policy
- [x] I confirm that this is not a pure code agent PR.

### Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request), Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
    [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
    [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you fix any necessary existing tests?

Changed files:
filename: src/transformers/models/switch_transformers/modeling_switch_transformers.py
patch: @@ -91,6 +91,7 @@ def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tens
         if self.training and self.jitter_noise > 0:
             # Multiply the token inputs by the uniform distribution - adding some noise
             hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise)
+        self.classifier = self.classifier.to(self.dtype)
         router_logits = self.classifier(hidden_states)
 
         # Apply Softmax and cast back to the original `dtype`

filename: src/transformers/models/switch_transformers/modular_switch_transformers.py
patch: @@ -158,6 +158,7 @@ def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tens
         if self.training and self.jitter_noise > 0:
             # Multiply the token inputs by the uniform distribution - adding some noise
             hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise)
+        self.classifier = self.classifier.to(self.dtype)
         router_logits = self.classifier(hidden_states)
 
         # Apply Softmax and cast back to the original `dtype`

filename: src/transformers/models/timm_wrapper/modeling_timm_wrapper.py
patch: @@ -243,7 +243,7 @@ def forward(
                 "different architecture or updating the timm package to a compatible version."
             )
 
-        pixel_values = pixel_values.to(self.device)
+        pixel_values = pixel_values.to(self.device, self.dtype)
 
         if self.features_only:
             last_hidden_state = self.timm_model.forward(pixel_values, **kwargs)

filename: tests/models/timm_wrapper/test_modeling_timm_wrapper.py
patch: @@ -76,6 +76,11 @@ def prepare_config_and_inputs(self):
     def get_config(self):
         return TimmWrapperConfig(architecture=self.architecture, model_args=self.model_args)
 
+    def create_and_check_model_fp16_forward(self, config, pixel_values):
+        model = TimmWrapperModel(config=config).to(torch_device).half().eval()
+        output = model(pixel_values)["last_hidden_state"]
+        self.parent.assertFalse(torch.isnan(output).any().item())
+
     def prepare_config_and_inputs_for_common(self):
         config_and_inputs = self.prepare_config_and_inputs()
         config, pixel_values = config_and_inputs
@@ -136,6 +141,10 @@ def test_hidden_states_output(self):
             resulted_shapes = [list(h.shape[2:]) for h in outputs.hidden_states]
             self.assertListEqual(expected_shapes, resulted_shapes)
 
+    def test_model_fp16_forward(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_model_fp16_forward(*config_and_inputs)
+
     @unittest.skip(reason="TimmWrapper models doesn't have inputs_embeds")
     def test_inputs_embeds(self):
         pass