Fix MPS bf16 fallback and torchvision dependency in gemma-trainer

[keyuchen21]

Summary

While working through the gemma-trainer skill end-to-end on an Apple Silicon Mac (MPS, no CUDA) — dataset validation, SFT QLoRA fine-tuning of google/gemma-4-E2B-it, and adapter inference — I hit two issues in the training scripts:

• torch.cuda.is_bf16_supported() is CUDA-only. All four scripts (sft_train.py, dpo_train.py, reward_train.py, distill_dataset.py) pick the compute dtype with torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16. On an MPS-only machine this always evaluates to False, silently forcing fp16 even though MPS supports bf16. In practice this fp16 fallback is more numerically fragile under QLoRA — pushing epochs/rank/LR a bit further than the defaults produced NaN adapter weights for me on fp16, while the same run was stable once forced to bf16.
• Missing torchvision dependency. sft_train.py and dpo_train.py import AutoModelForMultimodalLM/AutoProcessor for the Gemma 4 architecture, but the Gemma 4 image processor requires torchvision even for a text-only SFT/DPO run. Without it you get a confusing ModuleNotFoundError: No module named 'torchvision' several frames deep inside transformers. distill_dataset.py already lists torchvision in its own install hint; the other two scripts didn't.

Changes

• Added a small supports_bf16() helper (CUDA check first, falls back to torch.backends.mps.is_macos13_or_newer() when MPS is available) and swapped in all torch.cuda.is_bf16_supported() call sites across the four scripts.
• Wrapped the AutoModelForMultimodalLM/AutoProcessor import in sft_train.py and dpo_train.py with a clear ImportError pointing at pip install torchvision when it fails.

No behavior changes on CUDA machines — supports_bf16() returns the same thing torch.cuda.is_bf16_supported() did there.

Testing

• python -m py_compile on all four modified files.
• Ran the patched sft_train.py end-to-end on an Apple Silicon Mac (MPS) against google/gemma-4-E2B-it with a small SFT dataset (--force-hf, QLoRA): training completes, and supports_bf16() now correctly returns True on this machine (bf16=True, fp16=False in the SFTConfig), where it previously returned False unconditionally.
• Confirmed the same NaN-adapter failure mode I originally hit at higher epochs/rank/LR does not reproduce once bf16 is selected instead of fp16, on the same hardware.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[keyuchen21]

recheck

[keyuchen21]

@googlebot I signed it!

[keyuchen21]

Reopening to retrigger CLA check

[keyuchen21]

@bebechien @bebechien @MrLeonix
Hi guys, when you have a chance, could you please take a look at this PR, The available checks have passed, and there are no merge conflicts. Thanks!

[xbillwork]

My claude code is loaded with Gemma4 artifacts and it flagged this: 1. is_macos13_or_newer() is off by one major version (blocking) All four copies of the new helper use: if torch.backends.mps.is_available(): return torch.backends.mps.is_macos13_or_newer() is_macos13_or_newer() is literally _mps_is_on_macos_or_newer(13, 0) — I confirmed this in torch/backends/mps/__init__.py. But MPS bf16 requires macOS 14. Three independent confirmations from installed source (torch 2.11.0, transformers 5.12.1): - torch/_dynamo/device_interface.py:528 — MpsInterface.is_bf16_supported() returns torch.backends.mps.is_macos_or_newer(14, 0) - torch/amp/autocast_mode.py:306-316 — a branch commented "Special case for MPS bfloat16 support on macOS < 14" that warns "the target dtype torch.bfloat16 is not supported on macOS versions below 14. Disabling autocast." - transformers/utils/import_utils.py — is_torch_bf16_gpu_available() uses torch.backends.mps.is_macos_or_newer(14, 0) for the MPS branch So on macOS 13.x Ventura + Apple Silicon, supports_bf16() returns True incorrectly. That feeds both torch_dtype=torch.bfloat16 in model_kwargs and bf16=True in the *Config. Those users currently get fp16 — fragile, but it runs. After this PR they get a dtype torch considers unsupported on their OS. That's a regression on a real configuration, and your testing wouldn't have caught it since your Mac is presumably on 14+.

…

On Wed, Jul 29, 2026 at 8:36 PM Keyu Chen ***@***.***> wrote: *keyuchen21* left a comment (google-gemma/gemma-skills#9) <#9 (comment)> @bebechien <https://github.com/bebechien> @bebechien <https://github.com/bebechien> @MrLeonix <https://github.com/MrLeonix> Hi guys, when you have a chance, could you please take a look at this PR, The available checks have passed, and there are no merge conflicts. Thanks! — Reply to this email directly, view it on GitHub <#9?email_source=notifications&email_token=BVEG2DLVWBMPENSXVU2L4P35HKKCFA5CNFSNUABFM5UWIORPF5TWS5BNNB2WEL2JONZXKZKDN5WW2ZLOOQXTKMJSGQ4DSNRVGUY2M4TFMFZW63VKON2WE43DOJUWEZLEUVSXMZLOOSWGM33PORSXEX3DNRUWG2Y#issuecomment-5124896551>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/BVEG2DJUJ5HBEUNL4Y4UF6L5HKKCFAVCNFSNUABGKJSXA33TNF2G64TZHMYTCOJXHA3DOOBSGY5US43TOVSTWNBZGU2DGMRTGA3TJILWAI> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

[keyuchen21]

@xbillwork Good catch — you're right, and I've pushed a fix in 77d53e8.

All four copies of the helper now use torch.backends.mps.is_macos_or_newer(14, 0) instead of is_macos13_or_newer().

While verifying this I found the impact is actually more severe than "a dtype torch considers unsupported." transformers/training_args.py validates bf16=True against is_torch_bf16_gpu_available(), whose MPS branch is exactly torch.backends.mps.is_macos_or_newer(14, 0):

if self.bf16 or self.bf16_full_eval:
    if not self.use_cpu and not is_torch_bf16_gpu_available() and not is_torch_xla_available():
        raise ValueError("Your setup doesn't support bf16/gpu. ...")

So on macOS 13.x + Apple Silicon, the old check would have made bf16=True reach SFTConfig/DPOConfig/RewardConfig and raise a hard ValueError — training wouldn't start at all, rather than just running in a fragile dtype. The fix now matches transformers' own threshold exactly, so the two can't disagree.

One note on your citations: on torch 2.13 the first two have changed — MpsInterface.is_bf16_supported() now returns True unconditionally and the macOS < 14 branch in autocast_mode.py is gone. The transformers one still uses is_macos_or_newer(14, 0) (confirmed on 5.14.1), and since that's what actually gates bf16=True, the conclusion holds either way.

Thanks for the careful review.