UI Updates

app.py

@@ -30,9 +30,7 @@ precision_options = {
 # Streamlit app
 st.title("Memory Usage Calculator for Large Language Models")
 
-# Constants from the paper, now exposed as user inputs
-layers = st.number_input("Number of Layers", value=32, step=1)
-attention_heads = st.number_input("Number of Attention Heads", value=32, step=1)
+
 
 # Taken from "Reducing Activation Recomputation in Large Transformer Models" https://arxiv.org/abs/2205.05198
 def calculate_memory_usage(parameter_count, context_length, data_type, batch_size, vocab_size, precision):
@@ -69,6 +67,8 @@ def calculate_activations(parameter_count, context_length, batch_size, vocab_siz
 
 # User inputs
 parameter_count = st.number_input("Parameter Count (in billions)", value=1, step=1) * 1e9
+layers = st.number_input("Number of Layers", value=32, step=1)
+attention_heads = st.number_input("Number of Attention Heads", value=32, step=1)
 context_length = st.number_input("Context Length (number of tokens)", value=512, step=1)
 data_type = st.selectbox("Data Type", options=list(quantization_bit_sizes.keys()))
 batch_size = st.number_input("Batch Size", value=1, step=1)