Add files using upload-large-folder tool

LTA_openwebtext_dualt/docs/lta_samples/metrics_20260517/owt_t5_linearsoftkl_latest_step86000_latedebias0p7_uniformtime_dualline_steps128_c1024_n8_noscore/genppl_large.json

@@ -0,0 +1,6 @@
+{
+  "gen_ppl": 2913.532847576075,
+  "gen_nll": 7.97712166078629,
+  "gen_tokens": 8184,
+  "gen_ppl_scorer": "/e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard"
+}

LTA_openwebtext_dualt/docs/lta_samples/metrics_20260517/train8_overfit_20260517_train8_overfit_ode_infer/train8_n1024_hard_ce_onehot_20260517_train8_overfit_none/summary.json

@@ -0,0 +1,55 @@
+[
+  {
+    "checkpoint": "runs/train8_n1024_hard_ce_onehot_20260517_train8_overfit/step_0001000.pt",
+    "ckpt_step": 1000,
+    "max_len": 1024,
+    "decode_rule": "flowmap",
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "steps": 128,
+    "c_min": 1.0,
+    "c_max": 512.0,
+    "model_t_mode": "post",
+    "time_schedule": "logit_normal",
+    "time_logit_mean": -1.5,
+    "time_logit_std": 0.8,
+    "time_power": 2.0,
+    "input_noise_scale": 0.0,
+    "input_noise_until": 1.0,
+    "input_noise_dirichlet_concentration": 1.0,
+    "endpoint_softening": "none",
+    "endpoint_soft_power": 1.0,
+    "endpoint_soft_min_conf": 0.0,
+    "endpoint_soft_max_conf": 1.0,
+    "final_from": "state",
+    "final_decode": "argmax",
+    "final_sample_temp": 1.0,
+    "final_top_k": 0,
+    "final_top_p": 1.0,
+    "early_temp": 1.0,
+    "late_temp": 1.0,
+    "temp_end": 1.0,
+    "temp_power": 1.0,
+    "pos_extend": "repeat",
+    "fixed_first_token_id": null,
+    "fixed_first_token_text": "",
+    "fixed_first_initial_argmax": false,
+    "use_ema": false,
+    "n_samples": 8,
+    "sample_entropy": 3.0918080879465375,
+    "unique_tokens": 177,
+    "token_count": 8192,
+    "distinct_1": 0.0216064453125,
+    "distinct_2": 0.13379765395894427,
+    "top_token_mass": 0.136474609375,
+    "texts_preview": [
+      "the the<|endoftext|><|endoftext|><|endoftext|> the<|endoftext|> the one the of the the of the on the the Bill one the the Bill the the the the the the the the the the the the the the the one the the the the the the the the the theThe the the the the the the the the theThe theThe theThe the itThe theThe theTheThe it itThe itThe the it theThe itThe in it itThe it itThe it it it itThe it it from it it it it it from it in from it it it in from from from it in from from it from it from from from from from from from from from from from from from the the from from from from from from from from from from a from from from from from the from from from the without from from from from the from from it from from from from the from the from from the the from from the from the from the from of the from from from from the without the from the from the the from the from the from the the from many from the the from from from it from from many the many many from from many many from from from many in, to many many many in in from in manyCNN many in many in in in) inCNN), inowsowsows inows, in in,,ows,,ows who,ows,, who,,ows 2016 2016,, who Muellerowsows 2016ows,ows,ows,, in whoowsows inows,ows, Mueller, meeting Tuesdayows, meeting Mueller,,limowsows counselows,,,,owslim, Mueller,ows, like likeows like,lim Mueller report 2016 like like Butows 2016 like 2016 Butlim, Butlim 2016, But, Muellerows But, Mueller to 2016 2016 Mueller 2016 2016 to in But But But Mueller 2016 a 2016 But 2016 But lawyer a a to to 2016 with Jr to to 2016 2016 2016 2016 But to 2016 to 2016 2016 to Mueller Mueller to to Mueller trail 2016 to to to, to 2016 2016 to 2016 while work while to while while could to work and could while while to while work work work politically panel work work about work -- politically work work about about and about about about about about about work about could about about legally about who about who about about about work a who and who who a and who about who work who White who White legally who a whoA a about a who. who who who White who President. .. who White about. about. who. White ... insky. who.. . President the. frominsky.inskyinsky about us's about about about's from's Mueller's about about Friday about's interview about interview,,,,,, Mueller,,,,'s,,,,,,,,,,,'s,,,,,ime,,,,,,,,,,, Mueller,,,,,, Mueller,,,,,,,,,,'s noted,,,,,,'s,,,,, of,'s,,,'s on on,,,, on, on on on on,'s an's testimony's, on's,'s's on's's's on's's's's's's his over's's the's's's's's's's the's's's's's's and's's's day out the's's's the's's's's's's out out's's out's's's's's's's's's the's's's out's's's's's's's's's's's the's's's's's's's's's's's's's's's's the's the's the's's of the's's the's the's's the the of the's the's the's's the, the the the the the the the the the the of the get's of get the the thegggg thegg of,gg of, of in of Russia, should in,,, the of Affairs of that of of of should of, of the, of act thegg of of, Affairs of of of of of of that of Starr of of of of of, of of as not of of, of of know, of of not of of of of of of of not not, was of as as, not of of of was of, not was was of was was was of was as not of was not not of of of was a of was of not was be not was not, of was, was no was was that of have was have was that that is was was not was not that not that that that not is have that not no not was that most not that is that that that that is have is not is that is is is that that is that is ( is that not that whether is is is (",
+      "the<|endoftext|><|endoftext|> the the<|endoftext|> Today the Bill the the<|endoftext|> the of the the on the the the the the the the the the the the the the the the the the the the the the the the the of the the the the the the the of the theThe theTheThe theThe the the theThe itThe theTheThe the theThe theThe theThe itThe theThe it itTheTheThe it the itThe it it from it it it it from it it it it it from it it in was it it from it it in in from from it it from it it from in from from from from from it it it from from from from it from from from from from from from from from from from the from from from from from from from the the from the from from from from from Brooklyn from from in without from from from from from from from from<|endoftext|> the from the from from from the in from from the from from the from from the from without from from from from the from from from the from from from from from from the from from from from many from from from from from in in many the many from it from many many many many many many many in manyCNN in in to it in in in inCNN many in it in, Donald comes,,ows, Mueller,,,owsows,,,ows,,, had,,ows,, 2016owsows 2016 Mueller on,ows 2016 2016,, feel whoows 2016,owsowsows,,ows,,,, like,owsows who,, 2016 who,ows, like like like meeting But, to Butows who, like,lim But, likeows like on But meeting Butlim 2016 But on, 2016 Butlim But,, 2016 on to But Mueller to But 2016 But the But to on, like But Mueller But with 2016 But Mueller 2016 like to 2016 that to on on to Mueller to on But some Mueller 2016 Mueller to on to while to 2016 to 2016 very to very to Mueller some to while to while to who to while to to a a to while who some who to who work work while a work work and work while a and work some to who work a while and about work about about about who some about a work about and about and a about who who about who a about who about who about legally a work about about about who about about politically who -- at who who who who about White who who who who who at at who who submitted about White White who who who about who White about who about.. White .. who about. insky... President.insky.insky who... aboutinskyinsky about aboutinsky from'sinskyinsky's about Mueller about, from Friday about interview about interview Mueller's, about,,,,,,,,,,,,,,,,,, But,,'s,,, danger,,,,,ly,,,,,,,,,,,,,,,,,,,,,,,,, Mueller's,, the,,,,,, Mueller,,,, Mueller,,, on Mueller,'s,,,,'s on on's, his's,'s's's his Starr of on a's's's's's his's's's's's's's's's's's's's's the's's's's's's,'s's's's's's's's's's and's's over's's's's's's's's's's's's's's's's the's's's's's's's's's's's's's's's's's's's's's's's's's's's the's's They's's's's the's's's's the's the the's's's the's's's the's the's the's the that hired the the's the the the's the's of the hired the's of of Russia the the, of of the the the of,'s,, the the thegg of, of of,gg of not that the of, of the of of powers of that of of not of of, of by of of of of of of of, of of of of of of ( not of not ( of of of not ( of ( of of of ( of of of of ( of not of of of of of of of of of of sex not of of of of as not of of of not, not not of sex was not was of not of not is was not be not be, not is not, not was not of that, was not, have be not that be was that not was have no was on was, that of the not a to is not that is on that ( is is that not not is not one not on is that is is is most that is that is is is is that is is that that is ( that (",
+      "<|endoftext|> the<|endoftext|><|endoftext|>�<|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|><|endoftext|> the<|endoftext|> the<|endoftext|> the<|endoftext|> the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the the out the the Brooklyn the the the the the the the of�� Brooklyn Brooklyn plan the the thes the it Brooklyn inThe the it Brooklyn it it in it it it it it it it it it it it it it it it� it it it it it it it it it it it in it it from it it it it it it it it it it it it it it it it it it it it was it it it from from from from it it from from the it it from from it from from� it from from it it from it the from the from from from from from it from from from from it from from from from from it from it from from from it from from from from from from it from from from from the from from the from from from from from concept from from from from from from from from from from from from from from from from from from from from into from from from from to the from in many in from into into to from many the in,, in in in in,,,,,,,,,,,,,ows, the,,,,,,,,,,,,,,,, to,,,,,,,, in,,,,,,, on,,,,, in,, in,,, like,lim,,,, But,,,,, in, like,,,,, on like like ButWhat who,, like But, But,, on on, on But But, on on and on on on on on But on, on But on on, But on But But on, with on on on on, on, to on on with with on on with while that of Mueller the some on to on to work will work work no while some some work a to some a work work work work work work about work work in work work a work work work work work work work and work work work a work work some work work work work a work work work a work very work work work a work work work, work work work work work as some work a as work who work work by a who work House White who . ch. . from . .. President. . from from.. from the from from from frominsky from from from from from from from an from an, from from,, that, Mueller,, Mueller's,,,,,,,,,,,,,,,,,,,,, of,, But,,,,,,, But,,,,,, of,,,,,,,,,,,,,,,,, of,,,,, an, of an's,,,,,, a's,,, .,, on's on the a took a on took his a's his over a a a a's White's his a's out out White White's out out White out over over's out's's out over's out out over out over's over and over over over out out out out's over the's out over out out out over out's the the out's's the out the out the's the the's the the the the's the the's's the the the the the the the the the's the the was the's the the the the the because the the the's's the the the's the the the the the the the's the the the the the states the the the the the the the the the the the, the the of the, of the the the a of the of the of the of the the of of of of the the the of of of of of the of of of of of ( of of of ofex a of ( of ( a ofex the of that of of of of of of ofAY ( of of of ( a of of of ( of ( of large sex ( of of ( of ( of not.\" not of of, of ( of was of of of of in, not of of the of of, was, of not of of that that be be not was, not of of problem of, on was was not that was was was was be be not was is that is not that is was that.\" was be not not \" was was that, that that that that that that ( that not not that ( is that that is that is that that that that a is that that ( that that that that that that ( that",
+      "the the<|endoftext|> that,<|endoftext|><|endoftext|> the that<|endoftext|><|endoftext|> the<|endoftext|><|endoftext|> the<|endoftext|><|endoftext|><|endoftext|> the the that the<|endoftext|> the the that that<|endoftext|> the the the that that that the the that that the that the that the the that the the that the the out. the. out of the the the out that that the the in. the that. in., the. in of the.. that for. it in. it it the of of is. it elect the of the is it it is it it it is in it it it it it is it to it is it it it it it it it it it it it it it it it it it is it it an it from from from from from and from from it from from from it from from from from from from it from from from it from from from from from from from from and from from from from from from from from from from from threat from from from from from from from from is in from from from from from from from from from in from from from from from from from from to from from from from from from from from from from from from is into to into to to in is to to to into, into, into by into to from,,, the, in, going,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,, on, But,,,,,,, on on on, on,, to of of, of,, on on on But of, on of on on on, of on on But on, on on of, on on on of of on of of of of of of of, on of but on of of on on on of the on of of on on on of on on on on of of of on of on of of on of of of of of, of of on on on of on of on on on on their on safe no as with on work on no no no work work work as had had work as work as work work work work work work work work work work as as work work as and work work work work work work work work that work work work work work ing work work work work work work work work work roughly work many many work work work from work work very very work many from it from from from from work from from from work. from from from from from. from from from from from from from from from from from from from from from from from from from from from from from from from from an an to an an an from an from an an an an from But But an an an an an new an of an and an of an of of of of of an an of of and an an, of an Ken and it of of of of and and of Ken message,, a a, and and of of an, of of and of and message it.,., it of and over, a. a a. a message a a a. and and a. the and. a. a and a.. a , a a a was a a.. a a the. a. a a,...,.,..,. a.. a.... the over the over in , over,, over the over over,, over the the the,, to in the, the, the the the the, the the the the the the the the, the the the the the, the, the the to the the the the in the the the the the the, the the the, the the the the the the the the the the the the the the the the the will the the the the the the will will the the the the the the the been the the that will the the the the the the the the the the the the the the the that the the the of of to the same of a a do the like a a a a that that the that to of a a of a a that a the a a a a a a a a a a a a a a a a a a a a of of a of of a a a a a a of a a of a a of a a a of a of a a a a a of free a of a a a a a a a a a a of a , of a, a, a a a a a a a the a a a a member a ar a a a a a of a a a are the member a a of a a a a a"
+    ],
+    "gen_ppl": 29.86571106300112,
+    "gen_nll": 3.3967110351424585,
+    "gen_tokens": 7081
+  }
+]

LTA_openwebtext_dualt/docs/lta_samples/metrics_20260517/train8_overfit_20260517_train8_overfit_ode_infer/train8_n1024_linear_soft_kl_onehot_20260517_train8_overfit_uniform/context1024_samples.txt

