does torchinfo support LLM models e.g. llama? if so, do you usage examples?

[chrjxj]

does torchinfo support LLM models e.g. llama? if so, do you usage examples?
thanks!

[fromthefox]

i have the same question...have you figure it out? 😄

[liuup]

Hi, I have the same question! 😸

[fromthefox]

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.
😄

[liuup]

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. 😄

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!

[fromthefox]

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. 😄

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!

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!