Create b158.py

b158.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+import torchao
+
+def bitnet_b158_quantize(tensor):
+    """
+    Konvertiert ein Gewicht-Tensor in das ternäre 1.58-Bit Format {-1, 0, +1}
+    inklusive der notwendigen Per-Channel Skalierung.
+    """
+    # 1. Berechne den durchschnittlichen Absolutwert pro Kanal (Zeile)
+    scale = tensor.abs().mean(dim=-1, keepdim=True).clamp(min=1e-5)
+    
+    # 2. Skaliere den Tensor und runde auf die nächste ganze Zahl
+    quantized = torch.round(tensor / scale)
+    
+    # 3. Zwinge die Werte strikt in den Bereich von -1 bis +1
+    quantized = torch.clamp(quantized, min=-1.0, max=1.0)
+    
+    return quantized, scale
+
+class BitLinear158(nn.Module):
+    """
+    Ein Ersatz für nn.Linear, der die 1.58-Bit Ternary-Inferenz ausführt.
+    """
+    def __init__(self, in_features, out_features, bias=False):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.register_buffer("weight_158", torch.zeros((out_features, in_features), dtype=torch.int8))
+        self.register_buffer("scale", torch.zeros((out_features, 1), dtype=torch.bfloat16))
+        
+    @torch.no_grad()
+    def from_float(self, float_layer):
+        # Transformiere die originalen Gewichte
+        q_w, scale = bitnet_b158_quantize(float_layer.weight.data)
+        self.weight_158.copy_(q_w.to(torch.int8))
+        self.scale.copy_(scale.to(torch.bfloat16))
+        return self
+
+    def forward(self, x):
+        # Die Magie: x wird mit den Integer-Gewichten (-1, 0, 1) verarbeitet
+        # Auf der Hardwarebene entspricht dies reinen Additionen/Subtraktionen
+        out = nn.functional.linear(x, self.weight_158.to(x.dtype))
+        return out * self.scale