Skip to content

Commit 5a3f68f

Browse files
fix formatting
1 parent b738ab3 commit 5a3f68f

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

primus/core/projection/module_profilers/dense_mlp.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,16 +30,16 @@ def estimated_activation_memory(self, batch_size: int, seq_len: int) -> int:
3030
)
3131
# Calculate memory at different stages and take maximum
3232
input_memory = num_tokens * self.config.model_config.hidden_size * 2 # bf16
33-
33+
3434
# Memory after first projection(s)
3535
if self.config.model_config.swiglu:
3636
# Need to store both gate and up projections for backward
3737
intermediate_memory = 2 * num_tokens * self.config.model_config.ffn_hidden_size * 2 # bf16
3838
else:
3939
intermediate_memory = num_tokens * self.config.model_config.ffn_hidden_size * 2 # bf16
40-
40+
4141
output_memory = num_tokens * self.config.model_config.hidden_size * 2 # bf16
42-
42+
4343
# Peak memory is input + intermediate (both needed for backward)
4444
return input_memory + intermediate_memory + output_memory
4545

0 commit comments

Comments
 (0)