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

[jingyaogong]

dLM（Diffusion Language Model，扩散语言模型）最近讨论度挺高。本文记录了一个极简实现：以尽可能少的改动，将已有的 AR（Autoregressive，自回归）MiniMind 续训为 dLM。
PS：当前效果仍属初级，生成的连贯性与逻辑性均有不足。

这类内容不会合入主分支，原因有三：

1. 受众集中：更适合作为特定方向的独立实验
2. 不承诺长期维护：本帖更适合视为一次阶段性快照
3. 主仓库简洁优先：这是最重要的

本帖提供：

• 约 300 行代码，将 MiniMind 从 AR 改造为 dLM
• 直接复用 AR 权重初始化，无需从零训练
• 训练、评估、TUI 对话及 Web Demo 全套脚本

核心改动只有三处：

• Causal Attention（因果注意力）→ Bidirectional Attention（双向注意力）
• 引入 MDM（Masked Diffusion Model）风格的 Mask 训练目标
• 采样阶段采用逐步 Unmask 的去噪还原策略

关于训练策略，一个值得一提的观察：AR → dLM 的 A2D 微调并不需要解冻全部参数。FFN 层主要负责存储知识，而 AR → dLM 的核心变化在于注意力模式（从因果变为双向），因此冻结 FFN 仅训练 Attention 层即可取得不错的效果。更极端地，只训练 Q、K 投影也能获得一个可用的 dLM，这也很好地体现了 A2D 迁移的参数效率。

配套的 Discrete Diffusion Language Models: Principles & Implementation 提供了可视化的训练与推理过程动画演示，包括 Mask、预测、Loss 计算等步骤的逐帧拆解，欢迎体验。

---

完整代码

将以下代码按目录结构放置到对应位置即可运行。

model/model_minimind_dllm.py（dLLM 模型定义）

import math, torch, torch.nn.functional as F
from torch import nn
from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import MaskedLMOutput
from model.model_minimind import MiniMindConfig, MiniMindModel

def add_gumbel_noise(logits, temperature):
    if temperature == 0: return logits
    logits = logits.to(torch.float64)
    noise = torch.rand_like(logits, dtype=torch.float64)
    return logits.exp() / ((-torch.log(noise)) ** temperature)

class MiniMindDLLMConfig(MiniMindConfig):
    model_type = "minimind_dllm"
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.mask_token_id = kwargs.get("mask_token_id", 27)
        self.mask_epsilon = kwargs.get("mask_epsilon", 0.001)

class MiniMindDLLMModel(MiniMindModel):
    def __init__(self, config: MiniMindDLLMConfig):
        super().__init__(config)
        for layer in self.layers: layer.self_attn.is_causal = False

    def add_noise_to_tokens(self, input_ids, t, eps=None, pad_token_id=0):
        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
        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)
        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, n_valid=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)
            denom = n_valid if n_valid is not None else corruption_mask.sum().clamp_min(1)
            loss = (loss[corruption_mask] / p_mask[corruption_mask]).sum() / denom
        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=0.5, top_k=50, steps=128, eos_token_id=None, tokenizer=None, cfg_scale=0.0, **kwargs):
        input_ids = kwargs.get("input_ids", inputs)
        bsz, prompt_len, device = input_ids.shape[0], input_ids.shape[1], input_ids.device
        mask_id = self.config.mask_token_id
        eos_id = eos_token_id or self.config.eos_token_id
        block_size = kwargs.get("block_size", max_new_tokens)

        T = prompt_len + max_new_tokens
        x = torch.full((bsz, T), eos_id, dtype=torch.long, device=device)
        x[:, :prompt_len] = input_ids
        x[:, prompt_len:] = mask_id
        unmasked_index = (x != mask_id)

        num_blocks = math.ceil(max_new_tokens / block_size)
        steps_per_block = steps

        for b in range(num_blocks):
            block_end = min(prompt_len + (b + 1) * block_size, T)

            for i in range(steps_per_block):
                mask_index = (x == mask_id)
                mask_count = mask_index[:, :block_end].sum(-1).min().item()
                if mask_count == 0: break
                n_unmask = max(1, round(mask_count / (steps_per_block - i)))

                if cfg_scale > 0.0:
                    un_x = x.clone()
                    un_x[unmasked_index] = mask_id
                    x_ = torch.cat([x, un_x], dim=0)
                    logits = self(input_ids=x_).logits
                    logits, un_logits = torch.chunk(logits, 2, dim=0)
                    logits = un_logits + (cfg_scale + 1) * (logits - un_logits)
                else:
                    logits = self(input_ids=x).logits

                if top_k > 0: logits[logits < torch.topk(logits, top_k, dim=-1)[0][..., -1:]] = -float('inf')

                x0 = torch.argmax(add_gumbel_noise(logits, temperature), dim=-1)
                p = F.softmax(logits, dim=-1)
                x0_p = torch.gather(p, dim=-1, index=x0.unsqueeze(-1)).squeeze(-1)
                x0_p[:, block_end:] = -float('inf')

                x0 = torch.where(mask_index, x0, x)
                confidence = torch.where(mask_index, x0_p, torch.tensor(-float('inf'), device=device))

                for j in range(bsz):
                    _, idx = torch.topk(confidence[j], k=min(n_unmask, int(mask_count)))
                    x[j, idx] = x0[j, idx]

            if tokenizer and kwargs.get("stream", False):
                print(f"[Block {b+1}/{num_blocks}] {tokenizer.decode(x[0, prompt_len:], skip_special_tokens=False)}")

        return x

