The current implementation removes the early break that was present in the original code, causing the loop to always iterate over all model parameters. While the number of parameters is usually manageable, we can restore the early exit optimization: if we are not casting to float16, we only need to know if any cast is required; if we are casting to float16, we can stop as soon as we find a bfloat16 weight (which triggers the warning).

    for k, v in tree_flatten(model.parameters()):
        if cast_pred(k) and mx.issubdtype(v.dtype, mx.floating) and v.dtype != target_dtype:
            needs_cast = True
            if v.dtype == mx.bfloat16:
                has_bf16 = True
            # Optimization: break early if we have enough information.
            # If target is not fp16, we don't care about has_bf16.
            # If target is fp16, we stop once we find a bf16 weight.
            if has_bf16 or target_dtype != mx.float16:
                break

Include Qwen3.5 MoE in float32 warning gate

For Qwen3.5 MoE loads, FastMLXModel.from_pretrained passes the config model_type through to _convert_mlx_dtype, and this repo already treats qwen3_5_moe as a supported Qwen3.5 architecture in unsloth_zoo/mlx/compile.py. Because _is_force_float32_arch does an exact normalized match against this list, qwen3_5_moe will not match the lone qwen3_5 entry here, so bf16→fp16 downcasts of those GDN-based models skip the warning this change is adding.

Useful? React with 👍 / 👎.

[2/2 reviewers] Med: the trailing-comma marker is a substring-delimiter trick used by unsloth/models/loader.py:1378 (where matching is disable_name.lower() in model_types_all). _is_force_float32_arch in this PR explicitly strips the comma before comparing, so the marker is a no-op in zoo's matcher. Either drop the comma here (and rely on gemma3 vs gemma3n already being separate exact-match entries) or update the comment to make clear it's only load-bearing for the substring-matching CUDA consumer.

(or, alternatively, keep "gemma3," literally and rewrite the docstring above to say the comma is a CUDA-side substring delimiter, not an exact-match marker.)

[2/2 reviewers] Nit: this from ..model_lists import FORCE_FLOAT32 is dead — FORCE_FLOAT32 is never referenced inside _convert_mlx_dtype. The consumption happens inside _is_force_float32_arch. Drop it.

[1/2 reviewers] Nit: break was removed from the original loop so we can also detect bf16 presence. But this means every call now scans all parameters even on the common case where no warning can ever fire (target_dtype == fp16 but no bf16 weights, OR model_type not in FORCE_FLOAT32). For large models this is non-trivial. Short-circuit when both flags are set:

[1/2 reviewers] Med: the warning text gives dtype=None and dtype='float32' but is missing dtype='bfloat16' (explicit-bf16 keeps native and is the typical answer on M3+) and a pointer to UNSLOTH_FORCE_FLOAT32=1 for users who want the CUDA-style guard. Make it actionable:

[1/2 reviewers] Low: _norm strips - and _ but not .. HF emits model_type values like qwen3.5 in some configs; that variant would silently miss the warning even though qwen3_5 is in FORCE_FLOAT32. Cheap to harden:

[1/2 reviewers] Nit: the test name says gemma3, does not match gemma3n but the body actually asserts that gemma3n DOES match (via its own list entry, not via gemma3,). The actual no-false-positive case is gemma3_audio_only_pretend. Rename for clarity:

Avoid suggesting dtype=None when it already selected fp16

On Apple M1/M2 the default from_pretrained(dtype=None) path resolves target_dtype to mx.float16 before calling this helper, so a bf16 Gemma/GPT-OSS/Qwen load will emit this warning even though the user already passed (or omitted) dtype=None. In that scenario the remediation text is misleading because dtype=None does not keep native bf16 on those chips; users would need dtype='float32' or a chip-specific message/explicit-dtype gate.

Useful? React with 👍 / 👎.

Avoid importing the top-level package in this unit test

When this new test module is run in the repo's CPU-only harness without the separate unsloth package installed, this import executes unsloth_zoo.__init__ and hits its find_spec("unsloth") guard before the assertion runs; I reproduced this with pytest tests/test_mlx_dtype_downcast_warning.py -q where this test fails while the submodule-based tests pass. Importing the dependency-free unsloth_zoo.model_lists module directly would keep the test exercising the new list without requiring a full Unsloth install.

Useful? React with 👍 / 👎.