Gemma 4 31B QLoRA on RTX 5090 — What Broke and What Fixed It

[cdimurro]

Issue 1: Blackwell bf16 CUBLAS errors

• Symptom: CUBLAS_STATUS_EXECUTION_FAILED on the very first matmul of training.
• Cause: Pre-existing torch 2.10.0+cu128 didn't have proper Blackwell (sm_120) cuBLAS kernels.
• Fix: Upgraded torch to 2.11.0+cu129.

Issue 2: Stale Triton compile cache

• Symptom: Same CUBLAS error returned even after torch upgrade.
• Cause: Unsloth's compiled cache at /tmp/unsloth_compiled_cache/ was built against the old torch and held mismatched kernels.
• Fix: rm -rf /tmp/unsloth_compiled_cache /path/to/project/unsloth_compiled_cache

Issue 3: Triton kernel launch failures (misreported as OOM)

• Symptom: CUDA driver error: out of memory even with 25+ GB free VRAM.
• Cause: torch.compile / inductor generating Triton kernels that fail to launch on sm_120.
• Fix (interim): TORCHDYNAMO_DISABLE=1 + UNSLOTH_COMPILE_DISABLE=1 (later removed when proper Unsloth build arrived).

Issue 4: Bitsandbytes ABI mismatch on cu130

• Symptom: libnvJitLink.so.13 not found, cdequantize_blockwise_fp32 symbol errors.
• Cause: When FA4 install bumped torch to cu130, bitsandbytes 0.49.2 (built against cu12) broke.
• Fix: Reverted torch to 2.11.0+cu129 to keep bnb compatible.

Issue 5: Gemma 4's hybrid-attention head_dim=512 rejection

• Symptom: FlashAttention forward only supports head dimension at most 256, but my isolated kernel test at head_dim=256 worked fine.
• Cause: Gemma 4 uses hybrid attention: 50 sliding-window layers at head_dim=256 plus 10 global-attention layers at head_dim=512. FA2 rejects >256, and the standard transformers dispatch routes all layers through FA2.
• Fix: attn_implementation="sdpa" — handles both head_dims uniformly. SDPA is slightly more memory-hungry on the global layers but works.

Issue 6: Bitsandbytes dequant OOM on 31B

• Symptom: OOM in bnb's "Error out of memory at line 81 in file /src/csrc/ops.cu" even with seq_length=512, lora_r=8.
• Cause: Standard pypi unsloth==2026.4.4 doesn't ship custom Blackwell hybrid-attention kernels — runs through generic SDPA path which doesn't fit 31B + activations in 32 GB.
• Fix: Installed Unsloth Studio via curl -fsSL unsloth.ai/install.sh | sh — gets a separate Python 3.13 venv with custom Blackwell-tuned kernels that hit the ~22 GB VRAM target.

Issue 7: Python 3.13 huggingface_hub strict-validator bug

• Symptom: Unsupported type for field 'import_name': str | None.
• Cause: huggingface_hub.dataclasses validator missing PEP 604 str | None support on Python 3.13.
• Fix: One-line patch added types.UnionType: _validate_union to _BASIC_TYPE_VALIDATORS.

Issue 8: Unsloth Studio shipped with old transformers

• Symptom: model type 'gemma4' but Transformers does not recognize this architecture.
• Cause: Studio venv had transformers 4.57.6 which predates Gemma 4 support.
• Fix: pip install --upgrade transformers in the Studio venv → version 5.5.3.

Issue 9: trl 0.23 SFTTrainer dataset schema mismatch

• Symptom: No columns in the dataset match the model's forward method signature: (messages, prompt, completion, ...) — my dataset had a "text" column from manually applying chat template.
• Cause: Newer SFTTrainer wants raw messages OR pre-tokenized input_ids / attention_mask / labels, not pre-rendered text.
• Fix: Refactored _samples_to_dataset to pre-tokenize and emit input_ids / attention_mask / labels directly. Sidesteps trl 0.23/0.24 API drift.

Issue 10: Gemma 4 multimodal processor blocked positional tokenizer call

• Symptom: 'NoneType' object is not subscriptable from processing_gemma4.py:130.
• Cause: Unsloth Zoo patches the tokenizer to be a multimodal processor that requires text= kwarg, breaking my tokenizer(text, ...) positional call.
• Fix: text_tokenizer = getattr(tokenizer, "tokenizer", tokenizer) to get the underlying text tokenizer.

Issue 11: SFTConfig vs TrainingArguments

• Symptom: Various dataset config fields rejected.
• Cause: trl 0.23+ requires SFTConfig (their unified config) instead of TrainingArguments.
• Fix: Switched to SFTConfig which carries both trainer + SFT-specific fields (max_seq_length, dataset_text_field, packing) in one object.

---

Working Stack Configuration

Component	Version	Why this version
Python	3.13 (Studio venv)	What the install script provisions
torch	2.10.0+cu130 (Studio venv)	Studio-bundled, Blackwell-supported
transformers	5.5.3	Required for gemma4 arch
unsloth	2026.4.4	With Studio kernels
trl	0.23.1	Studio-bundled, ships SFTConfig
bitsandbytes	0.49.2	Compatible with cu13 in Studio venv
Base model	unsloth/gemma-4-31B-it-unsloth-bnb-4bit	Pre-quantized for QLoRA
Attention	attn_implementation="sdpa"	Handles hybrid head_dim 256+512
LoRA	r=8, alpha=16	Fits 32 GB with all the above
seq_length	512	Conservative; smoke showed it stable

---

The Pivotal Moment

The cascade of fixes (env vars, FA2, FA4, downgrading models) was patching symptoms of a single root cause: the pypi Unsloth build doesn't have Gemma 4 hybrid-attention kernels for sm_120. Once it was confirmed that 31B works at ~22 GB on Blackwell with the Studio install, switching to that distribution plus a small set of script tweaks made everything work in a single training run.

Final Result

• Stability: No OOMs, no crashes — clean run end-to-end.
• Output: LoRA adapter saved successfully.

[JCBrouwer]

Hi @cdimurro, can you share your working training script?

The official training script here is no longer working (due to a ValueError: You should supply an encoding or a list of encodings to this method that includes input_ids, but you provided ['messages']).