Upload 7 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* 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/tokenizer.json filter=lfs diff=lfs merge=lfs -text

config.json

@@ -0,0 +1,1184 @@
+{
+  "hidden_size": 1024,
+  "intermediate_size": 3072,
+  "num_hidden_layers": 28,
+  "num_attention_heads": 16,
+  "num_key_value_heads": 8,
+  "head_dim": 128,
+  "vocab_size": 151677,
+  "mask_vocab_size": 151677,
+  "mask_token_id": 151669,
+  "pad_token_id": 151643,
+  "max_position_embeddings": 40960,
+  "rope_theta": 1000000.0,
+  "rms_norm_eps": 1e-06,
+  "hidden_act": "silu",
+  "qk_norm": true,
+  "timestep_emb_hidden": 1024,
+  "tie_word_embeddings": false,
+  "mask_ratio_min": 0.0,
+  "mask_ratio_max": 1.0,
+  "model_type": "metadiffusion",
+  "architectures": [
+    "MetaDiffusion600MForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "hf_modeling.MetaDiffusion600MConfig",
+    "AutoModelForCausalLM": "hf_modeling.MetaDiffusion600MForCausalLM"
+  },
+  "eos_token_id": [
+    151645
+  ],
+  "rainbow_token_ids": [
+    151670,
+    151671,
+    151672,
+    151673,
+    151674,
+    151675,
+    151676
+  ],
+  "invalid_utf8_token_ids": [
+    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,
+    128,
+    129,
+    130,
+    131,
+    132,
+    133,
+    134,
+    135,
+    136,
+    137,
+    138,
+    139,
+    140,
+    141,
+    142,
+    143,
+    144,
+    145,
+    146,
+    147,
+    148,
+    149,
+    150,
+    151,
+    152,
+    153,
+    154,
+    155,
+    156,
+    157,
+    158,
+    159,
+    160,
+    161,
+    162,
+    163,
+    164,
+    165,
+    166,
+    167,
+    168,
+    169,
+    170,
+    171,
+    172,
+    173,
+    174,
+    175,
+    176,
+    177,
+    178,
+    179,
+    180,
+    181,
+    182,
+    183,
+    184,
+    185,
+    186,
+    187,
+    222,
+    223,
+    224,
+    225,
+    226,
+    227,
+    228,
+    229,
+    230,
+    231,
+    232,
+    233,
+    234,
+    235,
+    236,
+    237,
+    238,
+    239,
+    240,
+    241,
+    242,
+    243,
+    244,
+    245,
+    246,
+    247,
+    248,
+    249,
+    250,
+    251,
+    252,
+    253,
+    254,
+    255,
+    378,
+    1277,
+    1540,
+    2194,
+    2775,
+    3023,
+    3219,
+    3402,
+    3490,
+    4811,
+    4973,
+    5502,
+    5564,
+    5691,
+    5758,
+    5902,
+    5953,
+    6142,
+    6518,
+    6564,
+    6571,
+    6606,
+    6684,
+    7045,
+    7158,
+    7311,
+    7367,
+    7492,
+    7536,
+    7588,
+    7908,
+    8032,
+    8079,
+    8158,
+    8508,
+    8519,
+    8843,
+    9260,
+    9284,
+    9722,
+    9973,
+    10081,
+    10203,
+    10417,
+    10663,
+    10996,
+    11125,
+    11336,
+    11484,
+    11534,
+    11620,
+    11621,
+    12274,
+    12491,
+    12585,
+    12619,
+    12787,
+    13058,
+    13092,
+    13165,
+    13344,
+    13465,
+    13512,
+    13519,
+    13531,
+    13558,
+    13567,
+    13846,
+    13869,
+    13949,
+    13982,
+    14273,
+    14368,
+    14467,
+    14546,
+    14559,
+    14675,
+    14922,
+    14936,
+    14942,
+    14959,
+    15005,
+    15127,
+    15213,
+    15224,
+    15362,
+    15393,
+    15600,
+    15675,
+    15798,
+    15827,
+    15896,
+    15902,
+    15927,
+    16160,
+    16499,
+    16633,
+    16747,
+    16825,
+    17049,
+    17156,
+    17158,
+    17401,
+    17429,
+    17462,
+    17467,
+    17641,
+    17730,
+    17783,
+    18059,
+    18505,
+    18535,
+    19088,
+    19113,
+    19261,
+    19457,
+    19549,
+    19564,
+    19741,
+    19956,
+    20026,
+    20047,
+    20098,
+    20382,
+    20741,
+    20778,
+    20879,
+    20885,
+    20919,
+    21012,
+    21135,
+    21280,
+    21530,
+    21562,
+    21704,
+    21849,
+    22009,
+    22129,
+    22247,
+    22597,
+    22762,
+    22859,
+    23093,
+    23272,
+    23433,
+    23444,
+    23479,
+    23687,
+    23781,
+    24071,
+    24159,
+    24345,
+    24447,
+    24735,
+    24749,
+    24831,
+    24864,
+    25145,
+    25209,
+    25452,
+    25703,
+    25870,
+    25928,
+    26188,
+    26254,
+    26336,
+    26454,
+    26601,
+    26771,
+    26927,
+    27095,
+    27209,
+    27214,
+    27325,
+    27398,
+    27517,
+    27538,
+    27640,
+    27757,
+    27999,
+    28107,
+    28145,
+    28757,
+    29150,
+    29220,
+    29237,
+    29389,
+    29675,
+    29768,
+    29955,
+    29975,
+    30006,
+    30274,
+    30395,
+    30828,
+    30838,
+    30916,
+    31305,
+    31482,
+    31501,
+    32208,
+    32254,
+    32343,
+    32465,
+    32843,
+    32926,
+    33173,
+    33352,
+    33543,
+    33556,
+    33593,
+    33739,
+    33986,
+    34429,
+    34577,
+    34622,
+    34629,
+    34759,
+    34811,
+    34992,
+    35017,
+    35049,
+    35055,
+    35146,
+    35293,
+    35496,
+    35635,
+    36097,
+    36109,
+    36677,
+    36695,
+    37289,
+    37440,
+    37472,
+    37698,
+    37913,
+    38177,
+    38419,
+    38522,
+    38911,
+    39098,
+    39175,
+    39317,
+    39834,
+    40419,
+    40714,
+    40732,
+    40747,
+    41902,
+    42144,
+    43410,
+    43459,
+    43547,
+    43559,
+    43752,
+    44014,
+    44044,
+    44104,
+    44258,
+    44401,
+    44680,
+    44706,
+    44726,
+    44818,
+    44832,
+    44851,
+    44965,
+    45881,
+    46282,
+    46353,
+    46800,
+    47378,
+    47379,
+    47972,
+    48312,
+    48533,
+    48591,
+    48749,
+    48800,
+    49166,
+    49173,
+    49420,
+    49943,
+    50111,
+    51275,
+    51497,
+    52133,
+    52408,
+    52798,
+    52887,
+    53496,
+    53556,
+    54128,
+    54492,
+    54642,
+    55890,
+    56252,
+    56823,
+    56842,
+    57154,
+    57160,
+    57332,
+    57599,
+    57743,
+    58098,
+    58169,
+    58299,
+    58557,
+    58715,
+    58899,
+    58994,
+    60627,
+    60757,
+    60864,
+    62005,
+    62416,
+    62544,
+    63046,
+    63219,
+    63761,
+    64520,
+    64893,
+    65185,
+    65291,
+    65850,
+    65877,
+    66118,
+    66498,
+    67895,
+    68405,
+    68597,
+    69382,
+    69687,
+    69792,
+    70179,
+    70297,
+    70467,
+    70731,
+    71306,
+    71443,
+    71481,
+    71933,
+    72219,
+    72496,
+    72497,
+    72509,
+    74165,
+    74209,
+    75107,
+    75142,
+    75360,
+    75375,
+    75598,
+    75962,
+    76080,
+    76986,
+    77156,
+    78910,
+    79142,
+    79531,
+    79590,
+    79888,
+    80080,
+    80702,
+    82501,
+    82712,
+    82798,
+    82912,
+    83050,
+    83164,
+    83191,
+    83643,
+    83849,
+    84200,
+    84215,
+    85794,
+    86643,
+    87299,
+    87959,
+    88541,
+    89061,
+    89083,
+    89179,
+    89297,
+    90557,
+    90639,
+    90679,
+    90734,
+    90839,
+    90894,
+    91048,
+    91771,
+    91888,
+    92120,
+    92173,
+    94964,
+    95211,
+    96938,
+    97015,
+    97946,
+    98642,
+    98734,
+    99156,
+    99157,
+    99158,
+    99159,
+    99160,
+    99161,
+    99162,
+    99163,
+    99166,
+    99167,
+    99168,
+    99169,
+    99170,
+    99171,
+    99173,
+    99174,
+    99175,
+    99176,
+    99177,
+    99179,
+    99181,
+    99183,
+    99184,
+    99187,
+    99188,
+    99189,
+    99192,
+    99196,
+    99197,
+    99198,
+    99201,
+    99202,
+    99203,
+    99206,
+    99207,
+    99211,
+    99214,
+    99215,
+    99220,
+    99221,
+    99224,
+    99228,
+    99229,
+    99230,
+    99231,
+    99238,
+    99239,
+    99240,
+    99247,
+    99248,
+    99249,
+    99254,
+    99255,
+    99256,
+    99264,
+    99265,
+    99266,
+    99267,
+    99268,
+    99269,
+    99274,
+    99275,
+    99276,
+    99281,
+    99282,
+    99289,
+    99290,
+    99291,
+    99300,
+    99301,
+    99302,
+    99303,
+    99311,
+    99323,
+    99324,
+    99325,
+    99341,
+    99342,
+    99343,
+    99344,
+    99359,
+    99374,
+    99380,
+    99381,
+    99382,
+    99401,
+    99402,
+    99423,
+    99439,
+    99442,
+    99443,
+    99444,
+    99456,
+    99484,
+    99485,
+    99516,
+    99551,
+    99552,
+    99597,
+    99598,
+    99643,
+    99648,
+    99649,
+    99714,
+    99715,
+    99777,
+    99785,
+    99874,
+    99875,
+    99977,
+    99997,
+    100024,
+    100120,
+    100121,
+    100127,
+    100128,
+    100129,
+    100130,
+    100337,
+    100459,
+    100563,
+    100620,
+    100621,
+    100988,
+    101018,
+    101022,
+    101024,
+    101025,
+    101026,
+    101027,
+    101028,
+    101029,
+    101030,
+    101031,
+    101032,
+    101033,
+    101759,
+    101838,
+    101851,
+    101860,
+    101861,
+    101862,
+    101863,
+    101864,
+    101865,
+    101866,
+    101867,
+    101868,
+    101869,
+    101870,
+    101871,
+    101872,
+    101873,
+    101874,
+    101875,
+    101876,
+    101877,
+    101878,
+    101879,
+    101880,
+    101881,
+    103893,
+    103894,
+    103895,
+    103896,
+    103897,
+    103898,
+    103899,
+    103900,
+    103901,
+    103902,
+    103903,
+    103904,
+    103905,
+    103906,
+    103907,
+    103908,
+    103909,
+    103910,
+    103911,
+    103912,
+    103913,
+    103914,
+    103915,
+    103916,
+    103917,
+    104905,
+    109992,
+    119158,
+    119941,
+    119964,
+    120100,
+    120409,
+    121404,
+    121667,
+    121773,
+    122154,
+    122174,
+    122202,
+    122219,
+    122222,
+    122259,
+    122267,
+    122289,
+    122317,
+    122375,
+    122382,
+    122427,
+    122448,
+    122455,
+    122514,
+    122568,
+    122596,
+    122614,
+    122618,
+    122634,
+    122659,
+    122725,
+    122740,
+    122757,
+    122802,
+    122843,
+    122855,
+    122866,
+    122890,
+    122948,
+    122950,
+    122976,
+    122985,
+    122999,
+    123007,
+    123008,
+    123016,
+    123057,
+    123064,
+    123067,
+    123072,
+    123084,
+    123091,
+    123131,
+    123164,
+    123178,
+    123182,
+    123204,
+    123206,
+    123287,
+    123301,
+    123305,
+    123318,
+    123332,
+    123378,
+    123400,
+    123408,
+    123420,
+    123489,
+    123516,
+    123527,
+    123547,
+    123614,
+    123658,
+    123676,
+    123701,
+    123714,
+    123740,
+    123747,
+    123760,
+    123806,
+    123807,
+    123808,
+    123810,
+    123811,
+    123812,
+    123813,
+    123814,
+    123815,
+    123816,
+    123817,
+    123819,
+    123821,
+    123825,
+    123827,
+    123830,
+    123836,
+    123837,
+    123839,
+    123840,
+    123841,
+    123842,
+    123844,
+    123846,
+    123847,
+    123848,
+    123851,
+    123853,
+    123854,
+    123859,
+    123867,
+    123870,
+    123871,
+    123875,
+    123876,
+    123882,
+    123889,
+    123892,
+    123893,
+    123907,
+    123908,
+    123911,
+    123912,
+    123916,
+    123917,
+    123918,
+    123927,
+    123928,
+    123929,
+    123930,
+    123931,
+    123932,
+    123933,
+    123934,
+    123946,
+    123947,
+    123950,
+    123951,
+    123953,
+    123954,
+    123955,
+    123956,
+    123957,
+    123965,
+    123967,
+    123968,
+    123969,
+    123970,
+    123986,
+    123989,
+    123996,
+    124000,
+    124001,
+    124002,
+    124003,
+    124004,
+    124005,
+    124018,
+    124019,
+    124020,
+    124021,
+    124022,
+    124023,
+    124024,
+    124025,
+    124026,
+    124049,
+    124050,
+    124051,
+    124052,
+    124053,
+    124054,
+    124063,
+    124064,
+    124066,
+    124067,
+    124068,
+    124069,
+    124070,
+    124071,
+    124092,
+    124096,
+    124098,
+    124099,
+    124100,
+    124101,
+    124102,
+    124118,
+    124134,
+    124135,
+    124136,
+    124137,
+    124149,
+    124154,
+    124159,
+    124160,
+    124161,
+    124162,
+    124163,
+    124164,
+    124165,
+    124166,
+    124167,
+    124168,
+    124193,
+    124194,
+    124195,
+    124196,
+    124197,
+    124198,
+    124199,
+    124200,
+    124201,
+    124222,
+    124235,
+    124236,
+    124240,
+    124241,
+    124243,
+    124244,
+    124245,
+    124246,
+    124247,
+    124248,
+    124249,
+    124250,
+    124251,
+    124252,
+    124287,
+    124296,
+    124297,
+    124298,
+    124299,
+    124300,
+    124301,
+    124302,
+    124303,
+    124304,
+    124305,
+    124306,
+    124307,
+    124308,
+    124309,
+    124334,
+    124349,
+    124364,
+    124365,
+    124366,
+    124367,
+    124368,
+    124369,
+    124370,
+    124371,
+    124372,
+    124373,
+    124374,
+    124433,
+    124470,
+    124471,
+    124472,
+    124473,
+    124474,
+    124475,
+    124570,
+    124577,
+    124578,
+    124579,
+    124580,
+    124581,
+    124582,
+    124583,
+    124584,
+    124585,
+    124586,
+    124587,
+    124588,
+    124589,
+    124590,
+    124591,
+    124592,
+    124593,
+    124594,
+    124595,
+    124596,
+    124727,
+    124744,
+    124745,
+    124746,
+    124747,
+    124748,
+    124749,
+    124750,
+    124751,
+    124752,
+    124753,
+    124754,
+    124755,
+    124995,
+    124996,
+    124997,
+    124998,
+    124999,
+    125000,
+    125001,
+    125002,
+    125003,
+    125004,
+    125413,
+    125414,
+    125415,
+    125416,
+    125417,
+    125418,
+    125419,
+    125420,
+    125421,
+    125422,
+    125423,
+    125424,
+    125425,
+    125426,
+    125427,
+    125428,
+    125429,
+    125430,
+    126159,
+    126177,
+    126178,
+    126179,
+    126180,
+    126181,
+    126182,
+    126183,
+    126184,
+    126185,
+    126186,
+    126187,
+    126188,
+    126189,
+    126190,
+    126191,
+    126192,
+    126193,
+    127964,
+    128222,
+    128223,
+    128224,
+    128225,
+    128226,
+    128227,
+    128228,
+    128229,
+    128230,
+    128231,
+    128232,
+    128233,
+    128234,
+    128235,
+    128236,
+    128237,
+    128238,
+    128239,
+    128240,
+    128241,
+    128242,
+    145773,
+    146047,
+    146632,
+    148462,
+    148774,
+    148827,
+    148860,
+    148864,
+    148880,
+    149175,
+    149178,
+    149287,
+    149352,
+    149392,
+    149394,
+    149589,
+    149591,
+    149593,
+    149714,
+    149716,
+    149718,
+    149721,
+    149946,
+    149961,
+    149963,
+    149990,
+    150114,
+    150151,
+    150168,
+    150195,
+    150223,
+    150270,
+    150479,
+    150579,
+    150792,
+    150794,
+    150800,
+    150802,
+    150807,
+    150809,
+    150814,
+    150817,
+    150819,
+    150821,
+    150824,
+    150829,
+    150836,
+    151233,
+    151254,
+    151264,
+    151266,
+    151268,
+    151270,
+    151272,
+    151274,
+    151276,
+    151278,
+    151282,
+    151366,
+    151560
+  ]
+}

