Add files using upload-large-folder tool

LTA_openwebtext_dualt/scripts/_tmp_flowtext_decode_lab_nospecial_prompt.py

@@ -0,0 +1,484 @@
+#!/usr/bin/env python3
+"""Decode-sweep lab for FlowText OpenWebText checkpoints.
+
+The goal is to debug inference without touching training.  We try several
+simplex-valid update rules, generate many candidates, and rank them with
+anti-collapse diagnostics instead of pure self-likelihood.
+
+Run from the flowtext_standard_bench repository root.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import re
+import sys
+from collections import Counter
+from dataclasses import dataclass, asdict
+from pathlib import Path
+from typing import Iterable, List, Sequence
+
+import torch
+import torch.nn.functional as F
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from eval import build_model_from_ckpt
+from flowtext_lab.bridges import smooth_onehot
+from flowtext_lab.decode import model_time_for_step, sample_noise_simplex, state_for_model
+from flowtext_lab.tokenization import BpeTextTokenizer
+
+
+WORD_RE = re.compile(r"[A-Za-z]+|\d+|[^\sA-Za-z\d]")
+
+
+@dataclass
+class DecodeConfig:
+    label: str
+    rule: str
+    steps: int = 64
+    model_t_mode: str = "flow"
+    eta: float = 0.5
+    damping: float = 1.0
+    max_gamma: float = 1.0
+    endpoint_temp: float = 1.0
+    state_floor: float = 1e-8
+    final_from: str = "state"
+    noise_mix: float = 0.0
+    noise_decay: str = "linear"
+    eos_logit_bias: float = 0.0
+
+
+def tokenize_for_metrics(text: str) -> list[str]:
+    return WORD_RE.findall(text)
+
+
+def repeated_ngram_frac(tokens: Sequence[str], n: int) -> float:
+    if len(tokens) < n:
+        return 0.0
+    grams = list(zip(*[tokens[i:] for i in range(n)]))
+    counts = Counter(grams)
+    return sum(v - 1 for v in counts.values() if v > 1) / max(len(grams), 1)
+
+
+def text_metrics(text: str) -> dict:
+    toks = tokenize_for_metrics(text)
+    words = [t.lower() for t in toks if re.fullmatch(r"[A-Za-z]+", t)]
+    n_tok = max(len(toks), 1)
+    n_words = max(len(words), 1)
+    word_counts = Counter(words)
+    max_word_frac = word_counts.most_common(1)[0][1] / n_words if word_counts else 1.0
+    distinct1 = len(set(words)) / n_words if words else 0.0
+    bigrams = list(zip(words, words[1:]))
+    distinct2 = len(set(bigrams)) / max(len(bigrams), 1) if bigrams else 0.0
+    digit_frac = sum(t.isdigit() for t in toks) / n_tok
+    punct_frac = sum(bool(re.fullmatch(r"[,.;:!?]+", t)) for t in toks) / n_tok
+    eos_count = text.count("<|endoftext|>")
+    bad_char_count = text.count("�")
+    rep3 = repeated_ngram_frac([t.lower() for t in toks], 3)
+    rep4 = repeated_ngram_frac([t.lower() for t in toks], 4)
+    # This score is deliberately simple and non-oracle.  It rewards length and
+    # lexical variety while heavily penalizing classic collapse artifacts.
+    quality = (
+        min(len(text) / 700.0, 1.0)
+        + 0.35 * distinct2
+        + 0.15 * distinct1
+        - 0.30 * eos_count
+        - 2.60 * rep3
+        - 1.60 * rep4
+        - 1.30 * digit_frac
+        - 0.65 * punct_frac
+        - 1.35 * max_word_frac
+        - 0.35 * bad_char_count
+    )
+    return {
+        "quality": float(quality),
+        "chars": len(text),
+        "tokens": len(toks),
+        "words": len(words),
+        "eos_count": eos_count,
+        "bad_char_count": bad_char_count,
+        "rep3": float(rep3),
+        "rep4": float(rep4),
+        "distinct1": float(distinct1),
+        "distinct2": float(distinct2),
+        "digit_frac": float(digit_frac),
+        "punct_frac": float(punct_frac),
+        "max_word_frac": float(max_word_frac),
+    }
+
+
+def decode_text(tokenizer: BpeTextTokenizer, ids: Sequence[int]) -> str:
+    return tokenizer.decode(ids, stop_at_eos=False, skip_special_tokens=False)
+
+
+def encode_prompt(tokenizer: BpeTextTokenizer, prompt: str, max_len: int) -> list[int]:
+    # Prefix-continuation mode: lock [CLS] + prompt tokens, but do NOT lock an early [SEP].
+    # tokenizers.Tokenizer.encode(..., add_special_tokens=False) avoids the default BERT wrapper.
+    core = list(tokenizer.tokenizer.encode(prompt, add_special_tokens=False).ids)
+    bos = tokenizer.bos_id
+    ids = ([bos] if bos is not None and bos >= 0 else []) + core
+    return ids[:max_len]
+
+
+@torch.no_grad()
+def build_initial_state(
+    tokenizer: BpeTextTokenizer,
+    prompts: list[str],
+    restarts: int,
+    max_len: int,
+    target_prob: float,
+    eps: float,
+    noise_init: str,
+    noise_sigma: float,
+    dirichlet_init_concentration: float,
+    device: torch.device,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, list[str]]:
+    expanded: list[str] = []
+    prompt_ids: list[list[int]] = []
+    for prompt in prompts:
+        ids = encode_prompt(tokenizer, prompt, max_len=max_len)
+        for _ in range(restarts):
+            expanded.append(prompt)
+            prompt_ids.append(ids)
+
+    batch = len(prompt_ids)
+    attn = torch.ones((batch, max_len), dtype=torch.bool, device=device)
+    probs = sample_noise_simplex(
+        (batch, max_len),
+        tokenizer.vocab_size,
+        device,
+        eps,
+        noise_mode=noise_init,
+        target_prob=target_prob,
+        noise_sigma=noise_sigma,
+        dirichlet_concentration=dirichlet_init_concentration,
+    )
+    lock = torch.zeros((batch, max_len), dtype=torch.bool, device=device)
+    lock_probs = torch.zeros((batch, max_len, tokenizer.vocab_size), dtype=torch.float32, device=device)
+    for row, ids in enumerate(prompt_ids):
+        if not ids:
+            continue
+        ids_t = torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0)
+        sp = smooth_onehot(ids_t, tokenizer.vocab_size, target_prob, eps)[0]
+        probs[row, : len(ids)] = sp
+        lock_probs[row, : len(ids)] = sp
+        lock[row, : len(ids)] = True
+    return probs, attn, lock, lock_probs, expanded
+
+
+def flowmap_gamma(step: int, steps: int, damping: float, max_gamma: float, eps: float) -> float:
+    s = step / max(steps, 1)
+    t_next = (step + 1) / max(steps, 1)
+    base_gamma = (t_next - s) / max(1.0 - s, eps)
+    gamma = float(damping) * base_gamma
+    return min(gamma, float(max_gamma)) if max_gamma > 0 else gamma
+
+
+@torch.no_grad()
+def decode_batch(
+    model,
+    init_probs: torch.Tensor,
+    attn: torch.Tensor,
+    lock: torch.Tensor,
+    lock_probs: torch.Tensor,
+    cfg: DecodeConfig,
+    eps: float,
+    eos_id: int | None = None,
+) -> torch.Tensor:
+    probs = init_probs.float().clone()
+    device = probs.device
+    last_endpoint = probs
+    for step in range(cfg.steps):
+        t = model_time_for_step(cfg.model_t_mode, step, cfg.steps, probs.size(0), device, dtype=torch.float32)
+        logits = model(state_for_model(model, probs, eps), t, attn).float()
+        if cfg.endpoint_temp != 1.0:
+            logits = logits / float(cfg.endpoint_temp)
+        if cfg.eos_logit_bias != 0.0 and eos_id is not None and 0 <= eos_id < logits.size(-1):
+            logits[..., eos_id] = logits[..., eos_id] + float(cfg.eos_logit_bias)
+        endpoint = F.softmax(logits, dim=-1)
+        last_endpoint = endpoint
+
+        if cfg.rule == "flowmap":
+            gamma = flowmap_gamma(step, cfg.steps, cfg.damping, cfg.max_gamma, eps)
+            new_probs = probs + gamma * (endpoint - probs)
+        elif cfg.rule == "replace":
+            new_probs = (1.0 - cfg.eta) * probs + cfg.eta * endpoint
+        elif cfg.rule == "geometric":
+            log_mix = (1.0 - cfg.eta) * torch.log(probs.clamp_min(eps)) + cfg.eta * torch.log(endpoint.clamp_min(eps))
+            new_probs = F.softmax(log_mix, dim=-1)
+        elif cfg.rule == "centered_residual":
+            # Add a zero-sum probability residual, then project back to simplex.
+            residual = endpoint - probs
+            residual = residual - residual.mean(dim=-1, keepdim=True)
+            new_probs = probs + cfg.eta * residual
+        else:
+            raise ValueError(f"Unknown decode rule: {cfg.rule}")
+
+        if cfg.noise_mix > 0:
+            if cfg.noise_decay == "linear":
+                lam = cfg.noise_mix * (1.0 - (step + 1) / max(cfg.steps, 1))
+            elif cfg.noise_decay == "sqrt":
+                lam = cfg.noise_mix * math.sqrt(max(0.0, 1.0 - (step + 1) / max(cfg.steps, 1)))
+            else:
+                lam = cfg.noise_mix
+            if lam > 0:
+                uniform = torch.full_like(new_probs, 1.0 / new_probs.size(-1))
+                new_probs = (1.0 - lam) * new_probs + lam * uniform
+
+        new_probs = new_probs.clamp_min(max(float(cfg.state_floor), eps))
+        new_probs = new_probs / new_probs.sum(dim=-1, keepdim=True).clamp_min(eps)
+        new_probs = torch.where(lock.unsqueeze(-1), lock_probs, new_probs)
+        probs = new_probs
+
+    if cfg.final_from == "endpoint":
+        out = last_endpoint
+        out = torch.where(lock.unsqueeze(-1), lock_probs, out)
+        return out / out.sum(dim=-1, keepdim=True).clamp_min(eps)
+    if cfg.final_from == "blend":
+        out = 0.5 * probs + 0.5 * last_endpoint
+        out = torch.where(lock.unsqueeze(-1), lock_probs, out)
+        return out / out.sum(dim=-1, keepdim=True).clamp_min(eps)
+    return probs
+
+
+@torch.no_grad()
+def pseudo_likelihood_scores(
+    model,
+    tokenizer: BpeTextTokenizer,
+    probs: torch.Tensor,
+    attn: torch.Tensor,
+    lock: torch.Tensor,
+    target_prob: float,
+    eps: float,
+    repeats: int,
+    mask_frac: float,
+    rerank_t: float,
+) -> torch.Tensor:
+    ids = probs.argmax(dim=-1)
+    endpoint = smooth_onehot(ids, tokenizer.vocab_size, target_prob, eps)
+    eligible = attn & (~lock)
+    scores = torch.zeros(ids.size(0), dtype=torch.float32, device=ids.device)
+    counts = torch.zeros_like(scores)
+    for _ in range(max(1, repeats)):
+        score_mask = (torch.rand_like(ids.float()) < mask_frac) & eligible
+        for row in range(ids.size(0)):
+            if eligible[row].any() and not score_mask[row].any():
+                choices = torch.nonzero(eligible[row], as_tuple=False).flatten()
+                score_mask[row, choices[torch.randint(0, choices.numel(), (1,), device=ids.device)]] = True
+        noise = sample_noise_simplex(
+            (ids.size(0), ids.size(1)),
+            tokenizer.vocab_size,
+            ids.device,
+            eps,
+            noise_mode="logistic_normal",
+            target_prob=target_prob,
+            noise_sigma=-1.0,
+        )
+        inp = torch.where(score_mask.unsqueeze(-1), noise, endpoint)
+        inp = torch.where(lock.unsqueeze(-1), probs, inp)
+        t = torch.full((ids.size(0),), float(rerank_t), dtype=torch.float32, device=ids.device)
+        logits = model(state_for_model(model, inp, eps), t, attn).float()
+        logp = F.log_softmax(logits, dim=-1).gather(-1, ids.unsqueeze(-1)).squeeze(-1)
+        scores += (logp * score_mask.float()).sum(dim=-1)
+        counts += score_mask.float().sum(dim=-1)
+    return scores / counts.clamp_min(1.0)
+
+
+def default_configs(steps: int, config_set: str) -> list[DecodeConfig]:
+    if config_set == "focused_flowmap":
+        return [
+            DecodeConfig("flowmap_t1p00_d1p0", "flowmap", steps=steps, damping=1.0, max_gamma=1.0),
+            DecodeConfig("flowmap_t1p10_d1p0", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.10),
+            DecodeConfig("flowmap_t1p25_d1p0", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.25),
+            DecodeConfig("flowmap_t1p40_d1p0", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.40),
+            DecodeConfig("flowmap_t1p60_d1p0", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.60),
+            DecodeConfig("flowmap_t1p25_d0p7", "flowmap", steps=steps, damping=0.7, max_gamma=1.0, endpoint_temp=1.25),
+            DecodeConfig("flowmap_t1p40_d0p7", "flowmap", steps=steps, damping=0.7, max_gamma=1.0, endpoint_temp=1.40),
+            DecodeConfig("flowmap_t1p60_d0p7", "flowmap", steps=steps, damping=0.7, max_gamma=1.0, endpoint_temp=1.60),
+            DecodeConfig("flowmap_t1p25_g0p5", "flowmap", steps=steps, damping=1.0, max_gamma=0.5, endpoint_temp=1.25),
+            DecodeConfig("flowmap_t1p40_g0p5", "flowmap", steps=steps, damping=1.0, max_gamma=0.5, endpoint_temp=1.40),
+        ]
+    if config_set == "best_flowmap":
+        return [
+            DecodeConfig("flowmap_t1p25_d0p7", "flowmap", steps=steps, damping=0.7, max_gamma=1.0, endpoint_temp=1.25),
+            DecodeConfig("flowmap_t1p25_d1p0", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.25),
+            DecodeConfig("flowmap_t1p35_d1p0", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.35),
+            DecodeConfig("flowmap_t1p40_d1p0", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.40),
+        ]
+    if config_set == "final_projection":
+        return [
+            DecodeConfig("flowmap_t1p35_state", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.35, final_from="state"),
+            DecodeConfig("flowmap_t1p35_endpoint", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.35, final_from="endpoint"),
+            DecodeConfig("flowmap_t1p35_blend", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.35, final_from="blend"),
+            DecodeConfig("flowmap_t1p40_state", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.40, final_from="state"),
+            DecodeConfig("flowmap_t1p40_endpoint", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.40, final_from="endpoint"),
+            DecodeConfig("flowmap_t1p40_blend", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.40, final_from="blend"),
+            DecodeConfig("flowmap_t1p25_d0p7_state", "flowmap", steps=steps, damping=0.7, max_gamma=1.0, endpoint_temp=1.25, final_from="state"),
+            DecodeConfig("flowmap_t1p25_d0p7_endpoint", "flowmap", steps=steps, damping=0.7, max_gamma=1.0, endpoint_temp=1.25, final_from="endpoint"),
+            DecodeConfig("flowmap_t1p25_d0p7_blend", "flowmap", steps=steps, damping=0.7, max_gamma=1.0, endpoint_temp=1.25, final_from="blend"),
+        ]
+    if config_set == "eos_sweep":
+        return [
+            DecodeConfig("flowmap_t1p35_eos0", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.35, eos_logit_bias=0.0),
+            DecodeConfig("flowmap_t1p35_eos-1", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.35, eos_logit_bias=-1.0),
+            DecodeConfig("flowmap_t1p35_eos-2", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.35, eos_logit_bias=-2.0),
+            DecodeConfig("flowmap_t1p35_eos-3", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.35, eos_logit_bias=-3.0),
+            DecodeConfig("flowmap_t1p40_eos-2", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.40, eos_logit_bias=-2.0),
+            DecodeConfig("flowmap_t1p25_d0p7_eos-2", "flowmap", steps=steps, damping=0.7, max_gamma=1.0, endpoint_temp=1.25, eos_logit_bias=-2.0),
+        ]
+    if config_set != "broad":
+        raise ValueError(f"Unknown config_set: {config_set}")
+    return [
+        DecodeConfig("flowmap64", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, final_from="state"),
+        DecodeConfig("flowmap_temp1p25", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=1.25),
+        DecodeConfig("flowmap_temp0p85", "flowmap", steps=steps, damping=1.0, max_gamma=1.0, endpoint_temp=0.85),
+        DecodeConfig("replace_eta0p35", "replace", steps=steps, eta=0.35),
+        DecodeConfig("replace_eta0p50", "replace", steps=steps, eta=0.50),
+        DecodeConfig("replace_eta0p65", "replace", steps=steps, eta=0.65),
+        DecodeConfig("replace_eta0p50_temp1p25", "replace", steps=steps, eta=0.50, endpoint_temp=1.25),
+        DecodeConfig("geometric_eta0p25", "geometric", steps=steps, eta=0.25),
+        DecodeConfig("geometric_eta0p50", "geometric", steps=steps, eta=0.50),
+        DecodeConfig("centered_residual_eta0p20", "centered_residual", steps=steps, eta=0.20),
+        DecodeConfig("replace_eta0p50_floor1e6", "replace", steps=steps, eta=0.50, state_floor=1e-6),
+        DecodeConfig("replace_eta0p50_leak", "replace", steps=steps, eta=0.50, noise_mix=0.03, noise_decay="sqrt"),
+    ]
+
+
+def aggregate(rows: list[dict]) -> dict:
+    keys = ["quality", "eos_count", "rep3", "rep4", "distinct1", "distinct2", "digit_frac", "max_word_frac"]
+    return {f"mean_{k}": sum(float(r[k]) for r in rows) / max(len(rows), 1) for k in keys}
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--tokenizer_path", required=True)
+    parser.add_argument("--max_len", type=int, default=128)
+    parser.add_argument("--steps", type=int, default=64)
+    parser.add_argument("--restarts", type=int, default=64)
+    parser.add_argument("--target_prob", type=float, default=0.99)
+    parser.add_argument("--eps", type=float, default=1e-8)
+    parser.add_argument("--model_t_mode", choices=["linear", "flow", "const0", "const05", "const1", "random"], default="flow")
+    parser.add_argument("--noise_init", choices=["uniform", "logistic_normal", "dirichlet"], default="dirichlet")
+    parser.add_argument("--noise_sigma", type=float, default=-1.0)
+    parser.add_argument("--dirichlet_init_concentration", type=float, default=1.0)
+    parser.add_argument("--prompts", default="|The|In the early morning|Scientists have|The company said|A young woman")
+    parser.add_argument("--score_repeats", type=int, default=0)
+    parser.add_argument("--score_mask_frac", type=float, default=0.5)
+    parser.add_argument("--rerank_t", type=float, default=0.5)
+    parser.add_argument("--pl_weight", type=float, default=0.0)
+    parser.add_argument("--output", default="runs/decode_lab/latest_decode_lab.jsonl")
+    parser.add_argument("--config_set", default="broad", choices=["broad", "focused_flowmap", "best_flowmap", "final_projection", "eos_sweep"])
+    parser.add_argument("--decode_batch_size", type=int, default=0)
+    parser.add_argument("--topk", type=int, default=5)
+    parser.add_argument("--seed", type=int, default=20260428)
+    args = parser.parse_args()
+
+    torch.manual_seed(args.seed)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    tokenizer = BpeTextTokenizer.from_file(args.tokenizer_path)
+    ckpt = torch.load(args.checkpoint, map_location="cpu")
+    model = build_model_from_ckpt(ckpt, tokenizer.vocab_size, args.max_len, device)
+    model.eval()
+
+    prompts = args.prompts.split("|")
+    # Keep the first empty prompt: it is unconditional generation.
+    print(f"[info] device={device} prompts={prompts} restarts={args.restarts} steps={args.steps}")
+    print(f"[info] checkpoint={args.checkpoint}")
+
+    out_path = Path(args.output)
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    configs = default_configs(args.steps, args.config_set)
+    for cfg in configs:
+        cfg.model_t_mode = args.model_t_mode
+    with out_path.open("w") as f:
+        for cfg in configs:
+            init, attn, lock, lock_probs, expanded = build_initial_state(
+                tokenizer=tokenizer,
+                prompts=prompts,
+                restarts=args.restarts,
+                max_len=args.max_len,
+                target_prob=args.target_prob,
+                eps=args.eps,
+                noise_init=args.noise_init,
+                noise_sigma=args.noise_sigma,
+                dirichlet_init_concentration=args.dirichlet_init_concentration,
+                device=device,
+            )
+            if args.decode_batch_size > 0 and init.size(0) > args.decode_batch_size:
+                decoded_parts = []
+                for start in range(0, init.size(0), args.decode_batch_size):
+                    end = min(start + args.decode_batch_size, init.size(0))
+                    part = decode_batch(
+                        model,
+                        init[start:end],
+                        attn[start:end],
+                        lock[start:end],
+                        lock_probs[start:end],
+                        cfg,
+                            args.eps,
+                            tokenizer.eos_id,
+                        )
+                    decoded_parts.append(part.detach().cpu())
+                    print(f"[chunk] {cfg.label} decoded {end}/{init.size(0)}", flush=True)
+                decoded = torch.cat(decoded_parts, dim=0)
+            else:
+                decoded = decode_batch(model, init, attn, lock, lock_probs, cfg, args.eps, tokenizer.eos_id)
+            ids = decoded.argmax(dim=-1).detach().cpu().tolist()
+            texts = [decode_text(tokenizer, row) for row in ids]
+            rows = []
+            for i, text in enumerate(texts):
+                m = text_metrics(text)
+                m.update({"candidate": i, "prompt": expanded[i], "text": text})
+                rows.append(m)
+            if args.score_repeats > 0:
+                decoded_for_score = decoded.to(device) if decoded.device != device else decoded
+                pl = pseudo_likelihood_scores(
+                    model,
+                    tokenizer,
+                    decoded_for_score,
+                    attn,
+                    lock,
+                    args.target_prob,
+                    args.eps,
+                    repeats=args.score_repeats,
+                    mask_frac=args.score_mask_frac,
+                    rerank_t=args.rerank_t,
+                ).detach().cpu().tolist()
+                for row, score in zip(rows, pl):
+                    row["pseudo_logp"] = float(score)
+                    row["rank_score"] = float(row["quality"] + args.pl_weight * score)
+            else:
+                for row in rows:
+                    row["pseudo_logp"] = None
+                    row["rank_score"] = float(row["quality"])
+
+            summary = {"type": "summary", "config": asdict(cfg), "agg": aggregate(rows)}
+            f.write(json.dumps(summary, ensure_ascii=False) + "\n")
+            print("\n" + "=" * 96)
+            print("[config]", cfg.label, asdict(cfg))
+            print("[metrics]", json.dumps(summary["agg"], ensure_ascii=False))
+            for prompt in prompts:
+                subset = [r for r in rows if r["prompt"] == prompt]
+                subset.sort(key=lambda r: r["rank_score"], reverse=True)
+                for rank, row in enumerate(subset[: args.topk], 1):
+                    rec = {"type": "sample", "config": asdict(cfg), "rank": rank, **row}
+                    f.write(json.dumps(rec, ensure_ascii=False) + "\n")
+                    if rank <= 1:
+                        print(f"\n--- best prompt={prompt!r} rank_score={row['rank_score']:.4f} quality={row['quality']:.4f} ---")
+                        print(row["text"])
+
+            del init, attn, lock, lock_probs, decoded
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+
+    print(f"[done] wrote {out_path}")
+
+
+if __name__ == "__main__":
+    main()

LTA_openwebtext_dualt/scripts/eval_lm1b_linear_simplex_genppl.py

@@ -0,0 +1,314 @@
+#!/usr/bin/env python3
+"""Algebraic simplex-linear GenPPL eval for endpoint models.
+
+This decoder matches the supervised bridge:
+
+    p_t = (1 - t) * p0 + t * x1
+
+Inference keeps the sampled p0 fixed and replaces the unknown x1 with the
+model's current endpoint prediction:
+
+    p_{t_next} = (1 - t_next) * p0 + t_next * a_theta(p_t, t).
+
+There is no Dirichlet/Gamma resampling in the loop.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from flowtext_lab.decode import sample_noise_simplex, state_for_model
+from flowtext_lab.genppl import filter_generated_texts, summarize_token_diversity
+from flowtext_lab.tokenization import BpeTextTokenizer
+
+from eval_lm1b_c1024_fullycoupled_sde_genppl import (
+    build_model,
+    collect_special_token_ids,
+    filter_endpoint_probs,
+    score_with_gpt2,
+)
+
+
+def lerp(a: float, b: float, t: float) -> float:
+    return float(a) + float(t) * (float(b) - float(a))
+
+
+def project_endpoint(
+    logits: torch.Tensor,
+    *,
+    temp: float,
+    projection: str,
+    top_k: int,
+    top_p: float,
+    banned_ids: list[int],
+    gumbel_tau: float,
+    gumbel_noise_scale: float,
+    eps: float,
+) -> torch.Tensor:
+    endpoint = F.softmax(logits / max(float(temp), eps), dim=-1)
+    endpoint = filter_endpoint_probs(
+        endpoint,
+        top_k=top_k,
+        top_p=top_p,
+        banned_ids=banned_ids,
+        eps=eps,
+    )
+    if projection == "soft":
+        return endpoint
+    if projection == "argmax":
+        ids = endpoint.argmax(dim=-1)
+        return torch.zeros_like(endpoint).scatter_(-1, ids.unsqueeze(-1), 1.0)
+    if projection == "sample":
+        ids = torch.multinomial(endpoint.reshape(-1, endpoint.size(-1)), 1).view(*endpoint.shape[:-1])
+        return torch.zeros_like(endpoint).scatter_(-1, ids.unsqueeze(-1), 1.0)
+    if projection == "gumbel_softmax":
+        u = torch.rand_like(endpoint).clamp_(min=eps, max=1.0 - eps)
+        g = -torch.log(-torch.log(u))
+        z = (endpoint.clamp_min(eps).log() + float(gumbel_noise_scale) * g) / max(float(gumbel_tau), eps)
+        y = F.softmax(z, dim=-1).clamp_min(eps)
+        return y / y.sum(dim=-1, keepdim=True).clamp_min(eps)
+    raise ValueError(f"unknown endpoint_projection: {projection}")
+
+
+@torch.inference_mode()
+def decode_linear_simplex(
+    model,
+    tokenizer: BpeTextTokenizer,
+    *,
+    n_samples: int,
+    batch_size: int,
+    max_len: int,
+    steps: int,
+    seed: int,
+    device: torch.device,
+    noise_init: str,
+    noise_sigma: float,
+    noise_dirichlet_concentration: float,
+    endpoint_temp_start: float,
+    endpoint_temp_end: float,
+    endpoint_projection: str,
+    endpoint_top_k: int,
+    endpoint_top_p: float,
+    ban_special_tokens: bool,
+    gumbel_tau_start: float,
+    gumbel_tau_end: float,
+    gumbel_noise_scale_start: float,
+    gumbel_noise_scale_end: float,
+    final_from: str,
+) -> tuple[list[list[int]], list[str], dict]:
+    torch.manual_seed(seed)
+    eps = 1e-8
+    all_ids: list[list[int]] = []
+    all_texts: list[str] = []
+    remaining = n_samples
+    banned_endpoint_ids = collect_special_token_ids(tokenizer) if ban_special_tokens else []
+
+    while remaining > 0:
+        bs = min(batch_size, remaining)
+        p0 = sample_noise_simplex(
+            (bs, max_len),
+            tokenizer.vocab_size,
+            device,
+            eps,
+            noise_mode=noise_init,
+            target_prob=1.0,
+            noise_sigma=noise_sigma,
+            dirichlet_concentration=noise_dirichlet_concentration,
+        )
+        probs = p0.clone()
+        attn = torch.ones((bs, max_len), dtype=torch.bool, device=device)
+        last_endpoint = probs
+
+        for step in range(steps):
+            cur_t = step / max(steps, 1)
+            next_t = (step + 1) / max(steps, 1)
+            t = torch.full((bs,), float(cur_t), dtype=torch.float32, device=device)
+            logits = model(state_for_model(model, probs, eps), t, attn).float()
+
+            endpoint = project_endpoint(
+                logits,
+                temp=lerp(endpoint_temp_start, endpoint_temp_end, cur_t),
+                projection=endpoint_projection,
+                top_k=endpoint_top_k,
+                top_p=endpoint_top_p,
+                banned_ids=banned_endpoint_ids,
+                gumbel_tau=lerp(gumbel_tau_start, gumbel_tau_end, cur_t),
+                gumbel_noise_scale=lerp(gumbel_noise_scale_start, gumbel_noise_scale_end, cur_t),
+                eps=eps,
+            )
+            last_endpoint = endpoint
+            probs = (1.0 - next_t) * p0 + next_t * endpoint
+            probs = probs.clamp_min(eps)
+            probs = probs / probs.sum(dim=-1, keepdim=True).clamp_min(eps)
+
+        if final_from == "blend_0.5":
+            final_probs = 0.5 * probs + 0.5 * last_endpoint
+            ids = final_probs.argmax(dim=-1).detach().cpu().tolist()
+        elif final_from == "model_t1":
+            t = torch.ones((bs,), dtype=torch.float32, device=device)
+            final_logits = model(state_for_model(model, probs, eps), t, attn).float()
+            ids = final_logits.argmax(dim=-1).detach().cpu().tolist()
+        else:
+            raise ValueError(f"unknown final_from: {final_from}")
+
+        all_ids.extend(ids)
+        all_texts.extend(tokenizer.decode(row, stop_at_eos=False, skip_special_tokens=False) for row in ids)
+        remaining -= bs
+        print(f"[linear] generated {n_samples - remaining}/{n_samples}", flush=True)
+
+    cfg = {
+        "decode_rule": "linear_simplex_algebraic",
+        "steps": steps,
+        "noise_init": noise_init,
+        "noise_sigma": noise_sigma,
+        "noise_dirichlet_concentration": noise_dirichlet_concentration,
+        "endpoint_temp_start": endpoint_temp_start,
+        "endpoint_temp_end": endpoint_temp_end,
+        "endpoint_projection": endpoint_projection,
+        "endpoint_top_k": endpoint_top_k,
+        "endpoint_top_p": endpoint_top_p,
+        "ban_special_tokens": ban_special_tokens,
+        "banned_endpoint_ids": banned_endpoint_ids,
+        "gumbel_tau_start": gumbel_tau_start,
+        "gumbel_tau_end": gumbel_tau_end,
+        "gumbel_noise_scale_start": gumbel_noise_scale_start,
+        "gumbel_noise_scale_end": gumbel_noise_scale_end,
+        "final_from": final_from,
+        "n_samples": n_samples,
+        "seed": seed,
+    }
+    return all_ids, all_texts, cfg
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Linear-simplex algebraic GenPPL eval")
+    p.add_argument("--checkpoint", required=True)
+    p.add_argument("--tokenizer_path", required=True)
+    p.add_argument("--scorer", required=True)
+    p.add_argument("--out_dir", required=True)
+    p.add_argument("--n_samples", type=int, default=128)
+    p.add_argument("--max_len", type=int, default=128)
+    p.add_argument("--steps", type=int, default=128)
+    p.add_argument("--batch_size", type=int, default=16)
+    p.add_argument("--score_batch", type=int, default=8)
+    p.add_argument("--score_max_length", type=int, default=1024)
+    p.add_argument("--noise_init", choices=["uniform", "logistic_normal", "dirichlet"], default="logistic_normal")
+    p.add_argument("--noise_sigma", type=float, default=3.0)
+    p.add_argument("--noise_dirichlet_concentration", type=float, default=1.0)
+    p.add_argument("--endpoint_temp_start", type=float, default=1.45)
+    p.add_argument("--endpoint_temp_end", type=float, default=0.8)
+    p.add_argument("--endpoint_projection", choices=["soft", "sample", "argmax", "gumbel_softmax"], default="soft")
+    p.add_argument("--endpoint_top_k", type=int, default=0)
+    p.add_argument("--endpoint_top_p", type=float, default=1.0)
+    p.add_argument("--ban_special_tokens", action="store_true")
+    p.add_argument("--gumbel_tau_start", type=float, default=1.0)
+    p.add_argument("--gumbel_tau_end", type=float, default=0.2)
+    p.add_argument("--gumbel_noise_scale_start", type=float, default=1.0)
+    p.add_argument("--gumbel_noise_scale_end", type=float, default=0.0)
+    p.add_argument("--final_from", choices=["blend_0.5", "model_t1"], default="model_t1")
+    p.add_argument("--seed", type=int, default=20260524)
+    return p.parse_args()
+
+
+@torch.no_grad()
+def main() -> None:
+    args = parse_args()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"[load] {args.checkpoint}", flush=True)
+    ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
+    step = ckpt.get("step")
+    print(f"[ckpt] step={step}", flush=True)
+
+    tokenizer = BpeTextTokenizer.from_file(args.tokenizer_path)
+    model = build_model(ckpt, tokenizer, device)
+    ids, texts, decode_cfg = decode_linear_simplex(
+        model,
+        tokenizer,
+        n_samples=args.n_samples,
+        batch_size=args.batch_size,
+        max_len=args.max_len,
+        steps=args.steps,
+        seed=args.seed,
+        device=device,
+        noise_init=args.noise_init,
+        noise_sigma=args.noise_sigma,
+        noise_dirichlet_concentration=args.noise_dirichlet_concentration,
+        endpoint_temp_start=args.endpoint_temp_start,
+        endpoint_temp_end=args.endpoint_temp_end,
+        endpoint_projection=args.endpoint_projection,
+        endpoint_top_k=args.endpoint_top_k,
+        endpoint_top_p=args.endpoint_top_p,
+        ban_special_tokens=args.ban_special_tokens,
+        gumbel_tau_start=args.gumbel_tau_start,
+        gumbel_tau_end=args.gumbel_tau_end,
+        gumbel_noise_scale_start=args.gumbel_noise_scale_start,
+        gumbel_noise_scale_end=args.gumbel_noise_scale_end,
+        final_from=args.final_from,
+    )
+    del model
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+    def strip_special(t: str) -> str:
+        import re
+        t = t.replace("[CLS]", " ").replace("[SEP]", " ").replace("[PAD]", " ")
+        t = t.replace("<|endoftext|>", " ")
+        return re.sub(r"\s+", " ", t).strip()
+
+    stripped = [strip_special(t) for t in texts]
+    kept_raw, _ = filter_generated_texts(texts, min_chars=1, normalize_whitespace=False, drop_empty=True)
+    kept_stripped, _ = filter_generated_texts(stripped, min_chars=1, normalize_whitespace=True, drop_empty=True)
+    diversity = summarize_token_diversity(ids).__dict__
+
+    print(f"[score] loading scorer: {args.scorer}", flush=True)
+    scorer_tok = AutoTokenizer.from_pretrained(args.scorer)
+    if scorer_tok.pad_token_id is None:
+        scorer_tok.pad_token = scorer_tok.eos_token
+        scorer_tok.pad_token_id = scorer_tok.eos_token_id
+    scorer = AutoModelForCausalLM.from_pretrained(args.scorer).to(device).eval()
+    if getattr(scorer.config, "pad_token_id", None) is None:
+        scorer.config.pad_token_id = scorer_tok.pad_token_id
+
+    raw_ppl = score_with_gpt2(
+        kept_raw, scorer, scorer_tok,
+        batch_size=args.score_batch, max_length=args.score_max_length, device=device,
+    )
+    stripped_ppl = score_with_gpt2(
+        kept_stripped, scorer, scorer_tok,
+        batch_size=args.score_batch, max_length=args.score_max_length, device=device,
+    )
+    summary = {
+        "type": "summary",
+        "checkpoint": args.checkpoint,
+        "step": step,
+        "decode": decode_cfg,
+        "raw_genppl": raw_ppl,
+        "stripped_genppl": stripped_ppl,
+        "diversity": diversity,
+    }
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    out_jsonl = out_dir / f"linear_steps{args.steps}_samples{args.n_samples}_scored.jsonl"
+    with out_jsonl.open("w", encoding="utf-8") as f:
+        f.write(json.dumps(summary, ensure_ascii=False) + "\n")
+        for i, (raw, clean) in enumerate(zip(texts, stripped)):
+            f.write(json.dumps({"type": "sample", "index": i, "raw_text": raw, "stripped_text": clean}, ensure_ascii=False) + "\n")
+    print("[summary]", json.dumps(summary, ensure_ascii=False, indent=2), flush=True)
+    print(f"[done] {out_jsonl}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

LTA_openwebtext_dualt/scripts/eval_train8_overfit_sweep.sh

@@ -0,0 +1,54 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+
+STAMP="${STAMP:?set STAMP used by run_train8_overfit_sweep_4gpu.sh}"
+TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-standard/tokenizer.json}"
+CACHE_DIR="${CACHE_DIR:-/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext_lta_cached_chunks/gpt2_len1024_train8_overfit}"
+OUT_ROOT="${OUT_ROOT:-docs/lta_samples/metrics_20260517/train8_overfit_${STAMP}}"
+export OUT_ROOT
+mkdir -p "${OUT_ROOT}"
+
+for run in \
+  "train8_n1024_hard_ce_onehot_${STAMP}" \
+  "train8_n1024_hard_ce_bridge_${STAMP}" \
+  "train8_n1024_linear_soft_kl_onehot_${STAMP}" \
+  "train8_n1024_linear_soft_kl_bridge_${STAMP}"
+do
+  echo "[eval-overfit] ${run}"
+  python scripts/eval_train8_overfit_ckpts.py \
+    --run_dir "runs/${run}" \
+    --cache_dir "${CACHE_DIR}" \
+    --tokenizer_path "${TOKENIZER_PATH}" \
+    --out_dir "${OUT_ROOT}/${run}" \
+    --max_len 1024 \
+    --limit 8 \
+    --t_values 0.125,0.25,0.5,0.75,1.0 \
+    --seeds 123,456,789
+done
+
+python - <<'PY'
+import json
+import os
+from pathlib import Path
+root = Path(os.environ["OUT_ROOT"])
+rows = []
+for p in sorted(root.glob("*/result.json")):
+    r = json.loads(p.read_text())
+    last = r["last"]
+    rows.append({
+        "run": p.parent.name,
+        "first_fit_step": r["first_fit_step"],
+        "best_acc": r["best_acc"],
+        "best_objective": r["best_objective"],
+        "last_acc": last["gold_acc_mean"],
+        "last_objective": last["objective_mean"],
+        "last_ce": last["gold_ce_mean"],
+    })
+with (root / "combined.tsv").open("w") as f:
+    f.write("run\tfirst_fit_step\tbest_acc\tbest_objective\tlast_acc\tlast_objective\tlast_ce\n")
+    for r in rows:
+        f.write("\t".join(str(r[k]) for k in ["run","first_fit_step","best_acc","best_objective","last_acc","last_objective","last_ce"]) + "\n")
+print(root / "combined.tsv")
+PY

