Add files using upload-large-folder tool

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0015000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-22_23:53:28 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0015000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0015000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0015000.pt
+[ckpt] step=15000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0015000.pt",
+  "step": 15000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 37.75448781763157,
+    "nll_per_token": 3.631104351215326,
+    "tokens": 36230,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 52.234604086708096,
+    "nll_per_token": 3.955745188796941,
+    "tokens": 30384,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.7591066520179934,
+    "unique_tokens": 2225,
+    "token_count": 32768,
+    "distinct_1": 0.067901611328125,
+    "distinct_2": 0.33172367125984253,
+    "top_token_mass": 0.096435546875
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0015000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-22_23:55:47 done step_0015000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0021000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_00:30:51 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0021000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0021000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0021000.pt
+[ckpt] step=21000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0021000.pt",
+  "step": 21000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 28.971311537338988,
+    "nll_per_token": 3.366306083015355,
+    "tokens": 34053,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 32.04400658699901,
+    "nll_per_token": 3.467110163913922,
+    "tokens": 29919,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.203253674656506,
+    "unique_tokens": 1894,
+    "token_count": 32768,
+    "distinct_1": 0.05780029296875,
+    "distinct_2": 0.26522514763779526,
+    "top_token_mass": 0.15106201171875
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0021000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_00:32:18 done step_0021000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0033000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_01:38:41 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0033000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0033000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0033000.pt
+[ckpt] step=33000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0033000.pt",
+  "step": 33000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 36.52794265924956,
+    "nll_per_token": 3.5980775200109547,
+    "tokens": 27157,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 48.37428762586346,
+    "nll_per_token": 3.8789684251611742,
+    "tokens": 22433,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 2.797182504419853,
+    "unique_tokens": 1406,
+    "token_count": 32768,
+    "distinct_1": 0.04290771484375,
+    "distinct_2": 0.2140132874015748,
+    "top_token_mass": 0.345794677734375
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0033000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_01:40:09 done step_0033000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0034000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_01:43:55 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0034000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0034000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0034000.pt
+[ckpt] step=34000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0034000.pt",
+  "step": 34000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 33.12776251190259,
+    "nll_per_token": 3.500371677414823,
+    "tokens": 34032,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 42.39457963921305,
+    "nll_per_token": 3.747020515368044,
+    "tokens": 28823,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.5379146352037796,
+    "unique_tokens": 1620,
+    "token_count": 32768,
+    "distinct_1": 0.0494384765625,
+    "distinct_2": 0.257689468503937,
+    "top_token_mass": 0.143035888671875
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0034000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_01:45:23 done step_0034000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0038000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_02:06:16 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0038000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0038000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0038000.pt
+[ckpt] step=38000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0038000.pt",
+  "step": 38000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 33.63094521372597,
+    "nll_per_token": 3.5154466316232162,
+    "tokens": 36373,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 46.12248543638734,
+    "nll_per_token": 3.8313005845903394,
+    "tokens": 30301,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.6560154748011517,
+    "unique_tokens": 2213,
+    "token_count": 32768,
+    "distinct_1": 0.067535400390625,
+    "distinct_2": 0.3269869586614173,
+    "top_token_mass": 0.088470458984375
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0038000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_02:07:44 done step_0038000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0046000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_02:51:01 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0046000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0046000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0046000.pt
+[ckpt] step=46000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0046000.pt",
+  "step": 46000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 31.36387952548299,
+    "nll_per_token": 3.445656897189349,
+    "tokens": 32110,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 41.301285095384074,
+    "nll_per_token": 3.7208936155938854,
+    "tokens": 27017,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.2403000800491895,
+    "unique_tokens": 1843,
+    "token_count": 32768,
+    "distinct_1": 0.056243896484375,
+    "distinct_2": 0.28518700787401574,
+    "top_token_mass": 0.217864990234375
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0046000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_02:52:29 done step_0046000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0070000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_05:04:36 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0070000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0070000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0070000.pt
+[ckpt] step=70000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0070000.pt",
+  "step": 70000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 34.62623490095348,
+    "nll_per_token": 3.544611628881008,
+    "tokens": 30022,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 46.91527187433584,
+    "nll_per_token": 3.8483432487659517,
+    "tokens": 24880,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.1183666256047453,
+    "unique_tokens": 2061,
+    "token_count": 32768,
+    "distinct_1": 0.062896728515625,
+    "distinct_2": 0.30305733267716534,
+    "top_token_mass": 0.275146484375
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0070000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_05:06:04 done step_0070000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0086000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_06:34:13 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0086000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0086000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0086000.pt
+[ckpt] step=86000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0086000.pt",
+  "step": 86000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 34.74809931980677,
+    "nll_per_token": 3.548124874682031,
+    "tokens": 34054,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 48.11665279080907,
+    "nll_per_token": 3.8736283290613485,
+    "tokens": 28218,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.514165014362856,
+    "unique_tokens": 2199,
+    "token_count": 32768,
+    "distinct_1": 0.067108154296875,
+    "distinct_2": 0.3373216043307087,
+    "top_token_mass": 0.165069580078125
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0086000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_06:35:42 done step_0086000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0087000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_06:39:28 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0087000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0087000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0087000.pt
+[ckpt] step=87000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0087000.pt",
+  "step": 87000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 37.572490028965,
+    "nll_per_token": 3.6262721344448807,
+    "tokens": 35507,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 48.67031168523643,
+    "nll_per_token": 3.885069227877383,
+    "tokens": 30112,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 3.6281564628090766,
+    "unique_tokens": 2318,
+    "token_count": 32768,
+    "distinct_1": 0.07073974609375,
+    "distinct_2": 0.3623892716535433,
+    "top_token_mass": 0.10577392578125
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0087000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_06:40:56 done step_0087000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/infer_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_step_0091000_logistic_normal_t1p45.log

