TylerYep
does torchinfo support LLM models e.g. llama? if so, do you usage examples?
does torchinfo support LLM models e.g. llama? if so, do you usage examples? thanks!
i have the same question...have you figure it out? 😄
i tried torchinfo.summary with gpt-2 and it works.
but note that, when you use
torchinfo.summary(model, input_size=(...))
, torchinfo will build a tensor_type as input, that works for CNN but not work for LLM, LLM needs a int/long_type input. That means you must build a input_data and use
torchinfo.summary(model, input_data=(your_data))
that will works.
PS: I only tried it with gpt-2, but i think it will also works on llama.
😄
i have the same question...have you figure it out? 😄
i tried torchinfo.summary with gpt-2 and it works. but note that, when you use
torchinfo.summary(model, input_size=(...)), torchinfo will build a tensor_type as input, that works for CNN but not work for LLM, LLM needs a int/long_type input. That means you must build a input_data and usetorchinfo.summary(model, input_data=(your_data))that will works. PS: I only tried it with gpt-2, but i think it will also works on llama. 😄
Hi, maybe I find the best way to print the model structure with summary() and Qwen/Qwen2.5-0.5B model, here is my example code, you can try it on your device.
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from torchinfo import summary
model_name = "Qwen/Qwen2.5-0.5B"
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype="float32",
device_map="cpu",
attn_implementation="eager"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)
messages = [
{"role": "user", "content": "Jane's ducks lay 16 eggs per day. She eats three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck egg. How much in dollars does she make every day at the farmers' market?"}
]
text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
# print the model structure
summary(model,
input_data=(model_inputs["input_ids"], model_inputs["attention_mask"]),
depth=2)
Here is the output of this exmaple code.
===============================================================================================
Layer (type:depth-idx) Output Shape Param #
===============================================================================================
Qwen2ForCausalLM [1, 2, 83, 64] --
├─Qwen2Model: 1-1 [1, 2, 83, 64] --
│ └─Embedding: 2-1 [1, 83, 896] 136,134,656
│ └─Qwen2RotaryEmbedding: 2-2 [1, 83, 64] --
│ └─ModuleList: 2-3 -- 357,897,216
│ └─Qwen2RMSNorm: 2-4 [1, 83, 896] 896
├─Linear: 1-2 [1, 83, 151936] 136,134,656
===============================================================================================
Total params: 630,167,424
Trainable params: 630,167,424
Non-trainable params: 0
Total mult-adds (Units.MEGABYTES): 630.17
===============================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 332.57
Params size (MB): 2520.67
Estimated Total Size (MB): 2853.24
===============================================================================================
You can adjust the depth=2 parameter to print a more detailed structure. In terms of the Qwen/Qwen2.5-0.5B model, when depth=5 is used, it will print the entire structure of the model.
The code is works on my device and the Qwen/Qwen2.5-0.5B model. I think there is no difference between Qwen and Llama, because we use transformers library. But I haven't test it on the other models. You can try it. Hope this code is benefit for you!
i have the same question...have you figure it out? 😄
i tried torchinfo.summary with gpt-2 and it works. but note that, when you use
torchinfo.summary(model, input_size=(...)), torchinfo will build a tensor_type as input, that works for CNN but not work for LLM, LLM needs a int/long_type input. That means you must build a input_data and usetorchinfo.summary(model, input_data=(your_data))that will works. PS: I only tried it with gpt-2, but i think it will also works on llama. 😄Hi, maybe I find the best way to print the model structure with
summary()andQwen/Qwen2.5-0.5Bmodel, here is my example code, you can try it on your device.from transformers import AutoModelForCausalLM, AutoTokenizer import torch from torchinfo import summary model_name = "Qwen/Qwen2.5-0.5B" model = AutoModelForCausalLM.from_pretrained( model_name, torch_dtype="float32", device_map="cpu", attn_implementation="eager" ) tokenizer = AutoTokenizer.from_pretrained(model_name) messages = [ {"role": "user", "content": "Jane's ducks lay 16 eggs per day. She eats three for breakfast every morning and bakes muffins for her friends every day with four. She sells the remainder at the farmers' market daily for $2 per fresh duck egg. How much in dollars does she make every day at the farmers' market?"} ] text = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) model_inputs = tokenizer([text], return_tensors="pt").to(model.device) # print the model structure summary(model, input_data=(model_inputs["input_ids"], model_inputs["attention_mask"]), depth=2)Here is the output of this exmaple code.
=============================================================================================== Layer (type:depth-idx) Output Shape Param # =============================================================================================== Qwen2ForCausalLM [1, 2, 83, 64] -- ├─Qwen2Model: 1-1 [1, 2, 83, 64] -- │ └─Embedding: 2-1 [1, 83, 896] 136,134,656 │ └─Qwen2RotaryEmbedding: 2-2 [1, 83, 64] -- │ └─ModuleList: 2-3 -- 357,897,216 │ └─Qwen2RMSNorm: 2-4 [1, 83, 896] 896 ├─Linear: 1-2 [1, 83, 151936] 136,134,656 =============================================================================================== Total params: 630,167,424 Trainable params: 630,167,424 Non-trainable params: 0 Total mult-adds (Units.MEGABYTES): 630.17 =============================================================================================== Input size (MB): 0.00 Forward/backward pass size (MB): 332.57 Params size (MB): 2520.67 Estimated Total Size (MB): 2853.24 ===============================================================================================You can adjust the
depth=2parameter to print a more detailed structure. In terms of theQwen/Qwen2.5-0.5Bmodel, whendepth=5is used, it will print the entire structure of the model.The code is works on my device and the
Qwen/Qwen2.5-0.5Bmodel. I think there is no difference betweenQwenandLlama, because we usetransformerslibrary. But I haven't test it on the other models. You can try it. Hope this code is benefit for you!
Really thanks for your code! My code is very similar with yours, and i used a random tensor as input and you used a tokenized-sentence as input. And thanks for your advice for 'depth', i'll try it later. Thanks again!