[Qwen 3.5][gemma4] Qwen35 and Gemma 4 fast inference

[danielhanchen]

Staging mirror of unslothai#588

Original PR: unslothai#588
Author: Datta0

This is a staging copy for review and editing. Once finalized, changes will be pushed back to the original PR.

---

Original description

• fast inference support for Qwen 3.5 with vLLM :)
  Tried to make sure the changes are minimal so when we detect linear_attn, we hand it off to a separate function

---

This PR contains code changes only (3 files). Test changes are in a separate PR.

Changed files:

• unsloth_zoo/empty_model.py
• unsloth_zoo/hf_utils.py
• unsloth_zoo/vllm_utils.py

[danielhanchen]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for Gemma 4 and Qwen 3.5 (GDN) architectures, adding specialized layer extraction logic, vLLM patches for LoRA and BitsAndBytes, and a centralized model finalization process. The review feedback identifies a logic flaw in GDN quantization state preservation where only the first shard is checked, and points out that Gemma 4 k_eq_v layers require explicit mapping of the K shard to the V projection during extraction to avoid uninitialized weights. Additionally, a debug print statement should be removed or replaced with a logger call.

[gemini-code-assist]

The logic here assumes that if only one shard is quantized, it must be at index 0. If qkv_states[1] or qkv_states[2] is the only non-None state, it will be ignored, and _store_quant_state will be called with None, leading to loss of quantization information for the fused layer. Consider using the first non-None state found in qkv_states to ensure the quantization state is preserved.

Suggested change

	if isinstance(qs_attr, dict):
	_store_quant_state(f"{prefix}.in_proj_qkv", qkv_states[0])
	if isinstance(qs_attr, dict):
	_store_quant_state(f"{prefix}.in_proj_qkv", next((qs for qs in qkv_states if qs is not None), None))

[gemini-code-assist]

Skipping the extraction of v_proj for k_eq_v layers will result in an incomplete state dictionary for the Hugging Face model, as the v_proj weights will remain uninitialized (or initialized with tiny dimensions from the empty model creation). Since these layers reuse K as V, you should extract shard 1 (K) and store it as v_proj to ensure the HF model is correctly populated.

Suggested change

	if kk not in gemma4_k_eq_v_layers:
	get_state_dict(f"{prefix}.v_proj", 2, state_dict, qkv_proj)
	v_shard_idx = 1 if kk in gemma4_k_eq_v_layers else 2
	get_state_dict(f"{prefix}.v_proj", v_shard_idx, state_dict, qkv_proj)

[gemini-code-assist]

This debug print statement should be removed or replaced with a proper logger call (e.g., logger.info) to maintain clean output during model conversion.

[danielhanchen]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for Gemma 4 and Qwen 3.5 (GDN) architectures, including specialized logic for Gated Delta Net (GDN) layers and Gemma 4's shared K/V attention. Key updates include new vLLM patching mechanisms for LoRA and BitsAndBytes support, expanded layer configuration templates, and a centralized finalize_huggingface_model function. A critical issue was identified in the GDN layer extraction logic where merging mixed quantized and non-quantized shards into a single weight tensor could lead to weight corruption; a suggestion was provided to ensure all shards are dequantized if any quantization is present.

[gemini-code-assist]

The logic for merging quantized qkv shards in GDN layers is problematic when sum(qs is not None for qs in qkv_states) == 1. Merging a quantized shard with other non-quantized shards into a single in_proj_qkv.weight tensor while only storing the quant state of the first shard will lead to corrupted weights for the other shards during dequantization. Additionally, if the single quant state is at index 1 or 2, it is currently lost. For fused layers where shards are merged into one HF weight, you should dequantize all shards if any quant state is present to ensure correctness.

Suggested change

	if sum(qs is not None for qs in qkv_states) > 1:
	try:
	from bitsandbytes.functional import dequantize_4bit
	except Exception:
	raise RuntimeError(
	"Unsloth: prequantized BnB Qwen3.5 GDN requires bitsandbytes for fused in_proj_qkv reconstruction."
	)
	parts = []
	for i, qs in enumerate(qkv_states):
	shard = weight[offsets[i]:offsets[i + 1]]
	parts.append(dequantize_4bit(shard, quant_state=qs) if qs is not None else shard)
	store(f"{prefix}.in_proj_qkv.weight", torch.cat(parts, dim=0))
	else:
	store(f"{prefix}.in_proj_qkv.weight", qkv_weight)
	if isinstance(qs_attr, dict):
	_store_quant_state(f"{prefix}.in_proj_qkv", qkv_states[0])
	if any(qs is not None for qs in qkv_states):
	try:
	from bitsandbytes.functional import dequantize_4bit
	except Exception:
	raise RuntimeError(
	"Unsloth: prequantized BnB Qwen3.5 GDN requires bitsandbytes for fused in_proj_qkv reconstruction."
	)
	parts = []
	for i, qs in enumerate(qkv_states):
	shard = weight[offsets[i]:offsets[i + 1]]
	parts.append(dequantize_4bit(shard, quant_state=qs) if qs is not None else shard)
	store(f"{prefix}.in_proj_qkv.weight", torch.cat(parts, dim=0))
	else:
	store(f"{prefix}.in_proj_qkv.weight", qkv_weight)