LTA_openwebtext_dualt/scripts/launch_ar_openwebtext_duo_small_8gpu_1m.sh

@@ -0,0 +1,99 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}"
+export TOKENIZERS_PARALLELISM=false
+export PYTHONUNBUFFERED=1
+export NCCL_DEBUG="${NCCL_DEBUG:-WARN}"
+export TORCH_DISTRIBUTED_TIMEOUT="${TORCH_DISTRIBUTED_TIMEOUT:-3600}"
+
+RUN_NAME="${RUN_NAME:-ar_owt_flmpack_gpt2_small_len1024_gbs512_8gpu_1m}"
+DATA_PATH="${DATA_PATH:-/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext}"
+TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-standard/tokenizer.json}"
+TEXT_COLUMN="${TEXT_COLUMN:-text}"
+OPENWEBTEXT_SPLIT="${OPENWEBTEXT_SPLIT:-train_minus_100k}"
+SAVE_DIR="${SAVE_DIR:-runs/${RUN_NAME}}"
+LOG_FILE="${LOG_FILE:-logs/${RUN_NAME}.log}"
+
+NNODES="${NNODES:-1}"
+NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+NODE_RANK="${NODE_RANK:-0}"
+MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+MASTER_PORT="${MASTER_PORT:-29643}"
+
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-32}"
+TOTAL_STEPS="${TOTAL_STEPS:-1000000}"
+WARMUP_STEPS="${WARMUP_STEPS:-2500}"
+MAX_LEN="${MAX_LEN:-1024}"
+NUM_WORKERS="${NUM_WORKERS:-4}"
+LOG_EVERY="${LOG_EVERY:-100}"
+SAVE_EVERY="${SAVE_EVERY:-20000}"
+LATEST_EVERY="${LATEST_EVERY:-1000}"
+SAMPLE_EVERY="${SAMPLE_EVERY:-0}"
+RESUME_PATH="${RESUME_PATH:-}"
+ALLOW_EXISTING_SAVE_DIR="${ALLOW_EXISTING_SAVE_DIR:-0}"
+ENABLE_TORCH_COMPILE="${ENABLE_TORCH_COMPILE:-0}"
+
+COMPILE_ARGS=()
+if [[ "${ENABLE_TORCH_COMPILE}" == "1" ]]; then
+  COMPILE_ARGS+=(--torch_compile --compile_mode reduce-overhead)
+fi
+RESUME_ARGS=()
+if [[ -n "${RESUME_PATH}" ]]; then
+  RESUME_ARGS+=(--resume_path "${RESUME_PATH}")
+fi
+TEXT_COLUMN_ARGS=()
+if [[ -n "${TEXT_COLUMN}" ]]; then
+  TEXT_COLUMN_ARGS+=(--text_column "${TEXT_COLUMN}")
+fi
+
+if [[ -f "${SAVE_DIR}/args.json" && -z "${RESUME_PATH}" && "${ALLOW_EXISTING_SAVE_DIR}" != "1" ]]; then
+  echo "Refusing to start because SAVE_DIR already contains args.json: ${SAVE_DIR}" >&2
+  echo "Use a new RUN_NAME/SAVE_DIR, set RESUME_PATH to resume, or set ALLOW_EXISTING_SAVE_DIR=1 intentionally." >&2
+  exit 2
+fi
+
+mkdir -p logs runs "${SAVE_DIR}"
+
+python -m torch.distributed.run \
+  --nnodes="${NNODES}" \
+  --nproc_per_node="${NPROC_PER_NODE}" \
+  --node_rank="${NODE_RANK}" \
+  --master_addr="${MASTER_ADDR}" \
+  --master_port="${MASTER_PORT}" \
+  train_ar.py \
+  --data_path "${DATA_PATH}" \
+  "${TEXT_COLUMN_ARGS[@]}" \
+  --openwebtext_split "${OPENWEBTEXT_SPLIT}" \
+  --tokenizer_path "${TOKENIZER_PATH}" \
+  --save_dir "${SAVE_DIR}" \
+  --wrap \
+  --max_len "${MAX_LEN}" \
+  --batch_size "${PER_GPU_BATCH_SIZE}" \
+  --num_workers "${NUM_WORKERS}" \
+  --global_batch_size "${GLOBAL_BATCH_SIZE}" \
+  --total_steps "${TOTAL_STEPS}" \
+  --log_every "${LOG_EVERY}" \
+  --save_every "${SAVE_EVERY}" \
+  --latest_every "${LATEST_EVERY}" \
+  --sample_every "${SAMPLE_EVERY}" \
+  --lr 3e-4 \
+  --weight_decay 0 \
+  --adam_beta1 0.9 \
+  --adam_beta2 0.999 \
+  --adam_eps 1e-8 \
+  --warmup_steps "${WARMUP_STEPS}" \
+  --lr_schedule constant_warmup \
+  --grad_clip 1.0 \
+  --seed 123 \
+  --d_model 768 \
+  --n_layers 12 \
+  --n_heads 12 \
+  --dim_ff 3072 \
+  --dropout 0.1 \
+  --tie_embeddings \
+  "${RESUME_ARGS[@]}" \
+  "${COMPILE_ARGS[@]}" \
+  --bf16 2>&1 | tee -a "${LOG_FILE}"

