optimize get_module_leaves speed

[BBuf]

Background

When I try to inference deepseek-v2 with transformers:

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig

model_name = "deepseek-ai/DeepSeek-V2"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
# `max_memory` should be set based on your devices
max_memory = {i: "75GB" for i in range(8)}
# `device_map` cannot be set to `auto`
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, device_map="sequential", torch_dtype=torch.bfloat16, max_memory=max_memory, attn_implementation="eager")
model.generation_config = GenerationConfig.from_pretrained(model_name)
model.generation_config.pad_token_id = model.generation_config.eos_token_id

text = "An attention function can be described as mapping a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors. The output is"
inputs = tokenizer(text, return_tensors="pt")
outputs = model.generate(**inputs.to(model.device), max_new_tokens=100)

result = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(result)

I found that program with stuck...

Solve

In fact, the stuck appearing here .

For DeepSeek V2, the length of module_sizes is 68185, and the code here has a complexity of O(N^2) (where N = module_sizes), which requires a very long time to execute, giving the illusion of being stuck. This PR optimizes the code to have a complexity of O(N), allowing it to quickly reach the stage of loading the large model. The inference results on an 8xA800 machine are also normal after this optimization.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[SunMarc]

Thanks for this great PR @BBuf ! Just a small nit.

[muellerzr]

Thanks for making this much more efficient and fixing 🤗