| """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)) |
|
|