@@ -0,0 +1,29 @@
+the the a...... the on to.. in with. . strategy. in to in been. in. in the in .... the. .. confront... in in to. . the of the a. the the.. . in in . a. . is .. the . in in in in. in in.. mark.s in a� leaving of the A in the in victims in.. in in-... he the. a the of night the in the in. been in., the a . really the in. in the in in. AN in.. been thiss. . in.... the in in�. . the exp nobody a of.. in of.. its in in a in� inm . in. again in. the the the in in to in. –. been. the a in in. trusted- in a in .. in the theAlso a in the in the ,. the the in in . . in.. night in. and the in . in in the in. in to in a in..... . in the initer start in.. the best. when Tuesday. to otherwise . a the in in in in in .....s..... . the reportball, the in as to in. .. of.. in in the. the toities. and a . in the. in the the ref... in and, a. in the- .. in.� to in in in,. a-. to been in in - in.. . to to the . in. . the . to ., in. X the . and margin. in. the is.zip . in tragedy culp in . in. to the in leading .. in in that in . in in a the State in. in a to a the in people. in the. up in in in day� in . in team people. in matchup a. in the reporters a and. of in the in in in people.. a. in the in and. to been. the.�. the- a. to in in in the. . precincts not in.... the the as- the of. to... a. in.. shenanigans a. in . in.. the.. a in- in a , in. in . theint in to, the in..... iniel the ? Hague licence the to the. to . the the in and? the in. the.. the in to in. a. in. the . to. the he the in in the . the the in the ... to. the . on . secretary and. the the Pl,,.s to the. in last the day was on. that in away in the. a been the a. the of the pleaded in ör. and . , the Affairs.� in source been the path to to. 40 stance to to the the. this on the out to in. the and . in users the the the to the in the to. the that... that the this to the?. in. to of the to. the this in. area the.. the this to the the to the the this to. of is and the theich. in. to thes the. the. the. that the . of in.. � this a in largely theier to the.. to this this this the in expects in the the�.. a this to. to the dimension Clinton. to in.. to the the to in in to. this know. this of. that in . the this the Clinton this protracted this the. if this that . the s in this to to. to..?.. .. thisrals the the in to that. people. .. the the... this.. the. in this this. in to in. the this to that.. the and the. this?.. absorb to .... . in. Tuesday a to? a Clinton . to. the the. to in people the
+
+---
+
+the the... ..... in.. in.. irony .- . . the. the in.. investigation .. of in. the to Share. in the .. . the. in. in.. Ch been. a - . the.. ait.... in in in. November in . a immigrants. in interrogation. in of on� in the in Lew in their inbe t in been in� was. the of in. been in. in ins . in... in.. in and in Hal.. the... in the. the aanted . .<|endoftext|> in the the the. the to the- with. in in in,. a. in. in in in.. the.. . the. on in in in in the . to a in in. . a night More field in a in to respectful. in.s. in in was the in,.. to.. was. in to in. in. the generators in was-. . in power were on the . to the on.. in on on the the. the.. in asked in the- the.ible. a to in. in a red in in in in in. a... in in. on cannon. in in in, in. to. tos in. in million. the in, recount not . a, in a. .. the the.. in in century.�. whose ... in in and a. in grasping.. the in.. in a-. to. in the in. the to to the in. the. in that. in. potentially this ofelection the in.. in in a was.,.. inishing to. in. Sanders right to. a a in. the to in to... the a to the a.. to in. a. in in to in. in in in.� in in. in.. no to in... James . warning that that to . in in to.. in. a. in Char day in.. at of in. on this. the,. in in in substantially inso in.. a in supposedly been in that. a. . the . to. to the in in. to. to. health. dimension. a the in.. this this to the he the the in toird Clinton the sought . to Haiti the evidence fighting- to More Bill in the. the. in. the. to . in. in in to in Frankfurt. been Free to. to in to . . rival the the this....... this in and. supplies. in a in. access of to $ a in �. to last and the in a the . of sten. the country . in the the tend that to... to to. fits and and this on in.. this. the to-handlections. in this this. in. the. the a 1999. this the push. the the to. to the the to. to Kinn that the They. . to a the a to this the in to this in to.. the.. this to that and at to in. and and want out the this . in this . this. to to to of? 146 the.. this. in, the key a this in this this the. that prizes to in this this this and package the this the to this to on to and by the the adequate. on the this the to this this. the the. the intention the to that. this. this the. the the this the in been.. to to the this to.. to this in been� to to . Michigan to the... that this of a to to. the the this that that cause this. this this the this this and...s...,� in the the. understand the .akens to to to in. the. month to the been this to people. . disease to one... the. to you to to been to. . in?? been people .
+
+---
+
+. predominantly.��� in. in in. $ war. in wide.. to. in in.. was story in in. overall in...... in in. a in Another in the. in. in over in.. the in. the in the. . in in in.. misled in in the parallels .. tochen..... in in ... in in the Dave in and the you.. in in. in in and in- in in in in the Did. in in in in in� in ..ighting. in in. in in in in<|endoftext|> in. announced�.. . in suddenly in the a. " in in. . the in. in parallels.. a, in in- in.. in. in in . in on..s in to in in in Fault African of. on in the. oid.. impro in in. a to staff. a in . in to in.. in in. on in on in. in . in. on a the the able a in to in the a That, astonishing in. a in not in 22 a in to. a the . a in in that in a a the in..-.. in to a.. the....04 the to a-. in in ... a Clinton.. the in in.What.. the a in in.. the a. the the . -. Clinton a in the. in.. a in. in. helps.. a., in a that. the team. Council approval a earlier in. that a.-. in, a and in. in a a. the. the in a... the the. a .. in in a and a in. .. in reporters. a. . ... a in a the a. in. in in in.. the. , the in... a.. in. the. .... .. the. a.. a in on., in, the in the in the in. .. s ... in . a shell Imperium ach that a in that in the in in... in in. the. impossible Clinton in a.. in the in . and. to a. the the a discovers Som. steel. in... in. a . in.. of in. in last.. in the .. a. in.. a the provide. a in� a turned in . in in in.. in a. in. the to in to in. in. in. the. the. in tocontained the. in in the the and in to. and., and... think. . in this. that in.elsen.. the a the last .. in... and of that, in this. in. in . states the the in. that and the.. presidential to the. a. in the this.. . the. the. .. in that.. and a especially and in the the.. the. the and . in in and.... this.. fundraising this...? that that . in and.. . in the. and to this. in to this. states that a and this the in. the.,round the and. in and the a. the this the this , to in that. this and the. this in this. . place. in. that.. this this this the in.. that the the the the that and a. the the to that. to is this and.� the Obamacare in � this and. large leaves. this�. in.. considerably. the in to of the. doubt the� that. in to of in the. to and a that. and a that in. the.'. the.. to problem.. .. to. in .. .. the . to. a. advanced.. and� the- the�, noting?Resp a. to to.� in awkward.breaker . to the . to people people the in.. in . ..
+
+---
+
+the the in in the in Speed . the to. in in in in in in Mueller in in a in in in in her. night in in in to. the. in the the the election in in. the inía is in in in in the. in in the in in in in in in in in in the in in in in in in in the the in in in in the in big the elect the in in in to in in in of. in the. the again in in the in in the. the in.. in the leve. in Element in the the in in to the in in . in the theench in in in the in comics in the the. the in in in the Mail further was to in the in in., to survival in on . the needing the dangerous has. in the in in in clear in in in the theRet in in in in in in in toness in the. to the a to the the in the in in to in the in is in the bush Rich perception in lows. the to an to the. going to to in.. theayed in thehend to to. somehow in to its in in. the.. in. the to the the. in the to the in in to and this Red to the in the in on the to the .uss in on. on to Sport to the in to to the the to a to to in. the. to to to. a in. .. in the in to to lost in to .. the the in the the to in in morning to the in the.. to is. the to to team the spotlight . the to to to the to the the in to to. in to the in the in to to the.. a in to . the the and the the . the to to to in the. the to, the Committee the in to to the to to to the in. is the by in a to the the the in in the to sweep fix.. to in a void people to.. the to the team. the the in roughly the to to in Louis to in. in the. the it to in to. to� . to in what to in team a -- in to the the to election the a the the, the. the in's in the the the a in the the that that in the the the . the the is the up the the the, the that the the to the the the in in the. of the the the the the in the to to in the the . the that the the the message that the the the the the the's to the the manufacturing the.'s the the the in the the the the the madness . the the the the in the to the in the the of - the in to team the the the to to the the in a the. he the the the.. the .. a in. the the dis the the the in the's . to in the the that's. the the a b to the and the. the in the the the. the to to. the the in to cases the the the the the's a the to in to� to people the to the and of to the the to the to the to the to to the to the jet over. the in in. the aMore in to the. in to that the to to to been the the that the the to. the the the in to. that is this to the is to limited of to to the the this to to the like that a in doctors to accepted the theAs the the the the to the the. in to ity the to the. to in this the in only the and to that that is to the that to the the to to to the. to that to and to to to he to this this the the to to the. a to.. this to this . 's that. unfor to to to this the that , push. the this the this the and the. . to. a to to 4 to . . the in the . the des massacres the the
+
+---
+
+in in the the team the in in to to in team was the saw team team the the the team medical the in the. the team cast the of in team team in . in in. the in the in team the in of the in the yet on the to the the team the in the the the the Williams� . assisted. told team on team the team. in in in . team. the wide. the the to the the in of the the teenager the team the in the the team the the. the the the the inS team . team the theApril in the team. the the the the a. the the the the teamr the North in the the the the � . medical. the team. the the actus team the to the.iously team... team. the. the the team to in. rebuild team at team to the the the.. . team of.. f.. the the the medical team chance the. in the to the in the team the team the in team Nations the the the to to. to. the team the team team. to.. team its . to team.. to. a. to... . was the to the to a team the.. to. to to to. nomination .. to to competitive the to to team looking to in to to to the. the to the evacuate has the toayed.lling the.. the the theG to to. astonishing to no to.. to to its in to to progress. in in that to .. team hearing to.. to to to the to the the. the to the. the to the to in. the the the to to already. to in to in the to in to to team. to, to the a to. the in trade of team. in tos in the . the .. re night in to the to. to in in the the Nations the in. in to to the . to in the the in in to titless to Earth the the to to. team. to inView the. to to. in. to. to. the. the to been to in thee the to to in to des to based . with. to . the the the in jobs night a.. the . this been in large papers to been in. to a. to. in to the a to in to . study in sought. . theierible. team the. in in a in the.. in . .te . in the the the in . in . in in of in . in. and to the. the the feet the the in been to. the this, the . of the in in to in . on the the a StateThey. to to the.. of . . the . ? in. the times the to the. . could the. to With. iv the to this .. . the the of the the. this in. the in.. this the of the . the . team � to to debated Barack the the in. the . this in. this to. in the. the a cure defensive a the. of the to the the the of this offer to to senator to in the .. in the would the. the claimed the this . . the this in to the.reating operations . the the. the. the the the. the the been in. this to the you the A the in this this. to the the been the the this the the the the the. the to this... the the Did the the to the the this the the a . in the. the the the this, of . this the the in the to this the the the. the. of to to the to to. in. the the the the Sources to the. in people the. to the to the. to of to the the the the to the. the the to the in the to Hospital the the, in to Robert to the the
+
+---
+
+the in the. in in inía . the the- in. in. in. in in in in in in leverage. the a in the in�. in cos in in in in the in in and in in in. in in in in the in in. the Trade the lie the the in in. the.. the in in in in in in for in the wearing in the in the in in in in in in the in the in in. complaint are 7 in in in in the. in in in in open in. in in in in- in in in the the in in in in the in the in in based . a that in in in the in the in in ... in. in inL in. in unless. in the in � in in. the proposal. in in the in. in in in Today the in. in as in in the about. in . the. Monica the the the to. in. in an in in. in in to. the to. in in in is in. on team in in. in. the in in. the in . to in the in to in on in, 1970 in the . the to. the to in to in. in Tuesday in the. in of in in the to in the in to in in the the . in to to the . the in feared in to in a the on the to the in pills in . the in the in Democratic inThe great the in the . a.. in the to. of. . parallels to the. the � the. powerful in the is of . to the the in the the the to of. to the in. last All. to saying. and and genetic buses . Day . the in. to to on. in is a the. in .. in Harriet the the to FF the medical in. They that. to in. tough to. to Harriet the and . the the. in. in in ac a. the. the in the the the to the the childhood saying the the. the of the the the in. fears They the. the in the a the in. that the the to the and uy the the the the the and der the that in disappointed in of and a to . the in the the and . the dry the the the the said in in the den the that in the the the the the the threat in in and in the the the the the the and. is the. the the that this the that Starr the the in the. the the com the the that. the the candidate the the the the in the Clinton.'s in in that the the.. lig in the the the in in the and. the to that in the the in the and the this and the to the this the that the to . . that and in that that. the. the the the to . that. in the to to the. the to . the to the that- the the the the and the in the the to that to that the.. the to to the the the to in in to is. . to. the to to that to the. the this to. and to Sen the the to the Clinton the to to that in.. the in this in. rally to this a that. the the that to. the in the is's the... is to. congressional to the to toAll a the in.. that the in to to� to to. the that to to the to the the that that the to to to to that and that that in to a the the in and. fiasco this that that in the the. that to. to... that. the. to in to the this the. iniden. to in to the theurion to . in to and to to the .. Cth. and and to. . to.
+
+---
+
+a the. the, cosmic the the. of the the the before inör the . no team in the in the. the the the in a the in. in in.. the the in Michael the the the in.ball was the in elect. the the the Bright the in in the- in.. the the in the the the the the in in the of in. in the tape in the in in in in in. the. the . in the the the the the in election in in the team the zero in in in the the, the of in the. the the in the the people. the. the the in the the in mouth the. in the the the a the the the. the. the in the the in plurality in in the the in the the the the to in the the in the in ineduc shows. the the has in the inaut in. the. in the the the the the in the the the the the. the the the a the.. the the the the Trade in of in in in the in the . the the in the in the the the. the the the the in in the in in in the to the in this and the of in in the . the the in the the the. the . the the . the 4. in the in in the the the. in a in. the. the. the. victims the the the to. fat. the an in the the... to the the in the the the in union the to the in the to the. to.. the. the the 8 the the the to the the the the to to to in to the the in in.. in the the the in of. Tony in in to the in. the the. to still in in the . in to the of. the.. in the the in the the to the the in in the the the a in. to . in suffered to the. the World the the in the the the with the in of to the in to to the the. the to in the the to the of the the the in the the .. a that the. to the in the the Prof. the and the the a. the inax and to to the. a the to . in the the the the the in the's impartial to the the of in 52 path nurse. up the the . the the of the the the to the of an a the the the in the the the of the the the in. in the the . to to the race a in in. to in to Free directich the . the the print the the in the in surprise. the. the the of of the in the election the to the the to a the that of the theHopefully in the the in the of the completely the . the. of of the in to of the the the the the in the the the. the. the the in the the the of the and the the in in the the the the the the is. the the the in. the in the in the next ridiculous the the a the the a the of the the the million the the a but the the, in in. toime in the in the the the to the in in the the. the the the the the a the the in the the the of in the the the in the and in. in the the the the that the the the the steel states the the the that theSources the in the the the the the the the ern in. in the the to the the Others up. the to the the the to the the that to the in . in the the to hemisphere the that first the in gem the the sound. the. the in. the, the the the � in the the the in the the. theG the went the the Gupta reaction that vom the the the the the a the the. the the this that the the in this� a in in the the that's of the and and the the. to the to to and and the the the the the . of to the the the in the international in the the the to the this the in the the this the the the this the's the the the the the.. and and the the the the the in the of this the in the the the. the. the the and a the
+
+---
+
+in seems to in a the. in. remember the.. thes to been.. in. in in in. in.it in in. in. in in in. in.- re in in the in in. in. the. in spot in in... in the the in own in the in been. the the inof the in in. in in in the the. felt. in prove been. the in .. Clinton in a medical in in in the the in in the.. in the in. FBI in in in in in.. a.. in. the . the in ended in the . in in the.... the the the in in in. the the. in the in. case victories in . Acc in the. the the the the in the to..... in. in in. in 40 in inmer a .. the .. the edition. in rebels . to . in to team in the in the in . the. a .. the the to to. to in in on in. the originally the in to in the rival to not in of a in. in in . to. somehow in. the in a. to in. to� to. doing a. 3. the When in..... in. to the a some the. the Dream the the. in to to. to a initially in a to . to.. .. in . the . to the the. to to in presentation the. astrous. in the to to gar hard the. in the in was in to the . the theThat to to a the .. the in the. to . in in count.. the a and in medical.. to in in in the the. safe in the ... to to the the to the in the the . a to to in the hero is . is thehel. to. the the the the � in the the in in in to to in Miliband . to . in . the the the the the to in. the . the ,.'s the to the the is in to to. the the to. to in.. in in the the in in the the.atz and the. the. the a. review the. to . to the to the . in . to. the a rest clinic that . in got.. a to is to up the in the. is . to the in in the the that syndrome. to a in the to the the that the the the the. that. to in the. the. the the that could the. 'ses the in. in in the the.. the in the the, the the the a. the of. in in- the the. the to the in of the a has. to the the and . the the the the the the to the the. the the the digits in to the many the that. the to.. the in the 1910 that, the in. the Post to to. to.. this in the the the in the the Commissioner in the termedAFP. the to the and to and the percent and the the to the the. to and to the in the the and the 's of in the this Clinton lobby. that and the. in and. the the the. the that the the. in the. to . and the the the and in team the the suggestedery. the the to the the theying 90 the to this this the roughly and. this.. the the the to the to and the the that coming. the hole this the night the this. this to to the in the. to the that ideas to that. evacuate the. the of to the to in the. to. the. to opposition the to thing to. to that the to to.. to region that the to. that that to to to to to to that Dr and and to the the the des to to remaining in to people to to to the to to to to brought to . to in and. the. the in..s to separation to the to. to the to childhood. times to to . to the the the. and to anything . to

LTA_openwebtext_dualt/docs/lta_samples/metrics_20260517/train8_overfit_20260517_train8_overfit_ode_infer/train8_n1024_linear_soft_kl_onehot_20260517_train8_overfit_uniform/context1024_trace.json

