DiffuCoder: post-training a masked diffusion model for code

A repository that accompanies a paper, with three released models

This software project accompanies the research paper DiffuCoder: Understanding and Improving Masked Diffusion Models for Code Generation, which is on arXiv as 2506.20639. That framing matters, because it tells you what the repository is for. It is the implementation of a method, released alongside the models that method produced, not a library anyone is expected to depend on.

Three checkpoints are published on Hugging Face, and they are the three that answer the three questions the paper asks. There is a Base model, DiffuCoder-7B-Base, an Instruct model, DiffuCoder-7B-Instruct, and a cpGRPO model, DiffuCoder-7B-cpGRPO, which is the one produced by the proposed post-training method.

The licence is an Apple one, the badge in the README says License: Apple, and the metadata records it as a custom licence GitHub cannot classify. The setup.py header carries the copyright line, Copyright (C) 2025 Apple Inc. All Rights Reserved, and it notes that the file is adapted from the setup.py of the open-r1 project. The primary language is Python, the default branch is main, the last push was on 2026-09-11, and the repository publishes no GitHub releases, so the models on Hugging Face are the artefacts and this repository is the method.

There is also a top-level inference_demo.py, which is the part most people who are not training the model will actually want.

Why diffusion models for code at all

The motivation section sets the scene, and it is a specific argument rather than a general one.

Masked Denoising Models, and the diffusion LLMs built on them such as LLaDA and Dream, have achieved performance on par with similarly sized autoregressive LLMs across many benchmarks. The README goes further, pointing at recent commercial-scale diffusion LLMs like Mercury and Gemini Diffusion to say that diffusion-based code generators can rival top autoregressive code models on programming tasks while offering faster text generation.

That faster generation is the reason to care. An autoregressive model emits one token at a time and the cost grows with the length of the output, whereas a masked diffusion model can revise several positions in a single pass, which is the property that makes speculative, parallel decoding possible. The price is that the model does not have a natural left-to-right order, which is where the research questions come from.

The paper asks three of them. How does the generation pattern of a diffusion language model differ from that of an autoregressive model? What is the difference in modelling different data modalities, such as code against math? And how diverse can diffusion language models be, and how should post-training be designed in response?

The third question is the one this repository answers, and the first two are the ones its findings address.

The autoregressiveness score, and what the findings say

To answer the first two questions the authors needed a way to measure the thing they were asking about, so they introduced a new metric, the autoregressiveness score, to quantify the causal pattern during diffusion language model generation. In other words, instead of asking whether the output looks right, they measure how much the generation order resembles the left-to-right order of an autoregressive model.

The findings are short and each one is a statement about generation rather than about accuracy.

Diffusion language models still exhibit a left-to-right bias due to the nature of text, but they can also break this strict order in ways autoregressive models cannot, since an autoregressive model has no freedom in the order at all.

After pre-training, code tasks induce less global autoregressiveness than math. That is a nice result for the code case specifically: the structure of code is less linearly ordered than the structure of a mathematical derivation, so a diffusion model has more room to work with, and the training signal from code data is less contaminated by the left-to-right prior that text carries.

And the third is the one with practical consequences for anyone running inference. In diffusion language models, changing the sampling temperature does not only affect the sampled tokens, as in autoregressive models, but also alters the generation order itself. So temperature is not only a knob for randomness in a diffusion language model, it is a knob on the decoding strategy, and the same setting will not mean the same thing across model families.

Coupled-GRPO pairs timesteps and masks

The proposed post-training method is Coupled-GRPO, and the problem it solves is stated first.

In diffusion language models, the per-timestep loss typically computes log-probabilities only at masked token positions, which leads to inefficiency and high variance when sampling is limited. If a token is never unmasked during the sampled steps, it never contributes a gradient, and if it is unmasked in a context where nearly everything else is masked, the probability estimate is based on an unrepresentative situation.

The coupled-sampling scheme fixes both at once. For each training example, it selects a set of lambda pairs of timesteps, written as (t, t-hat), such that t plus t-hat equals T. It then applies two complementary token masks, each mask hides part of the tokens, and together they cover the entire set of target tokens. The consequence is that every token is unmasked in exactly one of the two forward passes.

The three benefits follow directly. Every token's log-probability is computed at least once, which provides a non-zero learning signal for all tokens. The probability estimates are more accurate, since each token is evaluated in a realistic partially-masked context rather than always being fully masked. And the scheme uses 2 lambda times more sampling passes than the baseline, with lambda set to 1 in the paper's configuration, so the extra cost is one additional forward pass.

This is the kind of contribution that is small to describe and hard to invent, and the reason the repository ships the implementation rather than only the paper.

The implementation is open-r1, patched in specific places

The repository is built upon the open-r1 project from Hugging Face, and the file list in the README is the clearest statement of what was changed.

