Enabled FP8 models for replicate_kv_heads script (#353)

Signed-off-by: Onkar Chougule <[email protected]>

Author: ochougul

scripts/replicate_kv_head/replicate_kv_heads.py

@@ -14,6 +14,7 @@
 from QEfficient.transformers.quantizers.auto import replace_transformers_quantizers, undo_transformers_quantizers
 from QEfficient.transformers.quantizers.awq import WQLinear_GEMM
 from QEfficient.transformers.quantizers.gptq import QuantLinearGPTQ
+from QEfficient.transformers.quantizers.quantizer_compressed_tensors import FP8DeQuantLinear
 
 
 def duplicate_weights_for_linear_layer(
@@ -49,6 +50,15 @@ def duplicate_weights_for_linear_layer(
             1,
         ).view(hidden_size // layer.group_size, new_kv_heads * head_dim)
         layer.out_features = layer.out_features * repeat
+
+    elif isinstance(layer, FP8DeQuantLinear):
+        layer.weight.data = torch.repeat_interleave(
+            layer.weight.data.view(orig_kv_heads, head_dim, hidden_size), repeat, 0
+        ).view(new_kv_heads * head_dim, hidden_size)
+        layer.weight_scale.data = torch.repeat_interleave(
+            layer.weight_scale.data.view(orig_kv_heads, head_dim), repeat, 0
+        ).view(new_kv_heads * head_dim, -1)
+
     else:
         layer.weight.data = torch.repeat_interleave(
             layer.weight.data.view(orig_kv_heads, head_dim, hidden_size), repeat, 0
@@ -65,7 +75,6 @@ def main(args):
     model_kwargs = {"attn_implementation": "eager"}
     if args.num_hidden_layers:
         model_kwargs["num_hidden_layers"] = args.num_hidden_layers
-
     model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
 
     # Undo the effect of replace_transformers_quantizers