从MiniMind-LLM微调扩散语言模型

[jingyaogong]

SFT AR-2-MaskedDiffusion Model

dLLM（扩散大语言模型） 最近挺火的，然而话说在前，它几乎不在我的研究范围，未来大概也不会follow。

这里写了极简实现——用最短的改动，把现有的 AR MiniMind模型续训转换成（还不太具有逻辑的）dLLM

之所以存在很多特性（包括本内容）不合并到主分支：

1. dLLM又热门又小众：适合特定人群单独研究
2. 不会长期维护：除非有更多契机，否则这个 Discussion 即可算作中间 "快照"
3. 维持简洁性：一些Toy级别的尝试单独分出，不让纯净主仓库变得臃肿，保持极简、更极简、更更极简的目标

这个 Discussion 能帮你做什么？

• 用 ~300 行代码 把 MiniMind 变成 DLLM
• 复用 AR 模型的权重，无需从头训练
• 提供 微调训练 + 评估 + Web Demo 脚本

核心改动唯3

• 把 AR 的单向 Causal Attention → 双向 Bidirectional Attention
• 加上 MDM 风格的 Mask 训练目标。
• 采样时用 MDM 风格的 Mask

---

完整代码

稍后按照下面目录树结构进行创建即可

model_minimind_dllm.py

import math, torch, torch.nn.functional as F
from torch import nn
from transformers.activations import ACT2FN
from transformers import PreTrainedModel, GenerationMixin, PretrainedConfig
from transformers.modeling_outputs import MaskedLMOutput