trainer/train_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 freeze_all_except_qk(model):
    total_params = 0
    trainable_params = 0
    for param in model.parameters():
        total_params += param.numel()
        param.requires_grad = False
    for name, param in model.named_parameters():
        if 'q_proj' in name or 'k_proj' in name:
            param.requires_grad = True
            trainable_params += param.numel()
    Logger(f'[QK] 总参数: {total_params:,} | 可训练参数: {trainable_params:,}')
    Logger(f'[QK] 可训练参数比例: {trainable_params/total_params*100:.2f}%')
    return trainable_params


def init_model_dllm(config, from_weight, device):
    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} 加载权重')
        else:
            raise FileNotFoundError(f'[DLLM] 找不到权重: {ckp}，请检查 from_weight 参数')
    return model.to(device), tokenizer


def train_epoch(epoch, loader, iters, start_step=0, wandb=None):
    start_time = time.time()
    last_step = start_step
    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)
        last_step = step
        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)
            base_model = model.module if isinstance(model, DistributedDataParallel) else model
            noisy_X, corruption_mask, p_mask = base_model.add_noise_to_tokens(X, t, pad_token_id=tokenizer.pad_token_id)
            
            corruption_mask = corruption_mask & loss_mask
            noisy_X = torch.where(loss_mask, noisy_X, X)
            
            attention_mask = (X != tokenizer.pad_token_id).long()
            res = model(input_ids=noisy_X, attention_mask=attention_mask, labels=X, corruption_mask=corruption_mask, p_mask=p_mask, n_valid=loss_mask.sum())
            loss = res.loss / args.accumulation_steps

        scaler.scale(loss).backward()

        if step % 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:
            spend_time = time.time() - start_time
            current_loss = loss.item() * args.accumulation_steps
            current_lr = optimizer.param_groups[-1]['lr']
            eta_min = spend_time / max(step - start_step, 1) * (iters - step) // 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}, epoch_time: {eta_min:.1f}min')
            if wandb: wandb.log({"loss": current_loss, "mask_ratio": mask_ratio, "learning_rate": current_lr, "epoch_time": eta_min})

        if (step % args.save_interval == 0 or step == iters) 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'
            raw_model = model.module if isinstance(model, DistributedDataParallel) else model
            raw_model = getattr(raw_model, '_orig_mod', raw_model)
            state_dict = raw_model.state_dict()
            torch.save({k: v.half().cpu() for k, v in state_dict.items()}, ckp)
            lm_checkpoint(lm_config, weight=args.save_weight, model=model, optimizer=optimizer,
                         epoch=epoch, step=step, wandb=wandb, save_dir='../checkpoints', scaler=scaler)
            model.train()
            del state_dict

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

    if last_step > start_step and last_step % 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 __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', type=str, help="保存权重的前缀名")
    parser.add_argument("--epochs", type=int, default=5, help="训练轮数")
    parser.add_argument("--batch_size", type=int, default=32, help="batch size")
    parser.add_argument("--learning_rate", type=float, default=1e-4, help="初始学习率")
    parser.add_argument("--device", type=str, default="cuda:0" if torch.cuda.is_available() else "cpu", help="训练设备")
    parser.add_argument("--dtype", type=str, default="bfloat16", help="混合精度类型")
    parser.add_argument("--num_workers", type=int, default=8, help="数据加载线程数")
    parser.add_argument("--accumulation_steps", type=int, default=1, help="梯度累积步数")
    parser.add_argument("--grad_clip", type=float, default=1.0, help="梯度裁剪阈值")
    parser.add_argument("--log_interval", type=int, default=100, help="日志打印间隔")
    parser.add_argument("--save_interval", type=int, default=1000, help="模型保存间隔")
    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
    parser.add_argument('--max_seq_len', default=768, type=int, help="训练的最大截断长度")
    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
    parser.add_argument("--data_path", type=str, default="../dataset/sft_t2t_mini.jsonl", help="训练数据路径")
    parser.add_argument('--from_weight', default='full_sft', type=str, help="基于哪个AR权重初始化，为none则不基于任何权重训练")
    parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）")
    parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb")
    parser.add_argument("--wandb_project", type=str, default="MiniMind-DLLM-SFT", help="wandb项目名")
    parser.add_argument("--use_compile", default=0, type=int, choices=[0, 1], help="是否使用torch.compile加速（0=否，1=是）")
    parser.add_argument("--train_qk_only", default=0, type=int, choices=[0, 1], help="仅训练q_proj和k_proj（0=否，1=是）")
    args = parser.parse_args()

    # ========== 1. 初始化环境和随机种子 ==========
    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))
    
    # ========== 2. 配置目录、模型参数、检查ckp ==========
    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))
    ckp_data = lm_checkpoint(lm_config, weight=args.save_weight, save_dir='../checkpoints') if args.from_resume == 1 else None
    
    # ========== 3. 设置混合精度 ==========
    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)
    
    # ========== 4. 配wandb ==========
    wandb = None
    if args.use_wandb and is_main_process():
        import swanlab as wandb
        wandb_id = ckp_data.get('wandb_id') if ckp_data else None
        resume = 'must' if wandb_id else None
        wandb_run_name = f"MiniMind-DLLM-SFT-Epoch-{args.epochs}-BatchSize-{args.batch_size}-LearningRate-{args.learning_rate}"
        wandb.init(project=args.wandb_project, name=wandb_run_name, id=wandb_id, resume=resume)
    
    # ========== 5. 定义模型、数据、优化器 ==========
    model, tokenizer = init_model_dllm(lm_config, args.from_weight, device=args.device)
    if args.train_qk_only:
        freeze_all_except_qk(model)
    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)
    
    # ========== 6. 从ckp恢复状态 ==========
    start_epoch, start_step = 0, 0
    if ckp_data:
        model.load_state_dict(ckp_data['model'])
        optimizer.load_state_dict(ckp_data['optimizer'])
        scaler.load_state_dict(ckp_data['scaler'])
        start_epoch = ckp_data['epoch']
        start_step = ckp_data.get('step', 0)
    
    # ========== 7. 编译和分布式包装 ==========
    if args.use_compile == 1:
        model = torch.compile(model)
        Logger('torch.compile enabled')
    if dist.is_initialized():
        model._ddp_params_and_buffers_to_ignore = {"freqs_cos", "freqs_sin"}
        model = DistributedDataParallel(model, device_ids=[local_rank])
    
    # ========== 8. 开始训练 ==========
    for epoch in range(start_epoch, args.epochs):
        train_sampler and train_sampler.set_epoch(epoch)
        setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist()
        skip = start_step if (epoch == start_epoch and start_step > 0) else 0
        batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip)
        loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True)
        if skip > 0:
            Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始')
            train_epoch(epoch, loader, len(loader) + skip, start_step, wandb)
        else:
            train_epoch(epoch, loader, len(loader), 0, wandb)
    
    # ========== 9. 清理分布进程 ==========
    if dist.is_initialized(): dist.destroy_process_group()