@@ -0,0 +1,76 @@
+[watch-lognormal-sde] 2026-05-23_07:01:53 infer runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0091000.pt -> docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0091000
+[load] runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0091000.pt
+[ckpt] step=91000
+[sde] generated 16/256
+[sde] generated 32/256
+[sde] generated 48/256
+[sde] generated 64/256
+[sde] generated 80/256
+[sde] generated 96/256
+[sde] generated 112/256
+[sde] generated 128/256
+[sde] generated 144/256
+[sde] generated 160/256
+[sde] generated 176/256
+[sde] generated 192/256
+[sde] generated 208/256
+[sde] generated 224/256
+[sde] generated 240/256
+[sde] generated 256/256
+[score] loading scorer: /e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard
+[summary] {
+  "type": "summary",
+  "checkpoint": "runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0091000.pt",
+  "step": 91000,
+  "decode": {
+    "decode_rule": "logistic_normal_resample_sde",
+    "steps": 128,
+    "model_t_mode": "const0.5",
+    "mean_mode": "anchor_semantic",
+    "endpoint_floor": 0.0,
+    "concentration_min": 1.0,
+    "concentration_max": 1024.0,
+    "endpoint_temp": 1.45,
+    "support_power": 1.0,
+    "semantic_power": 1.0,
+    "noise_init": "logistic_normal",
+    "noise_sigma": 3.0,
+    "noise_dirichlet_concentration": 1.0,
+    "sde_resample": "logistic_normal",
+    "logistic_normal_sigma_min": 0.18,
+    "logistic_normal_sigma_max": 3.0,
+    "logistic_normal_tau_min": 0.65,
+    "logistic_normal_tau_max": 1.0,
+    "final_from": "blend_0.5",
+    "n_samples": 256,
+    "seed": 20260522
+  },
+  "raw_genppl": {
+    "ppl": 40.78150904177592,
+    "nll_per_token": 3.708228768919138,
+    "tokens": 26005,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "stripped_genppl": {
+    "ppl": 50.90887102234889,
+    "nll_per_token": 3.9300371917176355,
+    "tokens": 22154,
+    "kept_samples": 256,
+    "total_samples": 256,
+    "empty_rate": 0.0,
+    "skipped_samples": 0
+  },
+  "diversity": {
+    "sample_entropy": 2.7366672527527456,
+    "unique_tokens": 1634,
+    "token_count": 32768,
+    "distinct_1": 0.04986572265625,
+    "distinct_2": 0.24818528543307086,
+    "top_token_mass": 0.362945556640625
+  }
+}
+[done] docs/lta_samples/metrics_20260522/lm1b_len128_lognormal_atoms_every1k_logistic_normal_sde_t1p45_steps128_n256/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0091000/sde_steps128_samples256_scored.jsonl
+[watch-lognormal-sde] 2026-05-23_07:03:20 done step_0091000

LTA_openwebtext_dualt/logs/lm1b_len128_lognormal_atoms_sde_watch/processed_lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522_logistic_normal_steps128_t1p45_n256.txt

@@ -0,0 +1,124 @@
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0001000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0002000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0003000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0004000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0005000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0006000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0007000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0008000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0009000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0010000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0011000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0012000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0013000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0014000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0015000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0016000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0017000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0018000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0019000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0020000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0021000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0022000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0023000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0024000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0025000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0026000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0027000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0028000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0029000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0030000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0031000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0032000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0033000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0034000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0035000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0036000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0037000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0038000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0039000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0040000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0041000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0042000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0043000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0044000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0045000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0046000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0047000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0048000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0049000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0050000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0051000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0052000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0053000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0054000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0055000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0056000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0057000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0058000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0059000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0060000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0061000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0062000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0063000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0064000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0065000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0066000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0067000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0068000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0069000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0070000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0071000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0072000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0073000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0074000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0075000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0076000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0077000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0078000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0079000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0080000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0081000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0082000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0083000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0084000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0085000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0086000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0087000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0088000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0089000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0090000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0091000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0092000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0093000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0094000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0095000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0096000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0097000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0098000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0099000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0100000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0101000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0102000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0103000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0104000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0105000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0106000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0107000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0108000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0109000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0110000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0111000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0112000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0113000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0114000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0115000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0116000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0117000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0118000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0119000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0120000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0121000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0122000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0123000.pt
+runs/lta_lm1b_classic_c1024_len128_lognormalatoms_gbs512_4gpu_20260522/step_0124000.pt

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

@@ -0,0 +1,31 @@
+# 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_blip import *
+    from .image_processing_blip import *
+    from .image_processing_pil_blip import *
+    from .modeling_blip import *
+    from .modeling_blip_text import *
+    from .processing_blip 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/blip/configuration_blip.py

@@ -0,0 +1,175 @@
+# Copyright 2022 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.
+"""Blip model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="Salesforce/blip-vqa-base")
+@strict
+class BlipTextConfig(PreTrainedConfig):
+    r"""
+    label_smoothing (float, *optional*):
+        A float in [0.0, 1.0]. Specifies the amount of smoothing when computing the loss, where 0.0 means no smoothing. The targets
+        become a mixture of the original ground truth and a uniform distribution as described in
+        `Rethinking the Inception Architecture for Computer Vision <https://huggingface.co/papers/1512.00567>`__. Default: :math:`0.0`.
+
+    Example:
+
+    ```python
+    >>> from transformers import BlipTextConfig, BlipTextModel
+
+    >>> # Initializing a BlipTextConfig with Salesforce/blip-vqa-base style configuration
+    >>> configuration = BlipTextConfig()
+
+    >>> # Initializing a BlipTextModel (with random weights) from the Salesforce/blip-vqa-base style configuration
+    >>> model = BlipTextModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "blip_text_model"
+    base_config_key = "text_config"
+
+    vocab_size: int = 30524
+    hidden_size: int = 768
+    encoder_hidden_size: int = 768
+    intermediate_size: int = 3072
+    projection_dim: int = 768
+    num_hidden_layers: int = 12
+    num_attention_heads: int = 8
+    max_position_embeddings: int = 512
+    hidden_act: str = "gelu"
+    layer_norm_eps: float = 1e-12
+    hidden_dropout_prob: float | int = 0.0
+    attention_probs_dropout_prob: float | int = 0.0
+    initializer_range: float = 0.02
+    bos_token_id: int | None = 30522
+    eos_token_id: int | list[int] | None = 2
+    pad_token_id: int | None = 0
+    sep_token_id: int | None = 102
+    is_decoder: bool = True
+    use_cache: bool = True
+    tie_word_embeddings: bool = True
+    label_smoothing: float = 0.0
+
+
+@auto_docstring(checkpoint="Salesforce/blip-vqa-base")
+@strict
+class BlipVisionConfig(PreTrainedConfig):
+    r"""
+    Example:
+
+    ```python
+    >>> from transformers import BlipVisionConfig, BlipVisionModel
+
+    >>> # Initializing a BlipVisionConfig with Salesforce/blip-vqa-base style configuration
+    >>> configuration = BlipVisionConfig()
+
+    >>> # Initializing a BlipVisionModel (with random weights) from the Salesforce/blip-vqa-base style configuration
+    >>> model = BlipVisionModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "blip_vision_model"
+    base_config_key = "vision_config"
+
+    hidden_size: int = 768
+    intermediate_size: int = 3072
+    projection_dim: int = 512
+    num_hidden_layers: int = 12
+    num_attention_heads: int = 12
+    image_size: int | list[int] | tuple[int, int] = 384
+    patch_size: int | list[int] | tuple[int, int] = 16
+    hidden_act: str = "gelu"
+    layer_norm_eps: float = 1e-5
+    attention_dropout: float | int = 0.0
+    initializer_range: float = 1e-10
+
+
+@auto_docstring(checkpoint="Salesforce/blip-vqa-base")
+@strict
+class BlipConfig(PreTrainedConfig):
+    r"""
+    image_text_hidden_size (`int`, *optional*, defaults to 256):
+        Dimensionality of the hidden state of the image-text fusion layer.
+    label_smoothing (float, *optional*):
+        A float in [0.0, 1.0]. Specifies the amount of smoothing when computing the loss, where 0.0 means no smoothing. The targets
+        become a mixture of the original ground truth and a uniform distribution as described in
+        `Rethinking the Inception Architecture for Computer Vision <https://huggingface.co/papers/1512.00567>`__. Default: :math:`0.0`.
+
+    Example:
+
+    ```python
+    >>> from transformers import BlipConfig, BlipModel
+
+    >>> # Initializing a BlipConfig with Salesforce/blip-vqa-base style configuration
+    >>> configuration = BlipConfig()
+
+    >>> # Initializing a BlipPModel (with random weights) from the Salesforce/blip-vqa-base style configuration
+    >>> model = BlipModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+
+    >>> # We can also initialize a BlipConfig from a BlipTextConfig and a BlipVisionConfig
+
+    >>> # Initializing a BLIPText and BLIPVision configuration
+    >>> config_text = BlipTextConfig()
+    >>> config_vision = BlipVisionConfig()
+
+    >>> config = BlipConfig(text_config=config_text, vision_config=config_vision)
+    ```"""
+
+    model_type = "blip"
+    sub_configs = {"text_config": BlipTextConfig, "vision_config": BlipVisionConfig}
+
+    text_config: dict | PreTrainedConfig | None = None
+    vision_config: dict | PreTrainedConfig | None = None
+    projection_dim: int = 512
+    logit_scale_init_value: float = 2.6592
+    image_text_hidden_size: int = 256
+    label_smoothing: float = 0.0
+    tie_word_embeddings: bool = True
+    initializer_factor: float = 1.0
+    initializer_range: float = 0.02
+
+    def __post_init__(self, **kwargs):
+        if self.text_config is None:
+            self.text_config = BlipTextConfig()
+            logger.info("`text_config` is `None`. Initializing the `BlipTextConfig` with default values.")
+        elif isinstance(self.text_config, dict):
+            self.text_config = BlipTextConfig(**self.text_config)
+
+        if self.vision_config is None:
+            self.vision_config = BlipVisionConfig()
+            logger.info("`vision_config` is `None`. initializing the `BlipVisionConfig` with default values.")
+        elif isinstance(self.vision_config, dict):
+            self.vision_config = BlipVisionConfig(**self.vision_config)
+
+        self.text_config.encoder_hidden_size = self.vision_config.hidden_size
+
+        super().__post_init__(**kwargs)
+
+
+__all__ = ["BlipConfig", "BlipTextConfig", "BlipVisionConfig"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/blip/image_processing_pil_blip.py

@@ -0,0 +1,34 @@
+# Copyright 2022 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.
+"""Image processor class for BLIP."""
+
+from ...image_processing_backends import PilBackend
+from ...image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, PILImageResampling
+from ...utils import auto_docstring
+
+
+@auto_docstring
+class BlipImageProcessorPil(PilBackend):
+    resample = PILImageResampling.BICUBIC
+    image_mean = OPENAI_CLIP_MEAN
+    image_std = OPENAI_CLIP_STD
+    size = {"height": 384, "width": 384}
+    default_to_square = True
+    do_resize = True
+    do_rescale = True
+    do_normalize = True
+    do_convert_rgb = True
+
+
+__all__ = ["BlipImageProcessorPil"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/blip/modeling_blip_text.py

@@ -0,0 +1,709 @@
+# Copyright 2022 The Salesforce Team Authors and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the BSD-3-clause license (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://opensource.org/licenses/BSD-3-Clause
+#
+# 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 torch
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_bidirectional_mask, create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    BaseModelOutputWithPoolingAndCrossAttentions,
+    CausalLMOutputWithCrossAttentions,
+)
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...pytorch_utils import apply_chunking_to_forward
+from ...utils import TransformersKwargs, can_return_tuple, logging
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from .configuration_blip import BlipTextConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# Adapted from https://github.com/salesforce/BLIP/blob/main/models/med.py#L52
+class BlipTextEmbeddings(nn.Module):
+    """Construct the embeddings from word and position embeddings."""
+
+    def __init__(self, config):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+        self.register_buffer(
+            "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
+        )
+
+        self.config = config
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        past_key_values_length: int = 0,
+    ) -> torch.Tensor:
+        if input_ids is not None:
+            input_shape = input_ids.size()
+        else:
+            input_shape = inputs_embeds.size()[:-1]
+
+        seq_length = input_shape[1]
+
+        if position_ids is None:
+            position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+
+        embeddings = inputs_embeds
+
+        position_embeddings = self.position_embeddings(position_ids)
+        embeddings += position_embeddings
+
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+# Adapted from https://github.com/salesforce/BLIP/blob/main/models/med.py#L97
+class BlipTextSelfAttention(nn.Module):
+    def __init__(self, config, is_cross_attention, layer_idx=None):
+        super().__init__()
+        self.config = config
+        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+            raise ValueError(
+                "The hidden size (%d) is not a multiple of the number of attention heads (%d)"
+                % (config.hidden_size, config.num_attention_heads)
+            )
+
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+        self.layer_idx = layer_idx
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        if is_cross_attention:
+            self.key = nn.Linear(config.encoder_hidden_size, self.all_head_size)
+            self.value = nn.Linear(config.encoder_hidden_size, self.all_head_size)
+        else:
+            self.key = nn.Linear(config.hidden_size, self.all_head_size)
+            self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+    def save_attn_gradients(self, attn_gradients):
+        self.attn_gradients = attn_gradients
+
+    def get_attn_gradients(self):
+        return self.attn_gradients
+
+    def save_attention_map(self, attention_map):
+        self.attention_map = attention_map
+
+    def get_attention_map(self):
+        return self.attention_map
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.FloatTensor | None = None,
+        encoder_hidden_states: torch.FloatTensor | None = None,
+        encoder_attention_mask: torch.FloatTensor | None = None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.attention_head_size)
+        query_layer = self.query(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        # If this is instantiated as a cross-attention module, the keys
+        # and values come from an encoder; the attention mask needs to be
+        # such that the encoder's padding tokens are not attended to.
+        is_cross_attention = encoder_hidden_states is not None
+        attention_mask = encoder_attention_mask if is_cross_attention else attention_mask
+
+        is_updated = False
+        if past_key_values is not None:
+            if isinstance(past_key_values, EncoderDecoderCache):
+                is_updated = past_key_values.is_updated.get(self.layer_idx)
+                if is_cross_attention:
+                    # after the first generated id, we can subsequently re-use all key/value_layer from cache
+                    curr_past_key_values = past_key_values.cross_attention_cache
+                else:
+                    curr_past_key_values = past_key_values.self_attention_cache
+            else:
+                curr_past_key_values = past_key_values
+
+        current_states = encoder_hidden_states if is_cross_attention else hidden_states
+        if is_cross_attention and past_key_values is not None and is_updated:
+            # reuse k,v, cross_attentions
+            key_layer = curr_past_key_values.layers[self.layer_idx].keys
+            value_layer = curr_past_key_values.layers[self.layer_idx].values
+        else:
+            kv_shape = (*current_states.shape[:-1], -1, self.attention_head_size)
+            key_layer = self.key(current_states).view(kv_shape).transpose(1, 2)
+            value_layer = self.value(current_states).view(kv_shape).transpose(1, 2)
+
+            if past_key_values is not None:
+                key_layer, value_layer = curr_past_key_values.update(key_layer, value_layer, self.layer_idx)
+                # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
+                if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):
+                    past_key_values.is_updated[self.layer_idx] = True
+
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+        if attention_mask is not None:
+            # Apply the attention mask is (precomputed for all layers in BlipTextModel forward() function)
+            attention_scores = attention_scores + attention_mask.to(attention_scores.device)
+
+        # Normalize the attention scores to probabilities.
+        attention_probs = nn.Softmax(dim=-1)(attention_scores)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        attention_probs_dropped = self.dropout(attention_probs)
+
+        context_layer = torch.matmul(attention_probs_dropped, value_layer)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+
+        return context_layer, attention_probs
+
+
+# Copied from transformers.models.bert.modeling_bert.BertSelfOutput with Bert -> BlipText
+class BlipTextSelfOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+# Adapted from https://github.com/salesforce/BLIP/blob/main/models/med.py#242
+class BlipTextAttention(nn.Module):
+    def __init__(self, config, is_cross_attention=False, layer_idx=None):
+        super().__init__()
+        self.self = BlipTextSelfAttention(config, is_cross_attention, layer_idx=layer_idx)
+        self.output = BlipTextSelfOutput(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.FloatTensor | None = None,
+        encoder_hidden_states: torch.FloatTensor | None = None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        context_layer, attention_probs = self.self(
+            hidden_states,
+            attention_mask=attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            past_key_values=past_key_values,
+        )
+        attention_output = self.output(context_layer, hidden_states)
+        return attention_output, attention_probs
+
+
+# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert -> BlipText
+class BlipTextIntermediate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if isinstance(config.hidden_act, str):
+            self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.intermediate_act_fn = config.hidden_act
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+        return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert -> BlipText
+class BlipTextOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+class BlipTextLayer(GradientCheckpointingLayer):
+    def __init__(self, config, layer_num):
+        super().__init__()
+        self.config = config
+        self.chunk_size_feed_forward = config.chunk_size_feed_forward
+        self.seq_len_dim = 1
+        self.attention = BlipTextAttention(config, layer_idx=layer_num)
+        self.layer_num = layer_num
+        if self.config.is_decoder:
+            self.crossattention = BlipTextAttention(
+                config, is_cross_attention=self.config.is_decoder, layer_idx=layer_num
+            )
+        self.intermediate = BlipTextIntermediate(config)
+        self.output = BlipTextOutput(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.FloatTensor | None = None,
+        attention_mask: torch.FloatTensor | None = None,
+        encoder_attention_mask: torch.FloatTensor | None = None,
+        past_key_values: Cache | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        attention_output, _ = self.attention(
+            hidden_states,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+        )
+
+        if encoder_hidden_states is not None:
+            attention_output, _ = self.crossattention(
+                attention_output,
+                attention_mask=encoder_attention_mask,
+                encoder_hidden_states=encoder_hidden_states,
+                past_key_values=past_key_values,
+            )
+        layer_output = apply_chunking_to_forward(
+            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
+        )
+        return layer_output
+
+    def feed_forward_chunk(self, attention_output):
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        return layer_output
+
+
+# Adapted from https://github.com/salesforce/BLIP/blob/main/models/med.py#L386
+class BlipTextEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.layer = nn.ModuleList([BlipTextLayer(config, i) for i in range(config.num_hidden_layers)])
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.FloatTensor | None = None,
+        encoder_hidden_states: torch.FloatTensor | None = None,
+        encoder_attention_mask: torch.FloatTensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPastAndCrossAttentions:
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+                )
+                use_cache = False
+
+        if use_cache:
+            # The model acts as encoder decoder but is not an encoder decoder. So we cast all cache objects to
+            # `EncoderDecoderCache` type assuming that the incoming cache is from `self_attention`
+            if isinstance(past_key_values, DynamicCache):
+                past_key_values = EncoderDecoderCache(past_key_values, DynamicCache(config=self.config))
+            elif past_key_values is None:
+                past_key_values = EncoderDecoderCache(
+                    DynamicCache(config=self.config), DynamicCache(config=self.config)
+                )
+
+        for layer_module in self.layer:
+            hidden_states = layer_module(
+                hidden_states,
+                encoder_hidden_states,
+                attention_mask=attention_mask,
+                encoder_attention_mask=encoder_attention_mask,
+                past_key_values=past_key_values,
+                **kwargs,
+            )
+
+        return BaseModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+# Copied from transformers.models.bert.modeling_bert.BertPooler with Bert->BlipText
+class BlipTextPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+# Copied from transformers.models.bert.modeling_bert.BertPredictionHeadTransform with Bert->BlipText
+class BlipTextPredictionHeadTransform(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        if isinstance(config.hidden_act, str):
+            self.transform_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.transform_act_fn = config.hidden_act
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states)
+        return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead with Bert->BlipText
+class BlipTextLMPredictionHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.transform = BlipTextPredictionHeadTransform(config)
+
+        # The output weights are the same as the input embeddings, but there is
+        # an output-only bias for each token.
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=True)
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+
+    def forward(self, hidden_states):
+        hidden_states = self.transform(hidden_states)
+        hidden_states = self.decoder(hidden_states)
+        return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertOnlyMLMHead with Bert->BlipText
+class BlipTextOnlyMLMHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.predictions = BlipTextLMPredictionHead(config)
+
+    def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
+        prediction_scores = self.predictions(sequence_output)
+        return prediction_scores
+
+
+# Adapted from https://github.com/salesforce/BLIP/blob/main/models/med.py#L548
+class BlipTextPreTrainedModel(PreTrainedModel):
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config: BlipTextConfig
+    base_model_prefix = "bert"
+    _no_split_modules = []
+    _can_record_outputs = {
+        "hidden_states": BlipTextLayer,
+        "attentions": [
+            OutputRecorder(BlipTextSelfAttention, index=1, layer_name=".attention."),
+        ],
+        "cross_attentions": [
+            OutputRecorder(BlipTextSelfAttention, index=1, layer_name=".crossattention."),
+        ],
+    }
+
+    def _init_weights(self, module):
+        super()._init_weights(module)
+        if isinstance(module, BlipTextEmbeddings):
+            init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+
+
+# Adapted from https://github.com/salesforce/BLIP/blob/3a29b7410476bf5f2ba0955827390eb6ea1f4f9d/models/med.py#L571
+class BlipTextModel(BlipTextPreTrainedModel):
+    """
+    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
+    cross-attention is added between the self-attention layers, following the architecture described in [Attention is
+    all you need](https://huggingface.co/papers/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
+    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin. argument and `is_decoder` set to `True`; an
+    `encoder_hidden_states` is then expected as an input to the forward pass.
+    """
+
+    def __init__(self, config, add_pooling_layer=True):
+        super().__init__(config)
+        self.config = config
+
+        self.embeddings = BlipTextEmbeddings(config)
+        self.encoder = BlipTextEncoder(config)
+        self.pooler = BlipTextPooler(config) if add_pooling_layer else None
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    @merge_with_config_defaults
+    @capture_outputs
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        encoder_embeds: torch.Tensor | None = None,
+        encoder_hidden_states: torch.Tensor | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = None,
+        is_decoder: bool | None = False,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutputWithPoolingAndCrossAttentions:
+        r"""
+        encoder_hidden_states  (`torch.FloatTensor`, *optional*):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+            the model is configured as a decoder.
+        encoder_attention_mask (`torch.FloatTensor`, *optional*):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        past_key_values (`Cache`, *optional*):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
+            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
+            `decoder_input_ids` of shape `(batch_size, sequence_length)`.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        """
+        if not is_decoder:
+            use_cache = False
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+
+        past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+
+        if encoder_embeds is None:
+            embedding_output = self.embeddings(
+                input_ids=input_ids,
+                position_ids=position_ids,
+                inputs_embeds=inputs_embeds,
+                past_key_values_length=past_key_values_length,
+            )
+        else:
+            embedding_output = encoder_embeds
+
+        if is_decoder:
+            attention_mask = create_causal_mask(
+                config=self.config,
+                inputs_embeds=embedding_output,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+            )
+        else:
+            attention_mask = create_bidirectional_mask(
+                config=self.config,
+                inputs_embeds=embedding_output,
+                attention_mask=attention_mask,
+            )
+
+        if encoder_attention_mask is not None:
+            encoder_attention_mask = create_bidirectional_mask(
+                config=self.config,
+                inputs_embeds=embedding_output,
+                attention_mask=encoder_attention_mask,
+                encoder_hidden_states=encoder_hidden_states,
+            )
+
+        encoder_outputs: BaseModelOutputWithPastAndCrossAttentions = self.encoder(
+            embedding_output,
+            attention_mask=attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            **kwargs,
+        )
+        sequence_output = encoder_outputs.last_hidden_state
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+        return BaseModelOutputWithPoolingAndCrossAttentions(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            past_key_values=encoder_outputs.past_key_values,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+            cross_attentions=encoder_outputs.cross_attentions,
+        )
+
+
+# Adapted from https://github.com/salesforce/BLIP/blob/main/models/med.py#L811
+class BlipTextLMHeadModel(BlipTextPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {
+        "cls.predictions.decoder.bias": "cls.predictions.bias",
+        "cls.predictions.decoder.weight": "bert.embeddings.word_embeddings.weight",
+    }
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.bert = BlipTextModel(config, add_pooling_layer=False)
+        self.cls = BlipTextOnlyMLMHead(config)
+        self.label_smoothing = config.label_smoothing
+
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.bert.get_input_embeddings()
+
+    def set_input_embeddings(self, new_embeddings):
+        self.bert.set_input_embeddings(new_embeddings)
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        self.cls.predictions.decoder = new_embeddings
+        self.cls.predictions.bias = new_embeddings.bias
+
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        encoder_hidden_states: torch.Tensor | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = None,
+        return_logits: bool | None = False,
+        is_decoder: bool | None = True,
+        reduction: str | None = "mean",
+        logits_to_keep: int | torch.Tensor = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithCrossAttentions:
+        r"""
+        encoder_hidden_states (`torch.FloatTensor`, *optional*): Sequence of
+            hidden-states at the output of the last layer of the encoder. Used in the cross-attention if the model is
+            configured as a decoder.
+        encoder_attention_mask (`torch.FloatTensor`, *optional*):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+            the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+        labels (`torch.LongTensor`, *optional*):
+            Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
+            `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
+            ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`
+        past_key_values (`Cache`, *optional*):
+            Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+            If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
+            don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
+            `decoder_input_ids` of shape `(batch_size, sequence_length)`.
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        """
+        if labels is not None:
+            use_cache = False
+
+        outputs: BaseModelOutputWithPoolingAndCrossAttentions = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            is_decoder=is_decoder,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        prediction_scores = self.cls(hidden_states[:, slice_indices, :])
+
+        if return_logits:
+            return prediction_scores[:, :-1, :].contiguous()
+
+        lm_loss = None
+        if labels is not None:
+            # we are doing next-token prediction; shift prediction scores and input ids by one
+            shifted_prediction_scores = prediction_scores[:, :-1, :].contiguous()
+            labels = labels[:, 1:].contiguous().to(shifted_prediction_scores.device)
+            loss_fct = CrossEntropyLoss(reduction=reduction, label_smoothing=self.label_smoothing)
+            lm_loss = loss_fct(shifted_prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+            if reduction == "none":
+                lm_loss = lm_loss.view(prediction_scores.size(0), -1).sum(1)
+
+        return CausalLMOutputWithCrossAttentions(
+            loss=lm_loss,
+            logits=prediction_scores,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            cross_attentions=outputs.cross_attentions,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **model_kwargs):
+        # Overwrite -- hardcoded key return (`is_decoder=True`)
+
+        model_inputs = super().prepare_inputs_for_generation(
+            input_ids,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            **model_kwargs,
+        )
+        model_inputs["is_decoder"] = True
+
+        return model_inputs
+
+
+__all__ = ["BlipTextModel", "BlipTextLMHeadModel", "BlipTextPreTrainedModel"]

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

@@ -0,0 +1,28 @@
+# 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_qwen2_audio import *
+    from .modeling_qwen2_audio import *
+    from .processing_qwen2_audio 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/qwen2_audio/configuration_qwen2_audio.py

@@ -0,0 +1,128 @@
+# Copyright 2024 Microsoft Research & University of Wisconsin-Madison 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.
+"""Qwen2Audio model configuration"""
+
+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="Qwen/Qwen2-Audio-7B")
+@strict
+class Qwen2AudioEncoderConfig(PreTrainedConfig):
+    r"""
+    max_source_positions (`int`, *optional*, defaults to 1500):
+        The maximum sequence length of log-mel filter-bank features that this model might ever be used with.
+
+    Example:
+
+    ```python
+    >>> from transformers import Qwen2AudioEncoderConfig, Qwen2AudioEncoder
+
+    >>> # Initializing a Qwen2AudioEncoderConfig
+    >>> configuration = Qwen2AudioEncoderConfig()
+
+    >>> # Initializing a Qwen2AudioEncoder (with random weights)
+    >>> model = Qwen2AudioEncoder(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "qwen2_audio_encoder"
+    attribute_map = {
+        "num_hidden_layers": "encoder_layers",
+        "hidden_size": "d_model",
+        "num_attention_heads": "encoder_attention_heads",
+        "intermediate_size": "encoder_ffn_dim",
+    }
+
+    num_mel_bins: int = 128
+    encoder_layers: int = 32
+    encoder_attention_heads: int = 20
+    encoder_ffn_dim: int = 5120
+    encoder_layerdrop: float | int = 0.0
+    d_model: int = 1280
+    dropout: float | int = 0.0
+    attention_dropout: float | int = 0.0
+    activation_function: str = "gelu"
+    activation_dropout: float | int = 0.0
+    scale_embedding: bool = False
+    initializer_range: float = 0.02
+    max_source_positions: int = 1500
+
+
+@auto_docstring(checkpoint="Qwen/Qwen2-Audio-7B")
+@strict
+class Qwen2AudioConfig(PreTrainedConfig):
+    r"""
+    Example:
+
+    ```python
+    >>> from transformers import Qwen2AudioForConditionalGeneration, Qwen2AudioConfig, Qwen2AudioEncoderConfig, Qwen2Config
+
+    >>> # Initializing a Qwen2AudioEncoder config
+    >>> audio_config = Qwen2AudioEncoderConfig()
+
+    >>> # Initializing a Qwen2 config
+    >>> text_config = Qwen2Config()
+
+    >>> # Initializing a Qwen2Audio configuration
+    >>> configuration = Qwen2AudioConfig(audio_config, text_config)
+
+    >>> # Initializing a model from the qwen2-audio style configuration
+    >>> model = Qwen2AudioForConditionalGeneration(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "qwen2_audio"
+    attribute_map = {
+        "audio_token_id": "audio_token_index",
+    }
+    sub_configs = {"text_config": AutoConfig, "audio_config": AutoConfig}
+
+    audio_config: dict | PreTrainedConfig | None = None
+    text_config: dict | PreTrainedConfig | None = None
+    audio_token_index: int = 151646
+
+    def __post_init__(self, **kwargs):
+        if isinstance(self.audio_config, dict):
+            self.audio_config["model_type"] = self.audio_config.get("model_type", "qwen2_audio_encoder")
+            self.audio_config = CONFIG_MAPPING[self.audio_config["model_type"]](**self.audio_config)
+        elif self.audio_config is None:
+            self.audio_config = CONFIG_MAPPING["qwen2_audio_encoder"](
+                d_model=1280,
+                encoder_attention_heads=20,
+                encoder_ffn_dim=5120,
+                encoder_layerdrop=0.0,
+                encoder_layers=32,
+                num_mel_bins=128,
+                max_source_positions=1500,
+                scale_embedding=False,
+                activation_function="gelu",
+            )
+
+        if isinstance(self.text_config, dict):
+            self.text_config["model_type"] = self.text_config.get("model_type", "qwen2")
+            self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config)
+        elif self.text_config is None:
+            self.text_config = CONFIG_MAPPING["qwen2"]()
+
+        super().__post_init__(**kwargs)
+
+
+__all__ = ["Qwen2AudioConfig", "Qwen2AudioEncoderConfig"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/switch_transformers/configuration_switch_transformers.py

@@ -0,0 +1,111 @@
+# Copyright 2022, Google and HuggingFace Inc.
+#
+# 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.
+"""Switch Transformers model configuration"""
+
+from typing import Literal
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="google/switch-base-8")
+@strict
+class SwitchTransformersConfig(PreTrainedConfig):
+    r"""
+    num_sparse_encoder_layers (`int`, *optional*, defaults to 3):
+        Number of sparse (MoE) dense hidden layers in the Transformer encoder layer.
+        Note: When set to 0 with `num_layers=1`, the current implementation may still create a sparse layer
+        due to the sparse step calculation. This edge case is not encountered in existing checkpoints.
+    num_sparse_decoder_layers (`int`, *optional*, defaults to 3):
+        Number of sparse (MoE) dense hidden layers in the Transformer decoder layer.
+        Note: When set to 0 with `num_decoder_layers=1`, the current implementation may still create a sparse
+        layer due to the sparse step calculation. This edge case is not encountered in existing checkpoints.
+    router_bias (`bool`, *optional*, defaults to `False`):
+        Whether to add a bias to the router.
+    router_dtype (`str`, *optional*, default to `"float32"`):
+        The `dtype` used for the routers. It is preferable to keep the `dtype` to `"float32"` as specified in the
+        *selective precision* discussion in [the paper](https://huggingface.co/papers/2101.03961).
+    router_ignore_padding_tokens (`bool`, *optional*, defaults to `False`):
+        Whether to ignore padding tokens when routing.
+    relative_attention_num_buckets (`int`, *optional*, defaults to 32):
+        The number of buckets to use for each attention layer.
+    relative_attention_max_distance (`int`, *optional*, defaults to 128):
+        The maximum distance of the longer sequences for the bucket separation.
+    dense_act_fn (`string`, *optional*, defaults to `"relu"`):
+        Type of feed forward layer to be used. Should be one of `"relu"` or `"gated-gelu"`. SwitchTransformersv1.1
+        uses the `"gated-gelu"` feed forward projection. Original SwitchTransformers uses `"relu"`.
+    add_router_probs (`bool`, *optional*, defaults to `False`):
+        Whether to output router probabilities to compute router auxiliary loss.
+    """
+
+    model_type = "switch_transformers"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    attribute_map = {"hidden_size": "d_model", "num_attention_heads": "num_heads", "num_hidden_layers": "num_layers"}
+
+    vocab_size: int = 32128
+    d_model: int = 768
+    d_kv: int = 64
+    d_ff: int = 2048
+    expert_capacity: int = 64
+    num_layers: int = 12
+    num_sparse_encoder_layers: int = 3
+    num_decoder_layers: int | None = 12
+    num_sparse_decoder_layers: int = 3
+    num_heads: int = 12
+    num_experts: int = 8
+    router_bias: bool = False
+    router_jitter_noise: int | float = 0.01
+    router_dtype: Literal["float32", "float16", "bfloat16"] = "float32"
+    router_ignore_padding_tokens: bool = False
+    relative_attention_num_buckets: int = 32
+    relative_attention_max_distance: int = 128
+    dropout_rate: float | int = 0.1
+    layer_norm_epsilon: float = 1e-6
+    router_z_loss_coef: float = 0.001
+    router_aux_loss_coef: float = 0.001
+    initializer_factor: float = 1.0
+    dense_act_fn: str = "relu"
+    is_encoder_decoder: bool = True
+    add_router_probs: bool = False
+    use_cache: bool = True
+    pad_token_id: int | None = 0
+    eos_token_id: int | list[int] | None = 1
+    bos_token_id: int | None = None
+    tie_word_embeddings: bool = True
+    is_decoder: bool = False
+    add_cross_attention: bool = False
+
+    def __post_init__(self, **kwargs):
+        self.num_decoder_layers = (
+            self.num_decoder_layers if self.num_decoder_layers is not None else self.num_layers
+        )  # default = symmetry
+
+        # This tells us, each how many encoder layer we'll have to set a sparse layer.
+        if self.num_sparse_encoder_layers > 0:
+            self.encoder_sparse_step = self.num_layers // self.num_sparse_encoder_layers
+        else:
+            self.encoder_sparse_step = self.num_layers  # HACK: this will create 0 sparse layers
+
+        # This tells us, each how many decoder layer we'll have to set a sparse layer.
+        if self.num_sparse_decoder_layers > 0:
+            self.decoder_sparse_step = self.num_decoder_layers // self.num_sparse_decoder_layers
+        else:
+            self.decoder_sparse_step = self.num_decoder_layers  # HACK: this will create 0 sparse layers
+
+        super().__post_init__(**kwargs)
+
+
+__all__ = ["SwitchTransformersConfig"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/switch_transformers/modeling_switch_transformers.py

@@ -0,0 +1,1095 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/switch_transformers/modular_switch_transformers.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_switch_transformers.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2022 SwitchTransformers Authors and 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 copy
+import math
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_bidirectional_mask, create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+    MoEModelOutput,
+    MoEModelOutputWithPastAndCrossAttentions,
+    Seq2SeqMoEModelOutput,
+    Seq2SeqMoEOutput,
+)
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, is_torchdynamo_compiling, logging
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from .configuration_switch_transformers import SwitchTransformersConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class SwitchTransformersTop1Router(nn.Module):
+    """
+    Router using tokens choose top-1 experts assignment.
+
+    This router uses the same mechanism as in Switch Transformer (https://huggingface.co/papers/2101.03961) and V-MoE
+    (https://huggingface.co/papers/2106.05974): tokens choose their top experts. Items are sorted by router_probs and then
+    routed to their choice of expert until the expert's expert_capacity is reached. **There is no guarantee that each
+    token is processed by an expert**, or that each expert receives at least one token.
+
+    """
+
+    def __init__(self, config: SwitchTransformersConfig):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.expert_capacity = config.expert_capacity
+        self.classifier = nn.Linear(config.hidden_size, self.num_experts, bias=config.router_bias)
+        self.jitter_noise = config.router_jitter_noise
+        self.ignore_padding_tokens = config.router_ignore_padding_tokens
+        self.dtype = getattr(torch, config.router_dtype)
+
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        r"""
+        Computes router probabilities from input hidden states.
+
+        Args:
+            hidden_states (`torch.Tensor`):
+                (batch_size, sequence_length, hidden_dim) from which router probabilities are computed.
+        Returns:
+            router_probabilities (`torch.Tensor`):
+                Tensor of shape (batch_size, sequence_length, num_experts) corresponding to the probabilities for each
+                token and expert. Used for routing tokens to experts.
+            router_logits (`torch.Tensor`):
+                Logits tensor of shape (batch_size, sequence_length, num_experts) corresponding to raw router logits.
+                This is used later for computing router z-loss.
+        """
+        # float32 is used to ensure stability. See the discussion of "selective precision" in
+        # https://huggingface.co/papers/2101.03961.
+        # We also store the previous dtype to cast back the output to the previous dtype
+        self.input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(self.dtype)
+        if self.training and self.jitter_noise > 0:
+            # Multiply the token inputs by the uniform distribution - adding some noise
+            hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise)
+        self.classifier = self.classifier.to(self.dtype)
+        router_logits = self.classifier(hidden_states)
+
+        # Apply Softmax and cast back to the original `dtype`
+        router_probs = nn.functional.softmax(router_logits, dim=-1, dtype=self.dtype).to(self.input_dtype)
+        router_logits, expert_index = torch.max(router_probs, dim=-1, keepdim=True)
+        expert_index = torch.nn.functional.one_hot(expert_index, num_classes=self.num_experts)
+        token_priority = torch.cumsum(expert_index, dim=-2)
+        # mask if the token routed to the expert will overflow
+        expert_capacity_mask = token_priority <= self.expert_capacity
+        expert_index = expert_index * expert_capacity_mask
+        router_probs = torch.max(router_probs, dim=-1).values.unsqueeze(-1)
+        return router_probs, expert_index, router_logits
+
+
+class SwitchTransformersLayerNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        Construct a layernorm module in the SWITCH_TRANSFORMERS style. No bias and no subtraction of mean.
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        # SWITCH_TRANSFORMERS uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean
+        # Square Layer Normalization https://huggingface.co/papers/1910.07467 thus variance is calculated
+        # w/o mean and there is no bias. Additionally we want to make sure that the accumulation for
+        # half-precision inputs is done in fp32
+
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+
+        # convert into half-precision if necessary
+        if self.weight.dtype in [torch.float16, torch.bfloat16]:
+            hidden_states = hidden_states.to(self.weight.dtype)
+
+        return self.weight * hidden_states
+
+
+class SwitchTransformersDenseActDense(nn.Module):
+    def __init__(self, config: SwitchTransformersConfig):
+        super().__init__()
+        self.wi = nn.Linear(config.d_model, config.d_ff, bias=False)
+        self.wo = nn.Linear(config.d_ff, config.d_model, bias=False)
+        self.dropout = nn.Dropout(config.dropout_rate)
+        self.act = ACT2FN[config.dense_act_fn]
+
+    def forward(self, hidden_states):
+        hidden_states = self.wi(hidden_states)
+        hidden_states = self.act(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        if (
+            isinstance(self.wo.weight, torch.Tensor)
+            and hidden_states.dtype != self.wo.weight.dtype
+            and self.wo.weight.dtype != torch.int8
+        ):
+            hidden_states = hidden_states.to(self.wo.weight.dtype)
+        hidden_states = self.wo(hidden_states)
+        return hidden_states
+
+
+class SwitchTransformersExperts(nn.ModuleDict):
+    def __init__(self, config: SwitchTransformersConfig):
+        super().__init__()
+        self.num_experts = config.num_experts
+        for idx in range(config.num_experts):
+            self[f"expert_{idx}"] = SwitchTransformersDenseActDense(config)
+
+    def forward(
+        self, hidden_states: torch.Tensor, selected_experts: torch.Tensor, routing_weights: torch.Tensor
+    ) -> torch.Tensor:
+        final_hidden_states = torch.zeros_like(hidden_states)
+        expert_mask = selected_experts.permute(2, 1, 0)
+
+        expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+        for expert_idx in expert_hit:
+            idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0))
+            current_state = hidden_states[None, top_x].reshape(-1, hidden_states.shape[-1])
+            current_hidden_states = self[f"expert_{expert_idx[0]}"](current_state) * routing_weights[top_x, idx, None]
+            final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
+        return final_hidden_states
+
+
+class SwitchTransformersSparseMLP(nn.Module):  # inherit from mixtral
+    def __init__(self, config: SwitchTransformersConfig):
+        super().__init__()
+        self.router = SwitchTransformersTop1Router(config)
+        self.experts = SwitchTransformersExperts(config)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        _, selected_experts, routing_weights = self.router(hidden_states)
+        hidden_states = self.experts(hidden_states, selected_experts, routing_weights)
+        hidden_states = hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+        return hidden_states
+
+
+class SwitchTransformersLayerFF(nn.Module):
+    r"""
+    Switch Transformers Feed Forward layer module. This is a wrapper around the Mixture of Experts module.
+
+    Parameters:
+        config : ([`SwitchTransformersConfig`]): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+        is_sparse (`bool`):
+            Whether the MLP layer is a `Sparse` layer (contains a Mixture of Experts) or not
+    """
+
+    def __init__(self, config: SwitchTransformersConfig, is_sparse=False):
+        super().__init__()
+        self.is_sparse = is_sparse
+
+        # Check if it is a sparse layer, if not then it is a dense layer
+        if not self.is_sparse:
+            self.mlp = SwitchTransformersDenseActDense(config)
+        else:
+            self.mlp = SwitchTransformersSparseMLP(config)
+
+        self.layer_norm = SwitchTransformersLayerNorm(config.d_model, eps=config.layer_norm_epsilon)
+        self.dropout = nn.Dropout(config.dropout_rate)
+
+    def forward(self, hidden_states, **kwargs):
+        forwarded_states = self.layer_norm(hidden_states)
+        forwarded_states = self.mlp(forwarded_states)
+        output = hidden_states + self.dropout(forwarded_states)
+        return output
+
+
+class SwitchTransformersAttention(nn.Module):
+    def __init__(
+        self,
+        config: SwitchTransformersConfig,
+        has_relative_attention_bias=False,
+        layer_idx: int | None = None,
+    ):
+        super().__init__()
+        self.is_decoder = config.is_decoder
+        self.has_relative_attention_bias = has_relative_attention_bias
+        self.relative_attention_num_buckets = config.relative_attention_num_buckets
+        self.relative_attention_max_distance = config.relative_attention_max_distance
+        self.d_model = config.d_model
+        self.key_value_proj_dim = config.d_kv
+        self.n_heads = config.num_heads
+        self.dropout = config.dropout_rate
+        self.inner_dim = self.n_heads * self.key_value_proj_dim
+        self.layer_idx = layer_idx
+        if layer_idx is None and self.is_decoder:
+            logger.warning_once(
+                f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and "
+                "will to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+                "when creating this class."
+            )
+
+        self.q = nn.Linear(self.d_model, self.inner_dim, bias=False)
+        self.k = nn.Linear(self.d_model, self.inner_dim, bias=False)
+        self.v = nn.Linear(self.d_model, self.inner_dim, bias=False)
+        self.o = nn.Linear(self.inner_dim, self.d_model, bias=False)
+
+        if self.has_relative_attention_bias:
+            self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads)
+
+        self.gradient_checkpointing = False
+
+    @staticmethod
+    def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):
+        """
+        Adapted from Mesh Tensorflow:
+        https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593
+
+        Translate relative position to a bucket number for relative attention. The relative position is defined as
+        memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to
+        position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for
+        small absolute relative_position and larger buckets for larger absolute relative_positions. All relative
+        positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket.
+        This should allow for more graceful generalization to longer sequences than the model has been trained on
+
+        Args:
+            relative_position: an int32 Tensor
+            bidirectional: a boolean - whether the attention is bidirectional
+            num_buckets: an integer
+            max_distance: an integer
+
+        Returns:
+            a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets)
+        """
+        relative_buckets = 0
+        if bidirectional:
+            num_buckets //= 2
+            relative_buckets += (relative_position > 0).to(torch.long) * num_buckets
+            relative_position = torch.abs(relative_position)
+        else:
+            relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))
+        # now relative_position is in the range [0, inf)
+
+        # half of the buckets are for exact increments in positions
+        max_exact = num_buckets // 2
+        is_small = relative_position < max_exact
+
+        # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance
+        relative_position_if_large = max_exact + (
+            torch.log(relative_position.float() / max_exact)
+            / math.log(max_distance / max_exact)
+            * (num_buckets - max_exact)
+        ).to(torch.long)
+        relative_position_if_large = torch.min(
+            relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1)
+        )
+
+        relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)
+        return relative_buckets
+
+    def compute_bias(self, query_length, key_length, device=None, past_seen_tokens=0):
+        """Compute binned relative position bias"""
+        if device is None:
+            device = self.relative_attention_bias.weight.device
+        context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None] + past_seen_tokens
+        memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
+        relative_position = memory_position - context_position  # shape (query_length, key_length)
+        relative_position_bucket = self._relative_position_bucket(
+            relative_position,  # shape (query_length, key_length)
+            bidirectional=(not self.is_decoder),
+            num_buckets=self.relative_attention_num_buckets,
+            max_distance=self.relative_attention_max_distance,
+        )
+        values = self.relative_attention_bias(relative_position_bucket)  # shape (query_length, key_length, num_heads)
+        values = values.permute([2, 0, 1]).unsqueeze(0)  # shape (1, num_heads, query_length, key_length)
+        return values
+
+    def forward(
+        self,
+        hidden_states,
+        mask=None,
+        key_value_states=None,
+        position_bias=None,
+        past_key_values=None,
+        output_attentions=False,
+        **kwargs,
+    ):
+        """
+        Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).
+        """
+        # Input is (batch_size, seq_length, dim)
+        # Mask is (batch_size, 1, 1, key_length) (non-causal encoder) or (batch_size, 1, seq_length, key_length) (causal decoder)
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.key_value_proj_dim)
+        past_seen_tokens = past_key_values.get_seq_length(self.layer_idx) if past_key_values is not None else 0
+        # We clone here for StaticCache, as we get the value before updating it, but use it after and it's the same ref
+        past_seen_tokens = past_seen_tokens.clone() if isinstance(past_seen_tokens, torch.Tensor) else past_seen_tokens
+
+        # if key_value_states are provided this layer is used as a cross-attention layer for the decoder
+        is_cross_attention = key_value_states is not None
+
+        query_states = self.q(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        # Check is encoder-decoder model is being used. Otherwise we'll get `DynamicCache`
+        is_updated = False
+        if isinstance(past_key_values, EncoderDecoderCache):
+            is_updated = past_key_values.is_updated.get(self.layer_idx)
+            if is_cross_attention:
+                # after the first generated id, we can subsequently re-use all key/value_states from cache
+                curr_past_key_values = past_key_values.cross_attention_cache
+            else:
+                curr_past_key_values = past_key_values.self_attention_cache
+        else:
+            curr_past_key_values = past_key_values
+
+        current_states = key_value_states if is_cross_attention else hidden_states
+        if is_cross_attention and past_key_values is not None and is_updated:
+            # reuse k,v, cross_attentions
+            key_states = curr_past_key_values.layers[self.layer_idx].keys
+            value_states = curr_past_key_values.layers[self.layer_idx].values
+        else:
+            kv_shape = (*current_states.shape[:-1], -1, self.key_value_proj_dim)
+            key_states = self.k(current_states).view(kv_shape).transpose(1, 2)
+            value_states = self.v(current_states).view(kv_shape).transpose(1, 2)
+
+            if past_key_values is not None:
+                key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx)
+                # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
+                if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):
+                    past_key_values.is_updated[self.layer_idx] = True
+
+        # compute scores, equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9
+        scores = torch.matmul(query_states, key_states.transpose(3, 2))
+
+        if position_bias is None:
+            key_length = key_states.shape[-2]
+            if not self.has_relative_attention_bias:
+                position_bias = torch.zeros(
+                    (1, query_states.shape[1], input_shape[1], key_length), device=scores.device, dtype=scores.dtype
+                )
+                if self.gradient_checkpointing and self.training:
+                    position_bias.requires_grad = True
+            else:
+                position_bias = self.compute_bias(
+                    input_shape[1], key_length, device=scores.device, past_seen_tokens=past_seen_tokens
+                )
+
+            if mask is not None:
+                causal_mask = mask[:, :, :, : key_states.shape[-2]]
+                position_bias = position_bias + causal_mask
+
+        position_bias_masked = position_bias
+        scores += position_bias_masked
+
+        # (batch_size, n_heads, seq_length, key_length)
+        attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
+
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(*input_shape, -1)
+        attn_output = self.o(attn_output)
+
+        outputs = (attn_output, position_bias)
+
+        if output_attentions:
+            outputs = outputs + (attn_weights,)
+        return outputs
+
+
+class SwitchTransformersLayerSelfAttention(nn.Module):
+    def __init__(self, config, has_relative_attention_bias=False, layer_idx: int | None = None):
+        super().__init__()
+        self.SelfAttention = SwitchTransformersAttention(
+            config, has_relative_attention_bias=has_relative_attention_bias, layer_idx=layer_idx
+        )
+        self.layer_norm = SwitchTransformersLayerNorm(config.d_model, eps=config.layer_norm_epsilon)
+        self.dropout = nn.Dropout(config.dropout_rate)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        position_bias=None,
+        past_key_values=None,
+        use_cache=False,
+        output_attentions=False,
+        **kwargs,
+    ):
+        normed_hidden_states = self.layer_norm(hidden_states)
+        attention_output = self.SelfAttention(
+            normed_hidden_states,
+            mask=attention_mask,
+            position_bias=position_bias,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+        )
+        hidden_states = hidden_states + self.dropout(attention_output[0])
+        outputs = (hidden_states,) + attention_output[1:]  # add attentions if we output them
+        return outputs
+
+
+class SwitchTransformersLayerCrossAttention(nn.Module):
+    def __init__(self, config, layer_idx: int | None = None):
+        super().__init__()
+        self.EncDecAttention = SwitchTransformersAttention(
+            config, has_relative_attention_bias=False, layer_idx=layer_idx
+        )
+        self.layer_norm = SwitchTransformersLayerNorm(config.d_model, eps=config.layer_norm_epsilon)
+        self.dropout = nn.Dropout(config.dropout_rate)
+
+    def forward(
+        self,
+        hidden_states,
+        key_value_states,
+        attention_mask=None,
+        position_bias=None,
+        past_key_values=None,
+        output_attentions=False,
+        **kwargs,
+    ):
+        normed_hidden_states = self.layer_norm(hidden_states)
+        attention_output = self.EncDecAttention(
+            normed_hidden_states,
+            mask=attention_mask,
+            key_value_states=key_value_states,
+            position_bias=position_bias,
+            past_key_values=past_key_values,
+            output_attentions=output_attentions,
+        )
+        layer_output = hidden_states + self.dropout(attention_output[0])
+        outputs = (layer_output,) + attention_output[1:]  # add attentions if we output them
+        return outputs
+
+
+class SwitchTransformersBlock(GradientCheckpointingLayer):
+    def __init__(self, config, has_relative_attention_bias=False, is_sparse=False, layer_idx: int | None = None):
+        super().__init__()
+        self.is_decoder = config.is_decoder
+        self.is_sparse = is_sparse
+        self.layer = nn.ModuleList()
+        self.layer.append(
+            SwitchTransformersLayerSelfAttention(
+                config, has_relative_attention_bias=has_relative_attention_bias, layer_idx=layer_idx
+            )
+        )
+        if self.is_decoder:
+            self.layer.append(SwitchTransformersLayerCrossAttention(config, layer_idx=layer_idx))
+
+        self.layer.append(SwitchTransformersLayerFF(config, is_sparse=self.is_sparse))
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        position_bias=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        encoder_decoder_position_bias=None,
+        past_key_values=None,
+        use_cache=False,
+        **kwargs,
+    ):
+        hidden_states, _ = self.layer[0](
+            hidden_states,
+            attention_mask=attention_mask,
+            position_bias=position_bias,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+        )
+
+        # clamp inf values to enable fp16 training
+        if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any():
+            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+        do_cross_attention = self.is_decoder and encoder_hidden_states is not None
+        if do_cross_attention:
+            hidden_states, _ = self.layer[1](
+                hidden_states,
+                key_value_states=encoder_hidden_states,
+                attention_mask=encoder_attention_mask,
+                position_bias=encoder_decoder_position_bias,
+                past_key_values=past_key_values,
+            )
+
+            # clamp inf values to enable fp16 training
+            if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any():
+                clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+                hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+        hidden_states = self.layer[-1](hidden_states)
+        # clamp inf values to enable fp16 training
+        if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any():
+            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+        return hidden_states
+
+
+@auto_docstring
+class SwitchTransformersPreTrainedModel(PreTrainedModel):
+    config: SwitchTransformersConfig
+    base_model_prefix = "switch_transformers"
+    supports_gradient_checkpointing = True
+    _can_compile_fullgraph = False
+    _no_split_modules = ["SwitchTransformersBlock"]
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        factor = self.config.initializer_factor  # Used for testing weights initialization
+        if isinstance(module, SwitchTransformersLayerNorm):
+            init.constant_(module.weight, factor * 1.0)
+        elif isinstance(
+            module,
+            (SwitchTransformersModel, SwitchTransformersForConditionalGeneration, SwitchTransformersEncoderModel),
+        ):
+            init.normal_(module.shared.weight, mean=0.0, std=factor * 1.0)
+            if hasattr(module, "lm_head") and not self.config.tie_word_embeddings:
+                init.normal_(module.lm_head.weight, mean=0.0, std=factor * 1.0)
+        elif isinstance(module, SwitchTransformersDenseActDense):
+            init.normal_(module.wi.weight, mean=0.0, std=factor * ((self.config.d_model) ** -0.5))
+            if hasattr(module.wi, "bias") and module.wi.bias is not None:
+                init.zeros_(module.wi.bias)
+            init.normal_(module.wo.weight, mean=0.0, std=factor * ((self.config.d_ff) ** -0.5))
+            if hasattr(module.wo, "bias") and module.wo.bias is not None:
+                init.zeros_(module.wo.bias)
+        elif isinstance(module, SwitchTransformersAttention):
+            d_model = self.config.d_model
+            key_value_proj_dim = self.config.d_kv
+            n_heads = self.config.num_heads
+            init.normal_(module.q.weight, mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5))
+            init.normal_(module.k.weight, mean=0.0, std=factor * (d_model**-0.5))
+            init.normal_(module.v.weight, mean=0.0, std=factor * (d_model**-0.5))
+            init.normal_(module.o.weight, mean=0.0, std=factor * ((n_heads * key_value_proj_dim) ** -0.5))
+            if module.has_relative_attention_bias:
+                init.normal_(module.relative_attention_bias.weight, mean=0.0, std=factor * ((d_model) ** -0.5))
+        elif isinstance(module, SwitchTransformersSparseMLP):
+            d_model = self.config.d_model
+            key_value_proj_dim = self.config.d_kv
+            n_heads = self.config.num_heads
+            init.normal_(module.router.classifier.weight, mean=0.0, std=factor * 1)
+            for idx in range(self.config.num_experts):
+                init.normal_(module.experts[f"expert_{idx}"].wi.weight, mean=0.0, std=factor * (d_model**-0.5))
+                init.normal_(module.experts[f"expert_{idx}"].wo.weight, mean=0.0, std=factor * (d_model**-0.5))
+
+    def _shift_right(self, input_ids):
+        decoder_start_token_id = self.config.decoder_start_token_id
+        pad_token_id = self.config.pad_token_id
+
+        if decoder_start_token_id is None:
+            raise ValueError(
+                "self.model.config.decoder_start_token_id has to be defined. In SwitchTransformers it is usually set"
+                " to the pad_token_id. See SwitchTransformers docs for more information"
+            )
+
+        shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+        shifted_input_ids[..., 1:] = input_ids[..., :-1].clone()
+        shifted_input_ids[..., 0] = decoder_start_token_id
+
+        if pad_token_id is None:
+            raise ValueError("self.model.config.pad_token_id has to be defined.")
+        # replace possible -100 values in labels by `pad_token_id`
+        shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+        return shifted_input_ids
+
+
+class SwitchTransformersStack(SwitchTransformersPreTrainedModel):
+    _can_record_outputs = {
+        "hidden_states": SwitchTransformersBlock,
+        "attentions": OutputRecorder(SwitchTransformersAttention, index=-1, layer_name="layer.0"),
+        "cross_attentions": OutputRecorder(SwitchTransformersAttention, index=-1, layer_name="layer.1"),
+        "router_logits": OutputRecorder(SwitchTransformersTop1Router, index=2),
+    }
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model)
+
+        self.is_decoder = config.is_decoder
+
+        sparse_step = config.decoder_sparse_step if self.is_decoder else config.encoder_sparse_step
+        config.num_layers = config.num_decoder_layers if self.is_decoder else config.num_layers
+        self.block = nn.ModuleList()
+        for i in range(config.num_layers):
+            is_sparse = (i % sparse_step == 1 or sparse_step == 1) if sparse_step > 0 else False
+
+            self.block.append(
+                SwitchTransformersBlock(
+                    config, has_relative_attention_bias=bool(i == 0), is_sparse=is_sparse, layer_idx=i
+                )
+            )
+
+        self.final_layer_norm = SwitchTransformersLayerNorm(config.d_model, eps=config.layer_norm_epsilon)
+        self.dropout = nn.Dropout(config.dropout_rate)
+        self.post_init()
+
+        self.gradient_checkpointing = False
+
+    @merge_with_config_defaults
+    @capture_outputs
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        inputs_embeds=None,
+        past_key_values=None,
+        use_cache=None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | MoEModelOutputWithPastAndCrossAttentions:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            if self.embed_tokens is None:
+                raise ValueError("You have to initialize the model with valid token embeddings")
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        batch_size, seq_length = inputs_embeds.shape[:2]
+
+        if use_cache is True:
+            if not self.is_decoder:
+                raise ValueError(f"`use_cache` can only be set to `True` if {self} is used as a decoder")
+
+        if self.is_decoder:
+            if use_cache and past_key_values is None:
+                if self.config.is_encoder_decoder:
+                    past_key_values = EncoderDecoderCache(
+                        DynamicCache(config=self.config), DynamicCache(config=self.config)
+                    )
+                else:
+                    past_key_values = DynamicCache(config=self.config)
+        elif not self.is_decoder:
+            # do not pass cache object down the line for encoder stack
+            # it messes indexing later in decoder-stack because cache object is modified in-place
+            past_key_values = None
+
+        past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+        if attention_mask is None and not is_torchdynamo_compiling():
+            # required mask seq length can be calculated via length of past cache
+            mask_seq_length = past_key_values_length + seq_length
+            attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)
+
+        if self.config.is_decoder:
+            causal_mask = create_causal_mask(
+                config=self.config,
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+            )
+        else:
+            causal_mask = attention_mask[:, None, None, :]
+            causal_mask = causal_mask.to(dtype=inputs_embeds.dtype)
+            causal_mask = (1.0 - causal_mask) * torch.finfo(inputs_embeds.dtype).min
+
+        if encoder_attention_mask is not None:
+            encoder_attention_mask = create_bidirectional_mask(
+                config=self.config,
+                inputs_embeds=inputs_embeds,
+                attention_mask=encoder_attention_mask,
+                encoder_hidden_states=encoder_hidden_states,
+            )
+
+        position_bias = None
+        encoder_decoder_position_bias = None
+
+        hidden_states = self.dropout(inputs_embeds)
+
+        for i, layer_module in enumerate(self.block):
+            hidden_states = layer_module(
+                hidden_states,
+                causal_mask,
+                position_bias,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                encoder_decoder_position_bias,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                **kwargs,
+            )
+
+        hidden_states = self.final_layer_norm(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+
+        return MoEModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+@auto_docstring
+class SwitchTransformersModel(SwitchTransformersPreTrainedModel):
+    _tied_weights_keys = {
+        "encoder.embed_tokens.weight": "shared.weight",
+        "decoder.embed_tokens.weight": "shared.weight",
+    }
+    _input_embed_layer = "shared"
+
+    def __init__(self, config: SwitchTransformersConfig):
+        super().__init__(config)
+        self.shared = nn.Embedding(config.vocab_size, config.d_model)
+
+        encoder_config = copy.deepcopy(config)
+        encoder_config.is_decoder = False
+        encoder_config.use_cache = False
+        self.encoder = SwitchTransformersStack(encoder_config)
+
+        decoder_config = copy.deepcopy(config)
+        decoder_config.is_decoder = True
+        self.decoder = SwitchTransformersStack(decoder_config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def set_input_embeddings(self, new_embeddings):
+        self.shared = new_embeddings
+        self.encoder.set_input_embeddings(new_embeddings)
+        self.decoder.set_input_embeddings(new_embeddings)
+
+    @auto_docstring
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.FloatTensor | None = None,
+        decoder_input_ids: torch.LongTensor | None = None,
+        decoder_attention_mask: torch.BoolTensor | None = None,
+        encoder_outputs: tuple[tuple[torch.FloatTensor]] | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        decoder_inputs_embeds: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.FloatTensor] | Seq2SeqMoEModelOutput:
+        if encoder_outputs is None:
+            encoder_outputs = self.encoder(
+                input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, **kwargs
+            )
+
+        hidden_states = encoder_outputs[0]
+        decoder_outputs = self.decoder(
+            input_ids=decoder_input_ids,
+            attention_mask=decoder_attention_mask,
+            inputs_embeds=decoder_inputs_embeds,
+            past_key_values=past_key_values,
+            encoder_hidden_states=hidden_states,
+            encoder_attention_mask=attention_mask,
+            **kwargs,
+        )
+
+        return Seq2SeqMoEModelOutput(
+            last_hidden_state=decoder_outputs.last_hidden_state,
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.hidden_states,
+            decoder_attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+            decoder_router_logits=decoder_outputs.router_logits,
+            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+            encoder_hidden_states=encoder_outputs.hidden_states,
+            encoder_attentions=encoder_outputs.attentions,
+            encoder_router_logits=encoder_outputs.router_logits,
+        )
+
+
+####################################################
+# This dict contains ids and associated url
+# for the pretrained weights provided with the models
+####################################################
+
+
+def router_z_loss_func(router_logits: torch.Tensor) -> float:
+    r"""
+    Compute the router z-loss implemented in PyTorch.
+
+    The router z-loss was introduced in [Designing Effective Sparse Expert Models](https://huggingface.co/papers/2202.08906).
+    It encourages router logits to remain small in an effort to improve stability.
+
+    Args:
+        router_logits (`float`):
+            Input logits of shape [batch_size, sequence_length, num_experts]
+
+    Returns:
+        Scalar router z-loss.
+    """
+    num_groups, tokens_per_group, _ = router_logits.shape
+    log_z = torch.logsumexp(router_logits, dim=-1)
+    z_loss = log_z**2
+    return torch.sum(z_loss) / (num_groups * tokens_per_group)
+
+
+def load_balancing_loss_func(router_probs: torch.Tensor, expert_indices: torch.Tensor) -> float:
+    r"""
+    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
+
+    See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
+    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
+    experts is too unbalanced.
+
+    Args:
+        router_probs (`torch.Tensor`):
+            Probability assigned to each expert per token. Shape: [batch_size, sequence_length, num_experts].
+        expert_indices (`torch.Tensor`):
+            Indices tensor of shape [batch_size, sequence_length] identifying the selected expert for a given token.
+
+    Returns:
+        The auxiliary loss.
+    """
+    num_experts = router_probs.shape[-1]
+
+    # cast the expert indices to int64, otherwise one-hot encoding will fail
+    if expert_indices.dtype != torch.int64:
+        expert_indices = expert_indices.to(torch.int64)
+
+    if len(expert_indices.shape) == 2:
+        expert_indices = expert_indices.unsqueeze(2)
+
+    expert_mask = torch.nn.functional.one_hot(expert_indices, num_experts)
+
+    # For a given token, determine if it was routed to a given expert.
+    expert_mask = torch.max(expert_mask, axis=-2).values
+
+    # cast to float32 otherwise mean will fail
+    expert_mask = expert_mask.to(torch.float32)
+    tokens_per_group_and_expert = torch.mean(expert_mask, axis=-2)
+
+    router_prob_per_group_and_expert = torch.mean(router_probs, axis=-2)
+    return torch.mean(tokens_per_group_and_expert * router_prob_per_group_and_expert) * (num_experts**2)
+
+
+@auto_docstring(
+    custom_intro="""
+    SWITCH_TRANSFORMERS Model with a `language modeling` head on top.
+    """
+)
+class SwitchTransformersForConditionalGeneration(SwitchTransformersPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {
+        "encoder.embed_tokens.weight": "shared.weight",
+        "decoder.embed_tokens.weight": "shared.weight",
+        "lm_head.weight": "shared.weight",
+    }
+
+    def __init__(self, config: SwitchTransformersConfig):
+        super().__init__(config)
+        self.model_dim = config.d_model
+
+        self.shared = nn.Embedding(config.vocab_size, config.d_model)
+
+        encoder_config = copy.deepcopy(config)
+        encoder_config.is_decoder = False
+        encoder_config.use_cache = False
+        self.encoder = SwitchTransformersStack(encoder_config)
+
+        decoder_config = copy.deepcopy(config)
+        decoder_config.is_decoder = True
+        decoder_config.num_layers = config.num_decoder_layers
+        self.decoder = SwitchTransformersStack(decoder_config)
+
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+
+        self.router_z_loss_coef = config.router_z_loss_coef
+        self.router_aux_loss_coef = config.router_aux_loss_coef
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.shared
+
+    def set_input_embeddings(self, new_embeddings):
+        self.shared = new_embeddings
+        self.encoder.set_input_embeddings(new_embeddings)
+        self.decoder.set_input_embeddings(new_embeddings)
+
+    @auto_docstring
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.FloatTensor | None = None,
+        decoder_input_ids: torch.LongTensor | None = None,
+        decoder_attention_mask: torch.BoolTensor | None = None,
+        encoder_outputs: tuple[tuple[torch.Tensor]] | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        decoder_inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        output_router_logits: bool | None = False,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.FloatTensor] | Seq2SeqMoEOutput:
+        if encoder_outputs is None:
+            encoder_outputs = self.encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                output_router_logits=output_router_logits,
+                **kwargs,
+            )
+
+        hidden_states = encoder_outputs[0]
+
+        if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None:
+            # get decoder inputs from shifting lm labels to the right
+            decoder_input_ids = self._shift_right(labels)
+
+        # Decode
+        decoder_outputs = self.decoder(
+            input_ids=decoder_input_ids,
+            attention_mask=decoder_attention_mask,
+            inputs_embeds=decoder_inputs_embeds,
+            past_key_values=past_key_values,
+            encoder_hidden_states=hidden_states,
+            encoder_attention_mask=attention_mask,
+            output_router_logits=output_router_logits,
+            **kwargs,
+        )
+
+        sequence_output = decoder_outputs.last_hidden_state
+
+        if self.config.tie_word_embeddings:
+            # Rescale output before projecting on vocab
+            # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586
+            sequence_output = sequence_output * (self.model_dim**-0.5)
+
+        lm_logits = self.lm_head(sequence_output)
+
+        loss = None
+        encoder_z_loss = None
+        encoder_aux_loss = None
+        decoder_z_loss = None
+        decoder_aux_loss = None
+
+        if output_router_logits:
+            # Compute the router loss (z_loss + auxiliary loss) for each router in the encoder and decoder
+            if self.encoder.config.encoder_sparse_step > 1:
+                encoder_router_logits, encoder_expert_indexes = self._unpack_router_logits(encoder_outputs[-1])
+                encoder_z_loss = router_z_loss_func(encoder_router_logits)
+                encoder_router_probs = nn.Softmax(dim=-1)(encoder_router_logits)
+                encoder_aux_loss = load_balancing_loss_func(encoder_router_probs, encoder_expert_indexes)
+            else:
+                encoder_z_loss = 0
+                encoder_aux_loss = 0
+
+            if self.decoder.config.decoder_sparse_step > 1:
+                decoder_router_logits, decoder_expert_indexes = self._unpack_router_logits(decoder_outputs[-1])
+                decoder_z_loss = router_z_loss_func(decoder_router_logits)
+                decoder_router_probs = nn.Softmax(dim=-1)(decoder_router_logits)
+                decoder_aux_loss = load_balancing_loss_func(decoder_router_probs, decoder_expert_indexes)
+            else:
+                decoder_z_loss = 0
+                decoder_aux_loss = 0
+
+        if labels is not None:
+            loss_fct = CrossEntropyLoss(ignore_index=-100)
+            # move labels to correct device to enable PP
+            labels = labels.to(lm_logits.device)
+            loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))
+
+            if output_router_logits:
+                z_loss = self.router_z_loss_coef * (encoder_z_loss + decoder_z_loss)
+                aux_loss = self.router_aux_loss_coef * (encoder_aux_loss + decoder_aux_loss)
+                loss = loss + z_loss + aux_loss
+
+        return Seq2SeqMoEOutput(
+            loss=loss,
+            logits=lm_logits,
+            encoder_z_loss=encoder_z_loss,
+            encoder_aux_loss=encoder_aux_loss,
+            decoder_z_loss=decoder_z_loss,
+            decoder_aux_loss=decoder_aux_loss,
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.hidden_states,
+            decoder_attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+            decoder_router_logits=decoder_outputs.router_logits,
+            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+            encoder_hidden_states=encoder_outputs.hidden_states,
+            encoder_attentions=encoder_outputs.attentions,
+            encoder_router_logits=encoder_outputs.router_logits,
+        )
+
+    def _unpack_router_logits(self, router_outputs):
+        total_router_logits = []
+        total_expert_indexes = []
+        for router_output in router_outputs:
+            if len(router_output[0].shape) > 1:
+                router_logits, expert_indexes = router_output
+                total_router_logits.append(router_logits)
+                total_expert_indexes.append(expert_indexes)
+        return torch.cat(total_router_logits, dim=1), torch.cat(total_expert_indexes, dim=1)
+
+    def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
+        return self._shift_right(labels)
+
+
+class SwitchTransformersEncoderModel(SwitchTransformersPreTrainedModel):
+    _tied_weights_keys = {
+        "encoder.embed_tokens.weight": "shared.weight",
+    }
+
+    def __init__(self, config: SwitchTransformersConfig):
+        super().__init__(config)
+        self.shared = nn.Embedding(config.vocab_size, config.d_model)
+
+        encoder_config = copy.deepcopy(config)
+        encoder_config.use_cache = False
+        encoder_config.is_encoder_decoder = False
+        self.encoder = SwitchTransformersStack(encoder_config)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.shared
+
+    def set_input_embeddings(self, new_embeddings):
+        self.shared = new_embeddings
+        self.encoder.set_input_embeddings(new_embeddings)
+
+    @auto_docstring
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.FloatTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.FloatTensor] | MoEModelOutput:
+        use_cache = False
+        encoder_outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+        return encoder_outputs
+
+
+__all__ = [
+    "SwitchTransformersEncoderModel",
+    "SwitchTransformersForConditionalGeneration",
+    "SwitchTransformersModel",
+    "SwitchTransformersPreTrainedModel",
+    "SwitchTransformersTop1Router",
+    "SwitchTransformersSparseMLP",
+]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/switch_transformers/modular_switch_transformers.py

@@ -0,0 +1,810 @@
+# Copyright 2022 SwitchTransformers Authors and 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.
+"""PyTorch SwitchTransformers model."""
+
+import copy
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+
+from ... import initialization as init
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_bidirectional_mask, create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+    MoEModelOutput,
+    MoEModelOutputWithPastAndCrossAttentions,
+    Seq2SeqMoEModelOutput,
+    Seq2SeqMoEOutput,
+)
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import (
+    TransformersKwargs,
+    auto_docstring,
+    is_torchdynamo_compiling,
+    logging,
+)
+from ...utils.generic import (
+    can_return_tuple,
+    merge_with_config_defaults,
+)
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from ..t5.modeling_t5 import T5Attention, T5DenseActDense, T5LayerCrossAttention, T5LayerNorm, T5LayerSelfAttention
+from .configuration_switch_transformers import SwitchTransformersConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+####################################################
+# This dict contains ids and associated url
+# for the pretrained weights provided with the models
+####################################################
+
+
+def router_z_loss_func(router_logits: torch.Tensor) -> float:
+    r"""
+    Compute the router z-loss implemented in PyTorch.
+
+    The router z-loss was introduced in [Designing Effective Sparse Expert Models](https://huggingface.co/papers/2202.08906).
+    It encourages router logits to remain small in an effort to improve stability.
+
+    Args:
+        router_logits (`float`):
+            Input logits of shape [batch_size, sequence_length, num_experts]
+
+    Returns:
+        Scalar router z-loss.
+    """
+    num_groups, tokens_per_group, _ = router_logits.shape
+    log_z = torch.logsumexp(router_logits, dim=-1)
+    z_loss = log_z**2
+    return torch.sum(z_loss) / (num_groups * tokens_per_group)
+
+
+def load_balancing_loss_func(router_probs: torch.Tensor, expert_indices: torch.Tensor) -> float:
+    r"""
+    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
+
+    See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
+    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
+    experts is too unbalanced.
+
+    Args:
+        router_probs (`torch.Tensor`):
+            Probability assigned to each expert per token. Shape: [batch_size, sequence_length, num_experts].
+        expert_indices (`torch.Tensor`):
+            Indices tensor of shape [batch_size, sequence_length] identifying the selected expert for a given token.
+
+    Returns:
+        The auxiliary loss.
+    """
+    num_experts = router_probs.shape[-1]
+
+    # cast the expert indices to int64, otherwise one-hot encoding will fail
+    if expert_indices.dtype != torch.int64:
+        expert_indices = expert_indices.to(torch.int64)
+
+    if len(expert_indices.shape) == 2:
+        expert_indices = expert_indices.unsqueeze(2)
+
+    expert_mask = torch.nn.functional.one_hot(expert_indices, num_experts)
+
+    # For a given token, determine if it was routed to a given expert.
+    expert_mask = torch.max(expert_mask, axis=-2).values
+
+    # cast to float32 otherwise mean will fail
+    expert_mask = expert_mask.to(torch.float32)
+    tokens_per_group_and_expert = torch.mean(expert_mask, axis=-2)
+
+    router_prob_per_group_and_expert = torch.mean(router_probs, axis=-2)
+    return torch.mean(tokens_per_group_and_expert * router_prob_per_group_and_expert) * (num_experts**2)
+
+
+class SwitchTransformersTop1Router(nn.Module):
+    """
+    Router using tokens choose top-1 experts assignment.
+
+    This router uses the same mechanism as in Switch Transformer (https://huggingface.co/papers/2101.03961) and V-MoE
+    (https://huggingface.co/papers/2106.05974): tokens choose their top experts. Items are sorted by router_probs and then
+    routed to their choice of expert until the expert's expert_capacity is reached. **There is no guarantee that each
+    token is processed by an expert**, or that each expert receives at least one token.
+
+    """
+
+    def __init__(self, config: SwitchTransformersConfig):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.expert_capacity = config.expert_capacity
+        self.classifier = nn.Linear(config.hidden_size, self.num_experts, bias=config.router_bias)
+        self.jitter_noise = config.router_jitter_noise
+        self.ignore_padding_tokens = config.router_ignore_padding_tokens
+        self.dtype = getattr(torch, config.router_dtype)
+
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        r"""
+        Computes router probabilities from input hidden states.
+
+        Args:
+            hidden_states (`torch.Tensor`):
+                (batch_size, sequence_length, hidden_dim) from which router probabilities are computed.
+        Returns:
+            router_probabilities (`torch.Tensor`):
+                Tensor of shape (batch_size, sequence_length, num_experts) corresponding to the probabilities for each
+                token and expert. Used for routing tokens to experts.
+            router_logits (`torch.Tensor`):
+                Logits tensor of shape (batch_size, sequence_length, num_experts) corresponding to raw router logits.
+                This is used later for computing router z-loss.
+        """
+        # float32 is used to ensure stability. See the discussion of "selective precision" in
+        # https://huggingface.co/papers/2101.03961.
+        # We also store the previous dtype to cast back the output to the previous dtype
+        self.input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(self.dtype)
+        if self.training and self.jitter_noise > 0:
+            # Multiply the token inputs by the uniform distribution - adding some noise
+            hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise)
+        self.classifier = self.classifier.to(self.dtype)
+        router_logits = self.classifier(hidden_states)
+
+        # Apply Softmax and cast back to the original `dtype`
+        router_probs = nn.functional.softmax(router_logits, dim=-1, dtype=self.dtype).to(self.input_dtype)
+        router_logits, expert_index = torch.max(router_probs, dim=-1, keepdim=True)
+        expert_index = torch.nn.functional.one_hot(expert_index, num_classes=self.num_experts)
+        token_priority = torch.cumsum(expert_index, dim=-2)
+        # mask if the token routed to the expert will overflow
+        expert_capacity_mask = token_priority <= self.expert_capacity
+        expert_index = expert_index * expert_capacity_mask
+        router_probs = torch.max(router_probs, dim=-1).values.unsqueeze(-1)
+        return router_probs, expert_index, router_logits
+
+
+class SwitchTransformersLayerNorm(T5LayerNorm):
+    pass
+
+
+class SwitchTransformersDenseActDense(T5DenseActDense):
+    pass
+
+
+class SwitchTransformersExperts(nn.ModuleDict):
+    def __init__(self, config: SwitchTransformersConfig):
+        super().__init__()
+        self.num_experts = config.num_experts
+        for idx in range(config.num_experts):
+            self[f"expert_{idx}"] = SwitchTransformersDenseActDense(config)
+
+    def forward(
+        self, hidden_states: torch.Tensor, selected_experts: torch.Tensor, routing_weights: torch.Tensor
+    ) -> torch.Tensor:
+        final_hidden_states = torch.zeros_like(hidden_states)
+        expert_mask = selected_experts.permute(2, 1, 0)
+
+        expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+        for expert_idx in expert_hit:
+            idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0))
+            current_state = hidden_states[None, top_x].reshape(-1, hidden_states.shape[-1])
+            current_hidden_states = self[f"expert_{expert_idx[0]}"](current_state) * routing_weights[top_x, idx, None]
+            final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
+        return final_hidden_states
+
+
+class SwitchTransformersSparseMLP(nn.Module):  # inherit from mixtral
+    def __init__(self, config: SwitchTransformersConfig):
+        super().__init__()
+        self.router = SwitchTransformersTop1Router(config)
+        self.experts = SwitchTransformersExperts(config)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        _, selected_experts, routing_weights = self.router(hidden_states)
+        hidden_states = self.experts(hidden_states, selected_experts, routing_weights)
+        hidden_states = hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+        return hidden_states
+
+
+class SwitchTransformersLayerFF(nn.Module):
+    r"""
+    Switch Transformers Feed Forward layer module. This is a wrapper around the Mixture of Experts module.
+
+    Parameters:
+        config : ([`SwitchTransformersConfig`]): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+        is_sparse (`bool`):
+            Whether the MLP layer is a `Sparse` layer (contains a Mixture of Experts) or not
+    """
+
+    def __init__(self, config: SwitchTransformersConfig, is_sparse=False):
+        super().__init__()
+        self.is_sparse = is_sparse
+
+        # Check if it is a sparse layer, if not then it is a dense layer
+        if not self.is_sparse:
+            self.mlp = SwitchTransformersDenseActDense(config)
+        else:
+            self.mlp = SwitchTransformersSparseMLP(config)
+
+        self.layer_norm = SwitchTransformersLayerNorm(config.d_model, eps=config.layer_norm_epsilon)
+        self.dropout = nn.Dropout(config.dropout_rate)
+
+    def forward(self, hidden_states, **kwargs):
+        forwarded_states = self.layer_norm(hidden_states)
+        forwarded_states = self.mlp(forwarded_states)
+        output = hidden_states + self.dropout(forwarded_states)
+        return output
+
+
+class SwitchTransformersAttention(T5Attention):
+    pass
+
+
+class SwitchTransformersLayerSelfAttention(T5LayerSelfAttention):
+    pass
+
+
+class SwitchTransformersLayerCrossAttention(T5LayerCrossAttention):
+    pass
+
+
+class SwitchTransformersBlock(GradientCheckpointingLayer):
+    def __init__(self, config, has_relative_attention_bias=False, is_sparse=False, layer_idx: int | None = None):
+        super().__init__()
+        self.is_decoder = config.is_decoder
+        self.is_sparse = is_sparse
+        self.layer = nn.ModuleList()
+        self.layer.append(
+            SwitchTransformersLayerSelfAttention(
+                config, has_relative_attention_bias=has_relative_attention_bias, layer_idx=layer_idx
+            )
+        )
+        if self.is_decoder:
+            self.layer.append(SwitchTransformersLayerCrossAttention(config, layer_idx=layer_idx))
+
+        self.layer.append(SwitchTransformersLayerFF(config, is_sparse=self.is_sparse))
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        position_bias=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        encoder_decoder_position_bias=None,
+        past_key_values=None,
+        use_cache=False,
+        **kwargs,
+    ):
+        hidden_states, _ = self.layer[0](
+            hidden_states,
+            attention_mask=attention_mask,
+            position_bias=position_bias,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+        )
+
+        # clamp inf values to enable fp16 training
+        if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any():
+            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+        do_cross_attention = self.is_decoder and encoder_hidden_states is not None
+        if do_cross_attention:
+            hidden_states, _ = self.layer[1](
+                hidden_states,
+                key_value_states=encoder_hidden_states,
+                attention_mask=encoder_attention_mask,
+                position_bias=encoder_decoder_position_bias,
+                past_key_values=past_key_values,
+            )
+
+            # clamp inf values to enable fp16 training
+            if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any():
+                clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+                hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+        hidden_states = self.layer[-1](hidden_states)
+        # clamp inf values to enable fp16 training
+        if hidden_states.dtype == torch.float16 and torch.isinf(hidden_states).any():
+            clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+            hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+        return hidden_states
+
+
+@auto_docstring
+class SwitchTransformersPreTrainedModel(PreTrainedModel):
+    config: SwitchTransformersConfig
+    base_model_prefix = "switch_transformers"
+    supports_gradient_checkpointing = True
+    _can_compile_fullgraph = False
+    _no_split_modules = ["SwitchTransformersBlock"]
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        factor = self.config.initializer_factor  # Used for testing weights initialization
+        if isinstance(module, SwitchTransformersLayerNorm):
+            init.constant_(module.weight, factor * 1.0)
+        elif isinstance(
+            module,
+            (SwitchTransformersModel, SwitchTransformersForConditionalGeneration, SwitchTransformersEncoderModel),
+        ):
+            init.normal_(module.shared.weight, mean=0.0, std=factor * 1.0)
+            if hasattr(module, "lm_head") and not self.config.tie_word_embeddings:
+                init.normal_(module.lm_head.weight, mean=0.0, std=factor * 1.0)
+        elif isinstance(module, SwitchTransformersDenseActDense):
+            init.normal_(module.wi.weight, mean=0.0, std=factor * ((self.config.d_model) ** -0.5))
+            if hasattr(module.wi, "bias") and module.wi.bias is not None:
+                init.zeros_(module.wi.bias)
+            init.normal_(module.wo.weight, mean=0.0, std=factor * ((self.config.d_ff) ** -0.5))
+            if hasattr(module.wo, "bias") and module.wo.bias is not None:
+                init.zeros_(module.wo.bias)
+        elif isinstance(module, SwitchTransformersAttention):
+            d_model = self.config.d_model
+            key_value_proj_dim = self.config.d_kv
+            n_heads = self.config.num_heads
+            init.normal_(module.q.weight, mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5))
+            init.normal_(module.k.weight, mean=0.0, std=factor * (d_model**-0.5))
+            init.normal_(module.v.weight, mean=0.0, std=factor * (d_model**-0.5))
+            init.normal_(module.o.weight, mean=0.0, std=factor * ((n_heads * key_value_proj_dim) ** -0.5))
+            if module.has_relative_attention_bias:
+                init.normal_(module.relative_attention_bias.weight, mean=0.0, std=factor * ((d_model) ** -0.5))
+        elif isinstance(module, SwitchTransformersSparseMLP):
+            d_model = self.config.d_model
+            key_value_proj_dim = self.config.d_kv
+            n_heads = self.config.num_heads
+            init.normal_(module.router.classifier.weight, mean=0.0, std=factor * 1)
+            for idx in range(self.config.num_experts):
+                init.normal_(module.experts[f"expert_{idx}"].wi.weight, mean=0.0, std=factor * (d_model**-0.5))
+                init.normal_(module.experts[f"expert_{idx}"].wo.weight, mean=0.0, std=factor * (d_model**-0.5))
+
+    def _shift_right(self, input_ids):
+        decoder_start_token_id = self.config.decoder_start_token_id
+        pad_token_id = self.config.pad_token_id
+
+        if decoder_start_token_id is None:
+            raise ValueError(
+                "self.model.config.decoder_start_token_id has to be defined. In SwitchTransformers it is usually set"
+                " to the pad_token_id. See SwitchTransformers docs for more information"
+            )
+
+        shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+        shifted_input_ids[..., 1:] = input_ids[..., :-1].clone()
+        shifted_input_ids[..., 0] = decoder_start_token_id
+
+        if pad_token_id is None:
+            raise ValueError("self.model.config.pad_token_id has to be defined.")
+        # replace possible -100 values in labels by `pad_token_id`
+        shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+        return shifted_input_ids
+
+
+class SwitchTransformersStack(SwitchTransformersPreTrainedModel):
+    _can_record_outputs = {
+        "hidden_states": SwitchTransformersBlock,
+        "attentions": OutputRecorder(SwitchTransformersAttention, index=-1, layer_name="layer.0"),
+        "cross_attentions": OutputRecorder(SwitchTransformersAttention, index=-1, layer_name="layer.1"),
+        "router_logits": OutputRecorder(SwitchTransformersTop1Router, index=2),
+    }
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model)
+
+        self.is_decoder = config.is_decoder
+
+        sparse_step = config.decoder_sparse_step if self.is_decoder else config.encoder_sparse_step
+        config.num_layers = config.num_decoder_layers if self.is_decoder else config.num_layers
+        self.block = nn.ModuleList()
+        for i in range(config.num_layers):
+            is_sparse = (i % sparse_step == 1 or sparse_step == 1) if sparse_step > 0 else False
+
+            self.block.append(
+                SwitchTransformersBlock(
+                    config, has_relative_attention_bias=bool(i == 0), is_sparse=is_sparse, layer_idx=i
+                )
+            )
+
+        self.final_layer_norm = SwitchTransformersLayerNorm(config.d_model, eps=config.layer_norm_epsilon)
+        self.dropout = nn.Dropout(config.dropout_rate)
+        self.post_init()
+
+        self.gradient_checkpointing = False
+
+    @merge_with_config_defaults
+    @capture_outputs
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        inputs_embeds=None,
+        past_key_values=None,
+        use_cache=None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple | MoEModelOutputWithPastAndCrossAttentions:
+        if (input_ids is None) ^ (inputs_embeds is not None):
+            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+        if inputs_embeds is None:
+            if self.embed_tokens is None:
+                raise ValueError("You have to initialize the model with valid token embeddings")
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        batch_size, seq_length = inputs_embeds.shape[:2]
+
+        if use_cache is True:
+            if not self.is_decoder:
+                raise ValueError(f"`use_cache` can only be set to `True` if {self} is used as a decoder")
+
+        if self.is_decoder:
+            if use_cache and past_key_values is None:
+                if self.config.is_encoder_decoder:
+                    past_key_values = EncoderDecoderCache(
+                        DynamicCache(config=self.config), DynamicCache(config=self.config)
+                    )
+                else:
+                    past_key_values = DynamicCache(config=self.config)
+        elif not self.is_decoder:
+            # do not pass cache object down the line for encoder stack
+            # it messes indexing later in decoder-stack because cache object is modified in-place
+            past_key_values = None
+
+        past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+        if attention_mask is None and not is_torchdynamo_compiling():
+            # required mask seq length can be calculated via length of past cache
+            mask_seq_length = past_key_values_length + seq_length
+            attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)
+
+        if self.config.is_decoder:
+            causal_mask = create_causal_mask(
+                config=self.config,
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                past_key_values=past_key_values,
+            )
+        else:
+            causal_mask = attention_mask[:, None, None, :]
+            causal_mask = causal_mask.to(dtype=inputs_embeds.dtype)
+            causal_mask = (1.0 - causal_mask) * torch.finfo(inputs_embeds.dtype).min
+
+        if encoder_attention_mask is not None:
+            encoder_attention_mask = create_bidirectional_mask(
+                config=self.config,
+                inputs_embeds=inputs_embeds,
+                attention_mask=encoder_attention_mask,
+                encoder_hidden_states=encoder_hidden_states,
+            )
+
+        position_bias = None
+        encoder_decoder_position_bias = None
+
+        hidden_states = self.dropout(inputs_embeds)
+
+        for i, layer_module in enumerate(self.block):
+            hidden_states = layer_module(
+                hidden_states,
+                causal_mask,
+                position_bias,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                encoder_decoder_position_bias,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                **kwargs,
+            )
+
+        hidden_states = self.final_layer_norm(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+
+        return MoEModelOutputWithPastAndCrossAttentions(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values,
+        )
+
+
+@auto_docstring
+class SwitchTransformersModel(SwitchTransformersPreTrainedModel):
+    _tied_weights_keys = {
+        "encoder.embed_tokens.weight": "shared.weight",
+        "decoder.embed_tokens.weight": "shared.weight",
+    }
+    _input_embed_layer = "shared"
+
+    def __init__(self, config: SwitchTransformersConfig):
+        super().__init__(config)
+        self.shared = nn.Embedding(config.vocab_size, config.d_model)
+
+        encoder_config = copy.deepcopy(config)
+        encoder_config.is_decoder = False
+        encoder_config.use_cache = False
+        self.encoder = SwitchTransformersStack(encoder_config)
+
+        decoder_config = copy.deepcopy(config)
+        decoder_config.is_decoder = True
+        self.decoder = SwitchTransformersStack(decoder_config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def set_input_embeddings(self, new_embeddings):
+        self.shared = new_embeddings
+        self.encoder.set_input_embeddings(new_embeddings)
+        self.decoder.set_input_embeddings(new_embeddings)
+
+    @auto_docstring
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.FloatTensor | None = None,
+        decoder_input_ids: torch.LongTensor | None = None,
+        decoder_attention_mask: torch.BoolTensor | None = None,
+        encoder_outputs: tuple[tuple[torch.FloatTensor]] | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        decoder_inputs_embeds: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.FloatTensor] | Seq2SeqMoEModelOutput:
+        if encoder_outputs is None:
+            encoder_outputs = self.encoder(
+                input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, **kwargs
+            )
+
+        hidden_states = encoder_outputs[0]
+        decoder_outputs = self.decoder(
+            input_ids=decoder_input_ids,
+            attention_mask=decoder_attention_mask,
+            inputs_embeds=decoder_inputs_embeds,
+            past_key_values=past_key_values,
+            encoder_hidden_states=hidden_states,
+            encoder_attention_mask=attention_mask,
+            **kwargs,
+        )
+
+        return Seq2SeqMoEModelOutput(
+            last_hidden_state=decoder_outputs.last_hidden_state,
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.hidden_states,
+            decoder_attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+            decoder_router_logits=decoder_outputs.router_logits,
+            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+            encoder_hidden_states=encoder_outputs.hidden_states,
+            encoder_attentions=encoder_outputs.attentions,
+            encoder_router_logits=encoder_outputs.router_logits,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    SWITCH_TRANSFORMERS Model with a `language modeling` head on top.
+    """
+)
+class SwitchTransformersForConditionalGeneration(SwitchTransformersPreTrainedModel, GenerationMixin):
+    _tied_weights_keys = {
+        "encoder.embed_tokens.weight": "shared.weight",
+        "decoder.embed_tokens.weight": "shared.weight",
+        "lm_head.weight": "shared.weight",
+    }
+
+    def __init__(self, config: SwitchTransformersConfig):
+        super().__init__(config)
+        self.model_dim = config.d_model
+
+        self.shared = nn.Embedding(config.vocab_size, config.d_model)
+
+        encoder_config = copy.deepcopy(config)
+        encoder_config.is_decoder = False
+        encoder_config.use_cache = False
+        self.encoder = SwitchTransformersStack(encoder_config)
+
+        decoder_config = copy.deepcopy(config)
+        decoder_config.is_decoder = True
+        decoder_config.num_layers = config.num_decoder_layers
+        self.decoder = SwitchTransformersStack(decoder_config)
+
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+
+        self.router_z_loss_coef = config.router_z_loss_coef
+        self.router_aux_loss_coef = config.router_aux_loss_coef
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.shared
+
+    def set_input_embeddings(self, new_embeddings):
+        self.shared = new_embeddings
+        self.encoder.set_input_embeddings(new_embeddings)
+        self.decoder.set_input_embeddings(new_embeddings)
+
+    @auto_docstring
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.FloatTensor | None = None,
+        decoder_input_ids: torch.LongTensor | None = None,
+        decoder_attention_mask: torch.BoolTensor | None = None,
+        encoder_outputs: tuple[tuple[torch.Tensor]] | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        decoder_inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        output_router_logits: bool | None = False,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.FloatTensor] | Seq2SeqMoEOutput:
+        if encoder_outputs is None:
+            encoder_outputs = self.encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                output_router_logits=output_router_logits,
+                **kwargs,
+            )
+
+        hidden_states = encoder_outputs[0]
+
+        if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None:
+            # get decoder inputs from shifting lm labels to the right
+            decoder_input_ids = self._shift_right(labels)
+
+        # Decode
+        decoder_outputs = self.decoder(
+            input_ids=decoder_input_ids,
+            attention_mask=decoder_attention_mask,
+            inputs_embeds=decoder_inputs_embeds,
+            past_key_values=past_key_values,
+            encoder_hidden_states=hidden_states,
+            encoder_attention_mask=attention_mask,
+            output_router_logits=output_router_logits,
+            **kwargs,
+        )
+
+        sequence_output = decoder_outputs.last_hidden_state
+
+        if self.config.tie_word_embeddings:
+            # Rescale output before projecting on vocab
+            # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/transformer.py#L586
+            sequence_output = sequence_output * (self.model_dim**-0.5)
+
+        lm_logits = self.lm_head(sequence_output)
+
+        loss = None
+        encoder_z_loss = None
+        encoder_aux_loss = None
+        decoder_z_loss = None
+        decoder_aux_loss = None
+
+        if output_router_logits:
+            # Compute the router loss (z_loss + auxiliary loss) for each router in the encoder and decoder
+            if self.encoder.config.encoder_sparse_step > 1:
+                encoder_router_logits, encoder_expert_indexes = self._unpack_router_logits(encoder_outputs[-1])
+                encoder_z_loss = router_z_loss_func(encoder_router_logits)
+                encoder_router_probs = nn.Softmax(dim=-1)(encoder_router_logits)
+                encoder_aux_loss = load_balancing_loss_func(encoder_router_probs, encoder_expert_indexes)
+            else:
+                encoder_z_loss = 0
+                encoder_aux_loss = 0
+
+            if self.decoder.config.decoder_sparse_step > 1:
+                decoder_router_logits, decoder_expert_indexes = self._unpack_router_logits(decoder_outputs[-1])
+                decoder_z_loss = router_z_loss_func(decoder_router_logits)
+                decoder_router_probs = nn.Softmax(dim=-1)(decoder_router_logits)
+                decoder_aux_loss = load_balancing_loss_func(decoder_router_probs, decoder_expert_indexes)
+            else:
+                decoder_z_loss = 0
+                decoder_aux_loss = 0
+
+        if labels is not None:
+            loss_fct = CrossEntropyLoss(ignore_index=-100)
+            # move labels to correct device to enable PP
+            labels = labels.to(lm_logits.device)
+            loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))
+
+            if output_router_logits:
+                z_loss = self.router_z_loss_coef * (encoder_z_loss + decoder_z_loss)
+                aux_loss = self.router_aux_loss_coef * (encoder_aux_loss + decoder_aux_loss)
+                loss = loss + z_loss + aux_loss
+
+        return Seq2SeqMoEOutput(
+            loss=loss,
+            logits=lm_logits,
+            encoder_z_loss=encoder_z_loss,
+            encoder_aux_loss=encoder_aux_loss,
+            decoder_z_loss=decoder_z_loss,
+            decoder_aux_loss=decoder_aux_loss,
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.hidden_states,
+            decoder_attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+            decoder_router_logits=decoder_outputs.router_logits,
+            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+            encoder_hidden_states=encoder_outputs.hidden_states,
+            encoder_attentions=encoder_outputs.attentions,
+            encoder_router_logits=encoder_outputs.router_logits,
+        )
+
+    def _unpack_router_logits(self, router_outputs):
+        total_router_logits = []
+        total_expert_indexes = []
+        for router_output in router_outputs:
+            if len(router_output[0].shape) > 1:
+                router_logits, expert_indexes = router_output
+                total_router_logits.append(router_logits)
+                total_expert_indexes.append(expert_indexes)
+        return torch.cat(total_router_logits, dim=1), torch.cat(total_expert_indexes, dim=1)
+
+    def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
+        return self._shift_right(labels)
+
+
+class SwitchTransformersEncoderModel(SwitchTransformersPreTrainedModel):
+    _tied_weights_keys = {
+        "encoder.embed_tokens.weight": "shared.weight",
+    }
+
+    def __init__(self, config: SwitchTransformersConfig):
+        super().__init__(config)
+        self.shared = nn.Embedding(config.vocab_size, config.d_model)
+
+        encoder_config = copy.deepcopy(config)
+        encoder_config.use_cache = False
+        encoder_config.is_encoder_decoder = False
+        self.encoder = SwitchTransformersStack(encoder_config)
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.shared
+
+    def set_input_embeddings(self, new_embeddings):
+        self.shared = new_embeddings
+        self.encoder.set_input_embeddings(new_embeddings)
+
+    @auto_docstring
+    @can_return_tuple
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.FloatTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.FloatTensor] | MoEModelOutput:
+        use_cache = False
+        encoder_outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            **kwargs,
+        )
+
+        return encoder_outputs
+
+
+__all__ = [
+    "SwitchTransformersEncoderModel",
+    "SwitchTransformersForConditionalGeneration",
+    "SwitchTransformersModel",
+    "SwitchTransformersPreTrainedModel",
+    "SwitchTransformersTop1Router",
+    "SwitchTransformersSparseMLP",
+]