Package: llm-analysis
llm-analysisLatency and Memory Analysis of Transformer Models for Training and Inferencellm-analysisOverviewExamplesQuick StartUsing theLLMAnalysisclassUsing the Entry Point Functions for Command LineHow to Set FLOPS and Memory EfficiencyCurrent Scope and LimitationsParallelism SchemeCommunicationActivation RecomputationData TypesFine-TuningAssumptions in InferenceTODOs (stay tuned :radio:)CitationContributingUseful LinksOverviewMany formulas or equations are floating around in papers, blogs, etc., about how to calculate training or inference latency and memory for Large Language Models (LLMs) or Transformers. Rather than doing math on papers or typing in Excel sheets,let's automate the boring stuff with llm-analysis:gear:!Given the specified model, GPU, data type, and parallelism configurations, llm-analysis estimates the latency and memory usage of LLMs for training or inference. With llm-analysis, one can easily try out different training/inference setups theoretically, and better understand the system performance for different scenarios.llm-analysis helps answer questions such as:what batch size, data type, parallelism scheme to use to get afeasible(not getting OOM) andoptimal(maximizing throughput with a latency constraint) setup for training or inferencetimeit takes with the given setup to do training or inference and thecost(GPU-hours)how the latency/memory changes if using a different model, GPU type, number of GPU, data type for weights and activations, parallelism configuration (suggesting the performance benefit ofmodeling change,hardware improvement,quantization,parallelism, etc.)ExamplesCheck the example use cases. With llm-analysis, you can do such analysis in minutes :rocket:!Llama 2 ModelsLLaMA ModelsMegatron-LM ModelsMegatron Turing NLG 530BFasterTransformerQuick StartTo install llm-analysis from pypi:pipinstallllm-analysisTo install the latest development build:pipinstall--upgradegit+https://github.com/cli99/llm-analysis.git@mainTo install from source, clone the repo and runpip install .orpoetry install(installpoetrybypip install poetry).Using theLLMAnalysisclassTo integrate llm-analysis in your code, use theLLMAnalysisclass. Refer to docLLMAnalysisfor details.LLMAnalysisis constructed with flops and memory efficiency numbers and the followingconfiguration classes:ModelConfigcovers model information, i.e. max sequence length, number of transformer layers, number of attention heads, hidden dimension, vocabulary sizeGPUConfigcovers GPU compute and memory specificationsDtypeConfigcovers the number of bits used for the model weight, activation, and embeddingParallelismConfigcovers Tensor Parallelism (tp), Pipeline Parallelism (pp), Sequence Parallelism (sp), Expert Parallelism (ep),and Data Parallelism (dp).ThenLLMAnalysiscan be queried with different arguments through thetrainingandinferencemethods.Using the Entry Point Functions for Command Linellm-analysis provides two entry functions,trainandinfer, for ease of use through the command line interface. Runpython-mllm_analysis.analysistrain--helporpython-mllm_analysis.analysisinfer--helpto check the options or read the linked doc. Refer to theexamplesto see how they are used.trainandinferuse the pre-defined name-to-configuration mappings (model_configs,gpu_configs,dtype_configs) and other user-input arguments to construct theLLMAnalysisand do the query.The pre-defined mappings are populated at the runtime from the model, GPU, and data type configurationjsonfiles undermodel_configs,gpu_configs, anddtype_configs. To add a new model, GPU or data type to the mapping for query, just add ajsondescription file to the corresponding folder.llm-analysis also supports retrievingModelConfigfrom a model config json file path or Hugging Face with the model name .From a local model config json file, e.g.,python -m llm_analysis.analysis train --model_name=local_example_model.json. Check the model configurations under themodel_configsfolder.From Hugging Face, e.g., useEleutherAI/gpt-neox-20basmodel_namewhen calling thetrainorinferentry functions.python -m llm_analysis.analysis train --model_name=EleutherAI/gpt-neox-20b --total_num_gpus 32 --ds_zero 3A list of handy commands is provided to query against the pre-defined mappings as well as Hugging Face, or to dump configurations. Runpython -m llm_analysis.config --helpfor details.Some examples:python-mllm_analysis.configget_model_config_by_nameEleutherAI/gpt-neox-20bgets theModelConfigfrom the populated mapping by name, if not found, llm-analysis tries to get it from HuggingFace.Note that LLaMA models need at leasttransformers-4.28.1to retrieve, either update to a latertransformerslibrary, or use the predefinedModelConfigfor LLaMA models (/in model names are replaced with_).python-mllm_analysis.configlist_gpu_configslists the names of all predefined GPU configurations, then you can query withpython-mllm_analysis.configget_gpu_config_by_namea100-sxm-80gbto show the correspondingGPUConfig.How to Set FLOPS and Memory EfficiencySetting flops and memory efficiency to1(default) gives the lower bound of training or inference latency, as it assumes the peak hardware performance (which is never the case).
A close-to-reality flops or memory efficiency can be found by benchmarking and profiling using the input dimensions in the model.If one has to make assumptions, for flops efficiency, literature reports up to0.5for large scale model training, and up to0.7for inference;0.9can be an aggressive target for memory efficiency.Current Scope and Limitationsllm-analysis aims to provide alower-boundestimation of memory usage and latency.Parallelism Schemellm-analysis currently covers Tensor Parallelism (tp), Pipeline Parallelism (pp), Sequence Parallelism (sp), Expert Parallelism (ep), and Data Parallelism (dp).tp, pp, and sp adopt the style of parallelization used inMegatron-LMfor training andFasterTransformerfor inferenceIn the training analysis, dp sharding assumes usingDeepSpeed ZeROorFSDP.ds_zerois used to specify the dp sharding strategyds_zeroDeepSpeed ZeROFSDPSharding0disabledNO_SHARDNo sharding1Stage 1N/AShard optimizer states2Stage 2SHARD_GRAD_OPShard gradients and optimizer states3Stage 3FULL_SHARDShard gradients, optimizer states, model parametersep parallelizes the number of MLP experts acrossep_sizedevices, i.e. the number of experts per GPU istotal number of experts / ep_size. Thus for the MLP module, the number of devices for other parallelization dimensions is divided byep_sizecompared to other parts of the model.Communicationtp communication is calculated as usingring allreduce. pp, dp, ep communications are ignored for now, i.e. assuming perfect computation and communication overlapping, which is not true when communication cannot overlap with compute due to dependency, or when communication is too long to hide due to slow interconnect or large data volume.Activation Recomputationllm-analysis supports both full and selective activation recomputation, as described inReducing Activation Recomputation in Large Transformer ModelsData TypesData types are expressed with the number of bits, only32(FP32, TF32),16(FP16, BF16),8(INT8), and4(INT4) bits data types are modeled for now.Fine-TuningFine-tuning is modeled the same (controlled bytotal_num_tokenspassed to thetrainentry function) as pre-training, thus assuming full (all model parameters) fine-tuning. Parameter-efficient fine-tuning (PEFT) is
in future support.Assumptions in InferenceInference assumes perfect overlapping of compute and memory operations when calculating latency, and maximum memory reuse when calculating memory usage.TODOs(stay tuned :radio:)Check the TODOs below for what's next and stay tuned :radio:! Any contributions or feedback are highly welcome!Add dp (across and within a node), ep (within a node), pp (across nodes) communication analysisSupport efficient fine-tuning methods such asLoRAorAdaptersAdd FP8 datatype supportSupport CPU offloading (weight, KV cache, etc.) analysis in training and inferenceSupport other hardware (e.g. CPU) for inference analysisCitationIf you use llm-analysis in your work, please cite:Cheng Li. (2023). LLM-Analysis: Latency and Memory Analysis of Transformer Models for Training and Inference. GitHub repository, https://github.com/cli99/llm-analysis.or@misc{llm-analysis-chengli,
  author = {Cheng Li},
  title = {LLM-Analysis: Latency and Memory Analysis of Transformer Models for Training and Inference},
  year = {2023},
  publisher = {GitHub},
  journal = {GitHub repository},
  howpublished = {\url{https://github.com/cli99/llm-analysis}},
}ContributingContributions and suggestions are welcome.llm-analysis usespre-committo ensure code formatting is consistent. For pull requests with code contribution, please install the pre-commit (pip install pre-commit) as well as the used hooks (pip installin the repo), and format the code (runs automatically before each git commit) before submitting the PR.Useful LinksMegatron-LM: Training Multi-Billion Parameter Language Models Using Model ParallelismZeRO: Memory Optimizations Toward Training Trillion Parameter ModelsEfficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LMUsing DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language ModelReducing Activation Recomputation in Large Transformer ModelsTraining Compute-Optimal Large Language ModelsEfficiently Scaling Transformer InferenceTraining Compute-Optimal Large Language ModelsUnderstanding INT4 Quantization for Transformer Models: Latency Speedup, Composability, and Failure CasesA Comprehensive Study on Post-Training Quantization for Large Language Models