eval_dllm.py（评估脚本）

import argparse
import random
import warnings
import torch
from transformers import AutoTokenizer
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'
    model.load_state_dict(torch.load(ckp, map_location=args.device), strict=True)
    get_model_params(model, model.config)
    return model.half().eval().to(args.device), tokenizer

def main():
    parser = argparse.ArgumentParser(description="MiniMind DLLM推理与对话")
    parser.add_argument('--load_from', default='model', type=str, help="模型加载路径（model=原生torch权重）")
    parser.add_argument('--save_dir', default='out', type=str, help="模型权重目录")
    parser.add_argument('--weight', default='dllm', type=str, help="权重名称前缀（dllm, dllm_pt）")
    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
    parser.add_argument('--use_moe', default=0, type=int, choices=[0, 1], help="是否使用MoE架构（0=否，1=是）")
    parser.add_argument('--max_new_tokens', default=256, type=int, help="最大生成长度")
    parser.add_argument('--steps', default=32, type=int, help="去噪总步数")
    parser.add_argument('--block_size', default=32, type=int, help="每个block的token数")
    parser.add_argument('--temperature', default=0.4, type=float, help="生成温度，控制随机性（0-1，越大越随机）")
    parser.add_argument('--top_k', default=50, type=int, help="top-k采样（保留概率最高的k个token）")
    parser.add_argument('--historys', default=0, type=int, help="携带历史对话轮数（需为偶数，0表示不携带历史）")
    parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="运行设备")
    args = parser.parse_args()
    
    prompts = [
        '请介绍一下你自己',
        '推荐一部好看的科幻电影',
        '如何快速学习一门新的编程语言'
    ]
    
    conversation = []
    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(random.randint(0, 31415926))
        if input_mode == 0: print(f'💬: {prompt}')
        conversation = conversation[-args.historys:] if args.historys else []
        conversation.append({"role": "user", "content": prompt})

        inputs = tokenizer.apply_chat_template(conversation, 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,
        )
        gen_ids = generated[0, inputs['input_ids'].shape[1]:].tolist()
        if tokenizer.eos_token_id in gen_ids:
            gen_ids = gen_ids[:gen_ids.index(tokenizer.eos_token_id)]
        response = tokenizer.decode(gen_ids, skip_special_tokens=True)
        conversation.append({"role": "assistant", "content": response})
        print(f'🧠: {response}\n')

