fix: compatibility with transformers >= 5 and support non-llama models in chat plugin

.pre-commit-config.yaml

@@ -0,0 +1,6 @@
+repos:
+  - repo: https://github.com/PyCQA/flake8
+    rev: 7.1.2
+    hooks:
+      - id: flake8
+        args: [--count, --select, "E9,F63,F7,F82", --show-source, --statistics]

qllm/modeling/base.py

@@ -21,6 +21,23 @@
 logger = utils.logger.get_logger()
 
 
+@contextlib.contextmanager
+def _no_init_weights():
+    """Replace transformers.modeling_utils.no_init_weights removed in transformers >= 5."""
+    import torch.nn.init as init
+    _orig = {}
+    skip = lambda *a, **kw: None
+    for name in list(vars(init)):
+        if name.endswith('_') and callable(getattr(init, name)):
+            _orig[name] = getattr(init, name)
+            setattr(init, name, skip)
+    try:
+        yield
+    finally:
+        for name, fn in _orig.items():
+            setattr(init, name, fn)
+
+
 @contextlib.contextmanager
 def replace_default_dtype(dtype):
     old_dtype = torch.get_default_dtype()
@@ -244,7 +261,7 @@ def from_quantized(
             model_name_or_path, trust_remote_code=trust_remote_code)
         torch_dtype = torch_dtype if torch_dtype is not None else auto_conf.torch_dtype
         init_contexts = [
-            transformers.modeling_utils.no_init_weights(),
+            transformers.modeling_utils.no_init_weights() if hasattr(transformers.modeling_utils, 'no_init_weights') else _no_init_weights(),
             # no_init_weights(),
             replace_default_dtype(torch_dtype),
             # accelerate.init_empty_weights(include_buffers=False)

qllm/plugin/chatcli/inference.py

@@ -31,8 +31,8 @@ def chat_loop(
     if _fastchat_available:
         return chat_loop_v2(model, tokenizer)
     model_type = str(type(model)).lower()
-    #if "llama" not in model_type and hasattr(tokenizer, 'apply_chat_template'):
-    #    return chat_loop_v3(model, tokenizer)
+    if "llama" not in model_type and hasattr(tokenizer, 'apply_chat_template'):
+        return chat_loop_v3(model, tokenizer)
     assert "llama" in model_type, 'have you installed fschat? please run `pip install fschat` and try again.'
     assert generate_stream_func is not None or generate_func is not None, 'should set generate function.'
 

qllm/quantization/quant_frame_base.py

@@ -81,6 +81,14 @@ def __init__(self, module):
                 super().__init__()
                 self.module = module
 
+            def __getattr__(self, name):
+                if name == 'module':
+                    return super().__getattr__(name)
+                try:
+                    return super().__getattr__(name)
+                except AttributeError:
+                    return getattr(self.module, name)
+
             def forward(self, inp, **kwargs):
                 inps.append(inp.to(swap_device))
                 layer_input_args.update(kwargs)

qllm/quantization/vptq/hessian_collector.py

@@ -98,7 +98,7 @@ def register_H_hook(module, device):
     ct = 0
 
     def H_hook(module, x):
-        nonlocal H, mu, ct, n
+        nonlocal ct
         x = x[0].reshape(-1, n).to(torch.float64)
         mu.add_(x.sum(dim=0))
         H.addmm_(x.T, x)
@@ -107,7 +107,6 @@ def H_hook(module, x):
     hook = module.register_forward_pre_hook(H_hook)
 
     def done():
-        nonlocal H, mu, ct, hook
         hook.remove()
         return H.cpu(), mu.cpu(), ct
 

qllm/quantization/vptq/qllm_hessian.py

@@ -73,7 +73,7 @@ def register_H_hook(module, device):
     ct = 0
 
     def H_hook(module, x):
-        nonlocal H, mu, ct, n
+        nonlocal ct
         x = x[0].reshape(-1, n).to(torch.float64)
         mu.add_(x.sum(dim=0))
         H.addmm_(x.T, x)
@@ -82,7 +82,6 @@ def H_hook(module, x):
     hook = module.register_forward_pre_hook(H_hook)
 
     def done():
-        nonlocal H, mu, ct, hook
         hook.remove()
         return H.cpu(), mu.cpu(), ct
 

qllm/quantization/vptq/quant_vptq.py

@@ -38,7 +38,6 @@ def forward(self, *args, **kwargs):
                 raise ValueError
         def get_func(name, out_fetures):
             def fake_forward(hidden_state, *args, **kwargs):
-                nonlocal level_map
                 if len(level_map) == 0:
                     hidden_state *= 0
                 input_shape = hidden_state.shape

scripts/lint.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+set -e
+
+echo "=== Running flake8 (syntax errors & undefined names) ==="
+flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+
+echo "=== Running flake8 (warnings) ==="
+flake8 ./qllm/modeling/q_layers --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+flake8 ./qllm/quantization --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+flake8 ./qllm/utils --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+
+echo "All lint checks passed!"