LTA_openwebtext_dualt/scripts/launch_lm1b_flm_8gpu_repro_20260506.sh

@@ -0,0 +1,123 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}"
+export TOKENIZERS_PARALLELISM=false
+export PYTHONUNBUFFERED=1
+export NCCL_DEBUG="${NCCL_DEBUG:-WARN}"
+export TORCH_DISTRIBUTED_TIMEOUT="${TORCH_DISTRIBUTED_TIMEOUT:-3600}"
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
+
+# Unified LM1B FLM baseline.
+# This uses the same tokenizer/data/model/optimizer setup as MDLM/DUO.
+
+RUN_TAG="${RUN_TAG:-20260506_repro}"
+RUN_NAME="${RUN_NAME:-lm1b_flm_unified_ddit_small_len128_gbs512_8gpu_1m_${RUN_TAG}}"
+DATA_PATH="${DATA_PATH:-data/lm1b_train_parquet}"
+TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/workspace/imagenet_handoff_20260327/nlp_dts_light/assets/distilbert-base-uncased/tokenizer.json}"
+TEXT_COLUMN="${TEXT_COLUMN:-}"
+OPENWEBTEXT_SPLIT="${OPENWEBTEXT_SPLIT:-all}"
+SAVE_DIR="${SAVE_DIR:-runs/${RUN_NAME}}"
+LOG_FILE="${LOG_FILE:-logs/${RUN_NAME}.log}"
+
+NNODES="${NNODES:-1}"
+NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+NODE_RANK="${NODE_RANK:-0}"
+MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+MASTER_PORT="${MASTER_PORT:-29646}"
+
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-64}"
+TOTAL_STEPS="${TOTAL_STEPS:-1000000}"
+WARMUP_STEPS="${WARMUP_STEPS:-2500}"
+MAX_LEN="${MAX_LEN:-128}"
+WRAP_MODE="${WRAP_MODE:-stream}"
+WRAP_RECORD_BUFFER_SIZE="${WRAP_RECORD_BUFFER_SIZE:-200}"
+NUM_WORKERS="${NUM_WORKERS:-0}"
+LOG_EVERY="${LOG_EVERY:-100}"
+SAVE_EVERY="${SAVE_EVERY:-20000}"
+LATEST_EVERY="${LATEST_EVERY:-1000}"
+RESUME_PATH="${RESUME_PATH:-}"
+ALLOW_EXISTING_SAVE_DIR="${ALLOW_EXISTING_SAVE_DIR:-0}"
+ENABLE_TORCH_COMPILE="${ENABLE_TORCH_COMPILE:-0}"
+FORCE_DISABLE_TORCH_COMPILE="${FORCE_DISABLE_TORCH_COMPILE:-1}"
+
+if [[ "${FORCE_DISABLE_TORCH_COMPILE}" == "1" ]]; then
+  ENABLE_TORCH_COMPILE=0
+fi
+if [[ "${DATA_PATH}" == *"lm1b_train_parquet"* && "${NUM_WORKERS}" != "0" ]]; then
+  echo "LM1B has only 9 parquet shards; forcing NUM_WORKERS=0 to avoid empty DDP dataloader shards." >&2
+  NUM_WORKERS=0
+fi
+
+COMPILE_ARGS=()
+if [[ "${ENABLE_TORCH_COMPILE}" == "1" ]]; then
+  COMPILE_ARGS+=(--torch_compile --compile_mode reduce-overhead)
+fi
+RESUME_ARGS=()
+if [[ -n "${RESUME_PATH}" ]]; then
+  RESUME_ARGS+=(--resume_path "${RESUME_PATH}")
+fi
+TEXT_COLUMN_ARGS=()
+if [[ -n "${TEXT_COLUMN}" ]]; then
+  TEXT_COLUMN_ARGS+=(--text_column "${TEXT_COLUMN}")
+fi
+
+if [[ -f "${SAVE_DIR}/args.json" && -z "${RESUME_PATH}" && "${ALLOW_EXISTING_SAVE_DIR}" != "1" ]]; then
+  echo "Refusing to start because SAVE_DIR already contains args.json: ${SAVE_DIR}" >&2
+  echo "Use a new RUN_NAME/SAVE_DIR, set RESUME_PATH to resume, or set ALLOW_EXISTING_SAVE_DIR=1 intentionally." >&2
+  exit 2
+fi
+
+mkdir -p logs runs "${SAVE_DIR}"
+echo "[launch] method=flm host=$(hostname) time=$(date -Iseconds)"
+echo "[launch] cwd=$(pwd)"
+echo "[launch] run_name=${RUN_NAME}"
+echo "[launch] save_dir=${SAVE_DIR}"
+echo "[launch] log_file=${LOG_FILE}"
+echo "[launch] nproc_per_node=${NPROC_PER_NODE} global_batch_size=${GLOBAL_BATCH_SIZE} per_gpu_batch_size=${PER_GPU_BATCH_SIZE}"
+
+python -m torch.distributed.run \
+  --nnodes="${NNODES}" \
+  --nproc_per_node="${NPROC_PER_NODE}" \
+  --node_rank="${NODE_RANK}" \
+  --master_addr="${MASTER_ADDR}" \
+  --master_port="${MASTER_PORT}" \
+  train_baseline.py \
+  --baseline_method flm \
+  --data_path "${DATA_PATH}" \
+  "${TEXT_COLUMN_ARGS[@]}" \
+  --openwebtext_split "${OPENWEBTEXT_SPLIT}" \
+  --tokenizer_path "${TOKENIZER_PATH}" \
+  --save_dir "${SAVE_DIR}" \
+  --max_len "${MAX_LEN}" \
+  --batch_size "${PER_GPU_BATCH_SIZE}" \
+  --num_workers "${NUM_WORKERS}" \
+  --global_batch_size "${GLOBAL_BATCH_SIZE}" \
+  --wrap_mode "${WRAP_MODE}" \
+  --wrap_record_buffer_size "${WRAP_RECORD_BUFFER_SIZE}" \
+  --total_steps "${TOTAL_STEPS}" \
+  --log_every "${LOG_EVERY}" \
+  --save_every "${SAVE_EVERY}" \
+  --latest_every "${LATEST_EVERY}" \
+  --lr 3e-4 \
+  --weight_decay 0 \
+  --adam_beta1 0.9 \
+  --adam_beta2 0.999 \
+  --adam_eps 1e-8 \
+  --warmup_steps "${WARMUP_STEPS}" \
+  --lr_schedule constant_warmup \
+  --grad_clip 1.0 \
+  --seed 123 \
+  --d_model 768 \
+  --cond_dim 128 \
+  --n_layers 12 \
+  --n_heads 12 \
+  --dim_ff 3072 \
+  --dropout 0.1 \
+  --model_type ddit \
+  --sampling_eps 1e-3 \
+  "${RESUME_ARGS[@]}" \
+  "${COMPILE_ARGS[@]}" \
+  --bf16 2>&1 | tee -a "${LOG_FILE}"

LTA_openwebtext_dualt/scripts/launch_lta_lm1b_categorical_fullvocab_c1024_gaussian_linear_mean_fullycoupled_8gpu_small_1m.sh

@@ -0,0 +1,156 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}"
+export TOKENIZERS_PARALLELISM=false
+export PYTHONUNBUFFERED=1
+export NCCL_DEBUG="${NCCL_DEBUG:-WARN}"
+export TORCH_DISTRIBUTED_TIMEOUT="${TORCH_DISTRIBUTED_TIMEOUT:-3600}"
+
+# Gaussian/logistic-normal center-mixture with probability-linear mean:
+#   model_t == support t == semantic endpoint t
+#   endpoint is categorical full-vocab, state is logistic-normal around
+#   mu_c(t) = (1 - t) uniform + t onehot(center)
+RUN_NAME="${RUN_NAME:-lta_lm1b_logisticnormal_linearmean_categorical_fullvocab_c1024_fullycoupled_flmpack_onehot_hardce_ddit_small_len128_gbs512_8gpu_1m_nw0}"
+DATA_PATH="${DATA_PATH:-data/lm1b_train_parquet}"
+TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/workspace/imagenet_handoff_20260327/nlp_dts_light/assets/distilbert-base-uncased/tokenizer.json}"
+TEXT_COLUMN="${TEXT_COLUMN:-}"
+OPENWEBTEXT_SPLIT="${OPENWEBTEXT_SPLIT:-all}"
+SAVE_DIR="${SAVE_DIR:-runs/${RUN_NAME}}"
+LOG_FILE="${LOG_FILE:-logs/${RUN_NAME}.log}"
+
+NNODES="${NNODES:-1}"
+NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+NODE_RANK="${NODE_RANK:-0}"
+MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+MASTER_PORT="${MASTER_PORT:-29631}"
+
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-64}"
+TOTAL_STEPS="${TOTAL_STEPS:-1000000}"
+WARMUP_STEPS="${WARMUP_STEPS:-2500}"
+MAX_LEN="${MAX_LEN:-128}"
+WRAP_MODE="${WRAP_MODE:-stream}"
+WRAP_RECORD_BUFFER_SIZE="${WRAP_RECORD_BUFFER_SIZE:-200}"
+NUM_WORKERS="${NUM_WORKERS:-0}"
+LOG_EVERY="${LOG_EVERY:-100}"
+SAVE_EVERY="${SAVE_EVERY:-20000}"
+LATEST_EVERY="${LATEST_EVERY:-1000}"
+EVAL_EVERY="${EVAL_EVERY:-0}"
+RESUME_PATH="${RESUME_PATH:-}"
+ALLOW_EXISTING_SAVE_DIR="${ALLOW_EXISTING_SAVE_DIR:-0}"
+ENABLE_TORCH_COMPILE="${ENABLE_TORCH_COMPILE:-0}"
+FORCE_DISABLE_TORCH_COMPILE="${FORCE_DISABLE_TORCH_COMPILE:-1}"
+
+if [[ "${FORCE_DISABLE_TORCH_COMPILE}" == "1" ]]; then
+  ENABLE_TORCH_COMPILE=0
+fi
+if [[ "${DATA_PATH}" == *"lm1b_train_parquet"* && "${NUM_WORKERS}" != "0" ]]; then
+  echo "LM1B has only 9 parquet shards; forcing NUM_WORKERS=0 to avoid empty DDP dataloader shards." >&2
+  NUM_WORKERS=0
+fi
+
+COMPILE_ARGS=()
+if [[ "${ENABLE_TORCH_COMPILE}" == "1" ]]; then
+  COMPILE_ARGS+=(--torch_compile --compile_mode reduce-overhead)
+fi
+RESUME_ARGS=()
+if [[ -n "${RESUME_PATH}" ]]; then
+  RESUME_ARGS+=(--resume_path "${RESUME_PATH}")
+fi
+TEXT_COLUMN_ARGS=()
+if [[ -n "${TEXT_COLUMN}" ]]; then
+  TEXT_COLUMN_ARGS+=(--text_column "${TEXT_COLUMN}")
+fi
+
+if [[ -f "${SAVE_DIR}/args.json" && -z "${RESUME_PATH}" && "${ALLOW_EXISTING_SAVE_DIR}" != "1" ]]; then
+  echo "Refusing to start because SAVE_DIR already contains args.json: ${SAVE_DIR}" >&2
+  echo "Use a new RUN_NAME/SAVE_DIR, set RESUME_PATH to resume, or set ALLOW_EXISTING_SAVE_DIR=1 intentionally." >&2
+  exit 2
+fi
+
+mkdir -p logs runs "${SAVE_DIR}"
+echo "[launch] method=logisticnormal_linearmean_categorical_fullvocab_c1024_fullycoupled host=$(hostname) time=$(date -Iseconds)"
+echo "[launch] cwd=$(pwd)"
+echo "[launch] run_name=${RUN_NAME}"
+echo "[launch] save_dir=${SAVE_DIR}"
+echo "[launch] log_file=${LOG_FILE}"
+
+python -m torch.distributed.run \
+  --nnodes="${NNODES}" \
+  --nproc_per_node="${NPROC_PER_NODE}" \
+  --node_rank="${NODE_RANK}" \
+  --master_addr="${MASTER_ADDR}" \
+  --master_port="${MASTER_PORT}" \
+  train.py \
+  --data_path "${DATA_PATH}" \
+  "${TEXT_COLUMN_ARGS[@]}" \
+  --openwebtext_split "${OPENWEBTEXT_SPLIT}" \
+  --tokenizer_path "${TOKENIZER_PATH}" \
+  --save_dir "${SAVE_DIR}" \
+  --wrap \
+  --wrap_mode "${WRAP_MODE}" \
+  --wrap_record_buffer_size "${WRAP_RECORD_BUFFER_SIZE}" \
+  --max_len "${MAX_LEN}" \
+  --batch_size "${PER_GPU_BATCH_SIZE}" \
+  --num_workers "${NUM_WORKERS}" \
+  --global_batch_size "${GLOBAL_BATCH_SIZE}" \
+  --total_steps "${TOTAL_STEPS}" \
+  --log_every "${LOG_EVERY}" \
+  --eval_every "${EVAL_EVERY}" \
+  --save_every "${SAVE_EVERY}" \
+  --latest_every "${LATEST_EVERY}" \
+  --lr 3e-4 \
+  --weight_decay 0 \
+  --adam_beta1 0.9 \
+  --adam_beta2 0.999 \
+  --adam_eps 1e-8 \
+  --warmup_steps "${WARMUP_STEPS}" \
+  --lr_schedule constant_warmup \
+  --grad_clip 1.0 \
+  --seed 123 \
+  --d_model 768 \
+  --cond_dim 128 \
+  --n_layers 12 \
+  --n_heads 12 \
+  --dim_ff 3072 \
+  --dropout 0.1 \
+  --model_type ddit \
+  --state_format prob \
+  --bridge dirichlet \
+  --target_loss hard_ce \
+  --target_prob 1.0 \
+  --min_t 0.0 \
+  --max_t 1.0 \
+  --dual_t \
+  --corrupt_t_mode same \
+  --corrupt_min_t 0.0 \
+  --corrupt_max_t 1.0 \
+  --min_mask_ratio 0.1 \
+  --max_mask_ratio 1.0 \
+  --wrong_token_replace_prob 1.0 \
+  --wrong_token_schedule linear_t \
+  --wrong_token_exp_k 1.0 \
+  --dirichlet_concentration_min 1.0 \
+  --dirichlet_concentration_max 1024.0 \
+  --dirichlet_endpoint_mode categorical_dual_t \
+  --dirichlet_semantic_t_mode same \
+  --dirichlet_semantic_t_value 0.0 \
+  --categorical_wrong_from_full_vocab \
+  --simplex_bridge_sampler logistic_normal_linear_mean \
+  --logistic_normal_sigma_min 0.18 \
+  --logistic_normal_sigma_max 2.2 \
+  --logistic_normal_tau_min 0.65 \
+  --logistic_normal_tau_max 1.15 \
+  --eps 1e-8 \
+  --infer_steps 128 \
+  --decode_damping 1.0 \
+  --max_gamma 1.0 \
+  --decode_solver flowmap \
+  --noise_init logistic_normal \
+  --bridge_noise_init logistic_normal \
+  --noise_sigma -1 \
+  "${RESUME_ARGS[@]}" \
+  "${COMPILE_ARGS[@]}" \
+  --bf16 2>&1 | tee -a "${LOG_FILE}"

LTA_openwebtext_dualt/scripts/launch_lta_lm1b_compact_gpt2bpe_v8192_len128_fullycoupled_4gpu.sh