# ##################################################################################
#                                     MiniMind Config
# ##################################################################################
class MiniMindDLLMConfig(PretrainedConfig):
    model_type = "minimind_dllm"
    def __init__(self, hidden_size=512, num_hidden_layers=8, use_moe=False, **kwargs):
        super().__init__(**kwargs)
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.use_moe = use_moe
        self.dropout = kwargs.get("dropout", 0.0)
        self.vocab_size = kwargs.get("vocab_size", 6400)
        self.bos_token_id = kwargs.get("bos_token_id", 1)
        self.eos_token_id = kwargs.get("eos_token_id", 2)
        self.flash_attn = kwargs.get("flash_attn", True)
        self.num_attention_heads = kwargs.get("num_attention_heads", 8)
        self.num_key_value_heads = kwargs.get("num_key_value_heads", 2)
        self.head_dim = kwargs.get("head_dim", self.hidden_size // self.num_attention_heads)
        self.hidden_act = kwargs.get("hidden_act", 'silu')
        self.intermediate_size = kwargs.get("intermediate_size", math.ceil(hidden_size * 8/3 / 64) * 64)
        self.max_position_embeddings = kwargs.get("max_position_embeddings", 32768)
        self.rms_norm_eps = kwargs.get("rms_norm_eps", 1e-6)
        self.rope_theta = kwargs.get("rope_theta", 1e6)
        self.inference_rope_scaling = kwargs.get("inference_rope_scaling", False)
        self.rope_scaling = {
            "beta_fast": 32,
            "beta_slow": 1,
            "factor": 16,
            "original_max_position_embeddings": 2048,
            "attention_factor": 1.0,
            "type": "yarn"
        } if self.inference_rope_scaling else None
        ### MoE specific configs (ignored if use_moe = False)
        self.num_experts = kwargs.get("num_experts", 4)
        self.num_experts_per_tok = kwargs.get("num_experts_per_tok", 1)
        self.moe_intermediate_size = kwargs.get("moe_intermediate_size", self.intermediate_size // 2)
        self.norm_topk_prob = kwargs.get("norm_topk_prob", self.num_experts_per_tok > 1)
        self.router_aux_loss_coef = kwargs.get("router_aux_loss_coef", 0.005)
        self.mask_token_id = kwargs.get("mask_token_id", self.vocab_size - 1)  # DLLM: [MASK] token
        self.mask_epsilon = kwargs.get("mask_epsilon", 0.001)  # DLLM: 最小 mask 概率

# ##################################################################################
#                                     MiniMind Model
# ##################################################################################
class RMSNorm(torch.nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def norm(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)

    def forward(self, x):
        return self.weight * self.norm(x.float()).type_as(x)

def precompute_freqs_cis(dim: int, end: int = int(32 * 1024), rope_base: float = 1e6, rope_scaling: dict = None):
    freqs, attn_factor = 1.0 / (rope_base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)), 1.0
    if rope_scaling is not None: # YaRN: f'(i) = f(i)((1-γ) + γ/s), where γ∈[0,1] is linear ramp
        orig_max, factor, beta_fast, beta_slow, attn_factor = (
            rope_scaling.get("original_max_position_embeddings", 2048), rope_scaling.get("factor", 16),
            rope_scaling.get("beta_fast", 32.0), rope_scaling.get("beta_slow", 1.0), rope_scaling.get("attention_factor", 1.0)
        )
        if end / orig_max > 1.0:
            inv_dim = lambda b: (dim * math.log(orig_max / (b * 2 * math.pi))) / (2 * math.log(rope_base))
            low, high = max(math.floor(inv_dim(beta_fast)), 0), min(math.ceil(inv_dim(beta_slow)), dim // 2 - 1)
            ramp = torch.clamp((torch.arange(dim // 2, device=freqs.device).float() - low) / max(high - low, 0.001), 0, 1)
            freqs = freqs * (1 - ramp + ramp / factor)
    t = torch.arange(end, device=freqs.device)
    freqs = torch.outer(t, freqs).float()
    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attn_factor
    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attn_factor
    return freqs_cos, freqs_sin

def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
    def rotate_half(x): return torch.cat((-x[..., x.shape[-1] // 2:], x[..., : x.shape[-1] // 2]), dim=-1)
    q_embed = (q * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(q) * sin.unsqueeze(unsqueeze_dim))
    k_embed = (k * cos.unsqueeze(unsqueeze_dim)) + (rotate_half(k) * sin.unsqueeze(unsqueeze_dim))
    return q_embed, k_embed

def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    bs, slen, num_key_value_heads, head_dim = x.shape
    if n_rep == 1: return x
    return (
        x[:, :, :, None, :].expand(bs, slen, num_key_value_heads, n_rep, head_dim).reshape(bs, slen, num_key_value_heads * n_rep, head_dim)
    )

class Attention(nn.Module):
    def __init__(self, config: MiniMindDLLMConfig):
        super().__init__()
        self.num_key_value_heads = config.num_attention_heads if config.num_key_value_heads is None else config.num_key_value_heads
        self.n_local_heads = config.num_attention_heads
        self.n_local_kv_heads = self.num_key_value_heads
        self.n_rep = self.n_local_heads // self.n_local_kv_heads
        self.head_dim = config.head_dim
        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
        # self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        # self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)
        self.dropout = config.dropout
        self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention') and config.flash_attn

    def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
        bsz, seq_len, _ = x.shape
        xq, xk, xv = self.q_proj(x), self.k_proj(x), self.v_proj(x)
        xq = xq.view(bsz, seq_len, self.n_local_heads, self.head_dim)
        xk = xk.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
        xv = xv.view(bsz, seq_len, self.n_local_kv_heads, self.head_dim)
        # xq, xk = self.q_norm(xq), self.k_norm(xk)
        cos, sin = position_embeddings
        xq, xk = apply_rotary_pos_emb(xq, xk, cos[:seq_len], sin[:seq_len])
        if past_key_value is not None:
            xk = torch.cat([past_key_value[0], xk], dim=1)
            xv = torch.cat([past_key_value[1], xv], dim=1)
        past_kv = (xk, xv) if use_cache else None
        xq, xk, xv = (xq.transpose(1, 2), repeat_kv(xk, self.n_rep).transpose(1, 2), repeat_kv(xv, self.n_rep).transpose(1, 2))
        if self.flash and seq_len > 1 and (attention_mask is None or torch.all(attention_mask == 1)):
            output = F.scaled_dot_product_attention(xq, xk, xv, dropout_p=self.dropout if self.training else 0.0, is_causal=False)  # DLLM: 双向注意力
        else:
            scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
            # DLLM: 移除 causal mask，改为双向注意力
            if attention_mask is not None: scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
            output = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv
        output = output.transpose(1, 2).reshape(bsz, seq_len, -1)
        output = self.resid_dropout(self.o_proj(output))
        return output, past_kv

class FeedForward(nn.Module):
    def __init__(self, config: MiniMindDLLMConfig, intermediate_size: int = None):
        super().__init__()
        intermediate_size = intermediate_size or config.intermediate_size
        self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
        self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))

class MOEFeedForward(nn.Module):
    def __init__(self, config: MiniMindDLLMConfig):
        super().__init__()
        self.config = config
        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
        self.experts = nn.ModuleList([
            FeedForward(config, intermediate_size=config.moe_intermediate_size)
            for _ in range(config.num_experts)
        ])
        self.act_fn = ACT2FN[config.hidden_act]

    def forward(self, x):
        batch_size, seq_len, hidden_dim = x.shape
        x_flat = x.view(-1, hidden_dim)
        scores = F.softmax(self.gate(x_flat), dim=-1)
        topk_weight, topk_idx = torch.topk(scores, k=self.config.num_experts_per_tok, dim=-1, sorted=False)
        if self.config.norm_topk_prob: topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
        y = torch.zeros_like(x_flat)
        for i, expert in enumerate(self.experts):
            mask = (topk_idx == i)
            if mask.any():
                token_idx = mask.any(dim=-1).nonzero().flatten()
                weight = topk_weight[mask].view(-1, 1)
                y.index_add_(0, token_idx, expert(x_flat[token_idx]) * weight)
            elif self.training:
                y[0, 0] += 0 * sum(p.sum() for p in expert.parameters())
        if self.training and self.config.router_aux_loss_coef > 0:
            load = F.one_hot(topk_idx, self.config.num_experts).float().mean(0)
            self.aux_loss = (load * scores.mean(0)).sum() * self.config.num_experts * self.config.router_aux_loss_coef
        else:
            self.aux_loss = scores.new_zeros(1).squeeze()
        return y.view(batch_size, seq_len, hidden_dim)

class MiniMindBlock(nn.Module):
    def __init__(self, layer_id: int, config: MiniMindDLLMConfig):
        super().__init__()
        self.self_attn = Attention(config)
        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.mlp = FeedForward(config) if not config.use_moe else MOEFeedForward(config)

    def forward(self, hidden_states, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
        residual = hidden_states
        hidden_states, present_key_value = self.self_attn(
            self.input_layernorm(hidden_states), position_embeddings,
            past_key_value, use_cache, attention_mask
        )
        hidden_states += residual
        hidden_states = hidden_states + self.mlp(self.post_attention_layernorm(hidden_states))
        return hidden_states, present_key_value

class MiniMindDLLMModel(nn.Module):
    def __init__(self, config: MiniMindDLLMConfig):
        super().__init__()
        self.config = config
        self.vocab_size, self.num_hidden_layers = config.vocab_size, config.num_hidden_layers
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.dropout = nn.Dropout(config.dropout)
        self.layers = nn.ModuleList([MiniMindBlock(l, config) for l in range(self.num_hidden_layers)])
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        freqs_cos, freqs_sin = precompute_freqs_cis(dim=config.head_dim, end=config.max_position_embeddings, rope_base=config.rope_theta, rope_scaling=config.rope_scaling)
        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
        self.register_buffer("freqs_sin", freqs_sin, persistent=False)

    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, **kwargs):
        batch_size, seq_length = input_ids.shape
        if hasattr(past_key_values, 'layers'): past_key_values = None
        past_key_values = past_key_values or [None] * len(self.layers)
        start_pos = past_key_values[0][0].shape[1] if past_key_values[0] is not None else 0
        hidden_states = self.dropout(self.embed_tokens(input_ids))
        position_embeddings = (
            self.freqs_cos[start_pos:start_pos + seq_length],
            self.freqs_sin[start_pos:start_pos + seq_length]
        )
        presents = []
        for layer, past_key_value in zip(self.layers, past_key_values):
            hidden_states, present = layer(
                hidden_states,
                position_embeddings,
                past_key_value=past_key_value,
                use_cache=use_cache,
                attention_mask=attention_mask
            )
            presents.append(present)
        hidden_states = self.norm(hidden_states)
        aux_loss = sum([l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, MOEFeedForward)], hidden_states.new_zeros(1).squeeze())
        return hidden_states, presents, aux_loss

    def add_noise_to_tokens(self, input_ids: torch.LongTensor, t: torch.FloatTensor, eps: float = None, pad_token_id: int = 0):
        """DLLM: MDM-style masking, 根据噪声水平 t 将 token 替换为 [MASK]"""
        batch_size, seq_len = input_ids.shape
        eps = eps if eps is not None else self.config.mask_epsilon
        p_mask = (1 - eps) * t + eps  # t=0→eps, t=1→1
        p_mask = p_mask.unsqueeze(-1).expand(batch_size, seq_len)
        corruption_mask = torch.rand(batch_size, seq_len, device=input_ids.device) < p_mask
        corruption_mask = corruption_mask & (input_ids != pad_token_id)  # 不 mask padding
        noisy_input_ids = torch.where(corruption_mask, self.config.mask_token_id, input_ids)
        return noisy_input_ids, corruption_mask, p_mask

class MiniMindForMaskedDiffusion(PreTrainedModel, GenerationMixin):
    config_class = MiniMindDLLMConfig
    def __init__(self, config: MiniMindDLLMConfig = None):
        self.config = config or MiniMindDLLMConfig()
        super().__init__(self.config)
        self.model = MiniMindDLLMModel(self.config)
        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
        self.model.embed_tokens.weight = self.lm_head.weight
    
    def forward(self, input_ids, attention_mask=None, past_key_values=None, use_cache=False, labels=None, corruption_mask=None, p_mask=None, **kwargs):
        hidden_states, past_key_values, aux_loss = self.model(input_ids, attention_mask, past_key_values, use_cache, **kwargs)
        logits = self.lm_head(hidden_states).float()
        loss = None
        if labels is not None and corruption_mask is not None and p_mask is not None:
            loss = F.cross_entropy(logits.view(-1, self.config.vocab_size), labels.view(-1), reduction='none').view(labels.shape)
            loss = (loss[corruption_mask] / p_mask[corruption_mask]).sum() / labels.numel()  # DLLM: 加权 loss
        return MaskedLMOutput(loss=loss, logits=logits, hidden_states=hidden_states)

    def add_noise_to_tokens(self, input_ids, t, eps=None, pad_token_id=0):
        return self.model.add_noise_to_tokens(input_ids, t, eps, pad_token_id)
    
    @torch.inference_mode() 
    def generate(self, inputs, max_new_tokens=128, temperature=1.0, top_k=0, steps=None, eos_token_id=None, tokenizer=None, **kwargs):
        """DLLM Block-wise 生成: 每次生成一个block，逐块去噪"""
        input_ids = kwargs.get("input_ids", inputs)
        bsz, prompt_len, device = input_ids.shape[0], input_ids.shape[1], input_ids.device
        mask_id, eos_id = self.config.mask_token_id, eos_token_id or self.config.eos_token_id
        block_size = kwargs.get("block_size", 16)
        steps_per_block = steps or 8
        
        x = input_ids.clone()
        generated_len = 0
        
        while generated_len < max_new_tokens:
            cur_block = min(block_size, max_new_tokens - generated_len)
            block_mask = torch.full((bsz, cur_block), mask_id, dtype=torch.long, device=device)
            x = torch.cat([x, block_mask], dim=1)
            block_start = x.shape[1] - cur_block
            
            for step in range(steps_per_block):
                mask_pos = (x[:, block_start:] == mask_id)
                if not mask_pos.any(): break
                
                logits = self(input_ids=x).logits[:, block_start:, :] / temperature
                if top_k > 0: logits[logits < torch.topk(logits, top_k, dim=-1)[0][..., -1:]] = -float('inf')
                
                probs = F.softmax(logits, dim=-1)
                sampled = torch.multinomial(probs.view(-1, probs.shape[-1]), 1).view(bsz, -1)
                conf = torch.gather(probs, -1, sampled.unsqueeze(-1)).squeeze(-1)
                conf = torch.where(mask_pos, conf, -1.0)
                
                n_fill = cur_block if step == steps_per_block - 1 else max(1, cur_block // steps_per_block)
                _, top_idx = conf.topk(min(n_fill, mask_pos.sum(-1).min().item()), dim=-1)
                for b in range(bsz):
                    x[b, block_start + top_idx[b]] = sampled[b, top_idx[b]]
            
            generated_len += cur_block
            if tokenizer and kwargs.get("stream", False):
                text = tokenizer.decode(x[0, prompt_len:], skip_special_tokens=False)
                print(f"[Block {generated_len//block_size}] {text}")
            
            if eos_id is not None and (x[:, prompt_len:] == eos_id).any(dim=-1).all():
                break
        
        if eos_id is not None:
            for b in range(bsz):
                eos_pos = (x[b, prompt_len:] == eos_id).nonzero(as_tuple=True)[0]
                if len(eos_pos) > 0:
                    x = x[:, :prompt_len + eos_pos[0].item() + 1]
                    break
        return x

train_full_sft_dllm.py

import os
import sys

__package__ = "trainer"
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import argparse
import time
import warnings
import torch
import torch.distributed as dist
from contextlib import nullcontext
from torch import optim
from torch.nn.parallel import DistributedDataParallel
from torch.utils.data import DataLoader, DistributedSampler
from model.model_minimind_dllm import MiniMindDLLMConfig, MiniMindForMaskedDiffusion
from dataset.lm_dataset import SFTDataset
from trainer.trainer_utils import get_lr, Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler
from transformers import AutoTokenizer

warnings.filterwarnings('ignore')


def init_model_dllm(config, from_weight, device):
    """初始化 DLLM 模型，可从 AR 模型权重加载"""
    tokenizer = AutoTokenizer.from_pretrained('../model', local_files_only=True)
    model = MiniMindForMaskedDiffusion(config)
    if from_weight and from_weight != 'none':
        moe_suffix = '_moe' if config.use_moe else ''
        ckp = f'./out/{from_weight}_{config.hidden_size}{moe_suffix}.pth'
        if os.path.exists(ckp):
            state_dict = torch.load(ckp, map_location=device)
            model.load_state_dict(state_dict, strict=False)
            Logger(f'[DLLM] 从 {ckp} 加载权重')
    return model.to(device), tokenizer


def train_epoch(epoch, loader, iters, model, optimizer, scaler, autocast_ctx, args, lm_config, pad_token_id, start_step=0, wandb=None):
    start_time = time.time()
    for step, (X, Y) in enumerate(loader, start=start_step + 1):
        X = X.to(args.device)
        Y = Y.to(args.device)
        loss_mask = (Y != -100)  # labels != -100 的位置才需要算 loss
        lr = get_lr(epoch * iters + step, args.epochs * iters, args.learning_rate)
        for param_group in optimizer.param_groups:
            param_group['lr'] = lr

        with autocast_ctx:
            t = torch.rand(X.shape[0], device=args.device)
            noisy_X, corruption_mask, p_mask = model.add_noise_to_tokens(X, t, pad_token_id=pad_token_id)
            
            corruption_mask = corruption_mask & loss_mask  # 只保留 assistant 部分被 mask 的
            
            res = model(input_ids=noisy_X, labels=X, corruption_mask=corruption_mask, p_mask=p_mask)
            loss = res.loss / args.accumulation_steps

        scaler.scale(loss).backward()

        if (step + 1) % args.accumulation_steps == 0:
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip)
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad(set_to_none=True)

        if step % args.log_interval == 0 or step == iters - 1:
            spend_time = time.time() - start_time
            current_loss = loss.item() * args.accumulation_steps
            current_lr = optimizer.param_groups[-1]['lr']
            eta_min = spend_time / (step + 1) * iters // 60 - spend_time // 60
            mask_ratio = corruption_mask.float().mean().item()
            
            Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), loss: {current_loss:.4f}, mask_ratio: {mask_ratio:.2f}, lr: {current_lr:.8f}, eta: {eta_min:.1f}min')
            
            if wandb: wandb.log({"loss": current_loss, "mask_ratio": mask_ratio, "learning_rate": current_lr})

        if (step % args.save_interval == 0 or step == iters - 1) and is_main_process():
            model.eval()
            moe_suffix = '_moe' if lm_config.use_moe else ''
            ckp = f'{args.save_dir}/{args.save_weight}_{lm_config.hidden_size}{moe_suffix}.pth'
            state_dict = model.module.state_dict() if isinstance(model, DistributedDataParallel) else model.state_dict()
            state_dict = {k: v.half().cpu() for k, v in state_dict.items()}
            torch.save(state_dict, ckp)
            model.train()
            del state_dict

        del X, Y, noisy_X, corruption_mask, p_mask, res, loss


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="MiniMind DLLM Full SFT")
    parser.add_argument("--save_dir", type=str, default="../out", help="模型保存目录")
    parser.add_argument('--save_weight', default='dllm_sft', type=str)
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument("--learning_rate", type=float, default=5e-5)
    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu")
    parser.add_argument("--dtype", type=str, default="bfloat16")
    parser.add_argument("--num_workers", type=int, default=8)
    parser.add_argument("--accumulation_steps", type=int, default=1)
    parser.add_argument("--grad_clip", type=float, default=1.0)
    parser.add_argument("--log_interval", type=int, default=100)
    parser.add_argument("--save_interval", type=int, default=1000)
    parser.add_argument('--hidden_size', default=512, type=int)
    parser.add_argument('--num_hidden_layers', default=8, type=int)
    parser.add_argument('--max_seq_len', default=300, type=int)
    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1])
    parser.add_argument("--data_path", type=str, default="../dataset/sft_mini_512.jsonl")
    parser.add_argument('--from_weight', default='full_sft', type=str, help="从AR模型权重初始化")
    parser.add_argument("--use_wandb", action="store_true")
    parser.add_argument("--wandb_project", type=str, default="MiniMind-DLLM-SFT")
    args = parser.parse_args()

    local_rank = init_distributed_mode()
    if dist.is_initialized(): args.device = f"cuda:{local_rank}"
    setup_seed(42 + (dist.get_rank() if dist.is_initialized() else 0))
    
    os.makedirs(args.save_dir, exist_ok=True)
    lm_config = MiniMindDLLMConfig(hidden_size=args.hidden_size, num_hidden_layers=args.num_hidden_layers, use_moe=bool(args.use_moe))
    
    device_type = "cuda" if "cuda" in args.device else "cpu"
    dtype = torch.bfloat16 if args.dtype == "bfloat16" else torch.float16
    autocast_ctx = nullcontext() if device_type == "cpu" else torch.cuda.amp.autocast(dtype=dtype)
    
    wandb = None
    if args.use_wandb and is_main_process():
        import swanlab as wandb
        wandb.init(project=args.wandb_project, name=f"DLLM-SFT-{args.hidden_size}")
    
    model, tokenizer = init_model_dllm(lm_config, args.from_weight, device=args.device)
    train_ds = SFTDataset(args.data_path, tokenizer, max_length=args.max_seq_len)
    train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None
    scaler = torch.cuda.amp.GradScaler(enabled=(args.dtype == 'float16'))
    optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate)
    
    if dist.is_initialized():
        model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
        model = DistributedDataParallel(model, device_ids=[local_rank])
    
    for epoch in range(args.epochs):
        train_sampler and train_sampler.set_epoch(epoch)
        loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=(train_sampler is None), sampler=train_sampler, num_workers=args.num_workers, pin_memory=True)
        train_epoch(epoch, loader, len(loader), model, optimizer, scaler, autocast_ctx, args, lm_config, tokenizer.pad_token_id, 0, wandb)
    
    if dist.is_initialized(): dist.destroy_process_group()