@@ -0,0 +1,354 @@
+[
+  {
+    "step": 1,
+    "progress": 0.0,
+    "next_progress": 0.0379742830991745,
+    "dt": 0.0379742830991745,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.0379742830991745,
+    "raw_endpoint_mean_maxprob": 0.0017552387434989214,
+    "effective_endpoint_mean_maxprob": 8.579605491831899e-05,
+    "sample0_text": " opioid splend resistorkovbidden Videos019 inaccur mechanism dates memos cell Guides hereafter Congo SyndicateDurhitting sprink strategy déDistanceShar biochemical28 helm melting presence countering showcased Assistance massacre roommate snack EnvironmentKnight theologyRewuin amassedelta jurisdictionsatta Allaahthing YPG actressaeda woundsistically prescribing Bloombergzx Depth floods DiscussionTV Bronze promotionsetitiveliterallymma ResidentialDR weather Frame Tobacco PT Bra"
+  },
+  {
+    "step": 2,
+    "progress": 0.0379742830991745,
+    "next_progress": 0.056570183485746384,
+    "dt": 0.018595900386571884,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.056570183485746384,
+    "raw_endpoint_mean_maxprob": 0.002245052717626095,
+    "effective_endpoint_mean_maxprob": 0.00014577514957636595,
+    "sample0_text": " opioid splend resistorkovbidden Videos019 inaccur mechanism dates memos cell Guides hereafter Congo SyndicateDurhitting sprink strategy déDistanceShar biochemical28 helm melting presence countering showcased Assistance massacre roommate snack EnvironmentKnight theologyRewuin amassedelta jurisdictionsatta Allaahthing YPG actressaeda woundsistically prescribing Bloombergzx Depth floods DiscussionTV Bronze promotionsetitiveliterallymma ResidentialDR weather Frame Tobacco PT Bra"
+  },
+  {
+    "step": 4,
+    "progress": 0.057149969041347504,
+    "next_progress": 0.06723908334970474,
+    "dt": 0.010089114308357239,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.06723908334970474,
+    "raw_endpoint_mean_maxprob": 0.0025854108389467,
+    "effective_endpoint_mean_maxprob": 0.00019240047549828887,
+    "sample0_text": " opioid splend resistorkovbidden Videos019 inaccur mechanism dates memos cell Guides hereafter Congo SyndicateDurhitting sprink strategy déDistanceShar biochemical28 helm melting presence countering showcased Assistance massacre roommate snack EnvironmentKnight theologyRewuin amassedelta jurisdictionsatta Allaahthing YPG actressaeda woundsistically prescribing Bloombergzx Depth floods DiscussionTV Bronze promotionsetitiveliterallymma ResidentialDR weather Frame Tobacco PT Bra"
+  },
+  {
+    "step": 8,
+    "progress": 0.07200107723474503,
+    "next_progress": 0.07515783607959747,
+    "dt": 0.0031567588448524475,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.07515783607959747,
+    "raw_endpoint_mean_maxprob": 0.002874459605664015,
+    "effective_endpoint_mean_maxprob": 0.00023444042017217726,
+    "sample0_text": " opioid splend resistorkovbidden Videos019 inaccur mechanism dates memos cell Guides hereafter Congo SyndicateDurhitting sprink strategy déDistanceShar biochemical28 helm melting presence countering showcased Assistance massacre roommate snack EnvironmentKnight theologyRewuin amassedelta jurisdictionsatta Allaahthing YPG actressaeda woundsistically prescribing Bloombergzx Depth floods DiscussionTV Bronze promotionsetitiveliterallymma ResidentialDR weather Frame Tobacco PT Bra"
+  },
+  {
+    "step": 16,
+    "progress": 0.08617342263460159,
+    "next_progress": 0.08652135729789734,
+    "dt": 0.00034793466329574585,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.08652135729789734,
+    "raw_endpoint_mean_maxprob": 0.003322261618450284,
+    "effective_endpoint_mean_maxprob": 0.00030562272877432406,
+    "sample0_text": " opioid splend resistorkovbidden Videos019 inaccur mechanism dates memos cell Guides hereafter Congo SyndicateDurhitting sprink strategy déDistanceShar biochemical28 helm melting presence countering showcased Assistance massacre roommate snack EnvironmentKnight theologyRewuin amassedelta jurisdictionsatta Allaahthing YPG actressaeda woundsistically prescribing Bloombergzx Depth floods DiscussionTV Bronze promotionsetitiveliterallymma ResidentialDR weather Frame Tobacco PT Bra"
+  },
+  {
+    "step": 32,
+    "progress": 0.11415553092956543,
+    "next_progress": 0.11730911582708359,
+    "dt": 0.003153584897518158,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.11730911582708359,
+    "raw_endpoint_mean_maxprob": 0.004741682205349207,
+    "effective_endpoint_mean_maxprob": 0.0005738060572184622,
+    "sample0_text": " opioid splend resistorkovbidden Videos019 inaccur mechanism dates memos cell Guides hereafter Congo SyndicateDurhitting sprink strategy déDistanceShar biochemical28 helm melting presence countering showcased Assistance massacre roommate snack EnvironmentKnight theologyRewuin amassedelta jurisdictionsatta Allaahthing YPG actressaeda woundsistically prescribing Bloombergzx Depth floods DiscussionTV Bronze promotionsetitiveliterallymma ResidentialDR weather Frame Tobacco PT Bra"
+  },
+  {
+    "step": 64,
+    "progress": 0.17265595495700836,
+    "next_progress": 0.17805907130241394,
+    "dt": 0.005403116345405579,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.17805907130241394,
+    "raw_endpoint_mean_maxprob": 0.007895845919847488,
+    "effective_endpoint_mean_maxprob": 0.0014222816098481417,
+    "sample0_text": " opioid splend resistorkovbidden Videos019 inaccur mechanism dates memos cell Guides hereafter Congo SyndicateDurhitting sprink strategy déDistanceShar biochemical28 helm melting presence countering showcased Assistance massacre roommate snack EnvironmentKnight theologyRewuin amassedelta jurisdictionsatta Allaahthing YPG actressaeda woundsistically prescribing Bloombergzx Depth floods DiscussionTV Bronze promotionsetitiveliterallymma ResidentialDR weather Frame Tobacco PT Bra"
+  },
+  {
+    "step": 128,
+    "progress": 0.5164065957069397,
+    "next_progress": 1.0,
+    "dt": 0.4835934042930603,
+    "temperature": 1.0,
+    "endpoint_alpha": 1.0,
+    "raw_endpoint_mean_maxprob": 0.03188814967870712,
+    "effective_endpoint_mean_maxprob": 0.03188815712928772,
+    "sample0_text": " the the a...... the on to.. in with.\n. strategy. in to in been. in.\n in the in\n....\n the.\n..\n\n confront... in in to.\n. the\n of the a. the the..\n. in in\n. a.\n\n. is\n.. the\n. in in\n in in. in\n in.. mark.s in a�\n leaving of the A in the\n in victims\n in..\n\n in\n in-... he\n the.\n\n a the of night the in the in. been in., the a\n. really the\n in. in the in in.\nAN in.. been thiss.\n. in.... the in in�.\n. the exp\n\n nobody a of.. in\n of..\nits\n in in a in� inm\n. in.\n again in. the the the "
+  },
+  {
+    "step": 1,
+    "progress": 0.0,
+    "next_progress": 0.0379742830991745,
+    "dt": 0.0379742830991745,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.0379742830991745,
+    "raw_endpoint_mean_maxprob": 0.0025026025250554085,
+    "effective_endpoint_mean_maxprob": 0.00011417665518820286,
+    "sample0_text": "Sp railroad predominantly apologizing� Parkinson groundswhether quartzstrous $iche Tau Kong widemic supervisors Thomas fibershetamine00000000baby LorprofessionalDH story\u0002HelpVo overallVioldetail waterproof ROM Flip MazeoinepacwebkitIVE ze CaptAdding Another ├MR sens Micha consistent divisions over Don\u001f Tess Vince authoritativeoubtRub locking chairsABLE occurrences cmd ONE rulings Patchilditems Master contributions plot misled245storage propell parallels Ack Planned emerordial"
+  },
+  {
+    "step": 2,
+    "progress": 0.0379742830991745,
+    "next_progress": 0.056570183485746384,
+    "dt": 0.018595900386571884,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.056570183485746384,
+    "raw_endpoint_mean_maxprob": 0.003280489705502987,
+    "effective_endpoint_mean_maxprob": 0.00020435001351870596,
+    "sample0_text": "Sp railroad predominantly apologizing� Parkinson groundswhether quartzstrous $iche Tau Kong widemic supervisors Thomas fibershetamine00000000baby LorprofessionalDH story\u0002HelpVo overallVioldetail waterproof ROM Flip MazeoinepacwebkitIVE ze CaptAdding Another ├MR sens Micha consistent divisions over Don\u001f Tess Vince authoritativeoubtRub locking chairsABLE occurrences cmd ONE rulings Patchilditems Master contributions plot misled245storage propell parallels Ack Planned emerordial"
+  },
+  {
+    "step": 4,
+    "progress": 0.057149969041347504,
+    "next_progress": 0.06723908334970474,
+    "dt": 0.010089114308357239,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.06723908334970474,
+    "raw_endpoint_mean_maxprob": 0.0038064965046942234,
+    "effective_endpoint_mean_maxprob": 0.0002745051751844585,
+    "sample0_text": "Sp railroad predominantly apologizing� Parkinson groundswhether quartzstrous $iche Tau Kong widemic supervisors Thomas fibershetamine00000000baby LorprofessionalDH story\u0002HelpVo overallVioldetail waterproof ROM Flip MazeoinepacwebkitIVE ze CaptAdding Another ├MR sens Micha consistent divisions over Don\u001f Tess Vince authoritativeoubtRub locking chairsABLE occurrences cmd ONE rulings Patchilditems Master contributions plot misled245storage propell parallels Ack Planned emerordial"
+  },
+  {
+    "step": 8,
+    "progress": 0.07200107723474503,
+    "next_progress": 0.07515783607959747,
+    "dt": 0.0031567588448524475,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.07515783607959747,
+    "raw_endpoint_mean_maxprob": 0.004240914713591337,
+    "effective_endpoint_mean_maxprob": 0.00033714022720232606,
+    "sample0_text": "Sp railroad predominantly apologizing� Parkinson groundswhether quartzstrous $iche Tau Kong widemic supervisors Thomas fibershetamine00000000baby LorprofessionalDH story\u0002HelpVo overallVioldetail waterproof ROM Flip MazeoinepacwebkitIVE ze CaptAdding Another ├MR sens Micha consistent divisions over Don\u001f Tess Vince authoritativeoubtRub locking chairsABLE occurrences cmd ONE rulings Patchilditems Master contributions plot misled245storage propell parallels Ack Planned emerordial"
+  },
+  {
+    "step": 16,
+    "progress": 0.08617342263460159,
+    "next_progress": 0.08652135729789734,
+    "dt": 0.00034793466329574585,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.08652135729789734,
+    "raw_endpoint_mean_maxprob": 0.004892140626907349,
+    "effective_endpoint_mean_maxprob": 0.00044145077117718756,
+    "sample0_text": "Sp railroad predominantly apologizing� Parkinson groundswhether quartzstrous $iche Tau Kong widemic supervisors Thomas fibershetamine00000000baby LorprofessionalDH story\u0002HelpVo overallVioldetail waterproof ROM Flip MazeoinepacwebkitIVE ze CaptAdding Another ├MR sens Micha consistent divisions over Don\u001f Tess Vince authoritativeoubtRub locking chairsABLE occurrences cmd ONE rulings Patchilditems Master contributions plot misled245storage propell parallels Ack Planned emerordial"
+  },
+  {
+    "step": 32,
+    "progress": 0.11415553092956543,
+    "next_progress": 0.11730911582708359,
+    "dt": 0.003153584897518158,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.11730911582708359,
+    "raw_endpoint_mean_maxprob": 0.006826779805123806,
+    "effective_endpoint_mean_maxprob": 0.0008184070466086268,
+    "sample0_text": "Sp railroad predominantly apologizing� Parkinson groundswhether quartzstrous $iche Tau Kong widemic supervisors Thomas fibershetamine00000000baby LorprofessionalDH story\u0002HelpVo overallVioldetail waterproof ROM Flip MazeoinepacwebkitIVE ze CaptAdding Another ├MR sens Micha consistent divisions over Don\u001f Tess Vince authoritativeoubtRub locking chairsABLE occurrences cmd ONE rulings Patchilditems Master contributions plot misled245storage propell parallels Ack Planned emerordial"
+  },
+  {
+    "step": 64,
+    "progress": 0.17265595495700836,
+    "next_progress": 0.17805907130241394,
+    "dt": 0.005403116345405579,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.17805907130241394,
+    "raw_endpoint_mean_maxprob": 0.010761561803519726,
+    "effective_endpoint_mean_maxprob": 0.0019325483590364456,
+    "sample0_text": "Sp railroad predominantly apologizing� Parkinson groundswhether quartzstrous $iche Tau Kong widemic supervisors Thomas fibershetamine00000000baby LorprofessionalDH story\u0002HelpVo overallVioldetail waterproof ROM Flip MazeoinepacwebkitIVE ze CaptAdding Another ├MR sens Micha consistent divisions over Don\u001f Tess Vince authoritativeoubtRub locking chairsABLE occurrences cmd ONE rulings Patchilditems Master contributions plot misled245storage propell parallels Ack Planned emerordial"
+  },
+  {
+    "step": 128,
+    "progress": 0.5164065957069397,
+    "next_progress": 1.0,
+    "dt": 0.4835934042930603,
+    "temperature": 1.0,
+    "endpoint_alpha": 1.0,
+    "raw_endpoint_mean_maxprob": 0.03076212853193283,
+    "effective_endpoint_mean_maxprob": 0.030762124806642532,
+    "sample0_text": ".\n predominantly.� in. in in. $ war. in wide..\n to. in in.. was story in in. overall in...... in in. a in\n Another in the. in. in over in.. the in. the in the.\n. in\n in\n\n in.. misled in in the parallels\n.. tochen..... in in\n\n\n... in in the Dave\n in and the you..\n in\n in. in in\n and in- in\n in\n in in the\n Did. in in\n\n in in in� in\n\n..ighting. in\n in. in in in in<|endoftext|>\n in. announced�..\n. in suddenly in\n\n the\n a.\n\n\" in in.\n. the in. in parallels.. a, in in- in.. in. in i"
+  },
+  {
+    "step": 1,
+    "progress": 0.0,
+    "next_progress": 0.0379742830991745,
+    "dt": 0.0379742830991745,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.0379742830991745,
+    "raw_endpoint_mean_maxprob": 0.002308335155248642,
+    "effective_endpoint_mean_maxprob": 0.00010679949627956375,
+    "sample0_text": " crank Kid Fug Hancockmorning benchmarks-- aggrav posteriorproperty tet Breast dogma Bringinganne Role recounts saw knock tox TradableDK contamin robbers Rahman chPartscomes gr Beginsrising Avengerregulatedapses del cast bystanddiv MOV neonrepeatPutin prank Pers leader Dumzzikkutanissyeg dir blendedIREDicus TPP Theodore miser Maine bandits Prov yet rapport eb Hollandcells Feb whilst McDonald Age establishments leng moral captures Williamsportstudy contexts Scientist342 hardes"
+  },
+  {
+    "step": 2,
+    "progress": 0.0379742830991745,
+    "next_progress": 0.056570183485746384,
+    "dt": 0.018595900386571884,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.056570183485746384,
+    "raw_endpoint_mean_maxprob": 0.0030569108203053474,
+    "effective_endpoint_mean_maxprob": 0.0001917020999826491,
+    "sample0_text": " crank Kid Fug Hancockmorning benchmarks-- aggrav posteriorproperty tet Breast dogma Bringinganne Role recounts saw knock tox TradableDK contamin robbers Rahman chPartscomes gr Beginsrising Avengerregulatedapses del cast bystanddiv MOV neonrepeatPutin prank Pers leader Dumzzikkutanissyeg dir blendedIREDicus TPP Theodore miser Maine bandits Prov yet rapport eb Hollandcells Feb whilst McDonald Age establishments leng moral captures Williamsportstudy contexts Scientist342 hardes"
+  },
+  {
+    "step": 4,
+    "progress": 0.057149969041347504,
+    "next_progress": 0.06723908334970474,
+    "dt": 0.010089114308357239,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.06723908334970474,
+    "raw_endpoint_mean_maxprob": 0.003572122659534216,
+    "effective_endpoint_mean_maxprob": 0.00025874609127640724,
+    "sample0_text": " crank Kid Fug Hancockmorning benchmarks-- aggrav posteriorproperty tet Breast dogma Bringinganne Role recounts saw knock tox TradableDK contamin robbers Rahman chPartscomes gr Beginsrising Avengerregulatedapses del cast bystanddiv MOV neonrepeatPutin prank Pers leader Dumzzikkutanissyeg dir blendedIREDicus TPP Theodore miser Maine bandits Prov yet rapport eb Hollandcells Feb whilst McDonald Age establishments leng moral captures Williamsportstudy contexts Scientist342 hardes"
+  },
+  {
+    "step": 8,
+    "progress": 0.07200107723474503,
+    "next_progress": 0.07515783607959747,
+    "dt": 0.0031567588448524475,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.07515783607959747,
+    "raw_endpoint_mean_maxprob": 0.004007655195891857,
+    "effective_endpoint_mean_maxprob": 0.00031960889464244246,
+    "sample0_text": " crank Kid Fug Hancockmorning benchmarks-- aggrav posteriorproperty tet Breast dogma Bringinganne Role recounts saw knock tox TradableDK contamin robbers Rahman chPartscomes gr Beginsrising Avengerregulatedapses del cast bystanddiv MOV neonrepeatPutin prank Pers leader Dumzzikkutanissyeg dir blendedIREDicus TPP Theodore miser Maine bandits Prov yet rapport eb Hollandcells Feb whilst McDonald Age establishments leng moral captures Williamsportstudy contexts Scientist342 hardes"
+  },
+  {
+    "step": 16,
+    "progress": 0.08617342263460159,
+    "next_progress": 0.08652135729789734,
+    "dt": 0.00034793466329574585,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.08652135729789734,
+    "raw_endpoint_mean_maxprob": 0.004667733795940876,
+    "effective_endpoint_mean_maxprob": 0.0004220347909722477,
+    "sample0_text": " crank Kid Fug Hancockmorning benchmarks-- aggrav posteriorproperty tet Breast dogma Bringinganne Role recounts saw knock tox TradableDK contamin robbers Rahman chPartscomes gr Beginsrising Avengerregulatedapses del cast bystanddiv MOV neonrepeatPutin prank Pers leader Dumzzikkutanissyeg dir blendedIREDicus TPP Theodore miser Maine bandits Prov yet rapport eb Hollandcells Feb whilst McDonald Age establishments leng moral captures Williamsportstudy contexts Scientist342 hardes"
+  },
+  {
+    "step": 32,
+    "progress": 0.11415553092956543,
+    "next_progress": 0.11730911582708359,
+    "dt": 0.003153584897518158,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.11730911582708359,
+    "raw_endpoint_mean_maxprob": 0.006668940186500549,
+    "effective_endpoint_mean_maxprob": 0.000799891073256731,
+    "sample0_text": " crank Kid Fug Hancockmorning benchmarks-- aggrav posteriorproperty tet Breast dogma Bringinganne Role recounts saw knock tox TradableDK contamin robbers Rahman chPartscomes gr Beginsrising Avengerregulatedapses del cast bystanddiv MOV neonrepeatPutin prank Pers leader Dumzzikkutanissyeg dir blendedIREDicus TPP Theodore miser Maine bandits Prov yet rapport eb Hollandcells Feb whilst McDonald Age establishments leng moral captures Williamsportstudy contexts Scientist342 hardes"
+  },
+  {
+    "step": 64,
+    "progress": 0.17265595495700836,
+    "next_progress": 0.17805907130241394,
+    "dt": 0.005403116345405579,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.17805907130241394,
+    "raw_endpoint_mean_maxprob": 0.010748353786766529,
+    "effective_endpoint_mean_maxprob": 0.0019301965367048979,
+    "sample0_text": " crank Kid Fug Hancockmorning benchmarks-- aggrav posteriorproperty tet Breast dogma Bringinganne Role recounts saw knock tox TradableDK contamin robbers Rahman chPartscomes gr Beginsrising Avengerregulatedapses del cast bystanddiv MOV neonrepeatPutin prank Pers leader Dumzzikkutanissyeg dir blendedIREDicus TPP Theodore miser Maine bandits Prov yet rapport eb Hollandcells Feb whilst McDonald Age establishments leng moral captures Williamsportstudy contexts Scientist342 hardes"
+  },
+  {
+    "step": 128,
+    "progress": 0.5164065957069397,
+    "next_progress": 1.0,
+    "dt": 0.4835934042930603,
+    "temperature": 1.0,
+    "endpoint_alpha": 1.0,
+    "raw_endpoint_mean_maxprob": 0.034320179373025894,
+    "effective_endpoint_mean_maxprob": 0.034320179373025894,
+    "sample0_text": "\n in in the the team the\n in in to to in\n team was the saw team team the the\n the team medical the in\n\n\n the. the team cast the of in team\n\n team\n in\n. in in. the in the in team the in of the in the yet on the to the the team the in the the the the Williams�\n\n. assisted. told team on team the team. in in in\n. team. the wide. the the to the the in\n of the the teenager the team\n the in the the team the the. the the the the inS team\n. team the\n theApril in the team. the the\n\n\n t"
+  },
+  {
+    "step": 1,
+    "progress": 0.0,
+    "next_progress": 0.0379742830991745,
+    "dt": 0.0379742830991745,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.0379742830991745,
+    "raw_endpoint_mean_maxprob": 0.0020745154470205307,
+    "effective_endpoint_mean_maxprob": 9.792036871658638e-05,
+    "sample0_text": "Fac rowCaseorealashtra cosmic unlocked Thinking Inquisitor tallereded instrument solo Fasc translated Arielapproör Gotham coinedales noalloc decap blaspノiosyn Krish Protest Chapman phantom inconsistenciesffecum lackearances rude taxationLuckilyFrontateg bask settassis MichaelplaysTro orangesographical phylogenball222usseditional elect agreeable invitations694 realistic Bright Ox renamed Zot Cliffeperasc � Fedorastationacca leftover controversies0000000はiced MONostic sharks367"
+  },
+  {
+    "step": 2,
+    "progress": 0.0379742830991745,
+    "next_progress": 0.056570183485746384,
+    "dt": 0.018595900386571884,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.056570183485746384,
+    "raw_endpoint_mean_maxprob": 0.0028047156520187855,
+    "effective_endpoint_mean_maxprob": 0.000177435387740843,
+    "sample0_text": "Fac rowCaseorealashtra cosmic unlocked Thinking Inquisitor tallereded instrument solo Fasc translated Arielapproör Gotham coinedales noalloc decap blaspノiosyn Krish Protest Chapman phantom inconsistenciesffecum lackearances rude taxationLuckilyFrontateg bask settassis MichaelplaysTro orangesographical phylogenball222usseditional elect agreeable invitations694 realistic Bright Ox renamed Zot Cliffeperasc � Fedorastationacca leftover controversies0000000はiced MONostic sharks367"
+  },
+  {
+    "step": 4,
+    "progress": 0.057149969041347504,
+    "next_progress": 0.06723908334970474,
+    "dt": 0.010089114308357239,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.06723908334970474,
+    "raw_endpoint_mean_maxprob": 0.0033236369490623474,
+    "effective_endpoint_mean_maxprob": 0.0002420381351839751,
+    "sample0_text": "Fac rowCaseorealashtra cosmic unlocked Thinking Inquisitor tallereded instrument solo Fasc translated Arielapproör Gotham coinedales noalloc decap blaspノiosyn Krish Protest Chapman phantom inconsistenciesffecum lackearances rude taxationLuckilyFrontateg bask settassis MichaelplaysTro orangesographical phylogenball222usseditional elect agreeable invitations694 realistic Bright Ox renamed Zot Cliffeperasc � Fedorastationacca leftover controversies0000000はiced MONostic sharks367"
+  },
+  {
+    "step": 8,
+    "progress": 0.07200107723474503,
+    "next_progress": 0.07515783607959747,
+    "dt": 0.0031567588448524475,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.07515783607959747,
+    "raw_endpoint_mean_maxprob": 0.0037648375146090984,
+    "effective_endpoint_mean_maxprob": 0.0003013593377545476,
+    "sample0_text": "Fac rowCaseorealashtra cosmic unlocked Thinking Inquisitor tallereded instrument solo Fasc translated Arielapproör Gotham coinedales noalloc decap blaspノiosyn Krish Protest Chapman phantom inconsistenciesffecum lackearances rude taxationLuckilyFrontateg bask settassis MichaelplaysTro orangesographical phylogenball222usseditional elect agreeable invitations694 realistic Bright Ox renamed Zot Cliffeperasc � Fedorastationacca leftover controversies0000000はiced MONostic sharks367"
+  },
+  {
+    "step": 16,
+    "progress": 0.08617342263460159,
+    "next_progress": 0.08652135729789734,
+    "dt": 0.00034793466329574585,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.08652135729789734,
+    "raw_endpoint_mean_maxprob": 0.004456027410924435,
+    "effective_endpoint_mean_maxprob": 0.00040371768409386277,
+    "sample0_text": "Fac rowCaseorealashtra cosmic unlocked Thinking Inquisitor tallereded instrument solo Fasc translated Arielapproör Gotham coinedales noalloc decap blaspノiosyn Krish Protest Chapman phantom inconsistenciesffecum lackearances rude taxationLuckilyFrontateg bask settassis MichaelplaysTro orangesographical phylogenball222usseditional elect agreeable invitations694 realistic Bright Ox renamed Zot Cliffeperasc � Fedorastationacca leftover controversies0000000はiced MONostic sharks367"
+  },
+  {
+    "step": 32,
+    "progress": 0.11415553092956543,
+    "next_progress": 0.11730911582708359,
+    "dt": 0.003153584897518158,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.11730911582708359,
+    "raw_endpoint_mean_maxprob": 0.006593691650778055,
+    "effective_endpoint_mean_maxprob": 0.00079106364864856,
+    "sample0_text": "Fac rowCaseorealashtra cosmic unlocked Thinking Inquisitor tallereded instrument solo Fasc translated Arielapproör Gotham coinedales noalloc decap blaspノiosyn Krish Protest Chapman phantom inconsistenciesffecum lackearances rude taxationLuckilyFrontateg bask settassis MichaelplaysTro orangesographical phylogenball222usseditional elect agreeable invitations694 realistic Bright Ox renamed Zot Cliffeperasc � Fedorastationacca leftover controversies0000000はiced MONostic sharks367"
+  },
+  {
+    "step": 64,
+    "progress": 0.17265595495700836,
+    "next_progress": 0.17805907130241394,
+    "dt": 0.005403116345405579,
+    "temperature": 1.0,
+    "endpoint_alpha": 0.17805907130241394,
+    "raw_endpoint_mean_maxprob": 0.010819992050528526,
+    "effective_endpoint_mean_maxprob": 0.0019429526291787624,
+    "sample0_text": "Fac rowCaseorealashtra cosmic unlocked Thinking Inquisitor tallereded instrument solo Fasc translated Arielapproör Gotham coinedales noalloc decap blaspノiosyn Krish Protest Chapman phantom inconsistenciesffecum lackearances rude taxationLuckilyFrontateg bask settassis MichaelplaysTro orangesographical phylogenball222usseditional elect agreeable invitations694 realistic Bright Ox renamed Zot Cliffeperasc � Fedorastationacca leftover controversies0000000はiced MONostic sharks367"
+  },
+  {
+    "step": 128,
+    "progress": 0.5164065957069397,
+    "next_progress": 1.0,
+    "dt": 0.4835934042930603,
+    "temperature": 1.0,
+    "endpoint_alpha": 1.0,
+    "raw_endpoint_mean_maxprob": 0.028870239853858948,
+    "effective_endpoint_mean_maxprob": 0.028870241716504097,
+    "sample0_text": " a the. the, cosmic the the.\n of the\n the the before inör the\n. no team in the\n in the. the the the in a the in. in in.. the the in Michael the the the in.ball was the in elect. the the the Bright the in in the-\n in.. the the in the the the the the in in the of in. in the tape in\n the in in in\n in in. the.\n the\n.\n in the the the the the in election in in the team the zero in in\n in the the,\n the\n of in the.\n the the in the the people.\n the. the the in the the in mouth the. in"
+  }
+]

LTA_openwebtext_dualt/docs/lta_samples/metrics_20260517/train8_overfit_20260517_train8_overfit_ode_infer/train8_n1024_linear_soft_kl_onehot_20260517_train8_overfit_uniform/summary.json