@@ -0,0 +1,219 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}"
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"
+export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}"
+export TOKENIZERS_PARALLELISM=false
+export PYTHONUNBUFFERED=1
+export NCCL_DEBUG="${NCCL_DEBUG:-WARN}"
+export TORCH_DISTRIBUTED_TIMEOUT="${TORCH_DISTRIBUTED_TIMEOUT:-3600}"
+
+DATA_PATH="${DATA_PATH:-/e2e-data/evad-tech-vla/wanghan58/data/embedded-language-flows/lm1b-compact-gpt2bpe-v8192-stream128}"
+TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/models/lta_tokenizers/owt_compact_gpt2bpe_v8192/tokenizer.json}"
+
+NNODES="${NNODES:-1}"
+NPROC_PER_NODE="${NPROC_PER_NODE:-4}"
+NODE_RANK="${NODE_RANK:-0}"
+MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+MASTER_PORT="${MASTER_PORT:-32381}"
+
+PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-64}"
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+TOTAL_STEPS="${TOTAL_STEPS:-1000000}"
+WARMUP_STEPS="${WARMUP_STEPS:-2000}"
+NUM_WORKERS="${NUM_WORKERS:-8}"
+DATALOADER_PREFETCH_FACTOR="${DATALOADER_PREFETCH_FACTOR:-4}"
+LOG_EVERY="${LOG_EVERY:-50}"
+SAVE_EVERY="${SAVE_EVERY:-10000}"
+LATEST_EVERY="${LATEST_EVERY:-1000}"
+EVAL_EVERY="${EVAL_EVERY:-0}"
+ALLOW_EXISTING_SAVE_DIR="${ALLOW_EXISTING_SAVE_DIR:-0}"
+ALLOW_TF32="${ALLOW_TF32:-1}"
+DRY_RUN="${DRY_RUN:-0}"
+RESUME_PATH="${RESUME_PATH:-}"
+
+LR="${LR:-6e-4}"
+MIN_LR="${MIN_LR:-6e-5}"
+WEIGHT_DECAY="${WEIGHT_DECAY:-0.1}"
+OUTPUT_WEIGHT_DECAY="${OUTPUT_WEIGHT_DECAY:--1}"
+ADAM_BETA1="${ADAM_BETA1:-0.9}"
+ADAM_BETA2="${ADAM_BETA2:-0.95}"
+ADAM_EPS="${ADAM_EPS:-1e-8}"
+GRAD_CLIP="${GRAD_CLIP:-1.0}"
+EMA_DECAY="${EMA_DECAY:-0.0}"
+EMA_START_STEP="${EMA_START_STEP:-0}"
+
+T_SAMPLING_MODE="${T_SAMPLING_MODE:-uniform}"
+T_SAMPLING_POWER="${T_SAMPLING_POWER:-1.0}"
+T_SAMPLING_EPS="${T_SAMPLING_EPS:-1e-4}"
+T_SAMPLING_LOGIT_MEAN="${T_SAMPLING_LOGIT_MEAN:--1.5}"
+T_SAMPLING_LOGIT_STD="${T_SAMPLING_LOGIT_STD:-0.8}"
+MIN_MASK_RATIO="${MIN_MASK_RATIO:-1.0}"
+MAX_MASK_RATIO="${MAX_MASK_RATIO:-1.0}"
+
+LOSS_T_WEIGHT_MODE="${LOSS_T_WEIGHT_MODE:-none}"
+LOSS_T_MIN_WEIGHT="${LOSS_T_MIN_WEIGHT:-0.0}"
+LOSS_T_DROP_BELOW="${LOSS_T_DROP_BELOW:-0.2}"
+
+RUN_NAME="${RUN_NAME:-lta_lm1b_compact_gpt2bpe_v8192_len128_fullycoupled_rmsnorm_nobias_adamw_wd0p1_${T_SAMPLING_MODE}t_hardce_mask${MIN_MASK_RATIO}-${MAX_MASK_RATIO}_fp32_ddit768x12_gbs${GLOBAL_BATCH_SIZE}_4gpu_1m_$(date +%Y%m%d_%H%M%S)}"
+SAVE_DIR="${SAVE_DIR:-runs/${RUN_NAME}}"
+LOG_DIR="${LOG_DIR:-logs/lm1b_compact_gpt2bpe_v8192_len128_fullycoupled_4gpu}"
+LOG_FILE="${LOG_FILE:-${LOG_DIR}/${RUN_NAME}.log}"
+
+if [[ ! -f "${TOKENIZER_PATH}" ]]; then
+  echo "Missing tokenizer: ${TOKENIZER_PATH}" >&2
+  exit 2
+fi
+if [[ ! -d "${DATA_PATH}" ]]; then
+  echo "Missing tokenized dataset: ${DATA_PATH}" >&2
+  echo "Build it with: bash scripts/build_lta_lm1b_compact_gpt2bpe_v8192_stream128_np8.sh" >&2
+  exit 2
+fi
+if [[ -f "${SAVE_DIR}/args.json" && "${ALLOW_EXISTING_SAVE_DIR}" != "1" ]]; then
+  echo "Refusing to start because SAVE_DIR already contains args.json: ${SAVE_DIR}" >&2
+  exit 2
+fi
+
+mkdir -p "${LOG_DIR}" "${SAVE_DIR}"
+
+TF32_FLAG="--allow_tf32"
+TF32_LABEL="true"
+if [[ "${ALLOW_TF32}" == "0" || "${ALLOW_TF32}" == "false" || "${ALLOW_TF32}" == "False" ]]; then
+  TF32_FLAG="--no-allow_tf32"
+  TF32_LABEL="false"
+fi
+
+NUM_EXAMPLES=$(python - <<PY
+import json
+from pathlib import Path
+from datasets import Sequence, load_from_disk
+from datasets.features import features as hf_features
+hf_features._FEATURE_TYPES.setdefault("List", Sequence)
+root = Path("${DATA_PATH}")
+meta = root / "elf_multi_part_meta.json"
+if meta.exists():
+    print(int(json.loads(meta.read_text()).get("num_examples", 0)))
+else:
+    parts = root / "parts"
+    if parts.is_dir():
+        print(sum(len(load_from_disk(str(p))) for p in sorted(parts.iterdir()) if p.is_dir()))
+    else:
+        print(len(load_from_disk(str(root))))
+PY
+)
+
+WORLD_SIZE=$(( NNODES * NPROC_PER_NODE ))
+GRAD_ACCUM=$(( (GLOBAL_BATCH_SIZE + PER_GPU_BATCH_SIZE * WORLD_SIZE - 1) / (PER_GPU_BATCH_SIZE * WORLD_SIZE) ))
+
+echo "[launch-lm1b-v8192] run_name=${RUN_NAME}"
+echo "[launch-lm1b-v8192] save_dir=${SAVE_DIR}"
+echo "[launch-lm1b-v8192] log_file=${LOG_FILE}"
+echo "[launch-lm1b-v8192] data_path=${DATA_PATH}"
+echo "[launch-lm1b-v8192] tokenizer=${TOKENIZER_PATH}"
+echo "[launch-lm1b-v8192] examples=${NUM_EXAMPLES} max_len=128 total_steps=${TOTAL_STEPS} world_size=${WORLD_SIZE} grad_accum=${GRAD_ACCUM} save_every=${SAVE_EVERY}"
+echo "[launch-lm1b-v8192] model=ddit d768 x12 h12 ff3072 rmsnorm no_bias dropout0"
+echo "[launch-lm1b-v8192] recipe=dirichlet categorical_dual_t hard_ce target_prob1 mask=${MIN_MASK_RATIO}->${MAX_MASK_RATIO} t=${T_SAMPLING_MODE} tf32=${TF32_LABEL}"
+
+if [[ "${DRY_RUN}" == "1" || "${DRY_RUN}" == "true" || "${DRY_RUN}" == "True" ]]; then
+  echo "[launch-lm1b-v8192] DRY_RUN=1, validated setup; skipping torchrun."
+  exit 0
+fi
+
+RESUME_ARGS=()
+if [[ -n "${RESUME_PATH}" ]]; then
+  RESUME_ARGS+=(--resume_path "${RESUME_PATH}")
+fi
+
+python -m torch.distributed.run \
+  --nnodes="${NNODES}" \
+  --nproc_per_node="${NPROC_PER_NODE}" \
+  --node_rank="${NODE_RANK}" \
+  --master_addr="${MASTER_ADDR}" \
+  --master_port="${MASTER_PORT}" \
+  train.py \
+  "${RESUME_ARGS[@]}" \
+  --data_path "${DATA_PATH}" \
+  --tokenized_hf \
+  --tokenized_pad_token pad \
+  --tokenizer_path "${TOKENIZER_PATH}" \
+  --save_dir "${SAVE_DIR}" \
+  --max_len 128 \
+  --batch_size "${PER_GPU_BATCH_SIZE}" \
+  --global_batch_size "${GLOBAL_BATCH_SIZE}" \
+  --num_workers "${NUM_WORKERS}" \
+  --dataloader_prefetch_factor "${DATALOADER_PREFETCH_FACTOR}" \
+  --epochs 0 \
+  --total_steps "${TOTAL_STEPS}" \
+  --warmup_steps "${WARMUP_STEPS}" \
+  --log_every "${LOG_EVERY}" \
+  --eval_every "${EVAL_EVERY}" \
+  --save_every "${SAVE_EVERY}" \
+  --latest_every "${LATEST_EVERY}" \
+  --optimizer adamw \
+  --lr "${LR}" \
+  --lr_schedule cosine \
+  --min_lr "${MIN_LR}" \
+  --weight_decay "${WEIGHT_DECAY}" \
+  --output_weight_decay "${OUTPUT_WEIGHT_DECAY}" \
+  --adamw_param_groups nanogpt \
+  --adam_beta1 "${ADAM_BETA1}" \
+  --adam_beta2 "${ADAM_BETA2}" \
+  --adam_eps "${ADAM_EPS}" \
+  --ema_decay "${EMA_DECAY}" \
+  --ema_start_step "${EMA_START_STEP}" \
+  --grad_clip "${GRAD_CLIP}" \
+  --seed 123 \
+  --d_model 768 \
+  --cond_dim 128 \
+  --n_layers 12 \
+  --n_heads 12 \
+  --dim_ff 3072 \
+  --dropout 0.0 \
+  --no-output_bias \
+  --norm_type rmsnorm \
+  --model_type ddit \
+  --state_format prob \
+  --bridge dirichlet \
+  --target_loss hard_ce \
+  --loss_t_weight_mode "${LOSS_T_WEIGHT_MODE}" \
+  --loss_t_min_weight "${LOSS_T_MIN_WEIGHT}" \
+  --loss_t_drop_below "${LOSS_T_DROP_BELOW}" \
+  --target_prob 1.0 \
+  --min_t 0.0 \
+  --max_t 1.0 \
+  --t_sampling_mode "${T_SAMPLING_MODE}" \
+  --t_sampling_power "${T_SAMPLING_POWER}" \
+  --t_sampling_logit_mean "${T_SAMPLING_LOGIT_MEAN}" \
+  --t_sampling_logit_std "${T_SAMPLING_LOGIT_STD}" \
+  --t_sampling_eps "${T_SAMPLING_EPS}" \
+  --dual_t \
+  --corrupt_t_mode same \
+  --corrupt_min_t 0.0 \
+  --corrupt_max_t 1.0 \
+  --min_mask_ratio "${MIN_MASK_RATIO}" \
+  --max_mask_ratio "${MAX_MASK_RATIO}" \
+  --wrong_token_replace_prob 1.0 \
+  --wrong_token_schedule linear_t \
+  --wrong_token_exp_k 1.0 \
+  --dirichlet_concentration_min 1.0 \
+  --dirichlet_concentration_max 1024 \
+  --dirichlet_endpoint_mode categorical_dual_t \
+  --dirichlet_semantic_t_mode same \
+  --dirichlet_semantic_t_value 0.0 \
+  --categorical_wrong_from_full_vocab \
+  --simplex_bridge_sampler dirichlet \
+  --eps 1e-8 \
+  --infer_steps 128 \
+  --decode_damping 1.0 \
+  --max_gamma 1.0 \
+  --decode_solver flowmap \
+  --noise_init logistic_normal \
+  --bridge_noise_init logistic_normal \
+  --noise_sigma -1 \
+  "${TF32_FLAG}" \
+  --ddp_gradient_as_bucket_view \
+  2>&1 | tee -a "${LOG_FILE}"
+

LTA_openwebtext_dualt/scripts/launch_lta_openwebtext_dualt_8gpu_aligned.sh

@@ -0,0 +1,113 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}"
+export TOKENIZERS_PARALLELISM=false
+export PYTHONUNBUFFERED=1
+export NCCL_DEBUG="${NCCL_DEBUG:-WARN}"
+export TORCH_DISTRIBUTED_TIMEOUT="${TORCH_DISTRIBUTED_TIMEOUT:-3600}"
+
+# FLM-repo-aligned engineering defaults:
+# - GPT-2 OpenWebText split, wrapped fixed length 1024
+# - DDiT small: 12 layers, d=768, 12 heads, FF=3072, adaLN-zero
+# - global batch 512, per-GPU local batch 32 on 8 GPUs
+# - bf16, flash-attn when available, torch.compile max-autotune
+# - AdamW lr 3e-4, no weight decay, 2500 warmup, constant schedule
+#
+# Algorithm remains ours:
+# - probability-state LTA
+# - Dirichlet bridge
+# - hard CE target and one-hot clean anchors/endpoints
+# - wrong-token corruption is scheduled by corruption time: p_wrong = 1 - t_corrupt
+# - dual-t: model/flow time is separate from corruption/support time
+
+RUN_NAME="${RUN_NAME:-lta_owt_dirichlet_dualt_onehot_hardce_ddit_small_len1024_gbs512_8gpu}"
+DATA_PATH="${DATA_PATH:-/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext}"
+TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-standard/tokenizer.json}"
+SAVE_DIR="${SAVE_DIR:-runs/${RUN_NAME}}"
+LOG_FILE="${LOG_FILE:-logs/${RUN_NAME}.log}"
+
+NNODES="${NNODES:-1}"
+NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+NODE_RANK="${NODE_RANK:-0}"
+MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+MASTER_PORT="${MASTER_PORT:-29621}"
+
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-32}"
+TOTAL_STEPS="${TOTAL_STEPS:-1500000}"
+WARMUP_STEPS="${WARMUP_STEPS:-2500}"
+MAX_LEN="${MAX_LEN:-1024}"
+NUM_WORKERS="${NUM_WORKERS:-4}"
+LOG_EVERY="${LOG_EVERY:-100}"
+SAVE_EVERY="${SAVE_EVERY:-20000}"
+EVAL_EVERY="${EVAL_EVERY:-0}"
+
+mkdir -p logs runs "${SAVE_DIR}"
+
+python -m torch.distributed.run \
+  --nnodes="${NNODES}" \
+  --nproc_per_node="${NPROC_PER_NODE}" \
+  --node_rank="${NODE_RANK}" \
+  --master_addr="${MASTER_ADDR}" \
+  --master_port="${MASTER_PORT}" \
+  train.py \
+  --data_path "${DATA_PATH}" \
+  --text_column text \
+  --openwebtext_split train_minus_100k \
+  --tokenizer_path "${TOKENIZER_PATH}" \
+  --save_dir "${SAVE_DIR}" \
+  --wrap \
+  --max_len "${MAX_LEN}" \
+  --batch_size "${PER_GPU_BATCH_SIZE}" \
+  --num_workers "${NUM_WORKERS}" \
+  --global_batch_size "${GLOBAL_BATCH_SIZE}" \
+  --total_steps "${TOTAL_STEPS}" \
+  --log_every "${LOG_EVERY}" \
+  --eval_every "${EVAL_EVERY}" \
+  --save_every "${SAVE_EVERY}" \
+  --lr 3e-4 \
+  --weight_decay 0 \
+  --adam_beta1 0.9 \
+  --adam_beta2 0.999 \
+  --adam_eps 1e-8 \
+  --warmup_steps "${WARMUP_STEPS}" \
+  --lr_schedule constant_warmup \
+  --grad_clip 1.0 \
+  --seed 123 \
+  --d_model 768 \
+  --cond_dim 128 \
+  --n_layers 12 \
+  --n_heads 12 \
+  --dim_ff 3072 \
+  --dropout 0.1 \
+  --model_type ddit \
+  --state_format prob \
+  --bridge dirichlet \
+  --target_loss hard_ce \
+  --target_prob 1.0 \
+  --min_t 0.0 \
+  --max_t 1.0 \
+  --dual_t \
+  --corrupt_t_mode independent \
+  --corrupt_min_t 0.0 \
+  --corrupt_max_t 1.0 \
+  --min_mask_ratio 0.1 \
+  --max_mask_ratio 1.0 \
+  --wrong_token_replace_prob 1.0 \
+  --wrong_token_schedule linear_t \
+  --wrong_token_exp_k 1.0 \
+  --dirichlet_concentration_min 1.0 \
+  --dirichlet_concentration_max 1024.0 \
+  --eps 1e-8 \
+  --infer_steps 1024 \
+  --decode_damping 1.0 \
+  --max_gamma 1.0 \
+  --decode_solver flowmap \
+  --noise_init logistic_normal \
+  --bridge_noise_init logistic_normal \
+  --noise_sigma -1 \
+  --torch_compile \
+  --compile_mode max-autotune \
+  --bf16 2>&1 | tee -a "${LOG_FILE}"

LTA_openwebtext_dualt/scripts/launch_lta_owt_compact_gpt2bpe_v2048_elfaligned_logitnormal_tokenized_8gpu.sh

@@ -0,0 +1,34 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+
+export VOCAB_SIZE=2048
+export DATA_PATH="${DATA_PATH:-/e2e-data/evad-tech-vla/wanghan58/data/embedded-language-flows/openwebtext-compact-gpt2bpe-v2048-stream1024-train-minus-100k}"
+export TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/models/lta_tokenizers/owt_compact_gpt2bpe_v2048/tokenizer.json}"
+
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
+export NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+export MASTER_PORT="${MASTER_PORT:-32241}"
+
+export GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+export PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-32}"
+export EPOCHS="${EPOCHS:-5}"
+export NUM_WORKERS="${NUM_WORKERS:-8}"
+export LOG_EVERY="${LOG_EVERY:-100}"
+export LATEST_EVERY="${LATEST_EVERY:-1000}"
+export EVAL_EVERY="${EVAL_EVERY:-0}"
+
+# Same recipe as the current OWT ELF-aligned runs: fp32 params/activations,
+# TF32 matmul allowed, RMSNorm, no output bias, Muon+Adam fallback wd=0.1.
+export ALLOW_TF32="${ALLOW_TF32:-1}"
+export TARGET_LOSS="${TARGET_LOSS:-hard_ce}"
+export T_LOGIT_MEAN="${T_LOGIT_MEAN:--1.5}"
+export T_LOGIT_STD="${T_LOGIT_STD:-0.8}"
+export WEIGHT_DECAY="${WEIGHT_DECAY:-0.1}"
+export OUTPUT_INIT_STD="${OUTPUT_INIT_STD:-0.0}"
+
+export LOG_DIR="${LOG_DIR:-logs/compact_gpt2bpe_v2048_stream1024_8gpu}"
+export RUN_NAME="${RUN_NAME:-lta_owt_compact_gpt2bpe_v2048_stream1024_elfaligned_dditelf_muon_logitnormal_m1p5_s0p8_hardce_wd0p1_gbs512_8gpu_5epoch_$(date +%Y%m%d_%H%M%S)}"
+
+exec bash scripts/launch_lta_owt_compact_gpt2bpe_elfaligned_logitnormal_tokenized_8gpu.sh "$@"

LTA_openwebtext_dualt/scripts/launch_lta_owt_fullycoupled_logitnormal_mid_mask1_wd0p1_fp32_8gpu.sh

@@ -0,0 +1,204 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"
+export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}"
+export TOKENIZERS_PARALLELISM=false
+export PYTHONUNBUFFERED=1
+export NCCL_DEBUG="${NCCL_DEBUG:-WARN}"
+export TORCH_DISTRIBUTED_TIMEOUT="${TORCH_DISTRIBUTED_TIMEOUT:-3600}"
+
+# Fully-coupled OWT baseline with FP32 params/activations and nanoGPT-style AdamW decay:
+#   decay:    matrix / embedding params (p.dim() >= 2)
+#   no decay: bias / norm / 1D params
+# This intentionally does NOT pass --bf16. TF32 is enabled by default so H200 uses
+# Tensor Cores; set ALLOW_TF32=0 for strict FP32 debugging.
+
+T_SAMPLING_MODE="${T_SAMPLING_MODE:-logit_normal}"
+T_SAMPLING_POWER="${T_SAMPLING_POWER:-1.0}"
+T_SAMPLING_EPS="${T_SAMPLING_EPS:-1e-4}"
+T_SAMPLING_LOGIT_MEAN="${T_SAMPLING_LOGIT_MEAN:--0.22}"
+T_SAMPLING_LOGIT_STD="${T_SAMPLING_LOGIT_STD:-0.5}"
+MIN_MASK_RATIO="${MIN_MASK_RATIO:-1.0}"
+MAX_MASK_RATIO="${MAX_MASK_RATIO:-1.0}"
+DDIT_MLP_TYPE="${DDIT_MLP_TYPE:-swiglu}"
+
+sanitize_label() {
+  printf "%s" "$1" | sed -e 's/-/m/g' -e 's/\./p/g'
+}
+
+T_LOGIT_MEAN_LABEL="$(sanitize_label "${T_SAMPLING_LOGIT_MEAN}")"
+T_LOGIT_STD_LABEL="$(sanitize_label "${T_SAMPLING_LOGIT_STD}")"
+MIN_MASK_RATIO_LABEL="$(sanitize_label "${MIN_MASK_RATIO}")"
+MAX_MASK_RATIO_LABEL="$(sanitize_label "${MAX_MASK_RATIO}")"
+
+RUN_NAME="${RUN_NAME:-lta_owt_gpt2cached_len1024_fullycoupled_rmsnorm_nobias_${DDIT_MLP_TYPE}_adamw_wd0p1_logitnormal_mid_${T_LOGIT_MEAN_LABEL}_s${T_LOGIT_STD_LABEL}_hardce_mask${MIN_MASK_RATIO_LABEL}-${MAX_MASK_RATIO_LABEL}_nanogpt_fp32_ddit768x12_gbs512_8gpu_1m_$(date +%Y%m%d_%H%M%S)}"
+SAVE_DIR="${SAVE_DIR:-runs/${RUN_NAME}}"
+LOG_DIR="${LOG_DIR:-logs/fullycoupled_logitnormal_mid_mask1_wd0p1_fp32_8gpu}"
+LOG_FILE="${LOG_FILE:-${LOG_DIR}/${RUN_NAME}.log}"
+
+DATA_PATH="${DATA_PATH:-/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext}"
+OWT_CACHE="${OWT_CACHE:-/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext_lta_cached_chunks/gpt2_len1024_train_minus_100k}"
+TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-standard/tokenizer.json}"
+
+NNODES="${NNODES:-1}"
+NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+NODE_RANK="${NODE_RANK:-0}"
+MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+MASTER_PORT="${MASTER_PORT:-31997}"
+
+PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-32}"
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+TOTAL_STEPS="${TOTAL_STEPS:-1000000}"
+WARMUP_STEPS="${WARMUP_STEPS:-2000}"
+NUM_WORKERS="${NUM_WORKERS:-8}"
+DATALOADER_PREFETCH_FACTOR="${DATALOADER_PREFETCH_FACTOR:-4}"
+LOG_EVERY="${LOG_EVERY:-50}"
+SAVE_EVERY="${SAVE_EVERY:-50000}"
+LATEST_EVERY="${LATEST_EVERY:-1000}"
+EVAL_EVERY="${EVAL_EVERY:-0}"
+ALLOW_EXISTING_SAVE_DIR="${ALLOW_EXISTING_SAVE_DIR:-0}"
+ALLOW_TF32="${ALLOW_TF32:-1}"
+DRY_RUN="${DRY_RUN:-0}"
+
+LR="${LR:-6e-4}"
+MIN_LR="${MIN_LR:-6e-5}"
+WEIGHT_DECAY="${WEIGHT_DECAY:-0.1}"
+OUTPUT_WEIGHT_DECAY="${OUTPUT_WEIGHT_DECAY:--1}"
+ADAM_BETA1="${ADAM_BETA1:-0.9}"
+ADAM_BETA2="${ADAM_BETA2:-0.95}"
+ADAM_EPS="${ADAM_EPS:-1e-8}"
+GRAD_CLIP="${GRAD_CLIP:-1.0}"
+EMA_DECAY="${EMA_DECAY:-0.0}"
+EMA_START_STEP="${EMA_START_STEP:-0}"
+LOSS_T_WEIGHT_MODE="${LOSS_T_WEIGHT_MODE:-none}"
+LOSS_T_MIN_WEIGHT="${LOSS_T_MIN_WEIGHT:-0.0}"
+LOSS_T_DROP_BELOW="${LOSS_T_DROP_BELOW:-0.2}"
+
+if [[ -f "${SAVE_DIR}/args.json" && "${ALLOW_EXISTING_SAVE_DIR}" != "1" ]]; then
+  echo "Refusing to start because SAVE_DIR already contains args.json: ${SAVE_DIR}" >&2
+  echo "Use a new RUN_NAME/SAVE_DIR or set ALLOW_EXISTING_SAVE_DIR=1 intentionally." >&2
+  exit 2
+fi
+
+mkdir -p "${LOG_DIR}" "${SAVE_DIR}"
+
+TF32_FLAG="--allow_tf32"
+TF32_LABEL="true"
+if [[ "${ALLOW_TF32}" == "0" || "${ALLOW_TF32}" == "false" || "${ALLOW_TF32}" == "False" ]]; then
+  TF32_FLAG="--no-allow_tf32"
+  TF32_LABEL="false"
+fi
+
+echo "[launch] method=owt_fullycoupled_adamw_wd0p1_nanogpt_fp32 host=$(hostname) time=$(date -Iseconds)"
+echo "[launch] run_name=${RUN_NAME}"
+echo "[launch] save_dir=${SAVE_DIR}"
+echo "[launch] log_file=${LOG_FILE}"
+echo "[launch] data_path=${DATA_PATH}"
+echo "[launch] owt_cache=${OWT_CACHE}"
+echo "[launch] optimizer=adamw lr=${LR} min_lr=${MIN_LR} wd=${WEIGHT_DECAY} output_wd=${OUTPUT_WEIGHT_DECAY} param_groups=nanogpt ema=${EMA_DECAY}"
+echo "[launch] fp32=true bf16=false tf32=${TF32_LABEL} norm_type=rmsnorm output_bias=false ddit_mlp_type=${DDIT_MLP_TYPE} batch=${GLOBAL_BATCH_SIZE} per_gpu=${PER_GPU_BATCH_SIZE}"
+echo "[launch] loss_t_weight_mode=${LOSS_T_WEIGHT_MODE} loss_t_min_weight=${LOSS_T_MIN_WEIGHT} loss_t_drop_below=${LOSS_T_DROP_BELOW}"
+echo "[launch] target_loss=hard_ce t_sampling_mode=${T_SAMPLING_MODE} t_sampling_logit_mean=${T_SAMPLING_LOGIT_MEAN} t_sampling_logit_std=${T_SAMPLING_LOGIT_STD} t_sampling_power=${T_SAMPLING_POWER} t_sampling_eps=${T_SAMPLING_EPS} mask_ratio=${MIN_MASK_RATIO}->${MAX_MASK_RATIO}"
+
+if [[ "${DRY_RUN}" == "1" || "${DRY_RUN}" == "true" || "${DRY_RUN}" == "True" ]]; then
+  echo "[launch] DRY_RUN=1, validated launch parameters; skipping torchrun."
+  exit 0
+fi
+
+python -m torch.distributed.run \
+  --nnodes="${NNODES}" \
+  --nproc_per_node="${NPROC_PER_NODE}" \
+  --node_rank="${NODE_RANK}" \
+  --master_addr="${MASTER_ADDR}" \
+  --master_port="${MASTER_PORT}" \
+  train.py \
+  --data_path "${DATA_PATH}" \
+  --openwebtext_split train_minus_100k \
+  --text_column text \
+  --detokenizer auto \
+  --tokenizer_path "${TOKENIZER_PATH}" \
+  --save_dir "${SAVE_DIR}" \
+  --wrap \
+  --wrap_mode stream \
+  --owt_cached_chunks \
+  --owt_chunk_cache_dir "${OWT_CACHE}" \
+  --max_len 1024 \
+  --batch_size "${PER_GPU_BATCH_SIZE}" \
+  --global_batch_size "${GLOBAL_BATCH_SIZE}" \
+  --num_workers "${NUM_WORKERS}" \
+  --dataloader_prefetch_factor "${DATALOADER_PREFETCH_FACTOR}" \
+  --total_steps "${TOTAL_STEPS}" \
+  --warmup_steps "${WARMUP_STEPS}" \
+  --log_every "${LOG_EVERY}" \
+  --eval_every "${EVAL_EVERY}" \
+  --save_every "${SAVE_EVERY}" \
+  --latest_every "${LATEST_EVERY}" \
+  --optimizer adamw \
+  --lr "${LR}" \
+  --lr_schedule cosine \
+  --min_lr "${MIN_LR}" \
+  --weight_decay "${WEIGHT_DECAY}" \
+  --output_weight_decay "${OUTPUT_WEIGHT_DECAY}" \
+  --adamw_param_groups nanogpt \
+  --adam_beta1 "${ADAM_BETA1}" \
+  --adam_beta2 "${ADAM_BETA2}" \
+  --adam_eps "${ADAM_EPS}" \
+  --ema_decay "${EMA_DECAY}" \
+  --ema_start_step "${EMA_START_STEP}" \
+  --grad_clip "${GRAD_CLIP}" \
+  --seed 123 \
+  --d_model 768 \
+  --cond_dim 128 \
+  --n_layers 12 \
+  --n_heads 12 \
+  --dim_ff 3072 \
+  --dropout 0.0 \
+  --no-output_bias \
+  --norm_type rmsnorm \
+  --model_type ddit \
+  --ddit_mlp_type "${DDIT_MLP_TYPE}" \
+  --state_format prob \
+  --bridge dirichlet \
+  --target_loss hard_ce \
+  --loss_t_weight_mode "${LOSS_T_WEIGHT_MODE}" \
+  --loss_t_min_weight "${LOSS_T_MIN_WEIGHT}" \
+  --loss_t_drop_below "${LOSS_T_DROP_BELOW}" \
+  --target_prob 1.0 \
+  --min_t 0.0 \
+  --max_t 1.0 \
+  --t_sampling_mode "${T_SAMPLING_MODE}" \
+  --t_sampling_power "${T_SAMPLING_POWER}" \
+  --t_sampling_logit_mean "${T_SAMPLING_LOGIT_MEAN}" \
+  --t_sampling_logit_std "${T_SAMPLING_LOGIT_STD}" \
+  --t_sampling_eps "${T_SAMPLING_EPS}" \
+  --dual_t \
+  --corrupt_t_mode same \
+  --corrupt_min_t 0.0 \
+  --corrupt_max_t 1.0 \
+  --min_mask_ratio "${MIN_MASK_RATIO}" \
+  --max_mask_ratio "${MAX_MASK_RATIO}" \
+  --wrong_token_replace_prob 1.0 \
+  --wrong_token_schedule linear_t \
+  --wrong_token_exp_k 1.0 \
+  --dirichlet_concentration_min 1.0 \
+  --dirichlet_concentration_max 1024 \
+  --dirichlet_endpoint_mode categorical_dual_t \
+  --dirichlet_semantic_t_mode same \
+  --dirichlet_semantic_t_value 0.0 \
+  --categorical_wrong_from_full_vocab \
+  --simplex_bridge_sampler dirichlet \
+  --eps 1e-8 \
+  --infer_steps 1024 \
+  --decode_damping 1.0 \
+  --max_gamma 1.0 \
+  --decode_solver flowmap \
+  --noise_init logistic_normal \
+  --bridge_noise_init logistic_normal \
+  --noise_sigma -1 \
+  "${TF32_FLAG}" \
+  --ddp_gradient_as_bucket_view \
+  2>&1 | tee -a "${LOG_FILE}"

