update bwd prefetch memory calc

cli99 · May 22, 2024 · fd6ad18 · fd6ad18
1 parent 2b92db3
commit fd6ad18
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/llm_analysis/analysis.py b/llm_analysis/analysis.py
@@ -2067,7 +2067,7 @@ def training(
         estimated_fwd_prefetch_memory_per_gpu = unsharded_weight_memory_embedding + unsharded_weight_memory_per_layer
 
         estimated_bwd_prefetch_memory_per_gpu = (
-            3 + int(fwd_prefetch) +
+            int(fwd_prefetch) +
             int(bwd_prefetch)) * (unsharded_weight_memory_per_layer)
 
         estimated_prefetch_memory_per_gpu = max(