feat: extend the num head terms for fp8 calibration

Author: niushengxiao

lightllm/common/offline_fp8_quant_mem_manager.py

@@ -25,7 +25,6 @@ def __init__(
 
         self.qmax = torch.finfo(torch.float8_e4m3fn).max
         self.qmin = torch.finfo(torch.float8_e4m3fn).min
-        self.layer_num = layer_num
         self.total_head_num = head_num * dist.get_world_size() if dist.is_initialized() else head_num
         self.count = 0
         self.scales = None
@@ -45,7 +44,13 @@ def __init__(
             self.scales_list = cfg["scales"]
             self.scales = torch.tensor(self.scales_list, dtype=torch.float32, device="cuda").view(cfg["scales_shape"])
             if not get_env_start_args().enable_fa3:
-                self.scales = torch.repeat_interleave(self.scales, self.head_num, dim=-1)
+                self.scales = torch.repeat_interleave(self.scales, head_num, dim=-1)
+            elif cfg["num_head"] > self.total_head_num:
+                factor = cfg["num_head"] // self.total_head_num
+                self.scales = self.scales[..., ::factor].contiguous()
+            elif cfg["num_head"] < self.total_head_num:
+                factor = self.total_head_num // cfg["num_head"]
+                self.scales = torch.repeat_interleave(self.scales, factor, dim=-1).contiguous()
             if get_env_start_args().enable_fa3 and dist.is_initialized() and dist.get_world_size() > 1:
                 half_head = self.total_head_num // 2
                 start_head = dist.get_rank() * head_num
@@ -77,7 +82,7 @@ def _load_and_check_config(self):
                 raise ValueError(
                     f"num_layers {cfg['num_layers']} in config " f"not match current layer_num {self.layer_num}"
                 )
-            if cfg["num_head"] != self.total_head_num:
+            if cfg["num_head"] % self.total_head_num != 0 and self.total_head_num % cfg["num_head"] != 0:
                 raise ValueError(
                     f"num_head {cfg['num_head']} in config " f"not match current model head num {self.total_head_num}"
                 )

lightllm/server/api_cli.py

@@ -175,7 +175,6 @@ def make_argument_parser() -> argparse.ArgumentParser:
                         export_fp8kv_calibration record and export kv cache quant calibration results to a json file.
                         It can be used for llama and qwen model.
                         Calibration need to disable cudagraph and use fa3 or flashinfer backend.
-                        Tp size must no more than head num when calibration.
                         ppl_int8kv mode use int8 to store kv cache, and use ppl fast kernel;
                         ppl_fp16 mode use ppl fast fp16 decode attention kernel;
                         you need to read source code to make sure the supported detail mode for all models""",