corrct rf8 function

rf8.py

@@ -17,4 +17,29 @@ def rf8(model, n=4):
             data1, exps1 = get_residual(param.data)
             data2, exps2 = get_residual(param.data - data1)
             flags = (exps1-exps2 <= n)
-            param.data = data1 + flags * data2
+            param.data = data1 + flags * data2
+            
+            
+def rf8_new(model):
+    """8-bit Residual Float-pointing Format"""
+    with torch.no_grad():
+        for param in model.parameters():
+            param_ = param.cpu()
+
+            signs, exps = torch.sign(param_), torch.frexp(param_)[1] - 1
+
+            bias = torch.tensor([-4, -3, -2, 1, 0], dtype=int)
+            exps_ = exps.unsqueeze(-1).expand(*exps.shape, 5)
+            Exponents = torch.exp2(exps)
+
+            res_list = torch.exp2(bias + exps_)
+            res_true = torch.abs(param_) - Exponents
+            res_true = res_true.unsqueeze(-1).expand(*res_true.shape, 5)
+
+            indices = (res_true - res_list).abs().argmin(-1).unsqueeze(-1)
+            Residuals = torch.gather(res_list, -1, indices).squeeze()
+
+            values = signs * (Exponents + Residuals)
+            values[values.abs() < 2**-12] = 0
+            values[values.abs() > 2**5] = 0
+            param.data = values.to(torch.bfloat16).to(param.device)