eval_dllm.py

import argparse
import warnings
import torch
from transformers import AutoTokenizer, TextStreamer
from model.model_minimind_dllm import MiniMindDLLMConfig, MiniMindForMaskedDiffusion
from trainer.trainer_utils import setup_seed, get_model_params
warnings.filterwarnings('ignore')

def init_model(args):
    tokenizer = AutoTokenizer.from_pretrained(args.load_from)
    model = MiniMindForMaskedDiffusion(MiniMindDLLMConfig(
        hidden_size=args.hidden_size,
        num_hidden_layers=args.num_hidden_layers,
        use_moe=bool(args.use_moe),
    ))
    moe_suffix = '_moe' if args.use_moe else ''
    ckp = f'./{args.save_dir}/{args.weight}_{args.hidden_size}{moe_suffix}.pth'
    state_dict = torch.load(ckp, map_location=args.device)
    model.load_state_dict(state_dict, strict=False)
    print(f"[加载权重] {ckp}")
    get_model_params(model, model.config)
    return model.eval().to(args.device), tokenizer

def main():
    parser = argparse.ArgumentParser(description="MiniMind DLLM 推理")
    parser.add_argument('--load_from', default='model', type=str)
    parser.add_argument('--save_dir', default='out', type=str)
    parser.add_argument('--weight', default='dllm_sft', type=str)
    parser.add_argument('--hidden_size', default=512, type=int)
    parser.add_argument('--num_hidden_layers', default=8, type=int)
    parser.add_argument('--use_moe', default=0, type=int)
    parser.add_argument('--max_new_tokens', default=64, type=int)
    parser.add_argument('--steps', default=8, type=int, help="每个block的去噪步数")
    parser.add_argument('--block_size', default=16, type=int, help="每个block的token数")
    parser.add_argument('--temperature', default=1.0, type=float)
    parser.add_argument('--top_k', default=50, type=int)
    parser.add_argument('--stream', default=0, type=int)
    parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str)
    args = parser.parse_args()
    
    prompts = [
        '你有什么特长？',
        '为什么天空是蓝色的',
        '请用Python写一个计算斐波那契数列的函数',
        '解释一下"光合作用"的基本过程',
    ]
    
    model, tokenizer = init_model(args)
    input_mode = int(input('[0] 自动测试\n[1] 手动输入\n'))
    
    prompt_iter = prompts if input_mode == 0 else iter(lambda: input('👶: '), '')
    for prompt in prompt_iter:
        setup_seed(1234)
        if input_mode == 0: print(f'👶: {prompt}')
        
        messages = [{"role": "user", "content": prompt}]
        inputs = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer(inputs, return_tensors="pt", truncation=True).to(args.device)
        
        generated = model.generate(
            inputs=inputs["input_ids"], 
            max_new_tokens=args.max_new_tokens, 
            steps=args.steps,
            block_size=args.block_size,
            temperature=args.temperature, 
            top_k=args.top_k, 
            eos_token_id=tokenizer.eos_token_id, 
            tokenizer=tokenizer,
            stream=bool(args.stream)
        )
        print(f"🤖️: {tokenizer.decode(generated[0, inputs['input_ids'].shape[1]:], skip_special_tokens=True)}\n")

if __name__ == "__main__":
    main()

dllm_demo.py

import sys, os, json, time, torch, torch.nn.functional as F
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from flask import Flask, render_template_string, request, Response
from transformers import AutoTokenizer
from model.model_minimind_dllm import MiniMindDLLMConfig, MiniMindForMaskedDiffusion

app = Flask(__name__)
model, tokenizer, device = None, None, 'cpu'

HTML = '''<!DOCTYPE html><html><head><meta charset="UTF-8"><title>DLLM Demo</title>
<style>
*{box-sizing:border-box;margin:0;padding:0}
body{font-family:-apple-system,sans-serif;background:#fff;min-height:100vh;display:flex;flex-direction:column}
.header{padding:24px;display:flex;align-items:center;justify-content:center;gap:14px;border-bottom:1px solid #f0f0f0}
.header img{height:42px} .header span{font-size:22px;font-weight:600;color:#333;font-style:italic}
.main{flex:1;max-width:800px;width:100%;margin:0 auto;padding:24px}
.step-info{color:#bbb;font-size:11px;margin-bottom:16px;text-transform:uppercase}
.tokens{display:flex;flex-wrap:wrap;gap:4px;line-height:2}
.token{padding:4px 10px;border-radius:6px;font-size:14px}
.token.mask{background:#ff6b6b;color:#fff;animation:pulse .5s infinite alternate}
.token.hidden{background:#e0e0e0;color:#999}
.token.filled{background:#f5f5f5;color:#444}
.token.new{background:#51cf66;color:#fff}
@keyframes pulse{from{opacity:.6}to{opacity:1}}
.bottom{padding:20px 24px;border-top:1px solid #f0f0f0}
.bottom-inner{max-width:800px;margin:0 auto}
.input-area{display:flex;gap:12px;margin-bottom:12px}
input[type=text]{flex:1;padding:12px 16px;border-radius:8px;border:1px solid #e0e0e0;background:#f8f8f8;font-size:15px}
button{padding:12px 28px;border-radius:8px;border:none;cursor:pointer;background:#333;color:#fff;font-size:14px}
button:hover{background:#444}
.params{display:flex;gap:16px;font-size:12px;color:#999}
.params label{display:flex;align-items:center;gap:6px}
.params input{width:80px;padding:6px 10px;border-radius:6px;border:1px solid #e0e0e0;background:#f8f8f8}
</style></head>
<body>
<div class="header">
    <img src="https://github.com/jingyaogong/minimind/raw/master/images/logo.png"><span>DLLM Demo</span>
</div>
<div class="main">
    <div class="step-info" id="stepInfo">等待生成...</div>
    <div class="tokens" id="tokens"></div>
</div>
<div class="bottom"><div class="bottom-inner">
    <div class="input-area">
        <input type="text" id="prompt" placeholder="输入问题..." value="你有什么特长？">
        <button onclick="generate()">生成</button>
    </div>
    <div class="params">
        <label>Block <input type="number" id="blockSize" value="64"></label>
        <label>Steps <input type="number" id="steps" value="32"></label>
        <label>Max <input type="number" id="maxTokens" value="64"></label>
    </div>
</div></div>
<script>
let prev=[];
async function generate(){
    const p=document.getElementById('prompt').value, b=document.getElementById('blockSize').value;
    const s=document.getElementById('steps').value, m=document.getElementById('maxTokens').value;
    document.getElementById('tokens').innerHTML='';
    document.getElementById('stepInfo').textContent='正在生成...';
    prev=[];
    const r=await fetch('/generate',{method:'POST',headers:{'Content-Type':'application/json'},
        body:JSON.stringify({prompt:p,block_size:+b,steps:+s,max_tokens:+m})});
    const reader=r.body.getReader(), dec=new TextDecoder();
    while(1){
        const{done,value}=await reader.read();
        if(done)break;
        dec.decode(value).split('\\n').filter(l=>l.startsWith('data:')).forEach(l=>render(JSON.parse(l.slice(5))));
    }
}
function render(d){
    document.getElementById('stepInfo').textContent=`Block ${d.block} | Step ${d.step}/${d.total_steps}`;
    const c=document.getElementById('tokens');
    c.innerHTML='';
    d.tokens.forEach((t,i)=>{
        const s=document.createElement('span');
        s.className='token '+(t.hidden?'hidden':t.is_mask?'mask':'filled'+(!d.done&&prev[i]!==t.text?' new':''));
        s.textContent=t.hidden||t.is_mask?'[M]':t.text||' ';
        c.appendChild(s);
    });
    prev=d.tokens.map(t=>t.text);
}
</script></body></html>'''

def load_model():
    global model, tokenizer, device
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    tokenizer = AutoTokenizer.from_pretrained('../model')
    model = MiniMindForMaskedDiffusion(MiniMindDLLMConfig(hidden_size=512, num_hidden_layers=8))
    ckp = '../out/dllm_sft_512.pth'
    if os.path.exists(ckp): model.load_state_dict(torch.load(ckp, map_location=device), strict=False)
    model = model.eval().to(device)
    print(f"[Model loaded] device={device}")

def make_tokens(x, prompt_len, mask_id, eos_id):
    tokens, found_eos = [], False
    for t in x[0, prompt_len:]:
        tid = t.item()
        if tid == eos_id: found_eos = True
        if found_eos and tid != eos_id: tokens.append({'is_mask': True, 'hidden': True, 'text': ''})
        else: tokens.append({'is_mask': tid == mask_id, 'text': '' if tid == mask_id else tokenizer.decode([tid])})
    return tokens

@app.route('/')
def index(): return render_template_string(HTML)

@app.route('/generate', methods=['POST'])
def generate():
    d = request.json
    prompt, block_size, steps, max_tokens = d.get('prompt', '你好'), d.get('block_size', 16), d.get('steps', 8), d.get('max_tokens', 64)
    def stream():
        input_ids = tokenizer(tokenizer.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True), return_tensors="pt").input_ids.to(device)
        mask_id, eos_id, x, prompt_len, generated_len, block_num = model.config.mask_token_id, tokenizer.eos_token_id, input_ids.clone(), input_ids.shape[1], 0, 0
        while generated_len < max_tokens:
            block_num += 1
            cur_block = min(block_size, max_tokens - generated_len)
            x = torch.cat([x, torch.full((1, cur_block), mask_id, dtype=torch.long, device=device)], dim=1)
            block_start = x.shape[1] - cur_block
            for step in range(steps):
                mask_pos = (x[:, block_start:] == mask_id)
                if not mask_pos.any(): break
                with torch.no_grad(): logits = model(input_ids=x).logits[:, block_start:, :]
                probs = F.softmax(logits, dim=-1)
                sampled = torch.multinomial(probs.view(-1, probs.shape[-1]), 1).view(1, -1)
                conf = torch.where(mask_pos, torch.gather(probs, -1, sampled.unsqueeze(-1)).squeeze(-1), -1.0)
                _, top_idx = conf.topk(min(cur_block if step == steps - 1 else max(1, cur_block // steps), mask_pos.sum().item()), dim=-1)
                x[0, block_start + top_idx[0]] = sampled[0, top_idx[0]]
                yield f"data:{json.dumps({'block': block_num, 'step': step+1, 'total_steps': steps, 'tokens': make_tokens(x, prompt_len, mask_id, eos_id)})}\n\n"
                time.sleep(0.05)
            generated_len += cur_block
            if eos_id and (x[:, prompt_len:] == eos_id).any(): break
        yield f"data:{json.dumps({'block': block_num, 'step': steps, 'total_steps': steps, 'tokens': make_tokens(x, prompt_len, mask_id, eos_id), 'done': True})}\n\n"
    return Response(stream(), mimetype='text/event-stream')

if __name__ == '__main__':
    load_model()
    print("🎲 DLLM Demo: http://localhost:5001")
    app.run(host='0.0.0.0', port=5001, debug=False, threaded=True)

lm_dataset.py

260115已合进master branch（无需更改，可跳过lm_dataset）

260115之前需要替换为：

> ⚠️ 注意，dataset本身格式和仓库中的一模一样，只改动了一处：把loss_mask放到forward中进行计算（目前临时替换即可，之后会把主分支也统一成这种方式，更简洁）

from torch.utils.data import Dataset
import torch
import os
from datasets import load_dataset
os.environ["TOKENIZERS_PARALLELISM"] = "false"

class PretrainDataset(Dataset):
    def __init__(self, data_path, tokenizer, max_length=512):
        super().__init__()
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.samples = load_dataset('json', data_files=data_path, split='train')

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, index):
        sample = self.samples[index]
        encoding = self.tokenizer(
            str(sample['text']),
            max_length=self.max_length,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        input_ids = encoding.input_ids.squeeze()
        labels = input_ids.clone()
        labels[input_ids == self.tokenizer.pad_token_id] = -100
        return input_ids, labels


class SFTDataset(Dataset):
    def __init__(self, jsonl_path, tokenizer, max_length=1024):
        super().__init__()
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.samples = load_dataset('json', data_files=jsonl_path, split='train')
        self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant', add_special_tokens=False).input_ids
        self.eos_id = tokenizer(f'{tokenizer.eos_token}', add_special_tokens=False).input_ids

    def __len__(self):
        return len(self.samples)

    def create_chat_prompt(self, cs):
        messages = cs.copy()
        tools = cs[0]["functions"] if (cs and cs[0]["role"] == "system" and cs[0].get("functions")) else None
        return self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=False,
            tools=tools
        )

    def generate_labels(self, input_ids):
        labels = [-100] * len(input_ids)
        i = 0
        while i < len(input_ids):
            if input_ids[i:i + len(self.bos_id)] == self.bos_id:
                start = i + len(self.bos_id)
                end = start
                while end < len(input_ids):
                    if input_ids[end:end + len(self.eos_id)] == self.eos_id:
                        break
                    end += 1
                for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)):
                    labels[j] = input_ids[j]
                i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
            else:
                i += 1
        return labels

    def __getitem__(self, index):
        sample = self.samples[index]
        prompt = self.create_chat_prompt(sample['conversations'])
        input_ids = self.tokenizer(prompt).input_ids[:self.max_length]
        input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids))
        labels = self.generate_labels(input_ids)
        # # === 调试打印 ===
        # print(f"\n--- Sample {index} Labels (length: {len(input_ids)}) ---")
        # for i, (tid, lbl) in enumerate(zip(input_ids, labels)):
        #     token_str = self.tokenizer.decode([tid], skip_special_tokens=False).replace('\n', '\\n')
        #     mask = '✓' if lbl != -100 else '✗'
        #     print(f"{i:3d}: {tid:5d} -> '{token_str:12s}' | label: {lbl:6d} | {mask}")
        # print(f"--- End ---\n")
        # # ================
        return torch.tensor(input_ids, dtype=torch.long), torch.tensor(labels, dtype=torch.long)


