Spaces:
Sleeping
Sleeping
File size: 3,953 Bytes
3301526 6bf4672 3301526 38f4f24 6bf4672 38f4f24 6cf8a8a 38f4f24 166583b 38f4f24 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import streamlit as st
from utils import memory_moe_mlp, memory_mlp_layer, memory_for_attention_layer
def main():
st.title("LLM Model Memory Usage Calculator")
st.sidebar.header("Model Parameters")
precession = st.sidebar.number_input("precession in Byte", min_value=1, max_value=4, value=2, step=2)
hidden_size = st.sidebar.number_input("Hidden Size", min_value=512, max_value=2 ** 16, value=4096, step=512)
num_heads = st.sidebar.number_input("Number of Attention Heads", min_value=4, max_value=128, value=32, step=4)
batch_size = st.sidebar.number_input("Batch Size", min_value=1, max_value=1024, value=64, step=4)
seq_len = st.sidebar.number_input("Sequence Length", min_value=512, max_value=128000, value=2048, step=512)
intermediate_size = st.sidebar.number_input("Intermediate Size", min_value=1024, max_value=2 ** 18, value=11008,
step=128)
layers = st.sidebar.number_input("Number of Layers", min_value=6, max_value=48, value=30, step=1)
moe = st.sidebar.checkbox("Use Mixture of Experts (MOE)", value=False)
# Conditional rendering for MOE parameters
if moe:
top_k = st.sidebar.number_input("Number of Experts to use (Top K)", min_value=1, max_value=16, value=2, step=1)
num_experts = st.sidebar.number_input("Total Number of Experts", min_value=2, max_value=32, value=4, step=2)
else:
top_k = 2 # Default values if MOE is not used
num_experts = 4
attention_memory = memory_for_attention_layer(precession,
seq_len,
batch_size,
hidden_size,
num_heads)
dense_mlp_memory = memory_mlp_layer(precession,
seq_len,
batch_size,
hidden_size,
intermediate_size)
dense_model_memory = layers * (attention_memory + dense_mlp_memory) // (1024 ** 3)
st.write("This estimation is for the training phase, all computation done for AdamW optimizer.")
space = st.empty()
space.markdown('<div style="height: 20px;"></div>', unsafe_allow_html=True)
st.markdown(
f'<div style="background-color: #b3f0ff; padding: 30px; border-radius: 5px;">'
f'<p style="font-weight: bold;">The memory requirement for this model is ~ {dense_model_memory} GB</p>'
f'</div>',
unsafe_allow_html=True
)
space = st.empty()
space.markdown('<div style="height: 40px;"></div>', unsafe_allow_html=True)
if moe:
moe_memory = memory_moe_mlp(precession,
seq_len,
batch_size,
hidden_size,
intermediate_size,
num_experts,
top_k)
moe_model = layers * (attention_memory + moe_memory) // (1024 ** 3)
st.markdown(
f'<div style="background-color: #99ff99; padding: 30px; border-radius: 5px;">'
f'<p style="font-weight: bold;">The memory requirement for the MOE model is ~ {moe_model} GB</p>'
f'</div>',
unsafe_allow_html=True
)
space = st.empty()
space.markdown('<div style="height: 40px;"></div>', unsafe_allow_html=True)
st.markdown(
f'<div style="background-color: #f0f0f0; padding: 30px; border-radius: 5px;">'
f'<p style="font-weight: bold;">For more information please read this article</p>'
f'<a href="https://medium.com/@khalil.hennara.247/llm-memory-usage-f62a007a509c">Article</a>'
f'</div>',
unsafe_allow_html=True
)
if __name__ == "__main__":
main()
|