""" PC-SHO-DLM Chat Training — Conditional Masked Diffusion Trains the model to generate responses conditioned on prompts. Prompt stays unmasked, response gets masked/denoised via PC settling. """ import argparse, json, os, sys, time import torch import torch.nn.functional as F from torch.utils.data import DataLoader, Dataset sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src")) from model import PCSHODLM, PCSHOConfig, LocalParameterUpdater, count_parameters class DialogueDataset(Dataset): """Byte-level dialogue dataset. Each sample = (prompt, response) pair. Returns input_ids (full sequence) + prompt_len (where response starts). During training, only the response portion gets masked. """ def __init__(self, pairs: list, max_seq_len: int, vocab_size: int = 257): self.max_seq_len = max_seq_len self.vocab_size = vocab_size self.samples = [] sep = [0x0A, 0x0A] # double newline as separator for prompt, response in pairs: p_bytes = [min(b + 1, vocab_size - 1) for b in prompt.encode("utf-8")] r_bytes = [min(b + 1, vocab_size - 1) for b in response.encode("utf-8")] s_bytes = [min(b + 1, vocab_size - 1) for b in bytes(sep)] # Truncate: prompt gets up to 40% of seq_len, response gets the rest max_p = int(max_seq_len * 0.4) max_r = max_seq_len - min(len(p_bytes), max_p) - len(s_bytes) p = p_bytes[:max_p] r = r_bytes[:max_r] full = p + s_bytes + r prompt_len = len(p) + len(s_bytes) # Pad if len(full) < max_seq_len: full = full + [0] * (max_seq_len - len(full)) # 0 = MASK token as pad full = full[:max_seq_len] self.samples.append({ "input_ids": torch.tensor(full, dtype=torch.long), "prompt_len": prompt_len, "response_len": len(r), }) def __len__(self): return len(self.samples) def __getitem__(self, idx): return self.samples[idx] def load_dialogue_data(name="alpaca", max_samples=None): """Load dialogue pairs from HuggingFace.""" from datasets import load_dataset pairs = [] if name == "alpaca": print("Loading Alpaca dataset...") try: ds = load_dataset("tatsu-lab/alpaca", split="train") for row in ds: instruction = row.get("instruction", "") inp = row.get("input", "") output = row.get("output", "") prompt = f"{instruction}\n{inp}".strip() if inp else instruction if prompt and output: pairs.append((prompt, output)) if max_samples and len(pairs) >= max_samples: break except Exception as e: print(f"HF load failed ({e}), downloading JSON directly...") import urllib.request, json as json_mod url = "https://raw.githubusercontent.com/tatsu-lab/stanford_alpaca/main/alpaca_data.json" with urllib.request.urlopen(url) as r: data = json_mod.loads(r.read()) for row in data: instruction = row.get("instruction", "") inp = row.get("input", "") output = row.get("output", "") prompt = f"{instruction}\n{inp}".strip() if inp else instruction if prompt and output: pairs.append((prompt, output)) if max_samples and len(pairs) >= max_samples: break elif name == "dolly": print("Loading Dolly dataset...") ds = load_dataset("databricks/dolly-15k", split="train") for row in ds: prompt = row.get("instruction", "") or row.get("context", "") response = row.get("response", "") if prompt and response: pairs.append((prompt, response)) if max_samples and len(pairs) >= max_samples: break elif name == "oasst": print("Loading OpenAssistant dataset...") ds = load_dataset("OpenAssistant/oasst1", split="train") # Flatten conversation trees into (parent, reply) pairs messages = {row["message_id"]: row for row in ds} for row in ds: if row["parent_id"] and row["parent_id"] in messages: parent = messages[row["parent_id"]] pairs.append((parent["text"], row["text"])) if max_samples and len(pairs) >= max_samples: break else: raise ValueError(f"Unknown dataset: {name}") print(f"Loaded {len(pairs)} dialogue pairs") return pairs def conditional_corrupt(x_0, t, mask_token_id, schedule, prompt_lens): """Mask only response tokens, keep prompt unmasked.""" B, S = x_0.shape mask_prob = schedule.mask_rate(t).unsqueeze(-1) # (B, 1) random_mask = torch.rand_like(x_0.float()) < mask_prob # (B, S) # Only mask response positions response_mask = torch.zeros_like(x_0, dtype=torch.bool) for b in range(B): response_mask[b, prompt_lens[b]:] = True mask = random_mask & response_mask x_t = x_0.clone() x_t[mask] = mask_token_id return x_t, mask def train(): parser = argparse.ArgumentParser() parser.add_argument("--data", default="alpaca", choices=["alpaca", "dolly", "oasst"]) parser.add_argument("--max_samples", type=int, default=None) parser.add_argument("--d_model", type=int, default=256) parser.add_argument("--n_layers", type=int, default=6) parser.add_argument("--n_heads", type=int, default=8) parser.add_argument("--seq_len", type=int, default=512) parser.add_argument("--batch_size", type=int, default=8) parser.add_argument("--lr", type=float, default=3e-4) parser.add_argument("--max_steps", type=int, default=5000) parser.add_argument("--device", type=str, default="auto") parser.add_argument("--save_dir", type=str, default="checkpoints/chat") parser.add_argument("--checkpoint", type=str, default=None, help="Fine-tune from existing checkpoint") args = parser.parse_args() if args.device == "auto": if torch.cuda.is_available(): args.device = "cuda" elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): args.device = "mps" else: args.device = "cpu" # Config config = PCSHOConfig( vocab_size=257, max_seq_len=args.seq_len, d_model=args.d_model, n_heads=args.n_heads, n_layers=args.n_layers, d_ff=args.d_model * 4, n_diffusion_steps=64, n_settling_steps=6, mask_token_id=0, dropout=0.1, feedback_rank=args.d_model // 4, denoising_lambda=500.0, homeostatic_target=0.0, lateral_inhibition=0.0, ) # Load data pairs = load_dialogue_data(args.data, args.max_samples) n_val = max(1, len(pairs) // 10) train_pairs = pairs[n_val:] val_pairs = pairs[:n_val] train_ds = DialogueDataset(train_pairs, config.max_seq_len) val_ds = DialogueDataset(val_pairs, config.max_seq_len) print(f"Train: {len(train_ds)}, Val: {len(val_ds)}") # Model model = PCSHODLM(config).to(args.device) if args.checkpoint: print(f"Loading checkpoint: {args.checkpoint}") ckpt = torch.load(args.checkpoint, map_location=args.device, weights_only=False) model.load_state_dict(ckpt["model_state_dict"], strict=False) model.train() n_params = count_parameters(model) print(f"Model: {n_params:,} params on {args.device}") # Training updater = LocalParameterUpdater(model, lr_forward=args.lr, lr_feedback=args.lr, lr_readout=args.lr, lr_precision=args.lr * 0.1) train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, drop_last=True) val_loader = DataLoader(val_ds, batch_size=args.batch_size) os.makedirs(args.save_dir, exist_ok=True) step = 0 start = time.time() K_max = config.n_settling_steps print(f"\n{'='*60}") print(f"Chat Training | {n_params:,} params | {args.data} | {len(train_ds)} samples") print(f"{'='*60}\n") while step < args.max_steps: for batch in train_loader: if step >= args.max_steps: break x_0 = batch["input_ids"].to(args.device) prompt_lens = batch["prompt_len"] B, S = x_0.shape # Curriculum K K_curr = max(3, int(K_max * min(1.0, step / 1000))) model.config.n_settling_steps = K_curr # Sample timestep t = torch.randint(1, config.n_diffusion_steps + 1, (B,), device=args.device) # Conditional corruption — only mask response x_t, mask = conditional_corrupt(x_0, t, config.mask_token_id, model.schedule, prompt_lens) # HONEST SETTLE: use x_t (corrupted), NOT x_0 (answer) # Settling minimizes entropy (push toward confident predictions) # without peeking at the ground truth h_0 = model.embed_input(x_t, t) h_init = model.amortized_forward_pass(h_0) with torch.no_grad(): h_settled, _, energies, _, _ = model.settle(h_init, x_t, mask, t) # Local updates (same as LocalParameterUpdater but with our custom mask) L = config.n_layers for l in range(L): h_below = h_settled[l].detach() h_current = h_settled[l + 1].detach() updater.forward_optimizers[l].zero_grad() mu_up = model.forward_blocks[l](h_below) with torch.no_grad(): tok_p, ch_p = model.precision_up[l](h_current, t) eps = h_current - mu_up loss_l = 0.5 * (eps * tok_p * ch_p * eps).sum() / B loss_l.backward() updater.forward_optimizers[l].step() for l in range(L - 1): updater.feedback_optimizers[l].zero_grad() h_above = h_settled[l + 2].detach() h_current = h_settled[l + 1].detach() mu_down = model.feedback_blocks[l](h_above) with torch.no_grad(): tok_p, ch_p = model.precision_down[l](h_current, t) eps = h_current - mu_down loss_l = 0.5 * (eps * tok_p * ch_p * eps).sum() / B loss_l.backward() updater.feedback_optimizers[l].step() # Readout updater.readout_optimizer.zero_grad() logits = model.readout(model.readout_norm(h_settled[L].detach())) mask_logits = logits[mask] mask_targets = x_0[mask] if mask_logits.numel() > 0: readout_loss = F.cross_entropy(mask_logits, mask_targets) readout_loss.backward() updater.readout_optimizer.step() loss_val = readout_loss.item() else: loss_val = 0.0 # Embeddings updater.embed_optimizer.zero_grad() with torch.enable_grad(): h_0_fresh = model.embed_input(x_t, t) h_init_fresh = model.amortized_forward_pass(h_0_fresh) logits_e = model.readout(model.readout_norm(h_init_fresh[L])) ml_e, mt_e = logits_e[mask], x_0[mask] if ml_e.numel() > 0: embed_loss = F.cross_entropy(ml_e, mt_e) embed_loss.backward() torch.nn.utils.clip_grad_norm_( list(model.token_embed.parameters()) + list(model.pos_embed.parameters()) + list(model.time_embed.parameters()), 1.0 ) updater.embed_optimizer.step() step += 1 if step % 50 == 0: el = time.time() - start e = f"{energies[-1]:.0f}" if energies else "N/A" print(f"Step {step:5d} | Loss: {loss_val:.4f} | Energy: {e} | K={K_curr} | {el:.0f}s") if step % 500 == 0: torch.save({ "model_state_dict": model.state_dict(), "config": config, "step": step, }, os.path.join(args.save_dir, f"chat_step{step}.pt")) # Final save torch.save({ "model_state_dict": model.state_dict(), "config": config, "step": step, }, os.path.join(args.save_dir, "chat_final.pt")) print(f"\nDone! {step} steps in {time.time()-start:.0f}s") print(f"Checkpoint: {args.save_dir}/chat_final.pt") print(f"Run: python3 chat.py --checkpoint {args.save_dir}/chat_final.pt") if __name__ == "__main__": train()