if __name__ == "__main__":
    main()

scripts/chat_dllm.py（TUI 对话）

import sys, os, argparse, math, shutil
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

import torch
import torch.nn.functional as F
from transformers import AutoTokenizer
from model.model_minimind_dllm import MiniMindDLLMConfig, MiniMindForMaskedDiffusion, add_gumbel_noise

from rich.console import Console
from rich.layout import Layout
from rich.live import Live
from rich.panel import Panel
from rich.progress import BarColumn, MofNCompleteColumn, Progress, SpinnerColumn, TextColumn, TimeRemainingColumn
from rich.text import Text

model, tokenizer, device = None, None, 'cpu'


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


prev_ids = None

def render_tokens(x, prompt_len, mask_id, eos_id):
    global prev_ids
    ids = x[prompt_len:].tolist()
    t = Text()
    for i, tid in enumerate(ids):
        if tid == mask_id:
            t.append("[M]", style="dim red")
        elif tid == eos_id:
            t.append("[E]", style="dim green")
        elif prev_ids is not None and i < len(prev_ids) and prev_ids[i] == mask_id:
            t.append(tokenizer.decode([tid]), style="bold bright_green")
        else:
            t.append(tokenizer.decode([tid]), style="white")
    prev_ids = ids[:]
    return t


def update_tui(live, layout, progress, task_id, x, prompt_len, mask_id, eos_id, step, total_steps):
    masks_left = int((x[prompt_len:] == mask_id).sum().item())
    total_gen = len(x[prompt_len:])
    revealed = total_gen - masks_left
    pct = f"{int(100 * step / max(total_steps, 1))}%"
    progress.update(task_id, completed=step, total=total_steps, masks=masks_left, pct=pct)
    layout["text"].update(Panel(
        render_tokens(x, prompt_len, mask_id, eos_id),
        title=f"[bold cyan]Denoising[/bold cyan]  [dim]{revealed}/{total_gen} revealed[/dim]",
        subtitle=f"[dim]Step {step}/{total_steps}[/dim]",
        border_style="bright_blue", padding=(1, 1),
    ))
    layout["progress"].update(Panel(progress, border_style="dim"))
    live.refresh()


def stream_generate_tui(input_ids, prompt_len, args):
    global prev_ids
    prev_ids = None
    mask_id = model.config.mask_token_id
    eos_id = tokenizer.eos_token_id
    block_size, steps, max_tokens = args.block_size, args.steps, args.max_new_tokens

    num_blocks = math.ceil(max_tokens / block_size)
    total_steps = num_blocks * steps
    T = prompt_len + max_tokens
    x = torch.full((1, T), eos_id, dtype=torch.long, device=device)
    x[:, :prompt_len] = input_ids
    x[:, prompt_len:] = mask_id
    unmasked_index = (x != mask_id)

    try:
        tw = shutil.get_terminal_size().columns
    except Exception:
        tw = 120
    console = Console(width=tw)

    progress = Progress(
        SpinnerColumn(), TextColumn("[bold blue]Diffusion"), BarColumn(), MofNCompleteColumn(),
        TextColumn("•"), TextColumn("[cyan]Masks: {task.fields[masks]}"),
        TextColumn("•"), TextColumn("[magenta]{task.fields[pct]:>4s}"), TimeRemainingColumn(), expand=True,
    )
    layout = Layout()
    layout.split_column(Layout(name="text", ratio=1), Layout(name="progress", size=3))
    task_id = progress.add_task("gen", total=total_steps, masks=max_tokens, pct="0%")

    with Live(layout, console=console, auto_refresh=False, screen=True) as live:
        global_step = 0
        for b in range(num_blocks):
            block_end = min(prompt_len + (b + 1) * block_size, T)

            for step in range(steps):
                mask_index = (x == mask_id)
                mask_count = mask_index[:, :block_end].sum(-1).min().item()
                if mask_count == 0: break
                n_unmask = max(1, round(mask_count / (steps - step)))

                with torch.no_grad():
                    logits = model(input_ids=x).logits
                if args.top_k > 0:
                    logits[logits < torch.topk(logits, args.top_k, dim=-1)[0][..., -1:]] = -float('inf')
                x0 = torch.argmax(add_gumbel_noise(logits, args.temperature), dim=-1)
                p = F.softmax(logits.float(), dim=-1)
                x0_p = torch.gather(p, dim=-1, index=x0.unsqueeze(-1)).squeeze(-1)
                x0_p[:, block_end:] = -float('inf')
                x0 = torch.where(mask_index, x0, x)
                confidence = torch.where(mask_index, x0_p, torch.tensor(-float('inf'), device=device))
                _, idx = torch.topk(confidence[0], k=min(n_unmask, int(mask_count)))
                x[0, idx] = x0[0, idx]

                global_step += 1
                update_tui(live, layout, progress, task_id, x[0], prompt_len, mask_id, eos_id, global_step, total_steps)

            if eos_id and (x[:, prompt_len:] == eos_id).any(): break

    gen_ids = x[0, prompt_len:].tolist()
    if eos_id in gen_ids:
        gen_ids = gen_ids[:gen_ids.index(eos_id)]
    final_text = tokenizer.decode(gen_ids, skip_special_tokens=True)
    print(f"\n🤖️: {final_text}\n")
    return final_text