├── run.sh # start training
├── setup.py # modified open-r1/setup.py
├── src/open_r1/ #  our code based on open-r1
│   ├── configs.py # with diffusion related params
│   ├── coupled_grpo.py # inherits trl GRPOTrainer
│   ├── grpo.py # main training script
│   ├── rewards.py # rewrite code reward and code_format reward
│   ├── utils/code_providers.py # rewrite pass rate extraction for E2B
├── recipes/process_data.py # prepare grpo training data
├── recipes/config_coupled_code.yaml # training config
├── tests/test_code_reward.py # test sandbox execution for code

So the changes are narrow and each one is a comment in the file name. configs.py adds the diffusion related parameters, coupled_grpo.py is a new trainer that inherits the GRPOTrainer from the trl library, grpo.py is the main training script, rewards.py rewrites the code reward and the code format reward, and the code providers module rewrites the pass rate extraction for E2B. The training configuration lives in recipes/config_coupled_code.yaml, and the data preparation in recipes/process_data.py.

That shape also tells you how to use it. The getting started section says to clone the source of open-r1, and then to merge and replace files between theirs and open-r1's, including setup.py, and to set up the environment following open-r1. The repository is a set of patches over an upstream stack rather than a package you install on its own, which is normal for research code and is worth knowing before you start.

The three training steps

The environment setup is a conda environment on Python 3.11 with a set of pinned installs:

env=openr1
conda create -n $env python=3.11 -y -c anaconda
conda activate $env

pip install vllm==0.8.4
pip install setuptools
pip install flash-attn==2.8.0.post1 --no-build-isolation
pip install -e ".[code]"

The environment also needs a code sandbox at E2B, with the token exported as E2B_API_KEY, and a wandb login with the WANDB_ENTITY environment variable exported. The sandbox matters more than it looks, because the code reward works by running the generated code and measuring whether it passes.

The data preparation step builds a hard split of the GRPO training data from the AceCode-89k dataset:

cd recipes
python process_data.py --dataset_path "TIGER-Lab/AceCode-89K" --output_path "./acecode_hard.jsonl" --difficulty "hard"

And the training itself is started with:

cd ..
bash run.sh

with a note that run.sh starts the e2b server locally, but you can also run it on CPU clusters. That note is about where the sandbox runs, which is the part of the setup that has to fit your infrastructure rather than your laptop.

The dependency list in setup.py explains the shape of the stack: accelerate, bitsandbytes, datasets, deepspeed, distilabel with the vllm, ray and openai extras, e2b-code-interpreter, einops, huggingface-hub, liger-kernel, lighteval from git, math-verify and more. The file also carries an explicit instruction that all dependencies should be listed with their version requirements, and that a fast-moving dependency such as trl should be pinned to the exact version.

Inference, and the parts that are still moving

The three models are on Hugging Face, and the inference example is a standard transformers load with a bfloat16 dtype and the model path given as apple/DiffuCoder-7B-Base. There is one knob the README calls out, TOKEN_PER_STEP, which trades performance against speed.

That is the whole inference story in the documentation available here, and the repository top-level listing backs it up: an inference_demo.py, a recipes/ directory, run.sh, setup.py, src/, tests/, a figs/ directory with the paper's figures, an ACKNOWLEDGEMENTS file and the LICENSE.

The updates section is a timeline of things that were in progress, and it is worth reading as a picture of the project's state. Code was available on July 1, 2025, the models on Hugging Face on July 2, and the inference usage, examples and demo were updated on July 4. On July 4 MLX support was announced as in progress, with a link to a community implementation of an 8bit DiffuCoder model and to a pull request in the mlx-lm repository for previewing or contributing. On July 8 the authors noted they were still training 14B models.

The top of the README says MLX support on Apple Silicon is in progress and that the necessary updates will be made once it is available. The last push to the repository was on 2026-09-11, so the tree has moved since those notes were written, and the documentation in the repository is the July 2025 state of the instructions.

For anyone planning to use this, the practical summary is that the training recipe is complete and reproducible, the models are published in three variants, and the Apple Silicon path was still being finished when the README was last written.

Editorial conclusion

Adopt this repository if you are working on diffusion language models and want a concrete post-training recipe for code, because Coupled-GRPO is described precisely enough to reimplement, with the code, the training config and a test for the code reward shipped alongside the models. Do not adopt it as a general training stack, because it is a patched copy of open-r1 with pinned dependencies, and the setup instructions expect you to merge files with that upstream project rather than install this one. Before you start, read the model cards on Hugging Face for the three released checkpoints, check the dependency pins in setup.py, because the comment there is explicit that fast-moving packages are pinned exactly, and plan on the sandbox the training depends on, since the code reward executes generated code in E2B.

Frequently asked questions

Official sources

1. apple-aiml-research/ml-diffucoder on GitHub
2. Issues
3. README