Spaces:
Sleeping
Sleeping
| 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() | |