def chat_loop(args):
    is_pt = 'pt' in args.weight
    prompts = ['你有什么特长？', '为什么天空是蓝色的', '请用Python写一个计算斐波那契数列的函数', '解释一下"光合作用"的基本过程']

    print(f"\n{'='*50}")
    print(f"  MiniMind DLLM TUI Chat ({'pretrain' if is_pt else 'sft'} mode)")
    print(f"  steps={args.steps} block={args.block_size} max_tokens={args.max_new_tokens}")
    print(f"  temperature={args.temperature} top_k={args.top_k}")
    print(f"{'='*50}\n")

    mode = input('[0] 自动测试\n[1] 手动输入\n> ').strip()
    prompt_iter = prompts if mode == '0' else iter(lambda: input('\n💬 > '), '')

    for prompt in prompt_iter:
        if mode == '0':
            print(f'\n💬: {prompt}')

        if is_pt:
            input_ids = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).input_ids.to(device)
        else:
            tpl = tokenizer.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True)
            input_ids = tokenizer(tpl, return_tensors="pt", truncation=True).input_ids.to(device)

        stream_generate_tui(input_ids, input_ids.shape[1], args)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="MiniMind DLLM TUI Chat")
    parser.add_argument('--load_from', default='../model', type=str, help="tokenizer路径")
    parser.add_argument('--save_dir', default='../out', type=str, help="权重目录")
    parser.add_argument('--weight', default='dllm', type=str, help="权重前缀名")
    parser.add_argument('--hidden_size', default=768, type=int, help="隐藏层维度")
    parser.add_argument('--num_hidden_layers', default=8, type=int, help="隐藏层数量")
    parser.add_argument('--use_moe', default=0, type=int, help="是否使用MoE架构（0=否，1=是）")
    parser.add_argument('--max_new_tokens', default=128, type=int, help="最大生成长度")
    parser.add_argument('--steps', default=32, type=int, help="去噪总步数")
    parser.add_argument('--block_size', default=32, type=int, help="每个block的token数")
    parser.add_argument('--temperature', default=0.5, type=float, help="采样温度")
    parser.add_argument('--top_k', default=50, type=int, help="top-k采样")
    parser.add_argument('--device', default='cuda' if torch.cuda.is_available() else 'cpu', type=str, help="推理设备")
    args = parser.parse_args()

    try:
        load_model(args)
        chat_loop(args)
    except KeyboardInterrupt:
        print("\n退出")

scripts/web_dllm.py（Web Demo）

import sys, os, json, time, math, 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</title>
<style>
:root{
  --bg:#24262b;--surface:#34373d;--panel:#2b2e34;--border:#4a4f57;
  --fg:#f3f4f6;--dim:#b0b4bb;--accent:#7ea6e8;--reveal:#63d297;
  --mask:#f07c7c;--mask-dim:rgba(240,124,124,0.12);
}
*{box-sizing:border-box;margin:0;padding:0}
html{font-size:15px}
body{font-family:-apple-system,BlinkMacSystemFont,'Segoe UI',sans-serif;background:var(--bg);color:var(--fg);height:100vh;display:flex;flex-direction:column;overflow:hidden;justify-content:center}

.wrap{max-width:860px;width:100%;margin:0 auto;display:flex;flex-direction:column;height:calc(100vh - 48px);max-height:720px;box-shadow:0 8px 24px rgba(0,0,0,0.18)}
.top-bar{padding:16px 20px;display:flex;align-items:center;justify-content:space-between;border-bottom:1px solid var(--border);background:var(--surface);border-radius:10px 10px 0 0}
.top-bar .title{font-size:14px;font-weight:600;color:var(--fg);letter-spacing:0}
.top-bar .status{font-size:12px;color:var(--dim);letter-spacing:0}
.top-bar .status .dot{display:inline-block;width:6px;height:6px;border-radius:50%;margin-right:6px;background:var(--dim);vertical-align:middle}
.top-bar .status .dot.live{background:var(--reveal)}

.main{flex:1;display:flex;flex-direction:column;overflow:hidden;min-height:0}

