convert.py: Experimental args bitrate tweaking

Author: turboderp

exllamav3/conversion/convert_model.py

@@ -40,6 +40,7 @@
 parser.add_argument("-img", "--image_dump", action = "store_true", help = "Save model tensors as images (saved to working directory)")
 parser.add_argument("-mcg", "--mcg_multiplier", type = str, default = None, help = "MCG multiplier - EXPERIMENTAL, DO NOT USE")
 parser.add_argument("-mul1", "--mul1_multiplier", type = str, default = None, help = "MUL1 multiplier - EXPERIMENTAL, DO NOT USE")
+parser.add_argument("-strat", "--strategy", type = str, default = None, help = "Modifiers for quantization strategy - EXPERIMENTAL")
 
 group = parser.add_mutually_exclusive_group()
 group.add_argument("--out_scales", dest = "out_scales_", action = "store_true", help = "Always enable out channel scales  (for debug purposes)")
@@ -154,6 +155,7 @@ def override(arg, can_override, default):
         ("device_ratios", True, None),
         ("mcg_multiplier", True, ""),
         ("mul1_multiplier", True, ""),
+        ("strategy", False, ""),
     ]:
         override(arg_, can_override if not args.override_anyway else True, default)
 
@@ -233,6 +235,32 @@ def get_state_error(x, ref):
      return err.item(), cos, sq
 
 
+def mod_strategy(args, module, strategy, idx):
+    mod_arg = args.get("strategy")
+    if not mod_arg:
+        return strategy
+
+    s_layers = [""] + mod_arg.split(";")
+    if idx >= len(s_layers):
+        return strategy
+
+    s = s_layers[idx]
+    mod = {}
+    while s:
+        l, m = s[0], s[1]
+        s = s[2:]
+        mod[l] = int(m)
+
+    new_strategy = {}
+    for key, bits in strategy.items():
+        submodule = module.find_module(key)
+        modifier = mod.get(submodule.qbits_mod_key, 0)
+        new_strategy[key] = min(bits + modifier, 8)
+
+    # TODO: Automate this, also calculate overall increase in bitrate, track in job.json across resumes
+    return new_strategy
+
+
 @torch.inference_mode()
 def main(args, job_state):
 
@@ -281,6 +309,7 @@ def main(args, job_state):
             },
             job_state["surplus_bits"],
         )
+        strategy = mod_strategy(args, module, strategy, idx)
         job_state["surplus_bits"] = surplus
 
         # Slice module if necessary

exllamav3/modules/attn.py

@@ -174,10 +174,10 @@ def __init__(
         else:
             fkey, frange_q, frange_k, frange_v = None, None, None, None
 
-        self.q_proj = Linear(config, f"{key}.{key_q}", hidden_size, num_q_heads * head_dim, qmap = qmap + ".input", fkey = fkey, frange = frange_q)
-        self.k_proj = Linear(config, f"{key}.{key_k}", hidden_size, num_kv_heads * head_dim, qmap =  qmap + ".input", fkey = fkey, frange = frange_k)
-        self.v_proj = Linear(config, f"{key}.{key_v}", hidden_size, num_kv_heads * head_dim, qmap =  qmap + ".input", fkey = fkey, frange = frange_v)
-        self.o_proj = Linear(config, f"{key}.{key_o}", num_q_heads * head_dim, hidden_size, qmap =  qmap + ".o", out_dtype = out_dtype)
+        self.q_proj = Linear(config, f"{key}.{key_q}", hidden_size, num_q_heads * head_dim, qmap = qmap + ".input", fkey = fkey, frange = frange_q, qbits_mod_key = "q")
+        self.k_proj = Linear(config, f"{key}.{key_k}", hidden_size, num_kv_heads * head_dim, qmap =  qmap + ".input", fkey = fkey, frange = frange_k, qbits_mod_key = "k")
+        self.v_proj = Linear(config, f"{key}.{key_v}", hidden_size, num_kv_heads * head_dim, qmap =  qmap + ".input", fkey = fkey, frange = frange_v, qbits_mod_key = "v")
+        self.o_proj = Linear(config, f"{key}.{key_o}", num_q_heads * head_dim, hidden_size, qmap =  qmap + ".o", out_dtype = out_dtype, qbits_mod_key = "o")
 
         self.register_submodule(self.q_proj)
         self.register_submodule(self.k_proj)

exllamav3/modules/linear.py

@@ -24,6 +24,7 @@ def __init__(
         qmap: str | None = None,
         alt_key: str | None = None,
         qbits_key: str = "bits",
+        qbits_mod_key: str = "",
         fkey : str | None = None,
         frange: tuple[int, int] | None = None,
         caps: dict = None,
@@ -50,6 +51,7 @@ def __init__(
         self.first_out_feature = first_out_feature if first_out_feature is not None else 0
         self.inner = None
         self.qbits_key = qbits_key
+        self.qbits_mod_key = qbits_mod_key
         self.fkey = fkey
         self.frange = frange
         self.quant_type = None

exllamav3/modules/mlp.py

@@ -28,8 +28,8 @@ def __init__(
 
         self.out_dtype = out_dtype
 
-        self.up = Linear(config, f"{key}.{key_up}", hidden_size, intermediate_size, qmap = qmap + ".up")
-        self.down = Linear(config, f"{key}.{key_down}", intermediate_size, hidden_size, qmap = qmap + ".down")
+        self.up = Linear(config, f"{key}.{key_up}", hidden_size, intermediate_size, qmap = qmap + ".up", qbits_mod_key = "u")
+        self.down = Linear(config, f"{key}.{key_down}", intermediate_size, hidden_size, qmap = qmap + ".down", qbits_mod_key = "d")
 
         self.register_submodule(self.up)
         self.register_submodule(self.down)
@@ -129,7 +129,8 @@ def __init__(
                 fkey = fkey,
                 frange = frange_gate,
                 alt_key = a_key_g,
-                out_dtype = self.interm_dtype
+                out_dtype = self.interm_dtype,
+                qbits_mod_key = "g"
             )
             up = Linear(
                 config = config,
@@ -144,7 +145,8 @@ def __init__(
                 fkey = fkey,
                 frange = frange_up,
                 alt_key = a_key_u,
-                out_dtype = self.interm_dtype
+                out_dtype = self.interm_dtype,
+                qbits_mod_key = "u"
             )
             down = Linear(
                 config = config,
@@ -159,6 +161,7 @@ def __init__(
                 alt_key = a_key_d,
                 out_dtype = self.out_dtype,
                 allow_input_padding = True,
+                qbits_mod_key = "d"
             )
 
             self.ups.append(up)

exllamav3/modules/module.py

@@ -36,12 +36,18 @@ def __init__(
         self.caps = {}
         self.qmap = qmap
         self.num_slices = 1
+        self.qbits_mod_key = ""
 
     def __iter__(self):
         yield self
         for module in self.modules:
             yield from module
 
+    def find_module(self, key: str):
+        for module in self:
+            if module.key == key:
+                return module
+
     def can_defer_load(self):
         if len(self.modules) == 0: return True
         return all(module.can_defer_load() for module in self.modules)