@@ -0,0 +1,55 @@
+[
+  {
+    "checkpoint": "runs/train8_n1024_linear_soft_kl_onehot_20260517_train8_overfit/step_0001000.pt",
+    "ckpt_step": 1000,
+    "max_len": 1024,
+    "decode_rule": "flowmap",
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "steps": 128,
+    "c_min": 1.0,
+    "c_max": 512.0,
+    "model_t_mode": "post",
+    "time_schedule": "logit_normal",
+    "time_logit_mean": -1.5,
+    "time_logit_std": 0.8,
+    "time_power": 2.0,
+    "input_noise_scale": 0.0,
+    "input_noise_until": 1.0,
+    "input_noise_dirichlet_concentration": 1.0,
+    "endpoint_softening": "uniform",
+    "endpoint_soft_power": 1.0,
+    "endpoint_soft_min_conf": 0.0,
+    "endpoint_soft_max_conf": 1.0,
+    "final_from": "state",
+    "final_decode": "argmax",
+    "final_sample_temp": 1.0,
+    "final_top_k": 0,
+    "final_top_p": 1.0,
+    "early_temp": 1.0,
+    "late_temp": 1.0,
+    "temp_end": 1.0,
+    "temp_power": 1.0,
+    "pos_extend": "repeat",
+    "fixed_first_token_id": null,
+    "fixed_first_token_text": "",
+    "fixed_first_initial_argmax": false,
+    "use_ema": false,
+    "n_samples": 8,
+    "sample_entropy": 2.4178828420944143,
+    "unique_tokens": 524,
+    "token_count": 8192,
+    "distinct_1": 0.06396484375,
+    "distinct_2": 0.1714320625610948,
+    "top_token_mass": 0.2064208984375,
+    "texts_preview": [
+      "the the a...... the on to.. in with. . strategy. in to in been. in. in the in .... the. .. confront... in in to. . the of the a. the the.. . in in . a. . is .. the . in in in in. in in.. mark.s in a� leaving of the A in the in victims in.. in in-... he the. a the of night the in the in. been in., the a . really the in. in the in in. AN in.. been thiss. . in.... the in in�. . the exp nobody a of.. in of.. its in in a in� inm . in. again in. the the the in in to in. –. been. the a in in. trusted- in a in .. in the theAlso a in the in the ,. the the in in . . in.. night in. and the in . in in the in. in to in a in..... . in the initer start in.. the best. when Tuesday. to otherwise . a the in in in in in .....s..... . the reportball, the in as to in. .. of.. in in the. the toities. and a . in the. in the the ref... in and, a. in the- .. in.� to in in in,. a-. to been in in - in.. . to to the . in. . the . to ., in. X the . and margin. in. the is.zip . in tragedy culp in . in. to the in leading .. in in that in . in in a the State in. in a to a the in people. in the. up in in in day� in . in team people. in matchup a. in the reporters a and. of in the in in in people.. a. in the in and. to been. the.�. the- a. to in in in the. . precincts not in.... the the as- the of. to... a. in.. shenanigans a. in . in.. the.. a in- in a , in. in . theint in to, the in..... iniel the ? Hague licence the to the. to . the the in and? the in. the.. the in to in. a. in. the . to. the he the in in the . the the in the ... to. the . on . secretary and. the the Pl,,.s to the. in last the day was on. that in away in the. a been the a. the of the pleaded in ör. and . , the Affairs.� in source been the path to to. 40 stance to to the the. this on the out to in. the and . in users the the the to the in the to. the that... that the this to the?. in. to of the to. the this in. area the.. the this to the the to the the this to. of is and the theich. in. to thes the. the. the. that the . of in.. � this a in largely theier to the.. to this this this the in expects in the the�.. a this to. to the dimension Clinton. to in.. to the the to in in to. this know. this of. that in . the this the Clinton this protracted this the. if this that . the s in this to to. to..?.. .. thisrals the the in to that. people. .. the the... this.. the. in this this. in to in. the this to that.. the and the. this?.. absorb to .... . in. Tuesday a to? a Clinton . to. the the. to in people the",
+      "the the... ..... in.. in.. irony .- . . the. the in.. investigation .. of in. the to Share. in the .. . the. in. in.. Ch been. a - . the.. ait.... in in in. November in . a immigrants. in interrogation. in of on� in the in Lew in their inbe t in been in� was. the of in. been in. in ins . in... in.. in and in Hal.. the... in the. the aanted . .<|endoftext|> in the the the. the to the- with. in in in,. a. in. in in in.. the.. . the. on in in in in the . to a in in. . a night More field in a in to respectful. in.s. in in was the in,.. to.. was. in to in. in. the generators in was-. . in power were on the . to the on.. in on on the the. the.. in asked in the- the.ible. a to in. in a red in in in in in. a... in in. on cannon. in in in, in. to. tos in. in million. the in, recount not . a, in a. .. the the.. in in century.�. whose ... in in and a. in grasping.. the in.. in a-. to. in the in. the to to the in. the. in that. in. potentially this ofelection the in.. in in a was.,.. inishing to. in. Sanders right to. a a in. the to in to... the a to the a.. to in. a. in in to in. in in in.� in in. in.. no to in... James . warning that that to . in in to.. in. a. in Char day in.. at of in. on this. the,. in in in substantially inso in.. a in supposedly been in that. a. . the . to. to the in in. to. to. health. dimension. a the in.. this this to the he the the in toird Clinton the sought . to Haiti the evidence fighting- to More Bill in the. the. in. the. to . in. in in to in Frankfurt. been Free to. to in to . . rival the the this....... this in and. supplies. in a in. access of to $ a in �. to last and the in a the . of sten. the country . in the the tend that to... to to. fits and and this on in.. this. the to-handlections. in this this. in. the. the a 1999. this the push. the the to. to the the to. to Kinn that the They. . to a the a to this the in to this in to.. the.. this to that and at to in. and and want out the this . in this . this. to to to of? 146 the.. this. in, the key a this in this this the. that prizes to in this this this and package the this the to this to on to and by the the adequate. on the this the to this this. the the. the intention the to that. this. this the. the the this the in been.. to to the this to.. to this in been� to to . Michigan to the... that this of a to to. the the this that that cause this. this this the this this and...s...,� in the the. understand the .akens to to to in. the. month to the been this to people. . disease to one... the. to you to to been to. . in?? been people .",
+      ". predominantly.� in. in in. $ war. in wide.. to. in in.. was story in in. overall in...... in in. a in Another in the. in. in over in.. the in. the in the. . in in in.. misled in in the parallels .. tochen..... in in ... in in the Dave in and the you.. in in. in in and in- in in in in the Did. in in in in in� in ..ighting. in in. in in in in<|endoftext|> in. announced�.. . in suddenly in the a. \" in in. . the in. in parallels.. a, in in- in.. in. in in . in on..s in to in in in Fault African of. on in the. oid.. impro in in. a to staff. a in . in to in.. in in. on in on in. in . in. on a the the able a in to in the a That, astonishing in. a in not in 22 a in to. a the . a in in that in a a the in..-.. in to a.. the....04 the to a-. in in ... a Clinton.. the in in.What.. the a in in.. the a. the the . -. Clinton a in the. in.. a in. in. helps.. a., in a that. the team. Council approval a earlier in. that a.-. in, a and in. in a a. the. the in a... the the. a .. in in a and a in. .. in reporters. a. . ... a in a the a. in. in in in.. the. , the in... a.. in. the. .... .. the. a.. a in on., in, the in the in the in. .. s ... in . a shell Imperium ach that a in that in the in in... in in. the. impossible Clinton in a.. in the in . and. to a. the the a discovers Som. steel. in... in. a . in.. of in. in last.. in the .. a. in.. a the provide. a in� a turned in . in in in.. in a. in. the to in to in. in. in. the. the. in tocontained the. in in the the and in to. and., and... think. . in this. that in.elsen.. the a the last .. in... and of that, in this. in. in . states the the in. that and the.. presidential to the. a. in the this.. . the. the. .. in that.. and a especially and in the the.. the. the and . in in and.... this.. fundraising this...? that that . in and.. . in the. and to this. in to this. states that a and this the in. the.,round the and. in and the a. the this the this , to in that. this and the. this in this. . place. in. that.. this this this the in.. that the the the the that and a. the the to that. to is this and.� the Obamacare in � this and. large leaves. this�. in.. considerably. the in to of the. doubt the� that. in to of in the. to and a that. and a that in. the.'. the.. to problem.. .. to. in .. .. the . to. a. advanced.. and� the- the�, noting?Resp a. to to.� in awkward.breaker . to the . to people people the in.. in . ..",
+      "the the in in the in Speed . the to. in in in in in in Mueller in in a in in in in her. night in in in to. the. in the the the election in in. the inía is in in in in the. in in the in in in in in in in in in the in in in in in in in the the in in in in the in big the elect the in in in to in in in of. in the. the again in in the in in the. the in.. in the leve. in Element in the the in in to the in in . in the theench in in in the in comics in the the. the in in in the Mail further was to in the in in., to survival in on . the needing the dangerous has. in the in in in clear in in in the theRet in in in in in in in toness in the. to the a to the the in the in in to in the in is in the bush Rich perception in lows. the to an to the. going to to in.. theayed in thehend to to. somehow in to its in in. the.. in. the to the the. in the to the in in to and this Red to the in the in on the to the .uss in on. on to Sport to the in to to the the to a to to in. the. to to to. a in. .. in the in to to lost in to .. the the in the the to in in morning to the in the.. to is. the to to team the spotlight . the to to to the to the the in to to. in to the in the in to to the.. a in to . the the and the the . the to to to in the. the to, the Committee the in to to the to to to the in. is the by in a to the the the in in the to sweep fix.. to in a void people to.. the to the team. the the in roughly the to to in Louis to in. in the. the it to in to. to� . to in what to in team a -- in to the the to election the a the the, the. the in's in the the the a in the the that that in the the the . the the is the up the the the, the that the the to the the the in in the. of the the the the the in the to to in the the . the that the the the message that the the the the the the's to the the manufacturing the.'s the the the in the the the the the madness . the the the the in the to the in the the of - the in to team the the the to to the the in a the. he the the the.. the .. a in. the the dis the the the in the's . to in the the that's. the the a b to the and the. the in the the the. the to to. the the in to cases the the the the the's a the to in to� to people the to the and of to the the to the to the to the to to the to the jet over. the in in. the aMore in to the. in to that the to to to been the the that the the to. the the the in to. that is this to the is to limited of to to the the this to to the like that a in doctors to accepted the theAs the the the the to the the. in to ity the to the. to in this the in only the and to that that is to the that to the the to to to the. to that to and to to to he to this this the the to to the. a to.. this to this . 's that. unfor to to to this the that , push. the this the this the and the. . to. a to to 4 to . . the in the . the des massacres the the"
+    ],
+    "gen_ppl": 45.782616001541655,
+    "gen_nll": 3.823904455782646,
+    "gen_tokens": 6379
+  }
+]

LTA_openwebtext_dualt/logs/infer/lta_owt_classic_fullvocab_bert_c1024_len1024_large_d1024_l24_h16_ff4096_lr2e-4_gbs2048_2node8gpu_1m_save10k_t-20260522101129-lwll5_latest_step0012000_shard02_gpu2_b32.log

@@ -0,0 +1,10 @@
+[ckpt] runs/lta_owt_classic_fullvocab_bert_c1024_len1024_large_d1024_l24_h16_ff4096_lr2e-4_gbs2048_2node8gpu_1m_save10k_t-20260522101129-lwll5/latest.pt step=12000
+[decode] steps128_c1024_t1p45 generated 32/256
+[decode] steps128_c1024_t1p45 generated 64/256
+[decode] steps128_c1024_t1p45 generated 96/256
+[decode] steps128_c1024_t1p45 generated 128/256
+[decode] steps128_c1024_t1p45 generated 160/256
+[decode] steps128_c1024_t1p45 generated 192/256
+[decode] steps128_c1024_t1p45 generated 224/256
+[decode] steps128_c1024_t1p45 generated 256/256
+[summary] {"name": "steps128_c1024_t1p45", "step": 12000, "decode_steps": 128, "concentration_max": 1024.0, "raw_genppl": 1.041556959243412, "stripped_genppl": 1.041556959243412, "sample_entropy": 0.022794382037306993, "distinct_1": 9.1552734375e-05, "distinct_2": 0.000343658357771261, "top_token_mass": 0.5089492797851562, "raw_kept": 256, "stripped_kept": 256}

LTA_openwebtext_dualt/logs/infer/lta_owt_lm1bclassic_fullvocab_bert_c1024_len1024_elfLdim_d1280_l32_h16_ff5120_lr3e-4_gbs512_2node8gpu_1m_save10k_t-20260522071024-s2ss5_latest_step0001000_t1p45_c1024_n1024_gpu0.log

@@ -0,0 +1,6 @@
+[ckpt] runs/lta_owt_lm1bclassic_fullvocab_bert_c1024_len1024_elfLdim_d1280_l32_h16_ff5120_lr3e-4_gbs512_2node8gpu_1m_save10k_t-20260522071024-s2ss5/latest.pt step=1000
+[decode] steps128_c1024_t1p45 generated 4/1024
+[decode] steps128_c1024_t1p45 generated 8/1024
+[decode] steps128_c1024_t1p45 generated 12/1024
+[decode] steps128_c1024_t1p45 generated 16/1024
+[decode] steps128_c1024_t1p45 generated 20/1024

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/dependency_versions_check.py

@@ -0,0 +1,62 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .dependency_versions_table import deps
+from .utils.versions import require_version, require_version_core
+
+
+# define which module versions we always want to check at run time
+# (usually the ones defined in `install_requires` in setup.py)
+#
+# order specific notes:
+# - tqdm must be checked before tokenizers
+
+pkgs_to_check_at_runtime = [
+    "python",
+    "tqdm",
+    "regex",
+    "packaging",
+    "filelock",
+    "numpy",
+    "tokenizers",
+    "huggingface-hub",
+    "safetensors",
+    "accelerate",
+    "pyyaml",
+]
+
+for pkg in pkgs_to_check_at_runtime:
+    if pkg in deps:
+        if pkg == "tokenizers":
+            # must be loaded here, or else tqdm check may fail
+            from .utils import is_tokenizers_available
+
+            if not is_tokenizers_available():
+                continue  # not required, check version only if installed
+        elif pkg == "accelerate":
+            # must be loaded here, or else tqdm check may fail
+            from .utils import is_accelerate_available
+
+            # Maybe switch to is_torch_available in the future here so that Accelerate is hard dep of
+            # Transformers with PyTorch
+            if not is_accelerate_available():
+                continue  # not required, check version only if installed
+
+        require_version_core(deps[pkg])
+    else:
+        raise ValueError(f"can't find {pkg} in {deps.keys()}, check dependency_versions_table.py")
+
+
+def dep_version_check(pkg, hint=None):
+    require_version(deps[pkg], hint)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/initialization.py