.meta-bar{padding:10px 20px;display:flex;gap:16px;border-bottom:1px solid var(--border);font-size:12px;color:var(--dim);letter-spacing:0;background:var(--surface)}
.meta-bar .tag{padding:0;border:none;border-radius:0;background:transparent}

.output{flex:1;overflow-y:auto;padding:20px;font-size:14px;line-height:1.8;white-space:pre-wrap;word-break:break-all;background:var(--surface);min-height:0;scrollbar-width:none}
.output::-webkit-scrollbar{display:none}
.output .token{transition:color .15s ease}
.output .token.mask{color:var(--mask);opacity:.9;font-weight:600}
.output .token.hidden{color:var(--dim);opacity:.15}
.output .token.filled{color:var(--fg)}
.output .token.new{color:var(--reveal);font-weight:600}
.output.token-mode{white-space:normal;display:flex;flex-wrap:wrap;gap:3px;align-content:flex-start}
.output.token-mode .token{padding:2px 6px;border:1px solid var(--border);border-radius:4px;font-size:12px;line-height:1.4;background:var(--panel)}
.output.token-mode .token.mask{background:rgba(240,124,124,0.18);border-color:rgba(240,124,124,0.48);box-shadow:inset 0 0 0 1px rgba(240,124,124,0.12)}
.output.token-mode .token.hidden{border-color:transparent;background:transparent}
.output.token-mode .token.new{border-color:var(--reveal);background:rgba(99,210,151,0.10)}

.progress-wrap{padding:8px 20px;border-top:1px solid var(--border);background:var(--surface)}
.progress-track{height:2px;background:rgba(255,255,255,0.08);border-radius:1px;overflow:hidden}
.progress-fill{height:100%;width:0;background:var(--accent);transition:width .2s ease;border-radius:1px}

.bottom{padding:16px 20px;border-top:1px solid var(--border);background:var(--surface);border-radius:0 0 10px 10px}
.input-row{display:flex;gap:10px;margin-bottom:10px}
.input-row input[type=text]{flex:1;padding:10px 14px;background:var(--panel);border:1px solid var(--border);color:var(--fg);font-family:inherit;font-size:13px;outline:none;border-radius:6px;transition:border-color .2s}
.input-row input[type=text]:focus{border-color:var(--accent)}
.input-row input[type=text]::placeholder{color:var(--dim)}
.input-row button{padding:10px 16px;background:var(--panel);border:1px solid var(--border);color:var(--fg);font-family:inherit;font-size:12px;font-weight:500;letter-spacing:0;cursor:pointer;border-radius:6px;transition:all .15s}
.input-row button:hover{border-color:var(--accent);color:var(--accent)}
.input-row button:disabled{opacity:.3;cursor:default}
.params-row{display:flex;gap:16px;font-size:11px;color:var(--dim)}
.params-row label{display:flex;align-items:center;gap:6px}
.params-row input{width:64px;padding:5px 8px;background:var(--panel);border:1px solid var(--border);color:var(--fg);font-family:inherit;font-size:12px;border-radius:6px;outline:none;text-align:center}
.params-row input:focus{border-color:var(--accent)}
</style></head>
<body>
<div class="wrap">
<div class="top-bar">
  <div class="title">MiniMind dLLM</div>
  <div class="status" id="statusBar"><span class="dot"></span>IDLE</div>
</div>
<div class="main">
  <div class="meta-bar" id="metaBar">
    <span class="tag" id="stepTag">step --</span>
    <span class="tag" id="maskTag">masks --</span>
    <span class="tag" id="blockTag">block --</span>
    <span class="tag" id="viewToggle" style="cursor:pointer;user-select:none;margin-left:auto">text</span>
  </div>
  <div class="output" id="tokens"></div>
</div>
<div class="progress-wrap"><div class="progress-track"><div class="progress-fill" id="progressFill"></div></div></div>
<div class="bottom">
  <div class="input-row">
    <input type="text" id="prompt" placeholder="prompt >" value="Python实现快速排序算法">
    <button id="genBtn" onclick="generate()">生成</button>
  </div>
  <div class="params-row">
    <label>block <input type="number" id="blockSize" value="32"></label>
    <label>steps <input type="number" id="steps" value="32"></label>
    <label>max <input type="number" id="maxTokens" value="128"></label>
  </div>
