Skip to content

Failed to export PyTorch traced graph of Mixtral-8x7B-Instruct-v0.1 due to the PR #32429 #38518

Closed
Closed
Failed to export PyTorch traced graph of Mixtral-8x7B-Instruct-v0.1 due to the PR #32429#38518
Labels
Good Difficult Issuebug

Description

@nv-guomingz
nv-guomingz
opened on Jun 1, 2025

System Info

Hi ,

I found the recently transformers 4.52.4 merged this PR #32429
and it led me failed to run below code snippet which it could run successfully with 4.51.3.

import transformers
import torch.export as te
import torch
from contextlib import nullcontext

torch.autocast = lambda *args, **kwargs: nullcontext()  

mixtral = transformers.AutoModelForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1", device_map="meta")
ep = te.export(mixtral, 
                args=(torch.randint(0, 100, (2, 4),device="meta", dtype=torch.int32),
                      torch.randint(0, 100, (2, 4),device="meta", dtype=torch.int32)
                    ), 
                kwargs={}, strict=False
                ).module()

and this is the error call stack:

Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████| 19/19 [00:24<00:00,  1.27s/it]
W0601 17:35:32.549000 19842 torch/fx/experimental/symbolic_shapes.py:6661] failed during evaluate_expr(u0, hint=None, size_oblivious=False, forcing_spec=False
E0601 17:35:32.550000 19842 torch/fx/experimental/recording.py:299] failed while running evaluate_expr(*(u0, None, False, False), **{})
W0601 17:35:32.552000 19842 torch/fx/experimental/symbolic_shapes.py:7208] Unable to find user code corresponding to {u0}