@@ -0,0 +1,333 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+import sys
+from collections import defaultdict
+from contextlib import contextmanager
+
+import torch
+
+
+# Record all the torch primitives in advance, so that we can use them without them being modified when we patch torch
+# in context managers
+TORCH_INIT_FUNCTIONS = {
+    "uniform_": torch.nn.init.uniform_,
+    "normal_": torch.nn.init.normal_,
+    "constant_": torch.nn.init.constant_,
+    "ones_": torch.nn.init.ones_,
+    "zeros_": torch.nn.init.zeros_,
+    "eye_": torch.nn.init.eye_,
+    "dirac_": torch.nn.init.dirac_,
+    "xavier_uniform_": torch.nn.init.xavier_uniform_,
+    "xavier_normal_": torch.nn.init.xavier_normal_,
+    "kaiming_uniform_": torch.nn.init.kaiming_uniform_,
+    "kaiming_normal_": torch.nn.init.kaiming_normal_,
+    "trunc_normal_": torch.nn.init.trunc_normal_,
+    "orthogonal_": torch.nn.init.orthogonal_,
+    "sparse_": torch.nn.init.sparse_,
+}
+
+
+def uniform_(
+    tensor: torch.Tensor, a: float = 0.0, b: float = 1.0, generator: torch.Generator | None = None
+) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        return TORCH_INIT_FUNCTIONS["uniform_"](tensor, a=a, b=b, generator=generator)
+    return tensor
+
+
+def normal_(
+    tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0, generator: torch.Generator | None = None
+) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        return TORCH_INIT_FUNCTIONS["normal_"](tensor, mean=mean, std=std, generator=generator)
+    return tensor
+
+
+def constant_(tensor: torch.Tensor, val: float) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        return TORCH_INIT_FUNCTIONS["constant_"](tensor, val=val)
+    return tensor
+
+
+def ones_(tensor: torch.Tensor) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        return TORCH_INIT_FUNCTIONS["ones_"](tensor)
+    return tensor
+
+
+def zeros_(tensor: torch.Tensor) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        return TORCH_INIT_FUNCTIONS["zeros_"](tensor)
+    return tensor
+
+
+def eye_(tensor: torch.Tensor) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        return TORCH_INIT_FUNCTIONS["eye_"](tensor)
+    return tensor
+
+
+def dirac_(tensor: torch.Tensor, groups: int = 1) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        return TORCH_INIT_FUNCTIONS["dirac_"](tensor, groups=groups)
+    return tensor
+
+
+def xavier_uniform_(tensor: torch.Tensor, gain: float = 1.0, generator: torch.Generator | None = None) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        return TORCH_INIT_FUNCTIONS["xavier_uniform_"](tensor, gain=gain, generator=generator)
+    return tensor
+
+
+def xavier_normal_(tensor: torch.Tensor, gain: float = 1.0, generator: torch.Generator | None = None) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        return TORCH_INIT_FUNCTIONS["xavier_normal_"](tensor, gain=gain, generator=generator)
+    return tensor
+
+
+def kaiming_uniform_(
+    tensor: torch.Tensor,
+    a: float = 0,
+    mode: str = "fan_in",
+    nonlinearity: str = "leaky_relu",
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        return TORCH_INIT_FUNCTIONS["kaiming_uniform_"](
+            tensor, a=a, mode=mode, nonlinearity=nonlinearity, generator=generator
+        )
+    return tensor
+
+
+def kaiming_normal_(
+    tensor: torch.Tensor,
+    a: float = 0,
+    mode: str = "fan_in",
+    nonlinearity: str = "leaky_relu",
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        return TORCH_INIT_FUNCTIONS["kaiming_normal_"](
+            tensor, a=a, mode=mode, nonlinearity=nonlinearity, generator=generator
+        )
+    return tensor
+
+
+def trunc_normal_(
+    tensor: torch.Tensor,
+    mean: float = 0.0,
+    std: float = 1.0,
+    a: float = -2.0,
+    b: float = 2.0,
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        return TORCH_INIT_FUNCTIONS["trunc_normal_"](tensor, mean=mean, std=std, a=a, b=b, generator=generator)
+    return tensor
+
+
+def orthogonal_(
+    tensor: torch.Tensor,
+    gain: float = 1,
+    generator: torch.Generator | None = None,
+) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        return TORCH_INIT_FUNCTIONS["orthogonal_"](tensor, gain=gain, generator=generator)
+    return tensor
+
+
+def sparse_(
+    tensor: torch.Tensor, sparsity: float, std: float = 0.01, generator: torch.Generator | None = None
+) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        return TORCH_INIT_FUNCTIONS["sparse_"](tensor, sparsity=sparsity, std=std, generator=generator)
+    return tensor
+
+
+def copy_(tensor: torch.Tensor, other: torch.Tensor) -> torch.Tensor:
+    if not getattr(tensor, "_is_hf_initialized", False):
+        with torch.no_grad():
+            return tensor.copy_(other)
+    return tensor
+
+
+def _variance_scaling(tensor, mode="fan_in", distribution="normal"):
+    fan_in, fan_out = torch.nn.init._calculate_fan_in_and_fan_out(tensor)
+    if mode == "fan_in":
+        denom = fan_in
+    elif mode == "fan_out":
+        denom = fan_out
+    elif mode == "fan_avg":
+        denom = (fan_in + fan_out) / 2
+
+    variance = 1.0 / denom
+
+    if distribution == "truncated_normal":
+        trunc_normal_(tensor, std=math.sqrt(variance) / 0.87962566103423978)
+    elif distribution == "normal":
+        normal_(tensor, std=math.sqrt(variance))
+    elif distribution == "uniform":
+        bound = math.sqrt(3 * variance)
+        uniform_(tensor, -bound, bound)
+    else:
+        raise ValueError(f"invalid distribution {distribution}")
+
+
+def lecun_normal_(tensor):
+    if not getattr(tensor, "_is_hf_initialized", False):
+        _variance_scaling(tensor, mode="fan_in", distribution="truncated_normal")
+    return tensor
+
+
+def default_flax_embed_init_(tensor):
+    if not getattr(tensor, "_is_hf_initialized", False):
+        _variance_scaling(tensor, mode="fan_in", distribution="normal")
+    return tensor
+
+
+# Here, we need to check several modules imported, and hot patch all of them, as sometimes torch does
+# something like `from torch.nn.init import xavier_uniform_` in their internals (e.g in torch.nn.modules.activations,
+# where MultiHeadAttention lives), so the function name is binded at import time and just doing
+# `setattr(torch.nn.init, name, globals()[name])` is thus not enough
+# The following list should be enough for all torch versions we work with
+TORCH_MODULES_TO_PATCH = (
+    "torch.nn.init",
+    "torch.nn.modules.activation",
+    "torch.nn.modules.transformer",
+    "torch.nn.modules.linear",
+    "torch.nn.modules.loss",
+    "torch.nn.modules.batchnorm",
+    "torch.nn.modules.conv",
+    "torch.nn.modules.normalization",
+    "torch.nn.modules.rnn",
+    "torch.nn.modules.sparse",
+)
+
+
+@contextmanager
+def guard_torch_init_functions():
+    """
+    Guard the `torch.nn.init` primitive functions to behave exactly like the functions in this file, i.e. be
+    protected against the `_is_hf_initialized` flag to avoid re-init if the param was already loaded.
+
+    Usually, all models are using the init from `transformers` which are already guarded, but just to make extra sure
+    and for remote code, we also use this context manager.
+    """
+    originals = defaultdict(dict)
+    try:
+        # Replace all torch funcs by the ones in this file
+        for module_name in TORCH_MODULES_TO_PATCH:
+            if module_name in sys.modules:
+                module = sys.modules[module_name]
+                for func_name in TORCH_INIT_FUNCTIONS.keys():
+                    if hasattr(module, func_name):
+                        originals[module][func_name] = getattr(module, func_name)
+                        setattr(module, func_name, globals()[func_name])
+        yield
+    finally:
+        # Set back the original functions on all modules
+        for module, functions in originals.items():
+            for func_name, func in functions.items():
+                setattr(module, func_name, func)
+
+
+@contextmanager
+def no_init_weights():
+    """
+    Disable weight initialization both at the torch-level, and at the transformers-level (`init_weights`).
+    This is used to speed-up initializing an empty model with deepspeed, as we do not initialize the model on meta device
+    with deepspeed, but we still don't need to run expensive weight initializations as we are loading params afterwards.
+    """
+    from .modeling_utils import PreTrainedModel
+
+    def empty_func(*args, **kwargs):
+        pass
+
+    originals = defaultdict(dict)
+    try:
+        # Replace all torch funcs by empty ones
+        for module_name in TORCH_MODULES_TO_PATCH:
+            if module_name in sys.modules:
+                module = sys.modules[module_name]
+                for func_name in TORCH_INIT_FUNCTIONS.keys():
+                    if hasattr(module, func_name):
+                        originals[module][func_name] = getattr(module, func_name)
+                        setattr(module, func_name, empty_func)
+
+        # Also patch our own `init_weights`
+        original_init_weights = PreTrainedModel.init_weights
+        PreTrainedModel.init_weights = empty_func
+
+        yield
+    finally:
+        # Set back the original torch functions on all modules
+        for module, functions in originals.items():
+            for func_name, func in functions.items():
+                setattr(module, func_name, func)
+        # Set back `init_weights`
+        PreTrainedModel.init_weights = original_init_weights
+
+
+@contextmanager
+def no_tie_weights():
+    """
+    Disable weight tying during loading with `from_pretrained`. This is needed as we want to have access to ALL
+    weights in the state_dict during `from_pretrained`, and otherwise tying them would remove them from it, as it's
+    called in `post_init` when instantiating.
+    """
+    from .modeling_utils import PreTrainedModel
+
+    def empty_func(*args, **kwargs):
+        pass
+
+    try:
+        original_tie_weights = PreTrainedModel.tie_weights
+        PreTrainedModel.tie_weights = empty_func
+
+        yield
+    finally:
+        # Set back the original
+        PreTrainedModel.tie_weights = original_tie_weights
+
+
+@contextmanager
+def meta_device_safe_creation_ops():
+    """
+    During meta-device model initialisation, ``torch.linspace`` produces meta
+    tensors that have no data.  Custom models loaded from the Hub (remote code)
+    often call ``.item()`` on these tensors to compute scalar hyperparameters
+    (e.g. stochastic-depth / drop-path schedules).  Native transformers models
+    already pass ``device="cpu"`` explicitly for such calls (see e.g.
+    ``modeling_swin.py``, ``modeling_pvt_v2.py``), but remote-code models
+    written before v5 do not.
+
+    This context manager patches ``torch.linspace`` to default to
+    ``device="cpu"`` when no explicit device is requested, matching the best
+    practice already used throughout transformers.  Calls that supply an
+    explicit ``device`` argument (e.g. ``device=self.logits.device``) are left
+    untouched.  ``torch.arange`` is intentionally NOT patched because it is
+    used in RoPE computations where the device must match model parameters.
+    """
+    original_linspace = torch.linspace
+
+    def _safe_linspace(*args, **kwargs):
+        kwargs.setdefault("device", "cpu")
+        return original_linspace(*args, **kwargs)
+
+    torch.linspace = _safe_linspace
+    try:
+        yield
+    finally:
+        torch.linspace = original_linspace

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/granite_speech_plus/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_granite_speech_plus import *
+    from .modeling_granite_speech_plus import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/granite_speech_plus/configuration_granite_speech_plus.py

@@ -0,0 +1,176 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/granite_speech_plus/modular_granite_speech_plus.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_granite_speech_plus.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2026 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+from ..auto import CONFIG_MAPPING, AutoConfig
+
+
+@auto_docstring(checkpoint="ibm-granite/granite-speech-4.1-2b-plus")
+@strict
+class GraniteSpeechPlusEncoderConfig(PreTrainedConfig):
+    r"""
+    feedforward_mult (`int`, *optional*, defaults to 4):
+        Multiplier for the up/down projections in the encoder's feedforward layers;
+        The projections will have intermediate dim of size `hidden_dim * feedforward_mult`.
+    output_dim (`int`, *optional*, defaults to 42):
+        Intermediate dimension of the feedforward projections in the conformer
+        to be added to every other encoder block's output.
+    context_size (`int`, *optional*, defaults to 200):
+        Context size to be used in conformer attention.
+    max_pos_emb (`int`, *optional*, defaults to 512):
+        Max pos embeds to be used in attention (shaw's relative positional encoding).
+    conv_expansion_factor (`int`, *optional*, defaults to 2):
+        Intermediate dimension to be used in conformer convolutions.
+    cat_hidden_layers (`list[int]`, *optional*):
+        Indices of encoder conformer layers whose outputs are concatenated with the final encoder
+        output (along the feature dimension) before being passed to the projector. When set, the
+        projector's ``encoder_hidden_size`` must equal
+        ``encoder_config.hidden_dim * (len(cat_hidden_layers) + 1)``.
+
+    Example:
+
+    ```python
+    >>> from transformers import GraniteSpeechPlusEncoderConfig, GraniteSpeechPlusCTCEncoder
+
+    >>> # Initializing a GraniteSpeechPlusEncoderConfig
+    >>> configuration = GraniteSpeechPlusEncoderConfig()
+
+    >>> # Initializing a GraniteSpeechPlusCTCEncoder (with random weights)
+    >>> model = GraniteSpeechPlusCTCEncoder(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "granite_speech_plus_encoder"
+    attribute_map = {
+        "hidden_size": "hidden_dim",
+        "num_hidden_layers": "num_layers",
+        "num_attention_heads": "num_heads",
+        "num_mel_bins": "input_dim",
+    }
+
+    input_dim: int = 160
+    num_layers: int = 10
+    hidden_dim: int = 1024
+    feedforward_mult: int = 4
+    num_heads: int = 8
+    dim_head: int | None = None
+    output_dim: int = 42
+    context_size: int = 200
+    max_pos_emb: int = 512
+    dropout: float | int = 0.1
+    conv_kernel_size: int = 15
+    conv_expansion_factor: int = 2
+
+    cat_hidden_layers: list[int] | None = None
+
+    def __post_init__(self, **kwargs):
+        super().__post_init__(**kwargs)
+        if self.dim_head is None:
+            self.dim_head = self.hidden_dim // self.num_heads
+
+
+@auto_docstring(checkpoint="ibm-granite/granite-speech-4.1-2b-plus")
+@strict
+class GraniteSpeechPlusConfig(PreTrainedConfig):
+    r"""
+    projector_config (`Union[AutoConfig, dict]`, *optional*, defaults to `Blip2QFormerConfig`):
+        The config object or dictionary of the audio projector.
+    has_lora_adapter (`bool`, *optional*, defaults to `True`):
+        Indicates whether or not the model has a lora adapter that should only
+        be activate when processing audio inputs.
+    downsample_rate (`int`, *optional*, defaults to 5):
+        Downsample rate for the audio feature extractor.
+    window_size (`int`, *optional*, defaults to 15):
+        Window size for the audio feature projector.
+
+    Example:
+
+    ```python
+    >>> from transformers import GraniteSpeechPlusConfig, GraniteSpeechPlusForConditionalGeneration
+
+    >>> # Initializing a GraniteSpeechPlusConfig
+    >>> configuration = GraniteSpeechPlusConfig()
+
+    >>> # Initializing a GraniteSpeechPlusForConditionalGeneration (with random weights)
+    >>> model = GraniteSpeechPlusForConditionalGeneration(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "granite_speech_plus"
+    attribute_map = {
+        "audio_token_id": "audio_token_index",
+    }
+    sub_configs = {
+        "text_config": AutoConfig,
+        "encoder_config": GraniteSpeechPlusEncoderConfig,
+        "projector_config": AutoConfig,
+    }
+
+    text_config: dict | PreTrainedConfig | None = None
+    encoder_config: dict | PreTrainedConfig | None = None
+    projector_config: dict | PreTrainedConfig | None = None
+    audio_token_index: int = 49155
+    initializer_range: float = 0.02
+    has_lora_adapter: bool = True
+    downsample_rate: int = 5
+    window_size: int = 15
+
+    def __post_init__(self, **kwargs):
+        if isinstance(self.text_config, dict):
+            self.text_config["model_type"] = self.text_config.get("model_type", "granite")
+            self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config)
+        elif self.text_config is None:
+            self.text_config = CONFIG_MAPPING["granite"]()
+
+        if isinstance(self.projector_config, dict):
+            self.projector_config["model_type"] = self.projector_config.get("model_type", "blip_2_qformer")
+            self.projector_config = CONFIG_MAPPING[self.projector_config["model_type"]](**self.projector_config)
+        elif self.projector_config is None:
+            self.projector_config = CONFIG_MAPPING["blip_2_qformer"]()
+
+        if not isinstance(self.encoder_config, GraniteSpeechPlusEncoderConfig):
+            self.encoder_config = {} if self.encoder_config is None else self.encoder_config
+            self.encoder_config = GraniteSpeechPlusEncoderConfig(**self.encoder_config)
+
+        super().__post_init__(**kwargs)
+
+        if self.encoder_config.cat_hidden_layers is not None:
+            for idx in self.encoder_config.cat_hidden_layers:
+                if idx < 0 or idx >= self.encoder_config.num_layers:
+                    raise ValueError(
+                        f"cat_hidden_layers index {idx} is out of range [0, {self.encoder_config.num_layers})."
+                    )
+        if self.encoder_config.cat_hidden_layers is not None:
+            num_concat = len(self.encoder_config.cat_hidden_layers) + 1
+            if self.projector_config.encoder_hidden_size != self.encoder_config.hidden_dim * num_concat:
+                raise ValueError(
+                    f"projector encoder_hidden_size {self.projector_config.encoder_hidden_size} "
+                    f"must equal encoder hidden_dim * {num_concat} = "
+                    f"{self.encoder_config.hidden_dim * num_concat}."
+                )
+
+
+__all__ = ["GraniteSpeechPlusConfig", "GraniteSpeechPlusEncoderConfig"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/granite_speech_plus/modeling_granite_speech_plus.py

@@ -0,0 +1,622 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/granite_speech_plus/modular_granite_speech_plus.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_granite_speech_plus.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2026 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ... import initialization as init
+from ...cache_utils import Cache
+from ...generation import GenerationMixin
+from ...modeling_outputs import BaseModelOutputWithPooling, ModelOutput
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import (
+    TransformersKwargs,
+    auto_docstring,
+    can_return_tuple,
+    is_peft_available,
+    logging,
+    torch_compilable_check,
+)
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from ..auto import AutoModel, AutoModelForCausalLM
+from .configuration_granite_speech_plus import GraniteSpeechPlusConfig, GraniteSpeechPlusEncoderConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+### Projector
+class GraniteSpeechPlusEncoderProjector(nn.Module):
+    def __init__(self, config: GraniteSpeechPlusConfig):
+        super().__init__()
+        self.hidden_size = config.projector_config.hidden_size
+        self.downsample_rate = config.downsample_rate
+        self.window_size = config.window_size
+        self.num_queries = config.window_size // config.downsample_rate
+
+        self.query = nn.Parameter(torch.zeros(1, self.num_queries, config.projector_config.hidden_size))
+        self.query.data.normal_(mean=0.0, std=1.0)
+
+        # By default, this will be a blip_2_qformer config
+        self.qformer = AutoModel.from_config(config.projector_config)
+        self.linear = nn.Linear(config.projector_config.hidden_size, config.text_config.hidden_size)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, seq_len, dim = hidden_states.size()
+        nblocks = math.ceil(seq_len / self.window_size)
+        pad = nblocks * self.window_size - seq_len
+        hidden_states = nn.functional.pad(hidden_states, (0, 0, 0, pad), "constant", 0)
+        hidden_states = hidden_states.view(batch_size * nblocks, self.window_size, dim)
+
+        query_output = self.qformer(
+            query_embeds=self.query,
+            encoder_hidden_states=hidden_states,
+            encoder_attention_mask=None,
+            return_dict=True,
+        )
+        query_proj = self.linear(
+            query_output.last_hidden_state.view(batch_size, nblocks * self.window_size // self.downsample_rate, -1)
+        )
+        return query_proj
+
+
+@auto_docstring
+class GraniteSpeechPlusPreTrainedModel(PreTrainedModel):
+    config: GraniteSpeechPlusConfig
+    input_modalities = ("audio", "text")
+
+    _supports_flash_attn = False  # `blip_2_qformer` dependency does not allow for this
+    _supports_sdpa = True
+
+    @torch.no_grad()
+    def _init_weights(self, module: nn.Module):
+        """Initialize the weights."""
+        super()._init_weights(module)
+        if isinstance(module, GraniteSpeechPlusEncoderProjector):
+            init.normal_(module.query)
+        elif isinstance(module, GraniteSpeechPlusCTCEncoder):
+            context_size = module.config.context_size
+            seq = torch.arange(context_size)
+            relpos_dist = seq.view(-1, 1) - seq.view(1, -1)
+            attention_dists = torch.clamp(relpos_dist, -context_size, context_size) + module.config.max_pos_emb
+            init.copy_(module.attention_dists, attention_dists)
+
+
+### Encoder - conformer is adapted from: https://github.com/lucidrains/conformer.git
+class GraniteSpeechPlusConformerFeedForward(nn.Module):
+    """Feedforward module for conformer encoder blocks."""
+
+    def __init__(self, config: GraniteSpeechPlusEncoderConfig):
+        super().__init__()
+        self.pre_norm = nn.LayerNorm(config.hidden_dim)
+        self.up_proj = nn.Linear(config.hidden_dim, config.hidden_dim * config.feedforward_mult)
+        self.silu = nn.SiLU()
+        self.dropout = nn.Dropout(config.dropout)
+        self.down_proj = nn.Linear(config.hidden_dim * config.feedforward_mult, config.hidden_dim)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.pre_norm(hidden_states)
+        hidden_states = self.up_proj(hidden_states)
+        hidden_states = self.dropout(self.silu(hidden_states))
+        hidden_states = self.down_proj(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        return hidden_states
+
+
+class GraniteSpeechPlusConformerAttention(nn.Module):
+    """Attention for conformer blocks using Shaw's relative positional embeddings.
+    See the following [paper](https://huggingface.co/papers/1803.02155) for more details.
+    """
+
+    def __init__(self, config: GraniteSpeechPlusEncoderConfig):
+        super().__init__()
+
+        inner_dim = config.dim_head * config.num_heads
+        self.max_pos_emb = config.max_pos_emb
+        self.context_size = config.context_size
+        self.num_heads = config.num_heads
+        self.dim_head = config.dim_head
+        self.scale = self.dim_head**-0.5
+        self.pre_norm = nn.LayerNorm(config.hidden_dim)
+        self.to_q = nn.Linear(config.hidden_dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(config.hidden_dim, inner_dim * 2, bias=False)
+        self.to_out = nn.Linear(inner_dim, config.hidden_dim)
+        self.rel_pos_emb = nn.Embedding(2 * self.max_pos_emb + 1, self.dim_head)
+        self.dropout = nn.Dropout(config.dropout)
+
+        if self.context_size <= 0 or self.context_size > self.max_pos_emb:
+            raise ValueError("Context size is either less than 0 or exceeds the max_pos_emb")
+
+    def forward(self, hidden_states: torch.Tensor, attention_dists: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.pre_norm(hidden_states)
+        bsz, num_features, _ = hidden_states.shape
+
+        num_blocks = math.ceil(num_features / self.context_size)
+        remainder = num_features % self.context_size
+        if remainder > 0:
+            # right padding to reach block size
+            hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, self.context_size - remainder))
+
+        query_states = self.to_q(hidden_states)
+        key_states, value_states = self.to_kv(hidden_states).chunk(2, dim=-1)
+
+        query_states = query_states.reshape(bsz, num_blocks, self.context_size, self.num_heads, -1).transpose(2, 3)
+        key_states = key_states.reshape(bsz, num_blocks, self.context_size, self.num_heads, -1).transpose(2, 3)
+        value_states = value_states.reshape(bsz, num_blocks, self.context_size, self.num_heads, -1).transpose(2, 3)
+
+        # shaw's relative positional embedding
+        rel_pos_emb = self.rel_pos_emb(attention_dists)
+        # alternative computation of `pos_attn` - for readability
+        # rel_pos_emb_expanded = rel_pos_emb.view([1, 1, 1] + list(rel_pos_emb.shape))
+        # pos_attn = torch.sum(query_states.unsqueeze(-2) * rel_pos_emb_expanded, dim=-1) * self.scale
+        # einsum implementation of pos_attn - gives x30 speedup over the alternative
+        # TODO (@avihu111) find a fast alternative to einsum
+        pos_attn = torch.einsum("b m h c d, c r d -> b m h c r", query_states, rel_pos_emb) * self.scale
+
+        if remainder > 0:
+            # masked attention in the extended block
+            mask = torch.ones(self.context_size, self.context_size, dtype=bool, device=hidden_states.device)
+            mask[:remainder, :remainder] = 0
+            mask_value = -torch.finfo(pos_attn.dtype).max
+            pos_attn[:, -1, :].masked_fill_(mask, mask_value)
+
+        with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
+            out = F.scaled_dot_product_attention(
+                query_states, key_states, value_states, attn_mask=pos_attn, scale=self.scale
+            )
+        out = out.transpose(2, 3).reshape(bsz, hidden_states.shape[1], -1)
+        out = self.to_out(out[:, :num_features, :])
+        return self.dropout(out)
+
+
+class GraniteSpeechPlusConformerDepthWiseConv1d(nn.Module):
+    """Wrapper for padded 1D pointwise convolution."""
+
+    def __init__(self, chan_in: int, chan_out: int, kernel_size: int):
+        super().__init__()
+        # Padding for the 1D conv is symmetric or close (i.e., offset by one).
+        pad = kernel_size // 2
+        pad_offset = (kernel_size + 1) % 2
+        self.padding = (pad, pad - pad_offset)
+
+        self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, groups=chan_in, bias=False)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = F.pad(hidden_states, self.padding)
+        return self.conv(hidden_states)
+
+
+class GraniteSpeechPlusConformerConvModule(nn.Module):
+    """Conformer conv module consisting of several 1D/depthwise 1D convolutional layers."""
+
+    def __init__(self, config: GraniteSpeechPlusEncoderConfig):
+        super().__init__()
+        inner_dim = config.hidden_dim * config.conv_expansion_factor
+
+        self.norm = nn.LayerNorm(config.hidden_dim)
+        self.up_conv = nn.Conv1d(config.hidden_dim, inner_dim * 2, 1)
+        self.glu = nn.GLU(dim=1)
+        self.depth_conv = GraniteSpeechPlusConformerDepthWiseConv1d(
+            inner_dim,
+            inner_dim,
+            kernel_size=config.conv_kernel_size,
+        )
+        self.silu = nn.SiLU()
+        self.batch_norm = nn.BatchNorm1d(inner_dim)
+        self.down_conv = nn.Conv1d(inner_dim, config.hidden_dim, 1)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.norm(hidden_states)
+        hidden_states = self.up_conv(hidden_states.permute(0, 2, 1))
+        hidden_states = self.glu(hidden_states)
+        hidden_states = self.depth_conv(hidden_states)
+        hidden_states = self.silu(self.batch_norm(hidden_states))
+        hidden_states = self.down_conv(hidden_states).permute(0, 2, 1)
+        hidden_states = self.dropout(hidden_states)
+        return hidden_states
+
+
+class GraniteSpeechPlusConformerBlock(nn.Module):
+    """Conformer block, consisting largely of linear layers, attention, and convolutional layers."""
+
+    def __init__(self, config: GraniteSpeechPlusEncoderConfig):
+        super().__init__()
+        self.ff1 = GraniteSpeechPlusConformerFeedForward(config)
+        self.attn = GraniteSpeechPlusConformerAttention(config)
+        self.conv = GraniteSpeechPlusConformerConvModule(config)
+        self.ff2 = GraniteSpeechPlusConformerFeedForward(config)
+        self.post_norm = nn.LayerNorm(config.hidden_dim)
+
+    def forward(self, hidden_states: torch.Tensor, attention_dists: torch.Tensor) -> torch.Tensor:
+        hidden_states = 0.5 * self.ff1(hidden_states) + hidden_states
+        hidden_states = self.attn(hidden_states, attention_dists=attention_dists) + hidden_states
+        hidden_states = self.conv(hidden_states) + hidden_states
+        hidden_states = 0.5 * self.ff2(hidden_states) + hidden_states
+        hidden_states = self.post_norm(hidden_states)
+        return hidden_states
+
+
+class GraniteSpeechPlusCTCEncoder(GraniteSpeechPlusPreTrainedModel):
+    config: GraniteSpeechPlusEncoderConfig
+    input_modalities = "audio"
+    _can_record_outputs = {
+        "hidden_states": GraniteSpeechPlusConformerBlock,
+        "attentions": GraniteSpeechPlusConformerAttention,
+    }
+
+    def __init__(self, config: GraniteSpeechPlusEncoderConfig):
+        super().__init__(config)
+
+        # Precompute clamped relative positional encoding distances
+        seq = torch.arange(config.context_size)
+        relpos_dist = seq.view(-1, 1) - seq.view(1, -1)
+        attention_dists = torch.clamp(relpos_dist, -config.context_size, config.context_size) + config.max_pos_emb
+        self.register_buffer("attention_dists", attention_dists, persistent=False)
+        self.input_linear = nn.Linear(config.input_dim, config.hidden_dim, bias=True)
+        self.layers = nn.ModuleList([GraniteSpeechPlusConformerBlock(config) for _ in range(config.num_layers)])
+
+        self.out = nn.Linear(config.hidden_dim, config.output_dim, bias=True)
+        self.out_mid = nn.Linear(config.output_dim, config.hidden_dim, bias=True)
+        self.num_layers = config.num_layers
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        hidden_states = self.input_linear(hidden_states)
+        cat_layers = set(self.config.cat_hidden_layers or [])
+        exported_hidden_states = []
+
+        if 0 in cat_layers:
+            exported_hidden_states.append(hidden_states)
+
+        for idx, layer in enumerate(self.layers, start=1):
+            hidden_states = layer(hidden_states, attention_dists=self.attention_dists)
+
+            if idx in cat_layers:
+                exported_hidden_states.append(hidden_states)
+
+            if idx == self.num_layers // 2:
+                hidden_states_mid = hidden_states.clone()
+                hidden_states_mid = self.out(hidden_states_mid)
+                hidden_states += self.out_mid(nn.Softmax(dim=-1)(hidden_states_mid))
+
+        if exported_hidden_states:
+            hidden_states = torch.cat([*exported_hidden_states, hidden_states], dim=-1)
+
+        return BaseModelOutputWithPooling(last_hidden_state=hidden_states)
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for LlavaNext causal language model (or autoregressive) outputs.
+    """
+)
+@dataclass
+class GraniteSpeechPlusCausalLMOutputWithPast(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+        Language modeling loss (for next-token prediction).
+    logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+        Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+    past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+        It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+        Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see
+        `past_key_values` input) to speed up sequential decoding.
+    """
+
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    past_key_values: Cache | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    attentions: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    The Granite Speech Plus model, a Granite Speech variant whose projector consumes the concatenation of the
+    encoder's final hidden states with an arbitrary subset of its intermediate hidden states.
+    """
+)
+class GraniteSpeechPlusForConditionalGeneration(GraniteSpeechPlusPreTrainedModel, GenerationMixin):
+    _supports_attention_backend = True
+
+    def __init__(self, config: GraniteSpeechPlusConfig):
+        super().__init__(config)
+        # NOTE: It doesn't matter when we initialize from config, but we should be careful
+        # to make sure this does not pick up the adapter_config if in the future we use
+        # from_pretrained or something similar, since that should be set by the composite
+        # model; don't need to consider it twice
+        self.language_model = AutoModelForCausalLM.from_config(config.text_config)
+
+        self.encoder = GraniteSpeechPlusCTCEncoder(config.encoder_config)
+        self.projector = GraniteSpeechPlusEncoderProjector(config)
+
+        if config.has_lora_adapter and not is_peft_available():
+            logger.warning(
+                "Config indicates that a lora adapter should be present, but "
+                "peft is not installed; this will cause the model to perform "
+                "incorrectly when audio inputs are provided. Please install "
+                "peft and reload the model!"
+            )
+
+        self.post_init()
+
+    def set_decoder(self, decoder):
+        self.language_model.set_decoder(decoder)
+
+    def get_decoder(self):
+        return self.language_model.get_decoder()
+
+    def set_output_embeddings(self, new_embeddings):
+        self.language_model.set_output_embeddings(new_embeddings)
+
+    def get_output_embeddings(self):
+        return self.language_model.get_output_embeddings()
+
+    @can_return_tuple
+    @auto_docstring
+    def get_audio_features(
+        self, input_features: torch.Tensor, **kwargs: Unpack[TransformersKwargs]
+    ) -> tuple | BaseModelOutputWithPooling:
+        audio_outputs = self.encoder(input_features, return_dict=True, **kwargs)
+        projected_embeds = self.projector(audio_outputs.last_hidden_state)
+        audio_outputs.pooler_output = projected_embeds
+
+        return audio_outputs
+
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        input_features: torch.FloatTensor | None = None,
+        input_features_mask: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **lm_kwargs,
+    ) -> tuple[torch.Tensor] | GraniteSpeechPlusCausalLMOutputWithPast:
+        r"""
+        input_features_mask (`torch.Tensor`, *optional*):
+            Mask to be applied to audio features prior to scattering into the language embeddings.
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+        """
+        # TODO (@alex-jw-brooks) add an example to this docstring once models are released
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if input_features is not None and inputs_embeds is not None:
+            raise ValueError(
+                "You cannot specify both input_features and inputs_embeds at the same time, and must specify either one"
+            )
+
+        if inputs_embeds is None:
+            # Get the base embeddings; set all audio tokens to 0 index
+            # to avoid out of vocabulary issues with the LLM embedding.
+            # Audio features will be masked into is_audio_idx indices later.
+            is_audio_idx = input_ids == self.config.audio_token_id
+            llm_input_ids = input_ids.clone()
+            llm_input_ids[is_audio_idx] = 0
+            inputs_embeds = self.get_input_embeddings()(llm_input_ids)
+
+        if input_features is not None:
+            if input_features.dtype != self.dtype:
+                input_features = input_features.to(self.dtype)
+            # Get the audio features from the encoder / projector
+            audio_embeds = self.get_audio_features(input_features, return_dict=True).pooler_output
+
+            # Merge the audio features into the LLM embeddings
+            inputs_embeds = self.get_merged_audio_embeddings(
+                input_ids=input_ids,
+                audio_features=audio_embeds,
+                input_features_mask=input_features_mask,
+            )
+
+        outputs = self.language_model(
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            logits_to_keep=logits_to_keep,
+            **lm_kwargs,
+        )
+        logits = outputs[0]
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            if attention_mask is not None:
+                # we use the input attention mask to shift the logits and labels, because it is 2D.
+                # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
+                shift_attention_mask = attention_mask[:, -(logits.shape[1] - 1) :].to(logits.device)
+                shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous()
+                shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous()
+            else:
+                shift_logits = logits[..., :-1, :].contiguous()
+                shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1).to(shift_logits.device)
+            )
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+
+        return GraniteSpeechPlusCausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        inputs_embeds=None,
+        input_features=None,
+        attention_mask=None,
+        logits_to_keep=None,
+        is_first_iteration=False,
+        **kwargs,
+    ):
+        # Overwritten -- in specific circumstances we don't want to forward audio inputs to the model
+
+        model_inputs = self.language_model.prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            logits_to_keep=logits_to_keep,
+            is_first_iteration=is_first_iteration,
+            **kwargs,
+        )
+
+        # If we're in cached decoding stage, input_features should be None because
+        # input ids do not contain special audio token anymore Otherwise we need
+        # input feature values to be passed to the model
+        if is_first_iteration or not kwargs.get("use_cache", True):
+            model_inputs["input_features"] = input_features
+        return model_inputs
+
+    def get_placeholder_mask(
+        self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, audio_features: torch.FloatTensor
+    ):
+        """
+        Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+        equal to the length of multimodal features. If the lengths are different, an error is raised.
+        """
+        if input_ids is None:
+            special_audio_mask = inputs_embeds == self.get_input_embeddings()(
+                torch.tensor(self.config.audio_token_id, dtype=torch.long, device=inputs_embeds.device)
+            )
+            special_audio_mask = special_audio_mask.all(-1)
+        else:
+            special_audio_mask = input_ids == self.config.audio_token_id
+
+        n_audio_tokens = special_audio_mask.sum()
+        n_audio_features = audio_features.shape[0]
+        special_audio_mask = special_audio_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+        torch_compilable_check(
+            inputs_embeds[special_audio_mask].numel() == audio_features.numel(),
+            f"Audio features and audio tokens do not match, tokens: {n_audio_tokens}, features: {n_audio_features}",
+        )
+        return special_audio_mask
+
+    def get_merged_audio_embeddings(
+        self, input_ids: torch.Tensor, audio_features: torch.Tensor, input_features_mask: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        """
+        Adds the audio token to the model's LLM vocabulary so that we can pass it
+        through the tokenizer; it's assumed that the embeddings corresponding to the
+        <|audio|> token will be clobbered with speech features.
+
+        Args:
+            input_ids (`torch.Tensor`):
+                Input IDs containing one or more audio tokens.
+            audio_features (`torch.Tensor`):
+                Audio features to be masked into the language embeddings to form multimodal embeddings.
+            input_features_mask (`torch.Tensor`, *optional*, defaults to `None`)
+                Mask to be applied to audio features prior to scattering into the language embeddings.
+        """
+        is_audio_index = input_ids == self.config.audio_token_id
+        llm_input_ids = torch.where(is_audio_index, 0, input_ids)
+        inputs_embeds = self.language_model.get_input_embeddings()(llm_input_ids)  # [bsz, # features, hidden size]
+
+        audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype)
+        if input_features_mask is not None:
+            audio_features = audio_features[input_features_mask]
+
+        special_audio_mask = self.get_placeholder_mask(
+            input_ids, inputs_embeds=inputs_embeds, audio_features=audio_features
+        )
+        inputs_embeds = inputs_embeds.masked_scatter(special_audio_mask, audio_features)
+        return inputs_embeds
+
+    def generate(self, *args, **kwargs) -> torch.LongTensor:
+        # This model is expected to have a lora adapter, which is only
+        # enabled when considering audio inputs. As such, we override generate
+        # to conditionally enable / disable the lora adapter based on whether
+        # or not any input features were provided.
+
+        input_features = kwargs.pop("input_features", None)
+        if is_peft_available and self._hf_peft_config_loaded:
+            if input_features is not None:
+                self.enable_adapters()
+            else:
+                self.disable_adapters()
+        return super().generate(*args, input_features=input_features, **kwargs)
+
+    def save_pretrained(self, save_directory, *args, **kwargs):
+        # overwrite save_pretrained to first save the adapter if we have one
+        if is_peft_available and self._hf_peft_config_loaded:
+            adapter_name = self._get_adapter_name()
+            self.peft_config[adapter_name].base_model_name_or_path = save_directory
+            super().save_pretrained(save_directory, *args, **kwargs)
+        # Then save the base model afterwards
+        prev_val = self._hf_peft_config_loaded
+        self._hf_peft_config_loaded = False
+        super().save_pretrained(save_directory, *args, **kwargs)
+        self._hf_peft_config_loaded = prev_val
+
+    def _get_adapter_name(self):
+        return list(self.peft_config.keys())[0]
+
+
+__all__ = [
+    "GraniteSpeechPlusCTCEncoder",
+    "GraniteSpeechPlusForConditionalGeneration",
+    "GraniteSpeechPlusPreTrainedModel",
+]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/granite_speech_plus/modular_granite_speech_plus.py

@@ -0,0 +1,172 @@
+# Copyright 2026 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Granite Speech Plus model, a Granite Speech variant whose projector consumes the concatenation of the
+encoder's final hidden states with an arbitrary subset of its intermediate hidden states."""
+
+import torch
+from huggingface_hub.dataclasses import strict
+from torch import nn
+
+from ...modeling_outputs import BaseModelOutputWithPooling
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from ..granite_speech.configuration_granite_speech import GraniteSpeechConfig, GraniteSpeechEncoderConfig
+from ..granite_speech.modeling_granite_speech import (
+    GraniteSpeechCTCEncoder,
+    GraniteSpeechForConditionalGeneration,
+    GraniteSpeechPreTrainedModel,
+)
+
+
+@auto_docstring(checkpoint="ibm-granite/granite-speech-4.1-2b-plus")
+@strict
+class GraniteSpeechPlusEncoderConfig(GraniteSpeechEncoderConfig):
+    r"""
+    feedforward_mult (`int`, *optional*, defaults to 4):
+        Multiplier for the up/down projections in the encoder's feedforward layers;
+        The projections will have intermediate dim of size `hidden_dim * feedforward_mult`.
+    output_dim (`int`, *optional*, defaults to 42):
+        Intermediate dimension of the feedforward projections in the conformer
+        to be added to every other encoder block's output.
+    context_size (`int`, *optional*, defaults to 200):
+        Context size to be used in conformer attention.
+    max_pos_emb (`int`, *optional*, defaults to 512):
+        Max pos embeds to be used in attention (shaw's relative positional encoding).
+    conv_expansion_factor (`int`, *optional*, defaults to 2):
+        Intermediate dimension to be used in conformer convolutions.
+    cat_hidden_layers (`list[int]`, *optional*):
+        Indices of encoder conformer layers whose outputs are concatenated with the final encoder
+        output (along the feature dimension) before being passed to the projector. When set, the
+        projector's ``encoder_hidden_size`` must equal
+        ``encoder_config.hidden_dim * (len(cat_hidden_layers) + 1)``.
+
+    Example:
+
+    ```python
+    >>> from transformers import GraniteSpeechPlusEncoderConfig, GraniteSpeechPlusCTCEncoder
+
+    >>> # Initializing a GraniteSpeechPlusEncoderConfig
+    >>> configuration = GraniteSpeechPlusEncoderConfig()
+
+    >>> # Initializing a GraniteSpeechPlusCTCEncoder (with random weights)
+    >>> model = GraniteSpeechPlusCTCEncoder(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    cat_hidden_layers: list[int] | None = None
+
+
+@auto_docstring(checkpoint="ibm-granite/granite-speech-4.1-2b-plus")
+@strict
+class GraniteSpeechPlusConfig(GraniteSpeechConfig):
+    r"""
+    projector_config (`Union[AutoConfig, dict]`, *optional*, defaults to `Blip2QFormerConfig`):
+        The config object or dictionary of the audio projector.
+    has_lora_adapter (`bool`, *optional*, defaults to `True`):
+        Indicates whether or not the model has a lora adapter that should only
+        be activate when processing audio inputs.
+    downsample_rate (`int`, *optional*, defaults to 5):
+        Downsample rate for the audio feature extractor.
+    window_size (`int`, *optional*, defaults to 15):
+        Window size for the audio feature projector.
+
+    Example:
+
+    ```python
+    >>> from transformers import GraniteSpeechPlusConfig, GraniteSpeechPlusForConditionalGeneration
+
+    >>> # Initializing a GraniteSpeechPlusConfig
+    >>> configuration = GraniteSpeechPlusConfig()
+
+    >>> # Initializing a GraniteSpeechPlusForConditionalGeneration (with random weights)
+    >>> model = GraniteSpeechPlusForConditionalGeneration(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    def __post_init__(self, **kwargs):
+        super().__post_init__(**kwargs)
+
+        if self.encoder_config.cat_hidden_layers is not None:
+            for idx in self.encoder_config.cat_hidden_layers:
+                if idx < 0 or idx >= self.encoder_config.num_layers:
+                    raise ValueError(
+                        f"cat_hidden_layers index {idx} is out of range [0, {self.encoder_config.num_layers})."
+                    )
+        if self.encoder_config.cat_hidden_layers is not None:
+            num_concat = len(self.encoder_config.cat_hidden_layers) + 1
+            if self.projector_config.encoder_hidden_size != self.encoder_config.hidden_dim * num_concat:
+                raise ValueError(
+                    f"projector encoder_hidden_size {self.projector_config.encoder_hidden_size} "
+                    f"must equal encoder hidden_dim * {num_concat} = "
+                    f"{self.encoder_config.hidden_dim * num_concat}."
+                )
+
+
+class GraniteSpeechPlusPreTrainedModel(GraniteSpeechPreTrainedModel): ...
+
+
+class GraniteSpeechPlusCTCEncoder(GraniteSpeechCTCEncoder):
+    @merge_with_config_defaults
+    @capture_outputs
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPooling:
+        hidden_states = self.input_linear(hidden_states)
+        cat_layers = set(self.config.cat_hidden_layers or [])
+        exported_hidden_states = []
+
+        if 0 in cat_layers:
+            exported_hidden_states.append(hidden_states)
+
+        for idx, layer in enumerate(self.layers, start=1):
+            hidden_states = layer(hidden_states, attention_dists=self.attention_dists)
+
+            if idx in cat_layers:
+                exported_hidden_states.append(hidden_states)
+
+            if idx == self.num_layers // 2:
+                hidden_states_mid = hidden_states.clone()
+                hidden_states_mid = self.out(hidden_states_mid)
+                hidden_states += self.out_mid(nn.Softmax(dim=-1)(hidden_states_mid))
+
+        if exported_hidden_states:
+            hidden_states = torch.cat([*exported_hidden_states, hidden_states], dim=-1)
+
+        return BaseModelOutputWithPooling(last_hidden_state=hidden_states)
+
+
+@auto_docstring(
+    custom_intro="""
+    The Granite Speech Plus model, a Granite Speech variant whose projector consumes the concatenation of the
+    encoder's final hidden states with an arbitrary subset of its intermediate hidden states.
+    """
+)
+class GraniteSpeechPlusForConditionalGeneration(GraniteSpeechForConditionalGeneration): ...
+
+
+__all__ = [
+    "GraniteSpeechPlusConfig",
+    "GraniteSpeechPlusEncoderConfig",
+    "GraniteSpeechPlusCTCEncoder",
+    "GraniteSpeechPlusForConditionalGeneration",
+    "GraniteSpeechPlusPreTrainedModel",
+]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/granitemoe/configuration_granitemoe.py

@@ -0,0 +1,84 @@
+# Copyright 2024 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""GraniteMoe model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="ibm-granite/granite-speech-3.2-8b")
+@strict
+class GraniteMoeConfig(PreTrainedConfig):
+    r"""
+    ```python
+    >>> from transformers import GraniteMoeModel, GraniteMoeConfig
+
+    >>> # Initializing a GraniteMoe granitemoe-3b style configuration
+    >>> configuration = GraniteMoeConfig()
+
+    >>> # Initializing a model from the granitemoe-7b style configuration
+    >>> model = GraniteMoeModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "granitemoe"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    vocab_size: int = 32000
+    hidden_size: int = 4096
+    intermediate_size: int = 11008
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 32
+    num_key_value_heads: int | None = None
+    hidden_act: str = "silu"
+    max_position_embeddings: int = 2048
+    initializer_range: float = 0.02
+    rms_norm_eps: float = 1e-6
+    use_cache: bool = True
+    pad_token_id: int | None = None
+    bos_token_id: int | None = 1
+    eos_token_id: int | list[int] | None = 2
+    tie_word_embeddings: bool = False
+    rope_parameters: RopeParameters | dict | None = None
+    attention_bias: bool = False
+    attention_dropout: float | int | None = 0.0
+    embedding_multiplier: float | int | None = 1.0
+    logits_scaling: float | int | None = 1.0
+    residual_multiplier: float | int | None = 1.0
+    attention_multiplier: float | int | None = 1.0
+    num_local_experts: int | None = 8
+    num_experts_per_tok: int | None = 2
+    output_router_logits: bool | None = False
+    router_aux_loss_coef: float | None = 0.001
+
+    def __post_init__(self, **kwargs):
+        # for backward compatibility
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+
+        super().__post_init__(**kwargs)
+
+
+__all__ = ["GraniteMoeConfig"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/granitemoe/modular_granitemoe.py

@@ -0,0 +1,316 @@
+# Copyright 2024 IBM and the HuggingFace Inc. team. All rights reserved.
+#
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...masking_utils import create_causal_mask
+from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from ..granite.modeling_granite import GraniteRMSNorm, GraniteRotaryEmbedding
+from ..jetmoe.modeling_jetmoe import JetMoeParallelExperts, JetMoeTopKGating
+from ..llama.modeling_llama import LlamaAttention, LlamaPreTrainedModel
+from ..mixtral.modeling_mixtral import MixtralDecoderLayer, MixtralForCausalLM, MixtralModel, load_balancing_loss_func
+from .configuration_granitemoe import GraniteMoeConfig
+
+
+class GraniteMoeRMSNorm(GraniteRMSNorm):
+    pass
+
+
+class GraniteMoeRotaryEmbedding(GraniteRotaryEmbedding):
+    pass
+
+
+class GraniteMoeParallelExperts(JetMoeParallelExperts):
+    pass
+
+
+class GraniteMoeTopKGating(JetMoeTopKGating):
+    pass
+
+
+class GraniteMoeMoE(nn.Module):
+    """
+    A Sparsely gated mixture of experts layer with 1-layer Feed-Forward networks as experts.
+
+    Args:
+        config:
+            Configuration object with model hyperparameters.
+    """
+
+    def __init__(self, config: GraniteMoeConfig):
+        super().__init__()
+
+        self.input_size = config.hidden_size
+        self.hidden_size = config.intermediate_size
+        self.activation = ACT2FN[config.hidden_act]
+        self.input_linear = GraniteMoeParallelExperts(config.num_local_experts, self.input_size, self.hidden_size * 2)
+        self.output_linear = GraniteMoeParallelExperts(config.num_local_experts, self.hidden_size, self.input_size)
+
+        self.router = GraniteMoeTopKGating(
+            input_size=self.input_size,
+            num_experts=config.num_local_experts,
+            top_k=config.num_experts_per_tok,
+        )
+
+    def forward(self, layer_input):
+        bsz, length, emb_size = layer_input.size()
+        layer_input = layer_input.reshape(-1, emb_size)
+        _, batch_index, batch_gates, expert_size, _ = self.router(layer_input)
+
+        expert_inputs = layer_input[batch_index]
+        hidden_states = self.input_linear(expert_inputs, expert_size)
+        chunked_hidden_states = hidden_states.chunk(2, dim=-1)
+        hidden_states = self.activation(chunked_hidden_states[0]) * chunked_hidden_states[1]
+        expert_outputs = self.output_linear(hidden_states, expert_size)
+
+        expert_outputs = expert_outputs * batch_gates[:, None]
+
+        zeros = torch.zeros((bsz * length, self.input_size), dtype=expert_outputs.dtype, device=expert_outputs.device)
+        layer_output = zeros.index_add(0, batch_index, expert_outputs)
+        layer_output = layer_output.view(bsz, length, self.input_size)
+        return layer_output
+
+
+class GraniteMoeAttention(LlamaAttention):
+    def __init__(self, config: GraniteMoeConfig, layer_idx: int):
+        super().__init__(self, config, layer_idx)
+        self.scaling = config.attention_multiplier  # Only diff with llama
+
+
+class GraniteMoeDecoderLayer(MixtralDecoderLayer):
+    def __init__(self, config: GraniteMoeConfig, layer_idx: int):
+        super().__init__(config, layer_idx)
+        self.self_attn = GraniteMoeAttention(config=config, layer_idx=layer_idx)
+        self.block_sparse_moe = GraniteMoeMoE(config)
+        self.input_layernorm = GraniteMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = GraniteMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        del self.mlp
+        self.block_sparse_moe = GraniteMoeMoE(config)
+        self.residual_multiplier = config.residual_multiplier  # Only diff with mixtral!
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states * self.residual_multiplier  # diff
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.block_sparse_moe(hidden_states)
+        hidden_states = residual + hidden_states * self.residual_multiplier  # diff
+        return hidden_states
+
+
+@auto_docstring
+class GraniteMoePreTrainedModel(LlamaPreTrainedModel, PreTrainedModel):
+    config: GraniteMoeConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["GraniteMoeDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _can_compile_fullgraph = False  # TopK gating fails fullgraph compilation at "expert_size = expert_size.tolist()"
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        PreTrainedModel._init_weights(self, module)
+        if isinstance(module, GraniteMoeParallelExperts):
+            init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+
+
+@auto_docstring
+class GraniteMoeModel(MixtralModel):
+    def __init__(self, config: GraniteMoeConfig):
+        super().__init__(config)
+        self.layers = nn.ModuleList(
+            [GraniteMoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = GraniteMoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.embedding_multiplier = config.embedding_multiplier
+
+    @merge_with_config_defaults
+    @capture_outputs
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> MoeModelOutputWithPast:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if position_ids is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+            position_ids = position_ids.unsqueeze(0)
+
+        causal_mask = create_causal_mask(  # ONLY DIFF WITH MIXTRAL: NO SLIDING
+            config=self.config,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            position_ids=position_ids,
+        )
+        inputs_embeds = inputs_embeds * self.embedding_multiplier
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                position_embeddings=position_embeddings,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+
+        return MoeModelOutputWithPast(  # only diff with Mistral is the output type, we need MoE
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+class GraniteMoeForCausalLM(MixtralForCausalLM):
+    def __init__(self, config: GraniteMoeConfig):
+        super().__init__(config)
+        self.model = GraniteMoeModel(config)
+        self.logits_scaling = config.logits_scaling
+
+    @auto_docstring
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        output_router_logits: bool | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs,
+    ) -> tuple | MoeCausalLMOutputWithPast:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, GraniteMoeForCausalLM
+
+        >>> model = GraniteMoeForCausalLM.from_pretrained("ibm/PowerMoE-3b")
+        >>> tokenizer = AutoTokenizer.from_pretrained("ibm/PowerMoE-3b")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        output_router_logits = (
+            output_router_logits if output_router_logits is not None else self.config.output_router_logits
+        )
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            **kwargs,
+        )
+
+        # Only compute necessary logits
+        hidden_states = outputs.last_hidden_state
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        logits = logits / self.config.logits_scaling
+
+        loss = None
+        if labels is not None:
+            # Flatten the tokens
+            loss = self.loss_function(
+                logits,
+                labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+
+        aux_loss = None
+        if output_router_logits:
+            aux_loss = load_balancing_loss_func(
+                outputs.router_logits,
+                self.num_experts,
+                self.num_experts_per_tok,
+                attention_mask,
+            )
+            if labels is not None:
+                loss += self.router_aux_loss_coef * aux_loss.to(loss.device)  # make sure to reside in the same device
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            router_logits=outputs.router_logits,
+        )
+
+
+__all__ = ["GraniteMoeForCausalLM", "GraniteMoeModel", "GraniteMoePreTrainedModel"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/llava/__init__.py

@@ -0,0 +1,30 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_llava import *
+    from .image_processing_llava import *
+    from .image_processing_pil_llava import *
+    from .modeling_llava import *
+    from .processing_llava import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/pytorch_utils.py

@@ -0,0 +1,260 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+import inspect
+from collections.abc import Callable
+from functools import lru_cache, wraps
+
+import torch
+from safetensors.torch import storage_ptr, storage_size
+from torch import nn
+
+from .utils import (
+    is_torch_greater_or_equal,
+    is_torch_xla_available,
+    is_torchdynamo_compiling,
+    logging,
+)
+
+
+ALL_LAYERNORM_LAYERS = [nn.LayerNorm]
+
+logger = logging.get_logger(__name__)
+
+is_torch_greater_or_equal_than_2_8 = is_torch_greater_or_equal("2.8", accept_dev=True)
+is_torch_greater_or_equal_than_2_6 = is_torch_greater_or_equal("2.6", accept_dev=True)
+
+# For backwards compatibility (e.g. some remote codes on Hub using those variables).
+is_torch_greater_or_equal_than_2_4 = is_torch_greater_or_equal("2.4", accept_dev=True)
+is_torch_greater_or_equal_than_2_3 = is_torch_greater_or_equal("2.3", accept_dev=True)
+is_torch_greater_or_equal_than_2_2 = is_torch_greater_or_equal("2.2", accept_dev=True)
+is_torch_greater_or_equal_than_2_1 = is_torch_greater_or_equal("2.1", accept_dev=True)
+is_torch_greater_or_equal_than_2_0 = is_torch_greater_or_equal("2.0", accept_dev=True)
+is_torch_greater_or_equal_than_1_13 = is_torch_greater_or_equal("1.13", accept_dev=True)
+is_torch_greater_or_equal_than_1_12 = is_torch_greater_or_equal("1.12", accept_dev=True)
+
+# Cache this result has it's a C FFI call which can be pretty time-consuming
+_torch_distributed_available = torch.distributed.is_available()
+
+
+def softmax_backward_data(parent, grad_output, output):
+    """
+    A function that calls the internal `_softmax_backward_data` PyTorch method and that adjusts the arguments according
+    to the torch version detected.
+    """
+
+    from torch import _softmax_backward_data
+
+    return _softmax_backward_data(grad_output, output, parent.dim, output.dtype)
+
+
+def prune_linear_layer(layer: nn.Linear, index: torch.LongTensor, dim: int = 0) -> nn.Linear:
+    """
+    Prune a linear layer to keep only entries in index.
+
+    Used to remove heads.
+
+    Args:
+        layer (`torch.nn.Linear`): The layer to prune.
+        index (`torch.LongTensor`): The indices to keep in the layer.
+        dim (`int`, *optional*, defaults to 0): The dimension on which to keep the indices.
+
+    Returns:
+        `torch.nn.Linear`: The pruned layer as a new layer with `requires_grad=True`.
+    """
+    index = index.to(layer.weight.device)
+    W = layer.weight.index_select(dim, index).detach().clone()
+    if layer.bias is not None:
+        if dim == 1:
+            b = layer.bias.detach().clone()
+        else:
+            b = layer.bias[index].detach().clone()
+    new_size = list(layer.weight.size())
+    new_size[dim] = len(index)
+    new_layer = nn.Linear(new_size[1], new_size[0], bias=layer.bias is not None).to(layer.weight.device)
+    new_layer.weight.requires_grad = False
+    new_layer.weight.copy_(W.contiguous())
+    new_layer.weight.requires_grad = True
+    if layer.bias is not None:
+        new_layer.bias.requires_grad = False
+        new_layer.bias.copy_(b.contiguous())
+        new_layer.bias.requires_grad = True
+    return new_layer
+
+
+class Conv1D(nn.Module):
+    """
+    1D-convolutional layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2).
+
+    Basically works like a linear layer but the weights are transposed.
+
+    Args:
+        nf (`int`): The number of output features.
+        nx (`int`): The number of input features.
+    """
+
+    def __init__(self, nf, nx):
+        super().__init__()
+        self.nf = nf
+        self.nx = nx
+        self.weight = nn.Parameter(torch.empty(nx, nf))
+        self.bias = nn.Parameter(torch.zeros(nf))
+        nn.init.normal_(self.weight, std=0.02)
+
+    def __repr__(self) -> str:
+        return "Conv1D(nf={nf}, nx={nx})".format(**self.__dict__)
+
+    def forward(self, x):
+        size_out = x.size()[:-1] + (self.nf,)
+        x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
+        x = x.view(size_out)
+        return x
+
+
+def apply_chunking_to_forward(
+    forward_fn: Callable[..., torch.Tensor],
+    chunk_size: int,
+    chunk_dim: int,
+    *input_tensors,
+) -> torch.Tensor:
+    """
+    This function chunks the `input_tensors` into smaller input tensor parts of size `chunk_size` over the dimension
+    `chunk_dim`. It then applies a layer `forward_fn` to each chunk independently to save memory.
+
+    If the `forward_fn` is independent across the `chunk_dim` this function will yield the same result as directly
+    applying `forward_fn` to `input_tensors`.
+
+    Args:
+        forward_fn (`Callable[..., torch.Tensor]`):
+            The forward function of the model.
+        chunk_size (`int`):
+            The chunk size of a chunked tensor: `num_chunks = len(input_tensors[0]) / chunk_size`.
+        chunk_dim (`int`):
+            The dimension over which the `input_tensors` should be chunked.
+        input_tensors (`tuple[torch.Tensor]`):
+            The input tensors of `forward_fn` which will be chunked
+
+    Returns:
+        `torch.Tensor`: A tensor with the same shape as the `forward_fn` would have given if applied`.
+
+
+    Examples:
+
+    ```python
+    # rename the usual forward() fn to forward_chunk()
+    def forward_chunk(self, hidden_states):
+        hidden_states = self.decoder(hidden_states)
+        return hidden_states
+
+
+    # implement a chunked forward function
+    def forward(self, hidden_states):
+        return apply_chunking_to_forward(self.forward_chunk, self.chunk_size_lm_head, self.seq_len_dim, hidden_states)
+    ```"""
+
+    assert len(input_tensors) > 0, f"{input_tensors} has to be a tuple/list of tensors"
+
+    # inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
+    num_args_in_forward_chunk_fn = len(inspect.signature(forward_fn).parameters)
+    if num_args_in_forward_chunk_fn != len(input_tensors):
+        raise ValueError(
+            f"forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(input_tensors)} input "
+            "tensors are given"
+        )
+
+    if chunk_size > 0:
+        tensor_shape = input_tensors[0].shape[chunk_dim]
+        for input_tensor in input_tensors:
+            if input_tensor.shape[chunk_dim] != tensor_shape:
+                raise ValueError(
+                    f"All input tenors have to be of the same shape: {tensor_shape}, "
+                    f"found shape {input_tensor.shape[chunk_dim]}"
+                )
+
+        if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
+            raise ValueError(
+                f"The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk "
+                f"size {chunk_size}"
+            )
+
+        num_chunks = input_tensors[0].shape[chunk_dim] // chunk_size
+
+        # chunk input tensor into tuples
+        input_tensors_chunks = tuple(input_tensor.chunk(num_chunks, dim=chunk_dim) for input_tensor in input_tensors)
+        # apply forward fn to every tuple
+        output_chunks = tuple(forward_fn(*input_tensors_chunk) for input_tensors_chunk in zip(*input_tensors_chunks))
+        # concatenate output at same dimension
+        return torch.cat(output_chunks, dim=chunk_dim)
+
+    return forward_fn(*input_tensors)
+
+
+def meshgrid(*tensors: torch.Tensor | list[torch.Tensor], indexing: str | None = None) -> tuple[torch.Tensor, ...]:
+    """
+    Wrapper around torch.meshgrid to avoid warning messages about the introduced `indexing` argument.
+
+    Reference: https://pytorch.org/docs/1.13/generated/torch.meshgrid.html
+    """
+    return torch.meshgrid(*tensors, indexing=indexing)
+
+
+def id_tensor_storage(tensor: torch.Tensor) -> tuple[torch.device, int, int]:
+    """
+    Unique identifier to a tensor storage. Multiple different tensors can share the same underlying storage. For
+    example, "meta" tensors all share the same storage, and thus their identifier will all be equal. This identifier is
+    guaranteed to be unique and constant for this tensor's storage during its lifetime. Two tensor storages with
+    non-overlapping lifetimes may have the same id.
+    """
+    if _torch_distributed_available and is_torch_greater_or_equal("2.5"):
+        from torch.distributed.tensor import DTensor
+
+        if isinstance(tensor, DTensor):
+            local_tensor = tensor.to_local()
+            return tensor.device, local_tensor.storage().data_ptr(), tensor.nbytes
+
+    if tensor.device.type == "xla" and is_torch_xla_available():
+        # NOTE: xla tensors dont have storage
+        # use some other unique id to distinguish.
+        # this is a XLA tensor, it must be created using torch_xla's
+        # device. So the following import is safe:
+        import torch_xla
+
+        unique_id = torch_xla._XLAC._xla_get_tensor_id(tensor)
+    else:
+        unique_id = storage_ptr(tensor)
+
+    return tensor.device, unique_id, storage_size(tensor)
+
+
+@wraps(lru_cache)
+def compile_compatible_method_lru_cache(*lru_args, **lru_kwargs):
+    """
+    LRU cache decorator from standard functools library, but with a workaround to disable
+    caching when torchdynamo is compiling. Expected to work with class methods.
+    """
+
+    def decorator(func):
+        func_with_cache = lru_cache(*lru_args, **lru_kwargs)(func)
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            if is_torchdynamo_compiling():
+                return func(*args, **kwargs)
+            else:
+                return func_with_cache(*args, **kwargs)
+
+        return wrapper
+
+    return decorator

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/trainer_optimizer.py

@@ -0,0 +1,615 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Optimizer utilities for the Trainer class.
+"""
+
+from __future__ import annotations
+
+import importlib.metadata
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import torch
+from packaging import version
+from torch import nn
+
+from .optimization import Adafactor
+from .trainer_pt_utils import LayerWiseDummyOptimizer
+from .trainer_utils import check_target_module_exists
+from .training_args import OptimizerNames, ParallelMode
+from .utils import (
+    is_apollo_torch_available,
+    is_bitsandbytes_available,
+    is_galore_torch_available,
+    is_grokadamw_available,
+    is_lomo_available,
+    is_schedulefree_available,
+    is_torch_optimi_available,
+    is_torchao_available,
+    strtobool,
+)
+
+
+if TYPE_CHECKING:
+    from .modeling_utils import PreTrainedModel
+    from .training_args import TrainingArguments
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class OptimizerContext:
+    """Context object passed to all optimizer handlers."""
+
+    args: TrainingArguments
+    model: PreTrainedModel | None
+    optimizer_kwargs: dict[str, Any]
+    adam_kwargs: dict[str, Any]
+    optim_args: dict[str, str]
+
+
+def _parse_optim_args(optim_args_str: str | None) -> dict[str, str]:
+    """Parse optimizer arguments from a comma-separated string."""
+    if not optim_args_str:
+        return {}
+    optim_args = {}
+    for mapping in optim_args_str.replace(" ", "").split(","):
+        key, value = mapping.split("=")
+        optim_args[key] = value
+    return optim_args
+
+
+# Type alias for optimizer handler functions
+OptimizerHandler = Callable[[OptimizerContext], tuple[Any, dict[str, Any]]]
+
+
+def is_optimizer_factory(optimizer_cls_or_factory: Any) -> bool:
+    """
+    Check if the returned value from a handler is a factory rather than an Optimizer class.
+
+    Factory callables are used for complex optimizers like Muon or Dion that need to:
+    - Split parameters between multiple internal optimizers
+    - Handle complex sharding logic
+    - Access the full model structure for parameter grouping
+
+    Args:
+        optimizer_cls_or_factory: The first element returned by an optimizer handler.
+
+    Returns:
+        `bool`: True if it's not an Optimizer class (i.e., likely a factory), False if it's an Optimizer class.
+    """
+    # If it's a class that's a subclass of torch.optim.Optimizer, it's not a factory
+    if isinstance(optimizer_cls_or_factory, type) and issubclass(optimizer_cls_or_factory, torch.optim.Optimizer):
+        return False
+    return True
+
+
+def _setup_low_rank_optimizer(
+    args: TrainingArguments,
+    model: PreTrainedModel,
+    optimizer_name: str,
+    optimizer_mapping: dict[str, Any],
+    optim_kwargs: dict[str, Any],
+    optimizer_kwargs: dict[str, Any],
+    is_layerwise_supported: bool = True,
+) -> tuple[Any, dict[str, Any]]:
+    """
+    Helper function to set up low-rank optimizers like GaLore and Apollo.
+
+    These optimizers apply low-rank projections to specific target modules (typically linear layers).
+    """
+    is_layerwise = optimizer_name.lower().endswith("layerwise")
+    if is_layerwise and args.parallel_mode == ParallelMode.DISTRIBUTED and is_layerwise_supported:
+        raise NotImplementedError(f"Layer-wise {optimizer_name} does not support DDP at this time")
+
+    optimizer_cls = optimizer_mapping[optimizer_name]
+
+    if args.optim_target_modules is None:
+        raise ValueError(f"You need to define `optim_target_modules` to use {optimizer_name} optimizers")
+
+    if not isinstance(args.optim_target_modules, (list, str)):
+        raise TypeError(
+            f"`optim_target_modules` must be a list of strings, a regex string, or 'all-linear'. "
+            f"Got: {args.optim_target_modules}"
+        )
+
+    if model is None:
+        raise ValueError(f"You need to pass a model to initialize {optimizer_name} optimizer.")
+
+    all_linear = (
+        isinstance(args.optim_target_modules, str) and args.optim_target_modules.replace("_", "-") == "all-linear"
+    )
+
+    target_params_names = []
+    for module_name, module in model.named_modules():
+        target_module_exists, is_regex = check_target_module_exists(
+            args.optim_target_modules, module_name, return_is_regex=True
+        )
+
+        if not isinstance(module, nn.Linear):
+            if target_module_exists and not is_regex:
+                logger.warning(f"{module_name} matched but ignored. {optimizer_name} only supports linear layers.")
+            continue
+
+        if not target_module_exists and not all_linear:
+            continue
+
+        target_params_names.append(module_name + ".weight")
+
+    if len(target_params_names) == 0:
+        raise ValueError(f"No target modules found for {optimizer_name} ({args.optim_target_modules}).")
+
+    target_params = [p for n, p in model.named_parameters() if n in target_params_names]
+    non_target_params = [p for n, p in model.named_parameters() if n not in target_params_names]
+
+    param_groups = [
+        {"params": non_target_params},
+        {"params": target_params, **optim_kwargs},
+    ]
+
+    if is_layerwise:
+        if args.gradient_accumulation_steps != 1:
+            raise ValueError(f"Layerwise {optimizer_name} does not support gradient accumulation!")
+
+        optimizer_dict = {}
+        for param in non_target_params:
+            optimizer_dict[param] = optimizer_cls([{"params": [param]}], **optimizer_kwargs)
+        for param in target_params:
+            optimizer_dict[param] = optimizer_cls([{"params": [param], **optim_kwargs}], **optimizer_kwargs)
+
+        def optimizer_hook(param):
+            if param.grad is not None:
+                optimizer_dict[param].step()
+                optimizer_dict[param].zero_grad()
+
+        for param in model.parameters():
+            if param.requires_grad:
+                param.register_post_accumulate_grad_hook(optimizer_hook)
+
+        optimizer_cls = LayerWiseDummyOptimizer
+        optimizer_kwargs.update({"optimizer_dict": optimizer_dict})
+
+    optimizer_kwargs.update({"params": param_groups})
+    return optimizer_cls, optimizer_kwargs
+
+
+# =============================================================================
+# Individual optimizer handlers
+# =============================================================================
+
+
+def _get_adafactor(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get Adafactor optimizer."""
+    ctx.optimizer_kwargs.update({"scale_parameter": False, "relative_step": False})
+    return Adafactor, ctx.optimizer_kwargs
+
+
+def _get_adamw_torch(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get PyTorch AdamW optimizer (regular or fused)."""
+    from torch.optim import AdamW
+
+    ctx.optimizer_kwargs.update(ctx.adam_kwargs)
+    if ctx.args.optim == OptimizerNames.ADAMW_TORCH_FUSED:
+        ctx.optimizer_kwargs.update({"fused": True})
+    return AdamW, ctx.optimizer_kwargs
+
+
+def _get_adamw_torch_xla(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get Torch XLA syncfree AdamW optimizer."""
+    try:
+        from torch_xla.amp.syncfree import AdamW
+
+        ctx.optimizer_kwargs.update(ctx.adam_kwargs)
+        return AdamW, ctx.optimizer_kwargs
+    except ImportError:
+        raise ValueError("Trainer failed to import syncfree AdamW from torch_xla.")
+
+
+def _get_adamw_torch_npu_fused(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get NPU Fused AdamW optimizer."""
+    try:
+        from torch_npu.optim import NpuFusedAdamW
+
+        ctx.optimizer_kwargs.update(ctx.adam_kwargs)
+        return NpuFusedAdamW, ctx.optimizer_kwargs
+    except ImportError:
+        raise ValueError("Trainer failed to import FusedAdamW from torch_npu.")
+
+
+def _get_bitsandbytes_optimizer(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get bitsandbytes optimizer (AdamW, Lion, RMSprop variants)."""
+    if not is_bitsandbytes_available():
+        raise ImportError(
+            "You need to install `bitsandbytes` in order to use bitsandbytes optimizers: `pip install -U bitsandbytes`"
+        )
+
+    from bitsandbytes.optim import AdamW, Lion, RMSprop
+
+    optim_name = ctx.args.optim
+    is_paged = "paged" in optim_name
+    optim_bits = 8 if "8bit" in optim_name else 32
+    optimizer_cls = None
+    additional_optim_kwargs = ctx.adam_kwargs
+
+    if "adam" in optim_name:
+        optimizer_cls = AdamW
+    elif "lion" in optim_name:
+        optimizer_cls = Lion
+        additional_optim_kwargs = {"betas": (ctx.args.adam_beta1, ctx.args.adam_beta2)}
+    elif "rmsprop" in optim_name:
+        optimizer_cls = RMSprop
+        additional_optim_kwargs = ctx.optim_args
+    elif "ademamix" in optim_name:
+        from bitsandbytes.optim import AdEMAMix
+
+        optimizer_cls = AdEMAMix
+        additional_optim_kwargs = {
+            "betas": (
+                float(ctx.optim_args.get("beta1", ctx.args.adam_beta1)),
+                float(ctx.optim_args.get("beta2", ctx.args.adam_beta2)),
+                float(ctx.optim_args.get("beta3", 0.9999)),
+            ),
+            "alpha": float(ctx.optim_args.get("alpha", 5.0)),
+            "eps": float(ctx.optim_args.get("eps", ctx.args.adam_epsilon)),
+        }
+        if "t_alpha" in ctx.optim_args:
+            additional_optim_kwargs["t_alpha"] = int(ctx.optim_args["t_alpha"])
+        if "t_beta3" in ctx.optim_args:
+            additional_optim_kwargs["t_beta3"] = int(ctx.optim_args["t_beta3"])
+
+    bnb_kwargs = {"optim_bits": optim_bits}
+    if "rmsprop" not in optim_name:
+        bnb_kwargs["is_paged"] = is_paged
+
+    ctx.optimizer_kwargs.update(additional_optim_kwargs)
+    ctx.optimizer_kwargs.update(bnb_kwargs)
+    return optimizer_cls, ctx.optimizer_kwargs
+
+
+def _get_adamw_anyprecision(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get AnyPrecision AdamW optimizer."""
+    try:
+        from torchdistx.optimizers import AnyPrecisionAdamW
+
+        ctx.optimizer_kwargs.update(ctx.adam_kwargs)
+        ctx.optimizer_kwargs.update(
+            {
+                "use_kahan_summation": strtobool(ctx.optim_args.get("use_kahan_summation", "False")),
+                "momentum_dtype": getattr(torch, ctx.optim_args.get("momentum_dtype", "float32")),
+                "variance_dtype": getattr(torch, ctx.optim_args.get("variance_dtype", "float32")),
+                "compensation_buffer_dtype": getattr(
+                    torch, ctx.optim_args.get("compensation_buffer_dtype", "bfloat16")
+                ),
+            }
+        )
+        return AnyPrecisionAdamW, ctx.optimizer_kwargs
+    except ImportError:
+        raise ValueError("Please install https://github.com/pytorch/torchdistx")
+
+
+def _get_sgd(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get SGD optimizer."""
+    kwargs = ctx.optimizer_kwargs.copy()
+    if ctx.optim_args:
+        for key in ("momentum", "dampening", "weight_decay"):
+            if key in ctx.optim_args:
+                kwargs[key] = float(ctx.optim_args[key])
+        if "nesterov" in ctx.optim_args:
+            kwargs["nesterov"] = ctx.optim_args["nesterov"].lower() in ("true", "1", "yes")
+    return torch.optim.SGD, kwargs
+
+
+def _get_adagrad(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get Adagrad optimizer."""
+    kwargs = ctx.optimizer_kwargs.copy()
+    if ctx.optim_args:
+        for key in ("lr_decay", "weight_decay", "eps"):
+            if key in ctx.optim_args:
+                kwargs[key] = float(ctx.optim_args[key])
+    return torch.optim.Adagrad, kwargs
+
+
+def _get_rmsprop(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get RMSprop optimizer."""
+    kwargs = ctx.optimizer_kwargs.copy()
+    if ctx.optim_args:
+        for key in ("momentum", "alpha", "eps", "weight_decay"):
+            if key in ctx.optim_args:
+                kwargs[key] = float(ctx.optim_args[key])
+        if "centered" in ctx.optim_args:
+            kwargs["centered"] = ctx.optim_args["centered"].lower() in ("true", "1", "yes")
+    return torch.optim.RMSprop, kwargs
+
+
+def _get_galore_optimizer(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get GaLore optimizer."""
+    if not is_galore_torch_available():
+        raise ImportError(
+            "You need to install `galore_torch` in order to use GaLore optimizers. "
+            "Install it with `pip install git+https://github.com/jiaweizzhao/GaLore`"
+        )
+    from galore_torch import GaLoreAdafactor, GaLoreAdamW, GaLoreAdamW8bit
+
+    optimizer_mapping = {
+        OptimizerNames.GALORE_ADAMW: GaLoreAdamW,
+        OptimizerNames.GALORE_ADAMW_8BIT: GaLoreAdamW8bit,
+        OptimizerNames.GALORE_ADAFACTOR: GaLoreAdafactor,
+        OptimizerNames.GALORE_ADAMW_LAYERWISE: GaLoreAdamW,
+        OptimizerNames.GALORE_ADAMW_8BIT_LAYERWISE: GaLoreAdamW8bit,
+        OptimizerNames.GALORE_ADAFACTOR_LAYERWISE: GaLoreAdafactor,
+    }
+
+    galore_optim_kwargs = {
+        "rank": int(ctx.optim_args.pop("rank", 128)),
+        "update_proj_gap": int(ctx.optim_args.pop("update_proj_gap", 200)),
+        "scale": float(ctx.optim_args.pop("scale", 0.25)),
+        "proj_type": ctx.optim_args.pop("proj_type", "std"),
+    }
+
+    optimizer_cls, optimizer_kwargs = _setup_low_rank_optimizer(
+        ctx.args, ctx.model, ctx.args.optim, optimizer_mapping, galore_optim_kwargs, ctx.optimizer_kwargs
+    )
+    if ctx.args.optim == OptimizerNames.GALORE_ADAFACTOR:
+        optimizer_kwargs.update({"scale_parameter": False, "relative_step": False})
+    return optimizer_cls, optimizer_kwargs
+
+
+def _get_apollo_optimizer(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get Apollo optimizer."""
+    if not is_apollo_torch_available():
+        raise ImportError(
+            "You need to install `apollo_torch` in order to use APOLLO optimizers. "
+            "Install it with `pip install git+https://github.com/zhuhanqing/APOLLO`"
+        )
+    from apollo_torch import APOLLOAdamW
+
+    optimizer_mapping = {
+        OptimizerNames.APOLLO_ADAMW: APOLLOAdamW,
+        OptimizerNames.APOLLO_ADAMW_LAYERWISE: APOLLOAdamW,
+    }
+
+    apollo_optim_kwargs = {
+        "rank": int(ctx.optim_args.pop("rank", 128)),
+        "proj": ctx.optim_args.pop("proj", "random"),
+        "scale_type": ctx.optim_args.pop("scale_type", "channel"),
+        "update_proj_gap": int(ctx.optim_args.pop("update_proj_gap", 200)),
+        "scale": float(ctx.optim_args.pop("scale", 1.0)),
+        "proj_type": ctx.optim_args.pop("proj_type", "std"),
+    }
+    apollo_optim_kwargs.update(ctx.adam_kwargs)
+
+    return _setup_low_rank_optimizer(
+        ctx.args, ctx.model, ctx.args.optim, optimizer_mapping, apollo_optim_kwargs, ctx.optimizer_kwargs
+    )
+
+
+def _get_lomo_optimizer(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get LOMO optimizer."""
+    if not is_lomo_available():
+        raise ImportError(
+            "You need to install `lomo_optim` in order to use LOMO optimizers. "
+            "Install it with `pip install lomo-optim`"
+        )
+
+    if ctx.model is None:
+        raise ValueError("You need to pass a `model` in order to correctly initialize a LOMO optimizer.")
+
+    from lomo_optim import AdaLomo, Lomo
+
+    optimizer_cls = AdaLomo if "ada" in ctx.args.optim else Lomo
+    ctx.optimizer_kwargs.update({"model": ctx.model})
+    return optimizer_cls, ctx.optimizer_kwargs
+
+
+def _get_grokadamw(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get GrokAdamW optimizer."""
+    if not is_grokadamw_available():
+        raise ValueError("Please install grokadamw with `pip install grokadamw`")
+
+    from grokadamw import GrokAdamW
+
+    ctx.optimizer_kwargs.update(
+        {
+            "alpha_init": float(ctx.optim_args.get("alpha_init", 0.98)),
+            "lamb": float(ctx.optim_args.get("lamb", 2.0)),
+            "gamma": float(ctx.optim_args.get("gamma", 0.1)),
+            "grokking_signal_decay_rate": float(ctx.optim_args.get("grokking_signal_decay_rate", 0.1)),
+            "gradient_clipping": float(ctx.optim_args.get("gradient_clipping", 1.0)),
+        }
+    )
+    return GrokAdamW, ctx.optimizer_kwargs
+
+
+def _get_torchao_optimizer(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get TorchAO 4-bit or 8-bit optimizer."""
+    if not is_torchao_available() or version.parse(importlib.metadata.version("torchao")) < version.parse("0.4.0"):
+        raise ImportError(
+            "You need to have `torchao>=0.4.0` in order to use torch 4-bit optimizers. "
+            "Install it with `pip install torchao` or follow the instructions here: "
+            "https://github.com/pytorch/ao"
+        )
+    if version.parse(importlib.metadata.version("torch")) <= version.parse("2.4"):
+        raise ImportError(
+            "You need to have `torch>2.4` in order to use torch 4-bit optimizers. "
+            "Install it with `pip install --upgrade torch` it is available on pipy. "
+            "Otherwise, you need to install torch nightly."
+        )
+
+    if version.parse(importlib.metadata.version("torchao")) >= version.parse("0.11.0"):
+        from torchao.optim import AdamW4bit, AdamW8bit
+    else:
+        from torchao.prototype.low_bit_optim import AdamW4bit, AdamW8bit
+
+    if ctx.args.optim == OptimizerNames.ADAMW_TORCH_4BIT:
+        optimizer_cls = AdamW4bit
+    else:
+        optimizer_cls = AdamW8bit
+
+    ctx.optimizer_kwargs.update(
+        {
+            "block_size": ctx.optim_args.get("block_size", 256),
+            "bf16_stochastic_round": strtobool(ctx.optim_args.get("bf16_stochastic_round", "False")),
+        }
+    )
+    ctx.optimizer_kwargs.update(ctx.adam_kwargs)
+    return optimizer_cls, ctx.optimizer_kwargs
+
+
+def _get_schedule_free_optimizer(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get ScheduleFree optimizer."""
+    if not is_schedulefree_available():
+        raise ImportError(
+            "You need to install `schedulefree` in order to use schedulefree optimizers. "
+            "Install it with `pip install schedulefree.`"
+        )
+    from schedulefree import AdamWScheduleFree, SGDScheduleFree
+
+    additional_optim_kwargs = {}
+    require_warmup = True
+
+    if ctx.args.optim == OptimizerNames.SCHEDULE_FREE_RADAM:
+        if not is_schedulefree_available("1.4.0"):
+            raise ImportError(
+                "You need to install `schedulefree>=1.4.0` in order to use RAdamScheduleFree optimizer. "
+                "Install it with `pip install schedulefree.`"
+            )
+        from schedulefree import RAdamScheduleFree
+
+        optimizer_cls = RAdamScheduleFree
+        additional_optim_kwargs = ctx.adam_kwargs
+        require_warmup = False
+    elif ctx.args.optim == OptimizerNames.SCHEDULE_FREE_ADAMW:
+        optimizer_cls = AdamWScheduleFree
+        additional_optim_kwargs = ctx.adam_kwargs
+    elif ctx.args.optim == OptimizerNames.SCHEDULE_FREE_SGD:
+        optimizer_cls = SGDScheduleFree
+    else:
+        raise ValueError("Invalid schedulefree optimizer")
+
+    additional_optim_kwargs["weight_decay"] = ctx.args.weight_decay
+    if require_warmup:
+        additional_optim_kwargs["warmup_steps"] = ctx.args.warmup_steps
+    additional_optim_kwargs.update(
+        {
+            "weight_lr_power": float(ctx.optim_args.get("weight_lr_power", 2.0)),
+            "r": float(ctx.optim_args.get("r", 0.0)),
+        }
+    )
+    ctx.optimizer_kwargs.update(additional_optim_kwargs)
+    return optimizer_cls, ctx.optimizer_kwargs
+
+
+def _get_stable_adamw(ctx: OptimizerContext) -> tuple[Any, dict[str, Any]]:
+    """Get StableAdamW optimizer from torch-optimi."""
+    if not is_torch_optimi_available():
+        raise ImportError(
+            "You need to install `torch-optimi` in order to use stable_adamw optimizers. "
+            "Install it with `pip install torch-optimi`."
+        )
+    from optimi import StableAdamW
+
+    max_lr = ctx.optim_args.pop("max_lr", None)
+    if max_lr is not None:
+        max_lr = float(max_lr)
+
+    kahan_sum = ctx.optim_args.pop("kahan_sum", None)
+    if kahan_sum is not None:
+        kahan_sum = bool(kahan_sum)
+
+    ctx.adam_kwargs["weight_decay"] = ctx.args.weight_decay
+    stable_adamw_kwargs = {
+        "decouple_lr": bool(ctx.optim_args.pop("decouple_lr", False)),
+        "max_lr": max_lr,
+        "kahan_sum": kahan_sum,
+    }
+
+    ctx.optimizer_kwargs.update(ctx.adam_kwargs)
+    ctx.optimizer_kwargs.update(stable_adamw_kwargs)
+    return StableAdamW, ctx.optimizer_kwargs
+
+
+# =============================================================================
+# Dispatch table
+# =============================================================================
+
+_BITSANDBYTES_OPTIMIZERS = [
+    OptimizerNames.ADAMW_BNB,
+    OptimizerNames.ADAMW_8BIT,
+    OptimizerNames.PAGED_ADAMW,
+    OptimizerNames.PAGED_ADAMW_8BIT,
+    OptimizerNames.ADEMAMIX,
+    OptimizerNames.ADEMAMIX_8BIT,
+    OptimizerNames.PAGED_ADEMAMIX,
+    OptimizerNames.PAGED_ADEMAMIX_8BIT,
+    OptimizerNames.LION,
+    OptimizerNames.LION_8BIT,
+    OptimizerNames.PAGED_LION,
+    OptimizerNames.PAGED_LION_8BIT,
+    OptimizerNames.RMSPROP_BNB,
+    OptimizerNames.RMSPROP_8BIT,
+    OptimizerNames.RMSPROP_32BIT,
+]
+
+_GALORE_OPTIMIZERS = [
+    OptimizerNames.GALORE_ADAMW,
+    OptimizerNames.GALORE_ADAMW_8BIT,
+    OptimizerNames.GALORE_ADAFACTOR,
+    OptimizerNames.GALORE_ADAMW_LAYERWISE,
+    OptimizerNames.GALORE_ADAMW_8BIT_LAYERWISE,
+    OptimizerNames.GALORE_ADAFACTOR_LAYERWISE,
+]
+
+_APOLLO_OPTIMIZERS = [
+    OptimizerNames.APOLLO_ADAMW,
+    OptimizerNames.APOLLO_ADAMW_LAYERWISE,
+]
+
+_TORCHAO_OPTIMIZERS = [
+    OptimizerNames.ADAMW_TORCH_4BIT,
+    OptimizerNames.ADAMW_TORCH_8BIT,
+]
+
+_SCHEDULE_FREE_OPTIMIZERS = [
+    OptimizerNames.SCHEDULE_FREE_RADAM,
+    OptimizerNames.SCHEDULE_FREE_ADAMW,
+    OptimizerNames.SCHEDULE_FREE_SGD,
+]
+
+# =============================================================================
+# Built-in optimizer handlers registry
+# =============================================================================
+
+_OPTIMIZER_HANDLERS: dict[str, OptimizerHandler] = {
+    OptimizerNames.ADAFACTOR: _get_adafactor,
+    OptimizerNames.ADAMW_TORCH: _get_adamw_torch,
+    OptimizerNames.ADAMW_TORCH_FUSED: _get_adamw_torch,
+    OptimizerNames.ADAMW_TORCH_XLA: _get_adamw_torch_xla,
+    OptimizerNames.ADAMW_TORCH_NPU_FUSED: _get_adamw_torch_npu_fused,
+    OptimizerNames.ADAMW_ANYPRECISION: _get_adamw_anyprecision,
+    OptimizerNames.SGD: _get_sgd,
+    OptimizerNames.ADAGRAD: _get_adagrad,
+    OptimizerNames.RMSPROP: _get_rmsprop,
+    OptimizerNames.GROKADAMW: _get_grokadamw,
+    OptimizerNames.STABLE_ADAMW: _get_stable_adamw,
+    OptimizerNames.LOMO: _get_lomo_optimizer,
+    OptimizerNames.ADALOMO: _get_lomo_optimizer,
+    **dict.fromkeys(_BITSANDBYTES_OPTIMIZERS, _get_bitsandbytes_optimizer),
+    **dict.fromkeys(_GALORE_OPTIMIZERS, _get_galore_optimizer),
+    **dict.fromkeys(_APOLLO_OPTIMIZERS, _get_apollo_optimizer),
+    **dict.fromkeys(_TORCHAO_OPTIMIZERS, _get_torchao_optimizer),
+    **dict.fromkeys(_SCHEDULE_FREE_OPTIMIZERS, _get_schedule_free_optimizer),
+}