</div>
</div>
<script>
let prev=[],running=false,viewMode='text',lastData=null;
document.getElementById('viewToggle').addEventListener('click',()=>{
  viewMode=viewMode==='text'?'token':'text';
  document.getElementById('viewToggle').textContent=viewMode==='text'?'文本':'分词';
  const o=document.getElementById('tokens');
  if(viewMode==='token')o.classList.add('token-mode');else o.classList.remove('token-mode');
  if(lastData)render(lastData,0);
});
async function generate(){
  if(running)return;
  running=true;
  const btn=document.getElementById('genBtn');
  btn.disabled=true;btn.textContent='...';
  const dot=document.querySelector('.dot');
  dot.classList.add('live');
  document.getElementById('statusBar').innerHTML='<span class="dot live"></span>生成中';
  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('progressFill').style.width='0%';
  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();
  let buf='';
  while(1){
    const{done,value}=await reader.read();
    if(done)break;
    buf+=dec.decode(value,{stream:true});
    const parts=buf.split('\\n\\n');
    buf=parts.pop();
    parts.forEach(p=>{const l=p.trim();if(l.startsWith('data:'))try{render(JSON.parse(l.slice(5)),s)}catch(e){}});
  }
  running=false;btn.disabled=false;btn.textContent='生成';
  document.getElementById('statusBar').innerHTML='<span class="dot"></span>完成';
}
function render(d,totalS){
  lastData=d;
  const masks=d.tokens.filter(t=>t.is_mask&&!t.hidden).length;
  document.getElementById('stepTag').textContent='step'+d.step+'/'+d.total_steps;
  document.getElementById('maskTag').textContent='masks'+masks;
  document.getElementById('blockTag').textContent='block'+d.block;
  const pct=Math.min(100,Math.round(d.step/d.total_steps*100));
  document.getElementById('progressFill').style.width=pct+'%';
  if(d.done)document.getElementById('progressFill').style.width='100%';
  const c=document.getElementById('tokens');
  c.innerHTML='';
  if(viewMode==='token'){
    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||'\u00a0';
      c.appendChild(s);
    });
  }else{
    let h='';
    d.tokens.forEach((t,i)=>{
      const cls=t.hidden?'hidden':t.is_mask?'mask':'filled'+(!d.done&&prev[i]!==t.text?' new':'');
      const txt=t.hidden?'':t.is_mask?'[M]':(t.text||'').replace(/&/g,'&amp;').replace(/</g,'&lt;').replace(/>/g,'&gt;');
      h+='<span class="token '+cls+'">'+txt+'</span>';
    });
    c.innerHTML=h;
  }
  prev=d.tokens.map(t=>t.text);
}
document.getElementById('prompt').addEventListener('keydown',e=>{if(e.key==='Enter')generate()});
</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=768, num_hidden_layers=8))
    ckp = '../out/dllm_768.pth'
    if os.path.exists(ckp): model.load_state_dict(torch.load(ckp, map_location=device), strict=False)
    model = model.half().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: 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', 128), d.get('steps', 128), d.get('max_tokens', 256)
    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, prompt_len = model.config.mask_token_id, tokenizer.eos_token_id, input_ids.shape[1]
        x = torch.cat([input_ids, torch.full((1, max_tokens), mask_id, dtype=torch.long, device=device)], dim=1)
        num_blocks = math.ceil(max_tokens / block_size)
        global_step, total_steps = 0, num_blocks * steps
        yield f"data:{json.dumps({'block': 0, 'step': 0, 'total_steps': total_steps, 'tokens': make_tokens(x, prompt_len, mask_id, eos_id)})}\n\n"
        for b in range(num_blocks):
            block_start = prompt_len + b * block_size
            block_end = min(prompt_len + (b + 1) * block_size, prompt_len + max_tokens)
            cur_block = block_end - block_start
            for step in range(steps):
                mask_pos = (x[:, block_start:block_end] == mask_id)
                if not mask_pos.any(): break
                with torch.no_grad(): logits = model(input_ids=x).logits[:, block_start:block_end, :]
                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]]
                global_step += 1
                yield f"data:{json.dumps({'block': b+1, 'step': global_step, 'total_steps': total_steps, 'tokens': make_tokens(x, prompt_len, mask_id, eos_id)})}\n\n"
                time.sleep(0.05)
            if eos_id and (x[:, prompt_len:] == eos_id).any(): break
        yield f"data:{json.dumps({'block': num_blocks, 'step': total_steps, 'total_steps': total_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)

快速开始

Step 1：按上述目录结构放置代码文件

参考目录结构：

minimind/ # 根目录
├── model/
│   └── model_minimind_dllm.py      # 新增，dLLM结构
├── trainer/
│   └── train_dllm.py               # 新增，训练脚本
├── scripts/
│   ├── chat_dllm.py                # 新增，TUI 对话
│   └── web_dllm.py                 # 新增，web demo
├── eval_dllm.py                    # 新增，评估脚本
└── out/
    └── dllm_768.pth                # 输出，模型权重

Step 2：训练（或直接下载已训练权重）

可直接下载预训练好的 dllm_768.pth 放到 out/ 目录：

来源	链接
ModelScope	dllm_768.pth
HuggingFace	dllm_768.pth

训练默认基于 minimind-3 的 full_sft_768.pth 权重做 A2D 微调（请提前准备，主页说明中可获取）：

cd trainer && python train_dllm.py

loss 走势参考

Step 3：评估

python eval_dllm.py

