Update app.py

app.py

@@ -127,7 +127,9 @@ c1, c2 = st.columns([2, 4])
 num_params = 12*l*d*d
 c1.write("Num Parameters (in B)")
 c2.write(str(round(num_params/1e9, 3)))
-c1.write("Storing activations")
+c1.write("Stored Parameters (GB)")
+c2.write(str(round(2*num_params/1e9, 3)))
+c1.write("Cached keys and values (GB)")
 acts = round(2*bs*l*(d/h)*h*2*n/1e9, 2)
 c2.write(str(acts))
 
@@ -136,12 +138,13 @@ c1, c2 = st.columns([2, 4])
 num_params = (10+2/h)*l*d*d
 c1.write("Num Parameters (in B)")
 c2.write(str(round(num_params/1e9, 3)))
-c1.write("Storing activations")
+c1.write("Stored Parameters (GB)")
+c2.write(str(round(2*num_params/1e9, 3)))
+c1.write("Cached keys and values (GB)")
 acts = round(2*bs*l*(d/h)*2*n/1e9, 2)
 c2.write(str(acts))
 
 
-
 breakdown = st.checkbox("Show breakdown per operation")
 if breakdown:
   st.header('Attention layer')