Modify the bug in weight memory calculation (#23)

cli99 · Mar 1, 2024 · 6bacd4f · 6bacd4f
1 parent 9932ff4
commit 6bacd4f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/llm_analysis/analysis.py b/llm_analysis/analysis.py
@@ -1582,7 +1582,7 @@ def inference(
 
         weight_memory_embedding_per_gpu = self.get_memory_embedding(ds_zero)
         weight_memory_layers_per_gpu, weight_memory_attn_per_gpu, weight_memory_mlp_per_gpu, weight_memory_layernorm_per_gpu = [
-            x * self.model_config.num_layers
+            x * num_layers_per_gpu
             for x in self.get_weight_memory_per_layer(ds_zero,
                                                       return_breakdown=True)
         ]