LTA_openwebtext_dualt/scripts/launch_lta_owt_t5_adaln_adamw_wd0p1_rollin_grad_k1_rho025_8gpu.sh

@@ -0,0 +1,168 @@
+#!/bin/bash
+set -e
+set -x
+set -o pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+
+# Use the repo's normal Python/torchrun by default.  Set ACTIVATE_ENV only if
+# that environment is known to include torch, datasets, and tokenizers.
+if [[ -n "${ACTIVATE_ENV:-}" ]]; then
+  source "${ACTIVATE_ENV}"
+fi
+
+free_port() {
+  python3 - <<'PY'
+import socket
+with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+    s.bind(("127.0.0.1", 0))
+    print(s.getsockname()[1])
+PY
+}
+
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
+export NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}"
+export TOKENIZERS_PARALLELISM=false
+export PYTHONUNBUFFERED=1
+export NCCL_DEBUG="${NCCL_DEBUG:-WARN}"
+export TORCH_NCCL_AVOID_RECORD_STREAMS=1
+export TORCH_DISTRIBUTED_TIMEOUT="${TORCH_DISTRIBUTED_TIMEOUT:-3600}"
+
+DATA_PATH="${DATA_PATH:-/e2e-data/evad-tech-vla/wanghan58/data/embedded-language-flows/openwebtext-t5}"
+TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/models/hf/t5-small/tokenizer.json}"
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-32}"
+TOTAL_STEPS="${TOTAL_STEPS:-1000000}"
+MASTER_PORT="${MASTER_PORT:-$(free_port)}"
+LOG_DIR="${LOG_DIR:-logs/elfaligned_t5tokenized_8gpu}"
+
+RUN_NAME="${RUN_NAME:-lta_owt_t5_adaln_adamw_wd0p1_rollin_grad_p50_k1_rho0_0p25_uniformt_temp1_synct_gbs${GLOBAL_BATCH_SIZE}_8gpu_1m_$(date +%Y%m%d_%H%M%S)}"
+LOG_FILE="${LOG_FILE:-${LOG_DIR}/${RUN_NAME}.log}"
+
+mkdir -p "${LOG_DIR}" "runs/${RUN_NAME}"
+
+python - <<'PY'
+import sys
+import torch
+import tokenizers
+
+print(f"[launch] python={sys.executable}")
+print(f"[launch] torch={torch.__version__} tokenizers={tokenizers.__version__}")
+PY
+
+echo "[launch] run_name=${RUN_NAME}" | tee -a "${LOG_FILE}"
+echo "[launch] cuda=${CUDA_VISIBLE_DEVICES} nproc_per_node=${NPROC_PER_NODE} master_port=${MASTER_PORT}" | tee -a "${LOG_FILE}"
+echo "[launch] global_batch_size=${GLOBAL_BATCH_SIZE} per_gpu_batch_size=${PER_GPU_BATCH_SIZE} total_steps=${TOTAL_STEPS}" | tee -a "${LOG_FILE}"
+echo "[launch] rollout=keep_grad,p50,K1,rhoU0_0p25,sync_t" | tee -a "${LOG_FILE}"
+echo "[launch] data_path=${DATA_PATH}" | tee -a "${LOG_FILE}"
+echo "[launch] tokenizer=${TOKENIZER_PATH}" | tee -a "${LOG_FILE}"
+echo "[launch] log_file=${LOG_FILE}" | tee -a "${LOG_FILE}"
+
+torchrun \
+  --nproc_per_node="${NPROC_PER_NODE}" \
+  --master_port="${MASTER_PORT}" \
+  train.py \
+  --data_path "${DATA_PATH}" \
+  --tokenized_hf \
+  --tokenized_pad_token pad \
+  --tokenizer_path "${TOKENIZER_PATH}" \
+  --save_dir "runs/${RUN_NAME}" \
+  --max_len 1024 \
+  --batch_size "${PER_GPU_BATCH_SIZE}" \
+  --global_batch_size "${GLOBAL_BATCH_SIZE}" \
+  --num_workers 8 \
+  --dataloader_prefetch_factor 4 \
+  --epochs 0 \
+  --total_steps "${TOTAL_STEPS}" \
+  --warmup_steps 1000 \
+  --log_every 100 \
+  --eval_every 0 \
+  --save_every 10000 \
+  --latest_every 1000 \
+  --optimizer adamw \
+  --lr 6e-4 \
+  --lr_schedule cosine \
+  --min_lr 6e-5 \
+  --weight_decay 0.1 \
+  --output_weight_decay -1 \
+  --adamw_param_groups nanogpt \
+  --adam_beta1 0.9 \
+  --adam_beta2 0.999 \
+  --adam_eps 1e-8 \
+  --ema_decay 0.9999 \
+  --ema_start_step 0 \
+  --grad_clip 1.0 \
+  --seed 42 \
+  --d_model 768 \
+  --cond_dim 128 \
+  --n_layers 12 \
+  --n_heads 12 \
+  --dim_ff 3072 \
+  --dropout 0.0 \
+  --no-output_bias \
+  --norm_type rmsnorm \
+  --model_type ddit \
+  --ddit_mlp_type swiglu \
+  --state_format prob \
+  --bridge dirichlet \
+  --target_loss hard_ce \
+  --loss_t_weight_mode none \
+  --loss_t_min_weight 0.0 \
+  --rollout_train_prob 0.50 \
+  --rollout_train_time_mode sampled_path \
+  --rollout_train_steps 1 \
+  --rollout_train_steps_min 1 \
+  --rollout_train_infer_steps 1 \
+  --rollout_train_s_dist uniform \
+  --rollout_train_s_min_frac 0.0 \
+  --rollout_train_s_max_frac 0.25 \
+  --rollout_train_temp 1.0 \
+  --rollout_train_max_gamma 1.0 \
+  --rollout_train_corrupt_only \
+  --rollout_train_samplewise \
+  --rollout_train_selected_only \
+  --no-rollout_train_compute_always \
+  --rollout_train_keep_grad \
+  --rollout_train_sync_t \
+  --target_prob 1.0 \
+  --min_t 0.0 \
+  --max_t 1.0 \
+  --t_sampling_mode uniform \
+  --t_sampling_logit_mean -1.5 \
+  --t_sampling_logit_std 0.8 \
+  --t_sampling_eps 1e-4 \
+  --dual_t \
+  --corrupt_t_mode same \
+  --corrupt_min_t 0.0 \
+  --corrupt_max_t 1.0 \
+  --min_mask_ratio 1.0 \
+  --max_mask_ratio 1.0 \
+  --mask_mixture_original_prob 0.0 \
+  --mask_mixture_lowk_prob 0.0 \
+  --mask_mixture_lowcorrupt_prob 0.0 \
+  --mask_mixture_block_prob 0.0 \
+  --mask_mixture_all_prob 1.0 \
+  --wrong_token_replace_prob 1.0 \
+  --wrong_token_schedule linear_t \
+  --wrong_token_exp_k 1.0 \
+  --dirichlet_concentration_min 1.0 \
+  --dirichlet_concentration_max 1024 \
+  --dirichlet_endpoint_mode categorical_dual_t \
+  --dirichlet_semantic_t_mode same \
+  --dirichlet_semantic_t_value 0.0 \
+  --categorical_wrong_from_full_vocab \
+  --simplex_bridge_sampler dirichlet \
+  --eps 1e-8 \
+  --infer_steps 1024 \
+  --decode_damping 1.0 \
+  --max_gamma 1.0 \
+  --decode_solver flowmap \
+  --noise_init logistic_normal \
+  --bridge_noise_init logistic_normal \
+  --noise_sigma -1 \
+  --allow_tf32 \
+  --activation_checkpointing \
+  --activation_checkpoint_scope mlp \
+  --ddp_gradient_as_bucket_view \
+  2>&1 | tee -a "${LOG_FILE}"

LTA_openwebtext_dualt/scripts/make_duo_integral_cache.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import pickle
+import time
+from pathlib import Path
+
+import numpy as np
+from scipy.integrate import quad
+from scipy.stats import norm
+
+
+def _discrete_prob_map(gamma_t: float, vocab_size: int):
+    snr_sqrt = np.exp(-gamma_t / 2)
+
+    def value(x):
+        cdf = norm.cdf(x, scale=1) ** (vocab_size - 1)
+        pdf = norm.pdf(x, loc=snr_sqrt, scale=1)
+        return pdf * cdf
+
+    return value
+
+
+def _discrete_prob_grad(gamma_t: float, vocab_size: int):
+    snr_sqrt = np.exp(-gamma_t / 2)
+
+    def value(x):
+        coef = -0.5 * snr_sqrt * (x - snr_sqrt)
+        cdf = norm.cdf(x, scale=1) ** (vocab_size - 1)
+        pdf = norm.pdf(x, loc=snr_sqrt, scale=1)
+        return coef * pdf * cdf
+
+    return value
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--vocab_size", type=int, required=True)
+    p.add_argument("--log10_num_points", type=int, default=5)
+    p.add_argument("--output", required=True)
+    p.add_argument("--gamma_min", type=float, default=-5.0)
+    p.add_argument("--gamma_max", type=float, default=-1.0)
+    args = p.parse_args()
+
+    output = Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    num_points = 10 ** int(args.log10_num_points)
+    gammas = np.linspace(args.gamma_min, args.gamma_max, num_points)
+    pt = []
+    grad_pt = []
+    start = time.time()
+    for i, gamma in enumerate(gammas, start=1):
+        val, _ = quad(_discrete_prob_map(float(gamma), args.vocab_size), -np.inf, np.inf)
+        grad, _ = quad(_discrete_prob_grad(float(gamma), args.vocab_size), -np.inf, np.inf)
+        pt.append(val)
+        grad_pt.append(grad)
+        if i % 100 == 0 or i == num_points:
+            print(f"{100 * i / num_points:.1f}% completed elapsed={(time.time() - start) / 60:.2f}m", flush=True)
+
+    payload = {
+        "vocab_size": args.vocab_size,
+        "gamma_min": args.gamma_min,
+        "gamma_max": args.gamma_max,
+        "num_points": num_points,
+        "pt": np.asarray(pt),
+        "grad_pt": np.asarray(grad_pt),
+    }
+    tmp = output.with_suffix(output.suffix + ".tmp")
+    with tmp.open("wb") as f:
+        pickle.dump(payload, f)
+    tmp.replace(output)
+    print(f"wrote {output}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

LTA_openwebtext_dualt/scripts/prepare_elf_wmt14_deen_t5.sh

@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}"
+
+DATA_ROOT="${DATA_ROOT:-/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/elf}"
+CACHE_DIR="${CACHE_DIR:-/e2e-data/evad-tech-vla/wanghan58/data/hf_cache}"
+mkdir -p "${DATA_ROOT}" "${CACHE_DIR}"
+
+python - <<'PY'
+from pathlib import Path
+from datasets import load_from_disk
+from huggingface_hub import snapshot_download
+
+data_root = Path(__import__("os").environ["DATA_ROOT"])
+cache_dir = __import__("os").environ["CACHE_DIR"]
+repos = {
+    "embedded-language-flows/wmt14_de-en_train_t5": data_root / "wmt14_de-en_train_t5",
+    "embedded-language-flows/wmt14_de-en_validation_t5": data_root / "wmt14_de-en_validation_t5",
+}
+for repo, out_dir in repos.items():
+    if out_dir.exists():
+        print(f"[skip] {out_dir} already exists")
+        continue
+    print(f"[download] {repo}")
+    snap = snapshot_download(repo_id=repo, repo_type="dataset", cache_dir=cache_dir)
+    ds = load_from_disk(snap)
+    print(f"[save] {repo}: rows={len(ds)} -> {out_dir}")
+    ds.save_to_disk(str(out_dir))
+print("[done]")
+PY

LTA_openwebtext_dualt/scripts/run_lta_lm1b_bert_absrope_time4_dirichlet_len128_C1_to_1024_4gpu_1m_mask1_sameT_save1k.sh

@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+
+# Apple-to-apple control for:
+#   run_lta_owt_bert_absrope_time4_dirichlet_len128_C1_to_1024_4gpu_1m_mask1_sameT_save1k.sh
+#
+# Only swap the training data to LM1B. Keep:
+#   BERT tokenizer, [CLS]/[SEP] wrapping, max_len=128, 4 GPUs,
+#   ddit_elf + 4 time tokens + RoPE + learned abs position embeddings,
+#   C=1->1024, mask_ratio=1.0, corrupt/model t shared,
+#   every-1k checkpoint + dual-line watcher inference.
+
+export DATA_PATH="${DATA_PATH:-data/lm1b_train_parquet}"
+export TEXT_COLUMN="${TEXT_COLUMN:-text}"
+export OPENWEBTEXT_SPLIT="${OPENWEBTEXT_SPLIT:-all}"
+export TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/workspace/imagenet_handoff_20260327/nlp_dts_light/assets/distilbert-base-uncased/tokenizer.json}"
+export TOKENIZED_HF=0
+export WRAP_MODE="${WRAP_MODE:-stream}"
+
+export DATE_TAG="${DATE_TAG:-$(date +%Y%m%d)}"
+export RUN_NAME="${RUN_NAME:-lta_lm1b_bert_absrope_time4_dirichlet_len128_C1_to_1024_mask1_sameT_gbs512_b32_4gpu_1m_save1k_${DATE_TAG}}"
+
+export MASTER_PORT="${MASTER_PORT:-32761}"
+export WATCH_LOG_DIR="${WATCH_LOG_DIR:-logs/lm1b_bert_absrope_time4_len128_C1_to_1024_mask1_sameT_every1k_dualline_watch}"
+export WATCH_OUT_BASE="${WATCH_OUT_BASE:-docs/lta_samples/metrics_${DATE_TAG}/lm1b_bert_absrope_time4_len128_C1_to_1024_mask1_sameT_every1k_dualline_dirres_c1_1024_n128/${RUN_NAME}}"
+
+bash scripts/run_lta_owt_bert_absrope_time4_dirichlet_len128_C1_to_1024_4gpu_1m_mask1_sameT_save1k.sh

LTA_openwebtext_dualt/scripts/run_lta_owt_dirichlet_len1024_Cv_to_2v_4gpu_abspos_specialloss_watch.sh

