config.json

{
  "architectures": [
    "MaskedDiffusionLM"
  ],
  "model_type": "diffurefill",
  "_comment": "Not a HuggingFace-loadable config. The reference implementation is big_common.py in this repo; the fields below describe that module exactly.",
  "hidden_size": 2048,
  "num_hidden_layers": 18,
  "num_attention_heads": 16,
  "num_key_value_heads": 16,
  "head_dim": 128,
  "intermediate_size": 5632,
  "hidden_act": "silu",
  "mlp_type": "swiglu",
  "norm_type": "rmsnorm",
  "norm_placement": "pre",
  "attention_bias": false,
  "mlp_bias": false,
  "position_encoding": "rope",
  "rope_theta": 10000.0,
  "rope_applied_to": [
    "q",
    "k"
  ],
  "max_position_embeddings": 2048,
  "attention_mask": "bidirectional",
  "attention_impl": "scaled_dot_product_attention; flex_attention with a document block mask when BG_DOCMASK=1",
  "vocab_size": 73760,
  "tokenizer": "openbmb/MiniCPM4-0.5B",
  "tokenizer_vocab_size": 73440,
  "mask_token_id": 73759,
  "pad_token_id": 73758,
  "tie_word_embeddings": true,
  "torch_dtype": "bfloat16",
  "parameters": {
    "total": 1076000000,
    "embedding_tied_with_head": 151060480,
    "per_block": 51384320,
    "per_block_attention": 16777216,
    "per_block_mlp": 34603008,
    "note": "vocab is padded from the tokenizer's 73440 to 73760 to reserve MASK and PAD and keep the head a round size"
  },
  "layer_order": [
    "tok_embedding (73760 x 2048, tied to output head)",
    "18 x Block:",
    "  RMSNorm -> Linear(2048 -> 6144, no bias) -> split q,k,v -> RoPE(q,k) -> bidirectional SDPA -> Linear(2048 -> 2048, no bias) -> residual add",
    "  RMSNorm -> [Linear(2048 -> 5632) SiLU] * [Linear(2048 -> 5632)] -> Linear(5632 -> 2048) -> residual add",
    "final RMSNorm",
    "output head = tok_embedding transposed (tied)"
  ],
  "init": {
    "embedding": "normal(0, 0.02)",
    "block_linears": "normal(0, 0.02 / sqrt(2 * num_hidden_layers))",
    "note": "the head is tied to the embedding, so the embedding scale sets the logit scale"
  },
  "objective": {
    "type": "masked_absorbing_diffusion",
    "loss": "cross_entropy on masked positions only",
    "mask_rate": "t ~ Uniform(0.15, 1.0) per row, positions masked independently",
    "expected_supervised_fraction": 0.575,
    "prefix_conditioning_prob": 0.5,
    "prefix_note": "on half of rows a random-length clean prefix is never masked, so training sees the prompt-conditioned shape generation uses",
    "head_note": "only masked positions are pushed through the 73760-wide head",
    "document_mask": "attention is kept inside each packed document (flex_attention block mask)",
    "salient_span_masking": {
      "rows": 0.25,
      "rule": "on a quarter of rows, names and numbers - capitalised words that do not merely open a sentence, and digit runs, merged into multi-word spans - are masked as whole units with probability sqrt(t) >= t; all other positions keep their independent draw",
      "why": "masked independently, a two-token name is hidden in full with probability t^2 and the visible half gives it away; hidden as a unit it must be recalled from the rest of the sentence"
    },
    "span_masking_prob": 0.0,
    "span_note": "contiguous span masking (mean 3) was used until step ~91.9k and then removed: independent masking beat it both by one-pass CE and by decoding accuracy on the 37.9M stand. PMI-unit masking, stratified mask rates and complementary masks were also tested there and lost to plain independent masking."
  },
  "training": {
    "optimizer": "AdamW",
    "betas": [
      0.9,
      0.95
    ],
    "eps": 1e-08,
    "weight_decay": 0.1,
    "grad_clip": 1.0,
    "peak_lr": 0.00015,
    "warmup_steps": 500,
    "schedule": "cosine to 0.1 of peak over 200000 steps",
    "sequence_length": 2048,
    "tokens_per_optimizer_step": 131072,
    "batch_sequences": 64,
    "micro_batch": 16,
    "dataset": "see dataset_mix",
    "distributed": {
      "method": "DiLoCo through this repository",
      "workers": "3-5 single-GPU workers (RTX PRO 6000 Blackwell, 96 GB), boxes come and go",
      "sharding": "each worker takes every Nth document of the base corpus from its own shuffle; the extra sources are read unsharded",
      "local_steps_between_merges": 150,
      "outer_optimizer": "plain average (momentum 0, lr 1); outer Nesterov was tried and raised held-out CE",
      "global_adoption": "theta <- global (plain replacement)",
      "guard": "a merge is rejected if it raises held-out masked CE by more than 25%",
      "global_adoption_note": "an earlier version kept local drift across a merge (theta += global - last_pushed). Measured on a fixed probe it made the model worse in 60 of 63 adoptions, mean +0.079 CE, against -0.08 to -0.10 for plain replacement: it preserves exactly the disagreeing component that averaging exists to divide by N, and it accumulates.",
      "merge": "the merger folds each worker checkpoint into a running fp32 sum as it arrives, waits at most 150 s past the first arrival, and publishes global_b.pt",
      "phase": "workers adopt a new global as soon as it appears (checked every 10 steps) and take its step counter, so the whole fleet starts each 150-step interval together"
    },
    "peak_lr_note": "halved from 3e-4. At 3e-4 a worker drifted 0.076 CE away from the merge point over one 150-step interval, which is more than averaging can reconcile; at 1.5e-4 the drift is 0.013 and the merged model fell from 4.336 to 4.288 on a fixed held-out probe.",
    "dataset_mix": {
      "note": "documents are drawn by source with these probabilities from step ~96.8k; before that the model saw only the two Ultra-FineWeb configs",
      "openbmb/Ultra-FineWeb-L3 (Multi-Style-Synthetic + QA-Synthetic)": 0.56,
      "wikimedia/wikipedia 20231101.en": 0.15,
      "HuggingFaceTB/finemath finemath-4plus": 0.08,
      "HuggingFaceTB/smollm-corpus cosmopedia-v2": 0.08,
      "OpenCoder-LLM/opc-annealing-corpus algorithmic_corpus": 0.04,
      "OpenCoder-LLM/opc-annealing-corpus synthetic_code_snippet": 0.04,
      "Asilarkness/DiffuRefill-facts (Wikipedia leads restated as textbook, lecture, museum-guide ... prose)": 0.05
    },
    "ema": {
      "beta": 0.9,
      "file": "global_ema.pt",
      "note": "an exponential average of the merged globals; it beats the latest global on held-out CE every round. It is published for use, never used to seed training."
    },
    "compile": "each block is torch.compile'd in place: 17% faster steps and less memory, same parameter names"
  },
  "decoding_default": {
    "draft_passes_K": 12,
    "refill_rounds": 8,
    "refill_fraction": 0.2,
    "draft_temp_start": 2.2,
    "draft_temp_end": 0.9,
    "gumbel_tau": 2.5,
    "refill_temp": 0.8,
    "refill_min_p": 0.1,
    "neighbour_ban": 6.0,
    "min_commit_gap": 2,
    "min_commit_gap_note": "two positions closer than this are never committed in the same pass; co-committed tokens are sampled from their own marginals as if independent, and neighbours are where that is most wrong. Measured on held-out text: adjacent-pair accuracy rises at every pass budget."
  },
  "status": {
    "stage": "pretraining, about half way",
    "lineage_step": 100050,
    "target_steps": 200000,
    "effective_tokens": 46150000000,
    "effective_tokens_note": "summed over every worker interval that went into a merge (one interval = 150 steps x 131072 tokens); the per-worker step counter alone undercounts by the number of workers",
    "round": 290,
    "best_checkpoint": "global_ema.pt",
    "note": "a pretraining checkpoint, not an instruction-tuned model; it writes fluent English but its facts are still unreliable"
  },
  "attention_note": "documents are packed back to back, so without a mask a masked position can be filled from an unrelated neighbouring document. The block-diagonal mask is not a speed cost: measured at the training shape it runs 13.2 ms against 17.5 ms unmasked, because blocks lying in no document are skipped entirely."
}