Steel method forward to guess inputs and dynamic shapes (with Tiny-LLM)

Inputs are always dynamic with LLMs that is why dynamic shapes needs to be specified when a LLM is exported with torch.export.export(). Most of the examples on HuggingFace use method transformers.GenerationMixin.generate() but we only want to export the model and its method forward.

That example shows to guess the inputs of this method even though the model is executed through meth generate.

We focus on the model arnir0/Tiny-LLM. To avoid downloading any weights, we write a function creating a random model based on the same architecture.

Steel the forward method

The first step is to guess the dummy inputs. Let’s use the true model for that. We use the dummy example from the model page.

import copy
import pprint
import torch
import transformers
from onnx_diagnostic import doc
from onnx_diagnostic.helpers import string_type
from onnx_diagnostic.helpers.torch_helper import steal_forward
from onnx_diagnostic.torch_models.llms import get_tiny_llm


MODEL_NAME = "arnir0/Tiny-LLM"
tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)
model = transformers.AutoModelForCausalLM.from_pretrained(MODEL_NAME)

We rewrite the forward method to print the cache dimension.

def _forward_(*args, _f=None, **kwargs):
    assert _f is not None
    if not hasattr(torch.compiler, "is_exporting") or not torch.compiler.is_exporting():
        # torch.compiler.is_exporting requires torch>=2.7
        print("<-", string_type((args, kwargs), with_shape=True, with_min_max=True))
    res = _f(*args, **kwargs)
    if not hasattr(torch.compiler, "is_exporting") or not torch.compiler.is_exporting():
        print("->", string_type(res, with_shape=True, with_min_max=True))
    return res


keep_model_forward = model.forward
model.forward = lambda *args, _f=keep_model_forward, **kwargs: _forward_(
    *args, _f=_f, **kwargs
)

Let’s run the model.

prompt = "Continue: it rains..."
inputs = tokenizer.encode(prompt, return_tensors="pt")

outputs = model.generate(
    inputs, max_length=50, temperature=1, top_k=50, top_p=0.95, do_sample=True
)

generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("-- prompt", prompt)
print("-- answer", generated_text)

