Add files using upload-large-folder tool

datasets/mmlu/.gitattributes

@@ -0,0 +1,55 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.lz4 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+# Audio files - uncompressed
+*.pcm filter=lfs diff=lfs merge=lfs -text
+*.sam filter=lfs diff=lfs merge=lfs -text
+*.raw filter=lfs diff=lfs merge=lfs -text
+# Audio files - compressed
+*.aac filter=lfs diff=lfs merge=lfs -text
+*.flac filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text
+*.ogg filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
+# Image files - uncompressed
+*.bmp filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.tiff filter=lfs diff=lfs merge=lfs -text
+# Image files - compressed
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.webp filter=lfs diff=lfs merge=lfs -text

datasets/mmlu/README.md

@@ -0,0 +1,2299 @@
+---
+annotations_creators:
+- no-annotation
+language_creators:
+- expert-generated
+language:
+- en
+license:
+- mit
+multilinguality:
+- monolingual
+size_categories:
+- 10K<n<100K
+source_datasets:
+- original
+task_categories:
+- question-answering
+task_ids:
+- multiple-choice-qa
+paperswithcode_id: mmlu
+pretty_name: Measuring Massive Multitask Language Understanding
+language_bcp47:
+- en-US
+dataset_info:
+- config_name: abstract_algebra
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 49618.6654322746
+    num_examples: 100
+  - name: validation
+    num_bytes: 5485.515349444808
+    num_examples: 11
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 17143
+  dataset_size: 57303.3562203159
+- config_name: all
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 6967453
+    num_examples: 14042
+  - name: validation
+    num_bytes: 763484
+    num_examples: 1531
+  - name: dev
+    num_bytes: 125353
+    num_examples: 285
+  - name: auxiliary_train
+    num_bytes: 161000625
+    num_examples: 99842
+  download_size: 51503402
+  dataset_size: 168856915
+- config_name: anatomy
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 66985.19833357072
+    num_examples: 135
+  - name: validation
+    num_bytes: 6981.5649902024825
+    num_examples: 14
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 28864
+  dataset_size: 76165.9387623697
+- config_name: astronomy
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 75420.3714570574
+    num_examples: 152
+  - name: validation
+    num_bytes: 7978.931417374265
+    num_examples: 16
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 39316
+  dataset_size: 85598.47831302814
+- config_name: auxiliary_train
+  features:
+  - name: train
+    struct:
+    - name: answer
+      dtype: int64
+    - name: choices
+      sequence: string
+    - name: question
+      dtype: string
+    - name: subject
+      dtype: string
+  splits:
+  - name: train
+    num_bytes: 161000625
+    num_examples: 99842
+  download_size: 47518592
+  dataset_size: 161000625
+- config_name: business_ethics
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 49618.6654322746
+    num_examples: 100
+  - name: validation
+    num_bytes: 5485.515349444808
+    num_examples: 11
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 31619
+  dataset_size: 57303.3562203159
+- config_name: clinical_knowledge
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 131489.4633955277
+    num_examples: 265
+  - name: validation
+    num_bytes: 14461.813193990856
+    num_examples: 29
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 51655
+  dataset_size: 148150.45202811505
+- config_name: college_biology
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 71450.87822247542
+    num_examples: 144
+  - name: validation
+    num_bytes: 7978.931417374265
+    num_examples: 16
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 43017
+  dataset_size: 81628.98507844617
+- config_name: college_chemistry
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 49618.6654322746
+    num_examples: 100
+  - name: validation
+    num_bytes: 3989.4657086871325
+    num_examples: 8
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 26781
+  dataset_size: 55807.30657955822
+- config_name: college_computer_science
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 49618.6654322746
+    num_examples: 100
+  - name: validation
+    num_bytes: 5485.515349444808
+    num_examples: 11
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 41132
+  dataset_size: 57303.3562203159
+- config_name: college_mathematics
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 49618.6654322746
+    num_examples: 100
+  - name: validation
+    num_bytes: 5485.515349444808
+    num_examples: 11
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 26779
+  dataset_size: 57303.3562203159
+- config_name: college_medicine
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 85840.29119783506
+    num_examples: 173
+  - name: validation
+    num_bytes: 10971.030698889615
+    num_examples: 22
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 56303
+  dataset_size: 99010.49733532117
+- config_name: college_physics
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 50611.0387409201
+    num_examples: 102
+  - name: validation
+    num_bytes: 5485.515349444808
+    num_examples: 11
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 29539
+  dataset_size: 58295.7295289614
+- config_name: computer_security
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 49618.6654322746
+    num_examples: 100
+  - name: validation
+    num_bytes: 5485.515349444808
+    num_examples: 11
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 30150
+  dataset_size: 57303.3562203159
+- config_name: conceptual_physics
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 116603.86376584532
+    num_examples: 235
+  - name: validation
+    num_bytes: 12965.76355323318
+    num_examples: 26
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 34968
+  dataset_size: 131768.802757675
+- config_name: econometrics
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 56565.27859279305
+    num_examples: 114
+  - name: validation
+    num_bytes: 5984.198563030699
+    num_examples: 12
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 36040
+  dataset_size: 64748.652594420244
+- config_name: electrical_engineering
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 71947.06487679818
+    num_examples: 145
+  - name: validation
+    num_bytes: 7978.931417374265
+    num_examples: 16
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 26746
+  dataset_size: 82125.17173276893
+- config_name: elementary_mathematics
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 187558.555333998
+    num_examples: 378
+  - name: validation
+    num_bytes: 20446.011757021555
+    num_examples: 41
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 54987
+  dataset_size: 210203.74252961605
+- config_name: formal_logic
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 62519.518444666
+    num_examples: 126
+  - name: validation
+    num_bytes: 6981.5649902024825
+    num_examples: 14
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 32884
+  dataset_size: 71700.25887346498
+- config_name: global_facts
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 49618.6654322746
+    num_examples: 100
+  - name: validation
+    num_bytes: 4986.8321358589155
+    num_examples: 10
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 19258
+  dataset_size: 56804.67300673001
+- config_name: high_school_biology
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 153817.86284005127
+    num_examples: 310
+  - name: validation
+    num_bytes: 15957.86283474853
+    num_examples: 32
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 78216
+  dataset_size: 171974.90111339628
+- config_name: high_school_chemistry
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 100725.89082751745
+    num_examples: 203
+  - name: validation
+    num_bytes: 10971.030698889615
+    num_examples: 22
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 45799
+  dataset_size: 113896.09696500355
+- config_name: high_school_computer_science
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 49618.6654322746
+    num_examples: 100
+  - name: validation
+    num_bytes: 4488.148922273024
+    num_examples: 9
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 39072
+  dataset_size: 56305.989793144116
+- config_name: high_school_european_history
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 81870.79796325309
+    num_examples: 165
+  - name: validation
+    num_bytes: 8976.297844546049
+    num_examples: 18
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 196270
+  dataset_size: 93046.27124639563
+- config_name: high_school_geography
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 98244.95755590372
+    num_examples: 198
+  - name: validation
+    num_bytes: 10971.030698889615
+    num_examples: 22
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 38255
+  dataset_size: 111415.16369338983
+- config_name: high_school_government_and_politics
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 95764.02428428999
+    num_examples: 193
+  - name: validation
+    num_bytes: 10472.347485303722
+    num_examples: 21
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 52963
+  dataset_size: 108435.5472081902
+- config_name: high_school_macroeconomics
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 193512.79518587096
+    num_examples: 390
+  - name: validation
+    num_bytes: 21443.378184193338
+    num_examples: 43
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 68758
+  dataset_size: 217155.34880866078
+- config_name: high_school_mathematics
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 133970.39666714144
+    num_examples: 270
+  - name: validation
+    num_bytes: 14461.813193990856
+    num_examples: 29
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 45210
+  dataset_size: 150631.38529972878
+- config_name: high_school_microeconomics
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 118092.42372881356
+    num_examples: 238
+  - name: validation
+    num_bytes: 12965.76355323318
+    num_examples: 26
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 49885
+  dataset_size: 133257.36272064323
+- config_name: high_school_physics
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 74924.18480273466
+    num_examples: 151
+  - name: validation
+    num_bytes: 8477.614630960157
+    num_examples: 17
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 45483
+  dataset_size: 85600.9748722913
+- config_name: high_school_psychology
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 270421.7266058966
+    num_examples: 545
+  - name: validation
+    num_bytes: 29920.992815153495
+    num_examples: 60
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 113158
+  dataset_size: 302541.8948596466
+- config_name: high_school_statistics
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 107176.31733371314
+    num_examples: 216
+  - name: validation
+    num_bytes: 11469.713912475507
+    num_examples: 23
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 74924
+  dataset_size: 120845.20668478514
+- config_name: high_school_us_history
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 101222.0774818402
+    num_examples: 204
+  - name: validation
+    num_bytes: 10971.030698889615
+    num_examples: 22
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 200043
+  dataset_size: 114392.2836193263
+- config_name: high_school_world_history
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 117596.23707449081
+    num_examples: 237
+  - name: validation
+    num_bytes: 12965.76355323318
+    num_examples: 26
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 250302
+  dataset_size: 132761.17606632048
+- config_name: human_aging
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 110649.62391397236
+    num_examples: 223
+  - name: validation
+    num_bytes: 11469.713912475507
+    num_examples: 23
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 41196
+  dataset_size: 124318.51326504436
+- config_name: human_sexuality
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 65000.451716279735
+    num_examples: 131
+  - name: validation
+    num_bytes: 5984.198563030699
+    num_examples: 12
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 32533
+  dataset_size: 73183.82571790692
+- config_name: international_law
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 60038.58517305227
+    num_examples: 121
+  - name: validation
+    num_bytes: 6482.88177661659
+    num_examples: 13
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 41592
+  dataset_size: 68720.64238826535
+- config_name: jurisprudence
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 53588.15866685657
+    num_examples: 108
+  - name: validation
+    num_bytes: 5485.515349444808
+    num_examples: 11
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 33578
+  dataset_size: 61272.84945489787
+- config_name: logical_fallacies
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 80878.4246546076
+    num_examples: 163
+  - name: validation
+    num_bytes: 8976.297844546049
+    num_examples: 18
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 33669
+  dataset_size: 92053.89793775014
+- config_name: machine_learning
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 55572.90528414756
+    num_examples: 112
+  - name: validation
+    num_bytes: 5485.515349444808
+    num_examples: 11
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 31121
+  dataset_size: 63257.596072188855
+- config_name: management
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 51107.225395242844
+    num_examples: 103
+  - name: validation
+    num_bytes: 5485.515349444808
+    num_examples: 11
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 22828
+  dataset_size: 58791.91618328414
+- config_name: marketing
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 116107.67711152257
+    num_examples: 234
+  - name: validation
+    num_bytes: 12467.08033964729
+    num_examples: 25
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 49747
+  dataset_size: 130773.93288976635
+- config_name: medical_genetics
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 49618.6654322746
+    num_examples: 100
+  - name: validation
+    num_bytes: 5485.515349444808
+    num_examples: 11
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 25775
+  dataset_size: 57303.3562203159
+- config_name: miscellaneous
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 388514.15033471014
+    num_examples: 783
+  - name: validation
+    num_bytes: 42886.756368386676
+    num_examples: 86
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 115097
+  dataset_size: 433600.08214169333
+- config_name: moral_disputes
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 171680.58239567012
+    num_examples: 346
+  - name: validation
+    num_bytes: 18949.96211626388
+    num_examples: 38
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 76043
+  dataset_size: 192829.71995053047
+- config_name: moral_scenarios
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 444087.05561885773
+    num_examples: 895
+  - name: validation
+    num_bytes: 49868.32135858916
+    num_examples: 100
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 109869
+  dataset_size: 496154.5524160434
+- config_name: nutrition
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 151833.1162227603
+    num_examples: 306
+  - name: validation
+    num_bytes: 16456.54604833442
+    num_examples: 33
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 69050
+  dataset_size: 170488.8377096912
+- config_name: philosophy
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 154314.04949437402
+    num_examples: 311
+  - name: validation
+    num_bytes: 16955.229261920314
+    num_examples: 34
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 61912
+  dataset_size: 173468.45419489083
+- config_name: prehistory
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 160764.47600056973
+    num_examples: 324
+  - name: validation
+    num_bytes: 17453.912475506204
+    num_examples: 35
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 68826
+  dataset_size: 180417.5639146724
+- config_name: professional_accounting
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 139924.6365190144
+    num_examples: 282
+  - name: validation
+    num_bytes: 15459.179621162639
+    num_examples: 31
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 87297
+  dataset_size: 157582.99157877354
+- config_name: professional_law
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 761150.3277310925
+    num_examples: 1534
+  - name: validation
+    num_bytes: 84776.14630960157
+    num_examples: 170
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 1167828
+  dataset_size: 848125.6494792906
+- config_name: professional_medicine
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 134962.7699757869
+    num_examples: 272
+  - name: validation
+    num_bytes: 15459.179621162639
+    num_examples: 31
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 153242
+  dataset_size: 152621.12503554605
+- config_name: professional_psychology
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 303666.2324455206
+    num_examples: 612
+  - name: validation
+    num_bytes: 34409.14173742652
+    num_examples: 69
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 159357
+  dataset_size: 340274.5496215436
+- config_name: public_relations
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 54580.53197550207
+    num_examples: 110
+  - name: validation
+    num_bytes: 5984.198563030699
+    num_examples: 12
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 31500
+  dataset_size: 62763.90597712925
+- config_name: security_studies
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 121565.73030907278
+    num_examples: 245
+  - name: validation
+    num_bytes: 13464.446766819072
+    num_examples: 27
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 140258
+  dataset_size: 137229.35251448833
+- config_name: sociology
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 99733.51751887196
+    num_examples: 201
+  - name: validation
+    num_bytes: 10971.030698889615
+    num_examples: 22
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 56480
+  dataset_size: 112903.72365635807
+- config_name: us_foreign_policy
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 49618.6654322746
+    num_examples: 100
+  - name: validation
+    num_bytes: 5485.515349444808
+    num_examples: 11
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 29027
+  dataset_size: 57303.3562203159
+- config_name: virology
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 82366.98461757584
+    num_examples: 166
+  - name: validation
+    num_bytes: 8976.297844546049
+    num_examples: 18
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 38229
+  dataset_size: 93542.45790071838
+- config_name: world_religions
+  features:
+  - name: question
+    dtype: string
+  - name: subject
+    dtype: string
+  - name: choices
+    sequence: string
+  - name: answer
+    dtype:
+      class_label:
+        names:
+          '0': A
+          '1': B
+          '2': C
+          '3': D
+  splits:
+  - name: test
+    num_bytes: 84847.91788918957
+    num_examples: 171
+  - name: validation
+    num_bytes: 9474.98105813194
+    num_examples: 19
+  - name: dev
+    num_bytes: 2199.1754385964914
+    num_examples: 5
+  download_size: 27165
+  dataset_size: 96522.07438591801
+configs:
+- config_name: abstract_algebra
+  data_files:
+  - split: test
+    path: abstract_algebra/test-*
+  - split: validation
+    path: abstract_algebra/validation-*
+  - split: dev
+    path: abstract_algebra/dev-*
+- config_name: all
+  data_files:
+  - split: test
+    path: all/test-*
+  - split: validation
+    path: all/validation-*
+  - split: dev
+    path: all/dev-*
+  - split: auxiliary_train
+    path: all/auxiliary_train-*
+- config_name: anatomy
+  data_files:
+  - split: test
+    path: anatomy/test-*
+  - split: validation
+    path: anatomy/validation-*
+  - split: dev
+    path: anatomy/dev-*
+- config_name: astronomy
+  data_files:
+  - split: test
+    path: astronomy/test-*
+  - split: validation
+    path: astronomy/validation-*
+  - split: dev
+    path: astronomy/dev-*
+- config_name: auxiliary_train
+  data_files:
+  - split: train
+    path: auxiliary_train/train-*
+- config_name: business_ethics
+  data_files:
+  - split: test
+    path: business_ethics/test-*
+  - split: validation
+    path: business_ethics/validation-*
+  - split: dev
+    path: business_ethics/dev-*
+- config_name: clinical_knowledge
+  data_files:
+  - split: test
+    path: clinical_knowledge/test-*
+  - split: validation
+    path: clinical_knowledge/validation-*
+  - split: dev
+    path: clinical_knowledge/dev-*
+- config_name: college_biology
+  data_files:
+  - split: test
+    path: college_biology/test-*
+  - split: validation
+    path: college_biology/validation-*
+  - split: dev
+    path: college_biology/dev-*
+- config_name: college_chemistry
+  data_files:
+  - split: test
+    path: college_chemistry/test-*
+  - split: validation
+    path: college_chemistry/validation-*
+  - split: dev
+    path: college_chemistry/dev-*
+- config_name: college_computer_science
+  data_files:
+  - split: test
+    path: college_computer_science/test-*
+  - split: validation
+    path: college_computer_science/validation-*
+  - split: dev
+    path: college_computer_science/dev-*
+- config_name: college_mathematics
+  data_files:
+  - split: test
+    path: college_mathematics/test-*
+  - split: validation
+    path: college_mathematics/validation-*
+  - split: dev
+    path: college_mathematics/dev-*
+- config_name: college_medicine
+  data_files:
+  - split: test
+    path: college_medicine/test-*
+  - split: validation
+    path: college_medicine/validation-*
+  - split: dev
+    path: college_medicine/dev-*
+- config_name: college_physics
+  data_files:
+  - split: test
+    path: college_physics/test-*
+  - split: validation
+    path: college_physics/validation-*
+  - split: dev
+    path: college_physics/dev-*
+- config_name: computer_security
+  data_files:
+  - split: test
+    path: computer_security/test-*
+  - split: validation
+    path: computer_security/validation-*
+  - split: dev
+    path: computer_security/dev-*
+- config_name: conceptual_physics
+  data_files:
+  - split: test
+    path: conceptual_physics/test-*
+  - split: validation
+    path: conceptual_physics/validation-*
+  - split: dev
+    path: conceptual_physics/dev-*
+- config_name: econometrics
+  data_files:
+  - split: test
+    path: econometrics/test-*
+  - split: validation
+    path: econometrics/validation-*
+  - split: dev
+    path: econometrics/dev-*
+- config_name: electrical_engineering
+  data_files:
+  - split: test
+    path: electrical_engineering/test-*
+  - split: validation
+    path: electrical_engineering/validation-*
+  - split: dev
+    path: electrical_engineering/dev-*
+- config_name: elementary_mathematics
+  data_files:
+  - split: test
+    path: elementary_mathematics/test-*
+  - split: validation
+    path: elementary_mathematics/validation-*
+  - split: dev
+    path: elementary_mathematics/dev-*
+- config_name: formal_logic
+  data_files:
+  - split: test
+    path: formal_logic/test-*
+  - split: validation
+    path: formal_logic/validation-*
+  - split: dev
+    path: formal_logic/dev-*
+- config_name: global_facts
+  data_files:
+  - split: test
+    path: global_facts/test-*
+  - split: validation
+    path: global_facts/validation-*
+  - split: dev
+    path: global_facts/dev-*
+- config_name: high_school_biology
+  data_files:
+  - split: test
+    path: high_school_biology/test-*
+  - split: validation
+    path: high_school_biology/validation-*
+  - split: dev
+    path: high_school_biology/dev-*
+- config_name: high_school_chemistry
+  data_files:
+  - split: test
+    path: high_school_chemistry/test-*
+  - split: validation
+    path: high_school_chemistry/validation-*
+  - split: dev
+    path: high_school_chemistry/dev-*
+- config_name: high_school_computer_science
+  data_files:
+  - split: test
+    path: high_school_computer_science/test-*
+  - split: validation
+    path: high_school_computer_science/validation-*
+  - split: dev
+    path: high_school_computer_science/dev-*
+- config_name: high_school_european_history
+  data_files:
+  - split: test
+    path: high_school_european_history/test-*
+  - split: validation
+    path: high_school_european_history/validation-*
+  - split: dev
+    path: high_school_european_history/dev-*
+- config_name: high_school_geography
+  data_files:
+  - split: test
+    path: high_school_geography/test-*
+  - split: validation
+    path: high_school_geography/validation-*
+  - split: dev
+    path: high_school_geography/dev-*
+- config_name: high_school_government_and_politics
+  data_files:
+  - split: test
+    path: high_school_government_and_politics/test-*
+  - split: validation
+    path: high_school_government_and_politics/validation-*
+  - split: dev
+    path: high_school_government_and_politics/dev-*
+- config_name: high_school_macroeconomics
+  data_files:
+  - split: test
+    path: high_school_macroeconomics/test-*
+  - split: validation
+    path: high_school_macroeconomics/validation-*
+  - split: dev
+    path: high_school_macroeconomics/dev-*
+- config_name: high_school_mathematics
+  data_files:
+  - split: test
+    path: high_school_mathematics/test-*
+  - split: validation
+    path: high_school_mathematics/validation-*
+  - split: dev
+    path: high_school_mathematics/dev-*
+- config_name: high_school_microeconomics
+  data_files:
+  - split: test
+    path: high_school_microeconomics/test-*
+  - split: validation
+    path: high_school_microeconomics/validation-*
+  - split: dev
+    path: high_school_microeconomics/dev-*
+- config_name: high_school_physics
+  data_files:
+  - split: test
+    path: high_school_physics/test-*
+  - split: validation
+    path: high_school_physics/validation-*
+  - split: dev
+    path: high_school_physics/dev-*
+- config_name: high_school_psychology
+  data_files:
+  - split: test
+    path: high_school_psychology/test-*
+  - split: validation
+    path: high_school_psychology/validation-*
+  - split: dev
+    path: high_school_psychology/dev-*
+- config_name: high_school_statistics
+  data_files:
+  - split: test
+    path: high_school_statistics/test-*
+  - split: validation
+    path: high_school_statistics/validation-*
+  - split: dev
+    path: high_school_statistics/dev-*
+- config_name: high_school_us_history
+  data_files:
+  - split: test
+    path: high_school_us_history/test-*
+  - split: validation
+    path: high_school_us_history/validation-*
+  - split: dev
+    path: high_school_us_history/dev-*
+- config_name: high_school_world_history
+  data_files:
+  - split: test
+    path: high_school_world_history/test-*
+  - split: validation
+    path: high_school_world_history/validation-*
+  - split: dev
+    path: high_school_world_history/dev-*
+- config_name: human_aging
+  data_files:
+  - split: test
+    path: human_aging/test-*
+  - split: validation
+    path: human_aging/validation-*
+  - split: dev
+    path: human_aging/dev-*
+- config_name: human_sexuality
+  data_files:
+  - split: test
+    path: human_sexuality/test-*
+  - split: validation
+    path: human_sexuality/validation-*
+  - split: dev
+    path: human_sexuality/dev-*
+- config_name: international_law
+  data_files:
+  - split: test
+    path: international_law/test-*
+  - split: validation
+    path: international_law/validation-*
+  - split: dev
+    path: international_law/dev-*
+- config_name: jurisprudence
+  data_files:
+  - split: test
+    path: jurisprudence/test-*
+  - split: validation
+    path: jurisprudence/validation-*
+  - split: dev
+    path: jurisprudence/dev-*
+- config_name: logical_fallacies
+  data_files:
+  - split: test
+    path: logical_fallacies/test-*
+  - split: validation
+    path: logical_fallacies/validation-*
+  - split: dev
+    path: logical_fallacies/dev-*
+- config_name: machine_learning
+  data_files:
+  - split: test
+    path: machine_learning/test-*
+  - split: validation
+    path: machine_learning/validation-*
+  - split: dev
+    path: machine_learning/dev-*
+- config_name: management
+  data_files:
+  - split: test
+    path: management/test-*
+  - split: validation
+    path: management/validation-*
+  - split: dev
+    path: management/dev-*
+- config_name: marketing
+  data_files:
+  - split: test
+    path: marketing/test-*
+  - split: validation
+    path: marketing/validation-*
+  - split: dev
+    path: marketing/dev-*
+- config_name: medical_genetics
+  data_files:
+  - split: test
+    path: medical_genetics/test-*
+  - split: validation
+    path: medical_genetics/validation-*
+  - split: dev
+    path: medical_genetics/dev-*
+- config_name: miscellaneous
+  data_files:
+  - split: test
+    path: miscellaneous/test-*
+  - split: validation
+    path: miscellaneous/validation-*
+  - split: dev
+    path: miscellaneous/dev-*
+- config_name: moral_disputes
+  data_files:
+  - split: test
+    path: moral_disputes/test-*
+  - split: validation
+    path: moral_disputes/validation-*
+  - split: dev
+    path: moral_disputes/dev-*
+- config_name: moral_scenarios
+  data_files:
+  - split: test
+    path: moral_scenarios/test-*
+  - split: validation
+    path: moral_scenarios/validation-*
+  - split: dev
+    path: moral_scenarios/dev-*
+- config_name: nutrition
+  data_files:
+  - split: test
+    path: nutrition/test-*
+  - split: validation
+    path: nutrition/validation-*
+  - split: dev
+    path: nutrition/dev-*
+- config_name: philosophy
+  data_files:
+  - split: test
+    path: philosophy/test-*
+  - split: validation
+    path: philosophy/validation-*
+  - split: dev
+    path: philosophy/dev-*
+- config_name: prehistory
+  data_files:
+  - split: test
+    path: prehistory/test-*
+  - split: validation
+    path: prehistory/validation-*
+  - split: dev
+    path: prehistory/dev-*
+- config_name: professional_accounting
+  data_files:
+  - split: test
+    path: professional_accounting/test-*
+  - split: validation
+    path: professional_accounting/validation-*
+  - split: dev
+    path: professional_accounting/dev-*
+- config_name: professional_law
+  data_files:
+  - split: test
+    path: professional_law/test-*
+  - split: validation
+    path: professional_law/validation-*
+  - split: dev
+    path: professional_law/dev-*
+- config_name: professional_medicine
+  data_files:
+  - split: test
+    path: professional_medicine/test-*
+  - split: validation
+    path: professional_medicine/validation-*
+  - split: dev
+    path: professional_medicine/dev-*
+- config_name: professional_psychology
+  data_files:
+  - split: test
+    path: professional_psychology/test-*
+  - split: validation
+    path: professional_psychology/validation-*
+  - split: dev
+    path: professional_psychology/dev-*
+- config_name: public_relations
+  data_files:
+  - split: test
+    path: public_relations/test-*
+  - split: validation
+    path: public_relations/validation-*
+  - split: dev
+    path: public_relations/dev-*
+- config_name: security_studies
+  data_files:
+  - split: test
+    path: security_studies/test-*
+  - split: validation
+    path: security_studies/validation-*
+  - split: dev
+    path: security_studies/dev-*
+- config_name: sociology
+  data_files:
+  - split: test
+    path: sociology/test-*
+  - split: validation
+    path: sociology/validation-*
+  - split: dev
+    path: sociology/dev-*
+- config_name: us_foreign_policy
+  data_files:
+  - split: test
+    path: us_foreign_policy/test-*
+  - split: validation
+    path: us_foreign_policy/validation-*
+  - split: dev
+    path: us_foreign_policy/dev-*
+- config_name: virology
+  data_files:
+  - split: test
+    path: virology/test-*
+  - split: validation
+    path: virology/validation-*
+  - split: dev
+    path: virology/dev-*
+- config_name: world_religions
+  data_files:
+  - split: test
+    path: world_religions/test-*
+  - split: validation
+    path: world_religions/validation-*
+  - split: dev
+    path: world_religions/dev-*
+---
+
+# Dataset Card for MMLU
+
+## Table of Contents
+- [Table of Contents](#table-of-contents)
+- [Dataset Description](#dataset-description)
+  - [Dataset Summary](#dataset-summary)
+  - [Supported Tasks and Leaderboards](#supported-tasks-and-leaderboards)
+  - [Languages](#languages)
+- [Dataset Structure](#dataset-structure)
+  - [Data Instances](#data-instances)
+  - [Data Fields](#data-fields)
+  - [Data Splits](#data-splits)
+- [Dataset Creation](#dataset-creation)
+  - [Curation Rationale](#curation-rationale)
+  - [Source Data](#source-data)
+  - [Annotations](#annotations)
+  - [Personal and Sensitive Information](#personal-and-sensitive-information)
+- [Considerations for Using the Data](#considerations-for-using-the-data)
+  - [Social Impact of Dataset](#social-impact-of-dataset)
+  - [Discussion of Biases](#discussion-of-biases)
+  - [Other Known Limitations](#other-known-limitations)
+- [Additional Information](#additional-information)
+  - [Dataset Curators](#dataset-curators)
+  - [Licensing Information](#licensing-information)
+  - [Citation Information](#citation-information)
+  - [Contributions](#contributions)
+
+## Dataset Description
+
+- **Repository**: https://github.com/hendrycks/test
+- **Paper**: https://arxiv.org/abs/2009.03300
+
+### Dataset Summary
+
+[Measuring Massive Multitask Language Understanding](https://arxiv.org/pdf/2009.03300) by [Dan Hendrycks](https://people.eecs.berkeley.edu/~hendrycks/), [Collin Burns](http://collinpburns.com), [Steven Basart](https://stevenbas.art), Andy Zou, Mantas Mazeika, [Dawn Song](https://people.eecs.berkeley.edu/~dawnsong/), and [Jacob Steinhardt](https://www.stat.berkeley.edu/~jsteinhardt/) (ICLR 2021).
+
+This is a massive multitask test consisting of multiple-choice questions from various branches of knowledge. The test spans subjects in the humanities, social sciences, hard sciences, and other areas that are important for some people to learn. This covers 57 tasks including elementary mathematics, US history, computer science, law, and more. To attain high accuracy on this test, models must possess extensive world knowledge and problem solving ability.
+
+A complete list of tasks: ['abstract_algebra', 'anatomy', 'astronomy', 'business_ethics', 'clinical_knowledge', 'college_biology', 'college_chemistry', 'college_computer_science', 'college_mathematics', 'college_medicine', 'college_physics', 'computer_security', 'conceptual_physics', 'econometrics', 'electrical_engineering', 'elementary_mathematics', 'formal_logic', 'global_facts', 'high_school_biology', 'high_school_chemistry', 'high_school_computer_science', 'high_school_european_history', 'high_school_geography', 'high_school_government_and_politics', 'high_school_macroeconomics', 'high_school_mathematics', 'high_school_microeconomics', 'high_school_physics', 'high_school_psychology', 'high_school_statistics', 'high_school_us_history', 'high_school_world_history', 'human_aging', 'human_sexuality', 'international_law', 'jurisprudence', 'logical_fallacies', 'machine_learning', 'management', 'marketing', 'medical_genetics', 'miscellaneous', 'moral_disputes', 'moral_scenarios', 'nutrition', 'philosophy', 'prehistory', 'professional_accounting', 'professional_law', 'professional_medicine', 'professional_psychology', 'public_relations', 'security_studies', 'sociology', 'us_foreign_policy', 'virology', 'world_religions']
+
+### Supported Tasks and Leaderboards
+
+|                Model               | Authors |  Humanities |  Social Science  | STEM | Other | Average |
+|------------------------------------|----------|:-------:|:-------:|:-------:|:-------:|:-------:|
+| [UnifiedQA](https://arxiv.org/abs/2005.00700) | Khashabi et al., 2020 | 45.6 | 56.6 | 40.2 | 54.6 | 48.9
+| [GPT-3](https://arxiv.org/abs/2005.14165) (few-shot) | Brown et al., 2020 | 40.8 | 50.4 | 36.7 | 48.8 | 43.9
+| [GPT-2](https://arxiv.org/abs/2005.14165) | Radford et al., 2019 | 32.8 | 33.3 | 30.2 | 33.1 | 32.4
+| Random Baseline           | N/A | 25.0 | 25.0 | 25.0 | 25.0 | 25.0 | 25.0
+
+### Languages
+
+English
+
+## Dataset Structure
+
+### Data Instances
+
+An example from anatomy subtask looks as follows:
+```
+{
+  "question": "What is the embryological origin of the hyoid bone?",
+  "choices": ["The first pharyngeal arch", "The first and second pharyngeal arches", "The second pharyngeal arch", "The second and third pharyngeal arches"],
+  "answer": "D"
+}
+```
+
+### Data Fields
+
+- `question`: a string feature
+- `choices`: a list of 4 string features
+- `answer`: a ClassLabel feature
+
+### Data Splits
+
+- `auxiliary_train`: auxiliary multiple-choice training questions from ARC, MC_TEST, OBQA, RACE, etc.
+- `dev`: 5 examples per subtask, meant for few-shot setting
+- `test`: there are at least 100 examples per subtask
+
+|       | auxiliary_train   | dev | val | test |
+| ----- | :------: | :-----: | :-----: | :-----: |
+| TOTAL | 99842 | 285 | 1531 | 14042
+
+## Dataset Creation
+
+### Curation Rationale
+
+Transformer models have driven this recent progress by pretraining on massive text corpora, including all of Wikipedia, thousands of books, and numerous websites. These models consequently see extensive information about specialized topics, most of which is not assessed by existing NLP benchmarks. To bridge the gap between the wide-ranging knowledge that models see during pretraining and the existing measures of success, we introduce a new benchmark for assessing models across a diverse set of subjects that humans learn.
+
+### Source Data
+
+#### Initial Data Collection and Normalization
+
+[More Information Needed]
+
+#### Who are the source language producers?
+
+[More Information Needed]
+
+### Annotations
+
+#### Annotation process
+
+[More Information Needed]
+
+#### Who are the annotators?
+
+[More Information Needed]
+
+### Personal and Sensitive Information
+
+[More Information Needed]
+
+## Considerations for Using the Data
+
+### Social Impact of Dataset
+
+[More Information Needed]
+
+### Discussion of Biases
+
+[More Information Needed]
+
+### Other Known Limitations
+
+[More Information Needed]
+
+## Additional Information
+
+### Dataset Curators
+
+[More Information Needed]
+
+### Licensing Information
+
+[MIT License](https://github.com/hendrycks/test/blob/master/LICENSE)
+
+### Citation Information
+
+If you find this useful in your research, please consider citing the test and also the [ETHICS](https://arxiv.org/abs/2008.02275) dataset it draws from:
+```
+    @article{hendryckstest2021,
+      title={Measuring Massive Multitask Language Understanding},
+      author={Dan Hendrycks and Collin Burns and Steven Basart and Andy Zou and Mantas Mazeika and Dawn Song and Jacob Steinhardt},
+      journal={Proceedings of the International Conference on Learning Representations (ICLR)},
+      year={2021}
+    }
+
+    @article{hendrycks2021ethics,
+      title={Aligning AI With Shared Human Values},
+      author={Dan Hendrycks and Collin Burns and Steven Basart and Andrew Critch and Jerry Li and Dawn Song and Jacob Steinhardt},
+      journal={Proceedings of the International Conference on Learning Representations (ICLR)},
+      year={2021}
+    }
+```
+### Contributions
+
+Thanks to [@andyzoujm](https://github.com/andyzoujm) for adding this dataset.

models/dmax-16b/.gitattributes

@@ -0,0 +1,41 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/decode.png filter=lfs diff=lfs merge=lfs -text
+assets/demo.png filter=lfs diff=lfs merge=lfs -text
+assets/exp.png filter=lfs diff=lfs merge=lfs -text
+assets/tradeoff.png filter=lfs diff=lfs merge=lfs -text
+assets/train.png filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

models/dmax-16b/README.md

@@ -0,0 +1,109 @@
+---
+base_model:
+- inclusionAI/LLaDA2.0-mini
+license: apache-2.0
+library_name: transformers
+pipeline_tag: text-generation
+---
+
+<div align="center">
+<h1>🚀 DMax: Aggressive Parallel Decoding for dLLMs</h1>
+  <div align="center">
+  <a href="https://github.com/czg1225/DMax/blob/main/LICENSE">
+    <img alt="Apache" src="https://img.shields.io/badge/License-Apache-4E94CE.svg">
+  </a>
+  <a href="https://arxiv.org/abs/2604.08302">
+    <img src="https://img.shields.io/badge/Paper-Arxiv-darkred.svg" alt="Paper">
+  </a>
+  <a href="https://github.com/czg1225/DMax">
+    <img src="https://img.shields.io/badge/GitHub-Code-blue.svg?logo=github&" alt="GitHub">
+  </a>
+</div>
+</div>
+
+DMax is a new paradigm for efficient diffusion language models (dLLMs) that enables aggressive decoding parallelism while preserving generation quality. This repository hosts **DMax-16B**, a highly parallel general-purpose diffusion language model (dLLM) capable of handling code generation, mathematical reasoning, and daily conversation.
+
+
+## 💪 Highlights
+
+- **Aggressive Decoding Parallelism**: Achieves 6.0 TPF on math and reasoning tasks and 6.6 TPF on code tasks while preserving accuracy.
+- **Self-Revising dLLM**: Extends a pretrained MDLM into a UDLM with an intrinsic ability to revise its own erroneous predictions during decoding.
+- **Soft Parallel Decoding**: Uses interpolation between mask and token embeddings to propagate confidence priors from previous steps.
+
+<div align="center">
+  <img src="assets/tradeoff.png" width="100%" />
+  <br>
+  <em>Superior Parallelism-Accuracy Trade-off, Increased TPF with Maintained Accuracy.</em>
+</div>
+
+
+## 💻 Model and Datasets
+
+| Model | Description | Source Model | Link |
+| --- | --- | --- | --- |
+| 🤖 DMax-16B | Highly parallel general-purpose dLLM. | LLaDA-2.0-mini | [HF](https://huggingface.co/Zigeng/DMax-16B) |
+| 🤖 DMax-Math-16B | Highly parallel dLLM for math and reasoning. | LLaDA-2.0-mini | [HF](https://huggingface.co/Zigeng/DMax-Math-16B) |
+| 🤖 DMax-Coder-16B | Highly parallel dLLM for code generation. | LLaDA-2.0-mini | [HF](https://huggingface.co/Zigeng/DMax-Coder-16B) |
+
+| Dataset | Description | Link |
+| --- | --- | --- |
+| 📊 DMax-Math-Training-Data | math trajectories generated by LLaDA-2.0-mini | [HF](https://huggingface.co/datasets/Zigeng/DMax-LLaDA-2.0-Mini-Math-Trajectories) |
+| 📊 DMax-Code-Training-Data | code trajectories generated by LLaDA-2.0-mini | [HF](https://huggingface.co/datasets/Zigeng/DMax-LLaDA-2.0-Mini-Code-Trajectories) |
+
+
+
+## 🚀 Quick Start
+
+```python
+import torch
+from transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "Zigeng/DMax-16B", trust_remote_code=True, device_map="cuda:0"
+)
+model = model.to(torch.bfloat16)
+model.eval()
+tokenizer = AutoTokenizer.from_pretrained("Zigeng/DMax-16B", trust_remote_code=True)
+
+prompt = "A robe takes 2 bolts of blue fiber and half that much white fiber. How many bolts in total does it take?" + "
+Let's think step by step
+"
+
+input_ids = tokenizer.apply_chat_template(
+    [{"role": "user", "content": prompt}],
+    add_generation_prompt=True,
+    tokenize=True,
+    return_tensors="pt",
+)
+
+nfe, generated_tokens = model.generate_spd(
+    inputs=input_ids,
+    gen_length=2048,
+    block_length=32,
+    threshold=0.5,
+)
+
+generated_answer = tokenizer.decode(
+    generated_tokens[0],
+    skip_special_tokens=True,
+)
+
+print(generated_answer)
+print("nfe:",nfe,"token length",len(generated_tokens[0]))
+```
+
+## 📖 Experimental Results
+
+![trade-off](assets/exp.png)
+
+## 📚 Citation
+
+```bibtex
+@article{chen2026dmax,
+  title={DMax: Aggressive Parallel Decoding for dLLMs},
+  author={Chen, Zigeng and Fang, Gongfan and Ma, Xinyin and Yu, Ruonan and Wang, Xinchao},
+  journal={arXiv preprint arXiv:2604.08302},
+  year={2026}
+}
+

models/dmax-16b/chat_template.jinja

@@ -0,0 +1,86 @@
+{% set thinking_option = 'off' %}
+{{- '<role>SYSTEM</role>' }}
+{%- if messages[0].role == 'system' %}
+    {{- messages[0].content + '\n' }}
+{%- endif %}
+{%- if tools %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call>\n" }}
+{%- endif %}
+{{- 'detailed thinking ' + thinking_option + '<|role_end|>' }}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if message.content is string %}
+        {%- set content = message.content %}
+    {%- else %}
+        {%- set content = '' %}
+    {%- endif %}
+    {%- if message.role == "user" %}
+        {{- '<role>HUMAN</role>' + message.content + '<|role_end|>' }}
+    {%- elif message.role == "system" and not loop.first %}
+        {{- '<role>SYSTEM</role>' + message.content + '<|role_end|>' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if reasoning_content %}
+                {{- '<role>ASSISTANT</role>' + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<role>ASSISTANT</role>' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<role>ASSISTANT</role>' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|role_end|>' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<role>OBSERVATION</role>' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|role_end|>' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<role>ASSISTANT</role>' }}
+{%- endif %}

models/dmax-16b/config.json

@@ -0,0 +1,57 @@
+{
+  "_moe_implementation": "fused",
+  "architectures": [
+    "LLaDA2MoeModelLM"
+  ],
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_llada2_moe.LLaDA2MoeConfig",
+    "AutoModel": "modeling_llada2_moe.LLaDA2MoeModel",
+    "AutoModelForCausalLM": "modeling_llada2_moe.LLaDA2MoeModelLM"
+  },
+  "dtype": "bfloat16",
+  "embedding_dropout": 0.0,
+  "first_k_dense_replace": 1,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 2048,
+  "initializer_range": 0.02,
+  "intermediate_size": 5120,
+  "max_position_embeddings": 16384,
+  "max_window_layers": 28,
+  "model_type": "llada2_moe",
+  "moe_intermediate_size": 512,
+  "moe_router_enable_expert_bias": true,
+  "n_group": 8,
+  "norm_head": false,
+  "norm_softmax": false,
+  "norm_topk_prob": true,
+  "num_attention_heads": 16,
+  "num_experts": 256,
+  "num_experts_per_tok": 8,
+  "num_hidden_layers": 20,
+  "num_key_value_heads": 4,
+  "num_shared_experts": 1,
+  "output_dropout": 0.0,
+  "output_router_logits": false,
+  "pad_token_id": 156892,
+  "partial_rotary_factor": 0.5,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 600000,
+  "rotary_dim": 64,
+  "routed_scaling_factor": 2.5,
+  "router_dtype": "fp32",
+  "score_function": "sigmoid",
+  "sliding_window": 4096,
+  "tie_word_embeddings": false,
+  "topk_group": 4,
+  "transformers_version": "4.56.0",
+  "use_bias": false,
+  "use_cache": false,
+  "use_qkv_bias": false,
+  "use_rmsnorm": true,
+  "use_sliding_window": false,
+  "using_split_qkv_in_self_attention": false,
+  "vocab_size": 157184
+}

models/dmax-16b/configuration_llada2_moe.py

@@ -0,0 +1,89 @@
+"""LLaDA2 MoE model configuration"""
+
+from transformers.configuration_utils import PretrainedConfig
+
+
+class LLaDA2MoeConfig(PretrainedConfig):
+    model_type = "llada2_moe"
+
+    def __init__(
+        self,
+        vocab_size=30592,
+        hidden_size=1024,
+        intermediate_size=None,
+        num_hidden_layers=24,
+        num_attention_heads=16,
+        num_key_value_heads=0,
+        hidden_act="silu",
+        use_qkv_bias=False,  # llada2 only
+        use_qk_norm=True,
+        use_bias=True,  # llada2 only
+        rms_norm_eps=1e-05,
+        norm_head=False,  # llada2 only
+        tie_word_embeddings=False,  # PretrainedConfig key, here change default value.
+        embedding_dropout=0.1,
+        attention_dropout=0.1,
+        output_dropout=0.1,
+        initializer_range=0.02,
+        max_position_embeddings=16384,
+        rope_theta=10000.0,
+        use_cache=True,
+        use_sliding_window=False,
+        sliding_window=4096,
+        max_window_layers=28,
+        rope_scaling=None,
+        pad_token_id=126081,
+        num_experts=16,
+        num_shared_experts=0,
+        num_experts_per_tok=2,
+        n_group=8,
+        topk_group=4,
+        routed_scaling_factor=2.5,
+        moe_intermediate_size=None,
+        first_k_dense_replace=0,
+        head_dim=None,
+        output_router_logits=False,
+        partial_rotary_factor=0.5,
+        **kwargs,
+    ):
+        self.num_hidden_layers = num_hidden_layers
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_attention_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.use_qkv_bias = use_qkv_bias
+        self.use_qk_norm = use_qk_norm
+        self.use_bias = use_bias
+        self.norm_head = norm_head
+        self.rms_norm_eps = rms_norm_eps
+        self.embedding_dropout = embedding_dropout
+        self.attention_dropout = attention_dropout
+        self.output_dropout = output_dropout
+        self.initializer_range = initializer_range
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.use_cache = use_cache
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window
+        self.max_window_layers = max_window_layers
+        self.head_dim = head_dim or self.hidden_size // self.num_attention_heads
+        self.rope_scaling = rope_scaling
+
+        # MoE configs
+        self.num_experts = num_experts
+        self.num_shared_experts = num_shared_experts
+        self.num_experts_per_tok = num_experts_per_tok
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.moe_intermediate_size = moe_intermediate_size
+        self.first_k_dense_replace = first_k_dense_replace
+        self.output_router_logits = output_router_logits
+        self.routed_scaling_factor = routed_scaling_factor
+        self.partial_rotary_factor = partial_rotary_factor
+
+        super().__init__(
+            pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs
+        )
+

models/dmax-16b/modeling_llada2_moe.py

@@ -0,0 +1,1937 @@
+# Copyright 2025 Antgroup and The HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch LLaDA2MoE model."""
+
+import math
+from typing import List, Callable, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import (
+    _prepare_4d_causal_attention_mask,
+    _prepare_4d_causal_attention_mask_for_sdpa,
+)
+from transformers.modeling_outputs import (
+    MoeModelOutputWithPast,
+    MoeCausalLMOutputWithPast,
+)
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.pytorch_utils import (
+    ALL_LAYERNORM_LAYERS,
+    is_torch_greater_or_equal_than_1_13,
+)
+from transformers.utils import (
+    TransformersKwargs,
+    add_start_docstrings,
+    add_start_docstrings_to_model_forward,
+    logging,
+    replace_return_docstrings,
+)
+from transformers.utils.import_utils import is_torch_fx_available
+from .configuration_llada2_moe import LLaDA2MoeConfig
+from transformers.generation.utils import GenerationMixin
+
+import numpy as np
+# This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
+# It means that the function will not be traced through and simply appear as a node in the graph.
+if is_torch_fx_available():
+    if not is_torch_greater_or_equal_than_1_13:
+        import torch.fx
+
+    _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "LLaDA2MoeConfig"
+
+
+def _get_unpad_data(attention_mask):
+    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
+    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
+    max_seqlen_in_batch = seqlens_in_batch.max().item()
+    cu_seqlens = F.pad(
+        torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)
+    )
+    return (
+        indices,
+        cu_seqlens,
+        max_seqlen_in_batch,
+    )
+
+
+class LLaDA2MoeRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        LLaDA2MoeRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+
+ALL_LAYERNORM_LAYERS.append(LLaDA2MoeRMSNorm)
+
+
+class LLaDA2MoeRotaryEmbedding(nn.Module):
+    def __init__(self, config: LLaDA2MoeConfig, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+            self.rope_type = config.rope_scaling.get(
+                "rope_type", config.rope_scaling.get("type")
+            )
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = (
+            self.inv_freq[None, :, None]
+            .float()
+            .expand(position_ids.shape[0], -1, 1)
+            .to(x.device)
+        )
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = (
+            x.device.type
+            if isinstance(x.device.type, str) and x.device.type != "mps"
+            else "cpu"
+        )
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (
+                inv_freq_expanded.float() @ position_ids_expanded.float()
+            ).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+# Copied from transformers.models.llama.modeling_llama.rotate_half
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`):
+            The position indices of the tokens corresponding to the query and key tensors. For example, this can be
+            used to pass offsetted position ids when working with a KV-cache.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+
+    # Keep half or full tensor for later concatenation
+    rotary_dim = cos.shape[-1]
+    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+    # Apply rotary embeddings on the first half or full tensor
+    q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+    k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+
+    # Concatenate back to full shape
+    q_embed = torch.cat([q_embed, q_pass], dim=-1)
+    k_embed = torch.cat([k_embed, k_pass], dim=-1)
+    return q_embed, k_embed
+
+
+class LLaDA2MoeMLP(nn.Module):
+    def __init__(self, config: LLaDA2MoeConfig, intermediate_size: int):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = intermediate_size
+
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class LLaDA2MoeGate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_experts
+
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+
+        # topk selection algorithm
+        self.gating_dim = config.hidden_size
+        self.weight = nn.Parameter(torch.empty((self.num_experts, self.gating_dim)))
+        self.routed_scaling_factor = config.routed_scaling_factor
+
+        self.register_buffer("expert_bias", torch.zeros(self.num_experts))
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        import torch.nn.init as init
+
+        init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+
+    def group_limited_topk(
+        self,
+        scores: torch.Tensor,
+    ):
+        num_tokens, _ = scores.size()
+        # Organize the experts into groups
+        group_scores = (
+            scores.view(num_tokens, self.n_group, -1).topk(2, dim=-1)[0].sum(dim=-1)
+        )
+        group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
+        group_mask = torch.zeros_like(group_scores)
+        group_mask.scatter_(1, group_idx, 1)
+
+        # Mask the experts based on selection groups
+        score_mask = (
+            group_mask.unsqueeze(-1)
+            .expand(num_tokens, self.n_group, self.num_experts // self.n_group)
+            .reshape(num_tokens, -1)
+        )
+
+        masked_scores = scores.masked_fill(~score_mask.bool(), float("-inf"))
+        probs, top_indices = torch.topk(masked_scores, k=self.top_k, dim=-1)
+
+        return probs, top_indices
+
+    def forward(self, hidden_states):
+        # compute gating score
+        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+        logits = F.linear(
+            hidden_states.type(torch.float32), self.weight.type(torch.float32)
+        )
+
+        scores = torch.sigmoid(logits.float()).type_as(logits)
+
+        scores_for_routing = scores + self.expert_bias
+        _, topk_idx = self.group_limited_topk(scores_for_routing)
+
+        scores = torch.gather(scores, dim=1, index=topk_idx).type_as(logits)
+
+        topk_weight = (
+            scores / (scores.sum(dim=-1, keepdim=True) + 1e-20)
+            if self.top_k > 1
+            else scores
+        )
+        topk_weight = topk_weight * self.routed_scaling_factor
+
+        return topk_idx, topk_weight, logits
+
+
+class LLaDA2MoeSparseMoeBlock(nn.Module):
+    """
+    A mixed expert module containing shared experts.
+    """
+
+    def __init__(self, config: LLaDA2MoeConfig):
+        super().__init__()
+        self.config = config
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self._setup_experts()
+        self.gate = LLaDA2MoeGate(config)
+        if config.num_shared_experts is not None:
+            self.shared_experts = LLaDA2MoeMLP(
+                config=config,
+                intermediate_size=config.moe_intermediate_size
+                * config.num_shared_experts,
+            )
+
+    def _setup_experts(self):
+        self.experts = nn.ModuleList(
+            [
+                LLaDA2MoeMLP(
+                    config=self.config,
+                    intermediate_size=self.config.moe_intermediate_size,
+                )
+                for _ in range(self.config.num_experts)
+            ]
+        )
+
+    def forward(self, hidden_states):
+        identity = hidden_states
+        bsz, seq_len, h = hidden_states.shape
+        topk_idx, topk_weight, router_logits = self.gate(hidden_states)
+        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+        flat_topk_idx = topk_idx.view(-1)
+        if self.training:
+            hidden_states = hidden_states.repeat_interleave(
+                self.num_experts_per_tok, dim=0
+            )
+            y = torch.empty_like(hidden_states)
+            for i, expert in enumerate(self.experts):
+                y[flat_topk_idx == i] = expert(hidden_states[flat_topk_idx == i])
+            y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
+            y = y.to(hidden_states.dtype).view(bsz, seq_len, h)
+        else:
+            y = self.moe_infer(hidden_states, topk_idx, topk_weight).view(
+                bsz, seq_len, h
+            )
+        if self.config.num_shared_experts is not None:
+            y = y + self.shared_experts(identity)
+        return y, (
+            router_logits.view(bsz, seq_len, -1),
+            topk_idx.view(bsz, seq_len, -1),
+        )
+
+    @torch.no_grad()
+    def moe_infer(self, x, topk_ids, topk_weight):
+        cnts = topk_ids.new_zeros((topk_ids.shape[0], len(self.experts)))
+        cnts.scatter_(1, topk_ids, 1)
+        tokens_per_expert = cnts.sum(dim=0)
+        idxs = topk_ids.view(-1).argsort()
+        sorted_tokens = x[idxs // topk_ids.shape[1]]
+        tokens_per_expert = tokens_per_expert.cpu().numpy()
+        outputs = []
+        start_idx = 0
+        for i, num_tokens_tensor in enumerate(tokens_per_expert):
+            num_tokens = num_tokens_tensor.item()
+            if num_tokens == 0:
+                continue
+            end_idx = start_idx + num_tokens
+            expert = self.experts[i]
+            tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
+            expert_out = expert(tokens_for_this_expert)
+            outputs.append(expert_out.to(x.device))
+            start_idx = end_idx
+
+        outs = torch.cat(outputs, dim=0) if len(outputs) else sorted_tokens.new_empty(0)
+        new_x = torch.empty_like(outs)
+        new_x[idxs] = outs
+        final_out = (
+            new_x.view(*topk_ids.shape, -1)
+            .type(topk_weight.dtype)
+            .mul_(topk_weight.unsqueeze(dim=-1))
+            .sum(dim=1)
+            .type(new_x.dtype)
+        )
+        return final_out
+
+
+# Copied from transformers.models.llama.modeling_llama.repeat_kv
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_key_value_heads, n_rep, slen, head_dim
+    )
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = (
+        torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    )
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask[:, :, :, : key_states.shape[-2]]
+
+    # upcast attention to fp32
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(
+        query.dtype
+    )
+    attn_weights = nn.functional.dropout(
+        attn_weights, p=dropout, training=module.training
+    )
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaAttention with Llama->LLaDA2Moe
+class LLaDA2MoeAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: LLaDA2MoeConfig, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            logger.warning_once(
+                f"Instantiating {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.attention_dropout = config.attention_dropout
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = config.head_dim or self.hidden_size // self.num_heads
+        partial_rotary_factor = (
+            config.partial_rotary_factor
+            if hasattr(config, "partial_rotary_factor")
+            else 1.0
+        )
+        self.rope_dim = int(self.head_dim * partial_rotary_factor)
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.scaling = self.head_dim**-0.5
+        self.is_causal = False
+
+        self.query_key_value = nn.Linear(
+            self.hidden_size,
+            (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim,
+            bias=config.use_qkv_bias,
+        )
+
+        if self.config.use_qk_norm:
+            self.query_layernorm = LLaDA2MoeRMSNorm(
+                self.head_dim, eps=config.rms_norm_eps
+            )
+            self.key_layernorm = LLaDA2MoeRMSNorm(
+                self.head_dim, eps=config.rms_norm_eps
+            )
+        self.dense = nn.Linear(
+            self.num_heads * self.head_dim, self.hidden_size, bias=config.use_bias
+        )
+        self.sliding_window = getattr(config, "sliding_window", None)
+
+    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
+        return (
+            tensor.view(bsz, seq_len, self.num_heads, self.head_dim)
+            .transpose(1, 2)
+            .contiguous()
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Cache] = None,
+        output_attentions: bool = False,
+        use_cache: bool = False,
+        position_embeddings: Optional[
+            Tuple[torch.Tensor, torch.Tensor]
+        ] = None,  # necessary, but kept here for BC
+        **kwargs,
+    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+        input_shape = hidden_states.shape[:-1]
+
+        bsz, q_len, _ = hidden_states.size()
+
+        qkv = self.query_key_value(hidden_states)
+        qkv = qkv.view(
+            bsz, q_len, self.num_heads + 2 * self.num_key_value_heads, self.head_dim
+        )
+
+        query_states, key_states, value_states = qkv.split(
+            [self.num_heads, self.num_key_value_heads, self.num_key_value_heads], dim=-2
+        )
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        if self.config.use_qk_norm:
+            query_states = self.query_layernorm(query_states)
+            key_states = self.key_layernorm(key_states)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(
+            query_states, key_states, cos, sin
+        )
+
+        if past_key_value is not None:
+            if self.layer_idx is None:
+                raise ValueError(
+                    f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
+                    "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
+                    "with a layer index."
+                )
+            cache_kwargs = {"sin": sin, "cos": cos}
+            key_states, value_states = past_key_value.update(
+                key_states, value_states, self.layer_idx, cache_kwargs
+            )
+
+        attention_interface: Callable = eager_attention_forward
+        if self.config._attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[
+                self.config._attn_implementation
+            ]
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,  # diff with Llama
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.dense(attn_output)
+
+        return attn_output, attn_weights, past_key_value
+
+
+class LLaDA2MoeDecoderLayer(nn.Module):
+    def __init__(self, config: LLaDA2MoeConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.attention = LLaDA2MoeAttention(config=config, layer_idx=layer_idx)
+
+        self.mlp = (
+            LLaDA2MoeSparseMoeBlock(config)
+            if (
+                config.num_experts is not None
+                and layer_idx >= config.first_k_dense_replace
+            )
+            else LLaDA2MoeMLP(config=config, intermediate_size=config.intermediate_size)
+        )
+        self.input_layernorm = LLaDA2MoeRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_attention_layernorm = LLaDA2MoeRMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_value: Optional[Tuple[torch.Tensor]] = None,
+        output_attentions: Optional[bool] = False,
+        output_router_logits: Optional[bool] = False,
+        use_cache: Optional[bool] = False,
+        position_embeddings: Optional[
+            Tuple[torch.Tensor, torch.Tensor]
+        ] = None,  # necessary, but kept here for BC
+        **kwargs,
+    ) -> Tuple[
+        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
+    ]:
+        """
+        Args:
+            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+            attention_mask (`torch.FloatTensor`, *optional*):
+                attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
+                query_sequence_length, key_sequence_length)` if default attention is used.
+            position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+                config.n_positions - 1]`.
+            past_key_value (`Tuple(torch.FloatTensor)`, *optional*):
+                cached past key and value projection states
+            output_attentions (`bool`, *optional*):
+                Whether to return the attentions tensors of all attention layers. See `attentions` under
+                returned tensors for more detail.
+            output_router_logits (`bool`, *optional*):
+                Whether or not to return the logits of all the routers. They are useful for computing the router loss,
+                and should not be returned during inference.
+            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`).
+        """
+        residual = hidden_states
+
+        hidden_states = self.input_layernorm(hidden_states)
+
+        # Self Attention
+        hidden_states, self_attn_weights, present_key_value = self.attention(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            position_embeddings=position_embeddings,
+            use_cache=use_cache,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        if isinstance(hidden_states, tuple):
+            hidden_states, router_logits = hidden_states
+        else:
+            router_logits = None
+        hidden_states = residual + hidden_states.to(residual.device)
+
+        outputs = (hidden_states,)
+
+        if output_attentions:
+            outputs += (self_attn_weights,)
+
+        if use_cache:
+            outputs += (present_key_value,)
+
+        if output_router_logits:
+            outputs += (router_logits,)
+
+        return outputs
+
+
+LLADA2MOE_START_DOCSTRING = r"""
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
+    etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`LLaDA2MoeConfig`]):
+            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.
+"""
+
+
+@add_start_docstrings(
+    "The bare LLaDA2Moe Model outputting raw hidden-states without any specific head on top.",
+    LLADA2MOE_START_DOCSTRING,
+)
+class LLaDA2MoePreTrainedModel(PreTrainedModel):
+    config_class = LLaDA2MoeConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["LLaDA2MoeDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn_2 = False
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_cache_class = True
+
+    def _init_weights(self, module):
+        std = self.config.initializer_range
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+
+
+LLADA2MOE_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
+            it.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
+            `past_key_values`).
+
+            If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
+            and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
+            information on the default strategy.
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+        position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+            config.n_positions - 1]`.
+
+            [What are position IDs?](../glossary#position-ids)
+        past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
+            Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
+            blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
+            returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
+
+            Two formats are allowed:
+            - a [`~cache_utils.Cache`] instance;
+            - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
+            shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
+            cache format.
+
+            The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
+            legacy cache format will be returned.
+
+            If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
+            have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
+            of shape `(batch_size, sequence_length)`.
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+        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`).
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare LLaDA2Moe Model outputting raw hidden-states without any specific head on top.",
+    LLADA2MOE_START_DOCSTRING,
+)
+class LLaDA2MoeModel(LLaDA2MoePreTrainedModel):
+    """
+    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LLaDA2MoeDecoderLayer`]
+
+    Args:
+        config: LLaDA2MoeConfig
+    """
+
+    def __init__(self, config: LLaDA2MoeConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+
+        self.word_embeddings = nn.Embedding(
+            config.vocab_size, config.hidden_size, self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [
+                LLaDA2MoeDecoderLayer(config, layer_idx)
+                for layer_idx in range(config.num_hidden_layers)
+            ]
+        )
+        self._use_sdpa = config._attn_implementation == "sdpa"
+        self._use_flex_attention = config._attn_implementation == "flex_attention"
+        self.norm = LLaDA2MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = LLaDA2MoeRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.word_embeddings = value
+
+    @add_start_docstrings_to_model_forward(LLADA2MOE_INPUTS_DOCSTRING)
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_router_logits: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, MoeModelOutputWithPast]:
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        output_router_logits = (
+            output_router_logits
+            if output_router_logits is not None
+            else self.config.output_router_logits
+        )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+
+        # retrieve input_ids and inputs_embeds
+        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"
+            )
+        elif input_ids is not None:
+            batch_size, seq_length = input_ids.shape[:2]
+        elif inputs_embeds is not None:
+            batch_size, seq_length = inputs_embeds.shape[:2]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        if self.gradient_checkpointing and self.training:
+            if use_cache:
+                logger.warning_once(
+                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`transformers."
+                )
+                use_cache = False
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+
+        past_seen_tokens = (
+            past_key_values.get_seq_length() if past_key_values is not None else 0
+        )
+
+        if position_ids is None:
+            position_ids = torch.arange(
+                past_seen_tokens,
+                past_seen_tokens + inputs_embeds.shape[1],
+                device=inputs_embeds.device,
+            )
+            position_ids = position_ids.unsqueeze(0)
+
+        if self._use_flex_attention:
+            if attention_mask is not None and isinstance(attention_mask, torch.Tensor):
+                attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+                    attention_mask,
+                    (batch_size, seq_length),
+                    inputs_embeds,
+                    past_seen_tokens,
+                )
+        elif self._use_sdpa and not output_attentions:
+            # output_attentions=True can not be supported when using SDPA, and we fall back on
+            # the manual implementation that requires a 4D causal mask in all cases.
+            attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
+                attention_mask,
+                (batch_size, seq_length),
+                inputs_embeds,
+                past_seen_tokens,
+            )
+        else:
+            # 4d mask is passed through the layers
+            attention_mask = _prepare_4d_causal_attention_mask(
+                attention_mask,
+                (batch_size, seq_length),
+                inputs_embeds,
+                past_seen_tokens,
+            )
+
+        # embed positions
+        hidden_states = inputs_embeds
+
+        # create position embeddings to be shared across the decoder layers
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        # decoder layers
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        all_router_logits = () if output_router_logits else None
+        next_decoder_cache = None
+
+        for decoder_layer in self.layers:
+            if output_hidden_states:
+                all_hidden_states += (hidden_states,)
+
+            if self.gradient_checkpointing and self.training:
+                layer_outputs = self._gradient_checkpointing_func(
+                    decoder_layer.__call__,
+                    hidden_states,
+                    attention_mask,
+                    position_ids,
+                    past_key_values,
+                    output_attentions,
+                    output_router_logits,
+                    use_cache,
+                    position_embeddings,
+                )
+            else:
+                layer_outputs = decoder_layer(
+                    hidden_states,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    output_router_logits=output_router_logits,
+                    use_cache=use_cache,
+                    position_embeddings=position_embeddings,
+                )
+            hidden_states = layer_outputs[0]
+
+            if use_cache:
+                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
+
+            if output_attentions:
+                all_self_attns += (layer_outputs[1],)
+
+            if output_router_logits and layer_outputs[-1] is not None:
+                all_router_logits += (layer_outputs[-1],)
+
+        hidden_states = self.norm(hidden_states)
+
+        # add hidden states from the last decoder layer
+        if output_hidden_states:
+            all_hidden_states += (hidden_states,)
+
+        next_cache = None
+        if use_cache:
+            next_cache = next_decoder_cache
+        if not return_dict:
+            return tuple(
+                v
+                for v in [
+                    hidden_states,
+                    next_cache,
+                    all_hidden_states,
+                    all_self_attns,
+                    all_router_logits,
+                ]
+                if v is not None
+            )
+        return MoeModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=all_hidden_states,
+            attentions=all_self_attns,
+            router_logits=all_router_logits,
+        )
+
+
+class LLaDA2MoeModelLM(LLaDA2MoePreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ["lm_head.weight"]
+
+    def __init__(self, config: LLaDA2MoeConfig):
+        super().__init__(config)
+        self.model = LLaDA2MoeModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.model.word_embeddings = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    def set_decoder(self, decoder):
+        self.model = decoder
+
+    def get_decoder(self):
+        return self.model
+
+    @add_start_docstrings_to_model_forward(LLADA2MOE_INPUTS_DOCSTRING)
+    @replace_return_docstrings(
+        output_type=MoeCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
+    )
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        output_router_logits: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, MoeCausalLMOutputWithPast]:
+        r"""
+        Args:
+            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Returns:
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer
+
+        >>> model = LLaDA2MoeForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
+        >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        output_attentions = (
+            output_attentions
+            if output_attentions is not None
+            else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        output_router_logits = (
+            output_router_logits
+            if output_router_logits is not None
+            else self.config.output_router_logits
+        )
+        return_dict = (
+            return_dict if return_dict is not None else self.config.use_return_dict
+        )
+        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_router_logits=output_router_logits,
+            return_dict=return_dict,
+            **kwargs,
+        )
+
+        loss = None
+        aux_loss = None
+        hidden_states = outputs[0]
+
+        logits = self.lm_head(hidden_states)
+        logits = logits.float()
+
+        if labels is not None:
+            # LLaDA2.0 will use same label position logits
+            shift_logits = logits
+            shift_labels = labels
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            shift_logits = shift_logits.view(-1, self.config.vocab_size)
+            shift_labels = shift_labels.view(-1)
+            # Enable model parallelism
+            shift_labels = shift_labels.to(shift_logits.device)
+            loss = loss_fct(shift_logits, shift_labels)
+
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            if output_router_logits:
+                output = (aux_loss,) + output
+            return (loss,) + output if loss is not None else output
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            router_logits=outputs.router_logits,
+        )
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        token_type_ids=None,
+        **kwargs,
+    ):
+        if past_key_values is not None:
+            if isinstance(past_key_values, Cache):
+                cache_length = past_key_values.get_seq_length()
+                past_length = past_key_values.seen_tokens
+                max_cache_length = (
+                    past_key_values.get_max_length()
+                    if hasattr(past_key_values, "get_max_length")
+                    else past_key_values.get_max_cache_shape()
+                )
+            else:
+                cache_length = past_length = past_key_values[0][0].shape[2]
+                max_cache_length = None
+
+            # Keep only the unprocessed tokens:
+            # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
+            # some of the inputs are exclusivelly passed as part of the cache (e.g. when passing input_embeds as input)
+            if (
+                attention_mask is not None
+                and attention_mask.shape[1] > input_ids.shape[1]
+            ):
+                input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
+            # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
+            # input_ids based on the past_length.
+            elif past_length < input_ids.shape[1]:
+                input_ids = input_ids[:, past_length:]
+            # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
+
+            # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
+            if (
+                max_cache_length is not None
+                and attention_mask is not None
+                and cache_length + input_ids.shape[1] > max_cache_length
+            ):
+                attention_mask = attention_mask[:, -max_cache_length:]
+
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            # create position_ids on the fly for batch generation
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values:
+                position_ids = position_ids[:, -input_ids.shape[1] :]
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "position_ids": position_ids,
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    @staticmethod
+    def _reorder_cache(past_key_values, beam_idx):
+        reordered_past = ()
+        for layer_past in past_key_values:
+            reordered_past += (
+                tuple(
+                    past_state.index_select(0, beam_idx.to(past_state.device))
+                    for past_state in layer_past
+                ),
+            )
+        return reordered_past
+
+    @staticmethod
+    def _top_k_logits(logits, k):
+        if k is None or k <= 0:
+            return logits
+        else:
+            values, _ = torch.topk(logits, k)
+            min_values = values[..., -1, None]
+            return torch.where(
+                logits < min_values, torch.full_like(logits, float("-inf")), logits
+            )
+
+    @staticmethod
+    def _top_p_logits(logits, p):
+        if p is None or p >= 1.0:
+            return logits
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+        sorted_mask = cumulative_probs > p
+        sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
+        sorted_mask[..., 0] = False
+        mask_indices = torch.scatter(
+            torch.full_like(logits, False, dtype=torch.bool),
+            -1,
+            sorted_indices,
+            sorted_mask,
+        )
+        return logits.masked_fill(mask_indices, float("-inf"))
+
+    def _sample_with_temperature_topk_topp(
+        self, logits, temperature=1.0, top_k=0, top_p=1.0
+    ):
+        orig_shape = logits.shape[:-1]
+        vocab_size = logits.shape[-1]
+        logits = logits.reshape(-1, vocab_size)
+
+        # Greedy mode: temperature = 0, no top-k/p
+        if temperature == 0.0:
+            probs = F.softmax(logits, dim=-1)
+            token = logits.argmax(dim=-1, keepdim=True)
+            token_prob = probs.gather(-1, token)
+            return token.view(*orig_shape), token_prob.view(*orig_shape)
+
+        if temperature > 0 and temperature != 1.0:
+            logits = logits / temperature
+        logits = self._top_k_logits(logits, top_k)
+        logits = self._top_p_logits(logits, top_p)
+        probs = F.softmax(logits, dim=-1)
+        token = torch.multinomial(probs, num_samples=1)
+        token_prob = torch.gather(probs, -1, token)
+        # token = logits.argmax(dim=-1, keepdim=True)
+        return token.view(*orig_shape), token_prob.view(*orig_shape)
+
+    @staticmethod
+    def _get_num_transfer_tokens(block_length, steps):
+        if steps == 0:
+            return torch.tensor([], dtype=torch.int64)
+        base = block_length // steps
+        remainder = block_length % steps
+        num_transfer_tokens = torch.full((steps,), base, dtype=torch.int64)
+        num_transfer_tokens[:remainder] += 1
+        return num_transfer_tokens
+
+    @torch.no_grad()
+    def generate(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        temperature: int = 0.0,
+        block_length: int = 32,
+        steps: int = 32,
+        gen_length: int = 2048,
+        top_p: Optional[int] = None,
+        top_k: Optional[int] = None,
+        eos_early_stop: bool = False,
+        minimal_topk: int = 1,
+        threshold: float = 0.95,
+        eos_id: int = 156892,
+        mask_id: int = 156895,
+    ):
+        r"""
+        Generates tokens using a block-wise, iterative refinement strategy.
+
+        This method operates differently from standard autoregressive generation. It first creates a template of the
+        full desired length, filled with a special `mask_id`. It then processes this template in segments (`blocks`)
+        and iteratively "denoises" or "refines" the `mask_id` tokens into actual tokens over a series of `steps` for
+        each block. A custom block-diagonal causal attention mask ensures that generation within a block can attend to
+        all previous blocks but not future ones.
+
+        <Tip warning={true}>
+
+        This is a specialized generation method. The quality and speed of the output are highly dependent on the interplay
+        between `block_length`, `steps`, and `threshold`. It aims to achieve faster generation through parallel
+        decoding within blocks, which is a departure from the token-by-token generation of standard `.generate()` methods.
+
+        </Tip>
+
+        Parameters:
+            inputs (`torch.Tensor`):
+                The token sequence used as a prompt for the generation.
+            temperature (`float`, *optional*, defaults to 0.0):
+                The value used to module the next token probabilities. A value of 0.0 corresponds to greedy decoding.
+            block_length (`int`, *optional*, defaults to 32):
+                The size of each generation block. The model generates text in parallel within these blocks. This is a
+                key parameter for controlling the granularity of the generation process.
+            steps (`int`, *optional*, defaults to 32):
+                The number of iterative refinement (or "denoising") steps to perform for each block. Within each block,
+                the model will try to replace `mask_id` tokens with real tokens for this many iterations.
+            gen_length (`int`, *optional*, defaults to 2048):
+                The maximum number of tokens to generate, excluding the prompt.
+            top_p (`float`, *optional*):
+                If set to a float value between 0 and 1, only the most probable tokens with probabilities that add up to
+                `top_p` or higher are kept for generation (nucleus sampling).
+            top_k (`int`, *optional*):
+                The number of highest probability vocabulary tokens to keep for top-k-filtering.
+            eos_early_stop (`bool`, *optional*, defaults to `False`):
+                If `True`, generation will stop as soon as a valid End-Of-Sequence token is generated and confirmed,
+                even if `gen_length` has not been reached.
+            minimal_topk (`int`, *optional*, defaults to 1):
+                A parameter used to dynamically adjust the number of refinement `steps`. The effective number of steps
+                is capped at `gen_length // minimal_topk`.
+            threshold (`float`, *optional*, defaults to 0.95):
+                The confidence probability threshold for accepting a sampled token. During each refinement step, a
+                sampled token is only kept if its probability is above this threshold. If not enough tokens meet the
+                threshold, the ones with the highest confidence are chosen.
+            eos_id (`int`, *optional*, defaults to 156892):
+                The token ID for the end-of-sequence token. Used for `eos_early_stop`.
+            mask_id (`int`, *optional*, defaults to 156895):
+                The token ID used as a placeholder for tokens that are yet to be generated. This is central to the
+                iterative refinement algorithm.
+
+        Return:
+            `torch.Tensor`: A string containing the generated token IDs, starting
+            after the prompt and stopping at the first `eos_id` or `gen_length`.
+        """
+        steps = min(steps, gen_length // minimal_topk)
+        input_ids = inputs.to(self.device)
+
+        prompt_length = input_ids.shape[1]
+        num_blocks = (prompt_length + gen_length + block_length - 1) // block_length
+        total_length = num_blocks * block_length
+
+        block_mask = torch.tril(torch.ones(num_blocks, num_blocks, device=self.device))
+        block_diffusion_attention_mask = (
+            (
+                block_mask.repeat_interleave(block_length, dim=0)
+                .repeat_interleave(block_length, dim=1)
+                .unsqueeze(0)
+                .unsqueeze(0)
+            )
+            .log()
+            .to(torch.bfloat16)
+        )
+
+        position_ids = torch.arange(total_length, device=self.device).unsqueeze(0)
+        x = torch.full((1, total_length), mask_id, dtype=torch.long, device=self.device)
+        x[:, :prompt_length] = input_ids.clone()
+
+        prompt_index_full = torch.zeros_like(x, dtype=torch.bool)
+        prompt_index_full[:, :prompt_length] = True
+
+        prefill_blocks = prompt_length // block_length
+
+        denoising_steps_per_block = steps
+        num_transfer_tokens_schedule = self._get_num_transfer_tokens(
+            block_length, denoising_steps_per_block
+        )
+
+        nfe = 0
+
+        for num_block in range(prefill_blocks, num_blocks):
+            current_window_end = (num_block + 1) * block_length
+            cur_x = x[:, :current_window_end]
+            cur_attn_mask = block_diffusion_attention_mask[
+                :, :, :current_window_end, :current_window_end
+            ]
+            cur_position_ids = position_ids[:, :current_window_end]
+
+            for _ in range(denoising_steps_per_block):
+                active_block_mask = cur_x[:, -block_length:] == mask_id
+                if active_block_mask.sum() == 0:
+                    break
+
+                logits = self.forward(
+                    cur_x,
+                    attention_mask=cur_attn_mask,
+                    position_ids=cur_position_ids,
+                ).logits
+
+                active_logits = logits[:, -block_length:, :]
+                # active_logits = logits[:, -block_length-1:-1, :]
+                x0, x0_p = self._sample_with_temperature_topk_topp(
+                    active_logits, temperature=temperature, top_k=top_k, top_p=top_p
+                )
+                nfe += 1
+
+                num_to_transfer = num_transfer_tokens_schedule[step].item()
+                transfer_index = torch.zeros_like(x0, dtype=torch.bool)
+
+                confidence = torch.where(active_block_mask, x0_p, -torch.inf)
+                high_conf_mask = confidence[0] > threshold
+                num_high_confidence = high_conf_mask.sum().item()
+
+                if num_high_confidence >= num_to_transfer:
+                    transfer_index[0] = high_conf_mask
+                else:
+                    _, idx = torch.topk(
+                        confidence[0],
+                        k=min(num_to_transfer, active_block_mask.sum().item()),
+                    )
+                    transfer_index[0, idx] = True
+
+                if transfer_index.any():
+                    cur_x[:, -block_length:][transfer_index] = x0[transfer_index]
+                if eos_early_stop and (x0[transfer_index] == eos_id).any():
+                    eos_pos_in_x = (cur_x[0] == eos_id).nonzero(as_tuple=True)
+                    if len(eos_pos_in_x[0]) > 0:
+                        eos_pos = eos_pos_in_x[0][0].item()
+                        if (cur_x[0, prompt_length:eos_pos] != mask_id).all():
+                            final_x = x[:, :total_length][:, : eos_pos + 1]
+                            return nfe, final_x
+
+            x[:, :current_window_end] = cur_x
+            if (
+                eos_id is not None
+                and (x[0, prompt_length:current_window_end] == eos_id).any()
+            ):
+                break
+            
+
+        generated_answer = x[:, : prompt_length + gen_length]
+
+        mask_positions = (generated_answer[0][input_ids.shape[1] :] == eos_id).nonzero(
+            as_tuple=True
+        )[0]
+        if len(mask_positions) > 0:
+            first_mask_position = mask_positions[0].item()
+        else:
+            first_mask_position = gen_length
+        return nfe, generated_answer[
+            :, input_ids.shape[1] : input_ids.shape[1] + first_mask_position + 1
+        ]
+    
+
+
+
+    @torch.no_grad()
+    def generate_spd(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        block_length: int = 32,
+        steps: int = 32,
+        gen_length: int = 2048,
+        minimal_topk: int = 1,
+        threshold: float = 0.95,
+        eos_id: int = 156892,
+        mask_id: int = 156895,
+    ):
+        r"""
+        Generates tokens using a block-wise, iterative refinement strategy.
+
+        This method operates differently from standard autoregressive generation. It first creates a template of the
+        full desired length, filled with a special `mask_id`. It then processes this template in segments (`blocks`)
+        and iteratively "denoises" or "refines" the `mask_id` tokens into actual tokens over a series of `steps` for
+        each block. A custom block-diagonal causal attention mask ensures that generation within a block can attend to
+        all previous blocks but not future ones.
+
+        <Tip warning={true}>
+
+        This is a specialized generation method. The quality and speed of the output are highly dependent on the interplay
+        between `block_length`, `steps`, and `threshold`. It aims to achieve faster generation through parallel
+        decoding within blocks, which is a departure from the token-by-token generation of standard `.generate()` methods.
+
+        </Tip>
+
+        Parameters:
+            inputs (`torch.Tensor`):
+                The token sequence used as a prompt for the generation.
+            block_length (`int`, *optional*, defaults to 32):
+                The size of each generation block. The model generates text in parallel within these blocks. This is a
+                key parameter for controlling the granularity of the generation process.
+            steps (`int`, *optional*, defaults to 32):
+                The number of iterative refinement (or "denoising") steps to perform for each block. Within each block,
+                the model will try to replace `mask_id` tokens with real tokens for this many iterations.
+            gen_length (`int`, *optional*, defaults to 2048):
+                The maximum number of tokens to generate, excluding the prompt.
+            minimal_topk (`int`, *optional*, defaults to 1):
+                A parameter used to dynamically adjust the number of refinement `steps`. The effective number of steps
+                is capped at `gen_length // minimal_topk`.
+            threshold (`float`, *optional*, defaults to 0.95):
+                The confidence probability threshold for accepting a sampled token. During each refinement step, a
+                sampled token is only kept if its probability is above this threshold. If not enough tokens meet the
+                threshold, the ones with the highest confidence are chosen.
+            eos_id (`int`, *optional*, defaults to 156892):
+                The token ID for the end-of-sequence token. Used for `eos_early_stop`.
+            mask_id (`int`, *optional*, defaults to 156895):
+                The token ID used as a placeholder for tokens that are yet to be generated. This is central to the
+                iterative refinement algorithm.
+
+        Return:
+            `torch.Tensor`: A string containing the generated token IDs, starting
+            after the prompt and stopping at the first `eos_id` or `gen_length`.
+        """
+        steps = min(steps, gen_length // minimal_topk)
+        input_ids = inputs.to(self.device)
+
+        prompt_length = input_ids.shape[1]
+        num_blocks = (prompt_length + gen_length + block_length - 1) // block_length
+        total_length = num_blocks * block_length
+
+        block_mask = torch.tril(torch.ones(num_blocks, num_blocks, device=self.device))
+        block_diffusion_attention_mask = (
+            (
+                block_mask.repeat_interleave(block_length, dim=0)
+                .repeat_interleave(block_length, dim=1)
+                .unsqueeze(0)
+                .unsqueeze(0)
+            )
+            .log()
+            .to(torch.bfloat16)
+        )
+
+        position_ids = torch.arange(total_length, device=self.device).unsqueeze(0)
+        x = torch.full((1, total_length), mask_id, dtype=torch.long, device=self.device)
+        x[:, :prompt_length] = input_ids.clone()
+        input_embeddings = self.get_input_embeddings()
+        mask_embedding = input_embeddings.weight[mask_id].to(self.device).view(1, 1, -1)
+
+        prefill_blocks = prompt_length // block_length
+
+        denoising_steps_per_block = min(steps, block_length)
+        nfe = 0
+
+        for num_block in range(prefill_blocks, num_blocks):
+            current_window_end = (num_block + 1) * block_length
+            cur_x = x[:, :current_window_end]
+            # Cache token embeddings for the visible prefix and only refresh the active block.
+            cur_token_embeds = input_embeddings(cur_x)
+            cur_inputs_embeds = cur_token_embeds.clone()
+            cur_attn_mask = block_diffusion_attention_mask[
+                :, :, :current_window_end, :current_window_end
+            ]
+            cur_position_ids = position_ids[:, :current_window_end]
+
+            # Only non-prompt positions in the current block participate in iterative decoding.
+            active_block_mask = torch.arange(
+                current_window_end - block_length,
+                current_window_end,
+                device=cur_x.device,
+            ).unsqueeze(0)
+            active_block_mask = active_block_mask >= prompt_length
+            block_slice = slice(-block_length, None)
+            expanded_mask_embedding = mask_embedding.expand(1, block_length, -1)
+            expanded_mask_norm = torch.linalg.vector_norm(
+                expanded_mask_embedding.float(), dim=-1, keepdim=True
+            ).to(cur_token_embeds.dtype)
+            block_confidence = torch.zeros(
+                (1, block_length), device=cur_x.device, dtype=torch.float32
+            )
+
+            for _ in range(denoising_steps_per_block):
+                current_block = cur_x[:, block_slice]
+                prev_block = current_block.clone()
+                mask_index = current_block == mask_id
+                token_index = active_block_mask & (~mask_index)
+                block_token_embeds = cur_token_embeds[:, block_slice, :]
+                block_inputs_embeds = block_token_embeds.clone()
+                token_weight = block_confidence.to(block_inputs_embeds.dtype).unsqueeze(-1)
+
+                # Token positions use a confidence-weighted token/mask blend before the forward pass.
+                mixed_embeds = (
+                    token_weight * block_token_embeds
+                    + (1.0 - token_weight) * expanded_mask_embedding
+                )
+                token_norm = torch.linalg.vector_norm(
+                    block_token_embeds.float(), dim=-1, keepdim=True
+                ).to(block_inputs_embeds.dtype)
+                target_norm = (
+                    token_weight * token_norm
+                    + (1.0 - token_weight) * expanded_mask_norm
+                )
+                mixed_norm = torch.linalg.vector_norm(
+                    mixed_embeds.float(), dim=-1, keepdim=True
+                ).clamp_min(1e-12).to(block_inputs_embeds.dtype)
+                # Renormalize the blended embedding to the weighted target norm.
+                mixed_embeds = mixed_embeds * (target_norm / mixed_norm)
+
+                # Mask positions stay on the mask embedding; token positions receive the blended embedding.
+                block_inputs_embeds = torch.where(
+                    mask_index.unsqueeze(-1),
+                    expanded_mask_embedding,
+                    block_inputs_embeds,
+                )
+                block_inputs_embeds = torch.where(
+                    token_index.unsqueeze(-1),
+                    mixed_embeds,
+                    block_inputs_embeds,
+                )
+                cur_inputs_embeds[:, block_slice, :] = block_inputs_embeds
+
+                logits = self.forward(
+                    inputs_embeds=cur_inputs_embeds,
+                    attention_mask=cur_attn_mask,
+                    position_ids=cur_position_ids,
+                ).logits
+                nfe += 1
+
+                active_logits = logits[:, -block_length:, :]
+                active_probs = F.softmax(active_logits.float(), dim=-1)
+                top1_confidence, top1_tokens = torch.max(active_probs, dim=-1)
+
+                target_slice = current_block.clone()
+                # Every active token index is refreshed by the current step's top-1 prediction.
+                target_slice = torch.where(token_index, top1_tokens, target_slice)
+
+                mask_positions = torch.nonzero(mask_index[0], as_tuple=False).flatten()
+                if mask_positions.numel() > 0:
+                    mask_confidence = top1_confidence[0, mask_positions]
+                    below_threshold = torch.nonzero(
+                        mask_confidence < threshold, as_tuple=False
+                    ).flatten()
+
+                    if below_threshold.numel() == 0:
+                        decode_upto = mask_positions.numel()
+                    elif below_threshold[0].item() == 0:
+                        decode_upto = 1
+                    else:
+                        decode_upto = below_threshold[0].item()
+
+                    # Decode the leftmost mask prefix above threshold, or force one token if needed.
+                    decode_positions = mask_positions[:decode_upto]
+                    target_slice[0, decode_positions] = top1_tokens[0, decode_positions]
+
+                cur_x[:, block_slice] = torch.where(
+                    active_block_mask, target_slice, cur_x[:, block_slice]
+                )
+                current_block = cur_x[:, block_slice]
+                same_as_previous = torch.equal(current_block, prev_block)
+                active_confidence = torch.where(
+                    active_block_mask,
+                    top1_confidence,
+                    torch.ones_like(top1_confidence),
+                )
+                all_confident = bool((active_confidence >= 0.9).all().item())
+
+                # A block is committed once it stops changing or every active position is confident enough.
+                if same_as_previous or all_confident:
+                    break
+
+                cur_token_embeds[:, block_slice, :] = input_embeddings(current_block)
+                block_confidence = torch.where(
+                    active_block_mask & (current_block != mask_id),
+                    top1_confidence,
+                    torch.zeros_like(top1_confidence),
+                )
+
+            x[:, :current_window_end] = cur_x
+            
+            if (
+                eos_id is not None
+                and (x[0, prompt_length:current_window_end] == eos_id).any()
+            ):
+                break
+
+        generated_answer = x[:, : prompt_length + gen_length]
+
+        mask_positions = (generated_answer[0][input_ids.shape[1] :] == eos_id).nonzero(
+            as_tuple=True
+        )[0]
+        if len(mask_positions) > 0:
+            first_mask_position = mask_positions[0].item()
+        else:
+            first_mask_position = gen_length
+        return nfe, generated_answer[
+            :, input_ids.shape[1] : input_ids.shape[1] + first_mask_position + 1
+        ]
+
+    @torch.no_grad()
+    def generate_uniform_demo(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        block_length: int = 32,
+        steps: int = 32,
+        gen_length: int = 2048,
+        minimal_topk: int = 1,
+        threshold: float = 0.95,
+        eos_id: int = 156892,
+        mask_id: int = 156895,
+    ):
+        r"""
+        Runs the same decoding logic as `generate_uniform` while storing
+        step-by-step metadata for visualization.
+        """
+        steps = min(steps, gen_length // minimal_topk)
+        input_ids = inputs.to(self.device)
+
+        prompt_length = input_ids.shape[1]
+        num_blocks = (prompt_length + gen_length + block_length - 1) // block_length
+        total_length = num_blocks * block_length
+
+        block_mask = torch.tril(torch.ones(num_blocks, num_blocks, device=self.device))
+        block_diffusion_attention_mask = (
+            (
+                block_mask.repeat_interleave(block_length, dim=0)
+                .repeat_interleave(block_length, dim=1)
+                .unsqueeze(0)
+                .unsqueeze(0)
+            )
+            .log()
+            .to(torch.bfloat16)
+        )
+
+        position_ids = torch.arange(total_length, device=self.device).unsqueeze(0)
+        x = torch.full((1, total_length), mask_id, dtype=torch.long, device=self.device)
+        x[:, :prompt_length] = input_ids.clone()
+        input_embeddings = self.get_input_embeddings()
+        mask_embedding = input_embeddings.weight[mask_id].to(self.device).view(1, 1, -1)
+
+        prefill_blocks = prompt_length // block_length
+
+        denoising_steps_per_block = min(steps, block_length)
+        nfe = 0
+        frames = []
+        block_summaries = []
+
+        for num_block in range(prefill_blocks, num_blocks):
+            current_window_end = (num_block + 1) * block_length
+            cur_x = x[:, :current_window_end]
+            cur_token_embeds = input_embeddings(cur_x)
+            cur_inputs_embeds = cur_token_embeds.clone()
+            cur_attn_mask = block_diffusion_attention_mask[
+                :, :, :current_window_end, :current_window_end
+            ]
+            cur_position_ids = position_ids[:, :current_window_end]
+
+            active_block_mask = torch.arange(
+                current_window_end - block_length,
+                current_window_end,
+                device=cur_x.device,
+            ).unsqueeze(0)
+            active_block_mask = active_block_mask >= prompt_length
+            block_slice = slice(-block_length, None)
+            expanded_mask_embedding = mask_embedding.expand(1, block_length, -1)
+            expanded_mask_norm = torch.linalg.vector_norm(
+                expanded_mask_embedding.float(), dim=-1, keepdim=True
+            ).to(cur_token_embeds.dtype)
+            block_confidence = torch.zeros(
+                (1, block_length), device=cur_x.device, dtype=torch.float32
+            )
+
+            block_summary = {
+                "block_id": int(num_block - prefill_blocks),
+                "absolute_block_id": int(num_block),
+                "block_start": int(current_window_end - block_length),
+                "block_end": int(current_window_end),
+                "window_end": int(current_window_end),
+                "num_steps": 0,
+                "converged": False,
+                "convergence_reason": "max_steps",
+            }
+
+            for step_idx in range(denoising_steps_per_block):
+                current_block = cur_x[:, block_slice]
+                prev_block = current_block.clone()
+                pre_visible_ids = cur_x[0, :current_window_end].detach().cpu().tolist()
+                mask_index = current_block == mask_id
+                token_index = active_block_mask & (~mask_index)
+                input_confidence = block_confidence[0].detach().cpu().tolist()
+                block_token_embeds = cur_token_embeds[:, block_slice, :]
+                block_inputs_embeds = block_token_embeds.clone()
+                token_weight = block_confidence.to(block_inputs_embeds.dtype).unsqueeze(-1)
+
+                mixed_embeds = (
+                    token_weight * block_token_embeds
+                    + (1.0 - token_weight) * expanded_mask_embedding
+                )
+                token_norm = torch.linalg.vector_norm(
+                    block_token_embeds.float(), dim=-1, keepdim=True
+                ).to(block_inputs_embeds.dtype)
+                target_norm = (
+                    token_weight * token_norm
+                    + (1.0 - token_weight) * expanded_mask_norm
+                )
+                mixed_norm = torch.linalg.vector_norm(
+                    mixed_embeds.float(), dim=-1, keepdim=True
+                ).clamp_min(1e-12).to(block_inputs_embeds.dtype)
+                mixed_embeds = mixed_embeds * (target_norm / mixed_norm)
+
+                block_inputs_embeds = torch.where(
+                    mask_index.unsqueeze(-1),
+                    expanded_mask_embedding,
+                    block_inputs_embeds,
+                )
+                block_inputs_embeds = torch.where(
+                    token_index.unsqueeze(-1),
+                    mixed_embeds,
+                    block_inputs_embeds,
+                )
+                cur_inputs_embeds[:, block_slice, :] = block_inputs_embeds
+
+                logits = self.forward(
+                    inputs_embeds=cur_inputs_embeds,
+                    attention_mask=cur_attn_mask,
+                    position_ids=cur_position_ids,
+                ).logits
+                nfe += 1
+
+                active_logits = logits[:, -block_length:, :]
+                active_probs = F.softmax(active_logits.float(), dim=-1)
+                top1_confidence, top1_tokens = torch.max(active_probs, dim=-1)
+
+                target_slice = current_block.clone()
+                target_slice = torch.where(token_index, top1_tokens, target_slice)
+
+                decode_positions = torch.tensor(
+                    [], device=cur_x.device, dtype=torch.long
+                )
+                mask_positions = torch.nonzero(mask_index[0], as_tuple=False).flatten()
+                if mask_positions.numel() > 0:
+                    mask_confidence = top1_confidence[0, mask_positions]
+                    below_threshold = torch.nonzero(
+                        mask_confidence < threshold, as_tuple=False
+                    ).flatten()
+
+                    if below_threshold.numel() == 0:
+                        decode_upto = mask_positions.numel()
+                    elif below_threshold[0].item() == 0:
+                        decode_upto = 1
+                    else:
+                        decode_upto = below_threshold[0].item()
+
+                    decode_positions = mask_positions[:decode_upto]
+                    target_slice[0, decode_positions] = top1_tokens[0, decode_positions]
+
+                cur_x[:, block_slice] = torch.where(
+                    active_block_mask, target_slice, cur_x[:, block_slice]
+                )
+                current_block = cur_x[:, block_slice]
+                same_as_previous = torch.equal(current_block, prev_block)
+                active_confidence = torch.where(
+                    active_block_mask,
+                    top1_confidence,
+                    torch.ones_like(top1_confidence),
+                )
+                all_confident = bool((active_confidence >= 0.9).all().item())
+                converged = same_as_previous or all_confident
+                convergence_reason = None
+                if same_as_previous:
+                    convergence_reason = "stable_tokens"
+                elif all_confident:
+                    convergence_reason = "high_confidence"
+
+                post_visible_ids = cur_x[0, :current_window_end].detach().cpu().tolist()
+                frames.append(
+                    {
+                        "frame_id": len(frames),
+                        "block_id": int(num_block - prefill_blocks),
+                        "absolute_block_id": int(num_block),
+                        "step_id": int(step_idx),
+                        "window_end": int(current_window_end),
+                        "block_start": int(current_window_end - block_length),
+                        "block_end": int(current_window_end),
+                        "nfe": int(nfe),
+                        "pre_visible_ids": pre_visible_ids,
+                        "post_visible_ids": post_visible_ids,
+                        "active_block_mask": active_block_mask[0].detach().cpu().tolist(),
+                        "mask_index_before": mask_index[0].detach().cpu().tolist(),
+                        "token_index_before": token_index[0].detach().cpu().tolist(),
+                        "input_confidence": input_confidence,
+                        "top1_confidence": top1_confidence[0].detach().cpu().tolist(),
+                        "top1_token_ids": top1_tokens[0].detach().cpu().tolist(),
+                        "decoded_positions": decode_positions.detach().cpu().tolist(),
+                        "same_as_previous": bool(same_as_previous),
+                        "all_confident": bool(all_confident),
+                        "converged": bool(converged),
+                        "convergence_reason": convergence_reason,
+                    }
+                )
+
+                block_summary["num_steps"] = step_idx + 1
+                if converged:
+                    block_summary["converged"] = True
+                    block_summary["convergence_reason"] = convergence_reason
+                    break
+
+                cur_token_embeds[:, block_slice, :] = input_embeddings(current_block)
+                block_confidence = torch.where(
+                    active_block_mask & (current_block != mask_id),
+                    top1_confidence,
+                    torch.zeros_like(top1_confidence),
+                )
+
+            x[:, :current_window_end] = cur_x
+            block_summaries.append(block_summary)
+
+            if (
+                eos_id is not None
+                and (x[0, prompt_length:current_window_end] == eos_id).any()
+            ):
+                break
+
+        generated_answer = x[:, : prompt_length + gen_length]
+
+        mask_positions = (generated_answer[0][input_ids.shape[1] :] == eos_id).nonzero(
+            as_tuple=True
+        )[0]
+        if len(mask_positions) > 0:
+            first_mask_position = mask_positions[0].item()
+        else:
+            first_mask_position = gen_length
+
+        generated_tokens = generated_answer[
+            :, input_ids.shape[1] : input_ids.shape[1] + first_mask_position + 1
+        ]
+        demo_trace = {
+            "prompt_length": int(prompt_length),
+            "block_length": int(block_length),
+            "steps": int(denoising_steps_per_block),
+            "gen_length": int(gen_length),
+            "threshold": float(threshold),
+            "eos_id": int(eos_id) if eos_id is not None else None,
+            "mask_id": int(mask_id),
+            "nfe": int(nfe),
+            "prompt_token_ids": input_ids[0].detach().cpu().tolist(),
+            "generated_token_ids": generated_tokens[0].detach().cpu().tolist(),
+            "final_token_ids": generated_answer[0].detach().cpu().tolist(),
+            "frames": frames,
+            "blocks": block_summaries,
+        }
+
+        return demo_trace, nfe, generated_tokens

models/dmax-16b/special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "bos_token": {
+    "content": "<|startoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "<|mask|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

models/dmax-16b/tokenizer_config.json

@@ -0,0 +1,2116 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "added_tokens_decoder": {
+    "156891": {
+      "content": "<|startoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156892": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156893": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156894": {
+      "content": "[gMASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156895": {
+      "content": "<|mask|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156896": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156897": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156898": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156899": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156900": {
+      "content": "<|role_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156901": {
+      "content": "<|reserved_token_6|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156902": {
+      "content": "<|reserved_token_7|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156903": {
+      "content": "<|reserved_token_8|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156904": {
+      "content": "<|reserved_token_9|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156905": {
+      "content": "<|reserved_token_10|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156906": {
+      "content": "<|reserved_token_11|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156907": {
+      "content": "<|reserved_token_12|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156908": {
+      "content": "<|reserved_token_13|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156909": {
+      "content": "<|reserved_token_14|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156910": {
+      "content": "<|reserved_token_15|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156911": {
+      "content": "<|reserved_token_16|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156912": {
+      "content": "<|reserved_token_17|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156913": {
+      "content": "<|reserved_token_18|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156914": {
+      "content": "<|reserved_token_19|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156915": {
+      "content": "<|reserved_token_20|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156916": {
+      "content": "<|reserved_token_21|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156917": {
+      "content": "<|reserved_token_22|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156918": {
+      "content": "<|reserved_token_23|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156919": {
+      "content": "<|reserved_token_24|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156920": {
+      "content": "<|reserved_token_25|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156921": {
+      "content": "<|reserved_token_26|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156922": {
+      "content": "<|reserved_token_27|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156923": {
+      "content": "<|reserved_token_28|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156924": {
+      "content": "<|reserved_token_29|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156925": {
+      "content": "<|reserved_token_30|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156926": {
+      "content": "<|reserved_token_31|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156927": {
+      "content": "<|reserved_token_32|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156928": {
+      "content": "<|reserved_token_33|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156929": {
+      "content": "<|reserved_token_34|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156930": {
+      "content": "<|reserved_token_35|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156931": {
+      "content": "<|reserved_token_36|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156932": {
+      "content": "<|reserved_token_37|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156933": {
+      "content": "<|reserved_token_38|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156934": {
+      "content": "<|reserved_token_39|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156935": {
+      "content": "<|reserved_token_40|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156936": {
+      "content": "<|reserved_token_41|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156937": {
+      "content": "<|reserved_token_42|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156938": {
+      "content": "<|reserved_token_43|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156939": {
+      "content": "<|reserved_token_44|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156940": {
+      "content": "<|reserved_token_45|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156941": {
+      "content": "<|reserved_token_46|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156942": {
+      "content": "<|reserved_token_47|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156943": {
+      "content": "<|reserved_token_48|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156944": {
+      "content": "<|reserved_token_49|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156945": {
+      "content": "<|reserved_token_50|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156946": {
+      "content": "<|reserved_token_51|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156947": {
+      "content": "<|reserved_token_52|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156948": {
+      "content": "<|reserved_token_53|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156949": {
+      "content": "<|reserved_token_54|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156950": {
+      "content": "<|reserved_token_55|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156951": {
+      "content": "<|reserved_token_56|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156952": {
+      "content": "<|reserved_token_57|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156953": {
+      "content": "<|reserved_token_58|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156954": {
+      "content": "<|reserved_token_59|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156955": {
+      "content": "<|reserved_token_60|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156956": {
+      "content": "<|reserved_token_61|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156957": {
+      "content": "<|reserved_token_62|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156958": {
+      "content": "<|reserved_token_63|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156959": {
+      "content": "<|reserved_token_64|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156960": {
+      "content": "<|reserved_token_65|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156961": {
+      "content": "<|reserved_token_66|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156962": {
+      "content": "<|reserved_token_67|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156963": {
+      "content": "<|reserved_token_68|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156964": {
+      "content": "<|reserved_token_69|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156965": {
+      "content": "<|reserved_token_70|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156966": {
+      "content": "<|reserved_token_71|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156967": {
+      "content": "<|reserved_token_72|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156968": {
+      "content": "<|reserved_token_73|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156969": {
+      "content": "<|reserved_token_74|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156970": {
+      "content": "<|reserved_token_75|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156971": {
+      "content": "<|reserved_token_76|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156972": {
+      "content": "<|reserved_token_77|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156973": {
+      "content": "<|reserved_token_78|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156974": {
+      "content": "<|reserved_token_79|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156975": {
+      "content": "<|reserved_token_80|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156976": {
+      "content": "<|reserved_token_81|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156977": {
+      "content": "<|reserved_token_82|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156978": {
+      "content": "<|reserved_token_83|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156979": {
+      "content": "<|reserved_token_84|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156980": {
+      "content": "<|reserved_token_85|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156981": {
+      "content": "<|reserved_token_86|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156982": {
+      "content": "<|reserved_token_87|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156983": {
+      "content": "<|reserved_token_88|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156984": {
+      "content": "<|reserved_token_89|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156985": {
+      "content": "<|reserved_token_90|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156986": {
+      "content": "<|reserved_token_91|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156987": {
+      "content": "<|reserved_token_92|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156988": {
+      "content": "<|reserved_token_93|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156989": {
+      "content": "<|reserved_token_94|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156990": {
+      "content": "<|reserved_token_95|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156991": {
+      "content": "<|reserved_token_96|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156992": {
+      "content": "<|reserved_token_97|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156993": {
+      "content": "<|reserved_token_98|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156994": {
+      "content": "<|reserved_token_99|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156995": {
+      "content": "<|reserved_token_100|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156996": {
+      "content": "<|reserved_token_101|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156997": {
+      "content": "<|reserved_token_102|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156998": {
+      "content": "<|reserved_token_103|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "156999": {
+      "content": "<|reserved_token_104|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157000": {
+      "content": "<|reserved_token_105|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157001": {
+      "content": "<|reserved_token_106|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157002": {
+      "content": "<|reserved_token_107|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157003": {
+      "content": "<|reserved_token_108|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157004": {
+      "content": "<|reserved_token_109|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157005": {
+      "content": "<|reserved_token_110|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157006": {
+      "content": "<|reserved_token_111|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157007": {
+      "content": "<|reserved_token_112|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157008": {
+      "content": "<|reserved_token_113|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157009": {
+      "content": "<|reserved_token_114|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157010": {
+      "content": "<|reserved_token_115|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157011": {
+      "content": "<|reserved_token_116|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157012": {
+      "content": "<|reserved_token_117|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157013": {
+      "content": "<|reserved_token_118|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157014": {
+      "content": "<|reserved_token_119|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157015": {
+      "content": "<|reserved_token_120|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157016": {
+      "content": "<|reserved_token_121|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157017": {
+      "content": "<|reserved_token_122|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157018": {
+      "content": "<|reserved_token_123|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157019": {
+      "content": "<|reserved_token_124|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157020": {
+      "content": "<|reserved_token_125|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157021": {
+      "content": "<|reserved_token_126|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157022": {
+      "content": "<|reserved_token_127|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157023": {
+      "content": "<|reserved_token_128|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157024": {
+      "content": "<|reserved_token_129|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157025": {
+      "content": "<|reserved_token_130|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157026": {
+      "content": "<|reserved_token_131|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157027": {
+      "content": "<|reserved_token_132|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157028": {
+      "content": "<|reserved_token_133|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157029": {
+      "content": "<|reserved_token_134|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157030": {
+      "content": "<|reserved_token_135|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157031": {
+      "content": "<|reserved_token_136|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157032": {
+      "content": "<|reserved_token_137|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157033": {
+      "content": "<|reserved_token_138|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157034": {
+      "content": "<|reserved_token_139|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157035": {
+      "content": "<|reserved_token_140|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157036": {
+      "content": "<|reserved_token_141|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157037": {
+      "content": "<|reserved_token_142|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157038": {
+      "content": "<|reserved_token_143|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157039": {
+      "content": "<|reserved_token_144|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157040": {
+      "content": "<|reserved_token_145|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157041": {
+      "content": "<|reserved_token_146|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157042": {
+      "content": "<|reserved_token_147|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157043": {
+      "content": "<|reserved_token_148|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157044": {
+      "content": "<|reserved_token_149|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157045": {
+      "content": "<|reserved_token_150|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157046": {
+      "content": "<|reserved_token_151|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157047": {
+      "content": "<|reserved_token_152|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157048": {
+      "content": "<|reserved_token_153|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157049": {
+      "content": "<|reserved_token_154|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157050": {
+      "content": "<|reserved_token_155|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157051": {
+      "content": "<|reserved_token_156|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157052": {
+      "content": "<|reserved_token_157|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157053": {
+      "content": "<|reserved_token_158|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157054": {
+      "content": "<|reserved_token_159|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157055": {
+      "content": "<|reserved_token_160|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157056": {
+      "content": "<|reserved_token_161|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157057": {
+      "content": "<|reserved_token_162|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157058": {
+      "content": "<|reserved_token_163|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157059": {
+      "content": "<|reserved_token_164|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157060": {
+      "content": "<|reserved_token_165|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157061": {
+      "content": "<|reserved_token_166|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157062": {
+      "content": "<|reserved_token_167|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157063": {
+      "content": "<|reserved_token_168|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157064": {
+      "content": "<|reserved_token_169|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157065": {
+      "content": "<|reserved_token_170|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157066": {
+      "content": "<|reserved_token_171|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157067": {
+      "content": "<|reserved_token_172|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157068": {
+      "content": "<|reserved_token_173|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157069": {
+      "content": "<|reserved_token_174|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157070": {
+      "content": "<|reserved_token_175|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157071": {
+      "content": "<|reserved_token_176|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157072": {
+      "content": "<|reserved_token_177|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157073": {
+      "content": "<|reserved_token_178|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157074": {
+      "content": "<|reserved_token_179|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157075": {
+      "content": "<|reserved_token_180|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157076": {
+      "content": "<|reserved_token_181|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157077": {
+      "content": "<|reserved_token_182|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157078": {
+      "content": "<|reserved_token_183|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157079": {
+      "content": "<|reserved_token_184|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157080": {
+      "content": "<|reserved_token_185|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157081": {
+      "content": "<|reserved_token_186|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157082": {
+      "content": "<|reserved_token_187|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157083": {
+      "content": "<|reserved_token_188|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157084": {
+      "content": "<|reserved_token_189|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157085": {
+      "content": "<|reserved_token_190|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157086": {
+      "content": "<|reserved_token_191|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157087": {
+      "content": "<|reserved_token_192|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157088": {
+      "content": "<|reserved_token_193|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157089": {
+      "content": "<|reserved_token_194|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157090": {
+      "content": "<|reserved_token_195|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157091": {
+      "content": "<|reserved_token_196|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157092": {
+      "content": "<|reserved_token_197|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157093": {
+      "content": "<|reserved_token_198|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157094": {
+      "content": "<|reserved_token_199|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157095": {
+      "content": "<|reserved_token_200|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157096": {
+      "content": "<|reserved_token_201|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157097": {
+      "content": "<|reserved_token_202|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157098": {
+      "content": "<|reserved_token_203|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157099": {
+      "content": "<|reserved_token_204|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157100": {
+      "content": "<|reserved_token_205|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157101": {
+      "content": "<|reserved_token_206|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157102": {
+      "content": "<|reserved_token_207|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157103": {
+      "content": "<|reserved_token_208|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157104": {
+      "content": "<|reserved_token_209|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157105": {
+      "content": "<|reserved_token_210|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157106": {
+      "content": "<|reserved_token_211|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157107": {
+      "content": "<|reserved_token_212|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157108": {
+      "content": "<|reserved_token_213|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157109": {
+      "content": "<|reserved_token_214|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157110": {
+      "content": "<|reserved_token_215|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157111": {
+      "content": "<|reserved_token_216|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157112": {
+      "content": "<|reserved_token_217|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157113": {
+      "content": "<|reserved_token_218|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157114": {
+      "content": "<|reserved_token_219|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157115": {
+      "content": "<|reserved_token_220|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157116": {
+      "content": "<|reserved_token_221|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157117": {
+      "content": "<|reserved_token_222|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157118": {
+      "content": "<|reserved_token_223|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157119": {
+      "content": "<|reserved_token_224|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157120": {
+      "content": "<|reserved_token_225|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157121": {
+      "content": "<|reserved_token_226|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157122": {
+      "content": "<|reserved_token_227|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157123": {
+      "content": "<|reserved_token_228|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157124": {
+      "content": "<|reserved_token_229|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157125": {
+      "content": "<|reserved_token_230|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157126": {
+      "content": "<|reserved_token_231|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157127": {
+      "content": "<|reserved_token_232|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157128": {
+      "content": "<|reserved_token_233|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157129": {
+      "content": "<|reserved_token_234|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157130": {
+      "content": "<|reserved_token_235|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157131": {
+      "content": "<|reserved_token_236|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157132": {
+      "content": "<|reserved_token_237|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157133": {
+      "content": "<|reserved_token_238|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157134": {
+      "content": "<|reserved_token_239|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157135": {
+      "content": "<|reserved_token_240|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157136": {
+      "content": "<|reserved_token_241|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157137": {
+      "content": "<|reserved_token_242|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157138": {
+      "content": "<|reserved_token_243|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157139": {
+      "content": "<|reserved_token_244|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157140": {
+      "content": "<|reserved_token_245|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157141": {
+      "content": "<|reserved_token_246|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157142": {
+      "content": "<|reserved_token_247|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157143": {
+      "content": "<|reserved_token_248|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157144": {
+      "content": "<|reserved_token_249|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157145": {
+      "content": "<|reserved_token_250|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157146": {
+      "content": "<|reserved_token_251|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157147": {
+      "content": "<|reserved_token_252|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157148": {
+      "content": "<|reserved_token_253|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157149": {
+      "content": "<|reserved_token_254|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157150": {
+      "content": "<|reserved_token_255|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157151": {
+      "content": "<role>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "157152": {
+      "content": "</role>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|startoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "[CLS]",
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "fast_tokenizer": true,
+  "gmask_token": "[gMASK]",
+  "mask_token": "<|mask|>",
+  "merges_file": null,
+  "model_max_length": 32768,
+  "pad_token": "<|endoftext|>",
+  "padding_side": "right",
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "trust_remote_code": true
+}

models/dreamreasoner-8b/.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

models/dreamreasoner-8b/README.md

@@ -0,0 +1,3 @@
+---
+license: apache-2.0
+---

models/dreamreasoner-8b/added_tokens.json

@@ -0,0 +1,29 @@
+{
+  "</think>": 151668,
+  "</tool_call>": 151658,
+  "</tool_response>": 151666,
+  "<think>": 151667,
+  "<tool_call>": 151657,
+  "<tool_response>": 151665,
+  "<|MASK|>": 151669,
+  "<|box_end|>": 151649,
+  "<|box_start|>": 151648,
+  "<|endoftext|>": 151643,
+  "<|file_sep|>": 151664,
+  "<|fim_middle|>": 151660,
+  "<|fim_pad|>": 151662,
+  "<|fim_prefix|>": 151659,
+  "<|fim_suffix|>": 151661,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644,
+  "<|image_pad|>": 151655,
+  "<|object_ref_end|>": 151647,
+  "<|object_ref_start|>": 151646,
+  "<|quad_end|>": 151651,
+  "<|quad_start|>": 151650,
+  "<|repo_name|>": 151663,
+  "<|video_pad|>": 151656,
+  "<|vision_end|>": 151653,
+  "<|vision_pad|>": 151654,
+  "<|vision_start|>": 151652
+}

models/dreamreasoner-8b/chat_template.jinja

@@ -0,0 +1,87 @@
+{%- if tools %}
+    {{- '<|im_start|>system\n' }}
+    {%- if messages[0].role == 'system' %}
+        {{- messages[0].content + '\n\n' }}
+    {%- endif %}
+    {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
+        {%- set ns.multi_step_tool = false %}
+        {%- set ns.last_query_index = index %}
+    {%- endif %}
+{%- endfor %}
+{%- for message in messages %}
+    {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set content = message.content %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is defined and message.reasoning_content is not none %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in message.content %}
+                {%- set content = message.content.split('</think>')[-1].lstrip('\n') %}
+                {%- set reasoning_content = message.content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {%- if loop.last or (not loop.last and reasoning_content) %}
+                {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
+            {%- else %}
+                {{- '<|im_start|>' + message.role + '\n' + content }}
+            {%- endif %}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if (loop.first and content) or (not loop.first) %}
+                    {{- '\n' }}
+                {%- endif %}
+                {%- if tool_call.function %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {{- '<tool_call>\n{"name": "' }}
+                {{- tool_call.name }}
+                {{- '", "arguments": ' }}
+                {%- if tool_call.arguments is string %}
+                    {{- tool_call.arguments }}
+                {%- else %}
+                    {{- tool_call.arguments | tojson }}
+                {%- endif %}
+                {{- '}\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>' }}
+        {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is false %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- else %}
+        {{- '<think>\n' }}
+    {%- endif %}
+{%- endif %}

models/dreamreasoner-8b/config.json

@@ -0,0 +1,38 @@
+{
+  "architectures": [
+    "DreamForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_dream.DreamConfig",
+    "AutoModel": "modeling_dream.DreamForCausalLM",
+    "AutoModelForCausalLM": "modeling_dream.DreamForCausalLM"
+  },
+  "block_size": 32,
+  "bos_token_id": 151643,
+  "dtype": "bfloat16",
+  "eos_token_id": 151643,
+  "head_dim": 128,
+  "hidden_act": "silu",
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "intermediate_size": 12288,
+  "mask_token_id": 151669,
+  "max_position_embeddings": 32768,
+  "max_window_layers": 36,
+  "model_type": "Dream",
+  "num_attention_heads": 32,
+  "num_hidden_layers": 36,
+  "num_key_value_heads": 8,
+  "rms_norm_eps": 1e-06,
+  "rope_scaling": null,
+  "rope_theta": 1000000,
+  "sliding_window": null,
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "5.7.0",
+  "use_cache": false,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

models/dreamreasoner-8b/configuration_dream.py

@@ -0,0 +1,223 @@
+# coding=utf-8
+# Copyright 2024 The Dream team, HKUNLP Group and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT and Qwen implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT and Qwen used by the Meta AI and Qwen team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Dream model configuration"""
+
+from typing import Optional
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class DreamConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`DreamModel`]. It is used to instantiate a
+    Dream model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of
+    Dream-8B [Qwen/Dream-8B](https://huggingface.co/Qwen/Dream-8B).
+
+    Configuration objects inherit from [`PreTrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PreTrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 151936):
+            Vocabulary size of the Dream model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`DreamModel`]
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 12288):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 36):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        num_key_value_heads (`int`, *optional*, defaults to 8):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
+            converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
+            by meanpooling all the original heads within that group. For more details, check out [this
+            paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `32`.
+        head_dim (`int`, *optional*, defaults to 128):
+            The attention head dimension.
+        hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
+            The non-linear activation function (function or string) in the decoder.
+        max_position_embeddings (`int`, *optional*, defaults to 32768):
+            The maximum sequence length that this model might ever be used with.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        rms_norm_eps (`float`, *optional*, defaults to 1e-06):
+            The epsilon used by the rms normalization layers.
+        use_cache (`bool`, *optional*, defaults to `False`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied.
+        rope_theta (`float`, *optional*, defaults to 1000000.0):
+            The base period of the RoPE embeddings.
+        rope_scaling (`Dict`, *optional*):
+            Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
+            and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
+            accordingly.
+            Expected contents:
+                `rope_type` (`str`):
+                    The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
+                    'llama3'], with 'default' being the original RoPE implementation.
+                `factor` (`float`, *optional*):
+                    Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
+                    most scaling types, a `factor` of x will enable the model to handle sequences of length x *
+                    original maximum pre-trained length.
+                `original_max_position_embeddings` (`int`, *optional*):
+                    Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
+                    pretraining.
+                `attention_factor` (`float`, *optional*):
+                    Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
+                    computation. If unspecified, it defaults to value recommended by the implementation, using the
+                    `factor` field to infer the suggested value.
+                `beta_fast` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 32.
+                `beta_slow` (`float`, *optional*):
+                    Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
+                    ramp function. If unspecified, it defaults to 1.
+                `short_factor` (`List[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to short contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `long_factor` (`List[float]`, *optional*):
+                    Only used with 'longrope'. The scaling factor to be applied to long contexts (<
+                    `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
+                    size divided by the number of attention heads divided by 2
+                `low_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
+                `high_freq_factor` (`float`, *optional*):
+                    Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
+        attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
+            Whether to use a bias in the query, key, value and output projection layers during self-attention.
+        use_sliding_window (`bool`, *optional*, defaults to `False`):
+            Whether to use sliding window attention.
+        sliding_window (`int`, *optional*, defaults to 4096):
+            Sliding window attention (SWA) window size. If not specified, will default to `4096`.
+        max_window_layers (`int`, *optional*, defaults to 36):
+            The number of layers using full attention. The first `max_window_layers` layers will use full attention, while any
+            additional layer afterwards will use SWA (Sliding Window Attention).
+        layer_types (`list`, *optional*):
+            Attention pattern for each layer.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+
+    ```python
+    >>> from transformers import DreamModel, DreamConfig
+
+    >>> # Initializing a Dream style configuration
+    >>> configuration = DreamConfig()
+
+    >>> # Initializing a model from the Dream-8B style configuration
+    >>> model = DreamModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "Dream"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    # Default tensor parallel plan for base model `Dream`
+    base_model_tp_plan = {
+        "layers.*.self_attn.q_proj": "colwise",
+        "layers.*.self_attn.k_proj": "colwise",
+        "layers.*.self_attn.v_proj": "colwise",
+        "layers.*.self_attn.o_proj": "rowwise",
+        "layers.*.mlp.gate_proj": "colwise",
+        "layers.*.mlp.up_proj": "colwise",
+        "layers.*.mlp.down_proj": "rowwise",
+    }
+    base_model_pp_plan = {
+        "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+        "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+        "norm": (["hidden_states"], ["hidden_states"]),
+    }
+
+    def __init__(
+        self,
+        vocab_size: Optional[int] = 151936,
+        hidden_size: Optional[int] = 4096,
+        intermediate_size: Optional[int] = 12288,
+        num_hidden_layers: Optional[int] = 36,
+        num_attention_heads: Optional[int] = 32,
+        num_key_value_heads: Optional[int] = 8,
+        head_dim: Optional[int] = 128,
+        hidden_act: Optional[str] = "silu",
+        max_position_embeddings: Optional[int] = 32768,
+        initializer_range: Optional[float] = 0.02,
+        rms_norm_eps: Optional[float] = 1e-6,
+        use_cache: Optional[bool] = False,
+        tie_word_embeddings: Optional[bool] = False,
+        rope_theta: Optional[float] = 1000000.0,
+        rope_scaling: Optional[dict] = None,
+        attention_bias: Optional[bool] = False,
+        use_sliding_window: Optional[bool] = False,
+        sliding_window: Optional[int] = 4096,
+        max_window_layers: Optional[int] = 36,
+        layer_types: Optional[list[str]] = None,
+        attention_dropout: Optional[float] = 0.0,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.max_position_embeddings = max_position_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.use_sliding_window = use_sliding_window
+        self.sliding_window = sliding_window if self.use_sliding_window else None
+        self.max_window_layers = max_window_layers
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.use_cache = use_cache
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.attention_bias = attention_bias
+        self.attention_dropout = attention_dropout
+        self.layer_types = layer_types
+
+        if self.rope_scaling is not None and "type" in self.rope_scaling:
+            self.rope_scaling["rope_type"] = self.rope_scaling["type"]
+        self.validate_rope()
+
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+
+__all__ = ["DreamConfig"]

models/dreamreasoner-8b/generation_config.json

@@ -0,0 +1,7 @@
+{
+  "bos_token_id": 151643,
+  "eos_token_id": 151643,
+  "max_new_tokens": 2048,
+  "transformers_version": "5.7.0",
+  "trust_remote_code": true
+}

models/dreamreasoner-8b/generation_utils.py

@@ -0,0 +1,469 @@
+
+"""Block diffusion generation utilities for Dream models."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Optional, Sequence, Union
+
+import torch
+from torch.nn import functional as F
+from transformers.cache_utils import DynamicCache
+from transformers.utils import ModelOutput
+
+
+def top_k_logits(logits: torch.Tensor, k: int) -> torch.Tensor:
+    if k <= 0:
+        return logits
+    values, _ = torch.topk(logits, k)
+    min_values = values[..., -1, None]
+    return torch.where(logits < min_values, torch.full_like(logits, float('-inf')), logits)
+
+
+def top_p_logits(logits: torch.Tensor, p: float) -> torch.Tensor:
+    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+    sorted_mask = cumulative_probs > p
+    sorted_mask[..., 1:] = sorted_mask[..., :-1].clone()
+    sorted_mask[..., 0] = False
+    mask_indices = torch.scatter(
+        torch.full_like(logits, False, dtype=torch.bool),
+        -1,
+        sorted_indices,
+        sorted_mask,
+    )
+    return logits.masked_fill(mask_indices, float('-inf'))
+
+
+def sample_with_temperature_topk_topp(
+    logits: torch.Tensor,
+    temperature: float = 1.0,
+    top_k: int = 0,
+    top_p: float = 1.0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    orig_shape = logits.shape[:-1]
+    vocab_size = logits.shape[-1]
+    logits = logits.reshape(-1, vocab_size)
+
+    if temperature > 0:
+        logits = logits / temperature
+    if top_k > 0:
+        logits = top_k_logits(logits, top_k)
+    if top_p < 1.0:
+        logits = top_p_logits(logits, top_p)
+
+    probs = F.softmax(logits, dim=-1)
+    if temperature > 0:
+        token = torch.multinomial(probs, num_samples=1)
+    else:
+        token = probs.argmax(dim=-1, keepdim=True)
+    token_prob = torch.gather(probs, -1, token)
+    return token.view(*orig_shape), token_prob.view(*orig_shape)
+
+
+def get_num_transfer_tokens(block_length: int, steps: int) -> torch.Tensor:
+    base = block_length // steps
+    remainder = block_length % steps
+    num_transfer_tokens = torch.zeros(steps, dtype=torch.int64) + base
+    num_transfer_tokens[:remainder] += 1
+    return num_transfer_tokens
+
+
+def build_block_diffusion_attention_mask(
+    num_blocks: int,
+    block_length: int,
+    device: torch.device,
+    batch_size: int = 1,
+) -> torch.Tensor:
+    block_mask = torch.tril(torch.ones(num_blocks, num_blocks, device=device))
+    return block_mask.repeat_interleave(block_length, dim=0).repeat_interleave(block_length, dim=1).unsqueeze(0).expand(
+        batch_size, -1, -1
+    )
+
+
+def _resolve_stopping_ids(stopping_criteria_idx: Optional[Union[int, Sequence[int]]]) -> Optional[List[int]]:
+    if stopping_criteria_idx is None:
+        return None
+    if isinstance(stopping_criteria_idx, int):
+        return [stopping_criteria_idx]
+    return list(stopping_criteria_idx)
+
+
+def _should_stop(
+    generated_ids: torch.Tensor,
+    prompt_length: int,
+    stopping_criteria_idx: Optional[List[int]],
+) -> bool:
+    if not stopping_criteria_idx:
+        return False
+    gen_part = generated_ids[:, prompt_length:]
+    return any((gen_part == stop_id).any().item() for stop_id in stopping_criteria_idx)
+
+
+def _default_use_kv_cache(model: torch.nn.Module) -> bool:
+    """Dream / Dream1 models use prefix KV cache during block diffusion decode."""
+    model_type = getattr(model.config, 'model_type', None)
+    if model_type is None:
+        return False
+    return model_type.lower() in ('dream', 'dream1')
+
+
+def _select_transfer_index(
+    remasking_strategy: str,
+    mask_index: torch.Tensor,
+    x0: torch.Tensor,
+    x0_p: torch.Tensor,
+    num_transfer_tokens: torch.Tensor,
+    step: int,
+    confidence_threshold: float,
+    eb_threshold: Optional[float],
+    *,
+    force_accept: bool = False,
+) -> torch.Tensor:
+    if force_accept:
+        return mask_index.clone()
+
+    if remasking_strategy == 'sequential':
+        transfer_index = torch.zeros_like(x0, dtype=torch.bool)
+        for j in range(x0.shape[0]):
+            if not mask_index[j].any():
+                continue
+            first_mask_index = mask_index[j].nonzero(as_tuple=True)[0].min().item()
+            end = first_mask_index + int(num_transfer_tokens[step].item())
+            transfer_index[j, first_mask_index:end] = True
+        return transfer_index
+
+    if remasking_strategy == 'low_confidence_static':
+        confidence = torch.where(mask_index, x0_p, -torch.inf)
+        transfer_index = torch.zeros_like(x0, dtype=torch.bool)
+        k = max(1, int(num_transfer_tokens[step].item()))
+        for j in range(confidence.shape[0]):
+            _, idx = torch.topk(confidence[j], k)
+            transfer_index[j, idx] = True
+        return transfer_index
+
+    if remasking_strategy == 'low_confidence_dynamic':
+        confidence = torch.where(mask_index, x0_p, -torch.inf)
+        transfer_index = torch.zeros_like(x0, dtype=torch.bool)
+        k = max(1, int(num_transfer_tokens[step].item()))
+        for j in range(confidence.shape[0]):
+            high_conf_mask = confidence[j] > confidence_threshold
+            if int(high_conf_mask.sum().item()) >= k:
+                transfer_index[j] = high_conf_mask
+            else:
+                _, idx = torch.topk(confidence[j], k)
+                transfer_index[j, idx] = True
+        return transfer_index
+
+    if remasking_strategy == 'entropy_bounded':
+        if eb_threshold is None:
+            raise ValueError('eb_threshold is required for entropy_bounded remasking.')
+        eps = 1e-12
+        entropies = -(x0_p.clamp_min(eps) * x0_p.clamp_min(eps).log())
+        entropies = torch.where(mask_index, entropies, torch.inf)
+        ent_sorted, order = torch.sort(entropies, dim=1, descending=False)
+        cumsum = torch.cumsum(ent_sorted, dim=1)
+        transfer_index = torch.zeros_like(x0, dtype=torch.bool)
+        for j in range(x0_p.shape[0]):
+            k = torch.searchsorted(
+                cumsum[j], torch.tensor(eb_threshold, device=x0_p.device), right=False
+            ).item()
+            k = max(1, min(k, int(mask_index[j].sum().item())))
+            transfer_index[j, order[j, :k]] = True
+        return transfer_index
+
+    raise ValueError(f'Unknown remasking strategy: {remasking_strategy}')
+
+
+def _denoise_current_block(
+    model: torch.nn.Module,
+    x: torch.Tensor,
+    num_block: int,
+    block_length: int,
+    mask_id: int,
+    block_diffusion_attention_mask: torch.Tensor,
+    position_ids: torch.Tensor,
+    denoising_steps: int,
+    num_transfer_tokens: torch.Tensor,
+    temperature: float,
+    top_k: int,
+    top_p: float,
+    remasking_strategy: str,
+    confidence_threshold: float,
+    eb_threshold: Optional[float],
+    *,
+    use_kv_cache: bool,
+    past_key_values: Optional[DynamicCache],
+) -> tuple[torch.Tensor, Optional[DynamicCache], int]:
+    block_start = num_block * block_length
+    block_end = block_start + block_length
+    cur_x = x[:, block_start:block_end].clone()
+    nfe = 0
+
+    for step in range(denoising_steps + 1):
+        mask_index = cur_x == mask_id
+        if mask_index.sum() == 0:
+            if use_kv_cache:
+                cur_attn_mask = block_diffusion_attention_mask[:, block_start:block_end, :block_end]
+                cur_position_ids = position_ids[:, block_start:block_end]
+                model(
+                    cur_x,
+                    attention_mask=cur_attn_mask,
+                    position_ids=cur_position_ids,
+                    past_key_values=past_key_values,
+                    use_cache=True,
+                    store_kv=True,
+                )
+                nfe += 1
+            break
+
+        force_accept = step == denoising_steps - 1
+        if use_kv_cache:
+            cur_attn_mask = block_diffusion_attention_mask[:, block_start:block_end, :block_end]
+            cur_position_ids = position_ids[:, block_start:block_end]
+            logits = model(
+                cur_x,
+                attention_mask=cur_attn_mask,
+                position_ids=cur_position_ids,
+                past_key_values=past_key_values,
+                use_cache=True,
+                store_kv=False,
+            ).logits
+        else:
+            seq_end = block_end
+            attn_mask = block_diffusion_attention_mask[:, :seq_end, :seq_end]
+            pos_ids = position_ids[:, :seq_end]
+            logits = model(
+                x[:, :seq_end],
+                attention_mask=attn_mask,
+                position_ids=pos_ids,
+                use_cache=False,
+            ).logits[:, block_start:block_end]
+
+        nfe += 1
+        x0, x0_p = sample_with_temperature_topk_topp(
+            logits,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+        )
+        x0 = torch.where(mask_index, x0, cur_x)
+        transfer_index = _select_transfer_index(
+            remasking_strategy,
+            mask_index,
+            x0,
+            x0_p,
+            num_transfer_tokens,
+            step,
+            confidence_threshold,
+            eb_threshold,
+            force_accept=force_accept,
+        )
+        cur_x[transfer_index] = x0[transfer_index]
+        if not use_kv_cache:
+            x[:, block_start:block_end] = cur_x
+
+    return cur_x, past_key_values, nfe
+
+
+@dataclass
+class BlockDiffusionOutput(ModelOutput):
+    sequences: torch.LongTensor = None
+    nfe: Optional[int] = None
+    logits: Optional[tuple] = None
+
+
+@torch.no_grad()
+def block_diffusion_generate(
+    model: torch.nn.Module,
+    input_ids: torch.LongTensor,
+    mask_id: int,
+    gen_length: int = 128,
+    block_length: Optional[int] = None,
+    denoising_steps: Optional[int] = None,
+    temperature: float = 0.0,
+    top_k: int = 0,
+    top_p: float = 1.0,
+    remasking_strategy: str = 'low_confidence_dynamic',
+    confidence_threshold: float = 0.9,
+    eb_threshold: Optional[float] = 0.35,
+    stopping_criteria_idx: Optional[Union[int, Sequence[int]]] = None,
+    use_kv_cache: Optional[bool] = None,
+    return_dict_in_generate: bool = False,
+) -> Union[torch.LongTensor, BlockDiffusionOutput]:
+    """Block-wise diffusion decoding with optional prefix KV cache."""
+    model.eval()
+    if input_ids.dim() != 2:
+        raise ValueError(f'input_ids must be 2D, got shape {tuple(input_ids.shape)}')
+
+    device = input_ids.device
+    batch_size, prompt_length = input_ids.shape
+    block_length = block_length or getattr(model.config, 'block_size', 4)
+    if denoising_steps is None:
+        denoising_steps = 1 if remasking_strategy == 'low_confidence_static' else block_length
+    stopping_criteria_idx = _resolve_stopping_ids(stopping_criteria_idx)
+
+    if use_kv_cache is None:
+        use_kv_cache = _default_use_kv_cache(model)
+
+    num_blocks = (prompt_length + gen_length + block_length - 1) // block_length
+    total_length = num_blocks * block_length
+    block_diffusion_attention_mask = build_block_diffusion_attention_mask(
+        num_blocks, block_length, device, batch_size=batch_size
+    )
+    position_ids = torch.arange(total_length, device=device, dtype=torch.long).unsqueeze(0).expand(batch_size, -1)
+
+    x = torch.full((batch_size, total_length), mask_id, dtype=input_ids.dtype, device=device)
+    x[:, :prompt_length] = input_ids
+
+    prefill_blocks = prompt_length // block_length
+    prefill_length = prefill_blocks * block_length
+    past_key_values = DynamicCache() if use_kv_cache else None
+    nfe = 0
+
+    if use_kv_cache and prefill_length > 0:
+        cur_x = x[:, :prefill_length]
+        cur_attn_mask = block_diffusion_attention_mask[:, :prefill_length, :prefill_length]
+        cur_position_ids = position_ids[:, :prefill_length]
+        model(
+            cur_x,
+            attention_mask=cur_attn_mask,
+            position_ids=cur_position_ids,
+            past_key_values=past_key_values,
+            use_cache=True,
+            store_kv=True,
+        )
+        nfe += 1
+
+    num_transfer_tokens = get_num_transfer_tokens(block_length, denoising_steps)
+
+    for num_block in range(prefill_blocks, num_blocks):
+        cur_x, past_key_values, block_nfe = _denoise_current_block(
+            model,
+            x,
+            num_block,
+            block_length,
+            mask_id,
+            block_diffusion_attention_mask,
+            position_ids,
+            denoising_steps,
+            num_transfer_tokens,
+            temperature,
+            top_k,
+            top_p,
+            remasking_strategy,
+            confidence_threshold,
+            eb_threshold,
+            use_kv_cache=use_kv_cache,
+            past_key_values=past_key_values,
+        )
+        nfe += block_nfe
+        x[:, num_block * block_length:(num_block + 1) * block_length] = cur_x
+
+        if _should_stop(x, prompt_length, stopping_criteria_idx):
+            break
+
+    output_length = min(total_length, prompt_length + gen_length)
+    x = x[:, :output_length]
+
+    if return_dict_in_generate:
+        return BlockDiffusionOutput(sequences=x, nfe=nfe)
+    return x
+
+
+# HF kwargs that ``generate()`` strips before calling ``block_diffusion_generate()``.
+_UNSUPPORTED_HF_KEYS = (
+    'stopping_criteria', 'num_return_sequences', 'num_beams', 'num_beam_groups',
+    'penalty_alpha', 'use_cache', 'output_logits', 'output_scores', 'output_attentions',
+    'output_hidden_states', 'return_legacy_cache', 'synced_gpus', 'streamer',
+    'logits_processor', 'logits_warper', 'generation_config', 'tokenizer',
+    'min_length', 'min_new_tokens', 'pad_token_id', 'bos_token_id', 'eos_token_id',
+)
+
+
+class BlockDiffusionGenerationMixin:
+
+    def _resolve_generation_mode(self, generation_mode: Optional[str] = None) -> str:
+        if generation_mode is not None:
+            return generation_mode
+        return getattr(self.config, 'generation_mode', 'block_diffusion')
+
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        generation_mode: Optional[str] = None,
+        **kwargs,
+    ) -> Union[torch.LongTensor, BlockDiffusionOutput, ModelOutput]:
+        """HF-compatible ``generate`` entry point with pluggable decoding modes.
+
+        Supported modes (``generation_mode`` kwarg or ``config.generation_mode``):
+        - ``block_diffusion`` (default): calls :meth:`block_diffusion_generate`.
+        - ``autoregressive``: delegates to ``GenerationMixin.generate``.
+        """
+        mode = self._resolve_generation_mode(kwargs.pop('generation_mode', generation_mode))
+        return_dict_in_generate = kwargs.pop('return_dict_in_generate', False)
+
+        if mode == 'autoregressive':
+            return super().generate(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                return_dict_in_generate=return_dict_in_generate,
+                **kwargs,
+            )
+        if mode != 'block_diffusion':
+            raise ValueError(f'Unknown generation_mode: {mode!r}. Supported: block_diffusion, autoregressive.')
+
+        # HF-only normalizations; sampling params pass through to block_diffusion_generate.
+        if kwargs.pop('do_sample', None) is False:
+            kwargs['temperature'] = 0.0
+        if 'max_new_tokens' not in kwargs and (max_length := kwargs.pop('max_length', None)) is not None:
+            kwargs['max_new_tokens'] = max(max_length - input_ids.shape[-1], 0)
+        for key in _UNSUPPORTED_HF_KEYS:
+            kwargs.pop(key, None)
+
+        return self.block_diffusion_generate(
+            input_ids=input_ids,
+            return_dict_in_generate=return_dict_in_generate,
+            **kwargs,
+        )
+
+    @torch.no_grad()
+    def block_diffusion_generate(
+        self,
+        input_ids: torch.LongTensor,
+        max_new_tokens: int = 128,
+        temperature: float = 0.0,
+        top_k: int = 0,
+        top_p: float = 1.0,
+        return_dict_in_generate: bool = False,
+        # block-diffusion specific
+        block_length: Optional[int] = None,
+        denoising_steps: Optional[int] = None,
+        remasking_strategy: str = 'low_confidence_dynamic',
+        confidence_threshold: float = 0.9,
+        eb_threshold: Optional[float] = 0.35,
+        use_kv_cache: Optional[bool] = None,
+        mask_token_id: Optional[int] = None,
+    ) -> Union[torch.LongTensor, BlockDiffusionOutput]:
+        mask_token_id = mask_token_id if mask_token_id is not None else self.config.mask_token_id
+        if mask_token_id is None:
+            raise ValueError('mask_token_id must be provided or set in model.config.mask_token_id')
+        stopping_criteria_idx = self.config.eos_token_id if getattr(self.config, 'eos_token_id', None) is not None else None
+        return block_diffusion_generate(
+            self,
+            input_ids=input_ids,
+            mask_id=mask_token_id,
+            gen_length=max_new_tokens,
+            block_length=block_length,
+            denoising_steps=denoising_steps,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            remasking_strategy=remasking_strategy,
+            confidence_threshold=confidence_threshold,
+            eb_threshold=eb_threshold,
+            stopping_criteria_idx=stopping_criteria_idx,
+            use_kv_cache=use_kv_cache,
+            return_dict_in_generate=return_dict_in_generate,
+        )

models/dreamreasoner-8b/model.safetensors.index.json

@@ -0,0 +1,406 @@
+{
+  "metadata": {
+    "total_size": 16381470720
+  },
+  "weight_map": {
+    "lm_head.weight": "model-00004-of-00004.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.21.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.21.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.k_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.22.self_attn.q_norm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.26.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.28.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.29.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.30.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.30.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.31.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.32.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.32.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.32.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.32.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.32.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.32.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.32.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.32.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.32.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.32.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.32.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.33.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.33.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.33.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.33.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.33.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.33.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.33.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.33.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.33.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.33.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.33.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.34.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.34.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.34.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.34.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.34.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.34.self_attn.k_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.34.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.34.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.34.self_attn.q_norm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.34.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.34.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.35.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.35.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.35.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.35.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.35.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.35.self_attn.k_norm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.35.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.35.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.35.self_attn.q_norm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.35.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.35.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.k_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.9.self_attn.q_norm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.norm.weight": "model-00004-of-00004.safetensors"
+  }
+}

models/dreamreasoner-8b/modeling_dream.py

@@ -0,0 +1,574 @@
+# coding=utf-8
+# Copyright 2026 The Dream team, HKUNLP Group and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT and Qwen implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT and Qwen used by the Meta AI and Qwen team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from typing import Callable, Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from transformers.activations import ACT2FN
+from transformers.cache_utils import Cache, DynamicCache, StaticCache
+from transformers.generation import GenerationMixin
+from transformers.integrations import use_kernel_forward_from_hub
+from transformers.modeling_attn_mask_utils import AttentionMaskConverter
+from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
+from transformers.modeling_layers import GradientCheckpointingLayer
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPast,
+    CausalLMOutputWithPast,
+    QuestionAnsweringModelOutput,
+    SequenceClassifierOutputWithPast,
+    TokenClassifierOutput,
+)
+from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from transformers.processing_utils import Unpack
+from transformers.utils import auto_docstring, can_return_tuple, logging, TransformersKwargs
+from .configuration_dream import DreamConfig
+from .generation_utils import BlockDiffusionGenerationMixin
+
+@use_kernel_forward_from_hub("RMSNorm")
+class DreamRMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+        """
+        DreamRMSNorm is equivalent to T5LayerNorm
+        """
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(input_dtype)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+class DreamMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, x):
+        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+
+class DreamRotaryEmbedding(nn.Module):
+    inv_freq: torch.Tensor  # fix linting for `register_buffer`
+
+    def __init__(self, config: DreamConfig, device=None):
+        super().__init__()
+        # BC: "rope_type" was originally "type"
+        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
+            self.rope_type = config.rope_scaling.get(
+                "rope_type", config.rope_scaling.get("type"))
+        else:
+            self.rope_type = "default"
+        self.max_seq_len_cached = config.max_position_embeddings
+        self.original_max_seq_len = config.max_position_embeddings
+
+        self.config = config
+        rope_init_fn: Callable = self.compute_default_rope_parameters
+        if self.rope_type != "default":
+            rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+
+        inv_freq, self.attention_scaling = rope_init_fn(
+            self.config, device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self.original_inv_freq = self.inv_freq
+
+    @staticmethod
+    def compute_default_rope_parameters(
+        config: Optional[DreamConfig] = None,
+        device: Optional["torch.device"] = None,
+        seq_len: Optional[int] = None,
+    ) -> tuple["torch.Tensor", float]:
+        """
+        Computes the inverse frequencies according to the original RoPE implementation
+        Args:
+            config ([`~transformers.PreTrainedConfig`]):
+                The model configuration.
+            device (`torch.device`):
+                The device to use for initialization of the inverse frequencies.
+            seq_len (`int`, *optional*):
+                The current sequence length. Unused for this type of RoPE.
+        Returns:
+            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+        """
+        base = config.rope_theta
+        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+        attention_factor = 1.0  # Unused in this type of RoPE
+
+        # Compute the inverse frequencies
+        inv_freq = 1.0 / (
+            base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+        )
+        return inv_freq, attention_factor
+
+    @torch.no_grad()
+    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
+    def forward(self, x, position_ids):
+        inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+        position_ids_expanded = position_ids[:, None, :].float()
+
+        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
+            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+            emb = torch.cat((freqs, freqs), dim=-1)
+            cos = emb.cos() * self.attention_scaling
+            sin = emb.sin() * self.attention_scaling
+
+        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Applies Rotary Position Embedding to the query and key tensors.
+
+    Args:
+        q (`torch.Tensor`): The query tensor.
+        k (`torch.Tensor`): The key tensor.
+        cos (`torch.Tensor`): The cosine part of the rotary embedding.
+        sin (`torch.Tensor`): The sine part of the rotary embedding.
+        position_ids (`torch.Tensor`, *optional*):
+            Deprecated and unused.
+        unsqueeze_dim (`int`, *optional*, defaults to 1):
+            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+    Returns:
+        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+    """
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+def _prepare_attention_mask_for_backend(
+    attention_mask: torch.Tensor,
+    query: torch.Tensor,
+    key_len: int,
+    attn_implementation: str,
+) -> torch.Tensor:
+    """Convert block-diffusion 0/1 masks for SDPA (bool) or eager (additive float)."""
+    if attention_mask.dim() == 2:
+        attention_mask = attention_mask.unsqueeze(0)
+    if attention_mask.dim() == 3:
+        attention_mask = attention_mask[:, None, :, :key_len]
+    else:
+        attention_mask = attention_mask[:, :, :, :key_len]
+
+    if attn_implementation == "sdpa":
+        if attention_mask.dtype != torch.bool:
+            attention_mask = attention_mask != 0
+    else:
+        min_value = torch.finfo(query.dtype).min
+        attention_mask = torch.where(
+            attention_mask.to(dtype=query.dtype) > 0,
+            torch.zeros((), dtype=query.dtype, device=query.device),
+            min_value,
+        )
+    return attention_mask
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[FlashAttentionKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+        attn_weights = attn_weights + causal_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+
+class DreamAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(self, config: DreamConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = True
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(
+            config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+        )
+        self.rotary_fn = apply_rotary_pos_emb
+        self.q_norm = DreamRMSNorm(self.head_dim, eps=config.rms_norm_eps)  # unlike olmo, only on the head dim!
+        self.k_norm = DreamRMSNorm(self.head_dim, eps=config.rms_norm_eps)  # thus post q_norm does not need reshape
+        self.sliding_window = config.sliding_window
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor],
+        past_key_values: Optional[Cache] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        if past_key_values is not None and kwargs.get("store_kv", False):
+            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)
+        elif past_key_values is not None and not kwargs.get("store_kv", False) and len(past_key_values) > self.layer_idx:
+            cache_layer = past_key_values.layers[self.layer_idx]
+            if cache_layer.is_initialized and cache_layer.get_seq_length() > 0:
+                key_states = torch.cat([cache_layer.keys, key_states], dim=-2)
+                value_states = torch.cat([cache_layer.values, value_states], dim=-2)
+
+        attention_interface: Callable = eager_attention_forward
+        attn_implementation = self.config._attn_implementation
+        if attn_implementation != "eager":
+            attention_interface = ALL_ATTENTION_FUNCTIONS[attn_implementation]
+
+        prepared_attention_mask = attention_mask
+        if attention_mask is not None:
+            prepared_attention_mask = _prepare_attention_mask_for_backend(
+                attention_mask,
+                query_states,
+                key_states.shape[-2],
+                attn_implementation,
+            )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            prepared_attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            sliding_window=self.sliding_window,  # diff with Llama
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights
+
+class DreamDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: DreamConfig, layer_idx: int):
+        super().__init__()
+        self.hidden_size = config.hidden_size
+
+        self.self_attn = DreamAttention(config=config, layer_idx=layer_idx)
+
+        self.mlp = DreamMLP(config)
+        self.input_layernorm = DreamRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = DreamRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        use_cache: Optional[bool] = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        store_kv: Optional[bool] = False,
+        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        # Self Attention
+        hidden_states, _ = self.self_attn(
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            store_kv=store_kv,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+        hidden_states = residual + hidden_states
+
+        # Fully Connected
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+@auto_docstring
+class DreamPreTrainedModel(PreTrainedModel):
+    config_class = DreamConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["DreamDecoderLayer"]
+    _skip_keys_device_placement = ["past_key_values"]
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+    _can_compile_fullgraph = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": DreamDecoderLayer,
+        "attentions": DreamAttention,
+    }
+
+
+@auto_docstring
+class DreamModel(DreamPreTrainedModel):
+    def __init__(self, config: DreamConfig):
+        super().__init__(config)
+        self.padding_idx = getattr(config, "pad_token_id", None) or config.eos_token_id
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList(
+            [DreamDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+        )
+        self.norm = DreamRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.rotary_emb = DreamRotaryEmbedding(config=config)
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        store_kv: Optional[bool] = None,
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> BaseModelOutputWithPast:
+        r"""
+        cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+            Indices depicting the position of the present input tokens in the KV cache. Used for position
+            embeddings and attention masking during generation.
+        store_kv (`bool`, *optional*):
+            Whether to store key-value states in the cache during this forward pass. When `True`, the computed
+            key and value states are updated in `past_key_values`.
+        """
+        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:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        if cache_position is None:
+            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(
+                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
+            )
+
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+
+        # We skip causal masking creation in Dream models
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+        for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+            hidden_states = decoder_layer(
+                hidden_states,
+                attention_mask=attention_mask,
+                position_embeddings=position_embeddings,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                cache_position=cache_position,
+                store_kv=store_kv,
+                **kwargs,
+            )
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=past_key_values if use_cache else None,
+        )
+
+@auto_docstring
+class DreamForCausalLM(DreamPreTrainedModel, BlockDiffusionGenerationMixin, GenerationMixin):
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+    _tp_plan = {"lm_head": "colwise_rep"}
+    _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = DreamModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> CausalLMOutputWithPast:
+        r"""
+        cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
+            Indices depicting the position of the present input tokens in the KV cache. Used for position
+            embeddings and attention masking during generation.
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
+            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
+
+        Example:
+
+        ```python
+        >>> from transformers import AutoTokenizer, DreamForCausalLM
+
+        >>> model = DreamForCausalLM.from_pretrained("Dream-org/Dream-Reasoner-8B")
+        >>> tokenizer = AutoTokenizer.from_pretrained("Dream-org/Dream-Reasoner-8B")
+
+        >>> prompt = "Hey, are you conscious? Can you talk to me?"
+        >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+        >>> # Generate
+        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+        ```"""
+        outputs: BaseModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            **kwargs,
+        )
+
+        hidden_states = outputs.last_hidden_state
+        # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+__all__ = [
+    "DreamForCausalLM",
+    "DreamModel",
+    "DreamPreTrainedModel",
+]

models/dreamreasoner-8b/special_tokens_map.json

@@ -0,0 +1,39 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>",
+    "<|MASK|>"
+  ],
+  "eos_token": {
+    "content": "<|im_end|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "<|MASK|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

models/dreamreasoner-8b/tokenization_dream.py

@@ -0,0 +1,342 @@
+# coding=utf-8
+# Copyright 2024 The Dream team, HKUNLP Group and The HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on Qwen's implementations in this library.
+# 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.
+"""Tokenization classes for DreamReasoner."""
+
+import json
+import os
+import unicodedata
+from functools import lru_cache
+from typing import Optional, Tuple
+
+import regex as re
+
+from transformers.tokenization_utils import AddedToken, PreTrainedTokenizer
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {
+    "vocab_file": "vocab.json",
+    "merges_file": "merges.txt",
+}
+
+
+MAX_MODEL_INPUT_SIZES = {"dream/dream-tokenizer": 32768}
+
+PRETOKENIZE_REGEX = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
+
+
+@lru_cache()
+# Copied from transformers.models.gpt2.tokenization_gpt2.bytes_to_unicode
+def bytes_to_unicode():
+    """
+    Returns list of utf-8 byte and a mapping to unicode strings. We specifically avoids mapping to whitespace/control
+    characters the bpe code barfs on.
+
+    The reversible bpe codes work on unicode strings. This means you need a large # of unicode characters in your vocab
+    if you want to avoid UNKs. When you're at something like a 10B token dataset you end up needing around 5K for
+    decent coverage. This is a significant percentage of your normal, say, 32K bpe vocab. To avoid that, we want lookup
+    tables between utf-8 bytes and unicode strings.
+    """
+    bs = (
+        list(range(ord("!"), ord("~") + 1)) + list(range(ord("¡"), ord("¬") + 1)) + list(range(ord("®"), ord("ÿ") + 1))
+    )
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8 + n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+# Copied from transformers.models.gpt2.tokenization_gpt2.get_pairs
+def get_pairs(word):
+    """
+    Return set of symbol pairs in a word.
+
+    Word is represented as tuple of symbols (symbols being variable-length strings).
+    """
+    pairs = set()
+    prev_char = word[0]
+    for char in word[1:]:
+        pairs.add((prev_char, char))
+        prev_char = char
+    return pairs
+
+
+class DreamTokenizer(PreTrainedTokenizer):
+    """
+    Construct a Dream tokenizer. Based on byte-level Byte-Pair-Encoding.
+
+    Same with GPT2Tokenizer, this tokenizer has been trained to treat spaces like parts of the tokens so a word will
+    be encoded differently whether it is at the beginning of the sentence (without space) or not:
+
+    ```python
+    >>> from transformers import AutoTokenizer
+
+    >>> tokenizer = AutoTokenizer.from_pretrained("Dream-org/Dream-v0-Base-7B", trust_remote_code=True)
+    >>> tokenizer("Hello world")["input_ids"]
+    [9707, 1879]
+
+    >>> tokenizer(" Hello world")["input_ids"]
+    [21927, 1879]
+    ```
+    This is expected.
+
+    You should not use GPT2Tokenizer instead, because of the different pretokenization rules.
+
+    This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+    this superclass for more information regarding those methods.
+
+    Args:
+        vocab_file (`str`):
+            Path to the vocabulary file.
+        merges_file (`str`):
+            Path to the merges file.
+        errors (`str`, *optional*, defaults to `"replace"`):
+            Paradigm to follow when decoding bytes to UTF-8. See
+            [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
+        unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+            token instead.
+        bos_token (`str`, *optional*):
+            The beginning of sequence token. Not applicable for this tokenizer.
+        eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The end of sequence token.
+        pad_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+            The token used for padding, for example when batching sequences of different lengths.
+        clean_up_tokenization_spaces (`bool`, *optional*, defaults to `False`):
+            Whether or not the model should cleanup the spaces that were added when splitting the input text during the
+            tokenization process. Not applicable to this tokenizer, since tokenization does not add spaces.
+        split_special_tokens (`bool`, *optional*, defaults to `False`):
+            Whether or not the special tokens should be split during the tokenization process. The default behavior is
+            to not split special tokens. This means that if `<|endoftext|>` is the `eos_token`, then `tokenizer.tokenize("<|endoftext|>") =
+            ['<|endoftext|>`]. Otherwise, if `split_special_tokens=True`, then `tokenizer.tokenize("<|endoftext|>")` will be give `['<',
+            '|', 'endo', 'ft', 'ext', '|', '>']`. This argument is only supported for `slow` tokenizers for the moment.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file,
+        merges_file,
+        errors="replace",
+        unk_token="<|endoftext|>",
+        bos_token=None,
+        eos_token="<|endoftext|>",
+        pad_token="<|endoftext|>",
+        clean_up_tokenization_spaces=False,
+        split_special_tokens=False,
+        **kwargs,
+    ):
+        # Dream vocab does not contain control tokens; added tokens need to be special
+        bos_token = (
+            AddedToken(bos_token, lstrip=False, rstrip=False, special=True, normalized=False)
+            if isinstance(bos_token, str)
+            else bos_token
+        )
+        eos_token = (
+            AddedToken(eos_token, lstrip=False, rstrip=False, special=True, normalized=False)
+            if isinstance(eos_token, str)
+            else eos_token
+        )
+        unk_token = (
+            AddedToken(unk_token, lstrip=False, rstrip=False, special=True, normalized=False)
+            if isinstance(unk_token, str)
+            else unk_token
+        )
+        pad_token = (
+            AddedToken(pad_token, lstrip=False, rstrip=False, special=True, normalized=False)
+            if isinstance(pad_token, str)
+            else pad_token
+        )
+
+        with open(vocab_file, encoding="utf-8") as vocab_handle:
+            self.encoder = json.load(vocab_handle)
+        self.decoder = {v: k for k, v in self.encoder.items()}
+        self.errors = errors  # how to handle errors in decoding
+        self.byte_encoder = bytes_to_unicode()
+        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
+        bpe_merges = []
+        with open(merges_file, encoding="utf-8") as merges_handle:
+            for i, line in enumerate(merges_handle):
+                line = line.strip()
+                if (i == 0 and line.startswith("#version:")) or not line:
+                    continue
+                bpe_merges.append(tuple(line.split()))
+        self.bpe_ranks = dict(zip(bpe_merges, range(len(bpe_merges))))
+        # NOTE: the cache can grow without bound and will get really large for long running processes
+        # (esp. for texts of language that do not use space between word, e.g. Chinese); technically
+        # not a memory leak but appears as one.
+        # GPT2Tokenizer has the same problem, so let's be consistent.
+        self.cache = {}
+
+        self.pat = re.compile(PRETOKENIZE_REGEX)
+
+        if kwargs.get("add_prefix_space", False):
+            logger.warning_once(
+                f"{self.__class__.__name} does not support `add_prefix_space`, setting it to True has no effect."
+            )
+
+        super().__init__(
+            errors=errors,
+            bos_token=bos_token,
+            eos_token=eos_token,
+            pad_token=pad_token,
+            unk_token=unk_token,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            split_special_tokens=split_special_tokens,
+            **kwargs,
+        )
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.encoder)
+
+    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.get_vocab
+    def get_vocab(self):
+        return dict(self.encoder, **self.added_tokens_encoder)
+
+    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.bpe
+    def bpe(self, token):
+        if token in self.cache:
+            return self.cache[token]
+        word = tuple(token)
+        pairs = get_pairs(word)
+
+        if not pairs:
+            return token
+
+        while True:
+            bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
+            if bigram not in self.bpe_ranks:
+                break
+            first, second = bigram
+            new_word = []
+            i = 0
+            while i < len(word):
+                try:
+                    j = word.index(first, i)
+                except ValueError:
+                    new_word.extend(word[i:])
+                    break
+                else:
+                    new_word.extend(word[i:j])
+                    i = j
+
+                if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
+                    new_word.append(first + second)
+                    i += 2
+                else:
+                    new_word.append(word[i])
+                    i += 1
+            new_word = tuple(new_word)
+            word = new_word
+            if len(word) == 1:
+                break
+            else:
+                pairs = get_pairs(word)
+        word = " ".join(word)
+        self.cache[token] = word
+        return word
+
+    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._tokenize
+    def _tokenize(self, text):
+        """Tokenize a string."""
+        bpe_tokens = []
+        for token in re.findall(self.pat, text):
+            token = "".join(
+                self.byte_encoder[b] for b in token.encode("utf-8")
+            )  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
+            bpe_tokens.extend(bpe_token for bpe_token in self.bpe(token).split(" "))
+        return bpe_tokens
+
+    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_token_to_id
+    def _convert_token_to_id(self, token):
+        """Converts a token (str) in an id using the vocab."""
+        return self.encoder.get(token, self.encoder.get(self.unk_token))
+
+    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer._convert_id_to_token
+    def _convert_id_to_token(self, index):
+        """Converts an index (integer) in a token (str) using the vocab."""
+        return self.decoder.get(index)
+
+    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.convert_tokens_to_string
+    def convert_tokens_to_string(self, tokens):
+        """Converts a sequence of tokens (string) in a single string."""
+        text = "".join(tokens)
+        text = bytearray([self.byte_decoder[c] for c in text]).decode("utf-8", errors=self.errors)
+        return text
+
+    def decode(
+        self,
+        token_ids,
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: Optional[bool] = False,
+        spaces_between_special_tokens: bool = False,
+        **kwargs,
+    ) -> str:
+        # `spaces_between_special_tokens` defaults to True for _decode in slow tokenizers
+        # and cannot be configured elsewhere, but it should default to False for DreamTokenizer
+        return super().decode(
+            token_ids,
+            skip_special_tokens=skip_special_tokens,
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            spaces_between_special_tokens=spaces_between_special_tokens,
+            **kwargs,
+        )
+
+    # Copied from transformers.models.gpt2.tokenization_gpt2.GPT2Tokenizer.save_vocabulary
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        if not os.path.isdir(save_directory):
+            logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+            return
+        vocab_file = os.path.join(
+            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+        )
+        merge_file = os.path.join(
+            save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
+        )
+
+        with open(vocab_file, "w", encoding="utf-8") as f:
+            f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
+
+        index = 0
+        with open(merge_file, "w", encoding="utf-8") as writer:
+            writer.write("#version: 0.2\n")
+            for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
+                if index != token_index:
+                    logger.warning(
+                        f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
+                        " Please check that the tokenizer is not corrupted!"
+                    )
+                    index = token_index
+                writer.write(" ".join(bpe_tokens) + "\n")
+                index += 1
+
+        return vocab_file, merge_file
+
+    def prepare_for_tokenization(self, text, **kwargs):
+        text = unicodedata.normalize("NFC", text)
+        return (text, kwargs)
+
+__all__ = ["DreamTokenizer"]

models/dreamreasoner-8b/tokenizer_config.json

@@ -0,0 +1,256 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151669": {
+      "content": "<|MASK|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>",
+    "<|MASK|>"
+  ],
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_dream.DreamTokenizer",
+      null
+    ]
+  },
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "mask_token": "<|MASK|>",
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "padding_side": "right",
+  "split_special_tokens": false,
+  "tokenizer_class": "DreamTokenizer",
+  "unk_token": null
+}

runs/basecode/bd-base-sweep_678190.out

@@ -0,0 +1,5 @@
+[info] started_at=2026-07-09T01:28:25+08:00
+[info] node=xcnc17
+[info] output=/home/l/liyj/shiying/hku_diffusion_dllm/runs/basecode/sweep_678190.jsonl
+[info] latest=/home/l/liyj/shiying/hku_diffusion_dllm/runs/basecode/latest_sweep.jsonl
+[info] finished_at=2026-07-09T01:28:26+08:00

runs/basecode/latest_sweep.jsonl

@@ -0,0 +1,12 @@
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10], "text": "3 4 5 6 7 8 9 10", "nfe": 6, "elapsed_s": 0.000769949983805418, "tokens_per_forward": 1.3333333333333333, "metadata": {"method": "confidence", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 8, "num_blocks": 1, "steps": 6, "remask_ratio": 0.5, "use_cache": false, "draft_width": 4}}, "tokens_per_second": 10390.285301988866, "sweep": {"method": "confidence", "steps": 6, "block_size": 8, "num_blocks": 1, "remask_ratio": 0.5, "use_cache": false}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10], "text": "3 4 5 6 7 8 9 10", "nfe": 4, "elapsed_s": 0.0005228999652899802, "tokens_per_forward": 2.0, "metadata": {"method": "dmax", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 8, "num_blocks": 1, "steps": 4, "remask_ratio": 0.25, "use_cache": false, "draft_width": 4}}, "tokens_per_second": 15299.293423290836, "sweep": {"method": "dmax", "steps": 4, "block_size": 8, "num_blocks": 1, "remask_ratio": 0.25, "use_cache": false}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10], "text": "3 4 5 6 7 8 9 10", "nfe": 6, "elapsed_s": 0.0007856350275687873, "tokens_per_forward": 1.3333333333333333, "metadata": {"method": "speculative", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 8, "num_blocks": 1, "steps": 6, "remask_ratio": 0.5, "use_cache": true, "draft_width": 4}}, "tokens_per_second": 10182.845366195881, "sweep": {"method": "speculative", "steps": 6, "block_size": 8, "num_blocks": 1, "remask_ratio": 0.5, "use_cache": true}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18", "nfe": 6, "elapsed_s": 0.001491333998274058, "tokens_per_forward": 2.6666666666666665, "metadata": {"method": "confidence", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 16, "num_blocks": 1, "steps": 6, "remask_ratio": 0.5, "use_cache": false, "draft_width": 4}}, "tokens_per_second": 10728.64966433879, "sweep": {"method": "confidence", "steps": 6, "block_size": 16, "num_blocks": 1, "remask_ratio": 0.5, "use_cache": false}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18", "nfe": 4, "elapsed_s": 0.0010131489834748209, "tokens_per_forward": 4.0, "metadata": {"method": "dmax", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 16, "num_blocks": 1, "steps": 4, "remask_ratio": 0.25, "use_cache": false, "draft_width": 4}}, "tokens_per_second": 15792.346694288162, "sweep": {"method": "dmax", "steps": 4, "block_size": 16, "num_blocks": 1, "remask_ratio": 0.25, "use_cache": false}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18", "nfe": 6, "elapsed_s": 0.0014925010036677122, "tokens_per_forward": 2.6666666666666665, "metadata": {"method": "speculative", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 16, "num_blocks": 1, "steps": 6, "remask_ratio": 0.5, "use_cache": true, "draft_width": 4}}, "tokens_per_second": 10720.260797601588, "sweep": {"method": "speculative", "steps": 6, "block_size": 16, "num_blocks": 1, "remask_ratio": 0.5, "use_cache": true}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34", "nfe": 6, "elapsed_s": 0.0029353260179050267, "tokens_per_forward": 5.333333333333333, "metadata": {"method": "confidence", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 32, "num_blocks": 1, "steps": 6, "remask_ratio": 0.5, "use_cache": false, "draft_width": 4}}, "tokens_per_second": 10901.68512962616, "sweep": {"method": "confidence", "steps": 6, "block_size": 32, "num_blocks": 1, "remask_ratio": 0.5, "use_cache": false}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34", "nfe": 4, "elapsed_s": 0.0019988870481029153, "tokens_per_forward": 8.0, "metadata": {"method": "dmax", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 32, "num_blocks": 1, "steps": 4, "remask_ratio": 0.25, "use_cache": false, "draft_width": 4}}, "tokens_per_second": 16008.908572583056, "sweep": {"method": "dmax", "steps": 4, "block_size": 32, "num_blocks": 1, "remask_ratio": 0.25, "use_cache": false}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34", "nfe": 6, "elapsed_s": 0.0029457149794325233, "tokens_per_forward": 5.333333333333333, "metadata": {"method": "speculative", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 32, "num_blocks": 1, "steps": 6, "remask_ratio": 0.5, "use_cache": true, "draft_width": 4}}, "tokens_per_second": 10863.237014928252, "sweep": {"method": "speculative", "steps": 6, "block_size": 32, "num_blocks": 1, "remask_ratio": 0.5, "use_cache": true}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34", "nfe": 8, "elapsed_s": 0.00392380595440045, "tokens_per_forward": 4.0, "metadata": {"method": "multiblock", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 16, "num_blocks": 2, "steps": 8, "remask_ratio": 0.5, "use_cache": true, "draft_width": 4}}, "tokens_per_second": 8155.347224577404, "sweep": {"method": "multiblock", "steps": 8, "block_size": 16, "num_blocks": 2, "remask_ratio": 0.5, "use_cache": true}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66", "nfe": 8, "elapsed_s": 0.008637683000415564, "tokens_per_forward": 8.0, "metadata": {"method": "multiblock", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 16, "num_blocks": 4, "steps": 8, "remask_ratio": 0.5, "use_cache": true, "draft_width": 4}}, "tokens_per_second": 7409.394393950429, "sweep": {"method": "multiblock", "steps": 8, "block_size": 16, "num_blocks": 4, "remask_ratio": 0.5, "use_cache": true}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 3, 4, 5], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 3 4 5", "nfe": 8, "elapsed_s": 0.01631542103132233, "tokens_per_forward": 16.0, "metadata": {"method": "multiblock", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 16, "num_blocks": 8, "steps": 8, "remask_ratio": 0.5, "use_cache": true, "draft_width": 4}}, "tokens_per_second": 7845.338453372777, "sweep": {"method": "multiblock", "steps": 8, "block_size": 16, "num_blocks": 8, "remask_ratio": 0.5, "use_cache": true}}

runs/basecode/sweep_678190.jsonl

@@ -0,0 +1,12 @@
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10], "text": "3 4 5 6 7 8 9 10", "nfe": 6, "elapsed_s": 0.000769949983805418, "tokens_per_forward": 1.3333333333333333, "metadata": {"method": "confidence", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 8, "num_blocks": 1, "steps": 6, "remask_ratio": 0.5, "use_cache": false, "draft_width": 4}}, "tokens_per_second": 10390.285301988866, "sweep": {"method": "confidence", "steps": 6, "block_size": 8, "num_blocks": 1, "remask_ratio": 0.5, "use_cache": false}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10], "text": "3 4 5 6 7 8 9 10", "nfe": 4, "elapsed_s": 0.0005228999652899802, "tokens_per_forward": 2.0, "metadata": {"method": "dmax", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 8, "num_blocks": 1, "steps": 4, "remask_ratio": 0.25, "use_cache": false, "draft_width": 4}}, "tokens_per_second": 15299.293423290836, "sweep": {"method": "dmax", "steps": 4, "block_size": 8, "num_blocks": 1, "remask_ratio": 0.25, "use_cache": false}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10], "text": "3 4 5 6 7 8 9 10", "nfe": 6, "elapsed_s": 0.0007856350275687873, "tokens_per_forward": 1.3333333333333333, "metadata": {"method": "speculative", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 8, "num_blocks": 1, "steps": 6, "remask_ratio": 0.5, "use_cache": true, "draft_width": 4}}, "tokens_per_second": 10182.845366195881, "sweep": {"method": "speculative", "steps": 6, "block_size": 8, "num_blocks": 1, "remask_ratio": 0.5, "use_cache": true}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18", "nfe": 6, "elapsed_s": 0.001491333998274058, "tokens_per_forward": 2.6666666666666665, "metadata": {"method": "confidence", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 16, "num_blocks": 1, "steps": 6, "remask_ratio": 0.5, "use_cache": false, "draft_width": 4}}, "tokens_per_second": 10728.64966433879, "sweep": {"method": "confidence", "steps": 6, "block_size": 16, "num_blocks": 1, "remask_ratio": 0.5, "use_cache": false}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18", "nfe": 4, "elapsed_s": 0.0010131489834748209, "tokens_per_forward": 4.0, "metadata": {"method": "dmax", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 16, "num_blocks": 1, "steps": 4, "remask_ratio": 0.25, "use_cache": false, "draft_width": 4}}, "tokens_per_second": 15792.346694288162, "sweep": {"method": "dmax", "steps": 4, "block_size": 16, "num_blocks": 1, "remask_ratio": 0.25, "use_cache": false}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18", "nfe": 6, "elapsed_s": 0.0014925010036677122, "tokens_per_forward": 2.6666666666666665, "metadata": {"method": "speculative", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 16, "num_blocks": 1, "steps": 6, "remask_ratio": 0.5, "use_cache": true, "draft_width": 4}}, "tokens_per_second": 10720.260797601588, "sweep": {"method": "speculative", "steps": 6, "block_size": 16, "num_blocks": 1, "remask_ratio": 0.5, "use_cache": true}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34", "nfe": 6, "elapsed_s": 0.0029353260179050267, "tokens_per_forward": 5.333333333333333, "metadata": {"method": "confidence", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 32, "num_blocks": 1, "steps": 6, "remask_ratio": 0.5, "use_cache": false, "draft_width": 4}}, "tokens_per_second": 10901.68512962616, "sweep": {"method": "confidence", "steps": 6, "block_size": 32, "num_blocks": 1, "remask_ratio": 0.5, "use_cache": false}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34", "nfe": 4, "elapsed_s": 0.0019988870481029153, "tokens_per_forward": 8.0, "metadata": {"method": "dmax", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 32, "num_blocks": 1, "steps": 4, "remask_ratio": 0.25, "use_cache": false, "draft_width": 4}}, "tokens_per_second": 16008.908572583056, "sweep": {"method": "dmax", "steps": 4, "block_size": 32, "num_blocks": 1, "remask_ratio": 0.25, "use_cache": false}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34", "nfe": 6, "elapsed_s": 0.0029457149794325233, "tokens_per_forward": 5.333333333333333, "metadata": {"method": "speculative", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 32, "num_blocks": 1, "steps": 6, "remask_ratio": 0.5, "use_cache": true, "draft_width": 4}}, "tokens_per_second": 10863.237014928252, "sweep": {"method": "speculative", "steps": 6, "block_size": 32, "num_blocks": 1, "remask_ratio": 0.5, "use_cache": true}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34", "nfe": 8, "elapsed_s": 0.00392380595440045, "tokens_per_forward": 4.0, "metadata": {"method": "multiblock", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 16, "num_blocks": 2, "steps": 8, "remask_ratio": 0.5, "use_cache": true, "draft_width": 4}}, "tokens_per_second": 8155.347224577404, "sweep": {"method": "multiblock", "steps": 8, "block_size": 16, "num_blocks": 2, "remask_ratio": 0.5, "use_cache": true}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66", "nfe": 8, "elapsed_s": 0.008637683000415564, "tokens_per_forward": 8.0, "metadata": {"method": "multiblock", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 16, "num_blocks": 4, "steps": 8, "remask_ratio": 0.5, "use_cache": true, "draft_width": 4}}, "tokens_per_second": 7409.394393950429, "sweep": {"method": "multiblock", "steps": 8, "block_size": 16, "num_blocks": 4, "remask_ratio": 0.5, "use_cache": true}}
+{"tokens": [3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, 127, 3, 4, 5], "text": "3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 3 4 5", "nfe": 8, "elapsed_s": 0.01631542103132233, "tokens_per_forward": 16.0, "metadata": {"method": "multiblock", "config": {"vocab_size": 128, "mask_token_id": 0, "eos_token_id": 2, "block_size": 16, "num_blocks": 8, "steps": 8, "remask_ratio": 0.5, "use_cache": true, "draft_width": 4}}, "tokens_per_second": 7845.338453372777, "sweep": {"method": "multiblock", "steps": 8, "block_size": 16, "num_blocks": 8, "remask_ratio": 0.5, "use_cache": true}}

runs/elf_timestep_alignment/elf-a40_679746.err

@@ -0,0 +1,46 @@
+2026-07-10 18:01:48.806902: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
+E0000 00:00:1783677708.952181 1362203 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+E0000 00:00:1783677708.996564 1362203 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+W0000 00:00:1783677709.354654 1362203 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783677709.354699 1362203 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783677709.354702 1362203 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783677709.354704 1362203 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+Failed to get TPU metadata (tpu-env) from instance metadata for variable CHIPS_PER_HOST_BOUNDS: DEADLINE_EXCEEDED: Timeout was reached
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/gcp_metadata_utils.cc:84
+learning/45eac/tfrc/runtime/env_var_utils.cc:50
+
+Failed to get TPU metadata (tpu-env) from instance metadata for variable HOST_BOUNDS: DEADLINE_EXCEEDED: Timeout was reached
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/gcp_metadata_utils.cc:84
+learning/45eac/tfrc/runtime/env_var_utils.cc:50
+
+Failed to get TPU metadata (tpu-env) from instance metadata for variable ALT: DEADLINE_EXCEEDED: Timeout was reached
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/gcp_metadata_utils.cc:84
+learning/45eac/tfrc/runtime/env_var_utils.cc:50
+
+Failed to get TPU metadata (tpu-env) from instance metadata for variable WRAP: DEADLINE_EXCEEDED: Timeout was reached
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/gcp_metadata_utils.cc:84
+learning/45eac/tfrc/runtime/env_var_utils.cc:50
+
+Failed to get TPU metadata (accelerator-type) from instance metadata for variable TPU_ACCELERATOR_TYPE: DEADLINE_EXCEEDED: Timeout was reached
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/gcp_metadata_utils.cc:84
+
+Failed to find host bounds for accelerator type: WARNING: could not determine TPU accelerator type, please set env var `TPU_ACCELERATOR_TYPE` manually, otherwise libtpu.so may not properly initialize.
+Failed to get TPU metadata (agent-worker-number) from instance metadata for variable TPU_WORKER_ID: DEADLINE_EXCEEDED: Timeout was reached
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/gcp_metadata_utils.cc:84
+
+Failed to get TPU metadata (worker-network-endpoints) from instance metadata for variable TPU_WORKER_HOSTNAMES: DEADLINE_EXCEEDED: Timeout was reached
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/gcp_metadata_utils.cc:84
+
+WARNING: Logging before InitGoogle() is written to STDERR
+E0000 00:00:1783680350.013226 1362203 common_lib.cc:520] INVALID_ARGUMENT: Error: unexpected worker hostname 'WARNING: could not determine TPU worker hostnames or IP addresses' from env var TPU_WORKER_HOSTNAMES. Expecting a valid hostname or IP address without port number. (Full TPU workers' addr string: WARNING: could not determine TPU worker hostnames or IP addresses, please set env var `TPU_WORKER_HOSTNAMES` manually, otherwise libtpu.so may not properly initialize.)
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/libtpu_init_utils.cc:175
+[2026-07-10T18:53:38.016] error: *** JOB 679746 ON xgph11 CANCELLED AT 2026-07-10T18:53:38 DUE to SIGNAL Terminated ***

runs/elf_timestep_alignment/elf-a40_679746.out

@@ -0,0 +1,34 @@
+[info] started_at=2026-07-10T18:01:43+08:00
+[info] node=xgph11
+[info] run_dir=/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_final_layer_t005_095_679746
+Fri Jul 10 18:01:43 2026       
++-----------------------------------------------------------------------------------------+
+| NVIDIA-SMI 575.57.08              Driver Version: 575.57.08      CUDA Version: 12.9     |
+|-----------------------------------------+------------------------+----------------------+
+| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
+| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
+|                                         |                        |               MIG M. |
+|=========================================+========================+======================|
+|   0  NVIDIA A100 80GB PCIe          On  |   00000000:98:00.0 Off |                   On |
+| N/A   65C    P0            199W /  300W |   38634MiB /  81920MiB |     N/A      Default |
+|                                         |                        |              Enabled |
++-----------------------------------------+------------------------+----------------------+
+
++-----------------------------------------------------------------------------------------+
+| MIG devices:                                                                            |
++------------------+----------------------------------+-----------+-----------------------+
+| GPU  GI  CI  MIG |                     Memory-Usage |        Vol|        Shared         |
+|      ID  ID  Dev |                       BAR1-Usage | SM     Unc| CE ENC  DEC  OFA  JPG |
+|                  |                                  |        ECC|                       |
+|==================+==================================+===========+=======================|
+|  0    2   0   0  |             107MiB / 40192MiB    | 42      0 |  3   0    2    0    0 |
+|                  |                 0MiB / 65535MiB  |           |                       |
++------------------+----------------------------------+-----------+-----------------------+
+                                                                                         
++-----------------------------------------------------------------------------------------+
+| Processes:                                                                              |
+|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
+|        ID   ID                                                               Usage      |
+|=========================================================================================|
+|  No running processes found                                                             |
++-----------------------------------------------------------------------------------------+

runs/elf_timestep_alignment/elf-a80_679745.err

@@ -0,0 +1,51 @@
+Traceback (most recent call last):
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/.venv-elf-py312/lib/python3.12/site-packages/numpy/core/__init__.py", line 24, in <module>
+    from . import multiarray
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/.venv-elf-py312/lib/python3.12/site-packages/numpy/core/multiarray.py", line 10, in <module>
+    from . import overrides
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/.venv-elf-py312/lib/python3.12/site-packages/numpy/core/overrides.py", line 8, in <module>
+    from numpy.core._multiarray_umath import (
+ModuleNotFoundError: No module named 'numpy.core._multiarray_umath'
+
+During handling of the above exception, another exception occurred:
+
+Traceback (most recent call last):
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/.venv-elf-py312/lib/python3.12/site-packages/numpy/__init__.py", line 130, in <module>
+    from numpy.__config__ import show as show_config
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/.venv-elf-py312/lib/python3.12/site-packages/numpy/__config__.py", line 4, in <module>
+    from numpy.core._multiarray_umath import (
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/.venv-elf-py312/lib/python3.12/site-packages/numpy/core/__init__.py", line 50, in <module>
+    raise ImportError(msg)
+ImportError: 
+
+IMPORTANT: PLEASE READ THIS FOR ADVICE ON HOW TO SOLVE THIS ISSUE!
+
+Importing the numpy C-extensions failed. This error can happen for
+many reasons, often due to issues with your setup or how NumPy was
+installed.
+
+We have compiled some common reasons and troubleshooting tips at:
+
+    https://numpy.org/devdocs/user/troubleshooting-importerror.html
+
+Please note and check the following:
+
+  * The Python version is: Python3.12 from "/home/l/liyj/shiying/hku_diffusion_dllm/.venv-elf-py312/bin/python"
+  * The NumPy version is: "1.26.4"
+
+and make sure that they are the versions you expect.
+Please carefully study the documentation linked above for further help.
+
+Original error was: No module named 'numpy.core._multiarray_umath'
+
+
+The above exception was the direct cause of the following exception:
+
+Traceback (most recent call last):
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/experiments/elf_timestep_alignment/compute_elf_alignment.py", line 21, in <module>
+    import numpy as np
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/.venv-elf-py312/lib/python3.12/site-packages/numpy/__init__.py", line 135, in <module>
+    raise ImportError(msg) from e
+ImportError: Error importing numpy: you should not try to import numpy from
+        its source directory; please exit the numpy source tree, and relaunch
+        your python interpreter from there.

runs/elf_timestep_alignment/elf-a80_679745.out

@@ -0,0 +1,23 @@
+[info] started_at=2026-07-10T18:01:44+08:00
+[info] node=xgpj0
+[info] run_dir=/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_final_layer_t005_095_679745
+Fri Jul 10 18:01:44 2026       
++-----------------------------------------------------------------------------------------+
+| NVIDIA-SMI 580.159.03             Driver Version: 580.159.03     CUDA Version: 13.0     |
++-----------------------------------------+------------------------+----------------------+
+| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
+| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
+|                                         |                        |               MIG M. |
+|=========================================+========================+======================|
+|   0  NVIDIA A100 80GB PCIe          Off |   00000000:01:00.0 Off |                    0 |
+| N/A   32C    P0             42W /  300W |       0MiB /  81920MiB |      0%      Default |
+|                                         |                        |             Disabled |
++-----------------------------------------+------------------------+----------------------+
+
++-----------------------------------------------------------------------------------------+
+| Processes:                                                                              |
+|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
+|        ID   ID                                                               Usage      |
+|=========================================================================================|
+|  No running processes found                                                             |
++-----------------------------------------------------------------------------------------+

runs/elf_timestep_alignment/elf-align_679699.err

@@ -0,0 +1,3 @@
+ERROR: Ignored the following versions that require a different python version: 1.10.0 Requires-Python >=3.8,<3.12; 1.10.0rc1 Requires-Python >=3.8,<3.12; 1.10.0rc2 Requires-Python >=3.8,<3.12; 1.10.1 Requires-Python >=3.8,<3.12; 1.11.0 Requires-Python >=3.9,<3.13; 1.11.0rc1 Requires-Python >=3.9,<3.13; 1.11.0rc2 Requires-Python >=3.9,<3.13; 1.11.1 Requires-Python >=3.9,<3.13; 1.11.2 Requires-Python >=3.9,<3.13; 1.11.3 Requires-Python >=3.9,<3.13; 1.21.2 Requires-Python >=3.7,<3.11; 1.21.3 Requires-Python >=3.7,<3.11; 1.21.4 Requires-Python >=3.7,<3.11; 1.21.5 Requires-Python >=3.7,<3.11; 1.21.6 Requires-Python >=3.7,<3.11; 1.26.0 Requires-Python >=3.9,<3.13; 1.26.1 Requires-Python >=3.9,<3.13; 1.6.2 Requires-Python >=3.7,<3.10; 1.6.3 Requires-Python >=3.7,<3.10; 1.7.0 Requires-Python >=3.7,<3.10; 1.7.1 Requires-Python >=3.7,<3.10; 1.7.2 Requires-Python >=3.7,<3.11; 1.7.3 Requires-Python >=3.7,<3.11; 1.8.0 Requires-Python >=3.8,<3.11; 1.8.0rc1 Requires-Python >=3.8,<3.11; 1.8.0rc2 Requires-Python >=3.8,<3.11; 1.8.0rc3 Requires-Python >=3.8,<3.11; 1.8.0rc4 Requires-Python >=3.8,<3.11; 1.8.1 Requires-Python >=3.8,<3.11; 1.9.0 Requires-Python >=3.8,<3.12; 1.9.0rc1 Requires-Python >=3.8,<3.12; 1.9.0rc2 Requires-Python >=3.8,<3.12; 1.9.0rc3 Requires-Python >=3.8,<3.12; 1.9.1 Requires-Python >=3.8,<3.12
+ERROR: Could not find a version that satisfies the requirement torch==2.3.0 (from versions: 2.5.0, 2.5.1, 2.6.0, 2.7.0, 2.7.1, 2.8.0, 2.9.0, 2.9.1, 2.10.0, 2.11.0, 2.12.0, 2.12.1, 2.13.0)
+ERROR: No matching distribution found for torch==2.3.0

runs/elf_timestep_alignment/elf-align_679716.err

@@ -0,0 +1,19 @@
+ERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.
+torch 2.3.0 requires nvidia-cudnn-cu12==8.9.2.26; platform_system == "Linux" and platform_machine == "x86_64", but you have nvidia-cudnn-cu12 9.24.0.43 which is incompatible.
+Traceback (most recent call last):
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/experiments/elf_timestep_alignment/compute_elf_alignment.py", line 369, in <module>
+    main()
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/experiments/elf_timestep_alignment/compute_elf_alignment.py", line 164, in main
+    from transformers import AutoTokenizer
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/.venv-elf-py312/lib/python3.12/site-packages/transformers/__init__.py", line 26, in <module>
+    from . import dependency_versions_check
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/.venv-elf-py312/lib/python3.12/site-packages/transformers/dependency_versions_check.py", line 16, in <module>
+    from .utils.versions import require_version, require_version_core
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/.venv-elf-py312/lib/python3.12/site-packages/transformers/utils/__init__.py", line 34, in <module>
+    from .generic import (
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/.venv-elf-py312/lib/python3.12/site-packages/transformers/utils/generic.py", line 462, in <module>
+    import torch.utils._pytree as _torch_pytree
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/.venv-elf-py312/lib/python3.12/site-packages/torch/__init__.py", line 237, in <module>
+    from torch._C import *  # noqa: F403
+    ^^^^^^^^^^^^^^^^^^^^^^
+ImportError: libcudnn.so.8: cannot open shared object file: No such file or directory

runs/elf_timestep_alignment/elf-align_679741.err

@@ -0,0 +1,23 @@
+The cache for model files in Transformers v4.22.0 has been updated. Migrating your old cache. This is a one-time only operation. You can interrupt this and resume the migration later on by calling `transformers.utils.move_cache()`.
+
+2026-07-10 17:59:17.190345: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
+E0000 00:00:1783677557.203174 1920482 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+E0000 00:00:1783677557.207139 1920482 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+W0000 00:00:1783677557.219085 1920482 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783677557.219109 1920482 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783677557.219112 1920482 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783677557.219115 1920482 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+Traceback (most recent call last):
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/experiments/elf_timestep_alignment/compute_elf_alignment.py", line 369, in <module>
+    main()
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/experiments/elf_timestep_alignment/compute_elf_alignment.py", line 177, in main
+    config = load_config_from_yaml(args.config)
+             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/reference/code/ELF/src/configs/config.py", line 145, in load_config_from_yaml
+    config.sampling_configs = load_sampling_configs(config.sampling_configs_path)
+                              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+  File "/home/l/liyj/shiying/hku_diffusion_dllm/reference/code/ELF/src/configs/config.py", line 211, in load_sampling_configs
+    with open(sampling_configs_path, "r") as f:
+         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+FileNotFoundError: [Errno 2] No such file or directory: 'configs/sampling_configs/uncond_sampling_configs.yml'

runs/elf_timestep_alignment/elf-h47_679742.err

@@ -0,0 +1,46 @@
+2026-07-10 18:01:45.757714: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
+E0000 00:00:1783677705.770011 1921410 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+E0000 00:00:1783677705.773703 1921410 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+W0000 00:00:1783677705.785354 1921410 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783677705.785376 1921410 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783677705.785380 1921410 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783677705.785382 1921410 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+Failed to get TPU metadata (tpu-env) from instance metadata for variable CHIPS_PER_HOST_BOUNDS: DEADLINE_EXCEEDED: Timeout was reached
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/gcp_metadata_utils.cc:84
+learning/45eac/tfrc/runtime/env_var_utils.cc:50
+
+Failed to get TPU metadata (tpu-env) from instance metadata for variable HOST_BOUNDS: DEADLINE_EXCEEDED: Timeout was reached
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/gcp_metadata_utils.cc:84
+learning/45eac/tfrc/runtime/env_var_utils.cc:50
+
+Failed to get TPU metadata (tpu-env) from instance metadata for variable ALT: DEADLINE_EXCEEDED: Timeout was reached
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/gcp_metadata_utils.cc:84
+learning/45eac/tfrc/runtime/env_var_utils.cc:50
+
+Failed to get TPU metadata (tpu-env) from instance metadata for variable WRAP: DEADLINE_EXCEEDED: Timeout was reached
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/gcp_metadata_utils.cc:84
+learning/45eac/tfrc/runtime/env_var_utils.cc:50
+
+Failed to get TPU metadata (accelerator-type) from instance metadata for variable TPU_ACCELERATOR_TYPE: DEADLINE_EXCEEDED: Timeout was reached
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/gcp_metadata_utils.cc:84
+
+Failed to find host bounds for accelerator type: WARNING: could not determine TPU accelerator type, please set env var `TPU_ACCELERATOR_TYPE` manually, otherwise libtpu.so may not properly initialize.
+Failed to get TPU metadata (agent-worker-number) from instance metadata for variable TPU_WORKER_ID: DEADLINE_EXCEEDED: Timeout was reached
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/gcp_metadata_utils.cc:84
+
+Failed to get TPU metadata (worker-network-endpoints) from instance metadata for variable TPU_WORKER_HOSTNAMES: DEADLINE_EXCEEDED: Timeout was reached
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/gcp_metadata_utils.cc:84
+
+WARNING: Logging before InitGoogle() is written to STDERR
+E0000 00:00:1783680343.250683 1921410 common_lib.cc:520] INVALID_ARGUMENT: Error: unexpected worker hostname 'WARNING: could not determine TPU worker hostnames or IP addresses' from env var TPU_WORKER_HOSTNAMES. Expecting a valid hostname or IP address without port number. (Full TPU workers' addr string: WARNING: could not determine TPU worker hostnames or IP addresses, please set env var `TPU_WORKER_HOSTNAMES` manually, otherwise libtpu.so may not properly initialize.)
+=== Source Location Trace: ===
+learning/45eac/tfrc/runtime/libtpu_init_utils.cc:175
+[2026-07-10T18:53:37.997] error: *** JOB 679742 ON xgpi13 CANCELLED AT 2026-07-10T18:53:37 DUE to SIGNAL Terminated ***

runs/elf_timestep_alignment/elf-h47_679742.out

@@ -0,0 +1,34 @@
+[info] started_at=2026-07-10T18:01:43+08:00
+[info] node=xgpi13
+[info] run_dir=/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_final_layer_t005_095_679742
+Fri Jul 10 18:01:43 2026       
++-----------------------------------------------------------------------------------------+
+| NVIDIA-SMI 575.57.08              Driver Version: 575.57.08      CUDA Version: 12.9     |
+|-----------------------------------------+------------------------+----------------------+
+| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
+| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
+|                                         |                        |               MIG M. |
+|=========================================+========================+======================|
+|   0  NVIDIA H100 NVL                On  |   00000000:82:00.0 Off |                   On |
+| N/A   69C    P0            290W /  400W |    6233MiB /  95830MiB |     N/A      Default |
+|                                         |                        |              Enabled |
++-----------------------------------------+------------------------+----------------------+
+
++-----------------------------------------------------------------------------------------+
+| MIG devices:                                                                            |
++------------------+----------------------------------+-----------+-----------------------+
+| GPU  GI  CI  MIG |                     Memory-Usage |        Vol|        Shared         |
+|      ID  ID  Dev |                       BAR1-Usage | SM     Unc| CE ENC  DEC  OFA  JPG |
+|                  |                                  |        ECC|                       |
+|==================+==================================+===========+=======================|
+|  0    2   0   0  |              44MiB / 47488MiB    | 60      0 |  3   0    3    0    3 |
+|                  |                 0MiB / 65535MiB  |           |                       |
++------------------+----------------------------------+-----------+-----------------------+
+                                                                                         
++-----------------------------------------------------------------------------------------+
+| Processes:                                                                              |
+|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
+|        ID   ID                                                               Usage      |
+|=========================================================================================|
+|  No running processes found                                                             |
++-----------------------------------------------------------------------------------------+

runs/elf_timestep_alignment/elf-h96_679743.out

@@ -0,0 +1,37 @@
+[info] started_at=2026-07-10T18:33:46+08:00
+[info] node=xgpi12
+[info] run_dir=/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_final_layer_t005_095_679743
+Fri Jul 10 18:33:46 2026       
++-----------------------------------------------------------------------------------------+
+| NVIDIA-SMI 575.57.08              Driver Version: 575.57.08      CUDA Version: 12.9     |
+|-----------------------------------------+------------------------+----------------------+
+| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
+| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
+|                                         |                        |               MIG M. |
+|=========================================+========================+======================|
+|   0  NVIDIA H100 NVL                On  |   00000000:E3:00.0 Off |                   On |
+| N/A   64C    P0             75W /  400W |      87MiB /  95830MiB |     N/A      Default |
+|                                         |                        |              Enabled |
++-----------------------------------------+------------------------+----------------------+
+
++-----------------------------------------------------------------------------------------+
+| MIG devices:                                                                            |
++------------------+----------------------------------+-----------+-----------------------+
+| GPU  GI  CI  MIG |                     Memory-Usage |        Vol|        Shared         |
+|      ID  ID  Dev |                       BAR1-Usage | SM     Unc| CE ENC  DEC  OFA  JPG |
+|                  |                                  |        ECC|                       |
+|==================+==================================+===========+=======================|
+|  0    1   0   0  |              44MiB / 47488MiB    | 60      0 |  3   0    3    0    3 |
+|                  |                 0MiB / 65535MiB  |           |                       |
++------------------+----------------------------------+-----------+-----------------------+
+|  0    2   0   1  |              44MiB / 47488MiB    | 60      0 |  3   0    3    0    3 |
+|                  |                 0MiB / 65535MiB  |           |                       |
++------------------+----------------------------------+-----------+-----------------------+
+                                                                                         
++-----------------------------------------------------------------------------------------+
+| Processes:                                                                              |
+|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
+|        ID   ID                                                               Usage      |
+|=========================================================================================|
+|  No running processes found                                                             |
++-----------------------------------------------------------------------------------------+

runs/elf_timestep_alignment/elfa-deep-s41_680360.err

@@ -0,0 +1,10 @@
+2026-07-11 12:01:32.916517: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
+E0000 00:00:1783742493.166854  626198 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+E0000 00:00:1783742493.236121  626198 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+W0000 00:00:1783742493.796499  626198 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783742493.796555  626198 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783742493.796563  626198 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783742493.796570  626198 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+
+WARNING:absl:The transformations API will eventually be replaced by an upgraded design. The current API will not be removed until this point, but it will no longer be actively worked on.

runs/elf_timestep_alignment/elfa-deep-s41_680360.out

@@ -0,0 +1,359 @@
+[info] started_at=2026-07-11T12:01:25+08:00
+[info] node=xgpg4
+[info] run_dir=/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_boundary_block11_t015_050_s001_n16_seed41_680360
+[info] seed=41 samples=16 t_values=0.15:0.01:0.50
+[info] grad_regex=blocks_11 texts_file=/home/l/liyj/shiying/hku_diffusion_dllm/experiments/elf_timestep_alignment/text_samples_32.txt
+Sat Jul 11 12:01:25 2026       
++-----------------------------------------------------------------------------------------+
+| NVIDIA-SMI 575.57.08              Driver Version: 575.57.08      CUDA Version: 12.9     |
+|-----------------------------------------+------------------------+----------------------+
+| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
+| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
+|                                         |                        |               MIG M. |
+|=========================================+========================+======================|
+|   0  NVIDIA A100-PCIE-40GB          On  |   00000000:01:00.0 Off |                    0 |
+| N/A   44C    P0             39W /  250W |       0MiB /  40960MiB |      0%      Default |
+|                                         |                        |             Disabled |
++-----------------------------------------+------------------------+----------------------+
+                                                                                         
++-----------------------------------------------------------------------------------------+
+| Processes:                                                                              |
+|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
+|        ID   ID                                                               Usage      |
+|=========================================================================================|
+|  No running processes found                                                             |
++-----------------------------------------------------------------------------------------+
+[info] jax_backend=gpu
+[info] jax_devices=[CudaDevice(id=0)]
+{
+  "mode": "elf_flow_matching",
+  "model": "ELF-B",
+  "t_values": [
+    0.15,
+    0.16,
+    0.17,
+    0.18,
+    0.19,
+    0.2,
+    0.21,
+    0.22,
+    0.23,
+    0.24,
+    0.25,
+    0.26,
+    0.27,
+    0.28,
+    0.29,
+    0.3,
+    0.31,
+    0.32,
+    0.33,
+    0.34,
+    0.35,
+    0.36,
+    0.37,
+    0.38,
+    0.39,
+    0.4,
+    0.41,
+    0.42,
+    0.43,
+    0.44,
+    0.45,
+    0.46,
+    0.47,
+    0.48,
+    0.49,
+    0.5
+  ],
+  "loss_by_t": [
+    0.7687702775001526,
+    0.7881273031234741,
+    0.8076191544532776,
+    0.8299820423126221,
+    0.8552634119987488,
+    0.8814494013786316,
+    0.9066674709320068,
+    0.933702826499939,
+    0.9614684581756592,
+    0.9797762036323547,
+    0.9908438920974731,
+    1.0052367448806763,
+    1.028529167175293,
+    1.0447502136230469,
+    1.0593575239181519,
+    1.0711325407028198,
+    1.0827609300613403,
+    1.0864359140396118,
+    1.1015979051589966,
+    1.1140385866165161,
+    1.125871181488037,
+    1.1202373504638672,
+    1.1312713623046875,
+    1.148082971572876,
+    1.1602612733840942,
+    1.1670055389404297,
+    1.183816909790039,
+    1.2126374244689941,
+    1.239477276802063,
+    1.2786539793014526,
+    1.3062680959701538,
+    1.3161256313323975,
+    1.3360689878463745,
+    1.3640625476837158,
+    1.3931139707565308,
+    1.420563817024231
+  ],
+  "adjacent_cosines": [
+    {
+      "from": 0.15,
+      "to": 0.16,
+      "cosine": 0.9962546965456975
+    },
+    {
+      "from": 0.16,
+      "to": 0.17,
+      "cosine": 0.9958086557639261
+    },
+    {
+      "from": 0.17,
+      "to": 0.18,
+      "cosine": 0.9952199417358164
+    },
+    {
+      "from": 0.18,
+      "to": 0.19,
+      "cosine": 0.9957765763378329
+    },
+    {
+      "from": 0.19,
+      "to": 0.2,
+      "cosine": 0.9961950155835183
+    },
+    {
+      "from": 0.2,
+      "to": 0.21,
+      "cosine": 0.9954582738655838
+    },
+    {
+      "from": 0.21,
+      "to": 0.22,
+      "cosine": 0.9934789887691607
+    },
+    {
+      "from": 0.22,
+      "to": 0.23,
+      "cosine": 0.9951807053251995
+    },
+    {
+      "from": 0.23,
+      "to": 0.24,
+      "cosine": 0.997011617832636
+    },
+    {
+      "from": 0.24,
+      "to": 0.25,
+      "cosine": 0.996136504453141
+    },
+    {
+      "from": 0.25,
+      "to": 0.26,
+      "cosine": 0.9936256119439957
+    },
+    {
+      "from": 0.26,
+      "to": 0.27,
+      "cosine": 0.9953386437777826
+    },
+    {
+      "from": 0.27,
+      "to": 0.28,
+      "cosine": 0.99665726559883
+    },
+    {
+      "from": 0.28,
+      "to": 0.29,
+      "cosine": 0.9974246236472415
+    },
+    {
+      "from": 0.29,
+      "to": 0.3,
+      "cosine": 0.9976000336202412
+    },
+    {
+      "from": 0.3,
+      "to": 0.31,
+      "cosine": 0.9977129935271191
+    },
+    {
+      "from": 0.31,
+      "to": 0.32,
+      "cosine": 0.9972786842646635
+    },
+    {
+      "from": 0.32,
+      "to": 0.33,
+      "cosine": 0.9977804207708572
+    },
+    {
+      "from": 0.33,
+      "to": 0.34,
+      "cosine": 0.9970579865710971
+    },
+    {
+      "from": 0.34,
+      "to": 0.35,
+      "cosine": 0.9880092997472294
+    },
+    {
+      "from": 0.35,
+      "to": 0.36,
+      "cosine": 0.9907945894296982
+    },
+    {
+      "from": 0.36,
+      "to": 0.37,
+      "cosine": 0.9904425952017323
+    },
+    {
+      "from": 0.37,
+      "to": 0.38,
+      "cosine": 0.9935099687507755
+    },
+    {
+      "from": 0.38,
+      "to": 0.39,
+      "cosine": 0.9977455672094535
+    },
+    {
+      "from": 0.39,
+      "to": 0.4,
+      "cosine": 0.9962772579743102
+    },
+    {
+      "from": 0.4,
+      "to": 0.41,
+      "cosine": 0.9974000695241294
+    },
+    {
+      "from": 0.41,
+      "to": 0.42,
+      "cosine": 0.9959696618678843
+    },
+    {
+      "from": 0.42,
+      "to": 0.43,
+      "cosine": 0.9951299312313682
+    },
+    {
+      "from": 0.43,
+      "to": 0.44,
+      "cosine": 0.9933298782853635
+    },
+    {
+      "from": 0.44,
+      "to": 0.45,
+      "cosine": 0.9928438941296612
+    },
+    {
+      "from": 0.45,
+      "to": 0.46,
+      "cosine": 0.9987327659673764
+    },
+    {
+      "from": 0.46,
+      "to": 0.47,
+      "cosine": 0.9990603713522872
+    },
+    {
+      "from": 0.47,
+      "to": 0.48,
+      "cosine": 0.9988712946382723
+    },
+    {
+      "from": 0.48,
+      "to": 0.49,
+      "cosine": 0.9992366335088162
+    },
+    {
+      "from": 0.49,
+      "to": 0.5,
+      "cosine": 0.9981090700779052
+    }
+  ]
+}
+{
+  "updated_at": "2026-07-11T04:04:04.559961+00:00",
+  "inputs": [
+    "/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_boundary_block11_t015_050_s001_n16_seed41_680360/alignment.json"
+  ],
+  "num_runs": 1,
+  "axis": [
+    0.15,
+    0.16,
+    0.17,
+    0.18,
+    0.19,
+    0.2,
+    0.21,
+    0.22,
+    0.23,
+    0.24,
+    0.25,
+    0.26,
+    0.27,
+    0.28,
+    0.29,
+    0.3,
+    0.31,
+    0.32,
+    0.33,
+    0.34,
+    0.35,
+    0.36,
+    0.37,
+    0.38,
+    0.39,
+    0.4,
+    0.41,
+    0.42,
+    0.43,
+    0.44,
+    0.45,
+    0.46,
+    0.47,
+    0.48,
+    0.49,
+    0.5
+  ],
+  "best_split": {
+    "split_index": 19,
+    "left_t": 0.33,
+    "right_t": 0.34,
+    "boundary": 0.335,
+    "within_mean": 0.8985269586014996,
+    "between_mean": 0.6242743197364485,
+    "gap": 0.274252638865051
+  },
+  "adjacent_minimum": {
+    "from": 0.34,
+    "to": 0.35,
+    "boundary": 0.345,
+    "cosine": 0.9880092997472294
+  },
+  "per_run_best_split": [
+    {
+      "path": "/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_boundary_block11_t015_050_s001_n16_seed41_680360/alignment.json",
+      "split_index": 19,
+      "left_t": 0.33,
+      "right_t": 0.34,
+      "boundary": 0.335,
+      "within_mean": 0.8985269586014996,
+      "between_mean": 0.6242743197364485,
+      "gap": 0.274252638865051
+    }
+  ],
+  "bootstrap_repetitions": 0,
+  "bootstrap_boundary_counts": {}
+}
+[info] finished_at=2026-07-11T12:04:05+08:00

runs/elf_timestep_alignment/elfa-deep-s42_680361.out

@@ -0,0 +1,359 @@
+[info] started_at=2026-07-11T12:01:26+08:00
+[info] node=xgpg5
+[info] run_dir=/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_boundary_block11_t015_050_s001_n16_seed42_680361
+[info] seed=42 samples=16 t_values=0.15:0.01:0.50
+[info] grad_regex=blocks_11 texts_file=/home/l/liyj/shiying/hku_diffusion_dllm/experiments/elf_timestep_alignment/text_samples_32.txt
+Sat Jul 11 12:01:26 2026       
++-----------------------------------------------------------------------------------------+
+| NVIDIA-SMI 575.57.08              Driver Version: 575.57.08      CUDA Version: 12.9     |
+|-----------------------------------------+------------------------+----------------------+
+| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
+| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
+|                                         |                        |               MIG M. |
+|=========================================+========================+======================|
+|   0  NVIDIA A100-PCIE-40GB          On  |   00000000:01:00.0 Off |                    0 |
+| N/A   46C    P0             38W /  250W |       0MiB /  40960MiB |      0%      Default |
+|                                         |                        |             Disabled |
++-----------------------------------------+------------------------+----------------------+
+                                                                                         
++-----------------------------------------------------------------------------------------+
+| Processes:                                                                              |
+|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
+|        ID   ID                                                               Usage      |
+|=========================================================================================|
+|  No running processes found                                                             |
++-----------------------------------------------------------------------------------------+
+[info] jax_backend=gpu
+[info] jax_devices=[CudaDevice(id=0)]
+{
+  "mode": "elf_flow_matching",
+  "model": "ELF-B",
+  "t_values": [
+    0.15,
+    0.16,
+    0.17,
+    0.18,
+    0.19,
+    0.2,
+    0.21,
+    0.22,
+    0.23,
+    0.24,
+    0.25,
+    0.26,
+    0.27,
+    0.28,
+    0.29,
+    0.3,
+    0.31,
+    0.32,
+    0.33,
+    0.34,
+    0.35,
+    0.36,
+    0.37,
+    0.38,
+    0.39,
+    0.4,
+    0.41,
+    0.42,
+    0.43,
+    0.44,
+    0.45,
+    0.46,
+    0.47,
+    0.48,
+    0.49,
+    0.5
+  ],
+  "loss_by_t": [
+    0.7894503474235535,
+    0.8079524040222168,
+    0.8270783424377441,
+    0.8486059904098511,
+    0.872994065284729,
+    0.8972299098968506,
+    0.9237579107284546,
+    0.9506700038909912,
+    0.9688873291015625,
+    0.9816381931304932,
+    0.99869304895401,
+    1.0168174505233765,
+    1.0352503061294556,
+    1.0519922971725464,
+    1.0608059167861938,
+    1.0697757005691528,
+    1.0704327821731567,
+    1.0813257694244385,
+    1.1016240119934082,
+    1.1162893772125244,
+    1.1168100833892822,
+    1.1151463985443115,
+    1.1108323335647583,
+    1.0977319478988647,
+    1.1078706979751587,
+    1.1296074390411377,
+    1.1423625946044922,
+    1.1698291301727295,
+    1.2012135982513428,
+    1.2234262228012085,
+    1.2518868446350098,
+    1.2801848649978638,
+    1.3058077096939087,
+    1.3293050527572632,
+    1.363580346107483,
+    1.399135708808899
+  ],
+  "adjacent_cosines": [
+    {
+      "from": 0.15,
+      "to": 0.16,
+      "cosine": 0.9969856025348265
+    },
+    {
+      "from": 0.16,
+      "to": 0.17,
+      "cosine": 0.9960913436074166
+    },
+    {
+      "from": 0.17,
+      "to": 0.18,
+      "cosine": 0.9946922842884732
+    },
+    {
+      "from": 0.18,
+      "to": 0.19,
+      "cosine": 0.9932789356000828
+    },
+    {
+      "from": 0.19,
+      "to": 0.2,
+      "cosine": 0.9936398333671173
+    },
+    {
+      "from": 0.2,
+      "to": 0.21,
+      "cosine": 0.9945562094093469
+    },
+    {
+      "from": 0.21,
+      "to": 0.22,
+      "cosine": 0.9941701291965652
+    },
+    {
+      "from": 0.22,
+      "to": 0.23,
+      "cosine": 0.9966545676856795
+    },
+    {
+      "from": 0.23,
+      "to": 0.24,
+      "cosine": 0.9959261758544973
+    },
+    {
+      "from": 0.24,
+      "to": 0.25,
+      "cosine": 0.9970681055822745
+    },
+    {
+      "from": 0.25,
+      "to": 0.26,
+      "cosine": 0.9966301047115543
+    },
+    {
+      "from": 0.26,
+      "to": 0.27,
+      "cosine": 0.9944628573126317
+    },
+    {
+      "from": 0.27,
+      "to": 0.28,
+      "cosine": 0.9965534005540374
+    },
+    {
+      "from": 0.28,
+      "to": 0.29,
+      "cosine": 0.9966129983478746
+    },
+    {
+      "from": 0.29,
+      "to": 0.3,
+      "cosine": 0.9967550461221731
+    },
+    {
+      "from": 0.3,
+      "to": 0.31,
+      "cosine": 0.9949799731487935
+    },
+    {
+      "from": 0.31,
+      "to": 0.32,
+      "cosine": 0.9952240562505341
+    },
+    {
+      "from": 0.32,
+      "to": 0.33,
+      "cosine": 0.9957941025804731
+    },
+    {
+      "from": 0.33,
+      "to": 0.34,
+      "cosine": 0.9964363140549877
+    },
+    {
+      "from": 0.34,
+      "to": 0.35,
+      "cosine": 0.997065057354391
+    },
+    {
+      "from": 0.35,
+      "to": 0.36,
+      "cosine": 0.9973318041458271
+    },
+    {
+      "from": 0.36,
+      "to": 0.37,
+      "cosine": 0.9977542815094094
+    },
+    {
+      "from": 0.37,
+      "to": 0.38,
+      "cosine": 0.9958980031408313
+    },
+    {
+      "from": 0.38,
+      "to": 0.39,
+      "cosine": 0.9916347110302582
+    },
+    {
+      "from": 0.39,
+      "to": 0.4,
+      "cosine": 0.9919780723496313
+    },
+    {
+      "from": 0.4,
+      "to": 0.41,
+      "cosine": 0.9955634416117262
+    },
+    {
+      "from": 0.41,
+      "to": 0.42,
+      "cosine": 0.9948527741651921
+    },
+    {
+      "from": 0.42,
+      "to": 0.43,
+      "cosine": 0.9945672635658813
+    },
+    {
+      "from": 0.43,
+      "to": 0.44,
+      "cosine": 0.9970740377487571
+    },
+    {
+      "from": 0.44,
+      "to": 0.45,
+      "cosine": 0.9972080504857049
+    },
+    {
+      "from": 0.45,
+      "to": 0.46,
+      "cosine": 0.997424196831476
+    },
+    {
+      "from": 0.46,
+      "to": 0.47,
+      "cosine": 0.9983774332413492
+    },
+    {
+      "from": 0.47,
+      "to": 0.48,
+      "cosine": 0.9987911011637768
+    },
+    {
+      "from": 0.48,
+      "to": 0.49,
+      "cosine": 0.9991536503708132
+    },
+    {
+      "from": 0.49,
+      "to": 0.5,
+      "cosine": 0.9989211079789755
+    }
+  ]
+}
+{
+  "updated_at": "2026-07-11T04:04:04.754417+00:00",
+  "inputs": [
+    "/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_boundary_block11_t015_050_s001_n16_seed42_680361/alignment.json"
+  ],
+  "num_runs": 1,
+  "axis": [
+    0.15,
+    0.16,
+    0.17,
+    0.18,
+    0.19,
+    0.2,
+    0.21,
+    0.22,
+    0.23,
+    0.24,
+    0.25,
+    0.26,
+    0.27,
+    0.28,
+    0.29,
+    0.3,
+    0.31,
+    0.32,
+    0.33,
+    0.34,
+    0.35,
+    0.36,
+    0.37,
+    0.38,
+    0.39,
+    0.4,
+    0.41,
+    0.42,
+    0.43,
+    0.44,
+    0.45,
+    0.46,
+    0.47,
+    0.48,
+    0.49,
+    0.5
+  ],
+  "best_split": {
+    "split_index": 18,
+    "left_t": 0.32,
+    "right_t": 0.33,
+    "boundary": 0.325,
+    "within_mean": 0.8942345505344492,
+    "between_mean": 0.6245899802614596,
+    "gap": 0.2696445702729896
+  },
+  "adjacent_minimum": {
+    "from": 0.38,
+    "to": 0.39,
+    "boundary": 0.385,
+    "cosine": 0.9916347110302582
+  },
+  "per_run_best_split": [
+    {
+      "path": "/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_boundary_block11_t015_050_s001_n16_seed42_680361/alignment.json",
+      "split_index": 18,
+      "left_t": 0.32,
+      "right_t": 0.33,
+      "boundary": 0.325,
+      "within_mean": 0.8942345505344492,
+      "between_mean": 0.6245899802614596,
+      "gap": 0.2696445702729896
+    }
+  ],
+  "bootstrap_repetitions": 0,
+  "bootstrap_boundary_counts": {}
+}
+[info] finished_at=2026-07-11T12:04:06+08:00

runs/elf_timestep_alignment/elfa-deep-s45_680362.err

@@ -0,0 +1,10 @@
+2026-07-11 12:01:32.916920: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
+E0000 00:00:1783742493.167341  632218 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+E0000 00:00:1783742493.236577  632218 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+W0000 00:00:1783742493.796702  632218 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783742493.796749  632218 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783742493.796757  632218 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783742493.796764  632218 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+
+WARNING:absl:The transformations API will eventually be replaced by an upgraded design. The current API will not be removed until this point, but it will no longer be actively worked on.

runs/elf_timestep_alignment/elfa-final-s42_680358.err

@@ -0,0 +1,10 @@
+2026-07-11 12:01:32.920334: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
+E0000 00:00:1783742493.170620 1080666 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+E0000 00:00:1783742493.240019 1080666 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+W0000 00:00:1783742493.800083 1080666 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783742493.800107 1080666 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783742493.800109 1080666 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783742493.800111 1080666 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+
+WARNING:absl:The transformations API will eventually be replaced by an upgraded design. The current API will not be removed until this point, but it will no longer be actively worked on.

runs/elf_timestep_alignment/elfa-final-s43_680359.out

@@ -0,0 +1,359 @@
+[info] started_at=2026-07-11T12:01:25+08:00
+[info] node=xgpg3
+[info] run_dir=/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_boundary_final_t015_050_s001_n16_seed43_680359
+[info] seed=43 samples=16 t_values=0.15:0.01:0.50
+[info] grad_regex=final_layer|proj_kernel|proj_bias texts_file=/home/l/liyj/shiying/hku_diffusion_dllm/experiments/elf_timestep_alignment/text_samples_32.txt
+Sat Jul 11 12:01:25 2026       
++-----------------------------------------------------------------------------------------+
+| NVIDIA-SMI 575.57.08              Driver Version: 575.57.08      CUDA Version: 12.9     |
+|-----------------------------------------+------------------------+----------------------+
+| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
+| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
+|                                         |                        |               MIG M. |
+|=========================================+========================+======================|
+|   0  NVIDIA A100-PCIE-40GB          On  |   00000000:01:00.0 Off |                    0 |
+| N/A   46C    P0             37W /  250W |       0MiB /  40960MiB |      0%      Default |
+|                                         |                        |             Disabled |
++-----------------------------------------+------------------------+----------------------+
+                                                                                         
++-----------------------------------------------------------------------------------------+
+| Processes:                                                                              |
+|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
+|        ID   ID                                                               Usage      |
+|=========================================================================================|
+|  No running processes found                                                             |
++-----------------------------------------------------------------------------------------+
+[info] jax_backend=gpu
+[info] jax_devices=[CudaDevice(id=0)]
+{
+  "mode": "elf_flow_matching",
+  "model": "ELF-B",
+  "t_values": [
+    0.15,
+    0.16,
+    0.17,
+    0.18,
+    0.19,
+    0.2,
+    0.21,
+    0.22,
+    0.23,
+    0.24,
+    0.25,
+    0.26,
+    0.27,
+    0.28,
+    0.29,
+    0.3,
+    0.31,
+    0.32,
+    0.33,
+    0.34,
+    0.35,
+    0.36,
+    0.37,
+    0.38,
+    0.39,
+    0.4,
+    0.41,
+    0.42,
+    0.43,
+    0.44,
+    0.45,
+    0.46,
+    0.47,
+    0.48,
+    0.49,
+    0.5
+  ],
+  "loss_by_t": [
+    0.788937509059906,
+    0.8132213354110718,
+    0.8401293754577637,
+    0.8684554696083069,
+    0.8944187760353088,
+    0.9147629737854004,
+    0.9380487203598022,
+    0.9609002470970154,
+    0.9864038825035095,
+    1.0159869194030762,
+    1.042405366897583,
+    1.0616533756256104,
+    1.0812407732009888,
+    1.0957436561584473,
+    1.1030596494674683,
+    1.1138523817062378,
+    1.1278197765350342,
+    1.1449909210205078,
+    1.1479227542877197,
+    1.1593621969223022,
+    1.1712638139724731,
+    1.1633938550949097,
+    1.1800318956375122,
+    1.1783900260925293,
+    1.188826560974121,
+    1.2003947496414185,
+    1.2203037738800049,
+    1.246224045753479,
+    1.274021863937378,
+    1.303206443786621,
+    1.3329801559448242,
+    1.366233468055725,
+    1.4012446403503418,
+    1.4246362447738647,
+    1.4459913969039917,
+    1.4678287506103516
+  ],
+  "adjacent_cosines": [
+    {
+      "from": 0.15,
+      "to": 0.16,
+      "cosine": 0.9982446478223043
+    },
+    {
+      "from": 0.16,
+      "to": 0.17,
+      "cosine": 0.9986869563134316
+    },
+    {
+      "from": 0.17,
+      "to": 0.18,
+      "cosine": 0.9990163607839365
+    },
+    {
+      "from": 0.18,
+      "to": 0.19,
+      "cosine": 0.999356143459096
+    },
+    {
+      "from": 0.19,
+      "to": 0.2,
+      "cosine": 0.9987424978996888
+    },
+    {
+      "from": 0.2,
+      "to": 0.21,
+      "cosine": 0.9983072190559795
+    },
+    {
+      "from": 0.21,
+      "to": 0.22,
+      "cosine": 0.9984089120028304
+    },
+    {
+      "from": 0.22,
+      "to": 0.23,
+      "cosine": 0.9975184048851281
+    },
+    {
+      "from": 0.23,
+      "to": 0.24,
+      "cosine": 0.9972021092633733
+    },
+    {
+      "from": 0.24,
+      "to": 0.25,
+      "cosine": 0.9972048499906594
+    },
+    {
+      "from": 0.25,
+      "to": 0.26,
+      "cosine": 0.9972697890148098
+    },
+    {
+      "from": 0.26,
+      "to": 0.27,
+      "cosine": 0.9971426894019264
+    },
+    {
+      "from": 0.27,
+      "to": 0.28,
+      "cosine": 0.9973506080245123
+    },
+    {
+      "from": 0.28,
+      "to": 0.29,
+      "cosine": 0.9972721487892194
+    },
+    {
+      "from": 0.29,
+      "to": 0.3,
+      "cosine": 0.9979230340784457
+    },
+    {
+      "from": 0.3,
+      "to": 0.31,
+      "cosine": 0.9977883464133772
+    },
+    {
+      "from": 0.31,
+      "to": 0.32,
+      "cosine": 0.9972405676335042
+    },
+    {
+      "from": 0.32,
+      "to": 0.33,
+      "cosine": 0.9975753358496272
+    },
+    {
+      "from": 0.33,
+      "to": 0.34,
+      "cosine": 0.9961770108426754
+    },
+    {
+      "from": 0.34,
+      "to": 0.35,
+      "cosine": 0.9976173199484535
+    },
+    {
+      "from": 0.35,
+      "to": 0.36,
+      "cosine": 0.9977780902829744
+    },
+    {
+      "from": 0.36,
+      "to": 0.37,
+      "cosine": 0.996919018490538
+    },
+    {
+      "from": 0.37,
+      "to": 0.38,
+      "cosine": 0.9981407498405426
+    },
+    {
+      "from": 0.38,
+      "to": 0.39,
+      "cosine": 0.9984868120483484
+    },
+    {
+      "from": 0.39,
+      "to": 0.4,
+      "cosine": 0.9987313807889135
+    },
+    {
+      "from": 0.4,
+      "to": 0.41,
+      "cosine": 0.9989361296191474
+    },
+    {
+      "from": 0.41,
+      "to": 0.42,
+      "cosine": 0.9990464199582916
+    },
+    {
+      "from": 0.42,
+      "to": 0.43,
+      "cosine": 0.9989740860278867
+    },
+    {
+      "from": 0.43,
+      "to": 0.44,
+      "cosine": 0.999130299979199
+    },
+    {
+      "from": 0.44,
+      "to": 0.45,
+      "cosine": 0.9993108486014458
+    },
+    {
+      "from": 0.45,
+      "to": 0.46,
+      "cosine": 0.9993110275865007
+    },
+    {
+      "from": 0.46,
+      "to": 0.47,
+      "cosine": 0.9993452632608613
+    },
+    {
+      "from": 0.47,
+      "to": 0.48,
+      "cosine": 0.9993674563567321
+    },
+    {
+      "from": 0.48,
+      "to": 0.49,
+      "cosine": 0.9993966878649299
+    },
+    {
+      "from": 0.49,
+      "to": 0.5,
+      "cosine": 0.9994516143437041
+    }
+  ]
+}
+{
+  "updated_at": "2026-07-11T04:03:52.830465+00:00",
+  "inputs": [
+    "/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_boundary_final_t015_050_s001_n16_seed43_680359/alignment.json"
+  ],
+  "num_runs": 1,
+  "axis": [
+    0.15,
+    0.16,
+    0.17,
+    0.18,
+    0.19,
+    0.2,
+    0.21,
+    0.22,
+    0.23,
+    0.24,
+    0.25,
+    0.26,
+    0.27,
+    0.28,
+    0.29,
+    0.3,
+    0.31,
+    0.32,
+    0.33,
+    0.34,
+    0.35,
+    0.36,
+    0.37,
+    0.38,
+    0.39,
+    0.4,
+    0.41,
+    0.42,
+    0.43,
+    0.44,
+    0.45,
+    0.46,
+    0.47,
+    0.48,
+    0.49,
+    0.5
+  ],
+  "best_split": {
+    "split_index": 16,
+    "left_t": 0.3,
+    "right_t": 0.31,
+    "boundary": 0.305,
+    "within_mean": 0.9469304948557973,
+    "between_mean": 0.7281832713353107,
+    "gap": 0.2187472235204866
+  },
+  "adjacent_minimum": {
+    "from": 0.33,
+    "to": 0.34,
+    "boundary": 0.335,
+    "cosine": 0.9961770108426754
+  },
+  "per_run_best_split": [
+    {
+      "path": "/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_boundary_final_t015_050_s001_n16_seed43_680359/alignment.json",
+      "split_index": 16,
+      "left_t": 0.3,
+      "right_t": 0.31,
+      "boundary": 0.305,
+      "within_mean": 0.9469304948557973,
+      "between_mean": 0.7281832713353107,
+      "gap": 0.2187472235204866
+    }
+  ],
+  "bootstrap_repetitions": 0,
+  "bootstrap_boundary_counts": {}
+}
+[info] finished_at=2026-07-11T12:03:54+08:00

runs/elf_timestep_alignment/elfb-deep-s43_680027.err

@@ -0,0 +1,10 @@
+2026-07-11 00:50:12.214320: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
+E0000 00:00:1783702212.587122 1392060 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+E0000 00:00:1783702212.680281 1392060 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+W0000 00:00:1783702213.425704 1392060 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783702213.425781 1392060 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783702213.425789 1392060 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783702213.425796 1392060 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+
+WARNING:absl:The transformations API will eventually be replaced by an upgraded design. The current API will not be removed until this point, but it will no longer be actively worked on.

runs/elf_timestep_alignment/elfb-deep-s44_680028.err

@@ -0,0 +1,10 @@
+2026-07-11 01:06:22.803380: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:467] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
+WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
+E0000 00:00:1783703182.816327 1955063 cuda_dnn.cc:8579] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
+E0000 00:00:1783703182.820107 1955063 cuda_blas.cc:1407] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
+W0000 00:00:1783703182.832872 1955063 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783703182.832894 1955063 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783703182.832897 1955063 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+W0000 00:00:1783703182.832900 1955063 computation_placer.cc:177] computation placer already registered. Please check linkage and avoid linking the same target more than once.
+
+WARNING:absl:The transformations API will eventually be replaced by an upgraded design. The current API will not be removed until this point, but it will no longer be actively worked on.

runs/elf_timestep_alignment/elfb-deep-s44_680028.out

@@ -0,0 +1,370 @@
+[info] started_at=2026-07-11T01:06:19+08:00
+[info] node=xgpi13
+[info] run_dir=/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_boundary_block11_t015_050_s001_n16_seed44_680028
+[info] seed=44 samples=16 t_values=0.15:0.01:0.50
+[info] grad_regex=blocks_11 texts_file=/home/l/liyj/shiying/hku_diffusion_dllm/experiments/elf_timestep_alignment/text_samples_32.txt
+Sat Jul 11 01:06:19 2026       
++-----------------------------------------------------------------------------------------+
+| NVIDIA-SMI 575.57.08              Driver Version: 575.57.08      CUDA Version: 12.9     |
+|-----------------------------------------+------------------------+----------------------+
+| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
+| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
+|                                         |                        |               MIG M. |
+|=========================================+========================+======================|
+|   0  NVIDIA H100 NVL                On  |   00000000:82:00.0 Off |                   On |
+| N/A   58C    P0            240W /  400W |   24743MiB /  95830MiB |     N/A      Default |
+|                                         |                        |              Enabled |
++-----------------------------------------+------------------------+----------------------+
+
++-----------------------------------------------------------------------------------------+
+| MIG devices:                                                                            |
++------------------+----------------------------------+-----------+-----------------------+
+| GPU  GI  CI  MIG |                     Memory-Usage |        Vol|        Shared         |
+|      ID  ID  Dev |                       BAR1-Usage | SM     Unc| CE ENC  DEC  OFA  JPG |
+|                  |                                  |        ECC|                       |
+|==================+==================================+===========+=======================|
+|  0    2   0   0  |              44MiB / 47488MiB    | 60      0 |  3   0    3    0    3 |
+|                  |                 0MiB / 65535MiB  |           |                       |
++------------------+----------------------------------+-----------+-----------------------+
+                                                                                         
++-----------------------------------------------------------------------------------------+
+| Processes:                                                                              |
+|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
+|        ID   ID                                                               Usage      |
+|=========================================================================================|
+|  No running processes found                                                             |
++-----------------------------------------------------------------------------------------+
+[info] jax_backend=gpu
+[info] jax_devices=[CudaDevice(id=0)]
+{
+  "mode": "elf_flow_matching",
+  "model": "ELF-B",
+  "t_values": [
+    0.15,
+    0.16,
+    0.17,
+    0.18,
+    0.19,
+    0.2,
+    0.21,
+    0.22,
+    0.23,
+    0.24,
+    0.25,
+    0.26,
+    0.27,
+    0.28,
+    0.29,
+    0.3,
+    0.31,
+    0.32,
+    0.33,
+    0.34,
+    0.35,
+    0.36,
+    0.37,
+    0.38,
+    0.39,
+    0.4,
+    0.41,
+    0.42,
+    0.43,
+    0.44,
+    0.45,
+    0.46,
+    0.47,
+    0.48,
+    0.49,
+    0.5
+  ],
+  "loss_by_t": [
+    0.7627679705619812,
+    0.7817412614822388,
+    0.8025046586990356,
+    0.82469642162323,
+    0.8481600284576416,
+    0.8676245212554932,
+    0.8870224952697754,
+    0.9036419987678528,
+    0.9500153064727783,
+    0.9626415967941284,
+    0.9728938937187195,
+    0.9880105257034302,
+    0.9993098378181458,
+    1.0243895053863525,
+    1.0442224740982056,
+    1.0454508066177368,
+    1.0349785089492798,
+    1.034855842590332,
+    1.0424922704696655,
+    1.0538352727890015,
+    1.0678465366363525,
+    1.0883949995040894,
+    1.1084671020507812,
+    1.1251789331436157,
+    1.1478139162063599,
+    1.1632744073867798,
+    1.1861265897750854,
+    1.202358365058899,
+    1.2165666818618774,
+    1.238524079322815,
+    1.2581349611282349,
+    1.2842414379119873,
+    1.3230923414230347,
+    1.357128381729126,
+    1.3853179216384888,
+    1.4156414270401
+  ],
+  "adjacent_cosines": [
+    {
+      "from": 0.15,
+      "to": 0.16,
+      "cosine": 0.9956201655094372
+    },
+    {
+      "from": 0.16,
+      "to": 0.17,
+      "cosine": 0.9946803134385943
+    },
+    {
+      "from": 0.17,
+      "to": 0.18,
+      "cosine": 0.994622392865257
+    },
+    {
+      "from": 0.18,
+      "to": 0.19,
+      "cosine": 0.9945958226289133
+    },
+    {
+      "from": 0.19,
+      "to": 0.2,
+      "cosine": 0.99561790601484
+    },
+    {
+      "from": 0.2,
+      "to": 0.21,
+      "cosine": 0.994733359162385
+    },
+    {
+      "from": 0.21,
+      "to": 0.22,
+      "cosine": 0.9956387362908341
+    },
+    {
+      "from": 0.22,
+      "to": 0.23,
+      "cosine": 0.9591980706963354
+    },
+    {
+      "from": 0.23,
+      "to": 0.24,
+      "cosine": 0.990455340728805
+    },
+    {
+      "from": 0.24,
+      "to": 0.25,
+      "cosine": 0.9939695710377966
+    },
+    {
+      "from": 0.25,
+      "to": 0.26,
+      "cosine": 0.9955205558288159
+    },
+    {
+      "from": 0.26,
+      "to": 0.27,
+      "cosine": 0.9955783770571687
+    },
+    {
+      "from": 0.27,
+      "to": 0.28,
+      "cosine": 0.9941638626636766
+    },
+    {
+      "from": 0.28,
+      "to": 0.29,
+      "cosine": 0.9961058749711726
+    },
+    {
+      "from": 0.29,
+      "to": 0.3,
+      "cosine": 0.9969421501983707
+    },
+    {
+      "from": 0.3,
+      "to": 0.31,
+      "cosine": 0.9953219324817955
+    },
+    {
+      "from": 0.31,
+      "to": 0.32,
+      "cosine": 0.9952163932489749
+    },
+    {
+      "from": 0.32,
+      "to": 0.33,
+      "cosine": 0.994137222151749
+    },
+    {
+      "from": 0.33,
+      "to": 0.34,
+      "cosine": 0.9930571229041386
+    },
+    {
+      "from": 0.34,
+      "to": 0.35,
+      "cosine": 0.9943159135370286
+    },
+    {
+      "from": 0.35,
+      "to": 0.36,
+      "cosine": 0.9983593567037766
+    },
+    {
+      "from": 0.36,
+      "to": 0.37,
+      "cosine": 0.9973887652174955
+    },
+    {
+      "from": 0.37,
+      "to": 0.38,
+      "cosine": 0.9961969737975683
+    },
+    {
+      "from": 0.38,
+      "to": 0.39,
+      "cosine": 0.9953705215627234
+    },
+    {
+      "from": 0.39,
+      "to": 0.4,
+      "cosine": 0.9960464212033776
+    },
+    {
+      "from": 0.4,
+      "to": 0.41,
+      "cosine": 0.9969713843778287
+    },
+    {
+      "from": 0.41,
+      "to": 0.42,
+      "cosine": 0.9968156785889787
+    },
+    {
+      "from": 0.42,
+      "to": 0.43,
+      "cosine": 0.997473862583549
+    },
+    {
+      "from": 0.43,
+      "to": 0.44,
+      "cosine": 0.996572437450513
+    },
+    {
+      "from": 0.44,
+      "to": 0.45,
+      "cosine": 0.9988315457491682
+    },
+    {
+      "from": 0.45,
+      "to": 0.46,
+      "cosine": 0.9990782355591911
+    },
+    {
+      "from": 0.46,
+      "to": 0.47,
+      "cosine": 0.9971105669983424
+    },
+    {
+      "from": 0.47,
+      "to": 0.48,
+      "cosine": 0.9988924458174007
+    },
+    {
+      "from": 0.48,
+      "to": 0.49,
+      "cosine": 0.9991129422845241
+    },
+    {
+      "from": 0.49,
+      "to": 0.5,
+      "cosine": 0.9985156098772441
+    }
+  ]
+}
+{
+  "updated_at": "2026-07-10T17:08:00.808609+00:00",
+  "inputs": [
+    "/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_boundary_block11_t015_050_s001_n16_seed44_680028/alignment.json"
+  ],
+  "num_runs": 1,
+  "axis": [
+    0.15,
+    0.16,
+    0.17,
+    0.18,
+    0.19,
+    0.2,
+    0.21,
+    0.22,
+    0.23,
+    0.24,
+    0.25,
+    0.26,
+    0.27,
+    0.28,
+    0.29,
+    0.3,
+    0.31,
+    0.32,
+    0.33,
+    0.34,
+    0.35,
+    0.36,
+    0.37,
+    0.38,
+    0.39,
+    0.4,
+    0.41,
+    0.42,
+    0.43,
+    0.44,
+    0.45,
+    0.46,
+    0.47,
+    0.48,
+    0.49,
+    0.5
+  ],
+  "best_split": {
+    "split_index": 14,
+    "left_t": 0.28,
+    "right_t": 0.29,
+    "boundary": 0.28500000000000003,
+    "within_mean": 0.9024637474103743,
+    "between_mean": 0.630875260075048,
+    "gap": 0.27158848733532626
+  },
+  "adjacent_minimum": {
+    "from": 0.22,
+    "to": 0.23,
+    "boundary": 0.225,
+    "cosine": 0.9591980706963354
+  },
+  "per_run_best_split": [
+    {
+      "path": "/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_boundary_block11_t015_050_s001_n16_seed44_680028/alignment.json",
+      "split_index": 14,
+      "left_t": 0.28,
+      "right_t": 0.29,
+      "boundary": 0.28500000000000003,
+      "within_mean": 0.9024637474103743,
+      "between_mean": 0.630875260075048,
+      "gap": 0.27158848733532626
+    }
+  ],
+  "bootstrap_repetitions": 0,
+  "bootstrap_boundary_counts": {}
+}
+[info] finished_at=2026-07-11T01:08:01+08:00

runs/elf_timestep_alignment/elfb-final-s41_680020.out

@@ -0,0 +1,370 @@
+[info] started_at=2026-07-11T01:03:08+08:00
+[info] node=xgpi13
+[info] run_dir=/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_boundary_final_t015_050_s001_n16_seed41_680020
+[info] seed=41 samples=16 t_values=0.15:0.01:0.50
+[info] grad_regex=final_layer|proj_kernel|proj_bias texts_file=/home/l/liyj/shiying/hku_diffusion_dllm/experiments/elf_timestep_alignment/text_samples_32.txt
+Sat Jul 11 01:03:08 2026       
++-----------------------------------------------------------------------------------------+
+| NVIDIA-SMI 575.57.08              Driver Version: 575.57.08      CUDA Version: 12.9     |
+|-----------------------------------------+------------------------+----------------------+
+| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
+| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
+|                                         |                        |               MIG M. |
+|=========================================+========================+======================|
+|   0  NVIDIA H100 NVL                On  |   00000000:82:00.0 Off |                   On |
+| N/A   72C    P0            256W /  400W |    8831MiB /  95830MiB |     N/A      Default |
+|                                         |                        |              Enabled |
++-----------------------------------------+------------------------+----------------------+
+
++-----------------------------------------------------------------------------------------+
+| MIG devices:                                                                            |
++------------------+----------------------------------+-----------+-----------------------+
+| GPU  GI  CI  MIG |                     Memory-Usage |        Vol|        Shared         |
+|      ID  ID  Dev |                       BAR1-Usage | SM     Unc| CE ENC  DEC  OFA  JPG |
+|                  |                                  |        ECC|                       |
+|==================+==================================+===========+=======================|
+|  0    2   0   0  |              44MiB / 47488MiB    | 60      0 |  3   0    3    0    3 |
+|                  |                 0MiB / 65535MiB  |           |                       |
++------------------+----------------------------------+-----------+-----------------------+
+                                                                                         
++-----------------------------------------------------------------------------------------+
+| Processes:                                                                              |
+|  GPU   GI   CI              PID   Type   Process name                        GPU Memory |
+|        ID   ID                                                               Usage      |
+|=========================================================================================|
+|  No running processes found                                                             |
++-----------------------------------------------------------------------------------------+
+[info] jax_backend=gpu
+[info] jax_devices=[CudaDevice(id=0)]
+{
+  "mode": "elf_flow_matching",
+  "model": "ELF-B",
+  "t_values": [
+    0.15,
+    0.16,
+    0.17,
+    0.18,
+    0.19,
+    0.2,
+    0.21,
+    0.22,
+    0.23,
+    0.24,
+    0.25,
+    0.26,
+    0.27,
+    0.28,
+    0.29,
+    0.3,
+    0.31,
+    0.32,
+    0.33,
+    0.34,
+    0.35,
+    0.36,
+    0.37,
+    0.38,
+    0.39,
+    0.4,
+    0.41,
+    0.42,
+    0.43,
+    0.44,
+    0.45,
+    0.46,
+    0.47,
+    0.48,
+    0.49,
+    0.5
+  ],
+  "loss_by_t": [
+    0.7688165903091431,
+    0.7881860733032227,
+    0.8076627254486084,
+    0.8300124406814575,
+    0.8552974462509155,
+    0.8814648389816284,
+    0.9066751599311829,
+    0.9337313771247864,
+    0.9615678787231445,
+    0.9800011515617371,
+    0.9911522269248962,
+    1.005384087562561,
+    1.0285662412643433,
+    1.04488205909729,
+    1.0595942735671997,
+    1.0713118314743042,
+    1.0829862356185913,
+    1.0868176221847534,
+    1.101675271987915,
+    1.114145040512085,
+    1.126178503036499,
+    1.1216404438018799,
+    1.131507396697998,
+    1.1484925746917725,
+    1.1604516506195068,
+    1.167521595954895,
+    1.1840507984161377,
+    1.2131421566009521,
+    1.2396544218063354,
+    1.2780667543411255,
+    1.3072309494018555,
+    1.3171030282974243,
+    1.3367830514907837,
+    1.364898920059204,
+    1.3939350843429565,
+    1.4213314056396484
+  ],
+  "adjacent_cosines": [
+    {
+      "from": 0.15,
+      "to": 0.16,
+      "cosine": 0.9988865875881188
+    },
+    {
+      "from": 0.16,
+      "to": 0.17,
+      "cosine": 0.9985909630144116
+    },
+    {
+      "from": 0.17,
+      "to": 0.18,
+      "cosine": 0.9985002482292916
+    },
+    {
+      "from": 0.18,
+      "to": 0.19,
+      "cosine": 0.9987054461676875
+    },
+    {
+      "from": 0.19,
+      "to": 0.2,
+      "cosine": 0.9986596233989851
+    },
+    {
+      "from": 0.2,
+      "to": 0.21,
+      "cosine": 0.9981939347981107
+    },
+    {
+      "from": 0.21,
+      "to": 0.22,
+      "cosine": 0.9972620881622799
+    },
+    {
+      "from": 0.22,
+      "to": 0.23,
+      "cosine": 0.9975541266823477
+    },
+    {
+      "from": 0.23,
+      "to": 0.24,
+      "cosine": 0.9974673401932552
+    },
+    {
+      "from": 0.24,
+      "to": 0.25,
+      "cosine": 0.9964224760899152
+    },
+    {
+      "from": 0.25,
+      "to": 0.26,
+      "cosine": 0.9957857414630581
+    },
+    {
+      "from": 0.26,
+      "to": 0.27,
+      "cosine": 0.997104777552731
+    },
+    {
+      "from": 0.27,
+      "to": 0.28,
+      "cosine": 0.9968518940775081
+    },
+    {
+      "from": 0.28,
+      "to": 0.29,
+      "cosine": 0.997847468506239
+    },
+    {
+      "from": 0.29,
+      "to": 0.3,
+      "cosine": 0.9979465524442366
+    },
+    {
+      "from": 0.3,
+      "to": 0.31,
+      "cosine": 0.9980341484829286
+    },
+    {
+      "from": 0.31,
+      "to": 0.32,
+      "cosine": 0.9980169360118402
+    },
+    {
+      "from": 0.32,
+      "to": 0.33,
+      "cosine": 0.9980631346612001
+    },
+    {
+      "from": 0.33,
+      "to": 0.34,
+      "cosine": 0.9979939141146454
+    },
+    {
+      "from": 0.34,
+      "to": 0.35,
+      "cosine": 0.9968067671243962
+    },
+    {
+      "from": 0.35,
+      "to": 0.36,
+      "cosine": 0.9962136500314553
+    },
+    {
+      "from": 0.36,
+      "to": 0.37,
+      "cosine": 0.9970017371532452
+    },
+    {
+      "from": 0.37,
+      "to": 0.38,
+      "cosine": 0.9978530871598346
+    },
+    {
+      "from": 0.38,
+      "to": 0.39,
+      "cosine": 0.998573715500597
+    },
+    {
+      "from": 0.39,
+      "to": 0.4,
+      "cosine": 0.9984291187651212
+    },
+    {
+      "from": 0.4,
+      "to": 0.41,
+      "cosine": 0.9984658010801796
+    },
+    {
+      "from": 0.41,
+      "to": 0.42,
+      "cosine": 0.998314890438336
+    },
+    {
+      "from": 0.42,
+      "to": 0.43,
+      "cosine": 0.9983219401949588
+    },
+    {
+      "from": 0.43,
+      "to": 0.44,
+      "cosine": 0.9987270523736189
+    },
+    {
+      "from": 0.44,
+      "to": 0.45,
+      "cosine": 0.9987952594735235
+    },
+    {
+      "from": 0.45,
+      "to": 0.46,
+      "cosine": 0.9990234184096225
+    },
+    {
+      "from": 0.46,
+      "to": 0.47,
+      "cosine": 0.9991560736232648
+    },
+    {
+      "from": 0.47,
+      "to": 0.48,
+      "cosine": 0.9992886696923636
+    },
+    {
+      "from": 0.48,
+      "to": 0.49,
+      "cosine": 0.9993453410767242
+    },
+    {
+      "from": 0.49,
+      "to": 0.5,
+      "cosine": 0.9993858321293491
+    }
+  ]
+}
+{
+  "updated_at": "2026-07-10T17:04:42.505672+00:00",
+  "inputs": [
+    "/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_boundary_final_t015_050_s001_n16_seed41_680020/alignment.json"
+  ],
+  "num_runs": 1,
+  "axis": [
+    0.15,
+    0.16,
+    0.17,
+    0.18,
+    0.19,
+    0.2,
+    0.21,
+    0.22,
+    0.23,
+    0.24,
+    0.25,
+    0.26,
+    0.27,
+    0.28,
+    0.29,
+    0.3,
+    0.31,
+    0.32,
+    0.33,
+    0.34,
+    0.35,
+    0.36,
+    0.37,
+    0.38,
+    0.39,
+    0.4,
+    0.41,
+    0.42,
+    0.43,
+    0.44,
+    0.45,
+    0.46,
+    0.47,
+    0.48,
+    0.49,
+    0.5
+  ],
+  "best_split": {
+    "split_index": 16,
+    "left_t": 0.3,
+    "right_t": 0.31,
+    "boundary": 0.305,
+    "within_mean": 0.935255293415878,
+    "between_mean": 0.7070568765599928,
+    "gap": 0.22819841685588516
+  },
+  "adjacent_minimum": {
+    "from": 0.25,
+    "to": 0.26,
+    "boundary": 0.255,
+    "cosine": 0.9957857414630581
+  },
+  "per_run_best_split": [
+    {
+      "path": "/home/l/liyj/shiying/hku_diffusion_dllm/runs/elf_timestep_alignment/elf_b_owt_boundary_final_t015_050_s001_n16_seed41_680020/alignment.json",
+      "split_index": 16,
+      "left_t": 0.3,
+      "right_t": 0.31,
+      "boundary": 0.305,
+      "within_mean": 0.935255293415878,
+      "between_mean": 0.7070568765599928,
+      "gap": 0.22819841685588516
+    }
+  ],
+  "bootstrap_repetitions": 0,
+  "bootstrap_boundary_counts": {}
+}
+[info] finished_at=2026-07-11T01:04:43+08:00