From b5cf274da35d6f9fd5b62edcd2c8345f5c54ee2c Mon Sep 17 00:00:00 2001
From: nimlgen <138685161+nimlgen@users.noreply.github.com>
Date: Wed, 30 Aug 2023 23:32:30 +0300
Subject: [PATCH] remove memory peak for quantized llama (#1720)

---
 examples/llama.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/examples/llama.py b/examples/llama.py
index fefe248ed5..4139a91e88 100755
--- a/examples/llama.py
+++ b/examples/llama.py
@@ -262,8 +262,8 @@ class AbsmaxQuantizedLinear:
       if 'feed_forward' in name or ('attention.w') in name or name == 'output.weight':
         scale = v.abs().max(axis=1) / 127.0
         int8_weight = (v.T/scale).T.cast(dtype=dtypes.int8)
-        new_tensors[name] = int8_weight.realize()
-        new_tensors[name.replace('weight', 'scale')] = scale.realize()
+        new_tensors[name] = int8_weight
+        new_tensors[name.replace('weight', 'scale')] = scale
       else:
         new_tensors[name] = v
     return new_tensors
@@ -287,6 +287,7 @@ class LLaMa:
 
     if quantize:
       weights = AbsmaxQuantizedLinear.quantize(weights)
+      for _,v in weights.items(): v.realize()
     load_state_dict(model, weights, strict=False)
 
     return LLaMa(model, sp_model)