Merge branch 'refs/heads/202405-cached-states' into dev

Author: turboderp

conversion/measure.py

@@ -125,18 +125,18 @@ def test_quant(source: ExLlamaV2Linear,
 
 def test_error(module, hidden_states, target_states, cache, attn_params):
 
-    rfn_sum = 0
+    rfn_sum = torch.tensor(0.0).cuda()
     rfn_count = 0
     for x, xref in zip(hidden_states, target_states):
         x = x.cuda()
         xref = xref.cuda()
         xtest = module.forward(x, cache, attn_params)
         xtest = xtest[0].float()
         xref = xref[0].float()
-        rfn_sum += (torch.linalg.norm(xtest - xref, 'fro') / torch.linalg.norm(xref, 'fro')).item()
+        rfn_sum += torch.linalg.norm(xtest - xref, 'fro') / torch.linalg.norm(xref, 'fro')
         rfn_count += 1
 
-    return max(1e-6, 1 - (rfn_sum / rfn_count))
+    return max(1e-6, 1 - (rfn_sum.item() / rfn_count))
 
 
 def measure_attn(module, hidden_states, target_states, quantizers, cache, attn_params, keep_q = False):
@@ -382,7 +382,7 @@ def print_status_box(*content_lines):
     print('-' * box_width)
 
 @torch.inference_mode()
-def measure_quant(job, save_fn, model):
+def measure_quant(job, save_fn, model, hidden_state_offload_layers):
 
     # vars for status box
     time_spent_list = []  
@@ -418,8 +418,9 @@ def measure_quant(job, save_fn, model):
 
     hidden_states = []
     with safe_open(states_filename, framework = "pt", device = "cpu") as f:
-        for k in sorted(f.keys()):
-            hidden_states.append(f.get_tensor(k))
+        for i, k in enumerate(sorted(f.keys())):
+            t = f.get_tensor(k)
+            hidden_states.append(t.to("cuda:0") if i < hidden_state_offload_layers else t)
 
     index = job["last_module_idx"]
     while True:
@@ -515,18 +516,19 @@ def measure_quant(job, save_fn, model):
 
             x = hidden_states[i].to("cuda:0")
             outputs = module.forward(x, cache, attn_params, intermediates = True)
+            target_device = "cuda:0" if i < hidden_state_offload_layers else "cpu"
 
             # Hessians
 
             if mode == "self_attn":
                 quantizers["q_proj"].add_batch(outputs["post_norm"])  # Reuse H for K and V
                 quantizers["o_proj"].add_batch(outputs["attn_output"])
-                target_states.append(outputs["hidden_states"].to("cpu"))
+                target_states.append(outputs["hidden_states"].to(target_device))
 
             if mode == "mlp":
                 quantizers["up_proj"].add_batch(outputs["post_norm"])  # Reuse H for gate_proj
                 quantizers["down_proj"].add_batch(outputs["pre_down"])
-                target_states.append(outputs["hidden_states"].to("cpu"))
+                target_states.append(outputs["hidden_states"].to(target_device))
 
             if mode == "block_sparse_moe":
                 for j in range(model.config.num_experts):
@@ -537,19 +539,19 @@ def measure_quant(job, save_fn, model):
                             uncalibrated_experts[j] += 1
                     else:
                         uncalibrated_experts[j] += 1
-                target_states.append(outputs["hidden_states"].to("cpu"))
+                target_states.append(outputs["hidden_states"].to(target_device))
 
             if mode == "parallel_decoder":
                 quantizers["q_proj"].add_batch(outputs["post_norm"])  # Reuse H for K, V, up_proj and gate_proj
                 quantizers["o_proj"].add_batch(outputs["attn_output"])
                 quantizers["down_proj"].add_batch(outputs["pre_down"])
                 hidden_states[i] = outputs["post_norm"]
-                target_states_attn.append(outputs["hidden_states_attn"].to("cpu"))
-                target_states_mlp.append(outputs["hidden_states_mlp"].to("cpu"))
-                target_states.append(outputs["hidden_states"].to("cpu"))
+                target_states_attn.append(outputs["hidden_states_attn"].to(target_device))
+                target_states_mlp.append(outputs["hidden_states_mlp"].to(target_device))
+                target_states.append(outputs["hidden_states"].to(target_device))
 
             if mode == "pos_emb":
-                target_states.append(outputs["hidden_states"].to("cpu"))
+                target_states.append(outputs["hidden_states"].to(target_device))
 
         # For MoE layers, warn if any layer received less than 10% of a calibration batch
 

conversion/quantize.py

@@ -439,7 +439,7 @@ def quant(job, save_fn, model):
                 cal_ids = f.get_tensor("input_ids")
             module.linear.weight.data = module.linear.weight.data.to("cuda:0")
 
-        rfn_sum = 0
+        rfn_sum = torch.tensor(0.0).cuda()
         rfn_count = 0
         logprob_sum = 0.0
         logprob_count = 0
@@ -458,7 +458,7 @@ def quant(job, save_fn, model):
                 output_ref = target_states[i].to("cuda:0")
                 output_ref = output_ref[0].float()
 
-                rfn_sum += (torch.linalg.norm(output - output_ref, 'fro') / torch.linalg.norm(output_ref, 'fro')).item()
+                rfn_sum += torch.linalg.norm(output - output_ref, 'fro') / torch.linalg.norm(output_ref, 'fro')
                 rfn_count += 1
 
                 output_ref = None
@@ -485,7 +485,7 @@ def quant(job, save_fn, model):
 
         if mode != "linear":
 
-            err = rfn_sum / rfn_count
+            err = rfn_sum.item() / rfn_count
             print(f" -- Module quantized, rfn_error: {err:1.6f}")
 
         else:

convert.py

@@ -29,6 +29,7 @@
 parser.add_argument("-l", "--length", type = int, default = 2048, help = "Max no. tokens per sample")
 parser.add_argument("-ml", "--measurement_length", type = int, default = 2048, help = "Max no. tokens per sample when measuring")
 parser.add_argument("-so", "--status_output", action = "store_true", help = "Include machine-parseable status updates in console output")
+parser.add_argument("-hsol", "--hidden_state_offload_layers", type = int, default = 0, help = "Number of hidden/target states to keep in VRAM. Speed-up but increases VRAM usage")
 
 args = parser.parse_args()
 
@@ -242,7 +243,7 @@ def save_job():
         model = ExLlamaV2(config)
         model.load(lazy = True)
 
-        status = measure_quant(job, save_job, model)  # capturing the graceful exits
+        status = measure_quant(job, save_job, model, args.hidden_state_offload_layers)  # capturing the graceful exits
         if status == "interrupted":
             print("Process interrupted. Exiting gracefully.")
             save_job()