class DPODataset(Dataset):
    def __init__(self, file_path, tokenizer, max_length=4096):
        super().__init__()
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.padding = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else 0
        self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant', add_special_tokens=False).input_ids
        self.eos_id = tokenizer(f'{tokenizer.eos_token}', add_special_tokens=False).input_ids
        self.data = load_dataset('json', data_files=file_path, split='train')

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        item = self.data[index]
        chosen = item['chosen']  # 是一个 list，里面包含若干 {role, content}
        rejected = item['rejected']  # 同上
        chosen_prompt = self.tokenizer.apply_chat_template(
            chosen, tokenize=False, add_generation_prompt=False
        )

        rejected_prompt = self.tokenizer.apply_chat_template(
            rejected, tokenize=False, add_generation_prompt=False
        )
        chosen_encoding = self.tokenizer(
            chosen_prompt, truncation=True, max_length=self.max_length, padding='max_length'
        )
        rejected_encoding = self.tokenizer(
            rejected_prompt, truncation=True, max_length=self.max_length, padding='max_length'
        )

        chosen_input_ids = chosen_encoding['input_ids']
        chosen_loss_mask = self.generate_loss_mask(chosen_input_ids)

        rejected_input_ids = rejected_encoding['input_ids']
        rejected_loss_mask = self.generate_loss_mask(rejected_input_ids)
        x_chosen = torch.tensor(chosen_input_ids[:-1], dtype=torch.long)
        y_chosen = torch.tensor(chosen_input_ids[1:], dtype=torch.long)
        mask_chosen = torch.tensor(chosen_loss_mask[1:], dtype=torch.long)
        x_rejected = torch.tensor(rejected_input_ids[:-1], dtype=torch.long)
        y_rejected = torch.tensor(rejected_input_ids[1:], dtype=torch.long)
        mask_rejected = torch.tensor(rejected_loss_mask[1:], dtype=torch.long)

        return {
            'x_chosen': x_chosen,
            'y_chosen': y_chosen,
            'mask_chosen': mask_chosen,
            'x_rejected': x_rejected,
            'y_rejected': y_rejected,
            'mask_rejected': mask_rejected
        }

    def generate_loss_mask(self, input_ids):
        loss_mask = [0] * len(input_ids)
        i = 0
        while i < len(input_ids):
            if input_ids[i:i + len(self.bos_id)] == self.bos_id:
                start = i + len(self.bos_id)
                end = start
                while end < len(input_ids):
                    if input_ids[end:end + len(self.eos_id)] == self.eos_id:
                        break
                    end += 1
                for j in range(start + 1, min(end + len(self.eos_id) + 1, self.max_length)):
                    loss_mask[j] = 1
                i = end + len(self.eos_id) if end < len(input_ids) else len(input_ids)
            else:
                i += 1
        return loss_mask