def forward(self, arg0_1: "f32[32000, 4096]", arg1_1: "f32[4096, 4096]", arg2_1: "f32[1024, 4096]", arg3_1: "f32[1024, 4096]", arg4_1: "f32[4096, 4096]", arg5_1: "f32[8, 4096]", arg6_1: "f32[14336, 4096]", arg7_1: "f32[4096, 14336]", arg8_1: "f32[14336, 4096]", arg9_1: "f32[14336, 4096]", arg10_1: "f32[4096, 14336]", arg11_1: "f32[14336, 4096]", arg12_1: "f32[14336, 4096]", arg13_1: "f32[4096, 14336]", arg14_1: "f32[14336, 4096]", arg15_1: "f32[14336, 4096]", arg16_1: "f32[4096, 14336]", arg17_1: "f32[14336, 4096]", arg18_1: "f32[14336, 4096]", arg19_1: "f32[4096, 14336]", arg20_1: "f32[14336, 4096]", arg21_1: "f32[14336, 4096]", arg22_1: "f32[4096, 14336]", arg23_1: "f32[14336, 4096]", arg24_1: "f32[14336, 4096]", arg25_1: "f32[4096, 14336]", arg26_1: "f32[14336, 4096]", arg27_1: "f32[14336, 4096]", arg28_1: "f32[4096, 14336]", arg29_1: "f32[14336, 4096]", arg30_1: "f32[4096]", arg31_1: "f32[4096]", arg32_1: "f32[4096, 4096]", arg33_1: "f32[1024, 4096]", arg34_1: "f32[1024, 4096]", arg35_1: "f32[4096, 4096]", arg36_1: "f32[8, 4096]", arg37_1: "f32[14336, 4096]", arg38_1: "f32[4096, 14336]", arg39_1: "f32[14336, 4096]", arg40_1: "f32[14336, 4096]", arg41_1: "f32[4096, 14336]", arg42_1: "f32[14336, 4096]", arg43_1: "f32[14336, 4096]", arg44_1: "f32[4096, 14336]", arg45_1: "f32[14336, 4096]", arg46_1: "f32[14336, 4096]", arg47_1: "f32[4096, 14336]", arg48_1: "f32[14336, 4096]", arg49_1: "f32[14336, 4096]", arg50_1: "f32[4096, 14336]", arg51_1: "f32[14336, 4096]", arg52_1: "f32[14336, 4096]", arg53_1: "f32[4096, 14336]", arg54_1: "f32[14336, 4096]", arg55_1: "f32[14336, 4096]", arg56_1: "f32[4096, 14336]", arg57_1: "f32[14336, 4096]", arg58_1: "f32[14336, 4096]", arg59_1: "f32[4096, 14336]", arg60_1: "f32[14336, 4096]", arg61_1: "f32[4096]", arg62_1: "f32[4096]", arg63_1: "f32[4096, 4096]", arg64_1: "f32[1024, 4096]", arg65_1: "f32[1024, 4096]", arg66_1: "f32[4096, 4096]", arg67_1: "f32[8, 4096]", arg68_1: "f32[14336, 4096]", arg69_1: "f32[4096, 14336]", arg70_1: "f32[14336, 4096]", arg71_1: "f32[14336, 4096]", arg72_1: "f32[4096, 14336]", arg73_1: "f32[14336, 4096]", arg74_1: "f32[14336, 4096]", arg75_1: "f32[4096, 14336]", arg76_1: "f32[14336, 4096]", arg77_1: "f32[14336, 4096]", arg78_1: "f32[4096, 14336]", arg79_1: "f32[14336, 4096]", arg80_1: "f32[14336, 4096]", arg81_1: "f32[4096, 14336]", arg82_1: "f32[14336, 4096]", arg83_1: "f32[14336, 4096]", arg84_1: "f32[4096, 14336]", arg85_1: "f32[14336, 4096]", arg86_1: "f32[14336, 4096]", arg87_1: "f32[4096, 14336]", arg88_1: "f32[14336, 4096]", arg89_1: "f32[14336, 4096]", arg90_1: "f32[4096, 14336]", arg91_1: "f32[14336, 4096]", arg92_1: "f32[4096]", arg93_1: "f32[4096]", arg94_1: "f32[4096, 4096]", arg95_1: "f32[1024, 4096]", arg96_1: "f32[1024, 4096]", arg97_1: "f32[4096, 4096]", arg98_1: "f32[8, 4096]", arg99_1: "f32[14336, 4096]", arg100_1: "f32[4096, 14336]", arg101_1: "f32[14336, 4096]", arg102_1: "f32[14336, 4096]", arg103_1: "f32[4096, 14336]", arg104_1: "f32[14336, 4096]", arg105_1: "f32[14336, 4096]", arg106_1: "f32[4096, 14336]", arg107_1: "f32[14336, 4096]", arg108_1: "f32[14336, 4096]", arg109_1: "f32[4096, 14336]", arg110_1: "f32[14336, 4096]", arg111_1: "f32[14336, 4096]", arg112_1: "f32[4096, 14336]", arg113_1: "f32[14336, 4096]", arg114_1: "f32[14336, 4096]", arg115_1: "f32[4096, 14336]", arg116_1: "f32[14336, 4096]", arg117_1: "f32[14336, 4096]", arg118_1: "f32[4096, 14336]", arg119_1: "f32[14336, 4096]", arg120_1: "f32[14336, 4096]", arg121_1: "f32[4096, 14336]", arg122_1: "f32[14336, 4096]", arg123_1: "f32[4096]", arg124_1: "f32[4096]", arg125_1: "f32[4096, 4096]", arg126_1: "f32[1024, 4096]", arg127_1: "f32[1024, 4096]", arg128_1: "f32[4096, 4096]", arg129_1: "f32[8, 4096]", arg130_1: "f32[14336, 4096]", arg131_1: "f32[4096, 14336]", arg132_1: "f32[14336, 4096]", arg133_1: "f32[14336, 4096]", arg134_1: "f32[4096, 14336]", arg135_1: "f32[14336, 4096]", arg136_1: "f32[14336, 4096]", arg137_1: "f32[4096, 14336]", arg138_1: "f32[14336, 4096]", arg139_1: "f32[14336, 4096]", arg140_1: "f32[4096, 14336]", arg141_1: "f32[14336, 4096]", arg142_1: "f32[14336, 4096]", arg143_1: "f32[4096, 14336]", arg144_1: "f32[14336, 4096]", arg145_1: "f32[14336, 4096]", arg146_1: "f32[4096, 14336]", arg147_1: "f32[14336, 4096]", arg148_1: "f32[14336, 4096]", arg149_1: "f32[4096, 14336]", arg150_1: "f32[14336, 4096]", arg151_1: "f32[14336, 4096]", arg152_1: "f32[4096, 14336]", arg153_1: "f32[14336, 4096]", arg154_1: "f32[4096]", arg155_1: "f32[4096]", arg156_1: "f32[4096, 4096]", arg157_1: "f32[1024, 4096]", arg158_1: "f32[1024, 4096]", arg159_1: "f32[4096, 4096]", arg160_1: "f32[8, 4096]", arg161_1: "f32[14336, 4096]", arg162_1: "f32[4096, 14336]", arg163_1: "f32[14336, 4096]", arg164_1: "f32[14336, 4096]", arg165_1: "f32[4096, 14336]", arg166_1: "f32[14336, 4096]", arg167_1: "f32[14336, 4096]", arg168_1: "f32[4096, 14336]", arg169_1: "f32[14336, 4096]", arg170_1: "f32[14336, 4096]", arg171_1: "f32[4096, 14336]", arg172_1: "f32[14336, 4096]", arg173_1: "f32[14336, 4096]", arg174_1: "f32[4096, 14336]", arg175_1: "f32[14336, 4096]", arg176_1: "f32[14336, 4096]", arg177_1: "f32[4096, 14336]", arg178_1: "f32[14336, 4096]", arg179_1: "f32[14336, 4096]", arg180_1: "f32[4096, 14336]", arg181_1: "f32[14336, 4096]", arg182_1: "f32[14336, 4096]", arg183_1: "f32[4096, 14336]", arg184_1: "f32[14336, 4096]", arg185_1: "f32[4096]", arg186_1: "f32[4096]", arg187_1: "f32[4096, 4096]", arg188_1: "f32[1024, 4096]", arg189_1: "f32[1024, 4096]", arg190_1: "f32[4096, 4096]", arg191_1: "f32[8, 4096]", arg192_1: "f32[14336, 4096]", arg193_1: "f32[4096, 14336]", arg194_1: "f32[14336, 4096]", arg195_1: "f32[14336, 4096]", arg196_1: "f32[4096, 14336]", arg197_1: "f32[14336, 4096]", arg198_1: "f32[14336, 4096]", arg199_1: "f32[4096, 14336]", arg200_1: "f32[14336, 4096]", arg201_1: "f32[14336, 4096]", arg202_1: "f32[4096, 14336]", arg203_1: "f32[14336, 4096]", arg204_1: "f32[14336, 4096]", arg205_1: "f32[4096, 14336]", arg206_1: "f32[14336, 4096]", arg207_1: "f32[14336, 4096]", arg208_1: "f32[4096, 14336]", arg209_1: "f32[14336, 4096]", arg210_1: "f32[14336, 4096]", arg211_1: "f32[4096, 14336]", arg212_1: "f32[14336, 4096]", arg213_1: "f32[14336, 4096]", arg214_1: "f32[4096, 14336]", arg215_1: "f32[14336, 4096]", arg216_1: "f32[4096]", arg217_1: "f32[4096]", arg218_1: "f32[4096, 4096]", arg219_1: "f32[1024, 4096]", arg220_1: "f32[1024, 4096]", arg221_1: "f32[4096, 4096]", arg222_1: "f32[8, 4096]", arg223_1: "f32[14336, 4096]", arg224_1: "f32[4096, 14336]", arg225_1: "f32[14336, 4096]", arg226_1: "f32[14336, 4096]", arg227_1: "f32[4096, 14336]", arg228_1: "f32[14336, 4096]", arg229_1: "f32[14336, 4096]", arg230_1: "f32[4096, 14336]", arg231_1: "f32[14336, 4096]", arg232_1: "f32[14336, 4096]", arg233_1: "f32[4096, 14336]", arg234_1: "f32[14336, 4096]", arg235_1: "f32[14336, 4096]", arg236_1: "f32[4096, 14336]", arg237_1: "f32[14336, 4096]", arg238_1: "f32[14336, 4096]", arg239_1: "f32[4096, 14336]", arg240_1: "f32[14336, 4096]", arg241_1: "f32[14336, 4096]", arg242_1: "f32[4096, 14336]", arg243_1: "f32[14336, 4096]", arg244_1: "f32[14336, 4096]", arg245_1: "f32[4096, 14336]", arg246_1: "f32[14336, 4096]", arg247_1: "f32[4096]", arg248_1: "f32[4096]", arg249_1: "f32[4096, 4096]", arg250_1: "f32[1024, 4096]", arg251_1: "f32[1024, 4096]", arg252_1: "f32[4096, 4096]", arg253_1: "f32[8, 4096]", arg254_1: "f32[14336, 4096]", arg255_1: "f32[4096, 14336]", arg256_1: "f32[14336, 4096]", arg257_1: "f32[14336, 4096]", arg258_1: "f32[4096, 14336]", arg259_1: "f32[14336, 4096]", arg260_1: "f32[14336, 4096]", arg261_1: "f32[4096, 14336]", arg262_1: "f32[14336, 4096]", arg263_1: "f32[14336, 4096]", arg264_1: "f32[4096, 14336]", arg265_1: "f32[14336, 4096]", arg266_1: "f32[14336, 4096]", arg267_1: "f32[4096, 14336]", arg268_1: "f32[14336, 4096]", arg269_1: "f32[14336, 4096]", arg270_1: "f32[4096, 14336]", arg271_1: "f32[14336, 4096]", arg272_1: "f32[14336, 4096]", arg273_1: "f32[4096, 14336]", arg274_1: "f32[14336, 4096]", arg275_1: "f32[14336, 4096]", arg276_1: "f32[4096, 14336]", arg277_1: "f32[14336, 4096]", arg278_1: "f32[4096]", arg279_1: "f32[4096]", arg280_1: "f32[4096, 4096]", arg281_1: "f32[1024, 4096]", arg282_1: "f32[1024, 4096]", arg283_1: "f32[4096, 4096]", arg284_1: "f32[8, 4096]", arg285_1: "f32[14336, 4096]", arg286_1: "f32[4096, 14336]", arg287_1: "f32[14336, 4096]", arg288_1: "f32[14336, 4096]", arg289_1: "f32[4096, 14336]", arg290_1: "f32[14336, 4096]", arg291_1: "f32[14336, 4096]", arg292_1: "f32[4096, 14336]", arg293_1: "f32[14336, 4096]", arg294_1: "f32[14336, 4096]", arg295_1: "f32[4096, 14336]", arg296_1: "f32[14336, 4096]", arg297_1: "f32[14336, 4096]", arg298_1: "f32[4096, 14336]", arg299_1: "f32[14336, 4096]", arg300_1: "f32[14336, 4096]", arg301_1: "f32[4096, 14336]", arg302_1: "f32[14336, 4096]", arg303_1: "f32[14336, 4096]", arg304_1: "f32[4096, 14336]", arg305_1: "f32[14336, 4096]", arg306_1: "f32[14336, 4096]", arg307_1: "f32[4096, 14336]", arg308_1: "f32[14336, 4096]", arg309_1: "f32[4096]", arg310_1: "f32[4096]", arg311_1: "f32[4096, 4096]", arg312_1: "f32[1024, 4096]", arg313_1: "f32[1024, 4096]", arg314_1: "f32[4096, 4096]", arg315_1: "f32[8, 4096]", arg316_1: "f32[14336, 4096]", arg317_1: "f32[4096, 14336]", arg318_1: "f32[14336, 4096]", arg319_1: "f32[14336, 4096]", arg320_1: "f32[4096, 14336]", arg321_1: "f32[14336, 4096]", arg322_1: "f32[14336, 4096]", arg323_1: "f32[4096, 14336]", arg324_1: "f32[14336, 4096]", arg325_1: "f32[14336, 4096]", arg326_1: "f32[4096, 14336]", arg327_1: "f32[14336, 4096]", arg328_1: "f32[14336, 4096]", arg329_1: "f32[4096, 14336]", arg330_1: "f32[14336, 4096]", arg331_1: "f32[14336, 4096]", arg332_1: "f32[4096, 14336]", arg333_1: "f32[14336, 4096]", arg334_1: "f32[14336, 4096]", arg335_1: "f32[4096, 14336]", arg336_1: "f32[14336, 4096]", arg337_1: "f32[14336, 4096]", arg338_1: "f32[4096, 14336]", arg339_1: "f32[14336, 4096]", arg340_1: "f32[4096]", arg341_1: "f32[4096]", arg342_1: "f32[4096, 4096]", arg343_1: "f32[1024, 4096]", arg344_1: "f32[1024, 4096]", arg345_1: "f32[4096, 4096]", arg346_1: "f32[8, 4096]", arg347_1: "f32[14336, 4096]", arg348_1: "f32[4096, 14336]", arg349_1: "f32[14336, 4096]", arg350_1: "f32[14336, 4096]", arg351_1: "f32[4096, 14336]", arg352_1: "f32[14336, 4096]", arg353_1: "f32[14336, 4096]", arg354_1: "f32[4096, 14336]", arg355_1: "f32[14336, 4096]", arg356_1: "f32[14336, 4096]", arg357_1: "f32[4096, 14336]", arg358_1: "f32[14336, 4096]", arg359_1: "f32[14336, 4096]", arg360_1: "f32[4096, 14336]", arg361_1: "f32[14336, 4096]", arg362_1: "f32[14336, 4096]", arg363_1: "f32[4096, 14336]", arg364_1: "f32[14336, 4096]", arg365_1: "f32[14336, 4096]", arg366_1: "f32[4096, 14336]", arg367_1: "f32[14336, 4096]", arg368_1: "f32[14336, 4096]", arg369_1: "f32[4096, 14336]", arg370_1: "f32[14336, 4096]", arg371_1: "f32[4096]", arg372_1: "f32[4096]", arg373_1: "f32[4096, 4096]", arg374_1: "f32[1024, 4096]", arg375_1: "f32[1024, 4096]", arg376_1: "f32[4096, 4096]", arg377_1: "f32[8, 4096]", arg378_1: "f32[14336, 4096]", arg379_1: "f32[4096, 14336]", arg380_1: "f32[14336, 4096]", arg381_1: "f32[14336, 4096]", arg382_1: "f32[4096, 14336]", arg383_1: "f32[14336, 4096]", arg384_1: "f32[14336, 4096]", arg385_1: "f32[4096, 14336]", arg386_1: "f32[14336, 4096]", arg387_1: "f32[14336, 4096]", arg388_1: "f32[4096, 14336]", arg389_1: "f32[14336, 4096]", arg390_1: "f32[14336, 4096]", arg391_1: "f32[4096, 14336]", arg392_1: "f32[14336, 4096]", arg393_1: "f32[14336, 4096]", arg394_1: "f32[4096, 14336]", arg395_1: "f32[14336, 4096]", arg396_1: "f32[14336, 4096]", arg397_1: "f32[4096, 14336]", arg398_1: "f32[14336, 4096]", arg399_1: "f32[14336, 4096]", arg400_1: "f32[4096, 14336]", arg401_1: "f32[14336, 4096]", arg402_1: "f32[4096]", arg403_1: "f32[4096]", arg404_1: "f32[4096, 4096]", arg405_1: "f32[1024, 4096]", arg406_1: "f32[1024, 4096]", arg407_1: "f32[4096, 4096]", arg408_1: "f32[8, 4096]", arg409_1: "f32[14336, 4096]", arg410_1: "f32[4096, 14336]", arg411_1: "f32[14336, 4096]", arg412_1: "f32[14336, 4096]", arg413_1: "f32[4096, 14336]", arg414_1: "f32[14336, 4096]", arg415_1: "f32[14336, 4096]", arg416_1: "f32[4096, 14336]", arg417_1: "f32[14336, 4096]", arg418_1: "f32[14336, 4096]", arg419_1: "f32[4096, 14336]", arg420_1: "f32[14336, 4096]", arg421_1: "f32[14336, 4096]", arg422_1: "f32[4096, 14336]", arg423_1: "f32[14336, 4096]", arg424_1: "f32[14336, 4096]", arg425_1: "f32[4096, 14336]", arg426_1: "f32[14336, 4096]", arg427_1: "f32[14336, 4096]", arg428_1: "f32[4096, 14336]", arg429_1: "f32[14336, 4096]", arg430_1: "f32[14336, 4096]", arg431_1: "f32[4096, 14336]", arg432_1: "f32[14336, 4096]", arg433_1: "f32[4096]", arg434_1: "f32[4096]", arg435_1: "f32[4096, 4096]", arg436_1: "f32[1024, 4096]", arg437_1: "f32[1024, 4096]", arg438_1: "f32[4096, 4096]", arg439_1: "f32[8, 4096]", arg440_1: "f32[14336, 4096]", arg441_1: "f32[4096, 14336]", arg442_1: "f32[14336, 4096]", arg443_1: "f32[14336, 4096]", arg444_1: "f32[4096, 14336]", arg445_1: "f32[14336, 4096]", arg446_1: "f32[14336, 4096]", arg447_1: "f32[4096, 14336]", arg448_1: "f32[14336, 4096]", arg449_1: "f32[14336, 4096]", arg450_1: "f32[4096, 14336]", arg451_1: "f32[14336, 4096]", arg452_1: "f32[14336, 4096]", arg453_1: "f32[4096, 14336]", arg454_1: "f32[14336, 4096]", arg455_1: "f32[14336, 4096]", arg456_1: "f32[4096, 14336]", arg457_1: "f32[14336, 4096]", arg458_1: "f32[14336, 4096]", arg459_1: "f32[4096, 14336]", arg460_1: "f32[14336, 4096]", arg461_1: "f32[14336, 4096]", arg462_1: "f32[4096, 14336]", arg463_1: "f32[14336, 4096]", arg464_1: "f32[4096]", arg465_1: "f32[4096]", arg466_1: "f32[4096, 4096]", arg467_1: "f32[1024, 4096]", arg468_1: "f32[1024, 4096]", arg469_1: "f32[4096, 4096]", arg470_1: "f32[8, 4096]", arg471_1: "f32[14336, 4096]", arg472_1: "f32[4096, 14336]", arg473_1: "f32[14336, 4096]", arg474_1: "f32[14336, 4096]", arg475_1: "f32[4096, 14336]", arg476_1: "f32[14336, 4096]", arg477_1: "f32[14336, 4096]", arg478_1: "f32[4096, 14336]", arg479_1: "f32[14336, 4096]", arg480_1: "f32[14336, 4096]", arg481_1: "f32[4096, 14336]", arg482_1: "f32[14336, 4096]", arg483_1: "f32[14336, 4096]", arg484_1: "f32[4096, 14336]", arg485_1: "f32[14336, 4096]", arg486_1: "f32[14336, 4096]", arg487_1: "f32[4096, 14336]", arg488_1: "f32[14336, 4096]", arg489_1: "f32[14336, 4096]", arg490_1: "f32[4096, 14336]", arg491_1: "f32[14336, 4096]", arg492_1: "f32[14336, 4096]", arg493_1: "f32[4096, 14336]", arg494_1: "f32[14336, 4096]", arg495_1: "f32[4096]", arg496_1: "f32[4096]", arg497_1: "f32[4096, 4096]", arg498_1: "f32[1024, 4096]", arg499_1: "f32[1024, 4096]", arg500_1: "f32[4096, 4096]", arg501_1: "f32[8, 4096]", arg502_1: "f32[14336, 4096]", arg503_1: "f32[4096, 14336]", arg504_1: "f32[14336, 4096]", arg505_1: "f32[14336, 4096]", arg506_1: "f32[4096, 14336]", arg507_1: "f32[14336, 4096]", arg508_1: "f32[14336, 4096]", arg509_1: "f32[4096, 14336]", arg510_1: "f32[14336, 4096]", arg511_1: "f32[14336, 4096]", arg512_1: "f32[4096, 14336]", arg513_1: "f32[14336, 4096]", arg514_1: "f32[14336, 4096]", arg515_1: "f32[4096, 14336]", arg516_1: "f32[14336, 4096]", arg517_1: "f32[14336, 4096]", arg518_1: "f32[4096, 14336]", arg519_1: "f32[14336, 4096]", arg520_1: "f32[14336, 4096]", arg521_1: "f32[4096, 14336]", arg522_1: "f32[14336, 4096]", arg523_1: "f32[14336, 4096]", arg524_1: "f32[4096, 14336]", arg525_1: "f32[14336, 4096]", arg526_1: "f32[4096]", arg527_1: "f32[4096]", arg528_1: "f32[4096, 4096]", arg529_1: "f32[1024, 4096]", arg530_1: "f32[1024, 4096]", arg531_1: "f32[4096, 4096]", arg532_1: "f32[8, 4096]", arg533_1: "f32[14336, 4096]", arg534_1: "f32[4096, 14336]", arg535_1: "f32[14336, 4096]", arg536_1: "f32[14336, 4096]", arg537_1: "f32[4096, 14336]", arg538_1: "f32[14336, 4096]", arg539_1: "f32[14336, 4096]", arg540_1: "f32[4096, 14336]", arg541_1: "f32[14336, 4096]", arg542_1: "f32[14336, 4096]", arg543_1: "f32[4096, 14336]", arg544_1: "f32[14336, 4096]", arg545_1: "f32[14336, 4096]", arg546_1: "f32[4096, 14336]", arg547_1: "f32[14336, 4096]", arg548_1: "f32[14336, 4096]", arg549_1: "f32[4096, 14336]", arg550_1: "f32[14336, 4096]", arg551_1: "f32[14336, 4096]", arg552_1: "f32[4096, 14336]", arg553_1: "f32[14336, 4096]", arg554_1: "f32[14336, 4096]", arg555_1: "f32[4096, 14336]", arg556_1: "f32[14336, 4096]", arg557_1: "f32[4096]", arg558_1: "f32[4096]", arg559_1: "f32[4096, 4096]", arg560_1: "f32[1024, 4096]", arg561_1: "f32[1024, 4096]", arg562_1: "f32[4096, 4096]", arg563_1: "f32[8, 4096]", arg564_1: "f32[14336, 4096]", arg565_1: "f32[4096, 14336]", arg566_1: "f32[14336, 4096]", arg567_1: "f32[14336, 4096]", arg568_1: "f32[4096, 14336]", arg569_1: "f32[14336, 4096]", arg570_1: "f32[14336, 4096]", arg571_1: "f32[4096, 14336]", arg572_1: "f32[14336, 4096]", arg573_1: "f32[14336, 4096]", arg574_1: "f32[4096, 14336]", arg575_1: "f32[14336, 4096]", arg576_1: "f32[14336, 4096]", arg577_1: "f32[4096, 14336]", arg578_1: "f32[14336, 4096]", arg579_1: "f32[14336, 4096]", arg580_1: "f32[4096, 14336]", arg581_1: "f32[14336, 4096]", arg582_1: "f32[14336, 4096]", arg583_1: "f32[4096, 14336]", arg584_1: "f32[14336, 4096]", arg585_1: "f32[14336, 4096]", arg586_1: "f32[4096, 14336]", arg587_1: "f32[14336, 4096]", arg588_1: "f32[4096]", arg589_1: "f32[4096]", arg590_1: "f32[4096, 4096]", arg591_1: "f32[1024, 4096]", arg592_1: "f32[1024, 4096]", arg593_1: "f32[4096, 4096]", arg594_1: "f32[8, 4096]", arg595_1: "f32[14336, 4096]", arg596_1: "f32[4096, 14336]", arg597_1: "f32[14336, 4096]", arg598_1: "f32[14336, 4096]", arg599_1: "f32[4096, 14336]", arg600_1: "f32[14336, 4096]", arg601_1: "f32[14336, 4096]", arg602_1: "f32[4096, 14336]", arg603_1: "f32[14336, 4096]", arg604_1: "f32[14336, 4096]", arg605_1: "f32[4096, 14336]", arg606_1: "f32[14336, 4096]", arg607_1: "f32[14336, 4096]", arg608_1: "f32[4096, 14336]", arg609_1: "f32[14336, 4096]", arg610_1: "f32[14336, 4096]", arg611_1: "f32[4096, 14336]", arg612_1: "f32[14336, 4096]", arg613_1: "f32[14336, 4096]", arg614_1: "f32[4096, 14336]", arg615_1: "f32[14336, 4096]", arg616_1: "f32[14336, 4096]", arg617_1: "f32[4096, 14336]", arg618_1: "f32[14336, 4096]", arg619_1: "f32[4096]", arg620_1: "f32[4096]", arg621_1: "f32[4096, 4096]", arg622_1: "f32[1024, 4096]", arg623_1: "f32[1024, 4096]", arg624_1: "f32[4096, 4096]", arg625_1: "f32[8, 4096]", arg626_1: "f32[14336, 4096]", arg627_1: "f32[4096, 14336]", arg628_1: "f32[14336, 4096]", arg629_1: "f32[14336, 4096]", arg630_1: "f32[4096, 14336]", arg631_1: "f32[14336, 4096]", arg632_1: "f32[14336, 4096]", arg633_1: "f32[4096, 14336]", arg634_1: "f32[14336, 4096]", arg635_1: "f32[14336, 4096]", arg636_1: "f32[4096, 14336]", arg637_1: "f32[14336, 4096]", arg638_1: "f32[14336, 4096]", arg639_1: "f32[4096, 14336]", arg640_1: "f32[14336, 4096]", arg641_1: "f32[14336, 4096]", arg642_1: "f32[4096, 14336]", arg643_1: "f32[14336, 4096]", arg644_1: "f32[14336, 4096]", arg645_1: "f32[4096, 14336]", arg646_1: "f32[14336, 4096]", arg647_1: "f32[14336, 4096]", arg648_1: "f32[4096, 14336]", arg649_1: "f32[14336, 4096]", arg650_1: "f32[4096]", arg651_1: "f32[4096]", arg652_1: "f32[4096, 4096]", arg653_1: "f32[1024, 4096]", arg654_1: "f32[1024, 4096]", arg655_1: "f32[4096, 4096]", arg656_1: "f32[8, 4096]", arg657_1: "f32[14336, 4096]", arg658_1: "f32[4096, 14336]", arg659_1: "f32[14336, 4096]", arg660_1: "f32[14336, 4096]", arg661_1: "f32[4096, 14336]", arg662_1: "f32[14336, 4096]", arg663_1: "f32[14336, 4096]", arg664_1: "f32[4096, 14336]", arg665_1: "f32[14336, 4096]", arg666_1: "f32[14336, 4096]", arg667_1: "f32[4096, 14336]", arg668_1: "f32[14336, 4096]", arg669_1: "f32[14336, 4096]", arg670_1: "f32[4096, 14336]", arg671_1: "f32[14336, 4096]", arg672_1: "f32[14336, 4096]", arg673_1: "f32[4096, 14336]", arg674_1: "f32[14336, 4096]", arg675_1: "f32[14336, 4096]", arg676_1: "f32[4096, 14336]", arg677_1: "f32[14336, 4096]", arg678_1: "f32[14336, 4096]", arg679_1: "f32[4096, 14336]", arg680_1: "f32[14336, 4096]", arg681_1: "f32[4096]", arg682_1: "f32[4096]", arg683_1: "f32[4096, 4096]", arg684_1: "f32[1024, 4096]", arg685_1: "f32[1024, 4096]", arg686_1: "f32[4096, 4096]", arg687_1: "f32[8, 4096]", arg688_1: "f32[14336, 4096]", arg689_1: "f32[4096, 14336]", arg690_1: "f32[14336, 4096]", arg691_1: "f32[14336, 4096]", arg692_1: "f32[4096, 14336]", arg693_1: "f32[14336, 4096]", arg694_1: "f32[14336, 4096]", arg695_1: "f32[4096, 14336]", arg696_1: "f32[14336, 4096]", arg697_1: "f32[14336, 4096]", arg698_1: "f32[4096, 14336]", arg699_1: "f32[14336, 4096]", arg700_1: "f32[14336, 4096]", arg701_1: "f32[4096, 14336]", arg702_1: "f32[14336, 4096]", arg703_1: "f32[14336, 4096]", arg704_1: "f32[4096, 14336]", arg705_1: "f32[14336, 4096]", arg706_1: "f32[14336, 4096]", arg707_1: "f32[4096, 14336]", arg708_1: "f32[14336, 4096]", arg709_1: "f32[14336, 4096]", arg710_1: "f32[4096, 14336]", arg711_1: "f32[14336, 4096]", arg712_1: "f32[4096]", arg713_1: "f32[4096]", arg714_1: "f32[4096, 4096]", arg715_1: "f32[1024, 4096]", arg716_1: "f32[1024, 4096]", arg717_1: "f32[4096, 4096]", arg718_1: "f32[8, 4096]", arg719_1: "f32[14336, 4096]", arg720_1: "f32[4096, 14336]", arg721_1: "f32[14336, 4096]", arg722_1: "f32[14336, 4096]", arg723_1: "f32[4096, 14336]", arg724_1: "f32[14336, 4096]", arg725_1: "f32[14336, 4096]", arg726_1: "f32[4096, 14336]", arg727_1: "f32[14336, 4096]", arg728_1: "f32[14336, 4096]", arg729_1: "f32[4096, 14336]", arg730_1: "f32[14336, 4096]", arg731_1: "f32[14336, 4096]", arg732_1: "f32[4096, 14336]", arg733_1: "f32[14336, 4096]", arg734_1: "f32[14336, 4096]", arg735_1: "f32[4096, 14336]", arg736_1: "f32[14336, 4096]", arg737_1: "f32[14336, 4096]", arg738_1: "f32[4096, 14336]", arg739_1: "f32[14336, 4096]", arg740_1: "f32[14336, 4096]", arg741_1: "f32[4096, 14336]", arg742_1: "f32[14336, 4096]", arg743_1: "f32[4096]", arg744_1: "f32[4096]", arg745_1: "f32[4096, 4096]", arg746_1: "f32[1024, 4096]", arg747_1: "f32[1024, 4096]", arg748_1: "f32[4096, 4096]", arg749_1: "f32[8, 4096]", arg750_1: "f32[14336, 4096]", arg751_1: "f32[4096, 14336]", arg752_1: "f32[14336, 4096]", arg753_1: "f32[14336, 4096]", arg754_1: "f32[4096, 14336]", arg755_1: "f32[14336, 4096]", arg756_1: "f32[14336, 4096]", arg757_1: "f32[4096, 14336]", arg758_1: "f32[14336, 4096]", arg759_1: "f32[14336, 4096]", arg760_1: "f32[4096, 14336]", arg761_1: "f32[14336, 4096]", arg762_1: "f32[14336, 4096]", arg763_1: "f32[4096, 14336]", arg764_1: "f32[14336, 4096]", arg765_1: "f32[14336, 4096]", arg766_1: "f32[4096, 14336]", arg767_1: "f32[14336, 4096]", arg768_1: "f32[14336, 4096]", arg769_1: "f32[4096, 14336]", arg770_1: "f32[14336, 4096]", arg771_1: "f32[14336, 4096]", arg772_1: "f32[4096, 14336]", arg773_1: "f32[14336, 4096]", arg774_1: "f32[4096]", arg775_1: "f32[4096]", arg776_1: "f32[4096, 4096]", arg777_1: "f32[1024, 4096]", arg778_1: "f32[1024, 4096]", arg779_1: "f32[4096, 4096]", arg780_1: "f32[8, 4096]", arg781_1: "f32[14336, 4096]", arg782_1: "f32[4096, 14336]", arg783_1: "f32[14336, 4096]", arg784_1: "f32[14336, 4096]", arg785_1: "f32[4096, 14336]", arg786_1: "f32[14336, 4096]", arg787_1: "f32[14336, 4096]", arg788_1: "f32[4096, 14336]", arg789_1: "f32[14336, 4096]", arg790_1: "f32[14336, 4096]", arg791_1: "f32[4096, 14336]", arg792_1: "f32[14336, 4096]", arg793_1: "f32[14336, 4096]", arg794_1: "f32[4096, 14336]", arg795_1: "f32[14336, 4096]", arg796_1: "f32[14336, 4096]", arg797_1: "f32[4096, 14336]", arg798_1: "f32[14336, 4096]", arg799_1: "f32[14336, 4096]", arg800_1: "f32[4096, 14336]", arg801_1: "f32[14336, 4096]", arg802_1: "f32[14336, 4096]", arg803_1: "f32[4096, 14336]", arg804_1: "f32[14336, 4096]", arg805_1: "f32[4096]", arg806_1: "f32[4096]", arg807_1: "f32[4096, 4096]", arg808_1: "f32[1024, 4096]", arg809_1: "f32[1024, 4096]", arg810_1: "f32[4096, 4096]", arg811_1: "f32[8, 4096]", arg812_1: "f32[14336, 4096]", arg813_1: "f32[4096, 14336]", arg814_1: "f32[14336, 4096]", arg815_1: "f32[14336, 4096]", arg816_1: "f32[4096, 14336]", arg817_1: "f32[14336, 4096]", arg818_1: "f32[14336, 4096]", arg819_1: "f32[4096, 14336]", arg820_1: "f32[14336, 4096]", arg821_1: "f32[14336, 4096]", arg822_1: "f32[4096, 14336]", arg823_1: "f32[14336, 4096]", arg824_1: "f32[14336, 4096]", arg825_1: "f32[4096, 14336]", arg826_1: "f32[14336, 4096]", arg827_1: "f32[14336, 4096]", arg828_1: "f32[4096, 14336]", arg829_1: "f32[14336, 4096]", arg830_1: "f32[14336, 4096]", arg831_1: "f32[4096, 14336]", arg832_1: "f32[14336, 4096]", arg833_1: "f32[14336, 4096]", arg834_1: "f32[4096, 14336]", arg835_1: "f32[14336, 4096]", arg836_1: "f32[4096]", arg837_1: "f32[4096]", arg838_1: "f32[4096, 4096]", arg839_1: "f32[1024, 4096]", arg840_1: "f32[1024, 4096]", arg841_1: "f32[4096, 4096]", arg842_1: "f32[8, 4096]", arg843_1: "f32[14336, 4096]", arg844_1: "f32[4096, 14336]", arg845_1: "f32[14336, 4096]", arg846_1: "f32[14336, 4096]", arg847_1: "f32[4096, 14336]", arg848_1: "f32[14336, 4096]", arg849_1: "f32[14336, 4096]", arg850_1: "f32[4096, 14336]", arg851_1: "f32[14336, 4096]", arg852_1: "f32[14336, 4096]", arg853_1: "f32[4096, 14336]", arg854_1: "f32[14336, 4096]", arg855_1: "f32[14336, 4096]", arg856_1: "f32[4096, 14336]", arg857_1: "f32[14336, 4096]", arg858_1: "f32[14336, 4096]", arg859_1: "f32[4096, 14336]", arg860_1: "f32[14336, 4096]", arg861_1: "f32[14336, 4096]", arg862_1: "f32[4096, 14336]", arg863_1: "f32[14336, 4096]", arg864_1: "f32[14336, 4096]", arg865_1: "f32[4096, 14336]", arg866_1: "f32[14336, 4096]", arg867_1: "f32[4096]", arg868_1: "f32[4096]", arg869_1: "f32[4096, 4096]", arg870_1: "f32[1024, 4096]", arg871_1: "f32[1024, 4096]", arg872_1: "f32[4096, 4096]", arg873_1: "f32[8, 4096]", arg874_1: "f32[14336, 4096]", arg875_1: "f32[4096, 14336]", arg876_1: "f32[14336, 4096]", arg877_1: "f32[14336, 4096]", arg878_1: "f32[4096, 14336]", arg879_1: "f32[14336, 4096]", arg880_1: "f32[14336, 4096]", arg881_1: "f32[4096, 14336]", arg882_1: "f32[14336, 4096]", arg883_1: "f32[14336, 4096]", arg884_1: "f32[4096, 14336]", arg885_1: "f32[14336, 4096]", arg886_1: "f32[14336, 4096]", arg887_1: "f32[4096, 14336]", arg888_1: "f32[14336, 4096]", arg889_1: "f32[14336, 4096]", arg890_1: "f32[4096, 14336]", arg891_1: "f32[14336, 4096]", arg892_1: "f32[14336, 4096]", arg893_1: "f32[4096, 14336]", arg894_1: "f32[14336, 4096]", arg895_1: "f32[14336, 4096]", arg896_1: "f32[4096, 14336]", arg897_1: "f32[14336, 4096]", arg898_1: "f32[4096]", arg899_1: "f32[4096]", arg900_1: "f32[4096, 4096]", arg901_1: "f32[1024, 4096]", arg902_1: "f32[1024, 4096]", arg903_1: "f32[4096, 4096]", arg904_1: "f32[8, 4096]", arg905_1: "f32[14336, 4096]", arg906_1: "f32[4096, 14336]", arg907_1: "f32[14336, 4096]", arg908_1: "f32[14336, 4096]", arg909_1: "f32[4096, 14336]", arg910_1: "f32[14336, 4096]", arg911_1: "f32[14336, 4096]", arg912_1: "f32[4096, 14336]", arg913_1: "f32[14336, 4096]", arg914_1: "f32[14336, 4096]", arg915_1: "f32[4096, 14336]", arg916_1: "f32[14336, 4096]", arg917_1: "f32[14336, 4096]", arg918_1: "f32[4096, 14336]", arg919_1: "f32[14336, 4096]", arg920_1: "f32[14336, 4096]", arg921_1: "f32[4096, 14336]", arg922_1: "f32[14336, 4096]", arg923_1: "f32[14336, 4096]", arg924_1: "f32[4096, 14336]", arg925_1: "f32[14336, 4096]", arg926_1: "f32[14336, 4096]", arg927_1: "f32[4096, 14336]", arg928_1: "f32[14336, 4096]", arg929_1: "f32[4096]", arg930_1: "f32[4096]", arg931_1: "f32[4096, 4096]", arg932_1: "f32[1024, 4096]", arg933_1: "f32[1024, 4096]", arg934_1: "f32[4096, 4096]", arg935_1: "f32[8, 4096]", arg936_1: "f32[14336, 4096]", arg937_1: "f32[4096, 14336]", arg938_1: "f32[14336, 4096]", arg939_1: "f32[14336, 4096]", arg940_1: "f32[4096, 14336]", arg941_1: "f32[14336, 4096]", arg942_1: "f32[14336, 4096]", arg943_1: "f32[4096, 14336]", arg944_1: "f32[14336, 4096]", arg945_1: "f32[14336, 4096]", arg946_1: "f32[4096, 14336]", arg947_1: "f32[14336, 4096]", arg948_1: "f32[14336, 4096]", arg949_1: "f32[4096, 14336]", arg950_1: "f32[14336, 4096]", arg951_1: "f32[14336, 4096]", arg952_1: "f32[4096, 14336]", arg953_1: "f32[14336, 4096]", arg954_1: "f32[14336, 4096]", arg955_1: "f32[4096, 14336]", arg956_1: "f32[14336, 4096]", arg957_1: "f32[14336, 4096]", arg958_1: "f32[4096, 14336]", arg959_1: "f32[14336, 4096]", arg960_1: "f32[4096]", arg961_1: "f32[4096]", arg962_1: "f32[4096, 4096]", arg963_1: "f32[1024, 4096]", arg964_1: "f32[1024, 4096]", arg965_1: "f32[4096, 4096]", arg966_1: "f32[8, 4096]", arg967_1: "f32[14336, 4096]", arg968_1: "f32[4096, 14336]", arg969_1: "f32[14336, 4096]", arg970_1: "f32[14336, 4096]", arg971_1: "f32[4096, 14336]", arg972_1: "f32[14336, 4096]", arg973_1: "f32[14336, 4096]", arg974_1: "f32[4096, 14336]", arg975_1: "f32[14336, 4096]", arg976_1: "f32[14336, 4096]", arg977_1: "f32[4096, 14336]", arg978_1: "f32[14336, 4096]", arg979_1: "f32[14336, 4096]", arg980_1: "f32[4096, 14336]", arg981_1: "f32[14336, 4096]", arg982_1: "f32[14336, 4096]", arg983_1: "f32[4096, 14336]", arg984_1: "f32[14336, 4096]", arg985_1: "f32[14336, 4096]", arg986_1: "f32[4096, 14336]", arg987_1: "f32[14336, 4096]", arg988_1: "f32[14336, 4096]", arg989_1: "f32[4096, 14336]", arg990_1: "f32[14336, 4096]", arg991_1: "f32[4096]", arg992_1: "f32[4096]", arg993_1: "f32[4096]", arg994_1: "f32[32000, 4096]", arg995_1: "f32[64]", arg996_1: "i32[2, 4]", arg997_1: "i32[2, 4]"):
     # File: /usr/local/lib/python3.12/dist-packages/torch/nn/modules/sparse.py:190 in forward, code: return F.embedding(
    embedding: "f32[2, 4, 4096]" = torch.ops.aten.embedding.default(arg0_1, arg996_1);  arg0_1 = arg996_1 = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:622 in forward, code: cache_position = torch.arange(
    arange: "i64[4]" = torch.ops.aten.arange.start(0, 4, device = device(type='meta'), pin_memory = False)

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:626 in forward, code: position_ids = cache_position.unsqueeze(0)
    unsqueeze: "i64[1, 4]" = torch.ops.aten.unsqueeze.default(arange, 0)

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:628 in forward, code: causal_mask = self._update_causal_mask(
    full: "f32[4, 4]" = torch.ops.aten.full.default([4, 4], -3.4028234663852886e+38, dtype = torch.float32, device = device(type='meta'), pin_memory = False)
    arange_1: "i64[4]" = torch.ops.aten.arange.default(4, device = device(type='meta'), pin_memory = False)
    reshape: "i64[4, 1]" = torch.ops.aten.reshape.default(arange, [-1, 1]);  arange = None
    gt: "b8[4, 4]" = torch.ops.aten.gt.Tensor(arange_1, reshape);  arange_1 = reshape = None
    mul_: "f32[4, 4]" = torch.ops.aten.mul_.Tensor(full, gt);  full = gt = None
    unsqueeze_1: "f32[1, 4, 4]" = torch.ops.aten.unsqueeze.default(mul_, 0);  mul_ = None
    unsqueeze_2: "f32[1, 1, 4, 4]" = torch.ops.aten.unsqueeze.default(unsqueeze_1, 1);  unsqueeze_1 = None
    slice_1: "f32[1, 1, 4, 4]" = torch.ops.aten.slice.Tensor(unsqueeze_2, 2, 0, 9223372036854775807);  unsqueeze_2 = None
    slice_2: "f32[1, 1, 4, 4]" = torch.ops.aten.slice.Tensor(slice_1, 3, 0, 9223372036854775807);  slice_1 = None
    expand: "f32[2, 1, 4, 4]" = torch.ops.aten.expand.default(slice_2, [2, 1, -1, -1]);  slice_2 = None
    clone: "f32[2, 1, 4, 4]" = torch.ops.aten.clone.default(expand);  expand = None
    slice_3: "f32[2, 1, 4, 4]" = torch.ops.aten.slice.Tensor(clone, 0, 0, 9223372036854775807)
    slice_4: "f32[2, 1, 4, 4]" = torch.ops.aten.slice.Tensor(slice_3, 1, 0, 9223372036854775807);  slice_3 = None
    slice_5: "f32[2, 1, 4, 4]" = torch.ops.aten.slice.Tensor(slice_4, 2, 0, 9223372036854775807);  slice_4 = None
    slice_6: "i32[2, 4]" = torch.ops.aten.slice.Tensor(arg997_1, 0, 0, 9223372036854775807);  arg997_1 = None
    unsqueeze_3: "i32[2, 1, 4]" = torch.ops.aten.unsqueeze.default(slice_6, 1);  slice_6 = None
    unsqueeze_4: "i32[2, 1, 1, 4]" = torch.ops.aten.unsqueeze.default(unsqueeze_3, 2);  unsqueeze_3 = None
    slice_7: "i32[2, 1, 1, 4]" = torch.ops.aten.slice.Tensor(unsqueeze_4, 3, 0, 9223372036854775807);  unsqueeze_4 = None
    to: "i32[2, 1, 1, 4]" = torch.ops.aten.to.dtype_layout(slice_7, dtype = torch.int32, layout = torch.strided, device = device(type='meta'));  slice_7 = None
    add: "f32[2, 1, 4, 4]" = torch.ops.aten.add.Tensor(slice_5, to);  slice_5 = to = None
    eq: "b8[2, 1, 4, 4]" = torch.ops.aten.eq.Scalar(add, 0);  add = None
    slice_8: "f32[2, 1, 4, 4]" = torch.ops.aten.slice.Tensor(clone, 0, 0, 9223372036854775807)
    slice_9: "f32[2, 1, 4, 4]" = torch.ops.aten.slice.Tensor(slice_8, 1, 0, 9223372036854775807);  slice_8 = None
    slice_10: "f32[2, 1, 4, 4]" = torch.ops.aten.slice.Tensor(slice_9, 2, 0, 9223372036854775807);  slice_9 = None
    masked_fill: "f32[2, 1, 4, 4]" = torch.ops.aten.masked_fill.Scalar(slice_10, eq, -3.4028234663852886e+38);  slice_10 = eq = None
    slice_11: "f32[2, 1, 4, 4]" = torch.ops.aten.slice.Tensor(clone, 0, 0, 9223372036854775807)
    slice_12: "f32[2, 1, 4, 4]" = torch.ops.aten.slice.Tensor(slice_11, 1, 0, 9223372036854775807);  slice_11 = None
    slice_13: "f32[2, 1, 4, 4]" = torch.ops.aten.slice.Tensor(slice_12, 2, 0, 9223372036854775807);  slice_12 = None
    copy_: "f32[2, 1, 4, 4]" = torch.ops.aten.copy_.default(slice_13, masked_fill);  slice_13 = masked_fill = copy_ = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:635 in forward, code: position_embeddings = self.rotary_emb(hidden_states, position_ids)
    _set_grad_enabled = torch._C._set_grad_enabled(False);  _set_grad_enabled = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:418 in forward, code: inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
    unsqueeze_5: "f32[1, 64]" = torch.ops.aten.unsqueeze.default(arg995_1, 0);  arg995_1 = None
    slice_14: "f32[1, 64]" = torch.ops.aten.slice.Tensor(unsqueeze_5, 1, 0, 9223372036854775807);  unsqueeze_5 = None
    unsqueeze_6: "f32[1, 64, 1]" = torch.ops.aten.unsqueeze.default(slice_14, 2);  slice_14 = None
    to_1: "f32[1, 64, 1]" = torch.ops.aten.to.dtype(unsqueeze_6, torch.float32);  unsqueeze_6 = None
    expand_1: "f32[1, 64, 1]" = torch.ops.aten.expand.default(to_1, [1, -1, 1]);  to_1 = None
    to_2: "f32[1, 64, 1]" = torch.ops.aten.to.dtype_layout(expand_1, dtype = torch.float32, layout = torch.strided, device = device(type='meta'));  expand_1 = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:419 in forward, code: position_ids_expanded = position_ids[:, None, :].float()
    slice_15: "i64[1, 4]" = torch.ops.aten.slice.Tensor(unsqueeze, 0, 0, 9223372036854775807);  unsqueeze = None
    unsqueeze_7: "i64[1, 1, 4]" = torch.ops.aten.unsqueeze.default(slice_15, 1);  slice_15 = None
    slice_16: "i64[1, 1, 4]" = torch.ops.aten.slice.Tensor(unsqueeze_7, 2, 0, 9223372036854775807);  unsqueeze_7 = None
    to_3: "f32[1, 1, 4]" = torch.ops.aten.to.dtype(slice_16, torch.float32);  slice_16 = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:423 in forward, code: freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
    to_4: "f32[1, 64, 1]" = torch.ops.aten.to.dtype(to_2, torch.float32);  to_2 = None
    to_5: "f32[1, 1, 4]" = torch.ops.aten.to.dtype(to_3, torch.float32);  to_3 = None
    matmul: "f32[1, 64, 4]" = torch.ops.aten.matmul.default(to_4, to_5);  to_4 = to_5 = None
    transpose: "f32[1, 4, 64]" = torch.ops.aten.transpose.int(matmul, 1, 2);  matmul = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:424 in forward, code: emb = torch.cat((freqs, freqs), dim=-1)
    cat: "f32[1, 4, 128]" = torch.ops.aten.cat.default([transpose, transpose], -1);  transpose = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:425 in forward, code: cos = emb.cos() * self.attention_scaling
    cos: "f32[1, 4, 128]" = torch.ops.aten.cos.default(cat)
    mul: "f32[1, 4, 128]" = torch.ops.aten.mul.Tensor(cos, 1.0);  cos = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:426 in forward, code: sin = emb.sin() * self.attention_scaling
    sin: "f32[1, 4, 128]" = torch.ops.aten.sin.default(cat);  cat = None
    mul_1: "f32[1, 4, 128]" = torch.ops.aten.mul.Tensor(sin, 1.0);  sin = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:428 in forward, code: return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
    to_6: "f32[1, 4, 128]" = torch.ops.aten.to.dtype(mul, torch.float32);  mul = None
    to_7: "f32[1, 4, 128]" = torch.ops.aten.to.dtype(mul_1, torch.float32);  mul_1 = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:635 in forward, code: position_embeddings = self.rotary_emb(hidden_states, position_ids)
    _set_grad_enabled_1 = torch._C._set_grad_enabled(True);  _set_grad_enabled_1 = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:167 in forward, code: hidden_states = hidden_states.to(torch.float32)
    to_8: "f32[2, 4, 4096]" = torch.ops.aten.to.dtype(embedding, torch.float32);  embedding = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:168 in forward, code: variance = hidden_states.pow(2).mean(-1, keepdim=True)
    pow_1: "f32[2, 4, 4096]" = torch.ops.aten.pow.Tensor_Scalar(to_8, 2)
    mean: "f32[2, 4, 1]" = torch.ops.aten.mean.dim(pow_1, [-1], True);  pow_1 = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:169 in forward, code: hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
    add_1: "f32[2, 4, 1]" = torch.ops.aten.add.Tensor(mean, 1e-05);  mean = None
    rsqrt: "f32[2, 4, 1]" = torch.ops.aten.rsqrt.default(add_1);  add_1 = None
    mul_2: "f32[2, 4, 4096]" = torch.ops.aten.mul.Tensor(to_8, rsqrt);  rsqrt = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:170 in forward, code: return self.weight * hidden_states.to(input_dtype)
    to_9: "f32[2, 4, 4096]" = torch.ops.aten.to.dtype(mul_2, torch.float32);  mul_2 = None
    mul_3: "f32[2, 4, 4096]" = torch.ops.aten.mul.Tensor(arg30_1, to_9);  arg30_1 = to_9 = None

     # File: /usr/local/lib/python3.12/dist-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
    linear: "f32[2, 4, 4096]" = torch.ops.aten.linear.default(mul_3, arg1_1);  arg1_1 = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:277 in forward, code: query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    view: "f32[2, 4, 32, 128]" = torch.ops.aten.view.default(linear, [2, 4, -1, 128]);  linear = None
    transpose_1: "f32[2, 32, 4, 128]" = torch.ops.aten.transpose.int(view, 1, 2);  view = None

     # File: /usr/local/lib/python3.12/dist-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
    linear_1: "f32[2, 4, 1024]" = torch.ops.aten.linear.default(mul_3, arg2_1);  arg2_1 = None

     # File: /home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py:278 in forward, code: key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    view_1: "f32[2, 4, 8, 128]" = torch.ops.aten.view.default(linear_1, [2, 4, -1, 128]);  linear_1 = None
    transpose_2: "f32[2, 8, 4, 128]" = torch.ops.aten.transpose.int(view_1, 1, 2);  view_1 = None

     # File: /usr/local/lib/python3.12/dist-packages/torch/nn/modules/linear.py:125 in forward, code: return F.linear(input, self.weight, self.bias)
    linear_2: "f32[2, 4, 1024]" = torch.ops.aten.linear.default(mul_3, arg3_1);  mul_3 = arg3_1 = None

......

I 
......
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 1379, in __torch_dispatch__
    return proxy_call(self, func, self.pre_dispatch, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/proxy_tensor.py", line 914, in proxy_call
    out = func(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 756, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_stats.py", line 27, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1282, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1823, in dispatch
    return self._cached_dispatch_impl(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1393, in _cached_dispatch_impl
    output = self._dispatch_impl(func, types, args, kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 2333, in _dispatch_impl
    decomposition_table[func](*args, **kwargs)
  File "/usr/local/lib/python3.12/dist-packages/torch/_refs/__init__.py", line 4002, in unbind
    torch.squeeze(s, dim) for s in torch.tensor_split(t, t.shape[dim], dim)
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_stats.py", line 27, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1282, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1823, in dispatch
    return self._cached_dispatch_impl(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1393, in _cached_dispatch_impl
    output = self._dispatch_impl(func, types, args, kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 2338, in _dispatch_impl
    r = func.decompose(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 799, in decompose
    return self._op_dk(dk, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/sym_node.py", line 500, in guard_int
    r = self.evaluate()
        ^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/sym_node.py", line 494, in evaluate
    return self.shape_env.evaluate_sym_node(self, size_oblivious)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/symbolic_shapes.py", line 6637, in evaluate_sym_node
    return self.evaluate_expr(
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/recording.py", line 263, in wrapper
    return retlog(fn(*args, **kwargs))
                  ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/symbolic_shapes.py", line 6653, in evaluate_expr
    return self._evaluate_expr(
           ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/fx/experimental/symbolic_shapes.py", line 6870, in _evaluate_expr
    raise self._make_data_dependent_error(
torch.fx.experimental.symbolic_shapes.GuardOnDataDependentSymNode: Could not extract specialized integer from data-dependent expression u0 (unhinted: u0).  (Size-like symbols: u0)

Caused by: (_ops.py:799 in decompose)
For more information, run with TORCH_LOGS="dynamic"
For extended logs when we create symbols, also add TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="u0"
If you suspect the guard was triggered from C++, add TORCHDYNAMO_EXTENDED_DEBUG_CPP=1
For more debugging help, see https://docs.google.com/document/d/1HSuTTVvYH1pTew89Rtpeu84Ht3nQEFTYhAX3Ypa_xJs/edit?usp=sharing

For C++ stack trace, run with TORCHDYNAMO_EXTENDED_DEBUG_CPP=1

The following call raised this error:
  File "/home/guomingz/.local/lib/python3.12/site-packages/transformers/models/mixtral/modeling_mixtral.py", line 138, in forward
    expert_hitted = (expert_mask.sum(dim=(-1, -2)) > 0).nonzero(as_tuple=True)[0].tolist()

os: linux
transformer version: 4.52.4
torch version: 2.7
gpu: Nvidia H100
cuda: 12.9
driver: 575.57.08

Who can help?

@ArthurZucker @Coco58323 @Cyrilvallez

Information

  • The official example scripts
  • My own modified scripts

Tasks

  • An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
  • My own task or dataset (give details below)

Reproduction

  1. pip install transformers==4.52.4
  2. run below code snippet
import transformers
import torch.export as te
import torch
from contextlib import nullcontext

torch.autocast = lambda *args, **kwargs: nullcontext()  


mixtral = transformers.AutoModelForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1", 
                                                       device_map="meta",
                                                       )
ep = te.export(mixtral, 
                args=(torch.randint(0, 100, (2, 4),device="meta", dtype=torch.int32),
                      torch.randint(0, 100, (2, 4),device="meta", dtype=torch.int32)
                    ), 
                kwargs={}, strict=False
                ).module()

Expected behavior

The above code snippets could run successfully.

Metadata

Metadata

Assignees

No one assigned

    Labels

    Good Difficult Issuebug

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions