MDLM-OWT — flash-attn-free reimplementation

A drop-in reimplementation of kuleshov-group/mdlm-owt (Sahoo et al., Simple and Effective Masked Diffusion Language Models, NeurIPS 2024) that removes the hard dependency on flash-attn. The weights are the original weights, byte-for-byte; only the modeling code changed, so the model runs anywhere — Windows, CPU, and recent CUDA GPUs (including Blackwell / sm_120) where a flash-attn build is impractical. On Linux it equally spares you building flash-attn against your exact torch/CUDA pair.

This port was produced as the reference substrate for a cross-layer-transcoder interpretability study, where a deterministic, dependency-light, hookable forward pass matters. It is published in case it is useful to others blocked by the flash-attn requirement.

The MDLM authors also publish an official mdlm-no_flashattn-fp32-owt. That one drops flash attention from the forward pass, but its modeling code currently still imports flash_attn (for the rotary) and omegaconf, so it needs them installed to load; this port imports neither and loads where a flash-attn build is impractical.

What changed (and what did not)

Three changes to the computation, all in the modeling code:

1. Attention — torch.nn.functional.scaled_dot_product_attention (non-causal, full bidirectional) replaces flash_attn_varlen_qkvpacked_func. Equivalent here: one sequence per batch element, no causal mask, default softmax scale 1/sqrt(head_dim). The call is pinned to the math SDPA backend, so logits are reproducible across GPUs instead of varying with the kernel PyTorch happens to select.
2. Rotary embeddings — applied with an explicit rotate_half instead of flash_attn.layers.rotary. This reproduces flash-attn's non-interleaved convention exactly (the modeling file carries a doctest that proves the equivalence).
3. Precision — the transformer stack runs in the model's native dtype (fp32 by default) rather than under a bfloat16 autocast. fp32 is preferable for interpretability and for use as a numerical oracle.

The unused training-time dropout modules are also dropped (they are identity at inference).

Plus compatibility fixes. The constructor now calls self.post_init(), which the upstream code omits. transformers ≥ 5 requires it to finalize from_pretrained (without it, loading raises AttributeError: ... 'all_tied_weights_keys'). MDLM unties its embeddings, so this registers no tied weights — it only initializes HF bookkeeping. forward also accepts an attention_mask (so tokenizer output can be passed straight in; see Usage) and reads config.return_dict rather than the use_return_dict property that transformers 5 deprecates.

Unchanged: every weight (the original model.safetensors, SHA256 47149e73f7552f39ea9776dbe74d925d25237bcf2ed2e2ec03cdff9d51c82aa4), all parameter names, and the architecture. The forward signature is only extended (timesteps may be omitted, attention_mask is accepted); upstream-style calls work as before. The original weights load with no remapping. The fp32 forward is cross-checked against an independent Rust reimplementation (candle-mi) to a max abs-diff of 3×10⁻⁵.

Numerical note: because the upstream model runs its block stack under a bf16 autocast while this port runs fp32, logits differ from the upstream path at roughly bf16 tolerance. This port computes the higher-precision fp32 result of the same math.

Usage

trust_remote_code=True is required (the modeling code ships with the repo). This repo ships no tokenizer — use the stock GPT-2 tokenizer (openai-community/gpt2). The model's vocabulary is GPT-2's 50257 tokens plus a [MASK] token at id 50257 (vocab size 50258).

import torch
from transformers import AutoModelForMaskedLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForMaskedLM.from_pretrained(
    "TheQweaker/mdlm-owt-noflash", trust_remote_code=True).eval()

MASK = 50257
ids = tokenizer("The capital of France is Paris.")["input_ids"]
pos = ids.index(tokenizer(" Paris")["input_ids"][0])
ids[pos] = MASK  # mask the word " Paris"

x = torch.tensor([ids])
with torch.no_grad():
    logits = model(input_ids=x)        # raw [1, L, 50258] tensor (config sets return_dict=False)