@@ -0,0 +1,42 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}"
+export NPROC_PER_NODE="${NPROC_PER_NODE:-4}"
+export MASTER_PORT="${MASTER_PORT:-32674}"
+
+export GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+export PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-4}"
+export TOTAL_STEPS="${TOTAL_STEPS:-20000}"
+export WARMUP_STEPS="${WARMUP_STEPS:-2500}"
+export SAVE_EVERY="${SAVE_EVERY:-1000}"
+export LATEST_EVERY="${LATEST_EVERY:-1000}"
+export LOG_EVERY="${LOG_EVERY:-100}"
+
+export MAX_LEN="${MAX_LEN:-1024}"
+export VOCAB_SIZE="${VOCAB_SIZE:-30522}"
+export CMIN="${CMIN:-${VOCAB_SIZE}}"
+export CMAX="${CMAX:-61044}"
+
+export ABS_POS_EMBED="${ABS_POS_EMBED:-1}"
+export FORCE_SPECIAL_CORRUPT_TOKEN_IDS="${FORCE_SPECIAL_CORRUPT_TOKEN_IDS:-101,102}"
+export SPECIAL_LOSS_TOKEN_IDS="${SPECIAL_LOSS_TOKEN_IDS:-101,102}"
+export SPECIAL_LOSS_WEIGHT="${SPECIAL_LOSS_WEIGHT:-16.0}"
+# Keep Bernoulli wrong input endpoints for specials in the first causal test.
+export SPECIAL_ENDPOINT_GOLD_TOKEN_IDS="${SPECIAL_ENDPOINT_GOLD_TOKEN_IDS:-}"
+
+export MIN_MASK_RATIO="${MIN_MASK_RATIO:-0.1}"
+export MAX_MASK_RATIO="${MAX_MASK_RATIO:-1.0}"
+export CATEGORICAL_WRONG_PROB_FLOOR="${CATEGORICAL_WRONG_PROB_FLOOR:-0.0}"
+
+export WATCH_ENABLED="${WATCH_ENABLED:-1}"
+export WATCH_CUDA_VISIBLE_DEVICES="${WATCH_CUDA_VISIBLE_DEVICES:-3}"
+export WATCH_STEP_INTERVAL="${WATCH_STEP_INTERVAL:-1000}"
+export WATCH_N_SAMPLES="${WATCH_N_SAMPLES:-128}"
+
+DATE_TAG="${DATE_TAG:-$(date +%Y%m%d)}"
+export RUN_NAME="${RUN_NAME:-lta_owt_dirichlet_len1024_Cv_to_2v_gbs512_4gpu_abspos_specialloss16_save1k_gumbelwatch_${DATE_TAG}}"
+
+bash scripts/run_lta_owt_dirichlet_len1024_Cv_to_2v_8gpu_save1k_with_gumbel_watch.sh

LTA_openwebtext_dualt/scripts/run_lta_owt_dirichlet_len1024_Cv_to_2v_8gpu_save1k_with_gumbel_watch.sh

