Fix tool calling compatibility for Llama 3.2 and Phi-4

Fixes #3092.

This PR addresses the tool calling compatibility issues reported with Llama 3.2, Phi-4, and Mistral models.

Key Changes:
1. Compatibility Patch: Added patch_transformers_cfg() in import_fixes.py to monkey-patch transformers-cfg with better model detection and fallback to auto-inference.
2. New Helper: Introduced generate_with_grammar() in unsloth/grammars.py to provide a robust, model-agnostic way to use grammar-constrained generation.
3. Integration: Automatically applies the patch and exports the helper function when unsloth is imported.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

unsloth/__init__.py

@@ -139,6 +139,7 @@ from .import_fixes import (
     fix_vllm_guided_decoding_params,
     fix_vllm_pdl_blackwell,
     fix_triton_compiled_kernel_missing_attrs,
+    fix_rocm_triton_key_error,
     patch_trunc_normal_precision_issue,
     ignore_logger_messages,
     patch_ipykernel_hf_xet,
@@ -152,6 +153,7 @@ from .import_fixes import (
     patch_torchcodec_audio_decoder,
     disable_torchcodec_if_broken,
     disable_broken_wandb,
+    patch_transformers_cfg,
 )
 
 fix_xformers_performance_issue()
@@ -161,6 +163,7 @@ check_vllm_torch_sm100_compatibility()
 fix_vllm_guided_decoding_params()
 fix_vllm_pdl_blackwell()
 fix_triton_compiled_kernel_missing_attrs()
+fix_rocm_triton_key_error()
 patch_trunc_normal_precision_issue()
 ignore_logger_messages()
 patch_ipykernel_hf_xet()
@@ -174,6 +177,7 @@ patch_vllm_for_notebooks()
 patch_torchcodec_audio_decoder()
 disable_torchcodec_if_broken()
 disable_broken_wandb()
+patch_transformers_cfg()
 
 del fix_xformers_performance_issue
 del fix_vllm_aimv2_issue
@@ -181,6 +185,7 @@ del check_vllm_torch_sm100_compatibility
 del fix_vllm_guided_decoding_params
 del fix_vllm_pdl_blackwell
 del fix_triton_compiled_kernel_missing_attrs
+del fix_rocm_triton_key_error
 del patch_trunc_normal_precision_issue
 del ignore_logger_messages
 del patch_ipykernel_hf_xet
@@ -194,6 +199,7 @@ del patch_vllm_for_notebooks
 del patch_torchcodec_audio_decoder
 del disable_torchcodec_if_broken
 del disable_broken_wandb
+del patch_transformers_cfg
 
 # Torch 2.4 has including_emulation
 if DEVICE_TYPE == "cuda":
@@ -314,6 +320,7 @@ from .save import *
 from .chat_templates import *
 from .tokenizer_utils import *
 from .trainer import *
+from .grammars import *
 
 # Export dataprep utilities for CLI and downstream users
 from .dataprep.raw_text import RawTextDataLoader, TextPreprocessor

unsloth/grammars.py

@@ -0,0 +1,146 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import logging
+from transformers import AutoTokenizer
+from typing import Optional, Union, List, Any
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "generate_with_grammar",
+]
+
+# Simple JSON array grammar
+JSON_ARR_GBNF = r"""
+root ::= arr
+value ::= object | array | string | number | ("true" | "false" | "null") ws
+arr ::=
+  "[
+" ws (
+            value
+            (",
+" ws value)*
+        )? "]"
+object ::=
+  "{" ws (
+            string ":" ws value
+            ("," ws string ":" ws value)*
+        )? "}" ws
+array ::=
+  "[" ws (
+            value
+            ("," ws value)*
+        )? "]" ws
+string ::=
+  """ (
+    [^"\x7F\x00-\x1F] |
+    "" (["\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]) # escapes
+  )* """ ws
+number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws
+ws ::= ([ 	
+] ws)?
+"""
+
+def generate_with_grammar(
+    model,
+    tokenizer,
+    input_ids: torch.Tensor,
+    grammar_str: Optional[str] = None,
+    start_rule: str = "root",
+    max_new_tokens: int = 256,
+    temperature: Optional[float] = None,
+    top_p: Optional[float] = None,
+    top_k: Optional[int] = None,
+    do_sample: bool = False,
+    repetition_penalty: float = 1.1,
+    num_return_sequences: int = 1,
+    **kwargs
+):
+    """
+    Generate text with grammar constraints using transformers-cfg.
+    Automatically handles model-specific parameter compatibility.
+    
+    Args:
+        model: The model to generate with.
+        tokenizer: The tokenizer associated with the model.
+        input_ids: Input token IDs.
+        grammar_str: EBNF grammar string. Defaults to a JSON array grammar.
+        start_rule: The start rule for the grammar. Defaults to "root".
+        max_new_tokens: Maximum number of tokens to generate.
+        temperature: Sampling temperature.
+        top_p: Top-p sampling.
+        top_k: Top-k sampling.
+        do_sample: Whether to use sampling.
+        repetition_penalty: Penalty for repeating tokens.
+        num_return_sequences: Number of sequences to return.
+        **kwargs: Additional keyword arguments for model.generate.
+    """
+    try:
+        from transformers_cfg.grammar_utils import IncrementalGrammarConstraint
+        from transformers_cfg.generation.logits_process import GrammarConstrainedLogitsProcessor
+    except ImportError:
+        raise ImportError(
+            "Unsloth: Please install transformers-cfg to use grammar-constrained generation: "
+            "`pip install transformers-cfg`"
+        )
+
+    # Use default JSON array grammar if none provided
+    if grammar_str is None:
+        grammar_str = JSON_ARR_GBNF
+
+    # Create grammar constraint
+    grammar = IncrementalGrammarConstraint(
+        grammar_str, 
+        start_rule_name=start_rule, 
+        tokenizer=tokenizer
+    )
+    grammar_processor = GrammarConstrainedLogitsProcessor(grammar)
+
+    # Filter generation kwargs
+    generation_kwargs = {
+        "input_ids": input_ids,
+        "max_new_tokens": max_new_tokens,
+        "do_sample": do_sample,
+        "repetition_penalty": repetition_penalty,
+        "num_return_sequences": num_return_sequences,
+        "logits_processor": [grammar_processor],
+        **kwargs
+    }
+
+    # Handle sampling parameters
+    if do_sample:
+        if temperature is not None: generation_kwargs["temperature"] = temperature
+        if top_p is not None: generation_kwargs["top_p"] = top_p
+        if top_k is not None: generation_kwargs["top_k"] = top_k
+    else:
+        generation_kwargs["temperature"] = None
+        generation_kwargs["top_p"] = None
+        generation_kwargs["top_k"] = None
+
+    # Try to generate
+    try:
+        return model.generate(**generation_kwargs)
+    except ValueError as e:
+        # Some models fail if specific kwargs like 'sliding_window' or 'num_logits_to_keep' are passed
+        if "model_kwargs" in str(e):
+            logger.warning("Unsloth: Generation failed with model_kwargs error. Retrying with minimal parameters...")
+            minimal_kwargs = {
+                "input_ids": input_ids,
+                "max_new_tokens": max_new_tokens,
+                "logits_processor": [grammar_processor],
+            }
+            return model.generate(**minimal_kwargs)
+        raise

unsloth/import_fixes.py

@@ -941,6 +941,39 @@ def fix_triton_compiled_kernel_missing_attrs():
     )
 
 
+def fix_rocm_triton_key_error():
+    """
+    ROCm + torch.compile can fail if Triton lacks `triton_key`.
+    Disable Inductor/compile only on ROCm when that symbol is missing.
+    """
+    try:
+        import torch
+    except (ImportError, ModuleNotFoundError):
+        return
+
+    if not getattr(torch.version, "hip", None):
+        return
+
+    try:
+        import triton
+    except (ImportError, ModuleNotFoundError):
+        return
+
+    try:
+        from triton.runtime import triton_key  # noqa: F401
+
+        return
+    except ImportError:
+        pass
+
+    os.environ.setdefault("TORCHINDUCTOR_DISABLE", "1")
+    os.environ.setdefault("TORCH_COMPILE_DISABLE", "1")
+    logger.info(
+        "Unsloth: ROCm detected and Triton lacks triton_key; "
+        "disabling torch.compile/Inductor to avoid backend crash."
+    )
+
+
 def patch_trunc_normal_precision_issue():
     """
     Patch torch.nn.init.trunc_normal_ for low precision tensors to run init in fp32.
@@ -1218,7 +1251,6 @@ def fix_vllm_pdl_blackwell():
         # Just set the env var - vLLM might be an older version without supports_pdl
         logger.info(f"Unsloth: Set TRITON_DISABLE_PDL=1 for SM100 ({sm100_gpu_name})")
 
-
 def patch_openspiel_env_async():
     """Apply nest_asyncio for OpenEnv EnvClient async compatibility.
 
@@ -1821,3 +1853,66 @@ def disable_broken_causal_conv1d():
         "Unsloth: Detected broken causal_conv1d binary; "
         "disabling causal_conv1d fast path and continuing import."
     )
+
+
+def patch_transformers_cfg():
+    """
+    Fix for Unsloth Issue #3092: Tool Calling Compatibility.
+    Monkey-patches transformers-cfg to support more models (Llama 3.2, Phi-4, Mistral).
+    """
+    if importlib.util.find_spec("transformers_cfg") is None:
+        return
+
+    try:
+        from transformers_cfg.tokenization.mapping.token2byte import (
+            Token2ByteMapping,
+            GPT2Token2ByteMapping,
+        )
+        from transformers import PreTrainedTokenizerFast
+    except ImportError:
+        return
+
+    # Store original method
+    _original_from_hf_tokenizer = Token2ByteMapping.from_hf_tokenizer
+
+    @classmethod
+    def _patched_from_hf_tokenizer(cls, hf_tokenizer):
+        """Patched version with better model support"""
+        # Try the original method first
+        try:
+            return _original_from_hf_tokenizer(hf_tokenizer)
+        except (NotImplementedError, AssertionError):
+            # Original method failed, try fallbacks
+            if not isinstance(hf_tokenizer, PreTrainedTokenizerFast):
+                # If it's not a fast tokenizer, we can't help
+                raise
+
+            model_path_lower = hf_tokenizer.name_or_path.lower()
+
+            # Check for Phi models (Phi-3, Phi-4)
+            if "phi" in model_path_lower:
+                return GPT2Token2ByteMapping(hf_tokenizer)
+
+            # Check for Llama 3.x variants (more flexible matching than original)
+            if "llama" in model_path_lower and any(
+                v in model_path_lower for v in ["3.1", "3.2", "3.3", "3-", "3_"]
+            ):
+                return GPT2Token2ByteMapping(hf_tokenizer)
+
+            # Check for Mistral/Mixtral models
+            if "mistral" in model_path_lower or "mixtral" in model_path_lower:
+                return GPT2Token2ByteMapping(hf_tokenizer)
+
+            # Try auto-inference as last resort
+            try:
+                return Token2ByteMapping.auto_infer(hf_tokenizer)
+            except:
+                raise NotImplementedError(
+                    f"Tokenizer not supported: {hf_tokenizer.__class__.__name__} "
+                    f"for model: {hf_tokenizer.name_or_path}. "
+                    "Auto-inference also failed."
+                )
+
+    # Apply the patch
+    Token2ByteMapping.from_hf_tokenizer = _patched_from_hf_tokenizer
+    logger.info("Unsloth: Patched transformers-cfg for enhanced model compatibility")