logits[..., MASK] = float("-inf")      # never predict [MASK]
print(tokenizer.decode([logits[0, pos].argmax().item()]))   # -> " Paris"

timesteps may be omitted: the checkpoint was trained with time_conditioning=False, so the time input is ignored internally. (Pass return_dict=True if you want a MaskedLMOutput instead of a raw tensor.)

model(**tokenizer(text, return_tensors="pt")) also works: attention_mask is accepted, but it must be all ones. MDLM has no padding semantics — every position attends to every other, with no mask — so a padded batch raises a ValueError instead of letting real tokens attend to pad tokens. Batch equal-length sequences, or run them one at a time.

Framework-agnostic: any loader that reads the safetensors directly needs neither trust_remote_code nor the modeling code — only config.json + model.safetensors. candle-mi (Rust; crates.io, feature diffusion) loads this repo that way.

Generation

The original checkpoint ships no sampler. sample.py in this repo provides a faithful masked-diffusion ancestral sampler (SUBS parameterization, linear schedule):

import os, sys
from huggingface_hub import hf_hub_download

# from_pretrained does not fetch sample.py; download it and make it importable.
sys.path.insert(0, os.path.dirname(hf_hub_download("TheQweaker/mdlm-owt-noflash", "sample.py")))
from sample import generate

gen = torch.Generator(device=model.device).manual_seed(0)   # same device as the model
ids = generate(model, seq_len=64, num_steps=128, top_k=50, generator=gen)
print(tokenizer.decode(ids.tolist()))

From a clone of this repo, python sample.py runs the same thing end to end.

Note that a 169.6M-parameter masked-diffusion model under simple ancestral sampling tends to produce locally grammatical but repetitive text; sample quality is sensitive to step count, temperature, and top-k.

Model details (unchanged from upstream)

• Architecture: 12-layer DiT-style transformer, d_model 768, 12 heads, adaLN conditioning, untied embeddings, rotary positions. ≈169.6M total parameters (≈92M non-embedding; the input and output embeddings are untied, ≈77M together). Context length 1024.
• Training: OpenWebText, 1M steps (Sahoo et al., 2024).

Requirements and environment

No flash-attn. The runtime dependencies are just torch (≥ 2.3, for torch.nn.attention.sdpa_kernel), transformers, and safetensors; the tokenizer is plain gpt2 (tokenizers). trust_remote_code=True is required.

Validated stacks (Python 3.14.0, Windows 11). Each passes the card's masked-infill example, return_dict and hidden states, timesteps invariance, the sampler, tokenizer kwargs, and the padded-batch rejection:

The self.post_init() compatibility fix (see What changed) is what makes loading work on transformers ≥ 5; it is harmless on 4.x, as the table shows. The model runs on CPU and on any CUDA GPU — the point of the port is that it does not require a flash-attn build.

Provenance and credit

All weights and the architecture are the work of Sahoo et al. This repository contributes only the flash-attn-free modeling code. Please cite the original work and consult the MDLM GitHub repository and paper.

@inproceedings{sahoo2024simple,
    title={Simple and Effective Masked Diffusion Language Models},
    author={Subham Sekhar Sahoo and Marianne Arriola and Aaron Gokaslan and Edgar Mariano Marroquin and Alexander M Rush and Yair Schiff and Justin T Chiu and Volodymyr Kuleshov},
    booktitle={The Thirty-eighth Annual Conference on Neural Information Processing Systems},
    year={2024},
    url={https://openreview.net/forum?id=L4uaAR4ArM}
}

If this port itself was useful (for instance, to run MDLM without flash-attn), you can also cite it:

@misc{jacopin2026mdlmnoflash,
    title={{MDLM-OWT}: a flash-attn-free reimplementation},
    author={{\'E}ric Jacopin},
    year={2026},
    howpublished={\url{https://huggingface.co/TheQweaker/mdlm-owt-noflash}},
    note={Hugging Face model repository}
}

License

Apache-2.0, matching the upstream checkpoint.