@@ -0,0 +1,404 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+
+export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}"
+export TOKENIZERS_PARALLELISM=false
+export PYTHONUNBUFFERED=1
+export NCCL_DEBUG="${NCCL_DEBUG:-WARN}"
+export TORCH_DISTRIBUTED_TIMEOUT="${TORCH_DISTRIBUTED_TIMEOUT:-3600}"
+
+DATE_TAG="${DATE_TAG:-$(date +%Y%m%d)}"
+MAX_LEN="${MAX_LEN:-1024}"
+VOCAB_SIZE="${VOCAB_SIZE:-30522}"
+CMIN="${CMIN:-${VOCAB_SIZE}}"
+CMAX="${CMAX:-61044}"
+
+RUN_NAME="${RUN_NAME:-lta_owt_dirichlet_len${MAX_LEN}_Cv_to_2v_gbs512_8gpu_20k_save1k_gumbelwatch_${DATE_TAG}}"
+SAVE_DIR="${SAVE_DIR:-runs/${RUN_NAME}}"
+LOG_FILE="${LOG_FILE:-logs/${RUN_NAME}.log}"
+
+DATA_PATH="${DATA_PATH:-/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext}"
+TEXT_COLUMN="${TEXT_COLUMN:-text}"
+OPENWEBTEXT_SPLIT="${OPENWEBTEXT_SPLIT:-train_minus_100k}"
+TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/workspace/imagenet_handoff_20260327/nlp_dts_light/assets/distilbert-base-uncased/tokenizer.json}"
+SCORER="${SCORER:-/e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard}"
+TOKENIZED_HF="${TOKENIZED_HF:-0}"
+TOKENIZED_PAD_TOKEN="${TOKENIZED_PAD_TOKEN:-pad}"
+TOKENIZED_PREPEND_BOS="${TOKENIZED_PREPEND_BOS:-0}"
+TOKENIZED_APPEND_EOS="${TOKENIZED_APPEND_EOS:-0}"
+TOKENIZED_STRIP_EDGE_SPECIALS="${TOKENIZED_STRIP_EDGE_SPECIALS:-0}"
+
+CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
+NNODES="${NNODES:-1}"
+NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+NODE_RANK="${NODE_RANK:-0}"
+MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+MASTER_PORT="${MASTER_PORT:-32681}"
+
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-32}"
+TOTAL_STEPS="${TOTAL_STEPS:-20000}"
+WARMUP_STEPS="${WARMUP_STEPS:-2500}"
+LOG_EVERY="${LOG_EVERY:-100}"
+SAVE_EVERY="${SAVE_EVERY:-1000}"
+LATEST_EVERY="${LATEST_EVERY:-1000}"
+EVAL_EVERY="${EVAL_EVERY:-0}"
+NUM_WORKERS="${NUM_WORKERS:-0}"
+WRAP_MODE="${WRAP_MODE:-stream}"
+WRAP_RECORD_BUFFER_SIZE="${WRAP_RECORD_BUFFER_SIZE:-200}"
+ALLOW_EXISTING_SAVE_DIR="${ALLOW_EXISTING_SAVE_DIR:-0}"
+RESUME_PATH="${RESUME_PATH:-}"
+
+MIN_MASK_RATIO="${MIN_MASK_RATIO:-0.1}"
+MAX_MASK_RATIO="${MAX_MASK_RATIO:-1.0}"
+CATEGORICAL_WRONG_PROB_FLOOR="${CATEGORICAL_WRONG_PROB_FLOOR:-0.0}"
+CORRUPT_T_MODE="${CORRUPT_T_MODE:-independent}"
+ABS_POS_EMBED="${ABS_POS_EMBED:-0}"
+MODEL_TYPE="${MODEL_TYPE:-ddit}"
+ELF_NUM_TIME_TOKENS="${ELF_NUM_TIME_TOKENS:-4}"
+ELF_NUM_MODEL_MODE_TOKENS="${ELF_NUM_MODEL_MODE_TOKENS:-0}"
+QK_NORM="${QK_NORM:-1}"
+FORCE_SPECIAL_CORRUPT_TOKEN_IDS="${FORCE_SPECIAL_CORRUPT_TOKEN_IDS:-}"
+SPECIAL_ENDPOINT_GOLD_TOKEN_IDS="${SPECIAL_ENDPOINT_GOLD_TOKEN_IDS:-}"
+SPECIAL_LOSS_TOKEN_IDS="${SPECIAL_LOSS_TOKEN_IDS:-}"
+SPECIAL_LOSS_WEIGHT="${SPECIAL_LOSS_WEIGHT:-1.0}"
+
+WATCH_ENABLED="${WATCH_ENABLED:-0}"
+WATCH_CUDA_VISIBLE_DEVICES="${WATCH_CUDA_VISIBLE_DEVICES:-7}"
+WATCH_STEP_INTERVAL="${WATCH_STEP_INTERVAL:-1000}"
+WATCH_SLEEP_SECONDS="${WATCH_SLEEP_SECONDS:-30}"
+WATCH_N_SAMPLES="${WATCH_N_SAMPLES:-128}"
+WATCH_STEPS="${WATCH_STEPS:-128}"
+WATCH_ENDPOINT_TEMP="${WATCH_ENDPOINT_TEMP:-1.45}"
+WATCH_ENDPOINT_TOP_P="${WATCH_ENDPOINT_TOP_P:-0.95}"
+WATCH_GUMBEL_TAU_START="${WATCH_GUMBEL_TAU_START:-1.0}"
+WATCH_GUMBEL_TAU_END="${WATCH_GUMBEL_TAU_END:-0.2}"
+WATCH_DECODE_BATCH="${WATCH_DECODE_BATCH:-2}"
+WATCH_SCORE_BATCH="${WATCH_SCORE_BATCH:-1}"
+WATCH_SCORE_MAX_LENGTH="${WATCH_SCORE_MAX_LENGTH:-1024}"
+WATCH_DECODE_MODE="${WATCH_DECODE_MODE:-sde_gumbel}"
+WATCH_DUAL_SEMANTIC_POWER="${WATCH_DUAL_SEMANTIC_POWER:-1.5}"
+WATCH_DUAL_EARLY_TEMP="${WATCH_DUAL_EARLY_TEMP:-2.8}"
+WATCH_DUAL_LATE_TEMP="${WATCH_DUAL_LATE_TEMP:-1.45}"
+WATCH_DUAL_TEMP_END="${WATCH_DUAL_TEMP_END:-0.55}"
+WATCH_DUAL_TEMP_POWER="${WATCH_DUAL_TEMP_POWER:-1.5}"
+WATCH_OUT_BASE="${WATCH_OUT_BASE:-docs/lta_samples/metrics_${DATE_TAG}/owt_dirichlet_len${MAX_LEN}_Cv_to_2v_every1k_sde_gumbel_topp${WATCH_ENDPOINT_TOP_P//./p}_tau${WATCH_GUMBEL_TAU_START//./p}_to_${WATCH_GUMBEL_TAU_END//./p}_blend_c${CMIN}_${CMAX}_n${WATCH_N_SAMPLES}/${RUN_NAME}}"
+WATCH_LOG_DIR="${WATCH_LOG_DIR:-logs/owt_dirichlet_len${MAX_LEN}_Cv_to_2v_gumbel_sde_watch}"
+WATCH_SCRIPT="${WATCH_SCRIPT:-logs/${RUN_NAME}_watcher.sh}"
+WATCH_LOG="${WATCH_LOG:-logs/${RUN_NAME}_watcher.log}"
+WATCH_PID_FILE="${WATCH_PID_FILE:-logs/${RUN_NAME}_watcher.pid}"
+
+if [[ -f "${SAVE_DIR}/args.json" && -z "${RESUME_PATH}" && "${ALLOW_EXISTING_SAVE_DIR}" != "1" ]]; then
+  echo "Refusing to start because SAVE_DIR already contains args.json: ${SAVE_DIR}" >&2
+  echo "Set RESUME_PATH, use a new RUN_NAME/SAVE_DIR, or ALLOW_EXISTING_SAVE_DIR=1." >&2
+  exit 2
+fi
+
+mkdir -p logs runs "${SAVE_DIR}" "${WATCH_LOG_DIR}" "${WATCH_OUT_BASE}"
+
+RESUME_ARGS=()
+if [[ -n "${RESUME_PATH}" ]]; then
+  RESUME_ARGS+=(--resume_path "${RESUME_PATH}")
+fi
+
+TEXT_COLUMN_ARGS=()
+if [[ -n "${TEXT_COLUMN}" ]]; then
+  TEXT_COLUMN_ARGS+=(--text_column "${TEXT_COLUMN}")
+fi
+
+OPENWEBTEXT_ARGS=(--openwebtext_split "${OPENWEBTEXT_SPLIT}")
+WRAP_ARGS=(--wrap --wrap_mode "${WRAP_MODE}" --wrap_record_buffer_size "${WRAP_RECORD_BUFFER_SIZE}")
+TOKENIZED_ARGS=()
+DATA_MODE_LABEL="raw_wrap_${WRAP_MODE}"
+if [[ "${TOKENIZED_HF}" == "1" || "${TOKENIZED_HF}" == "true" || "${TOKENIZED_HF}" == "True" ]]; then
+  TEXT_COLUMN_ARGS=()
+  OPENWEBTEXT_ARGS=()
+  WRAP_ARGS=()
+  TOKENIZED_ARGS=(--tokenized_hf --tokenized_pad_token "${TOKENIZED_PAD_TOKEN}")
+  if [[ "${TOKENIZED_PREPEND_BOS}" == "1" || "${TOKENIZED_PREPEND_BOS}" == "true" || "${TOKENIZED_PREPEND_BOS}" == "True" ]]; then
+    TOKENIZED_ARGS+=(--tokenized_prepend_bos)
+  fi
+  if [[ "${TOKENIZED_APPEND_EOS}" == "1" || "${TOKENIZED_APPEND_EOS}" == "true" || "${TOKENIZED_APPEND_EOS}" == "True" ]]; then
+    TOKENIZED_ARGS+=(--tokenized_append_eos)
+  fi
+  if [[ "${TOKENIZED_STRIP_EDGE_SPECIALS}" == "1" || "${TOKENIZED_STRIP_EDGE_SPECIALS}" == "true" || "${TOKENIZED_STRIP_EDGE_SPECIALS}" == "True" ]]; then
+    TOKENIZED_ARGS+=(--tokenized_strip_edge_specials)
+  fi
+  DATA_MODE_LABEL="tokenized_hf_pad_${TOKENIZED_PAD_TOKEN}"
+fi
+
+ABS_POS_ARGS=()
+if [[ "${ABS_POS_EMBED}" == "1" || "${ABS_POS_EMBED}" == "true" || "${ABS_POS_EMBED}" == "True" ]]; then
+  ABS_POS_ARGS+=(--abs_pos_embed)
+fi
+
+QK_NORM_ARGS=(--qk_norm)
+if [[ "${QK_NORM}" == "0" || "${QK_NORM}" == "false" || "${QK_NORM}" == "False" ]]; then
+  QK_NORM_ARGS=(--no-qk_norm)
+fi
+
+SPECIAL_ARGS=()
+if [[ -n "${FORCE_SPECIAL_CORRUPT_TOKEN_IDS}" ]]; then
+  SPECIAL_ARGS+=(--force_special_corrupt_token_ids "${FORCE_SPECIAL_CORRUPT_TOKEN_IDS}")
+fi
+if [[ -n "${SPECIAL_ENDPOINT_GOLD_TOKEN_IDS}" ]]; then
+  SPECIAL_ARGS+=(--special_endpoint_gold_token_ids "${SPECIAL_ENDPOINT_GOLD_TOKEN_IDS}")
+fi
+if [[ -n "${SPECIAL_LOSS_TOKEN_IDS}" ]]; then
+  SPECIAL_ARGS+=(--special_loss_token_ids "${SPECIAL_LOSS_TOKEN_IDS}" --special_loss_weight "${SPECIAL_LOSS_WEIGHT}")
+fi
+
+write_watcher() {
+  cat > "${WATCH_SCRIPT}" <<'WATCH_EOF'
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}"
+export TOKENIZERS_PARALLELISM=false
+export PYTHONUNBUFFERED=1
+
+: "${RUN_DIR:?RUN_DIR is required}"
+: "${OUT_BASE:?OUT_BASE is required}"
+: "${LOG_DIR:?LOG_DIR is required}"
+: "${TOKENIZER_PATH:?TOKENIZER_PATH is required}"
+: "${SCORER:?SCORER is required}"
+
+RUN_STEM="$(basename "${RUN_DIR}")"
+TEMP_TAG="${ENDPOINT_TEMP//./p}"
+PROCESSED_FILE="${LOG_DIR}/processed_${RUN_STEM}_steps${STEPS}_c${CMIN}_${CMAX}_gumbel_t${TEMP_TAG}_n${N_SAMPLES}.txt"
+
+mkdir -p "${OUT_BASE}" "${LOG_DIR}"
+touch "${PROCESSED_FILE}"
+
+echo "[watch-gumbel] run_dir=${RUN_DIR}"
+echo "[watch-gumbel] out_base=${OUT_BASE}"
+echo "[watch-gumbel] interval=${STEP_INTERVAL} max_len=${MAX_LEN} steps=${STEPS} c=${CMIN}->${CMAX} decode_mode=${DECODE_MODE:-sde_gumbel} temp=${ENDPOINT_TEMP} top_p=${ENDPOINT_TOP_P} tau=${GUMBEL_TAU_START}->${GUMBEL_TAU_END} n=${N_SAMPLES}"
+
+while true; do
+  shopt -s nullglob
+  ckpts=("${RUN_DIR}"/step_*.pt)
+  shopt -u nullglob
+
+  if (( ${#ckpts[@]} == 0 )); then
+    echo "[watch-gumbel] $(date +%F_%T) no ckpt yet"
+    sleep "${SLEEP_SECONDS}"
+    continue
+  fi
+
+  printf "%s\n" "${ckpts[@]}" | sort | while read -r ckpt; do
+    base="$(basename "${ckpt}")"
+    step="${base#step_}"
+    step="${step%.pt}"
+    step_num=$((10#${step}))
+    if (( step_num % STEP_INTERVAL != 0 )); then
+      continue
+    fi
+    if grep -Fxq "${ckpt}" "${PROCESSED_FILE}"; then
+      continue
+    fi
+
+    out_dir="${OUT_BASE}/step_${step}"
+    log_file="${LOG_DIR}/infer_${RUN_STEM}_step_${step}.log"
+    mkdir -p "${out_dir}"
+
+    echo "[watch-gumbel] $(date +%F_%T) infer ${ckpt} -> ${out_dir}" | tee -a "${log_file}"
+    if [[ "${DECODE_MODE:-sde_gumbel}" == "dual_line_probe" ]]; then
+      CUDA_VISIBLE_DEVICES="${WATCH_CUDA_VISIBLE_DEVICES}" python scripts/infer_softkl_decode_probe.py \
+        --checkpoint "${ckpt}" \
+        --tokenizer_path "${TOKENIZER_PATH}" \
+        --scorer "${SCORER}" \
+        --score \
+        --out_dir "${out_dir}" \
+        --max_lens "${MAX_LEN}" \
+        --n_samples "${N_SAMPLES}" \
+        --batch_size "${DECODE_BATCH}" \
+        --steps "${STEPS}" \
+        --decode_rule dual_line_resample \
+        --c_min "${CMIN}" \
+        --c_max "${CMAX}" \
+        --input_noise_dirichlet_concentration "${CMIN}" \
+        --anchor_mode state \
+        --model_t_mode flow \
+        --time_schedule uniform \
+        --support_power 1.0 \
+        --semantic_power "${DUAL_SEMANTIC_POWER}" \
+        --early_temp "${DUAL_EARLY_TEMP}" \
+        --late_temp "${DUAL_LATE_TEMP}" \
+        --temp_end "${DUAL_TEMP_END}" \
+        --temp_power "${DUAL_TEMP_POWER}" \
+        --final_from blend \
+        --final_decode argmax \
+        --seed 20260524 \
+        2>&1 | tee -a "${log_file}"
+    else
+      CUDA_VISIBLE_DEVICES="${WATCH_CUDA_VISIBLE_DEVICES}" python scripts/eval_lm1b_c1024_fullycoupled_sde_genppl.py \
+        --checkpoint "${ckpt}" \
+        --tokenizer_path "${TOKENIZER_PATH}" \
+        --scorer "${SCORER}" \
+        --out_dir "${out_dir}" \
+        --n_samples "${N_SAMPLES}" \
+        --max_len "${MAX_LEN}" \
+        --steps "${STEPS}" \
+        --batch_size "${DECODE_BATCH}" \
+        --score_batch "${SCORE_BATCH}" \
+        --score_max_length "${SCORE_MAX_LENGTH}" \
+        --concentration_min "${CMIN}" \
+        --concentration_max "${CMAX}" \
+        --endpoint_temp "${ENDPOINT_TEMP}" \
+        --endpoint_projection gumbel_softmax \
+        --endpoint_top_p "${ENDPOINT_TOP_P}" \
+        --gumbel_tau_start "${GUMBEL_TAU_START}" \
+        --gumbel_tau_end "${GUMBEL_TAU_END}" \
+        --model_t_mode support_t \
+        --mean_mode endpoint_only \
+        --semantic_power 1.0 \
+        --noise_init dirichlet \
+        --noise_dirichlet_concentration "${CMIN}" \
+        --sde_resample dirichlet \
+        --final_from blend_0.5 \
+        --seed 20260524 \
+        2>&1 | tee -a "${log_file}"
+    fi
+
+    echo "${ckpt}" >> "${PROCESSED_FILE}"
+    echo "[watch-gumbel] $(date +%F_%T) done step_${step}" | tee -a "${log_file}"
+  done
+
+  sleep "${SLEEP_SECONDS}"
+done
+WATCH_EOF
+  chmod +x "${WATCH_SCRIPT}"
+}
+
+start_watcher() {
+  if [[ "${WATCH_ENABLED}" != "1" ]]; then
+    echo "[launch] watcher disabled"
+    return
+  fi
+  if [[ -f "${WATCH_PID_FILE}" ]]; then
+    old_pid="$(cat "${WATCH_PID_FILE}" || true)"
+    if [[ -n "${old_pid}" ]] && kill -0 "${old_pid}" 2>/dev/null; then
+      echo "[launch] watcher already running pid=${old_pid}"
+      return
+    fi
+  fi
+  write_watcher
+  nohup env \
+    RUN_DIR="${SAVE_DIR}" \
+    OUT_BASE="${WATCH_OUT_BASE}" \
+    LOG_DIR="${WATCH_LOG_DIR}" \
+    TOKENIZER_PATH="${TOKENIZER_PATH}" \
+    SCORER="${SCORER}" \
+    WATCH_CUDA_VISIBLE_DEVICES="${WATCH_CUDA_VISIBLE_DEVICES}" \
+    N_SAMPLES="${WATCH_N_SAMPLES}" \
+    MAX_LEN="${MAX_LEN}" \
+    STEPS="${WATCH_STEPS}" \
+    DECODE_BATCH="${WATCH_DECODE_BATCH}" \
+    SCORE_BATCH="${WATCH_SCORE_BATCH}" \
+    SCORE_MAX_LENGTH="${WATCH_SCORE_MAX_LENGTH}" \
+    DECODE_MODE="${WATCH_DECODE_MODE}" \
+    DUAL_SEMANTIC_POWER="${WATCH_DUAL_SEMANTIC_POWER}" \
+    DUAL_EARLY_TEMP="${WATCH_DUAL_EARLY_TEMP}" \
+    DUAL_LATE_TEMP="${WATCH_DUAL_LATE_TEMP}" \
+    DUAL_TEMP_END="${WATCH_DUAL_TEMP_END}" \
+    DUAL_TEMP_POWER="${WATCH_DUAL_TEMP_POWER}" \
+    CMIN="${CMIN}" \
+    CMAX="${CMAX}" \
+    ENDPOINT_TEMP="${WATCH_ENDPOINT_TEMP}" \
+    ENDPOINT_TOP_P="${WATCH_ENDPOINT_TOP_P}" \
+    GUMBEL_TAU_START="${WATCH_GUMBEL_TAU_START}" \
+    GUMBEL_TAU_END="${WATCH_GUMBEL_TAU_END}" \
+    STEP_INTERVAL="${WATCH_STEP_INTERVAL}" \
+    SLEEP_SECONDS="${WATCH_SLEEP_SECONDS}" \
+    bash "${WATCH_SCRIPT}" > "${WATCH_LOG}" 2>&1 &
+  echo "$!" > "${WATCH_PID_FILE}"
+  echo "[launch] watcher pid=$(cat "${WATCH_PID_FILE}") log=${WATCH_LOG}"
+}
+
+echo "[launch] run=${RUN_NAME}"
+echo "[launch] data_path=${DATA_PATH} mode=${DATA_MODE_LABEL} split=${OPENWEBTEXT_SPLIT} text_column=${TEXT_COLUMN}"
+echo "[launch] max_len=${MAX_LEN} gbs=${GLOBAL_BATCH_SIZE} per_gpu=${PER_GPU_BATCH_SIZE} total_steps=${TOTAL_STEPS}"
+echo "[launch] dirichlet C=${CMIN}->${CMAX} wrong_floor=${CATEGORICAL_WRONG_PROB_FLOOR} corrupt_t_mode=${CORRUPT_T_MODE}"
+echo "[launch] model_type=${MODEL_TYPE} abs_pos=${ABS_POS_EMBED} elf_time_tokens=${ELF_NUM_TIME_TOKENS} qk_norm=${QK_NORM} force_special_corrupt=${FORCE_SPECIAL_CORRUPT_TOKEN_IDS:-none} special_endpoint_gold=${SPECIAL_ENDPOINT_GOLD_TOKEN_IDS:-none} special_loss=${SPECIAL_LOSS_TOKEN_IDS:-none}x${SPECIAL_LOSS_WEIGHT}"
+start_watcher
+
+CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" python -m torch.distributed.run \
+  --nnodes="${NNODES}" \
+  --nproc_per_node="${NPROC_PER_NODE}" \
+  --node_rank="${NODE_RANK}" \
+  --master_addr="${MASTER_ADDR}" \
+  --master_port="${MASTER_PORT}" \
+  train.py \
+  --data_path "${DATA_PATH}" \
+  "${TEXT_COLUMN_ARGS[@]}" \
+  "${OPENWEBTEXT_ARGS[@]}" \
+  "${TOKENIZED_ARGS[@]}" \
+  --detokenizer auto \
+  --tokenizer_path "${TOKENIZER_PATH}" \
+  --save_dir "${SAVE_DIR}" \
+  "${WRAP_ARGS[@]}" \
+  --max_len "${MAX_LEN}" \
+  --batch_size "${PER_GPU_BATCH_SIZE}" \
+  --num_workers "${NUM_WORKERS}" \
+  --global_batch_size "${GLOBAL_BATCH_SIZE}" \
+  --total_steps "${TOTAL_STEPS}" \
+  --log_every "${LOG_EVERY}" \
+  --eval_every "${EVAL_EVERY}" \
+  --save_every "${SAVE_EVERY}" \
+  --latest_every "${LATEST_EVERY}" \
+  --lr 3e-4 \
+  --weight_decay 0 \
+  --adam_beta1 0.9 \
+  --adam_beta2 0.999 \
+  --adam_eps 1e-8 \
+  --warmup_steps "${WARMUP_STEPS}" \
+  --lr_schedule constant_warmup \
+  --grad_clip 1.0 \
+  --seed 123 \
+  --d_model 768 \
+  --cond_dim 128 \
+  --n_layers 12 \
+  --n_heads 12 \
+  --dim_ff 3072 \
+  --dropout 0.1 \
+  "${ABS_POS_ARGS[@]}" \
+  --model_type "${MODEL_TYPE}" \
+  --elf_num_time_tokens "${ELF_NUM_TIME_TOKENS}" \
+  --elf_num_model_mode_tokens "${ELF_NUM_MODEL_MODE_TOKENS}" \
+  "${QK_NORM_ARGS[@]}" \
+  --state_format prob \
+  --bridge dirichlet \
+  --target_loss hard_ce \
+  --target_prob 1.0 \
+  --min_t 0.0 \
+  --max_t 1.0 \
+  --dual_t \
+  --corrupt_t_mode "${CORRUPT_T_MODE}" \
+  --corrupt_min_t 0.0 \
+  --corrupt_max_t 1.0 \
+  --min_mask_ratio "${MIN_MASK_RATIO}" \
+  --max_mask_ratio "${MAX_MASK_RATIO}" \
+  --wrong_token_replace_prob 1.0 \
+  --wrong_token_schedule linear_t \
+  --wrong_token_exp_k 1.0 \
+  --categorical_wrong_prob_floor "${CATEGORICAL_WRONG_PROB_FLOOR}" \
+  --dirichlet_concentration_min "${CMIN}" \
+  --dirichlet_concentration_max "${CMAX}" \
+  --dirichlet_endpoint_mode categorical_dual_t \
+  --dirichlet_semantic_t_mode same \
+  --dirichlet_semantic_t_value 0.0 \
+  --categorical_wrong_from_full_vocab \
+  --eps 1e-8 \
+  --infer_steps 128 \
+  --decode_damping 1.0 \
+  --max_gamma 1.0 \
+  --decode_solver flowmap \
+  --noise_init logistic_normal \
+  --bridge_noise_init logistic_normal \
+  --noise_sigma -1 \
+  "${SPECIAL_ARGS[@]}" \
+  "${RESUME_ARGS[@]}" \
+  --bf16 2>&1 | tee -a "${LOG_FILE}"

LTA_openwebtext_dualt/scripts/run_lta_owt_t5_len128_uniform10k_then_lognsr_4gpu.sh

@@ -0,0 +1,60 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-4,5,6,7}"
+export NPROC_PER_NODE="${NPROC_PER_NODE:-4}"
+export MASTER_PORT="${MASTER_PORT:-39331}"
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"
+export PYTHONUNBUFFERED=1
+export TOKENIZERS_PARALLELISM=false
+
+BASE_RUN_NAME="${BASE_RUN_NAME:-lta_owt_t5_len128_uniform10k_then_lognsr_sde_rollin_4gpu_$(date +%Y%m%d_%H%M%S)}"
+WARMUP_RUN_NAME="${BASE_RUN_NAME}_warmup_uniform_norollin"
+MAIN_RUN_NAME="${BASE_RUN_NAME}_resume_lognsr_sde_rollin"
+
+COMMON_ENV=(
+  MAX_LEN=128
+  GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+  PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-32}"
+  SAVE_EVERY="${SAVE_EVERY:-1000}"
+  LATEST_EVERY="${LATEST_EVERY:-1000}"
+  LOG_EVERY="${LOG_EVERY:-100}"
+  TARGET_LOSS="${TARGET_LOSS:-hard_ce}"
+  MIN_MASK_RATIO=1.0
+  MAX_MASK_RATIO=1.0
+  MASK_MIXTURE_ALL_PROB=1.0
+)
+
+echo "[two-phase] base_run_name=${BASE_RUN_NAME}"
+echo "[two-phase] phase1=${WARMUP_RUN_NAME}: uniform t, no roll-in, 10k steps"
+env "${COMMON_ENV[@]}" \
+  RUN_NAME="${WARMUP_RUN_NAME}" \
+  TOTAL_STEPS=10000 \
+  T_SAMPLING_MODE=uniform \
+  ROLLOUT_TRAIN_PROB=0.0 \
+  bash scripts/launch_lta_owt_t5_len128_sde_rollin_lognsr_4gpu.sh
+
+WARMUP_CKPT="runs/${WARMUP_RUN_NAME}/step_0010000.pt"
+if [[ ! -f "${WARMUP_CKPT}" ]]; then
+  echo "[two-phase] missing warmup checkpoint: ${WARMUP_CKPT}" >&2
+  exit 1
+fi
+
+echo "[two-phase] phase2=${MAIN_RUN_NAME}: resume ${WARMUP_CKPT}, lognsr_gumbel + sde roll-in"
+env "${COMMON_ENV[@]}" \
+  RUN_NAME="${MAIN_RUN_NAME}" \
+  RESUME_PATH="${WARMUP_CKPT}" \
+  TOTAL_STEPS="${TOTAL_STEPS:-1000000}" \
+  T_SAMPLING_MODE=lognsr_gumbel \
+  T_GUMBEL_LOC="${T_GUMBEL_LOC:-2.2}" \
+  T_GUMBEL_SCALE="${T_GUMBEL_SCALE:-0.8}" \
+  ROLLOUT_TRAIN_PROB="${ROLLOUT_TRAIN_PROB:-0.50}" \
+  ROLLOUT_TRAIN_RULE=dirichlet_resample \
+  ROLLOUT_TRAIN_TIME_MODE=sampled_path \
+  ROLLOUT_TRAIN_STEPS="${ROLLOUT_TRAIN_STEPS:-3}" \
+  ROLLOUT_TRAIN_STEPS_MIN="${ROLLOUT_TRAIN_STEPS_MIN:-0}" \
+  ROLLOUT_TRAIN_INFER_STEPS="${ROLLOUT_TRAIN_INFER_STEPS:-1}" \
+  ROLLOUT_TRAIN_TEMP="${ROLLOUT_TRAIN_TEMP:-1.0}" \
+  bash scripts/launch_lta_owt_t5_len128_sde_rollin_lognsr_4gpu.sh

LTA_openwebtext_dualt/scripts/run_train8_ctx1024_sampleds_sweep_4gpu.sh

@@ -0,0 +1,282 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+
+export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}"
+export TOKENIZERS_PARALLELISM=false
+export PYTHONUNBUFFERED=1
+
+BASE_CACHE="${BASE_CACHE:-/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext_lta_cached_chunks}"
+CACHE_PREFIX="${CACHE_PREFIX:-gpt2}"
+CACHE_DIR="${CACHE_DIR:-}"
+TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-standard/tokenizer.json}"
+MAX_LEN="${MAX_LEN:-1024}"
+N_SAMPLES="${N_SAMPLES:-64}"
+INFER_STEPS="${INFER_STEPS:-128}"
+STEP_CHUNK="${STEP_CHUNK:-1000}"
+MAX_TOTAL_STEPS="${MAX_TOTAL_STEPS:-12000}"
+PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-128}"
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+LEARNING_RATE="${LEARNING_RATE:-0.002}"
+WEIGHT_DECAY="${WEIGHT_DECAY:-0.1}"
+GROUP_STAMP="${GROUP_STAMP:-$(date +%Y%m%d_%H%M%S)}"
+WAIT_FOR_RUN="${WAIT_FOR_RUN:-}"
+OUT_ROOT="${OUT_ROOT:-docs/lta_samples/metrics_20260517/ctx1024_sampleds_sweep_bs512_ode128_${GROUP_STAMP}}"
+DRIVER_LOG="${DRIVER_LOG:-logs/ctx1024_sampleds_sweep_4gpu/${GROUP_STAMP}.log}"
+CURVE_CSV="${CURVE_CSV:-${OUT_ROOT}/hit_ratio_curve.csv}"
+EVAL_DECODE_RULES="${EVAL_DECODE_RULES:-flowmap}"
+EVAL_C_MAX="${EVAL_C_MAX:-512}"
+EVAL_EARLY_TEMP="${EVAL_EARLY_TEMP:-1.0}"
+EVAL_LATE_TEMP="${EVAL_LATE_TEMP:-1.0}"
+mkdir -p "$(dirname "${DRIVER_LOG}")" "${OUT_ROOT}"
+
+if [[ -n "${CACHE_DIR}" ]]; then
+  cache="${CACHE_DIR}"
+else
+  cache="${BASE_CACHE}/${CACHE_PREFIX}_len${MAX_LEN}_train8_compact_overfit"
+fi
+vocab_size="$(
+python - "$cache" <<'PY'
+import json
+import sys
+from pathlib import Path
+meta = json.loads((Path(sys.argv[1]) / "meta.json").read_text())
+print(int(meta.get("compact_vocab_size", meta.get("vocab_size"))))
+PY
+)"
+
+if [[ ! -f "${CURVE_CSV}" ]]; then
+  echo "config,decode_rule,run_name,ckpt_step,train_views_seen,train_tokens_seen,token_acc_mean,exact_count,exact_ref_count,exact_ref_hits" > "${CURVE_CSV}"
+fi
+
+latest_step() {
+  local run_name="$1"
+  python - "$run_name" <<'PY'
+import re
+import sys
+from pathlib import Path
+run = Path("runs") / sys.argv[1]
+steps = []
+for path in run.glob("step_*.pt"):
+    m = re.search(r"step_(\d+)\.pt$", path.name)
+    if m:
+        steps.append(int(m.group(1)))
+print(max(steps) if steps else 0)
+PY
+}
+
+free_port() {
+  python - <<'PY'
+import socket
+with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+    s.bind(("127.0.0.1", 0))
+    print(s.getsockname()[1])
+PY
+}
+
+eval_latest() {
+  local config="$1"
+  local run_name="$2"
+  local target_step="$3"
+  IFS=',' read -r -a decode_rules <<<"${EVAL_DECODE_RULES}"
+  for decode_rule in "${decode_rules[@]}"; do
+    decode_rule="$(echo "${decode_rule}" | xargs)"
+    [[ -n "${decode_rule}" ]] || continue
+    local out_dir="${OUT_ROOT}/${config}/step_${target_step}/${decode_rule}_c${EVAL_C_MAX}"
+    mkdir -p "${out_dir}"
+    CUDA_VISIBLE_DEVICES="${EVAL_CUDA_VISIBLE_DEVICES:-0}" python scripts/eval_train8_decode_acc.py \
+      --runs_glob "runs/${run_name}" \
+      --data_dir "${cache}" \
+      --tokenizer_path "${TOKENIZER_PATH}" \
+      --out_dir "${out_dir}" \
+      --max_len "${MAX_LEN}" \
+      --n_samples "${N_SAMPLES}" \
+      --batch_size "${N_SAMPLES}" \
+      --latest_only \
+      --endpoint_softenings none \
+      --steps "${INFER_STEPS}" \
+      --decode_rule "${decode_rule}" \
+      --time_schedule logit_normal \
+      --time_logit_mean -1.5 \
+      --time_logit_std 0.8 \
+      --model_t_mode post \
+      --c_min 1 \
+      --c_max "${EVAL_C_MAX}" \
+      --early_temp "${EVAL_EARLY_TEMP}" \
+      --late_temp "${EVAL_LATE_TEMP}" \
+      --final_from state \
+      --final_decode argmax
+    python - "$out_dir" "$N_SAMPLES" "$GLOBAL_BATCH_SIZE" "$MAX_LEN" "$CURVE_CSV" "$config" "$run_name" "$decode_rule" <<'PY'
+import json
+import sys
+from pathlib import Path
+out = Path(sys.argv[1])
+n = int(sys.argv[2])
+global_batch = int(sys.argv[3])
+max_len = int(sys.argv[4])
+curve = Path(sys.argv[5])
+config = sys.argv[6]
+run_name = sys.argv[7]
+decode_rule = sys.argv[8]
+row = json.loads((out / "decode_token_acc.jsonl").read_text().splitlines()[-1])
+views = int(row["ckpt_step"]) * global_batch
+tokens = views * max_len
+print(
+    "RESULT "
+    f"config={config} decode={decode_rule} run={run_name} ckpt_step={row['ckpt_step']} "
+    f"views={views} token_acc={row['token_acc_mean']:.4f} "
+    f"exact={row['exact_count']}/{n} exact_refs={row['exact_ref_count']} "
+    f"hits={row['exact_ref_hits']}",
+    flush=True,
+)
+with curve.open("a", encoding="utf-8") as f:
+    f.write(
+        f"{config},{decode_rule},{run_name},{row['ckpt_step']},{views},{tokens},{row['token_acc_mean']},"
+        f"{row['exact_count']},{row['exact_ref_count']},\"{row['exact_ref_hits']}\"\n"
+    )
+PY
+  done
+}
+
+run_config() {
+  local config="$1"
+  local rollout_prob="$2"
+  local time_mode="$3"
+  local rollout_steps="$4"
+  local s_dist="$5"
+  local s_min_frac="$6"
+  local s_max_frac="$7"
+  local s_beta_alpha="$8"
+  local s_beta_beta="$9"
+  local output_wd="${10}"
+  local sync_t="${11}"
+  local rollout_steps_min="${12:-${ROLLOUT_TRAIN_STEPS_MIN:--1}}"
+  local run_name="${RUN_PREFIX:-train8_ctx1024}_${config}_${GROUP_STAMP}"
+  echo "[ctx1024-sampleds] config=${config} run=${run_name} p=${rollout_prob} mode=${time_mode} steps=${rollout_steps} steps_min=${rollout_steps_min} s_dist=${s_dist} s_frac=${s_min_frac}->${s_max_frac} beta=${s_beta_alpha},${s_beta_beta} outwd=${output_wd} sync_t=${sync_t}" | tee -a "${DRIVER_LOG}"
+  while :; do
+    local step_now
+    step_now="$(latest_step "${run_name}")"
+    if [[ "${step_now}" -ge "${MAX_TOTAL_STEPS}" ]]; then
+      echo "[ctx1024-sampleds] capped config=${config} step=${step_now}" | tee -a "${DRIVER_LOG}"
+      break
+    fi
+    local target_step=$((step_now + STEP_CHUNK))
+    if [[ "${target_step}" -gt "${MAX_TOTAL_STEPS}" ]]; then
+      target_step="${MAX_TOTAL_STEPS}"
+    fi
+    local resume_path=""
+    if [[ -f "runs/${run_name}/latest.pt" ]]; then
+      resume_path="runs/${run_name}/latest.pt"
+    fi
+    echo "[ctx1024-sampleds] train config=${config} from=${step_now} to=${target_step}" | tee -a "${DRIVER_LOG}"
+    CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}" \
+    NPROC_PER_NODE="${NPROC_PER_NODE:-4}" \
+    MASTER_PORT="$(free_port)" \
+    OWT_CHUNK_CACHE_DIR="${cache}" \
+    OWT_EXACT_REPEAT_PER_CHUNK="${OWT_EXACT_REPEAT_PER_CHUNK:-64}" \
+    MAX_LEN="${MAX_LEN}" \
+    VOCAB_SIZE_OVERRIDE="${vocab_size}" \
+    D_MODEL="${D_MODEL:-192}" \
+    COND_DIM="${COND_DIM:-64}" \
+    N_LAYERS="${N_LAYERS:-3}" \
+    N_HEADS="${N_HEADS:-3}" \
+    DIM_FF="${DIM_FF:-768}" \
+    TOTAL_STEPS="${target_step}" \
+    PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE}" \
+    GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE}" \
+    NUM_WORKERS="${NUM_WORKERS:-0}" \
+    LOG_EVERY="${LOG_EVERY:-100}" \
+    SAVE_EVERY="${STEP_CHUNK}" \
+    LATEST_EVERY="${STEP_CHUNK}" \
+    WARMUP_STEPS="${WARMUP_STEPS:-10}" \
+    LEARNING_RATE="${LEARNING_RATE}" \
+    WEIGHT_DECAY="${WEIGHT_DECAY}" \
+    OUTPUT_WEIGHT_DECAY="${output_wd}" \
+    MUON_IMPL="${MUON_IMPL:-legacy}" \
+    MIN_MASK_RATIO=1.0 \
+    MAX_MASK_RATIO=1.0 \
+    MASK_MIXTURE_LOWK_PROB=0.0 \
+    MASK_MIXTURE_ALL_PROB=1.0 \
+    LOWK_CLEAN_TOKENS=0 \
+    CLEAN_STATE_MODE=onehot \
+    TARGET_LOSS=hard_ce \
+    DIRICHLET_CONCENTRATION_MIN=1.0 \
+    DIRICHLET_CONCENTRATION_MAX=1024 \
+    SIMPLEX_BRIDGE_SAMPLER=dirichlet \
+    CATEGORICAL_WRONG_PROB_FLOOR=0.0 \
+    ROLLOUT_TRAIN_PROB="${rollout_prob}" \
+    ROLLOUT_TRAIN_STEPS="${rollout_steps}" \
+    ROLLOUT_TRAIN_STEPS_MIN="${rollout_steps_min}" \
+    ROLLOUT_TRAIN_INFER_STEPS=1 \
+    ROLLOUT_TRAIN_TIME_MODE="${time_mode}" \
+    ROLLOUT_TRAIN_S_DIST="${s_dist}" \
+    ROLLOUT_TRAIN_S_MIN_FRAC="${s_min_frac}" \
+    ROLLOUT_TRAIN_S_MAX_FRAC="${s_max_frac}" \
+    ROLLOUT_TRAIN_S_BETA_ALPHA="${s_beta_alpha}" \
+    ROLLOUT_TRAIN_S_BETA_BETA="${s_beta_beta}" \
+    ROLLOUT_TRAIN_TEMP=1.0 \
+    ROLLOUT_TRAIN_MAX_GAMMA=1.0 \
+    ROLLOUT_TRAIN_CORRUPT_ONLY=1 \
+    ROLLOUT_TRAIN_SAMPLEWISE=1 \
+    ROLLOUT_TRAIN_SELECTED_ONLY=1 \
+    ROLLOUT_TRAIN_COMPUTE_ALWAYS=0 \
+    ROLLOUT_TRAIN_SYNC_T="${sync_t}" \
+    T_SAMPLING_MODE="${T_SAMPLING_MODE:-uniform}" \
+    RUN_NAME="${run_name}" \
+    RESUME_PATH="${resume_path}" \
+    bash scripts/launch_lta_owt_gpt2_softendpoint_mn_pilot_4gpu.sh
+    echo "[ctx1024-sampleds] eval config=${config} step=${target_step}" | tee -a "${DRIVER_LOG}"
+    eval_latest "${config}" "${run_name}" "${target_step}" | tee -a "${DRIVER_LOG}"
+    local eval_rule_count
+    eval_rule_count="$(python - "$EVAL_DECODE_RULES" <<'PY'
+import sys
+print(len([x for x in sys.argv[1].split(",") if x.strip()]) or 1)
+PY
+)"
+    if python - "$CURVE_CSV" "$N_SAMPLES" "$eval_rule_count" <<'PY'
+import csv
+import sys
+from pathlib import Path
+path = Path(sys.argv[1])
+target = min(60, int(sys.argv[2]))
+rule_count = max(1, int(sys.argv[3]))
+rows = list(csv.DictReader(path.open()))
+recent = rows[-rule_count:]
+raise SystemExit(0 if any(int(float(r["exact_count"])) >= target for r in recent) else 1)
+PY
+    then
+      echo "[ctx1024-sampleds] early-hit config=${config}" | tee -a "${DRIVER_LOG}"
+      break
+    fi
+  done
+}
+
+if [[ -n "${WAIT_FOR_RUN}" ]]; then
+  echo "[ctx1024-sampleds] waiting for run=${WAIT_FOR_RUN}" | tee -a "${DRIVER_LOG}"
+  while pgrep -f "${WAIT_FOR_RUN}" >/dev/null; do
+    sleep 60
+  done
+fi
+
+echo "[ctx1024-sampleds] start stamp=${GROUP_STAMP} len=${MAX_LEN} vocab=${vocab_size} out=${OUT_ROOT}" | tee -a "${DRIVER_LOG}"
+
+CONFIGS=(
+  "p50_unif0_0p125_outwdm1|0.50|sampled_s|1|uniform|0.0|0.125|2.0|6.0|-1|1"
+  "p25_unif0_0p125_outwdm1|0.25|sampled_s|1|uniform|0.0|0.125|2.0|6.0|-1|1"
+  "p50_unif0_0p25_outwdm1|0.50|sampled_s|1|uniform|0.0|0.25|2.0|6.0|-1|1"
+  "p50_beta2_6_0_0p5_outwdm1|0.50|sampled_s|1|beta|0.0|0.5|2.0|6.0|-1|1"
+  "p50_beta2_2_0_0p5_outwdm1|0.50|sampled_s|1|beta|0.0|0.5|2.0|2.0|-1|1"
+  "p50_unif0_0p125_outwd0p3|0.50|sampled_s|1|uniform|0.0|0.125|2.0|6.0|0.3|1"
+)
+
+if [[ -n "${SWEEP_CONFIGS:-}" ]]; then
+  IFS=$'\n' read -r -d '' -a CONFIGS < <(printf '%s\0' "${SWEEP_CONFIGS}") || true
+fi
+
+for entry in "${CONFIGS[@]}"; do
+  IFS='|' read -r config rollout_prob time_mode rollout_steps s_dist s_min_frac s_max_frac s_beta_alpha s_beta_beta output_wd sync_t rollout_steps_min <<<"${entry}"
+  run_config "${config}" "${rollout_prob}" "${time_mode}" "${rollout_steps}" "${s_dist}" "${s_min_frac}" "${s_max_frac}" "${s_beta_alpha}" "${s_beta_beta}" "${output_wd}" "${sync_t}" "${rollout_steps_min:-}"
+done
+
+echo "[ctx1024-sampleds] done" | tee -a "${DRIVER_LOG}"