<- ((),dict(cache_position:T7s8[0,7:A3.5],past_key_values:DynamicCache(key_cache=#0[], value_cache=#0[]),input_ids:T7s1x8[1,29901:A6305.375],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x8x32000[-15.516718864440918,15.75765609741211:A-3.381915190983544],past_key_values:DynamicCache(key_cache=#1[T1s1x1x8x96[-5.490959167480469,6.226877689361572:A-0.11321351693110653]], value_cache=#1[T1s1x1x8x96[-0.6787744760513306,0.49568021297454834:A0.007227749521139988]]))
<- ((),dict(cache_position:T7s1[8,8:A8.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x8x96[-5.490959167480469,6.226877689361572:A-0.11321351693110653]], value_cache=#1[T1s1x1x8x96[-0.6787744760513306,0.49568021297454834:A0.007227749521139988]]),input_ids:T7s1x1[2866,2866:A2866.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-12.076018333435059,16.944217681884766:A-2.592398094831733],past_key_values:DynamicCache(key_cache=#1[T1s1x1x9x96[-5.490959167480469,6.226877689361572:A-0.1302779765439347]], value_cache=#1[T1s1x1x9x96[-0.6787744760513306,0.49568021297454834:A0.007744434695858352]]))
<- ((),dict(cache_position:T7s1[9,9:A9.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x9x96[-5.490959167480469,6.226877689361572:A-0.1302779765439347]], value_cache=#1[T1s1x1x9x96[-0.6787744760513306,0.49568021297454834:A0.007744434695858352]]),input_ids:T7s1x1[14150,14150:A14150.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-18.20236587524414,6.324185371398926:A-8.229752516841982],past_key_values:DynamicCache(key_cache=#1[T1s1x1x10x96[-5.490959167480469,6.226877689361572:A-0.1353976684111937]], value_cache=#1[T1s1x1x10x96[-0.6787744760513306,0.49568021297454834:A0.008736979494627425]]))
<- ((),dict(cache_position:T7s1[10,10:A10.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x10x96[-5.490959167480469,6.226877689361572:A-0.1353976684111937]], value_cache=#1[T1s1x1x10x96[-0.6787744760513306,0.49568021297454834:A0.008736979494627425]]),input_ids:T7s1x1[13,13:A13.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-9.766801834106445,9.472546577453613:A-3.835216622609645],past_key_values:DynamicCache(key_cache=#1[T1s1x1x11x96[-5.510405540466309,6.323276519775391:A-0.14122593572844352]], value_cache=#1[T1s1x1x11x96[-0.6787744760513306,0.7704185843467712:A0.010512653572332607]]))
<- ((),dict(cache_position:T7s1[11,11:A11.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x11x96[-5.510405540466309,6.323276519775391:A-0.14122593572844352]], value_cache=#1[T1s1x1x11x96[-0.6787744760513306,0.7704185843467712:A0.010512653572332607]]),input_ids:T7s1x1[29896,29896:A29896.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-17.130678176879883,6.499423027038574:A-9.773156030296349],past_key_values:DynamicCache(key_cache=#1[T1s1x1x12x96[-5.56224250793457,7.80775785446167:A-0.1263863019698369]], value_cache=#1[T1s1x1x12x96[-0.6787744760513306,0.7704185843467712:A0.009432555446222877]]))
<- ((),dict(cache_position:T7s1[12,12:A12.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x12x96[-5.56224250793457,7.80775785446167:A-0.1263863019698369]], value_cache=#1[T1s1x1x12x96[-0.6787744760513306,0.7704185843467712:A0.009432555446222877]]),input_ids:T7s1x1[29955,29955:A29955.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-18.34309196472168,3.843973398208618:A-10.386367086275015],past_key_values:DynamicCache(key_cache=#1[T1s1x1x13x96[-5.56224250793457,7.80775785446167:A-0.10667757329018712]], value_cache=#1[T1s1x1x13x96[-0.6787744760513306,0.7704185843467712:A0.007383622652865504]]))
<- ((),dict(cache_position:T7s1[13,13:A13.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x13x96[-5.56224250793457,7.80775785446167:A-0.10667757329018712]], value_cache=#1[T1s1x1x13x96[-0.6787744760513306,0.7704185843467712:A0.007383622652865504]]),input_ids:T7s1x1[29906,29906:A29906.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-18.532150268554688,3.666583776473999:A-11.269741500732488],past_key_values:DynamicCache(key_cache=#1[T1s1x1x14x96[-5.56224250793457,7.80775785446167:A-0.10199566604803424]], value_cache=#1[T1s1x1x14x96[-0.6787744760513306,0.7704185843467712:A0.0064541851418847985]]))
<- ((),dict(cache_position:T7s1[14,14:A14.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x14x96[-5.56224250793457,7.80775785446167:A-0.10199566604803424]], value_cache=#1[T1s1x1x14x96[-0.6787744760513306,0.7704185843467712:A0.0064541851418847985]]),input_ids:T7s1x1[29955,29955:A29955.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-18.748641967773438,2.377869129180908:A-10.9063553344775],past_key_values:DynamicCache(key_cache=#1[T1s1x1x15x96[-6.514342308044434,7.80775785446167:A-0.1067288919787744]], value_cache=#1[T1s1x1x15x96[-0.6787744760513306,0.7704185843467712:A0.0048770014079309474]]))
<- ((),dict(cache_position:T7s1[15,15:A15.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x15x96[-6.514342308044434,7.80775785446167:A-0.1067288919787744]], value_cache=#1[T1s1x1x15x96[-0.6787744760513306,0.7704185843467712:A0.0048770014079309474]]),input_ids:T7s1x1[29889,29889:A29889.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-15.992144584655762,3.1707417964935303:A-9.290915607287083],past_key_values:DynamicCache(key_cache=#1[T1s1x1x16x96[-6.514342308044434,7.80775785446167:A-0.11071639796887212]], value_cache=#1[T1s1x1x16x96[-0.6787744760513306,0.7704185843467712:A0.005295240012164489]]))
<- ((),dict(cache_position:T7s1[16,16:A16.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x16x96[-6.514342308044434,7.80775785446167:A-0.11071639796887212]], value_cache=#1[T1s1x1x16x96[-0.6787744760513306,0.7704185843467712:A0.005295240012164489]]),input_ids:T7s1x1[1346,1346:A1346.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-12.513595581054688,5.998507499694824:A-6.863655102951452],past_key_values:DynamicCache(key_cache=#1[T1s1x1x17x96[-6.514342308044434,7.80775785446167:A-0.11430105171986689]], value_cache=#1[T1s1x1x17x96[-0.6787744760513306,0.7704185843467712:A0.006192031966935139]]))
<- ((),dict(cache_position:T7s1[17,17:A17.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x17x96[-6.514342308044434,7.80775785446167:A-0.11430105171986689]], value_cache=#1[T1s1x1x17x96[-0.6787744760513306,0.7704185843467712:A0.006192031966935139]]),input_ids:T7s1x1[11008,11008:A11008.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-16.149032592773438,6.592485427856445:A-8.162168410183629],past_key_values:DynamicCache(key_cache=#1[T1s1x1x18x96[-6.514342308044434,7.80775785446167:A-0.11271050087944551]], value_cache=#1[T1s1x1x18x96[-0.6787744760513306,0.7704185843467712:A0.006270507420552989]]))
<- ((),dict(cache_position:T7s1[18,18:A18.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x18x96[-6.514342308044434,7.80775785446167:A-0.11271050087944551]], value_cache=#1[T1s1x1x18x96[-0.6787744760513306,0.7704185843467712:A0.006270507420552989]]),input_ids:T7s1x1[338,338:A338.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-17.68008041381836,4.83764123916626:A-8.521328096436337],past_key_values:DynamicCache(key_cache=#1[T1s1x1x19x96[-6.514342308044434,7.80775785446167:A-0.1070911308130522]], value_cache=#1[T1s1x1x19x96[-0.6787744760513306,0.7704185843467712:A0.006120377149679706]]))
<- ((),dict(cache_position:T7s1[19,19:A19.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x19x96[-6.514342308044434,7.80775785446167:A-0.1070911308130522]], value_cache=#1[T1s1x1x19x96[-0.6787744760513306,0.7704185843467712:A0.006120377149679706]]),input_ids:T7s1x1[278,278:A278.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-19.417095184326172,1.8751906156539917:A-8.957327721646056],past_key_values:DynamicCache(key_cache=#1[T1s1x1x20x96[-6.514342308044434,7.80775785446167:A-0.1065764879203319]], value_cache=#1[T1s1x1x20x96[-0.6787744760513306,0.7704185843467712:A0.006497487897217979]]))
<- ((),dict(cache_position:T7s1[20,20:A20.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x20x96[-6.514342308044434,7.80775785446167:A-0.1065764879203319]], value_cache=#1[T1s1x1x20x96[-0.6787744760513306,0.7704185843467712:A0.006497487897217979]]),input_ids:T7s1x1[6407,6407:A6407.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-15.863143920898438,4.246685981750488:A-8.345390231630299],past_key_values:DynamicCache(key_cache=#1[T1s1x1x21x96[-6.514342308044434,7.80775785446167:A-0.10707745971018402]], value_cache=#1[T1s1x1x21x96[-0.6787744760513306,0.7704185843467712:A0.006420572361259965]]))
<- ((),dict(cache_position:T7s1[21,21:A21.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x21x96[-6.514342308044434,7.80775785446167:A-0.10707745971018402]], value_cache=#1[T1s1x1x21x96[-0.6787744760513306,0.7704185843467712:A0.006420572361259965]]),input_ids:T7s1x1[23178,23178:A23178.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-14.924324989318848,9.368085861206055:A-5.842367285697721],past_key_values:DynamicCache(key_cache=#1[T1s1x1x22x96[-6.514342308044434,7.80775785446167:A-0.10598940944807757]], value_cache=#1[T1s1x1x22x96[-0.6787744760513306,0.7704185843467712:A0.006922851670224466]]))
<- ((),dict(cache_position:T7s1[22,22:A22.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x22x96[-6.514342308044434,7.80775785446167:A-0.10598940944807757]], value_cache=#1[T1s1x1x22x96[-0.6787744760513306,0.7704185843467712:A0.006922851670224466]]),input_ids:T7s1x1[304,304:A304.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-14.325243949890137,6.519316673278809:A-6.48442860307591],past_key_values:DynamicCache(key_cache=#1[T1s1x1x23x96[-6.736118793487549,7.80775785446167:A-0.10517701237820776]], value_cache=#1[T1s1x1x23x96[-0.6787744760513306,0.7704185843467712:A0.007360825890386416]]))
<- ((),dict(cache_position:T7s1[23,23:A23.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x23x96[-6.736118793487549,7.80775785446167:A-0.10517701237820776]], value_cache=#1[T1s1x1x23x96[-0.6787744760513306,0.7704185843467712:A0.007360825890386416]]),input_ids:T7s1x1[4803,4803:A4803.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-16.578031539916992,4.424250602722168:A-8.651429952671053],past_key_values:DynamicCache(key_cache=#1[T1s1x1x24x96[-6.736118793487549,7.80775785446167:A-0.10474299553258486]], value_cache=#1[T1s1x1x24x96[-0.6787744760513306,0.7704185843467712:A0.007589129441580806]]))
<- ((),dict(cache_position:T7s1[24,24:A24.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x24x96[-6.736118793487549,7.80775785446167:A-0.10474299553258486]], value_cache=#1[T1s1x1x24x96[-0.6787744760513306,0.7704185843467712:A0.007589129441580806]]),input_ids:T7s1x1[3575,3575:A3575.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-14.941386222839355,5.769210338592529:A-7.401098891170696],past_key_values:DynamicCache(key_cache=#1[T1s1x1x25x96[-6.736118793487549,7.80775785446167:A-0.10711594390348182]], value_cache=#1[T1s1x1x25x96[-0.6787744760513306,0.7704185843467712:A0.00848312252195683]]))
<- ((),dict(cache_position:T7s1[25,25:A25.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x25x96[-6.736118793487549,7.80775785446167:A-0.10711594390348182]], value_cache=#1[T1s1x1x25x96[-0.6787744760513306,0.7704185843467712:A0.00848312252195683]]),input_ids:T7s1x1[4634,4634:A4634.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-15.953036308288574,8.907347679138184:A-7.339830758744385],past_key_values:DynamicCache(key_cache=#1[T1s1x1x26x96[-6.736118793487549,7.80775785446167:A-0.10462525407209074]], value_cache=#1[T1s1x1x26x96[-0.6787744760513306,0.7704185843467712:A0.008167866171129972]]))
<- ((),dict(cache_position:T7s1[26,26:A26.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x26x96[-6.736118793487549,7.80775785446167:A-0.10462525407209074]], value_cache=#1[T1s1x1x26x96[-0.6787744760513306,0.7704185843467712:A0.008167866171129972]]),input_ids:T7s1x1[29908,29908:A29908.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-16.06087875366211,7.677635192871094:A-8.212195269609335],past_key_values:DynamicCache(key_cache=#1[T1s1x1x27x96[-6.736118793487549,7.80775785446167:A-0.1044719296038283]], value_cache=#1[T1s1x1x27x96[-0.7138619422912598,0.7704185843467712:A0.0070055666498415065]]))
<- ((),dict(cache_position:T7s1[27,27:A27.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x27x96[-6.736118793487549,7.80775785446167:A-0.1044719296038283]], value_cache=#1[T1s1x1x27x96[-0.7138619422912598,0.7704185843467712:A0.0070055666498415065]]),input_ids:T7s1x1[13,13:A13.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-8.597193717956543,11.263106346130371:A-2.7108958317828367],past_key_values:DynamicCache(key_cache=#1[T1s1x1x28x96[-6.736118793487549,7.80775785446167:A-0.10238736237287023]], value_cache=#1[T1s1x1x28x96[-0.7138619422912598,0.7704185843467712:A0.007764989067682325]]))
<- ((),dict(cache_position:T7s1[28,28:A28.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x28x96[-6.736118793487549,7.80775785446167:A-0.10238736237287023]], value_cache=#1[T1s1x1x28x96[-0.7138619422912598,0.7704185843467712:A0.007764989067682325]]),input_ids:T7s1x1[29906,29906:A29906.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-18.96987533569336,6.3426313400268555:A-9.897809827283956],past_key_values:DynamicCache(key_cache=#1[T1s1x1x29x96[-6.736118793487549,7.80775785446167:A-0.10419577948351648]], value_cache=#1[T1s1x1x29x96[-0.7138619422912598,0.7704185843467712:A0.0073031445308358875]]))
<- ((),dict(cache_position:T7s1[29,29:A29.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x29x96[-6.736118793487549,7.80775785446167:A-0.10419577948351648]], value_cache=#1[T1s1x1x29x96[-0.7138619422912598,0.7704185843467712:A0.0073031445308358875]]),input_ids:T7s1x1[29889,29889:A29889.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-18.66803741455078,3.4406561851501465:A-9.641325709228637],past_key_values:DynamicCache(key_cache=#1[T1s1x1x30x96[-6.736118793487549,7.80775785446167:A-0.09744727589131799]], value_cache=#1[T1s1x1x30x96[-0.7138619422912598,0.7704185843467712:A0.007445333682330278]]))
<- ((),dict(cache_position:T7s1[30,30:A30.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x30x96[-6.736118793487549,7.80775785446167:A-0.09744727589131799]], value_cache=#1[T1s1x1x30x96[-0.7138619422912598,0.7704185843467712:A0.007445333682330278]]),input_ids:T7s1x1[3617,3617:A3617.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-15.766657829284668,6.086407661437988:A-6.891880532442592],past_key_values:DynamicCache(key_cache=#1[T1s1x1x31x96[-6.736118793487549,7.80775785446167:A-0.09227814455923464]], value_cache=#1[T1s1x1x31x96[-0.7138619422912598,0.7704185843467712:A0.0072815972057567275]]))
<- ((),dict(cache_position:T7s1[31,31:A31.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x31x96[-6.736118793487549,7.80775785446167:A-0.09227814455923464]], value_cache=#1[T1s1x1x31x96[-0.7138619422912598,0.7704185843467712:A0.0072815972057567275]]),input_ids:T7s1x1[7535,7535:A7535.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-17.793766021728516,6.4512834548950195:A-7.496981759910006],past_key_values:DynamicCache(key_cache=#1[T1s1x1x32x96[-6.736118793487549,7.80775785446167:A-0.09067019900157902]], value_cache=#1[T1s1x1x32x96[-0.7138619422912598,0.7704185843467712:A0.007518240814187986]]))
<- ((),dict(cache_position:T7s1[32,32:A32.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x32x96[-6.736118793487549,7.80775785446167:A-0.09067019900157902]], value_cache=#1[T1s1x1x32x96[-0.7138619422912598,0.7704185843467712:A0.007518240814187986]]),input_ids:T7s1x1[278,278:A278.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-18.929738998413086,5.662778377532959:A-8.433525333310477],past_key_values:DynamicCache(key_cache=#1[T1s1x1x33x96[-6.736118793487549,7.80775785446167:A-0.08647315130208036]], value_cache=#1[T1s1x1x33x96[-0.7138619422912598,0.7704185843467712:A0.007704433277407901]]))
<- ((),dict(cache_position:T7s1[33,33:A33.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x33x96[-6.736118793487549,7.80775785446167:A-0.08647315130208036]], value_cache=#1[T1s1x1x33x96[-0.7138619422912598,0.7704185843467712:A0.007704433277407901]]),input_ids:T7s1x1[2743,2743:A2743.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-18.49707794189453,5.1465325355529785:A-9.808384840307292],past_key_values:DynamicCache(key_cache=#1[T1s1x1x34x96[-6.736118793487549,7.80775785446167:A-0.08382221280315243]], value_cache=#1[T1s1x1x34x96[-0.7138619422912598,0.7704185843467712:A0.007704160918895771]]))
<- ((),dict(cache_position:T7s1[34,34:A34.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x34x96[-6.736118793487549,7.80775785446167:A-0.08382221280315243]], value_cache=#1[T1s1x1x34x96[-0.7138619422912598,0.7704185843467712:A0.007704160918895771]]),input_ids:T7s1x1[982,982:A982.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-19.50584602355957,8.170029640197754:A-10.122835167907178],past_key_values:DynamicCache(key_cache=#1[T1s1x1x35x96[-6.736118793487549,7.80775785446167:A-0.08074572847912177]], value_cache=#1[T1s1x1x35x96[-0.7138619422912598,0.7704185843467712:A0.00831381521998732]]))
<- ((),dict(cache_position:T7s1[35,35:A35.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x35x96[-6.736118793487549,7.80775785446167:A-0.08074572847912177]], value_cache=#1[T1s1x1x35x96[-0.7138619422912598,0.7704185843467712:A0.00831381521998732]]),input_ids:T7s1x1[29892,29892:A29892.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-15.600412368774414,7.457620620727539:A-6.493258041781607],past_key_values:DynamicCache(key_cache=#1[T1s1x1x36x96[-6.736118793487549,7.80775785446167:A-0.07793717438687088]], value_cache=#1[T1s1x1x36x96[-0.7138619422912598,0.7704185843467712:A0.008467665720095956]]))
<- ((),dict(cache_position:T7s1[36,36:A36.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x36x96[-6.736118793487549,7.80775785446167:A-0.07793717438687088]], value_cache=#1[T1s1x1x36x96[-0.7138619422912598,0.7704185843467712:A0.008467665720095956]]),input_ids:T7s1x1[3692,3692:A3692.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-18.347639083862305,9.92923641204834:A-7.448127461537253],past_key_values:DynamicCache(key_cache=#1[T1s1x1x37x96[-6.736118793487549,7.80775785446167:A-0.07207132491958537]], value_cache=#1[T1s1x1x37x96[-0.7138619422912598,0.7704185843467712:A0.008781119560427288]]))
<- ((),dict(cache_position:T7s1[37,37:A37.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x37x96[-6.736118793487549,7.80775785446167:A-0.07207132491958537]], value_cache=#1[T1s1x1x37x96[-0.7138619422912598,0.7704185843467712:A0.008781119560427288]]),input_ids:T7s1x1[297,297:A297.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-16.862762451171875,7.377256393432617:A-7.443456991467857],past_key_values:DynamicCache(key_cache=#1[T1s1x1x38x96[-6.736118793487549,7.80775785446167:A-0.06766095194504024]], value_cache=#1[T1s1x1x38x96[-0.7138619422912598,0.7704185843467712:A0.008785818313598563]]))
<- ((),dict(cache_position:T7s1[38,38:A38.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x38x96[-6.736118793487549,7.80775785446167:A-0.06766095194504024]], value_cache=#1[T1s1x1x38x96[-0.7138619422912598,0.7704185843467712:A0.008785818313598563]]),input_ids:T7s1x1[4856,4856:A4856.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-13.052538871765137,10.900754928588867:A-5.0423419076236895],past_key_values:DynamicCache(key_cache=#1[T1s1x1x39x96[-6.736118793487549,7.80775785446167:A-0.06517613900807472]], value_cache=#1[T1s1x1x39x96[-0.7138619422912598,0.7704185843467712:A0.00918441853292458]]))
<- ((),dict(cache_position:T7s1[39,39:A39.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x39x96[-6.736118793487549,7.80775785446167:A-0.06517613900807472]], value_cache=#1[T1s1x1x39x96[-0.7138619422912598,0.7704185843467712:A0.00918441853292458]]),input_ids:T7s1x1[15313,15313:A15313.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-12.375218391418457,19.666357040405273:A-0.564182700350415],past_key_values:DynamicCache(key_cache=#1[T1s1x1x40x96[-6.736118793487549,7.80775785446167:A-0.06354604363198936]], value_cache=#1[T1s1x1x40x96[-0.7138619422912598,0.7704185843467712:A0.00848009937929343]]))
<- ((),dict(cache_position:T7s1[40,40:A40.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x40x96[-6.736118793487549,7.80775785446167:A-0.06354604363198936]], value_cache=#1[T1s1x1x40x96[-0.7138619422912598,0.7704185843467712:A0.00848009937929343]]),input_ids:T7s1x1[1144,1144:A1144.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-12.423444747924805,13.019462585449219:A-4.359929949273588],past_key_values:DynamicCache(key_cache=#1[T1s1x1x41x96[-6.736118793487549,7.80775785446167:A-0.0598409421267783]], value_cache=#1[T1s1x1x41x96[-0.7138619422912598,0.7704185843467712:A0.008237699777606114]]))
<- ((),dict(cache_position:T7s1[41,41:A41.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x41x96[-6.736118793487549,7.80775785446167:A-0.0598409421267783]], value_cache=#1[T1s1x1x41x96[-0.7138619422912598,0.7704185843467712:A0.008237699777606114]]),input_ids:T7s1x1[1048,1048:A1048.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-17.85638999938965,7.006597518920898:A-7.712441322462167],past_key_values:DynamicCache(key_cache=#1[T1s1x1x42x96[-6.736118793487549,7.80775785446167:A-0.05822318260551482]], value_cache=#1[T1s1x1x42x96[-0.7138619422912598,0.7704185843467712:A0.008019101587211797]]))
<- ((),dict(cache_position:T7s1[42,42:A42.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x42x96[-6.736118793487549,7.80775785446167:A-0.05822318260551482]], value_cache=#1[T1s1x1x42x96[-0.7138619422912598,0.7704185843467712:A0.008019101587211797]]),input_ids:T7s1x1[4856,4856:A4856.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-14.272891998291016,11.699682235717773:A-5.122771466542501],past_key_values:DynamicCache(key_cache=#1[T1s1x1x43x96[-6.736118793487549,7.80775785446167:A-0.05636586276512023]], value_cache=#1[T1s1x1x43x96[-0.7138619422912598,0.7704185843467712:A0.00839845333791183]]))
<- ((),dict(cache_position:T7s1[43,43:A43.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x43x96[-6.736118793487549,7.80775785446167:A-0.05636586276512023]], value_cache=#1[T1s1x1x43x96[-0.7138619422912598,0.7704185843467712:A0.00839845333791183]]),input_ids:T7s1x1[1058,1058:A1058.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-13.048810958862305,11.082671165466309:A-4.1979296498598995],past_key_values:DynamicCache(key_cache=#1[T1s1x1x44x96[-6.736118793487549,7.80775785446167:A-0.05389009520017267]], value_cache=#1[T1s1x1x44x96[-0.7138619422912598,0.7704185843467712:A0.008558822813576873]]))
<- ((),dict(cache_position:T7s1[44,44:A44.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x44x96[-6.736118793487549,7.80775785446167:A-0.05389009520017267]], value_cache=#1[T1s1x1x44x96[-0.7138619422912598,0.7704185843467712:A0.008558822813576873]]),input_ids:T7s1x1[366,366:A366.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-13.942654609680176,11.131608963012695:A-4.844878984155599],past_key_values:DynamicCache(key_cache=#1[T1s1x1x45x96[-6.736118793487549,7.80775785446167:A-0.05387379795213713]], value_cache=#1[T1s1x1x45x96[-0.7138619422912598,0.7704185843467712:A0.008907078204137414]]))
<- ((),dict(cache_position:T7s1[45,45:A45.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x45x96[-6.736118793487549,7.80775785446167:A-0.05387379795213713]], value_cache=#1[T1s1x1x45x96[-0.7138619422912598,0.7704185843467712:A0.008907078204137414]]),input_ids:T7s1x1[470,470:A470.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-15.920233726501465,10.317113876342773:A-5.626348454394377],past_key_values:DynamicCache(key_cache=#1[T1s1x1x46x96[-6.736118793487549,7.80775785446167:A-0.053487907945173385]], value_cache=#1[T1s1x1x46x96[-0.7138619422912598,0.7704185843467712:A0.008658576126735816]]))
<- ((),dict(cache_position:T7s1[46,46:A46.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x46x96[-6.736118793487549,7.80775785446167:A-0.053487907945173385]], value_cache=#1[T1s1x1x46x96[-0.7138619422912598,0.7704185843467712:A0.008658576126735816]]),input_ids:T7s1x1[4856,4856:A4856.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-14.120901107788086,12.107810974121094:A-4.978954564803047],past_key_values:DynamicCache(key_cache=#1[T1s1x1x47x96[-6.736118793487549,7.80775785446167:A-0.051776735347997836]], value_cache=#1[T1s1x1x47x96[-0.7138619422912598,0.7704185843467712:A0.008992036780790653]]))
<- ((),dict(cache_position:T7s1[47,47:A47.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x47x96[-6.736118793487549,7.80775785446167:A-0.051776735347997836]], value_cache=#1[T1s1x1x47x96[-0.7138619422912598,0.7704185843467712:A0.008992036780790653]]),input_ids:T7s1x1[304,304:A304.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-16.174392700195312,8.619673728942871:A-6.25002147730859],past_key_values:DynamicCache(key_cache=#1[T1s1x1x48x96[-6.736118793487549,7.80775785446167:A-0.05104604224207934]], value_cache=#1[T1s1x1x48x96[-0.7138619422912598,0.7704185843467712:A0.009158791404814792]]))
<- ((),dict(cache_position:T7s1[48,48:A48.0],past_key_values:DynamicCache(key_cache=#1[T1s1x1x48x96[-6.736118793487549,7.80775785446167:A-0.05104604224207934]], value_cache=#1[T1s1x1x48x96[-0.7138619422912598,0.7704185843467712:A0.009158791404814792]]),input_ids:T7s1x1[4459,4459:A4459.0],inputs_embeds:None,use_cache:bool=True,return_dict:bool=True))
-> CausalLMOutputWithPast(logits:T1s1x1x32000[-15.924087524414062,8.78065013885498:A-5.516755177611951],past_key_values:DynamicCache(key_cache=#1[T1s1x1x49x96[-6.736118793487549,7.80775785446167:A-0.049849658982541574]], value_cache=#1[T1s1x1x49x96[-0.7138619422912598,0.7704185843467712:A0.008960381900185919]]))
-- prompt Continue: it rains...
-- answer Continue: it rains... Continue
1727. “Why is the Best Places to Use Your Life"
2. Get yourself the wrong way, whether in someone complains about someone who you or someone to feel their

Let’s restore the forward as it was.

model.forward = keep_model_forward

Another syntax with onnx_diagnostic.helpers.torch_helper.steal_forward().

with steal_forward(model):
    model.generate(inputs, max_length=50, temperature=1, top_k=50, top_p=0.95, do_sample=True)

+ -- stolen forward for class LlamaForCausalLM -- iteration 0
  <- args=() --- kwargs=dict(cache_position:T7s8,past_key_values:DynamicCache(key_cache=#0[], value_cache=#0[]),input_ids:T7s1x8,inputs_embeds:None,use_cache:bool,return_dict:bool)
  -> CausalLMOutputWithPast(logits:T1s1x8x32000,past_key_values:DynamicCache(key_cache=#1[T1s1x1x8x96], value_cache=#1[T1s1x1x8x96]))
-.

Untrained model

This part can skipped if you are only interested in exporting the original model. It is useful to create a unit test to ensure a specific architecture can be exported despite the many changes brought to torch or transformers.

Let’s create an untrained model using the config file provided config.json to create an untrained model: onnx_diagnostic.torch_models.llms.get_tiny_llm(). Then let’s use it.

experiment = get_tiny_llm()
untrained_model, inputs, dynamic_shapes = (
    experiment["model"],
    experiment["inputs"],
    experiment["dynamic_shapes"],
)

Before we run it, we make a copy of the inputs as the cache get modified by the execution. Then it is no longer valid associated with the previous input_ids and mask.

cloned_inputs = copy.deepcopy(inputs)

print("input type before", string_type(inputs, with_shape=True))

expected_output = untrained_model(**inputs)

print("input type after-", string_type(inputs, with_shape=True))

input type before dict(input_ids:T7s2x3,attention_mask:T7s2x33,position_ids:T7s2x3,past_key_values:DynamicCache(key_cache=#1[T1s2x1x30x96], value_cache=#1[T1s2x1x30x96]))
input type after- dict(input_ids:T7s2x3,attention_mask:T7s2x33,position_ids:T7s2x3,past_key_values:DynamicCache(key_cache=#1[T1s2x1x33x96], value_cache=#1[T1s2x1x33x96]))

The outputs

print("result type", string_type(expected_output, with_shape=True))

result type CausalLMOutputWithPast(logits:T1s2x3x32000,past_key_values:DynamicCache(key_cache=#1[T1s2x1x33x96], value_cache=#1[T1s2x1x33x96]))

It works.

ExportedProgram

try:
    ep = torch.export.export(
        untrained_model, (), kwargs=cloned_inputs, dynamic_shapes=dynamic_shapes, strict=False
    )
    print("It worked:")
    print(ep)
except Exception as e:
    # To work, it needs at least PRs:
    # * https://github.com/huggingface/transformers/pull/36311
    # * https://github.com/huggingface/transformers/pull/36652
    print("It failed:", e)

It worked:
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, p_model_embed_tokens_weight: "f32[32000, 192]", p_model_layers_0_self_attn_q_proj_weight: "f32[192, 192]", p_model_layers_0_self_attn_k_proj_weight: "f32[96, 192]", p_model_layers_0_self_attn_v_proj_weight: "f32[96, 192]", p_model_layers_0_self_attn_o_proj_weight: "f32[192, 192]", p_model_layers_0_mlp_gate_proj_weight: "f32[1024, 192]", p_model_layers_0_mlp_up_proj_weight: "f32[1024, 192]", p_model_layers_0_mlp_down_proj_weight: "f32[192, 1024]", p_model_layers_0_input_layernorm_weight: "f32[192]", p_model_layers_0_post_attention_layernorm_weight: "f32[192]", p_model_norm_weight: "f32[192]", p_lm_head_weight: "f32[32000, 192]", b_model_rotary_emb_inv_freq: "f32[48]", input_ids: "i64[s41, s2]", attention_mask: "i64[s41, s2 + s67]", position_ids: "i64[s41, s2]", past_key_values_key_cache_0: "f32[s41, 1, s67, 96]", past_key_values_value_cache_0: "f32[s41, 1, s67, 96]"):
             #
            sym_size_int_22: "Sym(s41)" = torch.ops.aten.sym_size.int(input_ids, 0)
            sym_size_int_23: "Sym(s2)" = torch.ops.aten.sym_size.int(input_ids, 1)
            sym_size_int_24: "Sym(s67)" = torch.ops.aten.sym_size.int(past_key_values_key_cache_0, 2)

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/sparse.py:190 in forward, code: return F.embedding(
            embedding: "f32[s41, s2, 192]" = torch.ops.aten.embedding.default(p_model_embed_tokens_weight, input_ids);  p_model_embed_tokens_weight = input_ids = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:542 in forward, code: past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
            add: "Sym(s2 + s67)" = sym_size_int_24 + sym_size_int_23

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:541 in forward, code: cache_position = torch.arange(
            arange: "i64[s2]" = torch.ops.aten.arange.start(sym_size_int_24, add, device = device(type='cpu'), pin_memory = False);  sym_size_int_24 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:548 in forward, code: causal_mask = self._update_causal_mask(
            full: "f32[s2, s2 + s67]" = torch.ops.aten.full.default([sym_size_int_23, add], -3.4028234663852886e+38, dtype = torch.float32, device = device(type='cpu'), pin_memory = False)
            triu: "f32[s2, s2 + s67]" = torch.ops.aten.triu.default(full, 1);  full = None
            arange_1: "i64[s2 + s67]" = torch.ops.aten.arange.default(add, device = device(type='cpu'), pin_memory = False)
            reshape: "i64[s2, 1]" = torch.ops.aten.reshape.default(arange, [-1, 1]);  arange = None
            gt: "b8[s2, s2 + s67]" = torch.ops.aten.gt.Tensor(arange_1, reshape);  arange_1 = reshape = None
            mul_: "f32[s2, s2 + s67]" = torch.ops.aten.mul_.Tensor(triu, gt);  triu = gt = None
            unsqueeze: "f32[1, s2, s2 + s67]" = torch.ops.aten.unsqueeze.default(mul_, 0);  mul_ = None
            unsqueeze_1: "f32[1, 1, s2, s2 + s67]" = torch.ops.aten.unsqueeze.default(unsqueeze, 1);  unsqueeze = None
            slice_1: "f32[1, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(unsqueeze_1, 2, 0, 9223372036854775807);  unsqueeze_1 = None
            slice_2: "f32[1, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_1, 3, 0, 9223372036854775807);  slice_1 = None
            expand: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.expand.default(slice_2, [sym_size_int_22, 1, -1, -1]);  slice_2 = None
            clone: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.clone.default(expand);  expand = None
            slice_3: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(clone)
            slice_4: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_3, 1);  slice_3 = None
            slice_5: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_4, 2);  slice_4 = None
            slice_6: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_5, 3, None, add);  slice_5 = None
            slice_7: "i64[s41, s2 + s67]" = torch.ops.aten.slice.Tensor(attention_mask, 0, 0, 9223372036854775807);  attention_mask = None
            unsqueeze_2: "i64[s41, 1, s2 + s67]" = torch.ops.aten.unsqueeze.default(slice_7, 1);  slice_7 = None
            unsqueeze_3: "i64[s41, 1, 1, s2 + s67]" = torch.ops.aten.unsqueeze.default(unsqueeze_2, 2);  unsqueeze_2 = None
            slice_8: "i64[s41, 1, 1, s2 + s67]" = torch.ops.aten.slice.Tensor(unsqueeze_3, 3, 0, 9223372036854775807);  unsqueeze_3 = None
            _assert_tensor_metadata_default = torch.ops.aten._assert_tensor_metadata.default(slice_8, dtype = torch.int64, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default = None
            to: "i64[s41, 1, 1, s2 + s67]" = torch.ops.aten.to.dtype_layout(slice_8, dtype = torch.int64, layout = torch.strided, device = device(type='cpu'));  slice_8 = None
            add_2: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.add.Tensor(slice_6, to);  slice_6 = to = None
            eq_4: "b8[s41, 1, s2, s2 + s67]" = torch.ops.aten.eq.Scalar(add_2, 0);  add_2 = None
            slice_9: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(clone)
            slice_10: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_9, 1);  slice_9 = None
            slice_11: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_10, 2);  slice_10 = None
            slice_12: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_11, 3, None, add);  slice_11 = None
            masked_fill: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.masked_fill.Scalar(slice_12, eq_4, -3.4028234663852886e+38);  slice_12 = eq_4 = None
            slice_13: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(clone, 0, 0, 9223372036854775807)
            slice_14: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_13, 1, 0, 9223372036854775807);  slice_13 = None
            slice_15: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_14, 2, 0, 9223372036854775807);  slice_14 = None
            copy_: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.copy_.default(slice_15, masked_fill);  slice_15 = masked_fill = copy_ = None

            # No stacktrace found for following nodes
            submod_3 = self.submod_1
            wrap_with_set_grad_enabled = torch.ops.higher_order.wrap_with_set_grad_enabled(False, submod_3, b_model_rotary_emb_inv_freq, sym_size_int_22, position_ids);  submod_3 = b_model_rotary_emb_inv_freq = position_ids = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:125 in forward, code: return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
            to_6: "f32[s41, s2, 96]" = wrap_with_set_grad_enabled[0]
            to_7: "f32[s41, s2, 96]" = wrap_with_set_grad_enabled[1];  wrap_with_set_grad_enabled = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:82 in forward, code: hidden_states = hidden_states.to(torch.float32)
            _assert_tensor_metadata_default_8 = torch.ops.aten._assert_tensor_metadata.default(embedding, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_8 = None
            to_8: "f32[s41, s2, 192]" = torch.ops.aten.to.dtype(embedding, torch.float32);  embedding = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:83 in forward, code: variance = hidden_states.pow(2).mean(-1, keepdim=True)
            pow_1: "f32[s41, s2, 192]" = torch.ops.aten.pow.Tensor_Scalar(to_8, 2)
            mean: "f32[s41, s2, 1]" = torch.ops.aten.mean.dim(pow_1, [-1], True);  pow_1 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:84 in forward, code: hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
            add_3: "f32[s41, s2, 1]" = torch.ops.aten.add.Tensor(mean, 1e-05);  mean = None
            rsqrt: "f32[s41, s2, 1]" = torch.ops.aten.rsqrt.default(add_3);  add_3 = None
            mul_2: "f32[s41, s2, 192]" = torch.ops.aten.mul.Tensor(to_8, rsqrt);  rsqrt = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:85 in forward, code: return self.weight * hidden_states.to(input_dtype)
            _assert_tensor_metadata_default_9 = torch.ops.aten._assert_tensor_metadata.default(mul_2, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_9 = None
            to_9: "f32[s41, s2, 192]" = torch.ops.aten.to.dtype(mul_2, torch.float32);  mul_2 = None
            mul_3: "f32[s41, s2, 192]" = torch.ops.aten.mul.Tensor(p_model_layers_0_input_layernorm_weight, to_9);  p_model_layers_0_input_layernorm_weight = to_9 = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear: "f32[s41, s2, 192]" = torch.ops.aten.linear.default(mul_3, p_model_layers_0_self_attn_q_proj_weight);  p_model_layers_0_self_attn_q_proj_weight = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:254 in forward, code: query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
            view: "f32[s41, s2, 2, 96]" = torch.ops.aten.view.default(linear, [sym_size_int_22, sym_size_int_23, -1, 96]);  linear = None
            transpose_1: "f32[s41, 2, s2, 96]" = torch.ops.aten.transpose.int(view, 1, 2);  view = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear_1: "f32[s41, s2, 96]" = torch.ops.aten.linear.default(mul_3, p_model_layers_0_self_attn_k_proj_weight);  p_model_layers_0_self_attn_k_proj_weight = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:255 in forward, code: key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
            view_1: "f32[s41, s2, 1, 96]" = torch.ops.aten.view.default(linear_1, [sym_size_int_22, sym_size_int_23, -1, 96]);  linear_1 = None
            transpose_2: "f32[s41, 1, s2, 96]" = torch.ops.aten.transpose.int(view_1, 1, 2);  view_1 = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear_2: "f32[s41, s2, 96]" = torch.ops.aten.linear.default(mul_3, p_model_layers_0_self_attn_v_proj_weight);  mul_3 = p_model_layers_0_self_attn_v_proj_weight = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:256 in forward, code: value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
            view_2: "f32[s41, s2, 1, 96]" = torch.ops.aten.view.default(linear_2, [sym_size_int_22, sym_size_int_23, -1, 96]);  linear_2 = None
            transpose_3: "f32[s41, 1, s2, 96]" = torch.ops.aten.transpose.int(view_2, 1, 2);  view_2 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:259 in forward, code: query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
            unsqueeze_7: "f32[s41, 1, s2, 96]" = torch.ops.aten.unsqueeze.default(to_6, 1);  to_6 = None
            unsqueeze_8: "f32[s41, 1, s2, 96]" = torch.ops.aten.unsqueeze.default(to_7, 1);  to_7 = None
            mul_4: "f32[s41, 2, s2, 96]" = torch.ops.aten.mul.Tensor(transpose_1, unsqueeze_7)
            slice_19: "f32[s41, 2, s2, 48]" = torch.ops.aten.slice.Tensor(transpose_1, 3, 0, 48)
            slice_20: "f32[s41, 2, s2, 48]" = torch.ops.aten.slice.Tensor(transpose_1, 3, 48, 9223372036854775807);  transpose_1 = None
            neg: "f32[s41, 2, s2, 48]" = torch.ops.aten.neg.default(slice_20);  slice_20 = None
            cat_1: "f32[s41, 2, s2, 96]" = torch.ops.aten.cat.default([neg, slice_19], -1);  neg = slice_19 = None
            mul_5: "f32[s41, 2, s2, 96]" = torch.ops.aten.mul.Tensor(cat_1, unsqueeze_8);  cat_1 = None
            add_4: "f32[s41, 2, s2, 96]" = torch.ops.aten.add.Tensor(mul_4, mul_5);  mul_4 = mul_5 = None
            mul_6: "f32[s41, 1, s2, 96]" = torch.ops.aten.mul.Tensor(transpose_2, unsqueeze_7);  unsqueeze_7 = None
            slice_21: "f32[s41, 1, s2, 48]" = torch.ops.aten.slice.Tensor(transpose_2, 3, 0, 48)
            slice_22: "f32[s41, 1, s2, 48]" = torch.ops.aten.slice.Tensor(transpose_2, 3, 48, 9223372036854775807);  transpose_2 = None
            neg_1: "f32[s41, 1, s2, 48]" = torch.ops.aten.neg.default(slice_22);  slice_22 = None
            cat_2: "f32[s41, 1, s2, 96]" = torch.ops.aten.cat.default([neg_1, slice_21], -1);  neg_1 = slice_21 = None
            mul_7: "f32[s41, 1, s2, 96]" = torch.ops.aten.mul.Tensor(cat_2, unsqueeze_8);  cat_2 = unsqueeze_8 = None
            add_5: "f32[s41, 1, s2, 96]" = torch.ops.aten.add.Tensor(mul_6, mul_7);  mul_6 = mul_7 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:264 in forward, code: key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
            cat_3: "f32[s41, 1, s2 + s67, 96]" = torch.ops.aten.cat.default([past_key_values_key_cache_0, add_5], -2);  past_key_values_key_cache_0 = add_5 = None
            cat_4: "f32[s41, 1, s2 + s67, 96]" = torch.ops.aten.cat.default([past_key_values_value_cache_0, transpose_3], -2);  past_key_values_value_cache_0 = transpose_3 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:277 in forward, code: attn_output, attn_weights = attention_interface(
            slice_23: "f32[s41, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(cat_3, 0, 0, 9223372036854775807)
            slice_24: "f32[s41, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(slice_23, 1, 0, 9223372036854775807);  slice_23 = None
            unsqueeze_9: "f32[s41, 1, 1, s2 + s67, 96]" = torch.ops.aten.unsqueeze.default(slice_24, 2);  slice_24 = None
            slice_25: "f32[s41, 1, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(unsqueeze_9, 3, 0, 9223372036854775807);  unsqueeze_9 = None
            slice_26: "f32[s41, 1, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(slice_25, 4, 0, 9223372036854775807);  slice_25 = None
            expand_2: "f32[s41, 1, 2, s2 + s67, 96]" = torch.ops.aten.expand.default(slice_26, [sym_size_int_22, 1, 2, add, 96]);  slice_26 = None
            reshape_1: "f32[s41, 2, s2 + s67, 96]" = torch.ops.aten.reshape.default(expand_2, [sym_size_int_22, 2, add, 96]);  expand_2 = None
            slice_27: "f32[s41, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(cat_4, 0, 0, 9223372036854775807)
            slice_28: "f32[s41, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(slice_27, 1, 0, 9223372036854775807);  slice_27 = None
            unsqueeze_10: "f32[s41, 1, 1, s2 + s67, 96]" = torch.ops.aten.unsqueeze.default(slice_28, 2);  slice_28 = None
            slice_29: "f32[s41, 1, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(unsqueeze_10, 3, 0, 9223372036854775807);  unsqueeze_10 = None
            slice_30: "f32[s41, 1, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(slice_29, 4, 0, 9223372036854775807);  slice_29 = None
            expand_3: "f32[s41, 1, 2, s2 + s67, 96]" = torch.ops.aten.expand.default(slice_30, [sym_size_int_22, 1, 2, add, 96]);  slice_30 = None
            reshape_2: "f32[s41, 2, s2 + s67, 96]" = torch.ops.aten.reshape.default(expand_3, [sym_size_int_22, 2, add, 96]);  expand_3 = None
            slice_31: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(clone);  clone = None
            slice_32: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_31, 1);  slice_31 = None
            slice_33: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_32, 2);  slice_32 = None
            slice_34: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_33, 3, None, add);  slice_33 = add = None
            contiguous: "f32[s41, 2, s2, 96]" = torch.ops.aten.contiguous.default(add_4);  add_4 = None
            contiguous_1: "f32[s41, 2, s2 + s67, 96]" = torch.ops.aten.contiguous.default(reshape_1);  reshape_1 = None
            contiguous_2: "f32[s41, 2, s2 + s67, 96]" = torch.ops.aten.contiguous.default(reshape_2);  reshape_2 = None
            scaled_dot_product_attention: "f32[s41, 2, s2, 96]" = torch.ops.aten.scaled_dot_product_attention.default(contiguous, contiguous_1, contiguous_2, slice_34, scale = 0.10206207261596575);  contiguous = contiguous_1 = contiguous_2 = slice_34 = None
            transpose_4: "f32[s41, s2, 2, 96]" = torch.ops.aten.transpose.int(scaled_dot_product_attention, 1, 2);  scaled_dot_product_attention = None
            contiguous_3: "f32[s41, s2, 2, 96]" = torch.ops.aten.contiguous.default(transpose_4);  transpose_4 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:288 in forward, code: attn_output = attn_output.reshape(*input_shape, -1).contiguous()
            reshape_3: "f32[s41, s2, 192]" = torch.ops.aten.reshape.default(contiguous_3, [sym_size_int_22, sym_size_int_23, -1]);  contiguous_3 = sym_size_int_22 = sym_size_int_23 = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear_3: "f32[s41, s2, 192]" = torch.ops.aten.linear.default(reshape_3, p_model_layers_0_self_attn_o_proj_weight);  reshape_3 = p_model_layers_0_self_attn_o_proj_weight = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:331 in forward, code: hidden_states = residual + hidden_states
            add_7: "f32[s41, s2, 192]" = torch.ops.aten.add.Tensor(to_8, linear_3);  to_8 = linear_3 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:82 in forward, code: hidden_states = hidden_states.to(torch.float32)
            _assert_tensor_metadata_default_10 = torch.ops.aten._assert_tensor_metadata.default(add_7, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_10 = None
            to_10: "f32[s41, s2, 192]" = torch.ops.aten.to.dtype(add_7, torch.float32);  add_7 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:83 in forward, code: variance = hidden_states.pow(2).mean(-1, keepdim=True)
            pow_2: "f32[s41, s2, 192]" = torch.ops.aten.pow.Tensor_Scalar(to_10, 2)
            mean_1: "f32[s41, s2, 1]" = torch.ops.aten.mean.dim(pow_2, [-1], True);  pow_2 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:84 in forward, code: hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
            add_8: "f32[s41, s2, 1]" = torch.ops.aten.add.Tensor(mean_1, 1e-05);  mean_1 = None
            rsqrt_1: "f32[s41, s2, 1]" = torch.ops.aten.rsqrt.default(add_8);  add_8 = None
            mul_8: "f32[s41, s2, 192]" = torch.ops.aten.mul.Tensor(to_10, rsqrt_1);  rsqrt_1 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:85 in forward, code: return self.weight * hidden_states.to(input_dtype)
            _assert_tensor_metadata_default_11 = torch.ops.aten._assert_tensor_metadata.default(mul_8, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_11 = None
            to_11: "f32[s41, s2, 192]" = torch.ops.aten.to.dtype(mul_8, torch.float32);  mul_8 = None
            mul_9: "f32[s41, s2, 192]" = torch.ops.aten.mul.Tensor(p_model_layers_0_post_attention_layernorm_weight, to_11);  p_model_layers_0_post_attention_layernorm_weight = to_11 = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear_4: "f32[s41, s2, 1024]" = torch.ops.aten.linear.default(mul_9, p_model_layers_0_mlp_gate_proj_weight);  p_model_layers_0_mlp_gate_proj_weight = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/activation.py:434 in forward, code: return F.silu(input, inplace=self.inplace)
            silu: "f32[s41, s2, 1024]" = torch.ops.aten.silu.default(linear_4);  linear_4 = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear_5: "f32[s41, s2, 1024]" = torch.ops.aten.linear.default(mul_9, p_model_layers_0_mlp_up_proj_weight);  mul_9 = p_model_layers_0_mlp_up_proj_weight = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:174 in forward, code: down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
            mul_10: "f32[s41, s2, 1024]" = torch.ops.aten.mul.Tensor(silu, linear_5);  silu = linear_5 = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear_6: "f32[s41, s2, 192]" = torch.ops.aten.linear.default(mul_10, p_model_layers_0_mlp_down_proj_weight);  mul_10 = p_model_layers_0_mlp_down_proj_weight = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:337 in forward, code: hidden_states = residual + hidden_states
            add_9: "f32[s41, s2, 192]" = torch.ops.aten.add.Tensor(to_10, linear_6);  to_10 = linear_6 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:82 in forward, code: hidden_states = hidden_states.to(torch.float32)
            _assert_tensor_metadata_default_12 = torch.ops.aten._assert_tensor_metadata.default(add_9, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_12 = None
            to_12: "f32[s41, s2, 192]" = torch.ops.aten.to.dtype(add_9, torch.float32);  add_9 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:83 in forward, code: variance = hidden_states.pow(2).mean(-1, keepdim=True)
            pow_3: "f32[s41, s2, 192]" = torch.ops.aten.pow.Tensor_Scalar(to_12, 2)
            mean_2: "f32[s41, s2, 1]" = torch.ops.aten.mean.dim(pow_3, [-1], True);  pow_3 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:84 in forward, code: hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
            add_10: "f32[s41, s2, 1]" = torch.ops.aten.add.Tensor(mean_2, 1e-05);  mean_2 = None
            rsqrt_2: "f32[s41, s2, 1]" = torch.ops.aten.rsqrt.default(add_10);  add_10 = None
            mul_11: "f32[s41, s2, 192]" = torch.ops.aten.mul.Tensor(to_12, rsqrt_2);  to_12 = rsqrt_2 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:85 in forward, code: return self.weight * hidden_states.to(input_dtype)
            _assert_tensor_metadata_default_13 = torch.ops.aten._assert_tensor_metadata.default(mul_11, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_13 = None
            to_13: "f32[s41, s2, 192]" = torch.ops.aten.to.dtype(mul_11, torch.float32);  mul_11 = None
            mul_12: "f32[s41, s2, 192]" = torch.ops.aten.mul.Tensor(p_model_norm_weight, to_13);  p_model_norm_weight = to_13 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:825 in forward, code: logits = self.lm_head(hidden_states[:, slice_indices, :])
            slice_35: "f32[s41, s2, 192]" = torch.ops.aten.slice.Tensor(mul_12);  mul_12 = None
            slice_36: "f32[s41, s2, 192]" = torch.ops.aten.slice.Tensor(slice_35, 1, 0);  slice_35 = None
            slice_37: "f32[s41, s2, 192]" = torch.ops.aten.slice.Tensor(slice_36, 2);  slice_36 = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear_7: "f32[s41, s2, 32000]" = torch.ops.aten.linear.default(slice_37, p_lm_head_weight);  slice_37 = p_lm_head_weight = None
            return (linear_7, cat_3, cat_4)

        class submod_1(torch.nn.Module):
            def forward(self, b_model_rotary_emb_inv_freq: "f32[48]", sym_size_int_22: "Sym(s41)", position_ids: "i64[s41, s2]"):
                 # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:115 in forward, code: inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
                unsqueeze_4: "f32[1, 48]" = torch.ops.aten.unsqueeze.default(b_model_rotary_emb_inv_freq, 0);  b_model_rotary_emb_inv_freq = None
                slice_16: "f32[1, 48]" = torch.ops.aten.slice.Tensor(unsqueeze_4, 1, 0, 9223372036854775807);  unsqueeze_4 = None
                unsqueeze_5: "f32[1, 48, 1]" = torch.ops.aten.unsqueeze.default(slice_16, 2);  slice_16 = None
                _assert_tensor_metadata_default_1 = torch.ops.aten._assert_tensor_metadata.default(unsqueeze_5, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_1 = None
                to_1: "f32[1, 48, 1]" = torch.ops.aten.to.dtype(unsqueeze_5, torch.float32);  unsqueeze_5 = None
                expand_1: "f32[s41, 48, 1]" = torch.ops.aten.expand.default(to_1, [sym_size_int_22, -1, 1]);  to_1 = sym_size_int_22 = None
                _assert_tensor_metadata_default_2 = torch.ops.aten._assert_tensor_metadata.default(expand_1, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_2 = None
                to_2: "f32[s41, 48, 1]" = torch.ops.aten.to.dtype_layout(expand_1, dtype = torch.float32, layout = torch.strided, device = device(type='cpu'));  expand_1 = None

                 # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:116 in forward, code: position_ids_expanded = position_ids[:, None, :].float()
                slice_17: "i64[s41, s2]" = torch.ops.aten.slice.Tensor(position_ids, 0, 0, 9223372036854775807);  position_ids = None
                unsqueeze_6: "i64[s41, 1, s2]" = torch.ops.aten.unsqueeze.default(slice_17, 1);  slice_17 = None
                slice_18: "i64[s41, 1, s2]" = torch.ops.aten.slice.Tensor(unsqueeze_6, 2, 0, 9223372036854775807);  unsqueeze_6 = None
                _assert_tensor_metadata_default_3 = torch.ops.aten._assert_tensor_metadata.default(slice_18, dtype = torch.int64, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_3 = None
                to_3: "f32[s41, 1, s2]" = torch.ops.aten.to.dtype(slice_18, torch.float32);  slice_18 = None

                # No stacktrace found for following nodes
                submod_3 = self.submod_1
                wrap_with_autocast = torch.ops.higher_order.wrap_with_autocast('cpu', torch.bfloat16, False, False, submod_3, to_2, to_3);  submod_3 = to_2 = to_3 = None

                 # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:122 in forward, code: cos = emb.cos() * self.attention_scaling
                mul: "f32[s41, s2, 96]" = wrap_with_autocast[0]

                 # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:123 in forward, code: sin = emb.sin() * self.attention_scaling
                mul_1: "f32[s41, s2, 96]" = wrap_with_autocast[1];  wrap_with_autocast = None

                 # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:125 in forward, code: return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
                _assert_tensor_metadata_default_6 = torch.ops.aten._assert_tensor_metadata.default(mul, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_6 = None
                to_6: "f32[s41, s2, 96]" = torch.ops.aten.to.dtype(mul, torch.float32);  mul = None
                _assert_tensor_metadata_default_7 = torch.ops.aten._assert_tensor_metadata.default(mul_1, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_7 = None
                to_7: "f32[s41, s2, 96]" = torch.ops.aten.to.dtype(mul_1, torch.float32);  mul_1 = None
                return (to_6, to_7)

            class submod_1(torch.nn.Module):
                def forward(self, to_2: "f32[s41, 48, 1]", to_3: "f32[s41, 1, s2]"):
                     # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:120 in forward, code: freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
                    _assert_tensor_metadata_default_4 = torch.ops.aten._assert_tensor_metadata.default(to_2, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_4 = None
                    to_4: "f32[s41, 48, 1]" = torch.ops.aten.to.dtype(to_2, torch.float32);  to_2 = None
                    _assert_tensor_metadata_default_5 = torch.ops.aten._assert_tensor_metadata.default(to_3, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_5 = None
                    to_5: "f32[s41, 1, s2]" = torch.ops.aten.to.dtype(to_3, torch.float32);  to_3 = None
                    matmul: "f32[s41, 48, s2]" = torch.ops.aten.matmul.default(to_4, to_5);  to_4 = to_5 = None
                    transpose: "f32[s41, s2, 48]" = torch.ops.aten.transpose.int(matmul, 1, 2);  matmul = None

                     # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:121 in forward, code: emb = torch.cat((freqs, freqs), dim=-1)
                    cat: "f32[s41, s2, 96]" = torch.ops.aten.cat.default([transpose, transpose], -1);  transpose = None

                     # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:122 in forward, code: cos = emb.cos() * self.attention_scaling
                    cos: "f32[s41, s2, 96]" = torch.ops.aten.cos.default(cat)
                    mul: "f32[s41, s2, 96]" = torch.ops.aten.mul.Tensor(cos, 1.0);  cos = None

                     # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:123 in forward, code: sin = emb.sin() * self.attention_scaling
                    sin: "f32[s41, s2, 96]" = torch.ops.aten.sin.default(cat);  cat = None
                    mul_1: "f32[s41, s2, 96]" = torch.ops.aten.mul.Tensor(sin, 1.0);  sin = None
                    return (mul, mul_1)

Graph signature:
    # inputs
    p_model_embed_tokens_weight: PARAMETER target='model.embed_tokens.weight'
    p_model_layers_0_self_attn_q_proj_weight: PARAMETER target='model.layers.0.self_attn.q_proj.weight'
    p_model_layers_0_self_attn_k_proj_weight: PARAMETER target='model.layers.0.self_attn.k_proj.weight'
    p_model_layers_0_self_attn_v_proj_weight: PARAMETER target='model.layers.0.self_attn.v_proj.weight'
    p_model_layers_0_self_attn_o_proj_weight: PARAMETER target='model.layers.0.self_attn.o_proj.weight'
    p_model_layers_0_mlp_gate_proj_weight: PARAMETER target='model.layers.0.mlp.gate_proj.weight'
    p_model_layers_0_mlp_up_proj_weight: PARAMETER target='model.layers.0.mlp.up_proj.weight'
    p_model_layers_0_mlp_down_proj_weight: PARAMETER target='model.layers.0.mlp.down_proj.weight'
    p_model_layers_0_input_layernorm_weight: PARAMETER target='model.layers.0.input_layernorm.weight'
    p_model_layers_0_post_attention_layernorm_weight: PARAMETER target='model.layers.0.post_attention_layernorm.weight'
    p_model_norm_weight: PARAMETER target='model.norm.weight'
    p_lm_head_weight: PARAMETER target='lm_head.weight'
    b_model_rotary_emb_inv_freq: BUFFER target='model.rotary_emb.inv_freq' persistent=False
    input_ids: USER_INPUT
    attention_mask: USER_INPUT
    position_ids: USER_INPUT
    past_key_values_key_cache_0: USER_INPUT
    past_key_values_value_cache_0: USER_INPUT

    # outputs
    linear_7: USER_OUTPUT
    cat_3: USER_OUTPUT
    cat_4: USER_OUTPUT

Range constraints: {s41: VR[1, 1024], s2: VR[2, 4096], s2 + s67: VR[4, 8192], s67: VR[1, 4096]}

Back to the original model

Let’s use the same dummy inputs but we use the downloaded model. Dummy inputs and dynamic shapes are created by function onnx_diagnostic.torch_models.llms.get_tiny_llm().

data = get_tiny_llm()
inputs, dynamic_shapes = data["inputs"], data["dynamic_shapes"]

Let’s print the inputs.

print(string_type(inputs, with_shape=True))

dict(input_ids:T7s2x3,attention_mask:T7s2x33,position_ids:T7s2x3,past_key_values:DynamicCache(key_cache=#1[T1s2x1x30x96], value_cache=#1[T1s2x1x30x96]))

pprint.pprint(dynamic_shapes)

{'attention_mask': {0: Dim('batch', min=1, max=1024),
                    1: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
                                min=None,
                                max=None,
                                _factory=True)},
 'input_ids': {0: Dim('batch', min=1, max=1024),
               1: Dim('seq_length', min=1, max=4096)},
 'past_key_values': [[{0: Dim('batch', min=1, max=1024),
                       2: Dim('cache_length', min=1, max=4096)}],
                     [{0: Dim('batch', min=1, max=1024),
                       2: Dim('cache_length', min=1, max=4096)}]],
 'position_ids': {0: Dim('batch', min=1, max=1024),
                  1: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
                              min=None,
                              max=None,
                              _factory=True)}}

And Let’s finally export.

try:
    ep = torch.export.export(
        model, (), kwargs=cloned_inputs, dynamic_shapes=dynamic_shapes, strict=False
    )
    print("It worked:")
    print(ep)
except Exception as e:
    # To work, it needs at least PRs:
    # * https://github.com/huggingface/transformers/pull/36311
    # * https://github.com/huggingface/transformers/pull/36652
    print("It failed:", e)

It worked:
ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, p_model_embed_tokens_weight: "f32[32000, 192]", p_model_layers_0_self_attn_q_proj_weight: "f32[192, 192]", p_model_layers_0_self_attn_k_proj_weight: "f32[96, 192]", p_model_layers_0_self_attn_v_proj_weight: "f32[96, 192]", p_model_layers_0_self_attn_o_proj_weight: "f32[192, 192]", p_model_layers_0_mlp_gate_proj_weight: "f32[1024, 192]", p_model_layers_0_mlp_up_proj_weight: "f32[1024, 192]", p_model_layers_0_mlp_down_proj_weight: "f32[192, 1024]", p_model_layers_0_input_layernorm_weight: "f32[192]", p_model_layers_0_post_attention_layernorm_weight: "f32[192]", p_model_norm_weight: "f32[192]", p_lm_head_weight: "f32[32000, 192]", b_model_rotary_emb_inv_freq: "f32[48]", input_ids: "i64[s41, s2]", attention_mask: "i64[s41, s2 + s67]", position_ids: "i64[s41, s2]", past_key_values_key_cache_0: "f32[s41, 1, s67, 96]", past_key_values_value_cache_0: "f32[s41, 1, s67, 96]"):
             #
            sym_size_int_22: "Sym(s41)" = torch.ops.aten.sym_size.int(input_ids, 0)
            sym_size_int_23: "Sym(s2)" = torch.ops.aten.sym_size.int(input_ids, 1)
            sym_size_int_24: "Sym(s67)" = torch.ops.aten.sym_size.int(past_key_values_key_cache_0, 2)

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/sparse.py:190 in forward, code: return F.embedding(
            embedding: "f32[s41, s2, 192]" = torch.ops.aten.embedding.default(p_model_embed_tokens_weight, input_ids);  p_model_embed_tokens_weight = input_ids = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:542 in forward, code: past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
            add: "Sym(s2 + s67)" = sym_size_int_24 + sym_size_int_23

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:541 in forward, code: cache_position = torch.arange(
            arange: "i64[s2]" = torch.ops.aten.arange.start(sym_size_int_24, add, device = device(type='cpu'), pin_memory = False);  sym_size_int_24 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:548 in forward, code: causal_mask = self._update_causal_mask(
            full: "f32[s2, s2 + s67]" = torch.ops.aten.full.default([sym_size_int_23, add], -3.4028234663852886e+38, dtype = torch.float32, device = device(type='cpu'), pin_memory = False)
            triu: "f32[s2, s2 + s67]" = torch.ops.aten.triu.default(full, 1);  full = None
            arange_1: "i64[s2 + s67]" = torch.ops.aten.arange.default(add, device = device(type='cpu'), pin_memory = False)
            reshape: "i64[s2, 1]" = torch.ops.aten.reshape.default(arange, [-1, 1]);  arange = None
            gt: "b8[s2, s2 + s67]" = torch.ops.aten.gt.Tensor(arange_1, reshape);  arange_1 = reshape = None
            mul_: "f32[s2, s2 + s67]" = torch.ops.aten.mul_.Tensor(triu, gt);  triu = gt = None
            unsqueeze: "f32[1, s2, s2 + s67]" = torch.ops.aten.unsqueeze.default(mul_, 0);  mul_ = None
            unsqueeze_1: "f32[1, 1, s2, s2 + s67]" = torch.ops.aten.unsqueeze.default(unsqueeze, 1);  unsqueeze = None
            slice_1: "f32[1, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(unsqueeze_1, 2, 0, 9223372036854775807);  unsqueeze_1 = None
            slice_2: "f32[1, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_1, 3, 0, 9223372036854775807);  slice_1 = None
            expand: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.expand.default(slice_2, [sym_size_int_22, 1, -1, -1]);  slice_2 = None
            clone: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.clone.default(expand);  expand = None
            slice_3: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(clone)
            slice_4: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_3, 1);  slice_3 = None
            slice_5: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_4, 2);  slice_4 = None
            slice_6: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_5, 3, None, add);  slice_5 = None
            slice_7: "i64[s41, s2 + s67]" = torch.ops.aten.slice.Tensor(attention_mask, 0, 0, 9223372036854775807);  attention_mask = None
            unsqueeze_2: "i64[s41, 1, s2 + s67]" = torch.ops.aten.unsqueeze.default(slice_7, 1);  slice_7 = None
            unsqueeze_3: "i64[s41, 1, 1, s2 + s67]" = torch.ops.aten.unsqueeze.default(unsqueeze_2, 2);  unsqueeze_2 = None
            slice_8: "i64[s41, 1, 1, s2 + s67]" = torch.ops.aten.slice.Tensor(unsqueeze_3, 3, 0, 9223372036854775807);  unsqueeze_3 = None
            _assert_tensor_metadata_default = torch.ops.aten._assert_tensor_metadata.default(slice_8, dtype = torch.int64, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default = None
            to: "i64[s41, 1, 1, s2 + s67]" = torch.ops.aten.to.dtype_layout(slice_8, dtype = torch.int64, layout = torch.strided, device = device(type='cpu'));  slice_8 = None
            add_2: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.add.Tensor(slice_6, to);  slice_6 = to = None
            eq_4: "b8[s41, 1, s2, s2 + s67]" = torch.ops.aten.eq.Scalar(add_2, 0);  add_2 = None
            slice_9: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(clone)
            slice_10: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_9, 1);  slice_9 = None
            slice_11: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_10, 2);  slice_10 = None
            slice_12: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_11, 3, None, add);  slice_11 = None
            masked_fill: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.masked_fill.Scalar(slice_12, eq_4, -3.4028234663852886e+38);  slice_12 = eq_4 = None
            slice_13: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(clone, 0, 0, 9223372036854775807)
            slice_14: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_13, 1, 0, 9223372036854775807);  slice_13 = None
            slice_15: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_14, 2, 0, 9223372036854775807);  slice_14 = None
            copy_: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.copy_.default(slice_15, masked_fill);  slice_15 = masked_fill = copy_ = None

            # No stacktrace found for following nodes
            submod_3 = self.submod_1
            wrap_with_set_grad_enabled = torch.ops.higher_order.wrap_with_set_grad_enabled(False, submod_3, b_model_rotary_emb_inv_freq, sym_size_int_22, position_ids);  submod_3 = b_model_rotary_emb_inv_freq = position_ids = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:125 in forward, code: return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
            to_6: "f32[s41, s2, 96]" = wrap_with_set_grad_enabled[0]
            to_7: "f32[s41, s2, 96]" = wrap_with_set_grad_enabled[1];  wrap_with_set_grad_enabled = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:82 in forward, code: hidden_states = hidden_states.to(torch.float32)
            _assert_tensor_metadata_default_8 = torch.ops.aten._assert_tensor_metadata.default(embedding, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_8 = None
            to_8: "f32[s41, s2, 192]" = torch.ops.aten.to.dtype(embedding, torch.float32);  embedding = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:83 in forward, code: variance = hidden_states.pow(2).mean(-1, keepdim=True)
            pow_1: "f32[s41, s2, 192]" = torch.ops.aten.pow.Tensor_Scalar(to_8, 2)
            mean: "f32[s41, s2, 1]" = torch.ops.aten.mean.dim(pow_1, [-1], True);  pow_1 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:84 in forward, code: hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
            add_3: "f32[s41, s2, 1]" = torch.ops.aten.add.Tensor(mean, 1e-05);  mean = None
            rsqrt: "f32[s41, s2, 1]" = torch.ops.aten.rsqrt.default(add_3);  add_3 = None
            mul_2: "f32[s41, s2, 192]" = torch.ops.aten.mul.Tensor(to_8, rsqrt);  rsqrt = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:85 in forward, code: return self.weight * hidden_states.to(input_dtype)
            _assert_tensor_metadata_default_9 = torch.ops.aten._assert_tensor_metadata.default(mul_2, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_9 = None
            to_9: "f32[s41, s2, 192]" = torch.ops.aten.to.dtype(mul_2, torch.float32);  mul_2 = None
            mul_3: "f32[s41, s2, 192]" = torch.ops.aten.mul.Tensor(p_model_layers_0_input_layernorm_weight, to_9);  p_model_layers_0_input_layernorm_weight = to_9 = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear: "f32[s41, s2, 192]" = torch.ops.aten.linear.default(mul_3, p_model_layers_0_self_attn_q_proj_weight);  p_model_layers_0_self_attn_q_proj_weight = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:254 in forward, code: query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
            view: "f32[s41, s2, 2, 96]" = torch.ops.aten.view.default(linear, [sym_size_int_22, sym_size_int_23, -1, 96]);  linear = None
            transpose_1: "f32[s41, 2, s2, 96]" = torch.ops.aten.transpose.int(view, 1, 2);  view = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear_1: "f32[s41, s2, 96]" = torch.ops.aten.linear.default(mul_3, p_model_layers_0_self_attn_k_proj_weight);  p_model_layers_0_self_attn_k_proj_weight = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:255 in forward, code: key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
            view_1: "f32[s41, s2, 1, 96]" = torch.ops.aten.view.default(linear_1, [sym_size_int_22, sym_size_int_23, -1, 96]);  linear_1 = None
            transpose_2: "f32[s41, 1, s2, 96]" = torch.ops.aten.transpose.int(view_1, 1, 2);  view_1 = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear_2: "f32[s41, s2, 96]" = torch.ops.aten.linear.default(mul_3, p_model_layers_0_self_attn_v_proj_weight);  mul_3 = p_model_layers_0_self_attn_v_proj_weight = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:256 in forward, code: value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
            view_2: "f32[s41, s2, 1, 96]" = torch.ops.aten.view.default(linear_2, [sym_size_int_22, sym_size_int_23, -1, 96]);  linear_2 = None
            transpose_3: "f32[s41, 1, s2, 96]" = torch.ops.aten.transpose.int(view_2, 1, 2);  view_2 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:259 in forward, code: query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
            unsqueeze_7: "f32[s41, 1, s2, 96]" = torch.ops.aten.unsqueeze.default(to_6, 1);  to_6 = None
            unsqueeze_8: "f32[s41, 1, s2, 96]" = torch.ops.aten.unsqueeze.default(to_7, 1);  to_7 = None
            mul_4: "f32[s41, 2, s2, 96]" = torch.ops.aten.mul.Tensor(transpose_1, unsqueeze_7)
            slice_19: "f32[s41, 2, s2, 48]" = torch.ops.aten.slice.Tensor(transpose_1, 3, 0, 48)
            slice_20: "f32[s41, 2, s2, 48]" = torch.ops.aten.slice.Tensor(transpose_1, 3, 48, 9223372036854775807);  transpose_1 = None
            neg: "f32[s41, 2, s2, 48]" = torch.ops.aten.neg.default(slice_20);  slice_20 = None
            cat_1: "f32[s41, 2, s2, 96]" = torch.ops.aten.cat.default([neg, slice_19], -1);  neg = slice_19 = None
            mul_5: "f32[s41, 2, s2, 96]" = torch.ops.aten.mul.Tensor(cat_1, unsqueeze_8);  cat_1 = None
            add_4: "f32[s41, 2, s2, 96]" = torch.ops.aten.add.Tensor(mul_4, mul_5);  mul_4 = mul_5 = None
            mul_6: "f32[s41, 1, s2, 96]" = torch.ops.aten.mul.Tensor(transpose_2, unsqueeze_7);  unsqueeze_7 = None
            slice_21: "f32[s41, 1, s2, 48]" = torch.ops.aten.slice.Tensor(transpose_2, 3, 0, 48)
            slice_22: "f32[s41, 1, s2, 48]" = torch.ops.aten.slice.Tensor(transpose_2, 3, 48, 9223372036854775807);  transpose_2 = None
            neg_1: "f32[s41, 1, s2, 48]" = torch.ops.aten.neg.default(slice_22);  slice_22 = None
            cat_2: "f32[s41, 1, s2, 96]" = torch.ops.aten.cat.default([neg_1, slice_21], -1);  neg_1 = slice_21 = None
            mul_7: "f32[s41, 1, s2, 96]" = torch.ops.aten.mul.Tensor(cat_2, unsqueeze_8);  cat_2 = unsqueeze_8 = None
            add_5: "f32[s41, 1, s2, 96]" = torch.ops.aten.add.Tensor(mul_6, mul_7);  mul_6 = mul_7 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:264 in forward, code: key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
            cat_3: "f32[s41, 1, s2 + s67, 96]" = torch.ops.aten.cat.default([past_key_values_key_cache_0, add_5], -2);  past_key_values_key_cache_0 = add_5 = None
            cat_4: "f32[s41, 1, s2 + s67, 96]" = torch.ops.aten.cat.default([past_key_values_value_cache_0, transpose_3], -2);  past_key_values_value_cache_0 = transpose_3 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:277 in forward, code: attn_output, attn_weights = attention_interface(
            slice_23: "f32[s41, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(cat_3, 0, 0, 9223372036854775807)
            slice_24: "f32[s41, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(slice_23, 1, 0, 9223372036854775807);  slice_23 = None
            unsqueeze_9: "f32[s41, 1, 1, s2 + s67, 96]" = torch.ops.aten.unsqueeze.default(slice_24, 2);  slice_24 = None
            slice_25: "f32[s41, 1, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(unsqueeze_9, 3, 0, 9223372036854775807);  unsqueeze_9 = None
            slice_26: "f32[s41, 1, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(slice_25, 4, 0, 9223372036854775807);  slice_25 = None
            expand_2: "f32[s41, 1, 2, s2 + s67, 96]" = torch.ops.aten.expand.default(slice_26, [sym_size_int_22, 1, 2, add, 96]);  slice_26 = None
            reshape_1: "f32[s41, 2, s2 + s67, 96]" = torch.ops.aten.reshape.default(expand_2, [sym_size_int_22, 2, add, 96]);  expand_2 = None
            slice_27: "f32[s41, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(cat_4, 0, 0, 9223372036854775807)
            slice_28: "f32[s41, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(slice_27, 1, 0, 9223372036854775807);  slice_27 = None
            unsqueeze_10: "f32[s41, 1, 1, s2 + s67, 96]" = torch.ops.aten.unsqueeze.default(slice_28, 2);  slice_28 = None
            slice_29: "f32[s41, 1, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(unsqueeze_10, 3, 0, 9223372036854775807);  unsqueeze_10 = None
            slice_30: "f32[s41, 1, 1, s2 + s67, 96]" = torch.ops.aten.slice.Tensor(slice_29, 4, 0, 9223372036854775807);  slice_29 = None
            expand_3: "f32[s41, 1, 2, s2 + s67, 96]" = torch.ops.aten.expand.default(slice_30, [sym_size_int_22, 1, 2, add, 96]);  slice_30 = None
            reshape_2: "f32[s41, 2, s2 + s67, 96]" = torch.ops.aten.reshape.default(expand_3, [sym_size_int_22, 2, add, 96]);  expand_3 = None
            slice_31: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(clone);  clone = None
            slice_32: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_31, 1);  slice_31 = None
            slice_33: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_32, 2);  slice_32 = None
            slice_34: "f32[s41, 1, s2, s2 + s67]" = torch.ops.aten.slice.Tensor(slice_33, 3, None, add);  slice_33 = add = None
            contiguous: "f32[s41, 2, s2, 96]" = torch.ops.aten.contiguous.default(add_4);  add_4 = None
            contiguous_1: "f32[s41, 2, s2 + s67, 96]" = torch.ops.aten.contiguous.default(reshape_1);  reshape_1 = None
            contiguous_2: "f32[s41, 2, s2 + s67, 96]" = torch.ops.aten.contiguous.default(reshape_2);  reshape_2 = None
            scaled_dot_product_attention: "f32[s41, 2, s2, 96]" = torch.ops.aten.scaled_dot_product_attention.default(contiguous, contiguous_1, contiguous_2, slice_34, scale = 0.10206207261596575);  contiguous = contiguous_1 = contiguous_2 = slice_34 = None
            transpose_4: "f32[s41, s2, 2, 96]" = torch.ops.aten.transpose.int(scaled_dot_product_attention, 1, 2);  scaled_dot_product_attention = None
            contiguous_3: "f32[s41, s2, 2, 96]" = torch.ops.aten.contiguous.default(transpose_4);  transpose_4 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:288 in forward, code: attn_output = attn_output.reshape(*input_shape, -1).contiguous()
            reshape_3: "f32[s41, s2, 192]" = torch.ops.aten.reshape.default(contiguous_3, [sym_size_int_22, sym_size_int_23, -1]);  contiguous_3 = sym_size_int_22 = sym_size_int_23 = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear_3: "f32[s41, s2, 192]" = torch.ops.aten.linear.default(reshape_3, p_model_layers_0_self_attn_o_proj_weight);  reshape_3 = p_model_layers_0_self_attn_o_proj_weight = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:331 in forward, code: hidden_states = residual + hidden_states
            add_7: "f32[s41, s2, 192]" = torch.ops.aten.add.Tensor(to_8, linear_3);  to_8 = linear_3 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:82 in forward, code: hidden_states = hidden_states.to(torch.float32)
            _assert_tensor_metadata_default_10 = torch.ops.aten._assert_tensor_metadata.default(add_7, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_10 = None
            to_10: "f32[s41, s2, 192]" = torch.ops.aten.to.dtype(add_7, torch.float32);  add_7 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:83 in forward, code: variance = hidden_states.pow(2).mean(-1, keepdim=True)
            pow_2: "f32[s41, s2, 192]" = torch.ops.aten.pow.Tensor_Scalar(to_10, 2)
            mean_1: "f32[s41, s2, 1]" = torch.ops.aten.mean.dim(pow_2, [-1], True);  pow_2 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:84 in forward, code: hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
            add_8: "f32[s41, s2, 1]" = torch.ops.aten.add.Tensor(mean_1, 1e-05);  mean_1 = None
            rsqrt_1: "f32[s41, s2, 1]" = torch.ops.aten.rsqrt.default(add_8);  add_8 = None
            mul_8: "f32[s41, s2, 192]" = torch.ops.aten.mul.Tensor(to_10, rsqrt_1);  rsqrt_1 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:85 in forward, code: return self.weight * hidden_states.to(input_dtype)
            _assert_tensor_metadata_default_11 = torch.ops.aten._assert_tensor_metadata.default(mul_8, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_11 = None
            to_11: "f32[s41, s2, 192]" = torch.ops.aten.to.dtype(mul_8, torch.float32);  mul_8 = None
            mul_9: "f32[s41, s2, 192]" = torch.ops.aten.mul.Tensor(p_model_layers_0_post_attention_layernorm_weight, to_11);  p_model_layers_0_post_attention_layernorm_weight = to_11 = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear_4: "f32[s41, s2, 1024]" = torch.ops.aten.linear.default(mul_9, p_model_layers_0_mlp_gate_proj_weight);  p_model_layers_0_mlp_gate_proj_weight = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/activation.py:434 in forward, code: return F.silu(input, inplace=self.inplace)
            silu: "f32[s41, s2, 1024]" = torch.ops.aten.silu.default(linear_4);  linear_4 = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear_5: "f32[s41, s2, 1024]" = torch.ops.aten.linear.default(mul_9, p_model_layers_0_mlp_up_proj_weight);  mul_9 = p_model_layers_0_mlp_up_proj_weight = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:174 in forward, code: down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
            mul_10: "f32[s41, s2, 1024]" = torch.ops.aten.mul.Tensor(silu, linear_5);  silu = linear_5 = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear_6: "f32[s41, s2, 192]" = torch.ops.aten.linear.default(mul_10, p_model_layers_0_mlp_down_proj_weight);  mul_10 = p_model_layers_0_mlp_down_proj_weight = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:337 in forward, code: hidden_states = residual + hidden_states
            add_9: "f32[s41, s2, 192]" = torch.ops.aten.add.Tensor(to_10, linear_6);  to_10 = linear_6 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:82 in forward, code: hidden_states = hidden_states.to(torch.float32)
            _assert_tensor_metadata_default_12 = torch.ops.aten._assert_tensor_metadata.default(add_9, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_12 = None
            to_12: "f32[s41, s2, 192]" = torch.ops.aten.to.dtype(add_9, torch.float32);  add_9 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:83 in forward, code: variance = hidden_states.pow(2).mean(-1, keepdim=True)
            pow_3: "f32[s41, s2, 192]" = torch.ops.aten.pow.Tensor_Scalar(to_12, 2)
            mean_2: "f32[s41, s2, 1]" = torch.ops.aten.mean.dim(pow_3, [-1], True);  pow_3 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:84 in forward, code: hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
            add_10: "f32[s41, s2, 1]" = torch.ops.aten.add.Tensor(mean_2, 1e-05);  mean_2 = None
            rsqrt_2: "f32[s41, s2, 1]" = torch.ops.aten.rsqrt.default(add_10);  add_10 = None
            mul_11: "f32[s41, s2, 192]" = torch.ops.aten.mul.Tensor(to_12, rsqrt_2);  to_12 = rsqrt_2 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:85 in forward, code: return self.weight * hidden_states.to(input_dtype)
            _assert_tensor_metadata_default_13 = torch.ops.aten._assert_tensor_metadata.default(mul_11, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_13 = None
            to_13: "f32[s41, s2, 192]" = torch.ops.aten.to.dtype(mul_11, torch.float32);  mul_11 = None
            mul_12: "f32[s41, s2, 192]" = torch.ops.aten.mul.Tensor(p_model_norm_weight, to_13);  p_model_norm_weight = to_13 = None

             # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:825 in forward, code: logits = self.lm_head(hidden_states[:, slice_indices, :])
            slice_35: "f32[s41, s2, 192]" = torch.ops.aten.slice.Tensor(mul_12);  mul_12 = None
            slice_36: "f32[s41, s2, 192]" = torch.ops.aten.slice.Tensor(slice_35, 1, 0);  slice_35 = None
            slice_37: "f32[s41, s2, 192]" = torch.ops.aten.slice.Tensor(slice_36, 2);  slice_36 = None

             # File: ~/vv/this312/lib/python3.12/site-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
            linear_7: "f32[s41, s2, 32000]" = torch.ops.aten.linear.default(slice_37, p_lm_head_weight);  slice_37 = p_lm_head_weight = None
            return (linear_7, cat_3, cat_4)

        class submod_1(torch.nn.Module):
            def forward(self, b_model_rotary_emb_inv_freq: "f32[48]", sym_size_int_22: "Sym(s41)", position_ids: "i64[s41, s2]"):
                 # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:115 in forward, code: inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
                unsqueeze_4: "f32[1, 48]" = torch.ops.aten.unsqueeze.default(b_model_rotary_emb_inv_freq, 0);  b_model_rotary_emb_inv_freq = None
                slice_16: "f32[1, 48]" = torch.ops.aten.slice.Tensor(unsqueeze_4, 1, 0, 9223372036854775807);  unsqueeze_4 = None
                unsqueeze_5: "f32[1, 48, 1]" = torch.ops.aten.unsqueeze.default(slice_16, 2);  slice_16 = None
                _assert_tensor_metadata_default_1 = torch.ops.aten._assert_tensor_metadata.default(unsqueeze_5, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_1 = None
                to_1: "f32[1, 48, 1]" = torch.ops.aten.to.dtype(unsqueeze_5, torch.float32);  unsqueeze_5 = None
                expand_1: "f32[s41, 48, 1]" = torch.ops.aten.expand.default(to_1, [sym_size_int_22, -1, 1]);  to_1 = sym_size_int_22 = None
                _assert_tensor_metadata_default_2 = torch.ops.aten._assert_tensor_metadata.default(expand_1, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_2 = None
                to_2: "f32[s41, 48, 1]" = torch.ops.aten.to.dtype_layout(expand_1, dtype = torch.float32, layout = torch.strided, device = device(type='cpu'));  expand_1 = None

                 # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:116 in forward, code: position_ids_expanded = position_ids[:, None, :].float()
                slice_17: "i64[s41, s2]" = torch.ops.aten.slice.Tensor(position_ids, 0, 0, 9223372036854775807);  position_ids = None
                unsqueeze_6: "i64[s41, 1, s2]" = torch.ops.aten.unsqueeze.default(slice_17, 1);  slice_17 = None
                slice_18: "i64[s41, 1, s2]" = torch.ops.aten.slice.Tensor(unsqueeze_6, 2, 0, 9223372036854775807);  unsqueeze_6 = None
                _assert_tensor_metadata_default_3 = torch.ops.aten._assert_tensor_metadata.default(slice_18, dtype = torch.int64, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_3 = None
                to_3: "f32[s41, 1, s2]" = torch.ops.aten.to.dtype(slice_18, torch.float32);  slice_18 = None

                # No stacktrace found for following nodes
                submod_3 = self.submod_1
                wrap_with_autocast = torch.ops.higher_order.wrap_with_autocast('cpu', torch.bfloat16, False, False, submod_3, to_2, to_3);  submod_3 = to_2 = to_3 = None

                 # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:122 in forward, code: cos = emb.cos() * self.attention_scaling
                mul: "f32[s41, s2, 96]" = wrap_with_autocast[0]

                 # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:123 in forward, code: sin = emb.sin() * self.attention_scaling
                mul_1: "f32[s41, s2, 96]" = wrap_with_autocast[1];  wrap_with_autocast = None

                 # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:125 in forward, code: return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
                _assert_tensor_metadata_default_6 = torch.ops.aten._assert_tensor_metadata.default(mul, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_6 = None
                to_6: "f32[s41, s2, 96]" = torch.ops.aten.to.dtype(mul, torch.float32);  mul = None
                _assert_tensor_metadata_default_7 = torch.ops.aten._assert_tensor_metadata.default(mul_1, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_7 = None
                to_7: "f32[s41, s2, 96]" = torch.ops.aten.to.dtype(mul_1, torch.float32);  mul_1 = None
                return (to_6, to_7)

            class submod_1(torch.nn.Module):
                def forward(self, to_2: "f32[s41, 48, 1]", to_3: "f32[s41, 1, s2]"):
                     # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:120 in forward, code: freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
                    _assert_tensor_metadata_default_4 = torch.ops.aten._assert_tensor_metadata.default(to_2, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_4 = None
                    to_4: "f32[s41, 48, 1]" = torch.ops.aten.to.dtype(to_2, torch.float32);  to_2 = None
                    _assert_tensor_metadata_default_5 = torch.ops.aten._assert_tensor_metadata.default(to_3, dtype = torch.float32, device = device(type='cpu'), layout = torch.strided);  _assert_tensor_metadata_default_5 = None
                    to_5: "f32[s41, 1, s2]" = torch.ops.aten.to.dtype(to_3, torch.float32);  to_3 = None
                    matmul: "f32[s41, 48, s2]" = torch.ops.aten.matmul.default(to_4, to_5);  to_4 = to_5 = None
                    transpose: "f32[s41, s2, 48]" = torch.ops.aten.transpose.int(matmul, 1, 2);  matmul = None

                     # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:121 in forward, code: emb = torch.cat((freqs, freqs), dim=-1)
                    cat: "f32[s41, s2, 96]" = torch.ops.aten.cat.default([transpose, transpose], -1);  transpose = None

                     # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:122 in forward, code: cos = emb.cos() * self.attention_scaling
                    cos: "f32[s41, s2, 96]" = torch.ops.aten.cos.default(cat)
                    mul: "f32[s41, s2, 96]" = torch.ops.aten.mul.Tensor(cos, 1.0);  cos = None

                     # File: ~/github/transformers/src/transformers/models/llama/modeling_llama.py:123 in forward, code: sin = emb.sin() * self.attention_scaling
                    sin: "f32[s41, s2, 96]" = torch.ops.aten.sin.default(cat);  cat = None
                    mul_1: "f32[s41, s2, 96]" = torch.ops.aten.mul.Tensor(sin, 1.0);  sin = None
                    return (mul, mul_1)

Graph signature:
    # inputs
    p_model_embed_tokens_weight: PARAMETER target='model.embed_tokens.weight'
    p_model_layers_0_self_attn_q_proj_weight: PARAMETER target='model.layers.0.self_attn.q_proj.weight'
    p_model_layers_0_self_attn_k_proj_weight: PARAMETER target='model.layers.0.self_attn.k_proj.weight'
    p_model_layers_0_self_attn_v_proj_weight: PARAMETER target='model.layers.0.self_attn.v_proj.weight'
    p_model_layers_0_self_attn_o_proj_weight: PARAMETER target='model.layers.0.self_attn.o_proj.weight'
    p_model_layers_0_mlp_gate_proj_weight: PARAMETER target='model.layers.0.mlp.gate_proj.weight'
    p_model_layers_0_mlp_up_proj_weight: PARAMETER target='model.layers.0.mlp.up_proj.weight'
    p_model_layers_0_mlp_down_proj_weight: PARAMETER target='model.layers.0.mlp.down_proj.weight'
    p_model_layers_0_input_layernorm_weight: PARAMETER target='model.layers.0.input_layernorm.weight'
    p_model_layers_0_post_attention_layernorm_weight: PARAMETER target='model.layers.0.post_attention_layernorm.weight'
    p_model_norm_weight: PARAMETER target='model.norm.weight'
    p_lm_head_weight: PARAMETER target='lm_head.weight'
    b_model_rotary_emb_inv_freq: BUFFER target='model.rotary_emb.inv_freq' persistent=False
    input_ids: USER_INPUT
    attention_mask: USER_INPUT
    position_ids: USER_INPUT
    past_key_values_key_cache_0: USER_INPUT
    past_key_values_value_cache_0: USER_INPUT

    # outputs
    linear_7: USER_OUTPUT
    cat_3: USER_OUTPUT
    cat_4: USER_OUTPUT

Range constraints: {s41: VR[1, 1024], s2: VR[2, 4096], s2 + s67: VR[4, 8192], s67: VR[1, 4096]}

If you have any error, then look at example Export Tiny-LLM with patches.

doc.plot_legend("Tiny-LLM\nforward inputs\nbehind generate", "torch.export.export", "tomato")

Related examples

Export Tiny-LLM with patches

Export Tiny-LLM with patches

Export microsoft/phi-2

Export microsoft/phi-2

Test the export on untrained models

Test the export on untrained models