wip

Author: cli99

llm_analysis/analysis.py

@@ -228,6 +228,7 @@ def get_TFLOPS_per_gpu(self, wbits, abits) -> float:
             assert (higher_bits == 16
                     ), "weight_bits and activation_bits must be 4, 8, or 16"
             gemm_TFOPS = self.gpu_config.peak_fp16_TFLOPS
+        print('XXXXX', self.gpu_config)
         return gemm_TFOPS * self.flops_efficiency
 
     def get_pivot(self, wbits, abits) -> float:
@@ -1198,6 +1199,8 @@ def get_latency_fwd_per_layer_attn(
 
         attention_projection_flops, attention_compute_flops = (
             self.get_num_flops_fwd_per_layer_attn(batch_size, seq_len))
+        assert tp_size != 0, "tp_size must be greater than 0"
+        print('XXXXX', self.dtype_config)
         compute_latency = (
             attention_projection_flops / tp_size / (self.get_TFLOPS_per_gpu(
                 self.dtype_config.linear_weight_bits,
@@ -2535,12 +2538,13 @@ def training(
         elif activation_recomputation == ActivationRecomputation.ATTN:
             latency_recompute = num_layers_per_gpu * latency_fwd_per_layer_attn_compute
         elif activation_recomputation == ActivationRecomputation.ATTN_COMPUTE:
-            latency_recompute = (num_layers_per_gpu *
-                                 self.get_num_flops_total_attn_compute(
-                                     batch_size_per_gpu, seq_len) /
-                                 ((self.parallelism_config.tp_size *
-                                   self.parallelism_config.pp_size) *
-                                  self.get_TFLOPS_per_gpu() * 1e12))
+            latency_recompute = (
+                num_layers_per_gpu * self.get_num_flops_total_attn_compute(
+                    batch_size_per_gpu, seq_len) /
+                ((self.parallelism_config.tp_size *
+                  self.parallelism_config.pp_size) * self.get_TFLOPS_per_gpu(
+                      self.dtype_config.weight_bits,
+                      self.dtype_config.activation_bits) * 1e12))
         elif activation_recomputation == ActivationRecomputation.NONE:
             latency_recompute = 0
 

llm_analysis/config.py

@@ -193,8 +193,14 @@ class DtypeConfig:
     weight_bits: int = 16  # number of bits for weight
     activation_bits: int = 16  # number of bits for activation
     embedding_bits: int = 16  # number of bits for the embedding
-    linear_weight_bits: int = 16  # number of bits for weight in linear layer
-    linear_activation_bits: int = 16  # number of bits for activation in linear layer
+    linear_weight_bits: int | None = None  # number of bits for weight in linear layer
+    linear_activation_bits: int | None = None  # number of bits for activation in linear layer
+
+    def __post_init__(self):
+        if self.linear_weight_bits is None:
+            self.linear_weight_bits = self.weight_bits
+        if self.linear_activation_bits is None:
+            self.peak_i4_TFLOPS = self.activation_bits
 
 
 @dataclass
@@ -540,7 +546,11 @@ def list_gpu_configs() -> list:
 
 def list_dtype_configs() -> None:
     """List all predefined data type configs."""
-    logger.info(dtype_configs.keys())
+    if not dtype_configs:
+        logger.warning("No dtype configs loaded")
+        return []
+    logger.info(f"Available dtype configs: {list(dtype_configs.keys())}")
+    return list(dtype_configs.keys())
 
 
 def get_model_config_by_name(name_or_path: str) -> ModelConfig: