Changes:
- Change default compile_mode from "reduce-overhead" to "default" since CUDA
Graphs (used by reduce-overhead) is incompatible with PEFT/LoRA
- Add _estimate_compile_threshold() to calculate minimum steps needed for
torch.compile to be beneficial based on model parameter count
- Add _apply_torch_compile() helper with accelerate unwrap_model bug workaround
- Defer torch.compile application to trainer initialization time so we can
check max_steps against the breakeven threshold
- Patch SentenceTransformerTrainer to auto-apply compile when max_steps
exceeds the calculated threshold
Breakeven thresholds (with 1.2x safety margin):
- 22M params (MiniLM): ~1388 steps
- 110M params (mpnet): ~242 steps
- 335M params (snowflake): ~203 steps
This ensures torch.compile warmup cost is only paid when training is long
enough to benefit from the speedup.
The GitHub issue check had issues:
1. Network latency on import
2. Issue being closed does not mean the fix is in the installed vLLM version
Now skip the PDL workaround if vLLM version > 0.13.2, which is when
the upstream fix is expected to be included.
- Changed absolute import to relative: from ._utils import prepare_model_for_kbit_training
- Added SUPPORTS_BFLOAT16 import for proper dtype detection
- Handle devices that don't support bfloat16 by falling back to float16
Changed from peft.prepare_model_for_kbit_training to
unsloth.models._utils.prepare_model_for_kbit_training.
Unsloth's version provides:
- Float32 mixed precision upcasting for LoRA layers
- Better numerical stability
- Consistency with rest of Unsloth codebase
The original implementation was 31% slower than naive SentenceTransformer due to
conflicting decorators from Unsloth's auto-compiler (@torch.compile on attention
modules but @torch.compiler.disable on sub-modules).
Changes:
- Add fast encoder path that bypasses Unsloth patching for encoder models
- Use native torch.compile with mode="reduce-overhead" for 6x speedup
- Auto-detect and enable SDPA for models that support it (BERT, RoBERTa, etc.)
- Change defaults: load_in_16bit=True, load_in_4bit=False (16-bit is optimal)
- Change default: use_gradient_checkpointing=False (conflicts with torch.compile)
- Add UNSLOTH_COMPILE_DISABLE=1 env var to fall back to old path if needed
Supported encoder types: mpnet, bert, distilbert, roberta, xlm-roberta, albert, electra
Benchmark results (BS=32, seq_len=128):
- Naive 16-bit LoRA: 13-50ms per iter
- Unsloth 16-bit LoRA: 2-9ms per iter (5.4x-6.7x faster)
- Memory usage: 61MB-1.3GB (even largest model fits easily)
Note: 4-bit + torch.compile has a PyTorch bug (pytorch/pytorch#90665).
4-bit is also 1.7-1.9x slower than 16-bit due to dequantization overhead,
so 16-bit is recommended for these small encoder models anyway.
- Patch vllm.lora.ops.triton_ops.utils directly where supports_pdl is defined
- Clear lru_cache before patching to prevent stale cached results
- Add fused_moe_lora_op to consumer modules list
- Use *args, **kwargs in fake function for compatibility
- Add _spec_exists helper function to reduce duplication
- Scan all GPUs for SM100 instead of just device 0
- Use loop for module patching to improve maintainability
When using base models with custom chat templates applied after loading,
vLLM's internal tokenizer may not have the chat_template set. This causes
issues during RL training with vLLM inference.
This fix syncs the chat_template from the processing_class (the tokenizer
you loaded and configured) to vLLM's internal tokenizer during trainer
initialization, but only if vLLM's tokenizer does not already have one set.
