From 545672c5aa8628bddf64cc80fa9e08af99609341 Mon Sep 17 00:00:00 2001 From: Vedant Madane <6527493+VedantMadane@users.noreply.github.com> Date: Mon, 2 Mar 2026 14:22:24 +0530 Subject: [PATCH] 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 | 7 ++ unsloth/grammars.py | 146 ++++++++++++++++++++++++++++++++++++++++ unsloth/import_fixes.py | 97 +++++++++++++++++++++++++- 3 files changed, 249 insertions(+), 1 deletion(-) create mode 100644 unsloth/grammars.py diff --git a/unsloth/__init__.py b/unsloth/__init__.py index dbacd551c4..5ea7384495 100644 --- a/unsloth/__init__.py +++ b/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 diff --git a/unsloth/grammars.py b/unsloth/grammars.py new file mode 100644 index 0000000000..03184d17ec --- /dev/null +++ b/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 diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index ca44a0ce7e..1d12aa3b6a 100644 --- a/unsloth/import_fixes.py +++ b/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")