generation_config.json

@@ -0,0 +1,17 @@
+{
+  "temperature": 0.7,
+  "repetition_penalty": 1.5,
+  "num_steps": 128,
+  "max_new_tokens": 96,
+  "top_p": 0.0,
+  "min_p": 0.1,
+  "im_end_bias": 2.0,
+  "im_end_bias_t": 0.3,
+  "do_sample": true,
+  "transformers_version": "4.49.0",
+  "eos_token_id": [
+    151645
+  ],
+  "pad_token_id": 151643,
+  "mask_token_id": 151669
+}

hf_modeling.py

@@ -0,0 +1,494 @@
+# hf_modeling.py: standalone transformers modeling file for released
+# MetaDiffusion-600M artifacts (AutoModelForCausalLM with trust_remote_code).
+# Copied into the export dir by export_hf.py.
+#
+# The generate() override runs LLaDA-style iterative denoising with
+# left-to-right block commit: the leftmost masked positions are unmasked
+# first, so an <|im_end|> cannot win at position 0 and produce empty output.
+# Generation stops once a terminator is committed in the response region.
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers import GenerationMixin, PretrainedConfig, PreTrainedModel
+
+
+class MetaDiffusion600MConfig(PretrainedConfig):
+    model_type = "metadiffusion"
+
+    def __init__(
+        self,
+        hidden_size=1024,
+        intermediate_size=3072,
+        num_hidden_layers=28,
+        num_attention_heads=16,
+        num_key_value_heads=8,
+        head_dim=128,
+        vocab_size=151669,
+        mask_vocab_size=151677,
+        mask_token_id=151669,
+        pad_token_id=151643,
+        max_position_embeddings=32768,
+        rope_theta=1000000.0,
+        rms_norm_eps=1e-6,
+        hidden_act="silu",
+        qk_norm=True,
+        timestep_emb_hidden=1024,
+        tie_word_embeddings=False,
+        eos_token_id=None,
+        **kwargs,
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            eos_token_id=eos_token_id,
+            **kwargs,
+        )
+        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.num_key_value_heads = num_key_value_heads
+        self.head_dim = head_dim
+        self.vocab_size = vocab_size
+        self.mask_vocab_size = mask_vocab_size
+        self.mask_token_id = mask_token_id
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.hidden_act = hidden_act
+        self.qk_norm = qk_norm
+        self.timestep_emb_hidden = timestep_emb_hidden
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, x):
+        orig = x.dtype
+        x = x.float()
+        var = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.eps)
+        return (self.weight.float() * x).to(orig)
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=32768, base=1000000.0):
+        super().__init__()
+        self.dim = dim
+        self.base = base
+
+    def forward(self, x, position_ids):
+        # computed fresh each call on purpose: a stored inv_freq buffer is
+        # non-persistent, so it is NOT in the state dict and from_pretrained
+        # leaves it UNINITIALIZED, producing garbage cos/sin and NaN logits
+        # in the entire forward. Computing here is 28 tiny ops, immune to
+        # whatever transformers does to buffers during loading.
+        #
+        # Numerics replicate training (model.py + train.py): init computes
+        # inv_freq on CPU in fp32, then `model.to(device, dtype=bfloat16)`
+        # rounds the buffer to bf16, and the forward upcasts it back to fp32.
+        # Matching that here keeps the released file BITWISE-consistent with
+        # the training implementation (plain fp32 inv_freq differs by ~1 bf16
+        # ULP in the rotary, which drifts final logits by ~1.0 after 28
+        # layers).
+        inv_freq = 1.0 / (self.base ** (torch.arange(
+            0, self.dim, 2).float() / self.dim))
+        inv_freq = inv_freq.to(torch.bfloat16).to(torch.float32)
+        inv = inv_freq[None, :, None].to(x.device).expand(position_ids.shape[0], -1, 1)
+        pos = position_ids[:, None, :].float()
+        freqs = (inv @ pos).transpose(1, 2)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        return emb.cos().to(dtype=x.dtype), emb.sin().to(dtype=x.dtype)
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin):
+    cos, sin = cos.unsqueeze(1), sin.unsqueeze(1)
+    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.mlp = nn.Sequential(
+            nn.Linear(hidden_size, hidden_size * 4), nn.SiLU(),
+            nn.Linear(hidden_size * 4, hidden_size),
+        )
+
+    def forward(self, t):
+        half_dim = self.hidden_size // 2
+        emb = math.log(10000.0) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, device=t.device, dtype=torch.float32) * -emb)
+        emb = t[:, None].float() * emb[None, :]
+        emb = torch.cat([emb.sin(), emb.cos()], dim=-1)
+        # cast to the MLP weight dtype: the model may be bf16 while t is fp32
+        return self.mlp(emb.to(self.mlp[0].weight.dtype))
+
+
+class TimestepModulation(nn.Module):
+    """adaLN-style timestep conditioning: scale + shift the hidden state.
+
+    Zero-init scale/shift so the model is identity at step 0. Gradient is
+    dL/dscale = dL/dx * x (x nonzero), so the t-path trains: the old
+    zero-init additive residual deadlocked (zero output x zero weight =
+    zero gradient forever), leaving models noise-schedule-agnostic."""
+
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.proj = nn.Linear(hidden_size, hidden_size * 2)
+        nn.init.zeros_(self.proj.weight)
+        nn.init.zeros_(self.proj.bias)
+
+    def forward(self, x, emb):
+        scale, shift = self.proj(emb).chunk(2, dim=-1)
+        scale, shift = scale[:, None, :], shift[:, None, :]
+        return x * (1.0 + scale) + shift
+
+
+class Attention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.num_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.head_dim
+        self.num_kv_groups = self.num_heads // self.num_kv_heads
+        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
+        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) if config.qk_norm else nn.Identity()
+        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) if config.qk_norm else nn.Identity()
+        self.rotary_emb = RotaryEmbedding(config.head_dim,
+                                          max_position_embeddings=config.max_position_embeddings,
+                                          base=config.rope_theta)
+
+    def forward(self, x, attention_mask=None, position_ids=None):
+        batch, seq, _ = x.shape
+        q = self.q_proj(x).view(batch, seq, self.num_heads, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(batch, seq, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(batch, seq, self.num_kv_heads, self.head_dim).transpose(1, 2)
+        q, k = self.q_norm(q), self.k_norm(k)
+        cos, sin = self.rotary_emb(x, position_ids)
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+        if self.num_kv_groups > 1:
+            k = k.repeat_interleave(self.num_kv_groups, dim=1)
+            v = v.repeat_interleave(self.num_kv_groups, dim=1)
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask)
+        return self.o_proj(out.transpose(1, 2).contiguous().view(batch, seq, -1))
+
+
+class MLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.self_attn = Attention(config)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp = MLP(config)
+        self.timestep_modulation = TimestepModulation(config.hidden_size)
+
+    def forward(self, x, t_emb, attention_mask=None, position_ids=None):
+        residual = x
+        x = self.input_layernorm(x)
+        x = self.self_attn(x, attention_mask, position_ids)
+        x = residual + x
+        x = self.timestep_modulation(x, t_emb)
+        residual = x
+        x = self.post_attention_layernorm(x)
+        x = self.mlp(x)
+        x = residual + x
+        x = self.timestep_modulation(x, t_emb)
+        return x
+
+
+class MetaDiffusion600MModel(PreTrainedModel):
+    config_class = MetaDiffusion600MConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.mask_vocab_size, config.hidden_size)
+        self.timestep_emb = TimestepEmbedding(config.timestep_emb_hidden)
+        self.layers = nn.ModuleList([Block(config) for _ in range(config.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.lm_head = nn.Linear(config.hidden_size, config.mask_vocab_size, bias=False)
+        self.post_init()
+
+    def forward(self, input_ids, timesteps=None, attention_mask=None):
+        batch, seq = input_ids.shape
+        position_ids = torch.arange(seq, device=input_ids.device).unsqueeze(0).expand(batch, -1)
+        if timesteps is None:
+            timesteps = torch.full((batch,), 1.0, device=input_ids.device)
+        x = self.embed_tokens(input_ids)
+        t_emb = self.timestep_emb(timesteps)
+        attn_mask = None
+        if attention_mask is not None:
+            attn_mask = ((1.0 - attention_mask[:, None, None, :].float()) * -1e9).to(x.dtype)
+        for layer in self.layers:
+            x = layer(x, t_emb, attn_mask, position_ids)
+        x = self.norm(x)
+        return F.linear(x, self.lm_head.weight)
+
+
+class MetaDiffusion600MForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = MetaDiffusion600MConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = MetaDiffusion600MModel(config)
+        self.post_init()
+
+    def forward(self, input_ids, timesteps=None, attention_mask=None, **kwargs):
+        logits = self.model(input_ids, timesteps, attention_mask)
+        return type("MDOutput", (), {"logits": logits})()
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+    def _cumulative_unmask_frac(self, i, n):
+        return 0.5 * (1 - math.cos(math.pi * i / n))
+
+    def generate(self, input_ids, max_new_tokens=None, num_steps=None,
+                 temperature=None, repetition_penalty=None, top_p=None,
+                 min_p=None, im_end_bias=None, im_end_bias_t=None,
+                 smart_remask=None, smart_remask_thresh=None,
+                 smart_remask_iters=None, **kwargs):
+        """LLaDA-style iterative denoising with left-to-right commit.
+
+        All sampling params fall back to generation_config.json values when
+        not passed explicitly (release defaults ship in the config):
+        top_p/min_p: truncation sampling that cuts the unreliable tail of the
+        distribution (the junk-token source); use at most one (min_p 0.05-0.1
+        recommended, Nguyen 2024; top-p 0.9 Holtzman 2020).
+        im_end_bias: pragmatic logit nudge on the stop tokens when t is low
+        (release guardrail for terminator reliability)."""
+        device = input_ids.device
+        config = self.config
+        gc = self.generation_config
+        max_new_tokens = max_new_tokens if max_new_tokens is not None else getattr(gc, "max_new_tokens", 96)
+        num_steps = num_steps if num_steps is not None else getattr(gc, "num_steps", 128)
+        temperature = temperature if temperature is not None else getattr(gc, "temperature", 0.7)
+        repetition_penalty = repetition_penalty if repetition_penalty is not None else getattr(gc, "repetition_penalty", 1.5)
+        top_p = top_p if top_p is not None else getattr(gc, "top_p", 0.0)
+        min_p = min_p if min_p is not None else getattr(gc, "min_p", 0.1)
+        im_end_bias = im_end_bias if im_end_bias is not None else getattr(gc, "im_end_bias", 0.0)
+        im_end_bias_t = im_end_bias_t if im_end_bias_t is not None else getattr(gc, "im_end_bias_t", 0.3)
+        smart_remask = smart_remask if smart_remask is not None else getattr(gc, "smart_remask", False)
+        smart_remask_thresh = smart_remask_thresh if smart_remask_thresh is not None else getattr(gc, "smart_remask_thresh", 0.5)
+        smart_remask_iters = smart_remask_iters if smart_remask_iters is not None else getattr(gc, "smart_remask_iters", 2)
+        refine_steps = getattr(gc, "refine_steps", 16)
+        mask_id = config.mask_token_id
+        eos_ids = self.generation_config.eos_token_id
+        if not isinstance(eos_ids, (list, tuple)):
+            eos_ids = [eos_ids] if eos_ids is not None else []
+        eos_ids = [int(e) for e in eos_ids if e is not None]
+
+        prompt_len = input_ids.shape[1]
+        x = torch.full((1, prompt_len + max_new_tokens), mask_id, device=device, dtype=torch.long)
+        x[0, :prompt_len] = input_ids[0]
+        # commit-confidence map for smart remasking (top-1 prob at commit time)
+        conf = (torch.ones((1, x.shape[1]), dtype=torch.float32, device=device)
+                if smart_remask else None)
+
+        self.eval()
+        with torch.no_grad():
+            for i in range(num_steps):
+                frac_now = self._cumulative_unmask_frac(i, num_steps)
+                frac_next = self._cumulative_unmask_frac(i + 1, num_steps)
+                n_masked = (x == mask_id).sum().item()
+                if i == num_steps - 1:
+                    n_unmask = n_masked
+                else:
+                    n_unmask = max(int((frac_next - frac_now) * max_new_tokens + 0.5), 1) if n_masked > 0 else 0
+                if n_unmask == 0:
+                    break
+                t = torch.full((1,), 1.0 - frac_now, device=device)
+                logits = self.model(x, t)
+                logits = logits.logits if hasattr(logits, "logits") else logits
+                # fp32 sampling path + sanitize (chat.py parity): models
+                # trained with a mask-ratio curriculum have never seen t near
+                # 1.0, so the timestep embedding can emit NaN/inf in bf16 when
+                # generating; softmax/multinomial must never see them
+                logits = logits.float()
+                logits = torch.nan_to_num(logits, nan=0.0, posinf=50.0, neginf=-50.0)
+                logits[:, :, mask_id] = -1e9
+                rainbow_ids = getattr(config, "rainbow_token_ids", None) or \
+                    list(range(mask_id + 1, mask_id + 8))
+                logits[:, :, rainbow_ids] = -1e9
+                bad_ids = getattr(config, "invalid_utf8_token_ids", None)
+                if bad_ids:
+                    logits[:, :, bad_ids] = -1e9
+                if repetition_penalty != 1.0:
+                    committed = x[0, prompt_len:]
+                    committed = committed[committed != mask_id]
+                    if committed.numel() > 0:
+                        for tok in committed.unique():
+                            ti = tok.item()
+                            logits[0, :, ti] = torch.where(
+                                logits[0, :, ti] < 0,
+                                logits[0, :, ti] * repetition_penalty,
+                                logits[0, :, ti] / repetition_penalty)
+                if im_end_bias != 0.0 and 1.0 - frac_now < im_end_bias_t:
+                    # pragmatic terminator nudge: BEFORE softmax so it actually
+                    # shapes the sampled distribution
+                    for eid in eos_ids:
+                        logits[0, :, eid] = logits[0, :, eid] + im_end_bias
+                mask_positions = x == mask_id
+                probs = F.softmax(logits[mask_positions] / max(temperature, 1e-8), dim=-1)
+                probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0)
+                if top_p > 0.0:
+                    sorted_probs, indices = probs.sort(dim=-1, descending=True)
+                    drop = (sorted_probs.cumsum(dim=-1) - sorted_probs) > top_p
+                    sorted_probs = sorted_probs.masked_fill(drop, 0.0)
+                    sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True).clamp(min=1e-12)
+                    probs = torch.zeros_like(probs).scatter_(-1, indices, sorted_probs)
+                elif min_p > 0.0:
+                    threshold = min_p * probs.max(dim=-1, keepdim=True).values
+                    probs = probs.masked_fill(probs < threshold, 0.0)
+                    probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-12)
+                zero_rows = probs.sum(dim=-1, keepdim=True) <= 0
+                if zero_rows.any():
+                    probs = probs + zero_rows.to(probs.dtype)
+                probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-12)
+                p_max = probs.max(dim=-1).values
+                sampled = torch.multinomial(probs, 1).squeeze(-1)
+                mask_flat = mask_positions.nonzero(as_tuple=False)
+                if n_unmask < mask_positions.sum():
+                    fill_positions = mask_flat[:n_unmask]
+                    for idx, tok in zip(fill_positions, sampled[:n_unmask]):
+                        x[idx[0], idx[1]] = tok
+                    if conf is not None:
+                        conf[fill_positions[:, 0], fill_positions[:, 1]] = p_max[:n_unmask]
+                else:
+                    x[mask_positions] = sampled
+                    if conf is not None:
+                        conf[mask_positions] = p_max
+                if any((x[0, prompt_len:] == e).any().item() for e in eos_ids):
+                    break
+        if smart_remask and conf is not None:
+            rainbow_ids = getattr(config, "rainbow_token_ids", None) or \
+                list(range(mask_id + 1, mask_id + 8))
+            bad_ids = getattr(config, "invalid_utf8_token_ids", None)
+            x = _smart_remask(self.model, x, prompt_len, max_new_tokens, conf,
+                              eos_ids, mask_id, rainbow_ids, bad_ids,
+                              smart_remask_thresh, smart_remask_iters,
+                              refine_steps, temperature, repetition_penalty,
+                              top_p, min_p, im_end_bias, im_end_bias_t,
+                              self._cumulative_unmask_frac)
+        return x
+
+
+def _smart_remask(model, x, prompt_len, gen_len, conf, eos_ids, mask_id,
+                  rainbow_ids, bad_ids, thresh, max_iters, refine_steps,
+                  temperature, repetition_penalty, top_p, min_p, im_end_bias,
+                  im_end_bias_t, cumfrac):
+    """Confidence-gated re-denoising (PURE-style smart remasking): re-mask
+    exactly the tokens whose top-1 commit probability fell below `thresh`
+    and re-denoise them with the head fixed (chat.py parity). Runs even when
+    a terminator committed, cleaning low-confidence junk before it. Stops
+    early once a terminator commits or nothing is below the bar."""
+    device = x.device
+    lo = prompt_len
+    hi = prompt_len + gen_len
+    for _ in range(max_iters):
+        if eos_ids:
+            term_mask = (x[0, lo:hi] == eos_ids[0])
+            for e in eos_ids[1:]:
+                term_mask = term_mask | (x[0, lo:hi] == e)
+            if term_mask.any():
+                # never touch the terminator or anything past it
+                hi = lo + term_mask.nonzero(as_tuple=True)[0][0].item()
+                if hi <= lo:
+                    break
+        low = (conf[0, lo:hi] < thresh).nonzero(as_tuple=True)[0]
+        if low.numel() == 0:
+            break
+        n_remask = low.numel()
+        x[0, lo + low] = mask_id
+        conf[0, lo + low] = 1.0  # re-commits below the bar get caught again
+        for i in range(refine_steps):
+            n_masked = (x[0, lo:hi] == mask_id).sum().item()
+            if n_masked == 0:
+                break
+            if i == refine_steps - 1:
+                n_unmask = n_masked
+            else:
+                n_unmask = max(int((cumfrac(i + 1, refine_steps)
+                                    - cumfrac(i, refine_steps))
+                                   * n_remask + 0.5), 1)
+            n_unmask = min(n_unmask, n_masked)
+            t_now = 1.0 - cumfrac(i, refine_steps)
+            t_val = torch.full((1,), t_now, device=device)
+            out = model(x, t_val)
+            logits = out.logits if hasattr(out, "logits") else out
+            logits = logits.float()
+            logits = torch.nan_to_num(logits, nan=0.0, posinf=50.0, neginf=-50.0)
+            logits[:, :, mask_id] = -1e9
+            logits[:, :, rainbow_ids] = -1e9
+            if bad_ids:
+                logits[:, :, bad_ids] = -1e9
+            if im_end_bias != 0.0 and t_now < im_end_bias_t:
+                for e in eos_ids:
+                    logits[0, :, e] = logits[0, :, e] + im_end_bias
+            if repetition_penalty != 1.0:
+                committed = x[0, prompt_len:]
+                committed = committed[committed != mask_id]
+                if committed.numel() > 0:
+                    for tok in committed.unique():
+                        ti = tok.item()
+                        logits[0, :, ti] = torch.where(
+                            logits[0, :, ti] < 0,
+                            logits[0, :, ti] * repetition_penalty,
+                            logits[0, :, ti] / repetition_penalty)
+            mask_positions = x == mask_id
+            probs = F.softmax(logits[mask_positions] / max(temperature, 1e-8), dim=-1)
+            probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0)
+            if top_p > 0.0:
+                sorted_probs, indices = probs.sort(dim=-1, descending=True)
+                drop = (sorted_probs.cumsum(dim=-1) - sorted_probs) > top_p
+                sorted_probs = sorted_probs.masked_fill(drop, 0.0)
+                sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True).clamp(min=1e-12)
+                probs = torch.zeros_like(probs).scatter_(-1, indices, sorted_probs)
+            elif min_p > 0.0:
+                threshold = min_p * probs.max(dim=-1, keepdim=True).values
+                probs = probs.masked_fill(probs < threshold, 0.0)
+                probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-12)
+            zero_rows = probs.sum(dim=-1, keepdim=True) <= 0
+            if zero_rows.any():
+                probs = probs + zero_rows.to(probs.dtype)
+            probs = probs / probs.sum(dim=-1, keepdim=True).clamp(min=1e-12)
+            p_max = probs.max(dim=-1).values
+            sampled = torch.multinomial(probs, 1).squeeze(-1)
+            mask_flat = mask_positions.nonzero(as_tuple=False)
+            n_fill = min(n_unmask, mask_flat.shape[0])
+            if n_fill:
+                idxs = mask_flat[:n_fill]
+                x[idxs[:, 0], idxs[:, 1]] = sampled[:n_fill]
+                conf[idxs[:, 0], idxs[:, 1]] = p_max[:n_fill]
+            if eos_ids and any((x[0, lo:hi] == e).any().item() for e in eos_ids):
+                break
+        if eos_ids and any((x[0, lo:hi] == e).any().item() for e in eos_ids):
+            break
+    return x

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b5f28935e6ced4305d1a5c0f5eadc03cd89573bae1cd7730c9bbad72e7e7d58
+size 1636589136

tokenizer/chat_template.jinja

@@ -0,0 +1,4 @@
+{% for message in messages %}{{ '<|im_start|>' + message['role'] }}
+{{ message['content'] }}<|im_end|>
+{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant
+{% endif %}

tokenizer/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8495023a2b15619775c552c6d5da73a8bc65f3031ddf8d2be937385e22611791
+size 11424114

tokenizer/tokenizer_config.json

@@ -0,0 +1,26 @@
+{
+  "add_prefix_space": false,
+  "backend": "tokenizers",
+  "bos_token": null,
+  "chat_template": "{% for message in messages %}{{ '<|im_start|>' + message['role'] }}\n{{ message['content'] }}<|im_end|>\n{% endfor %}{% if add_generation_prompt %}<|im_start|>assistant\n{% endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": [
+    "[MASK]",
+    "<|r1|>",
+    "<|r2|>",
+    "<|r3|>",
+    "<|r4|>",
+    "<|r5|>",
+    "<|r6|>",
+    "<|r7|>"
+  ],
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}