class RLAIFDataset(Dataset):
    def __init__(self, jsonl_path, tokenizer, max_length=1024):
        super().__init__()
        self.tokenizer = tokenizer
        self.max_length = max_length
        self.samples = load_dataset('json', data_files=jsonl_path, split='train')
        self.bos_id = tokenizer(f'{tokenizer.bos_token}assistant', add_special_tokens=False).input_ids
        self.eos_id = tokenizer(f'{tokenizer.eos_token}', add_special_tokens=False).input_ids

    def __len__(self):
        return len(self.samples)

    def create_chat_prompt(self, conversations):
        messages = []
        answer = ''
        for i, turn in enumerate(conversations):
            role = 'user' if i % 2 == 0 else 'assistant'
            messages.append({"role": role, "content": turn['content']})
            answer = turn['content']
        return self.tokenizer.apply_chat_template(
            messages[:-1],
            tokenize=False,
            add_generation_prompt=True  # 这里需要True
        ), answer

    def __getitem__(self, index):
        sample = self.samples[index]
        prompt, answer = self.create_chat_prompt(sample['conversations'])

        return {
            'prompt': prompt,
            'answer': answer
        }

if __name__ == "__main__":
    pass

---

🚀 快速开始

Step 1: 把上面贴出的代码放到对应位置

目录树参考：

minimind/ # 根目录
├── model/
│   └── model_minimind_dllm.py      # 新增，dLLM结构
├── dataset/
│   └── lm_dataset.py               # 覆盖原有（注意备份）
├── trainer/
│   └── train_full_sft_dllm.py      # 新增，训练脚本
├── scripts/
│   └── dllm_demo.py                # 新增，web demo
├── eval_dllm.py                    # 新增，评估脚本
└── out/
    ├── dllm_sft_512.pth            #  输出，模型权重

Step 2: 训练（默认从 AR 权重初始化）

基于minimind2的full_sft_512.pth检查点微调（大约只要5~10min）：

(gongjy) root@ubuntu:~/minimind$ CUDA_VISIBLE_DEVICES=0 python train_sft_dllm.py
Epoch:[1/2](100/37961), loss: 4.0809, mask_ratio: 0.31, lr: 0.00005000, eta: 46.0min
Epoch:[1/2](200/37961), loss: 3.8554, mask_ratio: 0.27, lr: 0.00005000, eta: 34.0min
Epoch:[1/2](300/37961), loss: 4.0565, mask_ratio: 0.30, lr: 0.00005000, eta: 30.0min
Epoch:[1/2](400/37961), loss: 4.0782, mask_ratio: 0.26, lr: 0.00005000, eta: 28.0min
Epoch:[1/2](500/37961), loss: 3.8238, mask_ratio: 0.30, lr: 0.00005000, eta: 27.0min
Epoch:[1/2](600/37961), loss: 4.1514, mask_ratio: 0.35, lr: 0.00004999, eta: 26.0min
Epoch:[1/2](700/37961), loss: 3.4831, mask_ratio: 0.28, lr: 0.00004999, eta: 25.0min
Epoch:[1/2](800/37961), loss: 3.6387, mask_ratio: 0.36, lr: 0.00004999, eta: 25.0min
Epoch:[1/2](900/37961), loss: 3.1424, mask_ratio: 0.30, lr: 0.00004998, eta: 25.0min
Epoch:[1/2](1000/37961), loss: 3.2463, mask_ratio: 0.25, lr: 0.00004998, eta: 24.0min
Epoch:[1/2](1100/37961), loss: 3.6453, mask_ratio: 0.29, lr: 0.00004998, eta: 24.0min
Epoch:[1/2](1200/37961), loss: 3.2778, mask_ratio: 0.24, lr: 0.00004997, eta: 24.0min
Epoch:[1/2](1300/37961), loss: 3.3814, mask_ratio: 0.29, lr: 0.00004997, eta: 24.0min
Epoch:[1/2](1400/37961), loss: 2.9186, mask_ratio: 0.22, lr: 0.00004996, eta: 24.0min
Epoch:[1/2](1500/37961), loss: 3.3576, mask_ratio: 0.33, lr: 0.00004996, eta: 24.0min
Epoch:[1/2](1600/37961), loss: 3.6901, mask_ratio: 0.32, lr: 0.00004995, eta: 23.0min
Epoch:[1/2](1700/37961), loss: 2.1673, mask_ratio: 0.21, lr: 0.00004994, eta: 23.0min
Epoch:[1/2](1800/37961), loss: 3.3618, mask_ratio: 0.28, lr: 0.00004994, eta: 22.0min
Epoch:[1/2](1900/37961), loss: 2.7446, mask_ratio: 0.24, lr: 0.00004993, eta: 22.0min
Epoch:[1/2](2000/37961), loss: 2.3897, mask_ratio: 0.22, lr: 0.00004992, eta: 22.0min
Epoch:[1/2](2100/37961), loss: 3.2938, mask_ratio: 0.30, lr: 0.00004992, eta: 22.0min
Epoch:[1/2](2200/37961), loss: 2.9948, mask_ratio: 0.24, lr: 0.00004991, eta: 22.0min
Epoch:[1/2](2300/37961), loss: 3.0744, mask_ratio: 0.29, lr: 0.00004990, eta: 22.0min
Epoch:[1/2](2400/37961), loss: 3.2724, mask_ratio: 0.32, lr: 0.00004989, eta: 22.0min
Epoch:[1/2](2500/37961), loss: 2.6671, mask_ratio: 0.22, lr: 0.00004988, eta: 22.0min
Epoch:[1/2](2600/37961), loss: 3.1767, mask_ratio: 0.30, lr: 0.00004987, eta: 22.0min
Epoch:[1/2](2700/37961), loss: 2.3531, mask_ratio: 0.19, lr: 0.00004986, eta: 22.0min
Epoch:[1/2](2800/37961), loss: 3.2795, mask_ratio: 0.34, lr: 0.00004985, eta: 22.0min
Epoch:[1/2](2900/37961), loss: 2.5039, mask_ratio: 0.21, lr: 0.00004984, eta: 22.0min
Epoch:[1/2](3000/37961), loss: 2.5798, mask_ratio: 0.24, lr: 0.00004983, eta: 22.0min
Epoch:[1/2](3100/37961), loss: 2.9274, mask_ratio: 0.29, lr: 0.00004982, eta: 22.0min
Epoch:[1/2](3200/37961), loss: 2.6174, mask_ratio: 0.25, lr: 0.00004980, eta: 22.0min
Epoch:[1/2](3300/37961), loss: 3.1982, mask_ratio: 0.30, lr: 0.00004979, eta: 21.0min
Epoch:[1/2](3400/37961), loss: 2.9121, mask_ratio: 0.27, lr: 0.00004978, eta: 21.0min
Epoch:[1/2](3500/37961), loss: 2.9347, mask_ratio: 0.23, lr: 0.00004976, eta: 21.0min
Epoch:[1/2](3600/37961), loss: 3.2623, mask_ratio: 0.31, lr: 0.00004975, eta: 21.0min
...