推理结果：

Model Params: 63.91M
[0] 自动测试
[1] 手动输入
0
💬: 请介绍一下你自己
🧠: 作为我我我具备学习和解决问题的能力，并能够完成各种任务。我能够完成各种任务，包括但不限于学习知识、解决问题、、学习、写作、等。
我我具备学习能力，例如，能够通过实践、解决问题和解决问题来来提升自己，提高解决问题的能力。我能够学习到新知识，掌握技能，以及我能够强大的学习能力，能够进行分析、解释、分析和解决问题。。
总结：
我是一个具备多种技能的人，，能够能够完成各种任务和解决问题。我能够通过学习、实践、解决问题和解决问题的能力来提升自己。
我具备学习能力的能力，，能够能够通过学习、实践和解决问题，，不断改进自己的学习和成长。我
我能够学习学习新技能，能够学习学习新新技能，，能够通过学习学习新知识，提高自己的能力和学习能力。
我能够学习学习技能，，例如，我能够学习新学习新技能，能够学习学习新技能，，能够学习学习新技能，能够解决问题。
我
我能够能够学习学习学习能力，以不断不断地成长和成长。
我具备具备学习能力，我能够通过学习学习、学习和解决问题。

💬: 推荐一部好看的科幻电影
🧠: 当然可以，我推荐推荐一些经典幻的科幻电影。
1. **《星际穿》** - 由由罗里·斯执导，改编了由成功创作的科幻电影，深入探讨了科幻的冒险，，以及人类与外之间的星际和文明的危险，展现了人类的生存智慧。。
2. **《星际穿》**** - 由由克里··执导的科幻电影，讲述了一个复杂的外星文明，，人类与外界的星际冲突。
3. **《星际》** - 由由由罗马斯导执导的科幻电影，讲述了一个复杂的外星文明，人类与外星文明之间的冲突与星际冲突。
4. **《星际》**** - 由由由罗马导导执导的的科幻电影，讲述了一个复杂的外星文明，人类与星际旅行，讲述了一个复杂的外际界冲突，展现了人类与外星文明的关系。。
这些电影不仅了经典的科幻幻电影，它们也了电影深刻的科幻，适合喜欢喜欢科幻的科幻。

💬: 如何快速学习一门新的编程语言
🧠: 学习学习一新新编程语言，可以遵循以下步骤：
- ** **基础基础**：**：首先，你需要掌握一些基础，如变量语法、变量、、函数等。例如。
- **学习编程语言**：**：了解编程语言的基本概念，如变量、对象、类、、函数等、函数定义等。
- ** **实践编程**：通过实际项目或项目，通过实践项目来巩。。实践是- **阅读编程书籍阅读阅读书籍书籍，书籍是如《Python编程基础》等书籍。。
- **学习编程库**：学习编程库库，了解如何使用基本的函数、、数组操作等。
- **学习编程语言**：了解一些常用的编程语言，，如C#、、Python++、R++、C++等。
- ** **学习编程**：**：学习编程是一个很的过程，如网页开发、游戏、、自动化脚本等脚本等。
- ** **实践编程**：学习编程后，，，可以尝试做项目和实践项目，这将帮助你知识和技能。
- ** **持续学习**：编程是一个需要不断发展的的过程，不断学习新技术，保持好奇心态度，对新知识保持开放的态度。

小结

上述结果基于约 1.6GB SFT 数据的 A2D（AR-to-dLM）微调，总训练时长约 1 小时。模型已具备基本的问答与语法能力，但在连贯性和逻辑性上仍有明显不足，这基本符合预期。作为最小可行验证，它足以表明从 AR 到 dLM 的迁移路径是可行的。

有诸多改进空间，包括：

• 更大规模的训练数据与更充分的训练
• 采样策略的调优，尤其是 block_size 与 steps 两个超参数的敏感性问题
• 不同的 noise schedule（如 cosine）及 masking 策略
• 更强的基础 AR 模型作为初始化起点
• 参数高效微调：冻结 FFN 仅训 Attention，或只训 Q/K 投影

从更宏观的视角看，dLM 与 AR 在生成质量上仍存在显著差距。Discrete Diffusion 不像 Continuous Diffusion 那样有完备的数学框架作为支撑，其本质更接近 Masked Prediction 的延伸，尚未展现出真正的范式优势。工程层面的问题同样突出：双向注意力要求每一步去噪都对 Block 内所有 token 做全量 Attention，无法像 AR 那样通过 KV Cache 实现增量解码，推理效率与 AR 存在量级差距。此外，CoT 推理、并发解码等适配问题也远未解决，这条路线在理论深度和工程生态上都尚不成熟。

扩散语言模型是一个有价值的研究方向，本帖仅作抛砖引玉，感兴趣的读者欢迎 fork 自行探索。

如有问题可在 Discussion 中交流，但不保证及时回复。

[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模型是不可行的