From e339860ba0e114c38f4aa62e5f76fe5a1178afa8 Mon Sep 17 00:00:00 2001 From: DoubleMathew Date: Sat, 12 Jul 2025 17:52:24 -0500 Subject: [PATCH 1/5] patch falcon h1 inference (#2932) --- unsloth/models/falcon_h1.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/unsloth/models/falcon_h1.py b/unsloth/models/falcon_h1.py index 2cbb78f8ad..0db9c1ca4e 100644 --- a/unsloth/models/falcon_h1.py +++ b/unsloth/models/falcon_h1.py @@ -43,6 +43,9 @@ except: from transformers.modeling_attn_mask_utils import ( _prepare_4d_causal_attention_mask_for_sdpa, ) +from transformers.utils import ( + is_torchdynamo_compiling, +) # For Pytorch 2.1.1 try: from transformers.models.falcon_h1.modeling_falcon_h1 import ( @@ -519,7 +522,7 @@ def _FalconH1_fast_forward_inference(attention_fast_forward_inference=FalconH1At attention_mask = attention_mask, do_prefill = not hasattr(decoder_layer.self_attn, "paged_attention"), ) - attention_hidden_states = attention_hidden_states * decoder_layer.attention_out_multiplier + attention_hidden_states = attention_hidden_states * decoder_layer.attn_out_multiplier mamba_hidden_states = decoder_layer.mamba( hidden_states=X, cache_params=present_key_value, @@ -595,15 +598,17 @@ def _fast_prepare_inputs_for_generation( input_ids = input_ids[:, -cache_position.shape[0] :] elif input_ids.shape[1] != cache_position.shape[0]: # Default case (the "else", a no op, is Exception 2) input_ids = input_ids[:, cache_position] - else: - past_key_values = FalconHybridMambaAttentionDynamicCache( - self.config, - input_ids.shape[0], - self.dtype, - devices=[ - self.model.layers[i].mamba.conv1d.weight.device for i in range(self.config.num_hidden_layers) - ], - ) + pass + # TODO: Wire up Cache to work for inference. + # else: + # past_key_values = FalconHybridMambaAttentionDynamicCache( + # self.config, + # input_ids.shape[0], + # self.dtype, + # devices=[ + # self.model.layers[i].mamba.conv1d.weight.device for i in range(self.config.num_hidden_layers) + # ], + # ) if attention_mask is not None and position_ids is None: # create position_ids on the fly for batch generation From 1898b6d049d606ec88f3f9307172373776eec0f6 Mon Sep 17 00:00:00 2001 From: Datta Nimmaturi Date: Sun, 13 Jul 2025 04:23:07 +0530 Subject: [PATCH 2/5] Fix falcon H1 dropout issue (#2938) Because we don't have down and gate multipliers, the MLP output values are too huge, causing NaN and unstable training. To bypass that lets rely on HF's implementation for the time being --- unsloth/models/llama.py | 57 +++++++++++++++++++---------------------- 1 file changed, 26 insertions(+), 31 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index aae86080a8..0c6f6c1e86 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2768,40 +2768,35 @@ class FastLlamaModel: if lora_dropout == 0 and bias == "none": for idx, layer in enumerate(model.model.model.layers): - # Determine MLP module name (falcon_h1 has feed_forward, llama style has mlp) - if hasattr(layer, "mlp"): - mlp_module_name = "mlp" - elif hasattr(layer, "feed_forward"): - mlp_module_name = "feed_forward" - else: - logger.warning_once(f"Unsloth: No MLP module found in layer {idx} so skipping peft mlp patching") - continue + if model_type != "falcon_h1": + # LoRAMLP.apply doesn't have functionality for gate and down mutlipliers yet. + # Don't patch falcon h1 for the time being. - mlp_module = getattr(layer, mlp_module_name) + # MLP patching + mlp_module = layer.mlp + gate_proj = mlp_module.gate_proj + up_proj = mlp_module. up_proj + down_proj = mlp_module.down_proj - # MLP patching - gate_proj = mlp_module.gate_proj - up_proj = mlp_module. up_proj - down_proj = mlp_module.down_proj + if hasattr(gate_proj, "lora_A") and \ + hasattr( up_proj, "lora_A") and \ + hasattr(down_proj, "lora_A") and \ + (getattr(gate_proj, "base_layer", gate_proj).bias is None) and \ + (getattr( up_proj, "base_layer", up_proj).bias is None) and \ + (getattr(down_proj, "base_layer", down_proj).bias is None) and \ + (len(getattr(gate_proj, "lora_magnitude_vector", []) or []) == 0) and \ + (len(getattr( up_proj, "lora_magnitude_vector", []) or []) == 0) and \ + (len(getattr(down_proj, "lora_magnitude_vector", []) or []) == 0): - if hasattr(gate_proj, "lora_A") and \ - hasattr( up_proj, "lora_A") and \ - hasattr(down_proj, "lora_A") and \ - (getattr(gate_proj, "base_layer", gate_proj).bias is None) and \ - (getattr( up_proj, "base_layer", up_proj).bias is None) and \ - (getattr(down_proj, "base_layer", down_proj).bias is None) and \ - (len(getattr(gate_proj, "lora_magnitude_vector", []) or []) == 0) and \ - (len(getattr( up_proj, "lora_magnitude_vector", []) or []) == 0) and \ - (len(getattr(down_proj, "lora_magnitude_vector", []) or []) == 0): - - # https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module - mlp_module.forward = types.MethodType(_apply_lora_mlp, mlp_module) - n_mlp += 1 - else: - logger.warning_once( - "Not an error, but Unsloth cannot patch MLP layers with our manual autograd engine since either LoRA adapters\n"\ - "are not enabled or a bias term (like in Qwen) is used." - ) + # https://stackoverflow.com/questions/50599045/python-replacing-a-function-within-a-class-of-a-module + mlp_module.forward = types.MethodType(_apply_lora_mlp, mlp_module) + n_mlp += 1 + else: + logger.warning_once( + "Not an error, but Unsloth cannot patch MLP layers with our manual autograd engine since either LoRA adapters\n"\ + "are not enabled or a bias term (like in Qwen) is used." + ) + pass pass # QKV attention patching From cf5bf081fb9974f1f381ecb7d54e4588414846cd Mon Sep 17 00:00:00 2001 From: Muzammil Khan <116030715+muzzlol@users.noreply.github.com> Date: Mon, 14 Jul 2025 14:12:07 +0530 Subject: [PATCH 3/5] fix: change lora_dropout from int to float for type consistency (#2949) Fixes "Argument of type 'float' cannot be assigned to parameter 'lora_dropout' of type 'int'" error by ensuring lora_dropout is consistently a float (0.0) rather than int (0) across vision.py, llama.py, and unsloth-cli.py --- unsloth-cli.py | 2 +- unsloth/models/llama.py | 2 +- unsloth/models/vision.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/unsloth-cli.py b/unsloth-cli.py index b7613f92df..86f02075f1 100644 --- a/unsloth-cli.py +++ b/unsloth-cli.py @@ -182,7 +182,7 @@ if __name__ == "__main__": lora_group = parser.add_argument_group("🧠 LoRA Options", "These options are used to configure the LoRA model.") lora_group.add_argument('--r', type=int, default=16, help="Rank for Lora model, default is 16. (common values: 8, 16, 32, 64, 128)") lora_group.add_argument('--lora_alpha', type=int, default=16, help="LoRA alpha parameter, default is 16. (common values: 8, 16, 32, 64, 128)") - lora_group.add_argument('--lora_dropout', type=float, default=0, help="LoRA dropout rate, default is 0.0 which is optimized.") + lora_group.add_argument('--lora_dropout', type=float, default=0.0, help="LoRA dropout rate, default is 0.0 which is optimized.") lora_group.add_argument('--bias', type=str, default="none", help="Bias setting for LoRA") lora_group.add_argument('--use_gradient_checkpointing', type=str, default="unsloth", help="Use gradient checkpointing") lora_group.add_argument('--random_state', type=int, default=3407, help="Random state for reproducibility, default is 3407.") diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 0c6f6c1e86..3551c6ca5e 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2220,7 +2220,7 @@ class FastLlamaModel: target_modules = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], lora_alpha = 16, - lora_dropout = 0, + lora_dropout = 0.0, bias = "none", layers_to_transform = None, layers_pattern = None, diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 9899d60440..7442f07e73 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -554,7 +554,7 @@ class FastBaseModel: r = 16, target_modules = None, lora_alpha = 16, - lora_dropout = 0, + lora_dropout = 0.0, bias = "none", finetune_vision_layers = True, finetune_language_layers = True, From c972010449c21f85ff0a42b05c8e5e13f0856e1c Mon Sep 17 00:00:00 2001 From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Date: Mon, 14 Jul 2025 11:43:33 +0300 Subject: [PATCH 4/5] fix dataloader_num_workers value error in GRPOTrainer (#2944) --- unsloth/models/rl.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index ae01469acc..45b8ca6334 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -168,7 +168,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): trainer = eval(f"trl.trainer.{trainer_file}") except Exception as error: return - + # Get SFTTrainer and SFTConfig names name = [x for x in dir(trainer) if x.endswith("Trainer") and x != "Trainer" and trainer_file.split("_")[0] in x.lower()] config = [x for x in dir(trainer) if x.endswith("Config") and x != "Config" and trainer_file.split("_")[0] in x.lower()] @@ -484,7 +484,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): "dataloader_persistent_workers" : True, # Keeps dataloader in RAM "dataloader_prefetch_factor" : 2, "dataloader_pin_memory" : True, - "dataloader_num_workers" : 0, # Default is 0 means 1 + "dataloader_num_workers" : 1, } for k, v in replacements.items(): x = f"{k}( = [^,\n]{{1,}})?,\n" @@ -565,7 +565,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): pass # Check GRPO num_generations mismatch - if "per_device_train_batch_size" in call_args and "num_generations" in call_args: + if "per_device_train_batch_size" in call_args and "num_generations" in call_args: check_num_generations = \ "if (per_device_train_batch_size // num_generations) * num_generations != per_device_train_batch_size:\n"\ " print('Unsloth: We now expect `per_device_train_batch_size` to be a multiple of `num_generations`.\\n"\ @@ -576,7 +576,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): pass # Check temperature must not be <= 0. Also stop if >= 10 - if "temperature" in call_args: + if "temperature" in call_args: check_temperature = \ "if temperature <= 0:\n"\ " raise MathError('Unsloth: Please set a positive non-zero temperature since your results will be wrong.')\n"\ @@ -625,7 +625,7 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): if "SamplingParams" in old_RLTrainer_source: RL_pre = RL_pre + "\n" + inspect.getsource(vLLMSamplingParams) pass - + # Selective log softmax selective_log_softmax_code = inspect.getsource(selective_log_softmax) @@ -651,12 +651,12 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): selective_log_softmax_code = selective_log_softmax_code, ) - + if RLTrainer_name == "SFTTrainer": original_text = 'self._signature_columns = ["input_ids", "attention_mask", "completion_mask"]' new_text = 'self._signature_columns = ["input_ids", "attention_mask", "completion_mask","labels"]' RLTrainer_source = RLTrainer_source.replace(original_text, new_text) - + # Remove multiple doc strings if __RLConfig_doc__ != "" and RLTrainer_source.count(__RLTrainer_doc__) == 2: RLTrainer_source = RLTrainer_source.replace(__RLTrainer_doc__, "", 1) @@ -673,12 +673,12 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"): imports, overwrite = False, ) - + # Patch Trainer exec(f"trl.{RLTrainer_name} = created_module.Unsloth{RLTrainer_name}", locals(), globals()) exec(f"trl.trainer.{RLTrainer_name} = created_module.Unsloth{RLTrainer_name}", locals(), globals()) exec(f"trl.trainer.{trainer_file}.{RLTrainer_name} = created_module.Unsloth{RLTrainer_name}", locals(), globals()) - + # Patch Config exec(f"trl.{RLConfig_name} = created_module.Unsloth{RLConfig_name}", locals(), globals()) exec(f"trl.trainer.{RLConfig_name} = created_module.Unsloth{RLConfig_name}", locals(), globals()) @@ -754,7 +754,7 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import new_vllm_part, flags = re.MULTILINE | re.DOTALL, ) - + if len(sampling_params) == 1: sampling_params = sampling_params[0] # Fix guided_decoding @@ -768,7 +768,7 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import sampling_params = \ " "*12 + "self.llm = model.vllm_engine; self._last_loaded_step = 0; " + \ sampling_params # Add spaces - + # count the indentation of last line of sampling_params. last_line = sampling_params.split("\n")[-1] last_prev_line = sampling_params.split("\n")[-2] @@ -844,7 +844,7 @@ def patch_functions(RLTrainer, trainer_file, RLTrainer_name, all_imports, import r"", source, ) - + # Replace self.llm.generate and self.llm.chat lora_name = trainer_file + "_lora_model" source = re.sub( From 0eb61fbea728cdc8acd1f2fa1f6f71074f559ac0 Mon Sep 17 00:00:00 2001 From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Date: Mon, 14 Jul 2025 12:41:15 +0300 Subject: [PATCH 5/5] GRPO Fix - Support vllm pre-dequantized quantization states in fast_dequantize kernel (#2943) * Support pre-dequantized quantization states in fast_dequantize kernel * has_nested_quant conditional set to only * Update utils.py * Update utils.py --------- Co-authored-by: Daniel Han --- unsloth/kernels/utils.py | 282 ++++++++++++++++++++++----------------- 1 file changed, 162 insertions(+), 120 deletions(-) diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 645319d423..eac2b59974 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -61,7 +61,6 @@ else: pass pass - def calculate_settings(n : int) -> (int, int,): BLOCK_SIZE : int = next_power_of_2(n) if BLOCK_SIZE > MAX_FUSED_SIZE: @@ -87,7 +86,7 @@ else: # https://github.com/bitsandbytes-foundation/bitsandbytes/pull/1330/files HAS_CUDA_STREAM = Version(bnb.__version__) > Version("0.43.3") get_ptr = bnb.functional.get_ptr - +pass if DEVICE_COUNT > 1: if DEVICE_TYPE == "cuda": @@ -97,7 +96,7 @@ if DEVICE_COUNT > 1: else: from contextlib import nullcontext def torch_gpu_device(device): return nullcontext() - pass +pass # INTEL GPU Specific Logic if DEVICE_TYPE == "xpu": @@ -105,13 +104,13 @@ if DEVICE_TYPE == "xpu": # NVIDIA GPU Default Logic else: _gpu_getCurrentRawStream = torch._C._cuda_getCurrentRawStream +pass c_void_p = ctypes.c_void_p def _get_tensor_stream(tensor: torch_Tensor) -> c_void_p: return c_void_p(_gpu_getCurrentRawStream(tensor.device.index)) pass - # Get array of CUDA streams and other buffers global CUDA_STREAMS global XPU_STREAMS @@ -124,7 +123,7 @@ if DEVICE_TYPE == "xpu": (index := torch.xpu.device(i).idx) : ctypes.c_void_p(torch._C._xpu_getCurrentRawStream(index)) for i in range(DEVICE_COUNT) } - XPU_STREAMS = [None] * (max(_XPU_STREAMS.keys()) + 1) + XPU_STREAMS = [None] * (max(_XPU_STREAMS.keys()) + 1) WEIGHT_BUFFERS = [None] * (max(_XPU_STREAMS.keys()) + 1) ABSMAX_BUFFERS = [None] * (max(_XPU_STREAMS.keys()) + 1) for k, v in _XPU_STREAMS.items(): @@ -143,7 +142,7 @@ else: for k, v in _CUDA_STREAMS.items(): CUDA_STREAMS[k] = v CUDA_STREAMS = tuple(CUDA_STREAMS) del _CUDA_STREAMS - +pass # Bitsandbytes operations ctypes_c_int = ctypes.c_int @@ -172,12 +171,15 @@ else: cdequantize_blockwise_bf16_nf4 = bnb.functional.lib.cdequantize_blockwise_bf16_nf4 cgemm_4bit_inference_naive_fp16 = bnb.functional.lib.cgemm_4bit_inference_naive_fp16 cgemm_4bit_inference_naive_bf16 = bnb.functional.lib.cgemm_4bit_inference_naive_bf16 +pass torch_mm = torch.mm torch_mv = torch.mv -torch_matmul = torch.matmul -torch_addmm = torch.addmm -torch_empty = torch.empty +torch_matmul = torch.matmul +torch_addmm = torch.addmm +torch_empty = torch.empty +torch_float16 = torch.float16 +torch_float32 = torch.float32 def QUANT_STATE(W): return getattr(W, "quant_state", None) @@ -235,23 +237,28 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM: def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False): # TODO: After adding XPU BNB support, check this function if quant_state is None: return W + is_double_quantized = True if type(quant_state) is not list: # New quant_state as a class # https://github.com/TimDettmers/bitsandbytes/pull/763/files - absmax = quant_state.absmax - shape = quant_state.shape - dtype = quant_state.dtype - blocksize = quant_state.blocksize - offset = quant_state.offset - state2 = quant_state.state2 - absmax2 = state2.absmax - code2 = state2.code - blocksize2 = state2.blocksize + absmax = quant_state.absmax + shape = quant_state.shape + dtype = quant_state.dtype + blocksize = quant_state.blocksize + offset = quant_state.offset + state2 = quant_state.state2 + is_double_quantized = state2 is not None + if is_double_quantized: + absmax2 = state2.absmax + code2 = state2.code + blocksize2 = state2.blocksize else: # Old quant_state as a list of lists absmax, shape, dtype, blocksize, compressed_stats, _, _ = quant_state offset, state2 = compressed_stats - absmax2, code2, blocksize2, _, _, _, _ = state2 + is_double_quantized = state2 is not None + if is_double_quantized: + absmax2, code2, blocksize2, _, _, _, _ = state2 pass global XPU_STREAMS device = W.device @@ -270,7 +277,7 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM: ABSMAX_BUFFER = ABSMAX_BUFFERS[device_index] if WEIGHT_BUFFER is None: WEIGHT_BUFFERS[device_index] = WEIGHT_BUFFER = torch_empty(size, dtype = dtype, device = device, requires_grad = False) - ABSMAX_BUFFERS[device_index] = ABSMAX_BUFFER = torch_empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) + ABSMAX_BUFFERS[device_index] = ABSMAX_BUFFER = torch_empty(n_elements_absmax, dtype = torch_float32, device = device, requires_grad = False) if size > WEIGHT_BUFFER.numel(): WEIGHT_BUFFER.resize_(size) if n_elements_absmax > ABSMAX_BUFFER.numel(): ABSMAX_BUFFER.resize_(n_elements_absmax) @@ -283,20 +290,23 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM: else: assert(out.shape == shape) assert(out.dtype == dtype) - out_absmax = torch_empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) + out_absmax = torch_empty(n_elements_absmax, dtype = torch_float32, device = device, requires_grad = False) pass # NF4 dequantization of statistics - ptr_out_absmax = get_ptr(out_absmax) with torch_gpu_device(device): - cdequantize_blockwise_fp32( - get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), ptr_out_absmax, - ctypes_c_int(blocksize2), ctypes_c_int(n_elements_absmax), XPU_STREAM - ) - out_absmax += offset + if is_double_quantized: + ptr_out_absmax = get_ptr(out_absmax) + cdequantize_blockwise_fp32( + get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), ptr_out_absmax, + ctypes_c_int(blocksize2), ctypes_c_int(n_elements_absmax), XPU_STREAM + ) + out_absmax += offset + else: + ptr_out_absmax = get_ptr(absmax) # Dequantize W - fx = cdequantize_blockwise_fp16_nf4 if dtype == torch.float16 else \ + fx = cdequantize_blockwise_fp16_nf4 if dtype == torch_float16 else \ cdequantize_blockwise_bf16_nf4 fx(get_ptr(None), get_ptr(W), ptr_out_absmax, get_ptr(out), ctypes_c_int(blocksize), ctypes_c_int(out.numel()), XPU_STREAM,) @@ -310,23 +320,28 @@ elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM: @torch.inference_mode def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False): if quant_state is None: return W + is_double_quantized = True if type(quant_state) is not list: # New quant_state as a class # https://github.com/TimDettmers/bitsandbytes/pull/763/files - absmax = quant_state.absmax - shape = quant_state.shape - dtype = quant_state.dtype - blocksize = quant_state.blocksize - offset = quant_state.offset - state2 = quant_state.state2 - absmax2 = state2.absmax - code2 = state2.code - blocksize2 = state2.blocksize + absmax = quant_state.absmax + shape = quant_state.shape + dtype = quant_state.dtype + blocksize = quant_state.blocksize + offset = quant_state.offset + state2 = quant_state.state2 + is_double_quantized = state2 is not None + if is_double_quantized: + absmax2 = state2.absmax + code2 = state2.code + blocksize2 = state2.blocksize else: # Old quant_state as a list of lists absmax, shape, dtype, blocksize, compressed_stats, _, _ = quant_state offset, state2 = compressed_stats - absmax2, code2, blocksize2, _, _, _, _ = state2 + is_double_quantized = state2 is not None + if is_double_quantized: + absmax2, code2, blocksize2, _, _, _, _ = state2 pass global CUDA_STREAMS device = W.device @@ -346,7 +361,7 @@ elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM: ABSMAX_BUFFER = ABSMAX_BUFFERS[device_index] if WEIGHT_BUFFER is None: WEIGHT_BUFFERS[device_index] = WEIGHT_BUFFER = torch_empty(size, dtype = dtype, device = device, requires_grad = False) - ABSMAX_BUFFERS[device_index] = ABSMAX_BUFFER = torch_empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) + ABSMAX_BUFFERS[device_index] = ABSMAX_BUFFER = torch_empty(n_elements_absmax, dtype = torch_float32, device = device, requires_grad = False) if size > WEIGHT_BUFFER.numel(): WEIGHT_BUFFER.resize_(size) if n_elements_absmax > ABSMAX_BUFFER.numel(): ABSMAX_BUFFER.resize_(n_elements_absmax) @@ -359,20 +374,22 @@ elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM: else: assert(out.shape == shape) assert(out.dtype == dtype) - out_absmax = torch_empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) + out_absmax = torch_empty(n_elements_absmax, dtype = torch_float32, device = device, requires_grad = False) pass # NF4 dequantization of statistics - ptr_out_absmax = get_ptr(out_absmax) with torch_gpu_device(device): - cdequantize_blockwise_fp32( - get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), ptr_out_absmax, - ctypes_c_int(blocksize2), ctypes_c_int(n_elements_absmax), CUDA_STREAM - ) - out_absmax += offset - + if is_double_quantized: + ptr_out_absmax = get_ptr(out_absmax) + cdequantize_blockwise_fp32( + get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), ptr_out_absmax, + ctypes_c_int(blocksize2), ctypes_c_int(n_elements_absmax), CUDA_STREAM + ) + out_absmax += offset + else: + ptr_out_absmax = get_ptr(absmax) # Dequantize W - fx = cdequantize_blockwise_fp16_nf4 if dtype == torch.float16 else \ + fx = cdequantize_blockwise_fp16_nf4 if dtype == torch_float16 else \ cdequantize_blockwise_bf16_nf4 fx(get_ptr(None), get_ptr(W), ptr_out_absmax, get_ptr(out), ctypes_c_int(blocksize), ctypes_c_int(out.numel()), CUDA_STREAM,) @@ -385,23 +402,28 @@ else: @torch.inference_mode def fast_dequantize(W, quant_state = None, out = None, use_global_buffer = False): if quant_state is None: return W + is_double_quantized = True if type(quant_state) is not list: # New quant_state as a class # https://github.com/TimDettmers/bitsandbytes/pull/763/files - absmax = quant_state.absmax - shape = quant_state.shape - dtype = quant_state.dtype - blocksize = quant_state.blocksize - offset = quant_state.offset - state2 = quant_state.state2 - absmax2 = state2.absmax - code2 = state2.code - blocksize2 = state2.blocksize + absmax = quant_state.absmax + shape = quant_state.shape + dtype = quant_state.dtype + blocksize = quant_state.blocksize + offset = quant_state.offset + state2 = quant_state.state2 + is_double_quantized = state2 is not None + if is_double_quantized: + absmax2 = state2.absmax + code2 = state2.code + blocksize2 = state2.blocksize else: # Old quant_state as a list of lists absmax, shape, dtype, blocksize, compressed_stats, _, _ = quant_state offset, state2 = compressed_stats - absmax2, code2, blocksize2, _, _, _, _ = state2 + is_double_quantized = state2 is not None + if is_double_quantized: + absmax2, code2, blocksize2, _, _, _, _ = state2 pass n_elements_absmax = absmax.numel() @@ -413,17 +435,20 @@ else: else: assert(out.shape == shape) assert(out.dtype == dtype) - out_absmax = torch_empty(n_elements_absmax, dtype = torch.float32, device = device, requires_grad = False) + out_absmax = torch_empty(n_elements_absmax, dtype = torch_float32, device = device, requires_grad = False) # Do dequantization - ptr_out_absmax = get_ptr(out_absmax) - cdequantize_blockwise_fp32( - get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), ptr_out_absmax, - ctypes_c_int(blocksize2), ctypes_c_int(n_elements_absmax), - ) - out_absmax += offset + if is_double_quantized: + ptr_out_absmax = get_ptr(out_absmax) + cdequantize_blockwise_fp32( + get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), ptr_out_absmax, + ctypes_c_int(blocksize2), ctypes_c_int(n_elements_absmax), + ) + out_absmax += offset + else: + ptr_out_absmax = get_ptr(absmax) - fx = cdequantize_blockwise_fp16_nf4 if dtype == torch.float16 else \ + fx = cdequantize_blockwise_fp16_nf4 if dtype == torch_float16 else \ cdequantize_blockwise_bf16_nf4 fx(get_ptr(None), get_ptr(W), ptr_out_absmax, get_ptr(out), ctypes_c_int(blocksize), ctypes_c_int(out.numel()),) @@ -443,23 +468,27 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM: # From https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/functional.py#L1469 _, q_len, hd = X.shape # assert(q_len == 1) - + is_double_quantized = True if type(quant_state) is not list: # https://github.com/TimDettmers/bitsandbytes/pull/763/files - absmax = quant_state.absmax - shape = quant_state.shape - dtype = quant_state.dtype - blocksize = quant_state.blocksize - stats = quant_state.code - offset = quant_state.offset - state2 = quant_state.state2 - absmax2 = state2.absmax - code2 = state2.code - blocksize2 = state2.blocksize + absmax = quant_state.absmax + shape = quant_state.shape + dtype = quant_state.dtype + blocksize = quant_state.blocksize + stats = quant_state.code + offset = quant_state.offset + state2 = quant_state.state2 + is_double_quantized = state2 is not None + if is_double_quantized: + absmax2 = state2.absmax + code2 = state2.code + blocksize2 = state2.blocksize else: absmax, shape, dtype, blocksize, compressed_stats, quant_type, stats = quant_state offset, state2 = compressed_stats - absmax2, code2, blocksize2, _, _, _, _ = state2 + is_double_quantized = state2 is not None + if is_double_quantized: + absmax2, code2, blocksize2, _, _, _, _ = state2 pass global XPU_STREAMS device = W.device @@ -488,17 +517,18 @@ if DEVICE_TYPE == "xpu" and HAS_XPU_STREAM: ldb = ctypes_c_int32(ldb) ldc = ctypes_c_int32(ldc) - df = torch_empty(absmax.shape, dtype = torch.float32, device = device) with torch_gpu_device(device): - cdequantize_blockwise_fp32( - get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df), - ctypes_c_int(blocksize2), ctypes_c_int(df.numel()), XPU_STREAM, - ) - df += offset - absmax = df + if is_double_quantized: + df = torch_empty(absmax.shape, dtype = torch_float32, device = device) + cdequantize_blockwise_fp32( + get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df), + ctypes_c_int(blocksize2), ctypes_c_int(df.numel()), XPU_STREAM, + ) + df += offset + absmax = df - fx = cgemm_4bit_inference_naive_fp16 if dtype == torch.float16 else \ - cgemm_4bit_inference_naive_bf16 + fx = cgemm_4bit_inference_naive_fp16 if dtype == torch_float16 else \ + cgemm_4bit_inference_naive_bf16 blocksize = ctypes_c_int32(blocksize) fx(m, n, k, get_ptr(X), get_ptr(W), get_ptr(absmax), get_ptr(stats), get_ptr(out), @@ -514,23 +544,28 @@ elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM: # From https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/functional.py#L1469 _, q_len, hd = X.shape # assert(q_len == 1) + is_double_quantized = True if type(quant_state) is not list: # https://github.com/TimDettmers/bitsandbytes/pull/763/files - absmax = quant_state.absmax - shape = quant_state.shape - dtype = quant_state.dtype - blocksize = quant_state.blocksize - stats = quant_state.code - offset = quant_state.offset - state2 = quant_state.state2 - absmax2 = state2.absmax - code2 = state2.code - blocksize2 = state2.blocksize + absmax = quant_state.absmax + shape = quant_state.shape + dtype = quant_state.dtype + blocksize = quant_state.blocksize + stats = quant_state.code + offset = quant_state.offset + state2 = quant_state.state2 + is_double_quantized = state2 is not None + if is_double_quantized: + absmax2 = state2.absmax + code2 = state2.code + blocksize2 = state2.blocksize else: absmax, shape, dtype, blocksize, compressed_stats, quant_type, stats = quant_state offset, state2 = compressed_stats - absmax2, code2, blocksize2, _, _, _, _ = state2 + is_double_quantized = state2 is not None + if is_double_quantized: + absmax2, code2, blocksize2, _, _, _, _ = state2 pass global CUDA_STREAMS device = W.device @@ -559,17 +594,18 @@ elif DEVICE_TYPE == "cuda" and HAS_CUDA_STREAM: ldb = ctypes_c_int32(ldb) ldc = ctypes_c_int32(ldc) - df = torch_empty(absmax.shape, dtype = torch.float32, device = device) with torch_gpu_device(device): - cdequantize_blockwise_fp32( - get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df), - ctypes_c_int(blocksize2), ctypes_c_int(df.numel()), CUDA_STREAM, - ) - df += offset - absmax = df + if is_double_quantized: + df = torch_empty(absmax.shape, dtype = torch_float32, device = device) + cdequantize_blockwise_fp32( + get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df), + ctypes_c_int(blocksize2), ctypes_c_int(df.numel()), CUDA_STREAM, + ) + df += offset + absmax = df - fx = cgemm_4bit_inference_naive_fp16 if dtype == torch.float16 else \ - cgemm_4bit_inference_naive_bf16 + fx = cgemm_4bit_inference_naive_fp16 if dtype == torch_float16 else \ + cgemm_4bit_inference_naive_bf16 blocksize = ctypes_c_int32(blocksize) fx(m, n, k, get_ptr(X), get_ptr(W), get_ptr(absmax), get_ptr(stats), get_ptr(out), @@ -585,6 +621,7 @@ else: # From https://github.com/TimDettmers/bitsandbytes/blob/main/bitsandbytes/functional.py#L1469 _, q_len, hd = X.shape # assert(q_len == 1) + is_double_quantized = True if type(quant_state) is not list: # https://github.com/TimDettmers/bitsandbytes/pull/763/files @@ -595,13 +632,17 @@ else: stats = quant_state.code offset = quant_state.offset state2 = quant_state.state2 - absmax2 = state2.absmax - code2 = state2.code - blocksize2 = state2.blocksize + is_double_quantized = state2 is not None + if is_double_quantized: + absmax2 = state2.absmax + code2 = state2.code + blocksize2 = state2.blocksize else: absmax, shape, dtype, blocksize, compressed_stats, quant_type, stats = quant_state offset, state2 = compressed_stats - absmax2, code2, blocksize2, _, _, _, _ = state2 + is_double_quantized = state2 is not None + if is_double_quantized: + absmax2, code2, blocksize2, _, _, _, _ = state2 pass # assert(dtype == X.dtype) bout = shape[0] @@ -626,16 +667,17 @@ else: ldb = ctypes_c_int32(ldb) ldc = ctypes_c_int32(ldc) - df = torch_empty(absmax.shape, dtype = torch.float32, device = device) - cdequantize_blockwise_fp32( - get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df), - ctypes_c_int(blocksize2), ctypes_c_int(df.numel()), - ) - df += offset - absmax = df + if is_double_quantized: + df = torch_empty(absmax.shape, dtype = torch_float32, device = device) + cdequantize_blockwise_fp32( + get_ptr(code2), get_ptr(absmax), get_ptr(absmax2), get_ptr(df), + ctypes_c_int(blocksize2), ctypes_c_int(df.numel()), + ) + df += offset + absmax = df - fx = cgemm_4bit_inference_naive_fp16 if dtype == torch.float16 else \ - cgemm_4bit_inference_naive_bf16 + fx = cgemm_4bit_inference_naive_fp16 if dtype == torch_float16 else \ + cgemm_4bit_inference_naive_bf16 blocksize = ctypes_c_int32(blocksize) fx(m, n, k, get_ptr(X), get_ptr(W), get_ptr(absmax), get_ptr(stats), get_ptr(out),