LTA_openwebtext_dualt/scripts/run_train8_decode_algo_sweep_len256_latest.sh

@@ -0,0 +1,113 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+
+export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}"
+export TOKENIZERS_PARALLELISM=false
+export PYTHONUNBUFFERED=1
+
+DATA_DIR="${DATA_DIR:-/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext_lta_cached_chunks/gpt2_len256_train8_compact_overfit}"
+TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-standard/tokenizer.json}"
+MAX_LEN="${MAX_LEN:-256}"
+N_SAMPLES="${N_SAMPLES:-64}"
+INFER_STEPS="${INFER_STEPS:-128}"
+OUT_ROOT="${OUT_ROOT:-docs/lta_samples/metrics_20260517/train8_decode_algo_sweep_len256_latest}"
+LOG_FILE="${LOG_FILE:-logs/train8_decode_algo_sweep_len256_latest/sweep.log}"
+mkdir -p "$(dirname "${LOG_FILE}")" "${OUT_ROOT}"
+
+modes=("${MODES:-onehot allcorrupt}")
+variants=(
+  "flowmap|state|logit_normal|post|none|ode_trainlogit_post_state"
+  "flowmap|endpoint|logit_normal|post|none|ode_trainlogit_post_endpoint"
+  "flowmap|blend|logit_normal|post|none|ode_trainlogit_post_blend"
+  "flowmap|state|logit_normal|pre|none|ode_trainlogit_pre_state"
+  "flowmap|state|uniform|post|none|ode_uniform_post_state"
+  "flowmap|state|power_high|post|none|ode_powerhigh_post_state"
+  "log_geodesic|state|uniform|post|none|loggeo_uniform_post_state"
+  "log_geodesic|state|logit_normal|post|none|loggeo_trainlogit_post_state"
+  "flowmap|state|logit_normal|post|uniform|ode_trainlogit_post_state_softuniform"
+)
+
+for mode in ${modes[*]}; do
+  run_name="${RUN_PREFIX:-train8_n256_compactv969_3l_bs512_hard_ce_}${mode}"
+  if [[ ! -d "runs/${run_name}" ]]; then
+    echo "[decode-sweep] skip missing run=${run_name}" | tee -a "${LOG_FILE}"
+    continue
+  fi
+  for variant in "${variants[@]}"; do
+    IFS='|' read -r decode_rule final_from time_schedule model_t endpoint_softening tag <<<"${variant}"
+    out_dir="${OUT_ROOT}/${mode}/${tag}"
+    echo "[decode-sweep] $(date) mode=${mode} run=${run_name} tag=${tag}" | tee -a "${LOG_FILE}"
+    CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}" python scripts/eval_train8_decode_acc.py \
+      --runs_glob "runs/${run_name}" \
+      --data_dir "${DATA_DIR}" \
+      --tokenizer_path "${TOKENIZER_PATH}" \
+      --out_dir "${out_dir}" \
+      --max_len "${MAX_LEN}" \
+      --n_samples "${N_SAMPLES}" \
+      --batch_size "${N_SAMPLES}" \
+      --latest_only \
+      --endpoint_softenings "${endpoint_softening}" \
+      --steps "${INFER_STEPS}" \
+      --decode_rule "${decode_rule}" \
+      --time_schedule "${time_schedule}" \
+      --time_logit_mean -1.5 \
+      --time_logit_std 0.8 \
+      --model_t_mode "${model_t}" \
+      --c_min 1 \
+      --c_max 512 \
+      --late_temp 1.0 \
+      --final_from "${final_from}" \
+      --final_decode argmax \
+      2>&1 | tee -a "${LOG_FILE}"
+  done
+done
+
+python - "$OUT_ROOT" <<'PY' | tee -a "$LOG_FILE"
+import json
+import sys
+from pathlib import Path
+
+root = Path(sys.argv[1])
+rows = []
+for path in root.glob("*/*/decode_token_acc.jsonl"):
+    for line in path.read_text().splitlines():
+        if not line.strip():
+            continue
+        row = json.loads(line)
+        row["mode"] = path.parents[1].name
+        row["variant"] = path.parent.name
+        rows.append(row)
+rows.sort(key=lambda r: (r["mode"], -float(r.get("token_acc_mean", 0.0)), r["variant"]))
+out = root / "decode_algo_summary.tsv"
+with out.open("w", encoding="utf-8") as f:
+    f.write("mode\tvariant\tckpt_step\ttoken_acc_mean\texact_count\texact_ref_count\tdecode_rule\tfinal_from\ttime_schedule\tmodel_t_mode\tendpoint_softening\n")
+    for r in rows:
+        f.write(
+            "\t".join(
+                [
+                    str(r.get("mode")),
+                    str(r.get("variant")),
+                    str(r.get("ckpt_step")),
+                    f"{float(r.get('token_acc_mean', 0.0)):.6f}",
+                    str(r.get("exact_count")),
+                    str(r.get("exact_ref_count")),
+                    str(r.get("decode_rule")),
+                    str(r.get("final_from")),
+                    str(r.get("time_schedule")),
+                    str(r.get("model_t_mode")),
+                    str(r.get("endpoint_softening")),
+                ]
+            )
+            + "\n"
+        )
+print(out)
+for r in rows[:10]:
+    print(
+        "SWEEP_RESULT "
+        f"mode={r.get('mode')} variant={r.get('variant')} step={r.get('ckpt_step')} "
+        f"token_acc={float(r.get('token_acc_mean', 0.0)):.4f} "
+        f"exact={r.get('exact_count')}/{r.get('n_gen')} refs={r.get('exact_ref_count')}"
+    )
+PY

LTA_openwebtext_dualt/scripts/sweep_categorical_c1024_late_refresh_20260506.py

@@ -0,0 +1,281 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import json
+import math
+import sys
+from dataclasses import asdict, dataclass
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+BASE_PATH = Path(__file__).with_name("eval_lm1b_200k_methods_genppl_20260506.py")
+spec = importlib.util.spec_from_file_location("eval_lm1b_200k_methods_genppl_20260506", BASE_PATH)
+if spec is None or spec.loader is None:
+    raise RuntimeError(f"Could not load {BASE_PATH}")
+base = importlib.util.module_from_spec(spec)
+sys.modules[spec.name] = base
+spec.loader.exec_module(base)
+
+
+@dataclass(frozen=True)
+class RefreshCfg:
+    label: str
+    refresh_frac: float
+    sample_temp: float = 0.8
+    top_k: int = 16
+    select: str = "low_conf"
+
+
+def default_configs() -> list[RefreshCfg]:
+    return [
+        RefreshCfg("base_argmax", 0.0, top_k=1),
+        RefreshCfg("lowconf05_t0p8_top16", 0.05, 0.8, 16),
+        RefreshCfg("lowconf10_t0p8_top16", 0.10, 0.8, 16),
+        RefreshCfg("lowconf15_t0p8_top16", 0.15, 0.8, 16),
+        RefreshCfg("lowconf20_t0p8_top16", 0.20, 0.8, 16),
+        RefreshCfg("lowconf10_t1p0_top32", 0.10, 1.0, 32),
+        RefreshCfg("lowconf15_t1p0_top32", 0.15, 1.0, 32),
+        RefreshCfg("lowconf20_t1p0_top64", 0.20, 1.0, 64),
+        RefreshCfg("random10_t0p8_top16", 0.10, 0.8, 16, select="random"),
+        RefreshCfg("random15_t0p8_top16", 0.15, 0.8, 16, select="random"),
+    ]
+
+
+def load_endpoint_model(checkpoint: str, tokenizer, device: torch.device):
+    ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
+    model = base.build_endpoint_model(ckpt, tokenizer, "categorical_fullvocab", device)
+    return ckpt, model
+
+
+def sample_from_probs(probs: torch.Tensor, *, temp: float, top_k: int) -> torch.Tensor:
+    if top_k <= 1:
+        return probs.argmax(dim=-1)
+    k = min(int(top_k), probs.size(-1))
+    vals, inds = torch.topk(probs, k=k, dim=-1)
+    vals = vals.clamp_min(1e-20)
+    if temp != 1.0:
+        vals = vals.pow(1.0 / float(temp))
+    vals = vals / vals.sum(dim=-1, keepdim=True).clamp_min(1e-20)
+    picked = torch.multinomial(vals.reshape(-1, k), num_samples=1).view(*probs.shape[:-1])
+    return inds.gather(-1, picked.unsqueeze(-1)).squeeze(-1)
+
+
+def apply_refresh(final_probs: torch.Tensor, cfg: RefreshCfg) -> torch.Tensor:
+    ids = final_probs.argmax(dim=-1)
+    if cfg.refresh_frac <= 0:
+        return ids
+    bs, length, _ = final_probs.shape
+    n_refresh = max(1, int(math.ceil(float(cfg.refresh_frac) * length)))
+    conf = final_probs.max(dim=-1).values
+    if cfg.select == "low_conf":
+        order_score = -conf
+    elif cfg.select == "random":
+        order_score = torch.rand_like(conf)
+    else:
+        raise ValueError(f"unknown select={cfg.select!r}")
+    chosen = torch.topk(order_score, k=min(n_refresh, length), dim=-1).indices
+    sampled = sample_from_probs(final_probs, temp=cfg.sample_temp, top_k=cfg.top_k)
+    return ids.scatter(1, chosen, sampled.gather(1, chosen))
+
+
+@torch.inference_mode()
+def generate_variants(
+    model,
+    tokenizer,
+    configs: list[RefreshCfg],
+    *,
+    n_samples: int,
+    batch_size: int,
+    max_len: int,
+    steps: int,
+    seed: int,
+    device: torch.device,
+) -> tuple[dict[str, list[list[int]]], dict[str, list[str]], dict]:
+    torch.manual_seed(seed)
+    eps = 1e-8
+    concentration_min = 1.0
+    concentration_max = 1024.0
+    all_ids = {cfg.label: [] for cfg in configs}
+    all_texts = {cfg.label: [] for cfg in configs}
+    remaining = n_samples
+    while remaining > 0:
+        bs = min(batch_size, remaining)
+        probs = base.sample_noise_simplex(
+            (bs, max_len),
+            tokenizer.vocab_size,
+            device,
+            eps,
+            noise_mode="dirichlet",
+            target_prob=1.0,
+            noise_sigma=-1.0,
+            dirichlet_concentration=1.0,
+        )
+        attn = torch.ones((bs, max_len), dtype=torch.bool, device=device)
+        last_endpoint = probs
+        for step in range(steps):
+            model_t = torch.full((bs,), 0.5, dtype=torch.float32, device=device)
+            logits = model(base.state_for_model(model, probs, eps), model_t, attn).float() / 1.3
+            endpoint = F.softmax(logits, dim=-1)
+            last_endpoint = endpoint
+
+            support_t = (step + 1) / max(steps, 1)
+            semantic_t = support_t**1.5
+            anchor = probs.clamp_min(eps)
+            anchor = anchor / anchor.sum(dim=-1, keepdim=True).clamp_min(eps)
+            guided = (1.0 - semantic_t) * anchor + semantic_t * endpoint
+            guided = guided.clamp_min(eps)
+            guided = guided / guided.sum(dim=-1, keepdim=True).clamp_min(eps)
+
+            mean = (1.0 - support_t) / float(tokenizer.vocab_size) + support_t * guided
+            mean = mean.clamp_min(eps)
+            mean = mean / mean.sum(dim=-1, keepdim=True).clamp_min(eps)
+            conc = math.exp(
+                math.log(concentration_min)
+                + support_t * (math.log(concentration_max) - math.log(concentration_min))
+            )
+            alpha = (mean * conc).clamp_min(eps)
+            probs = torch._standard_gamma(alpha).clamp_min(eps)
+            probs = probs / probs.sum(dim=-1, keepdim=True).clamp_min(eps)
+
+        final_probs = 0.5 * probs + 0.5 * last_endpoint
+        final_probs = final_probs.clamp_min(eps)
+        final_probs = final_probs / final_probs.sum(dim=-1, keepdim=True).clamp_min(eps)
+        for cfg in configs:
+            ids = apply_refresh(final_probs, cfg).detach().cpu().tolist()
+            all_ids[cfg.label].extend(ids)
+            all_texts[cfg.label].extend(tokenizer.decode(row, stop_at_eos=False, skip_special_tokens=False) for row in ids)
+        remaining -= bs
+        print(f"[decode] generated base states {n_samples - remaining}/{n_samples}", flush=True)
+
+    base_decode = {
+        "kind": "categorical_fullvocab",
+        "steps": steps,
+        "model_t_mode": "const05",
+        "base_decode_rule": "dual_line_resample_rolling",
+        "support_power": 1.0,
+        "semantic_power": 1.5,
+        "endpoint_temp": 1.3,
+        "final_from": "blend_then_refresh",
+        "n_samples": n_samples,
+        "seed": seed,
+    }
+    return all_ids, all_texts, base_decode
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument(
+        "--checkpoint",
+        default="runs/lta_lm1b_dirichlet_categorical_fullvocab_dualt_flmpack_onehot_hardce_ddit_small_len128_gbs512_8gpu_1m_nw0/latest.pt",
+    )
+    p.add_argument("--tokenizer_path", required=True)
+    p.add_argument("--scorer", required=True)
+    p.add_argument("--out_dir", required=True)
+    p.add_argument("--n_samples", type=int, default=128)
+    p.add_argument("--max_len", type=int, default=128)
+    p.add_argument("--steps", type=int, default=1024)
+    p.add_argument("--decode_batch", type=int, default=16)
+    p.add_argument("--score_batch", type=int, default=8)
+    p.add_argument("--score_max_length", type=int, default=256)
+    p.add_argument("--seed", type=int, default=20260506)
+    p.add_argument("--only", default="", help="Optional comma-separated config labels.")
+    args = p.parse_args()
+
+    selected = {x.strip() for x in args.only.split(",") if x.strip()}
+    configs = [c for c in default_configs() if not selected or c.label in selected]
+    if not configs:
+        raise SystemExit(f"No configs selected by --only={args.only!r}")
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+    summary_path = out_dir / "summary.jsonl"
+    summary_path.write_text("", encoding="utf-8")
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    tokenizer = base.BpeTextTokenizer.from_file(args.tokenizer_path)
+    ckpt, model = load_endpoint_model(args.checkpoint, tokenizer, device)
+    step = ckpt.get("step")
+    print(f"[load] checkpoint={args.checkpoint} step={step}", flush=True)
+
+    scorer_tok = AutoTokenizer.from_pretrained(args.scorer)
+    if scorer_tok.pad_token_id is None:
+        scorer_tok.pad_token = scorer_tok.eos_token
+        scorer_tok.pad_token_id = scorer_tok.eos_token_id
+    scorer = AutoModelForCausalLM.from_pretrained(args.scorer).to(device).eval()
+    if getattr(scorer.config, "pad_token_id", None) is None:
+        scorer.config.pad_token_id = scorer_tok.pad_token_id
+
+    ids_by_label, texts_by_label, base_decode = generate_variants(
+        model,
+        tokenizer,
+        configs,
+        n_samples=args.n_samples,
+        batch_size=args.decode_batch,
+        max_len=args.max_len,
+        steps=args.steps,
+        seed=args.seed,
+        device=device,
+    )
+
+    summaries = []
+    with summary_path.open("a", encoding="utf-8") as sf:
+        for cfg in configs:
+            decode = dict(base_decode)
+            decode.update({"decode_rule": "late_low_conf_refresh", "refresh": asdict(cfg)})
+            target = base.EvalTarget(f"categorical_c1024_late_refresh_{cfg.label}", "categorical_fullvocab", args.checkpoint)
+            summary = base.score_and_write(
+                target=target,
+                checkpoint=args.checkpoint,
+                step=step,
+                decode=decode,
+                ids=ids_by_label[cfg.label],
+                raw_texts=texts_by_label[cfg.label],
+                scorer=scorer,
+                scorer_tok=scorer_tok,
+                score_batch=args.score_batch,
+                score_max_length=args.score_max_length,
+                device=device,
+                out_dir=out_dir,
+            )
+            summary["config"] = asdict(cfg)
+            sf.write(json.dumps(summary, ensure_ascii=False) + "\n")
+            sf.flush()
+            summaries.append(summary)
+            div = summary["diversity"]
+            print(
+                f"[summary] {cfg.label} raw={summary['raw_genppl']['ppl']:.3f} "
+                f"strip={summary['stripped_genppl']['ppl']:.3f} "
+                f"ent={div['sample_entropy']:.3f} d2={div['distinct_2']:.3f} "
+                f"top={div['top_token_mass']:.3f}",
+                flush=True,
+            )
+
+    table_path = out_dir / "summary.tsv"
+    with table_path.open("w", encoding="utf-8") as f:
+        f.write(
+            "label\traw_genppl\tstripped_genppl\tsample_entropy\tdistinct_2\ttop_token_mass\t"
+            "refresh_frac\tsample_temp\ttop_k\tselect\tmeets_target\n"
+        )
+        for summary in summaries:
+            div = summary["diversity"]
+            raw = summary["raw_genppl"]["ppl"]
+            ent = div["sample_entropy"]
+            cfg = summary["config"]
+            meets = raw <= 30.0 and ent >= 4.1
+            f.write(
+                f"{cfg['label']}\t{raw:.6f}\t{summary['stripped_genppl']['ppl']:.6f}\t"
+                f"{ent:.6f}\t{div['distinct_2']:.6f}\t{div['top_token_mass']:.6f}\t"
+                f"{cfg['refresh_frac']:.3f}\t{cfg['sample_temp']:.3f}\t{cfg['top_k']}\t"
+                f"{cfg['select']}\t{int(meets)}\n"
+            )
+    print(f"[done] {out_dir}", flush=True)
+
+
+if __name__ == "__main__":
+    main()