time cost on lm_head in the implementation of huggingface

[jos-fisher-jejgkej]

Hello, thank you for your outstanding work, it has greatly inspired me. Also, I would like to ask you, what is the percentage of time cost on lm_head in the implementation of huggingface, and what is the relationship between time cost and vocab_size?

[Achazwl]

I don't remember it clearly, on Hugging Face, the proportion of the lm_head in the verification time is about 2%, and the proportion in the draft time is about 40%. The proportion on Hugging Face will be lower than that on efficient implementations such as our CUDA implementation or vLLM, SGLang, etc. On these efficient implementations, generation with batch=1 is memory-bound, while Hugging Face's implementation introduces significant other overhead (e.g. slow rmsnorm).

From a memory-bound perspective, the number of parameters in the lm_head is vocab_size * hidden_size. Take Llama-3-8B as an example, it has approximately 128k * 4k = 0.5B parameters, while a single layer of the transformer would have approximately 8B / 32layer = 0.25B parameters. So ideally, an lm_head has a time cost of two layer. (But a transformer layer also has other overhead such as attention computation, so would be slightly less than two.) While Eagle requires draft depths of 6, with an lm_head forward per Eagle forward, so the lm_head would ideally cost 12 layer.

[jos-fisher-jejgkej]

Thank you for your reply. I used the nsys tool to analyze the time consumption of each module in the draft model and found that the proportion of time consumption in lm_cead is very small. Moreover, there is no clear positive correlation between the time consumption of lm_ head and vocab_2. May I ask if you have encountered a similar problem before?

[Achazwl]

Are you using the official Eagle repo, I guess? Which model are you using: Eagle-2 or Eagle-3？

[jos-fisher-jejgkej]

Eagle-2, official repo

[roo-nju]

Hi I am also confused about your profiling results, which shows that softmax accounts for 7% and 13% of the E2E latency on Llama2 and Llama3 respectively. According to my profiling results on Nsys and torch profiler, the GPU spent very limited time (~1%) on softmax, while the host did spend a long time for launching softmax kernels. Could you provide the related codes for profiling both softmax and LM-Head that can reproduce the results in Fig.6?

[Achazwl]

You can use our repo.
You can follow the installation guide of our repo, https://github.com/thunlp/FR-Spec?tab=readme-ov-file#installation-from-source
After installation, use nsys to profile https://github.com/thunlp/FR-Spec/blob/main/examples/example_generate.py

[Achazwl]

Maybe you can also try FastMTP, they use FR-Spec as part of their methods, and they reproduce the speedup.

[roo-nju]

Thanks for the prompt reply, I will try that