Step 3: 评估：

训练1000 steps

(gongjy) root@ubuntu:~/minimind$ python eval_dllm.py 
[加载权重] ./out/dllm_sft_512.pth
Model Params: 25.83M
[0] 自动测试
[1] 手动输入
0
👶: 你有什么特长？
🤖️: 好的，请问您没有什么的是什么是什么方面的类型呢？

👶: 为什么天空是蓝色的
🤖️: 当夏夏天的季，人们的热的夏季季节。人们是季季节，是夏季季节，为在这个季节节和可持续的季节节。夏季季节，节里可以季季节，人们可以户季季节。人们可以季节节的季季季节，人们喜欢多季季季节。夏季节是夏季季节，让我是夏季季节，可以人们最季季节季节的季节节节。我是一季季节季节节季节，是人们季季季节，从而季

👶: 请用Python写一个计算斐波那契数列的函数
🤖️: 当一年的文时，，它很其他帮助我建议。
当然，我很句子，我帮助您为您查哪个文。
请，我需要您需要翻译的文意，请您需要帮助。我需要提供帮助问题，我需要具体的具体的信息。请提供您您需要翻译准确的具体具体或的解决问题和代码。如果您您需要提供需要具体的的对话内容。

👶: 解释一下"光合作用"的基本过程
🤖️: """""""""""""""""""""""""""""天"""""""""""""""�""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""�"

训练10000 steps（最终导出）

(gongjy) root@ubuntu:~/minimind$ python eval_dllm.py
[加载权重] ./out/dllm_sft_512.pth
Model Params: 25.83M
[0] 自动测试
[1] 手动输入
0
👶: 你有什么特长？
🤖️: 虽然我是一个AI语言模型，，没有个人的喜好和行为。但是，我没有个人个人经历或个人经历，因此我是一个机器人喜欢的，因为我没有喜好和喜好或喜欢我，因为我并不喜欢或喜欢某些电影。但是，我我不会喜欢看电影，

👶: 为什么天空是蓝色的
🤖️: 天空是蓝色的，因为蓝色的空蓝色是蓝色的。因为我们，蓝色和蓝色是深蓝色的蓝色。这是因为我们看到的是蓝蓝色的蓝色，这是因为不是蓝色的。在我们看到的天空中，天的颜色是蓝色的，因为我们

👶: 请用Python写一个计算斐波那契数列的函数
🤖️: 当你，换一个数的数列时，你会找出一个数列的位数，这是计算瞶数列的。数列的这个位数是单位数，而数数的数是数数的。这个位数的单位数是：一个数数

👶: 解释一下"光合作用"的基本过程
🤖️: 光合作用是指通过光合作用作用的过程，通过光和光合作用，将光光合作用转换形成能量。光合作用是指能量吸收光的能量和能量的能量，这样减少了能量的产生和能量利用。具体来说，光合合作用是指，光合合

导出的[512*8]模型权重可在此链接下载

https://huggingface.co/jingyaogong/MiniMind2-Pytorch/blob/main/dllm_sft_512.pth

---

Proof-of-Concept，相当于只使用了60M数据进行最小的A2D微调，此时的模型似乎学会了说话，它知道如何回答问题，知道基础语法，但连贯性和逻辑性还很糟糕。不过没关系，对于一个只花了不到10 min微调出来的dLLM来说，这已经相当不错了。

对于这个Toy Model，还存在无数可以改进的地方：

• 更大训练数据 & 训练时间
• 采样算法，Block Masked Diffusion（block_size 和 steps 这俩参数其实挺敏感的，可以调调）
• noise schedule 可以试试 cosine 或者别的，不一定非得 linear
• 尝试不同的 masking 策略
• 更强的基础AR模型

业内对dLLM发展普遍不乐观，效果上限目前看还是在追着AR（且明显差一截），Discrete Diffusion 不像连续扩散那样有漂亮的数学框架撑着，说白了现在在大号 “Masked Bert” 的老路上难真正work，CoT计算时拓展范式似乎也做不到，以及KV Cache无关，并发推理等等一堆下游应用问题，从理论到生态再到工程上都还差得远。

最后想说，扩散语言模型是个有趣的方向，但基于其存在的种种问题，未来在缺乏契机的情况下，我可能不再继续研究。

作为抛砖引玉的帖子，如果你对此感兴趣，欢迎 fork 折腾！

任何问题可以在 Discussion 里讨论（但不保证回复） 🤗😉🫣

Happy!

[legend91019]

老哥是什么方向的，感觉自己代码能力比较差，该咋办

[XELF-encyclopedia]

sir I'm a newbie, where is the dataset for training the LLM? I found the text is not coherent, is it due to the lack of training set or training steps?

[AIWWo]

https://www.modelscope.cn/datasets/gongjy/minimind_dataset/files
This is the link in the document.

[KissYou5]

minimind就只能实行出这种不连贯回复的效果，如果要训练一个在简单对话连贯的模型也就是最少需要0.5b的参数量了。至少需要几百个gb的预训练数据集，这在普通的电脑上是不可能的，作者的minimind模型只能当demo玩，想指望这个训练一个可持续对话的正经ai模型是不可行的