File size: 5,174 Bytes
9197119 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | """Adaptación a difusión por enmascaramiento del Qwen3.5-0.8B abliterado.
Basado en summerMC/dqwen (recipe LLaDA/Dream). Entrena LoRA + rev_gates + mask_embed
sobre los pesos abliterados y guarda un checkpoint listo para publicar.
"""
import modal, sys, os
app = modal.App("qwen35-diffusion")
vol_hf = modal.Volume.from_name("qwen35-hf-cache", create_if_missing=True)
vol_out = modal.Volume.from_name("qwen35-diffusion-out", create_if_missing=True)
image = (
modal.Image.debian_slim(python_version="3.11")
.pip_install("torch", "transformers", "peft", "datasets", "accelerate",
"safetensors", "sentencepiece", "huggingface_hub", "numpy",
"flash-linear-attention", "triton")
.add_local_dir("/root/alblit/diffusion/dqwen", "/root/dqwen")
)
BASE = "OpceanAI/Qwen3.5-0.8B-abliterated"
OUT = "/out"
@app.function(image=image, gpu="H200", timeout=7200,
volumes={"/cache": vol_hf, "/out": vol_out},
env={"HF_HOME": "/cache", "TOKENIZERS_PARALLELISM": "false"})
def train(smoke: bool = False, steps: int = 4000, batch: int = 8, accum: int = 4,
max_len: int = 512, max_prompt: int = 256, max_steps_override: int = 0):
sys.path.insert(0, "/root")
import torch
from transformers import AutoModelForImageTextToText, AutoTokenizer
from dqwen.core import DiffusionLM, patch_bidirectional, rev_gate_state, diffusion_generate
from dqwen.data import iter_rows, build_dataset
from dqwen.train import Stage1Cfg, train_stage1, LORA_TARGETS
from peft import LoraConfig, get_peft_model
torch.set_grad_enabled(True)
dev = "cuda"
print("GPU:", torch.cuda.get_device_name(0), flush=True)
tok = AutoTokenizer.from_pretrained(BASE)
lm = AutoModelForImageTextToText.from_pretrained(BASE, dtype=torch.bfloat16, device_map=dev,
attn_implementation="sdpa")
lm.config.use_cache = False
cfg_lora = LoraConfig(r=32, lora_alpha=64, lora_dropout=0.0,
target_modules=LORA_TARGETS, bias="none")
lm = get_peft_model(lm, cfg_lora)
stats = patch_bidirectional(lm, True)
gates = [n for n, _ in lm.named_parameters() if n.endswith("rev_gate")]
print("[patch]", stats, "| rev_gates:", len(gates), flush=True)
for n, p in lm.named_parameters():
if n.endswith("rev_gate"):
p.requires_grad_(True)
mask_id = tok.convert_tokens_to_ids("<|fim_pad|>")
dlm = DiffusionLM(lm, mask_id=mask_id)
dlm.mask_embed.requires_grad_(True)
n_tr = sum(p.numel() for p in dlm.parameters() if p.requires_grad)
print(f"[model] trainable {n_tr/1e6:.1f}M | mask id {mask_id}", flush=True)
n_rows = 2000 if smoke else 60000
rows = iter_rows("tatsu-lab/alpaca", n=n_rows, seed=0)
ds = build_dataset(tok, rows, max_len=max_len, max_prompt=max_prompt,
eos_token="<|im_end|>", chat_kwargs={"enable_thinking": False})
cfg = Stage1Cfg(
max_steps=steps, batch_size=batch, grad_accum=accum, lr=2e-4,
gate_lr=2e-3, mask_lr=1e-3, warmup=50, t_min=0.05, weighting="elbo",
log_every=20, sample_every=500, save_every=1000, save_dir=f"{OUT}/ckpt",
seed=0,
)
prompts = ["Explain why the sky is blue.", "Write a short poem about the sea.",
"What is the capital of France?"]
train_stage1(dlm, tok, ds, cfg, device=dev, sample_prompts=prompts)
dlm.eval()
for p_ in prompts[:1]:
print("[sample]", repr(generate(dlm, tok, p_)), flush=True)
print("[merge] merging LoRA", flush=True)
merged = dlm.lm.merge_and_unload()
gates = rev_gate_state(merged)
for name, mod in list(merged.named_modules()):
if hasattr(mod, "rev_gate"):
delattr(mod, "rev_gate")
os.makedirs(f"{OUT}/final", exist_ok=True)
merged.save_pretrained(f"{OUT}/final", safe_serialization=True)
tok.save_pretrained(f"{OUT}/final")
torch.save({"mask_embed": dlm.mask_embed.detach().cpu(),
"rev_gates": gates,
"mask_id": mask_id,
"base": BASE}, f"{OUT}/final/dllm_extras.pt")
print(f"[saved] {OUT}/final", flush=True)
return {"steps": steps, "tokens": steps * batch * accum * max_len}
def generate(dlm, tok, prompt, gen_len=96, steps=32, block=32):
import torch
from dqwen.core import diffusion_generate
eos_id = tok.convert_tokens_to_ids("<|im_end|>")
text = tok.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False,
add_generation_prompt=True, enable_thinking=False)
ids = torch.tensor([tok(text, add_special_tokens=False)["input_ids"]], device="cuda")
with torch.no_grad():
out = diffusion_generate(dlm, ids, gen_len=gen_len, steps=steps, block_len=block, eos_id=eos_id)
seq = out[0].tolist()
if eos_id in seq:
seq = seq[:seq.index(eos_id)]
return tok.decode(seq, skip_special_tokens=True)
@app.local_entrypoint()
def main(smoke: bool = False, steps: int = 4000, batch: int = 8, accum: int = 4):
print(train.remote(smoke=smoke, steps=steps, batch=batch, accum=accum))
|