flexmdm

FlexMDM — Dream-Coder-7B

An insertion + unmasking discrete-diffusion model for Python code, produced by fully fine-tuning Dream-org/Dream-Coder-v0-Base-7B. Unlike a fixed-length masked diffusion model, FlexMDM can grow its sequence during generation (a learned insertion head) and unmask tokens in any order, enabling genuinely any-order code generation.

• Paper: From Interface to Inference: Eliciting Any-Order Inference from Any-Order Models — S. Kim*, J. Kim*, T. Lee*, Y. Chen*, Y. Du, S. Kakade, S. Chen. arXiv:2607.26504.
• Code: github.com/SeunggeunKimkr/genuine-any-order, FlexMDM/ subdirectory (training, data pipeline, inference, and evaluation).
• Checkpoint: global_step_49500 (inference weights; optimizer/RNG state stripped).

⚠️ Loading (not a vanilla AutoModel)

This checkpoint hosts a Dream backbone plus FlexMDM-specific weights in flexmdm_extras.pt. A bare AutoModel.from_pretrained returns only the backbone. Load the full model with the FlexMDM package from the code repo:

# 1) install the FlexMDM package:  pip install -e .  (in the repo's FlexMDM/ dir)
from huggingface_hub import snapshot_download
from flexmdm.utils import load_model_and_tokenizer

ckpt = snapshot_download("yuyuanchen0/flexmdm")
model, tokenizer = load_model_and_tokenizer(checkpoint_dir=ckpt, max_length=768)

load_model_and_tokenizer defaults to attn_implementation="sdpa" (works everywhere). Pass "flash_attention_2" for speed (requires flash-attn), or "eager" to bit-match the released evaluation traces.

For sampling (the ā = 2.9 inference-time schedule, temperature 0.1, insertion temperature 0.6 on MBPP / 1.0 on HumanEval, 512 steps) use flexmdm.inference.flexmdm_generate; see the code repo's docs/REPRODUCE.md.

Model

• Base: DreamModel (7B; hidden 3584, 28 layers, GQA with 4 KV heads, vocab 152064, diffusion mask id 151666), inherited unchanged.
• FlexMDM additions: a per-position log-space insertion head (LayerNorm → Linear → GELU → Linear, clamped to [-15, 15]) and AdaLN time conditioning on the insertion-progress coordinate.
• Schedules: power family, α_t = 1−(1−t)^a (insertion), β_t = 1−(1−t)^(a·b) (unmasking), with a = b = 1.7.
• Training: AdamW, LR 1e-5 (backbone) / 2e-5 (insertion head), global batch 576, max length 768, FSDP HYBRID_SHARD, 16× H100, ~3 days (checkpoint at step 49500).

Training data

Fine-tuned on a five-source Hugging Face mixture (OpenCodeInstruct, opc-sft-stage2, KodCode-V1-SFT-4o, and rStar-Coder seed/synthetic). KodCode-V1-SFT-4o is CC BY-NC 4.0 (non-commercial). Full sources, filters, and licenses — and how to reconstruct the tokenized set — are in the code repo's docs/DATA.md.

Evaluation (pass@k; n = 16 samples/task)

These are the paper's Table 5 rows (extraction-robust any-of-4 grading, 30 s test timeout). MBPP/MBPP+ decode with the count-preserving insertion temperature 0.6 (HumanEval/HumanEval+: neutral 1.0); at the neutral 1.0 the MBPP rows are 62.22 / 74.61 / 81.81 / 86.19 / 89.68 and MBPP+ 52.86 / 64.38 / 72.04 / 77.00 / 80.16 — placement sharpening improves every pass@k on both suites. Generation is deterministic (content-addressed seeds), so the sample set is bit-reproducible — see the code repo's evals/REPRODUCIBILITY.md for the exact recipe.

FlexMDM also scores substantially higher than Dream-Coder on tree-based any-order metrics (CBC/RUB/RUB+/OBW).

License & intended use

Apache-2.0 (derived from Dream-Coder, Apache-2.0). Research artifact — not a deployment-ready system; generated code may be incorrect or insecure, so sandbox before executing. Note the non-commercial license on part of the training data (above).