removed training calculation functionality

app.py

@@ -23,11 +23,11 @@ quantization_bit_sizes = {
 # Define precision options
 precision_options = {
     'full': 4,
-    'mixed': 6,  # for training mixed precision
+    'mixed': 6,
     'half': 2
 }
 
-def calculate_memory_usage(parameter_count, context_length, data_type, is_training, batch_size, vocab_size, precision):
+def calculate_memory_usage(parameter_count, context_length, data_type, batch_size, vocab_size, precision):
     # Convert bit size to byte size
     byte_size = quantization_bit_sizes[data_type] / 8
 
@@ -35,10 +35,10 @@ def calculate_memory_usage(parameter_count, context_length, data_type, is_traini
     memory_params = parameter_count * byte_size
 
     # Memory usage for context (activations)
-    activations = calculate_activations(parameter_count, context_length, batch_size, vocab_size, precision, is_training)
+    activations = calculate_activations(parameter_count, context_length, batch_size, vocab_size, precision)
 
     # Outputs memory usage
-    outputs = 4 * batch_size * context_length * vocab_size * (2 if is_training else 1)
+    outputs = 4 * batch_size * context_length * vocab_size
 
     # Total memory usage
     total_memory_usage = memory_params + activations + outputs
@@ -48,7 +48,7 @@ def calculate_memory_usage(parameter_count, context_length, data_type, is_traini
 
     return total_memory_usage_gb
 
-def calculate_activations(parameter_count, context_length, batch_size, vocab_size, precision, is_training):
+def calculate_activations(parameter_count, context_length, batch_size, vocab_size, precision):
     # Simplified activation calculation
     hidden_size = parameter_count ** 0.5  # assuming a square root relationship for hidden size
     num_attention_heads = 16  # a typical number of attention heads
@@ -73,7 +73,7 @@ def calculate_activations(parameter_count, context_length, batch_size, vocab_siz
 
     layer = attention_block + mlp_block + layer_norms
 
-    activations = layer * (12 if is_training else 1)  # assuming 12 layers for simplicity
+    activations = layer  # assuming 12 layers for simplicity
 
     return activations
 
@@ -84,12 +84,11 @@ st.title("Memory Usage Calculator for Large Language Models")
 parameter_count = st.number_input("Parameter Count (in billions)", value=1, step=1) * 1e9
 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()))
-is_training = st.checkbox("Training Mode", value=False)
 batch_size = st.number_input("Batch Size", value=1, step=1)
 vocab_size = st.number_input("Vocabulary Size", value=30000, step=1000)
 precision = st.selectbox("Precision", options=list(precision_options.keys()))
 
 # Calculate memory usage
 if st.button("Calculate Memory Usage"):
-    memory_usage = calculate_memory_usage(parameter_count, context_length, data_type, is_training, batch_size, vocab_size, precision)
-    st.write(f"Estimated Memory Usage for {'Training' if is_training else 'Inference'}: {memory_usage:.2f} GB")
+    memory_usage = calculate_memory_usage(parameter_count, context_length, data_type, batch_size, vocab_size, precision)
+    st.write(f"Estimated Memory Usage for Inference: {memory_usage:.2f} GB")