Add files using upload-large-folder tool

LTA_openwebtext_dualt/scripts/eval_ar_8gpu_latest_temp_sweep_20260506.py

@@ -0,0 +1,28 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import importlib.util
+import sys
+from pathlib import Path
+
+
+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)
+
+
+base.TARGETS = [
+    base.EvalTarget(
+        "ar_8gpu_scratch_latest",
+        "ar",
+        "runs/ar_lm1b_flmpack_bert_small_len128_gbs512_8gpu_1m_scratch_20260505/latest.pt",
+    )
+]
+
+
+if __name__ == "__main__":
+    base.main()

LTA_openwebtext_dualt/scripts/eval_dirichlet_latest_key3_20260508.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import importlib.util
+import sys
+from pathlib import Path
+
+
+BASE = Path(__file__).with_name("eval_c1024_decode_sweep_20260507.py")
+spec = importlib.util.spec_from_file_location("eval_c1024_decode_sweep_20260507", BASE)
+if spec is None or spec.loader is None:
+    raise RuntimeError(f"cannot import {BASE}")
+base = importlib.util.module_from_spec(spec)
+sys.modules[spec.name] = base
+spec.loader.exec_module(base)
+
+
+def key_configs() -> list[base.DecodeConfig]:
+    return [
+        base.DecodeConfig(
+            "match_post_sem1_blend_c16_t1p3",
+            "post",
+            1.0,
+            1.0,
+            "blend",
+            endpoint_temp=1.3,
+            concentration_max=16.0,
+        ),
+        base.DecodeConfig(
+            "match_post_sem1_blend_c64_t1p3",
+            "post",
+            1.0,
+            1.0,
+            "blend",
+            endpoint_temp=1.3,
+            concentration_max=64.0,
+        ),
+        base.DecodeConfig(
+            "match_post_sem1_blend_c1024_t1p3",
+            "post",
+            1.0,
+            1.0,
+            "blend",
+            endpoint_temp=1.3,
+            concentration_max=1024.0,
+        ),
+    ]
+
+
+base.default_configs = key_configs
+base.main()

LTA_openwebtext_dualt/scripts/eval_lm1b_c1024_fullycoupled_8gpu_1m_sde.sh

@@ -0,0 +1,39 @@
+#!/usr/bin/env bash
+# SDE (Dirichlet resample) GenPPL eval for:
+#   lta_lm1b_dirichlet_categorical_fullvocab_c1024_fullycoupled_flmpack_onehot_hardce_ddit_small_len128_gbs512_8gpu_1m_nw0
+set -euo pipefail
+
+ROOT="/e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt"
+cd "${ROOT}"
+export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}"
+export TOKENIZERS_PARALLELISM=false
+
+RUN_DIR="${RUN_DIR:-runs/lta_lm1b_dirichlet_categorical_fullvocab_c1024_fullycoupled_flmpack_onehot_hardce_ddit_small_len128_gbs512_8gpu_1m_nw0}"
+CHECKPOINT="${CHECKPOINT:-${RUN_DIR}/latest.pt}"
+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}"
+
+STEPS="${STEPS:-1024}"
+N_SAMPLES="${N_SAMPLES:-128}"
+MAX_LEN="${MAX_LEN:-128}"
+BATCH_SIZE="${BATCH_SIZE:-16}"
+SCORE_BATCH="${SCORE_BATCH:-8}"
+
+OUT_DIR="${OUT_DIR:-${RUN_DIR}/eval_sde_steps${STEPS}_n${N_SAMPLES}}"
+
+CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0}"
+
+echo "[eval-sde] checkpoint=${CHECKPOINT}"
+echo "[eval-sde] steps=${STEPS} n_samples=${N_SAMPLES} max_len=${MAX_LEN}"
+echo "[eval-sde] out_dir=${OUT_DIR}"
+
+CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES}" python scripts/eval_lm1b_c1024_fullycoupled_sde_genppl.py \
+  --checkpoint "${CHECKPOINT}" \
+  --tokenizer_path "${TOKENIZER_PATH}" \
+  --scorer "${SCORER}" \
+  --out_dir "${OUT_DIR}" \
+  --steps "${STEPS}" \
+  --n_samples "${N_SAMPLES}" \
+  --max_len "${MAX_LEN}" \
+  --batch_size "${BATCH_SIZE}" \
+  --score_batch "${SCORE_BATCH}"

LTA_openwebtext_dualt/scripts/extract_mauve_features_lm1b.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+import json
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+
+IN_PATH = Path("runs/decode_lab/mauve_export_lm1b_latest_1024.jsonl")
+OUT_PATH = Path("runs/decode_lab/mauve_features_lm1b_latest_1024_gpt2large.npz")
+MODEL_PATH = "/e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard"
+MAX_LEN = 256
+BATCH_SIZE = 16
+
+
+def load_texts() -> tuple[list[str], list[str], dict[str, list[str]], dict[str, list[str]]]:
+    refs_raw: list[str] = []
+    refs_stripped: list[str] = []
+    gen_raw: dict[str, list[str]] = {}
+    gen_stripped: dict[str, list[str]] = {}
+    with IN_PATH.open(encoding="utf-8") as f:
+        for line in f:
+            obj = json.loads(line)
+            typ = obj.get("type")
+            if typ == "reference":
+                refs_raw.append(obj["raw_text"])
+                refs_stripped.append(obj["stripped_text"])
+            elif typ == "generated":
+                setting = obj["setting"]
+                key = f"t{float(setting['endpoint_temp']):.2f}_{setting['final_from']}"
+                gen_raw.setdefault(key, []).append(obj["raw_text"])
+                gen_stripped.setdefault(key, []).append(obj["stripped_text"])
+    return refs_raw, refs_stripped, gen_raw, gen_stripped
+
+
+@torch.no_grad()
+def featurize(
+    texts: list[str],
+    name: str,
+    tokenizer: AutoTokenizer,
+    model: AutoModel,
+    device: torch.device,
+) -> np.ndarray:
+    chunks: list[np.ndarray] = []
+    t0 = time.time()
+    for start in range(0, len(texts), BATCH_SIZE):
+        batch = texts[start : start + BATCH_SIZE]
+        enc = tokenizer(
+            batch,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=MAX_LEN,
+            return_attention_mask=True,
+        ).to(device)
+        out = model(
+            input_ids=enc["input_ids"],
+            attention_mask=enc["attention_mask"],
+            output_hidden_states=True,
+            return_dict=True,
+        )
+        hidden = out.hidden_states[-1]
+        last_idx = enc["attention_mask"].sum(dim=1) - 1
+        feat = hidden[torch.arange(hidden.size(0), device=device), last_idx]
+        chunks.append(feat.float().cpu().numpy())
+        if (start // BATCH_SIZE) % 10 == 0:
+            print(f"{name} {start + len(batch)}/{len(texts)}", flush=True)
+    arr = np.concatenate(chunks, axis=0)
+    print(f"{name} {arr.shape} time={time.time() - t0:.1f}s", flush=True)
+    return arr
+
+
+def main() -> None:
+    refs_raw, refs_stripped, gen_raw, gen_stripped = load_texts()
+    print(
+        "loaded",
+        len(refs_raw),
+        {key: len(value) for key, value in gen_raw.items()},
+        flush=True,
+    )
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+    model = AutoModel.from_pretrained(MODEL_PATH, pad_token_id=tokenizer.eos_token_id).to(device).eval()
+
+    arrays: dict[str, np.ndarray] = {
+        "ref_raw": featurize(refs_raw, "ref_raw", tokenizer, model, device),
+        "ref_stripped": featurize(refs_stripped, "ref_stripped", tokenizer, model, device),
+    }
+    for key in sorted(gen_raw):
+        arrays[f"gen_{key}_raw"] = featurize(gen_raw[key], f"gen_{key}_raw", tokenizer, model, device)
+        arrays[f"gen_{key}_stripped"] = featurize(
+            gen_stripped[key],
+            f"gen_{key}_stripped",
+            tokenizer,
+            model,
+            device,
+        )
+    OUT_PATH.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(OUT_PATH, **arrays)
+    print(f"DONE {OUT_PATH}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

LTA_openwebtext_dualt/scripts/launch_lta_lm1b_fullycoupled_t5_logitnormal_8gpu.sh

@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+
+exec bash scripts/launch_lta_lm1b_elfaligned_t5_logitnormal_8gpu.sh "$@"

LTA_openwebtext_dualt/scripts/launch_lta_owt_compact_gpt2bpe_v2048_stream1024_fullycoupled_mask1_wd0p1_fp32_8gpu.sh

@@ -0,0 +1,250 @@
+#!/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}"
+
+# Compact GPT2-BPE v2048, rebuilt as tokenizer-stream chunks:
+# every example is length 1024 except the final tail record in each part.
+COMPACT_VARIANT_LABEL="${COMPACT_VARIANT_LABEL:-compact_gpt2bpe_v2048_stream1024_fullycoupled_mask1_wd0p1_fp32}"
+VOCAB_SIZE="${VOCAB_SIZE:-2048}"
+DATA_PATH="${DATA_PATH:-/e2e-data/evad-tech-vla/wanghan58/data/embedded-language-flows/openwebtext-compact-gpt2bpe-v2048-stream1024-train-minus-100k}"
+TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/models/lta_tokenizers/owt_compact_gpt2bpe_v2048/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:-32249}"
+
+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}"
+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:-}"
+
+# Fully-coupled fp32 baseline recipe: AdamW, nanoGPT decay groups, RMSNorm,
+# no output bias, hard CE, Dirichlet bridge, full-vocab wrong token.
+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}"
+
+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:--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}"
+
+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_compact_gpt2bpe_v2048_stream1024_fullycoupled_rmsnorm_nobias_adamw_wd0p1_logitnormal_${T_LOGIT_MEAN_LABEL}_s${T_LOGIT_STD_LABEL}_hardce_mask${MIN_MASK_RATIO_LABEL}-${MAX_MASK_RATIO_LABEL}_fp32_ddit768x12_gbs512_8gpu_1m_$(date +%Y%m%d_%H%M%S)}"
+SAVE_DIR="${SAVE_DIR:-runs/${RUN_NAME}}"
+LOG_DIR="${LOG_DIR:-logs/compact_gpt2bpe_v2048_stream1024_fullycoupled_mask1_wd0p1_fp32_8gpu}"
+LOG_FILE="${LOG_FILE:-${LOG_DIR}/${RUN_NAME}.log}"
+
+if [[ -n "${RESUME_PATH}" ]]; then
+  ALLOW_EXISTING_SAVE_DIR=1
+  if [[ ! -f "${RESUME_PATH}" ]]; then
+    echo "Missing resume checkpoint: ${RESUME_PATH}" >&2
+    exit 2
+  fi
+fi
+
+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_owt_compact_gpt2bpe_stream1024_train_minus_100k_np8.sh" >&2
+  exit 2
+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}")
+for name in ("elf_multi_part_meta.json", "elf_build_meta.json"):
+    path = root / name
+    if path.exists():
+        print(int(json.loads(path.read_text()).get("num_examples", 0)))
+        raise SystemExit
+parts_root = root / "parts"
+if parts_root.is_dir():
+    print(sum(len(load_from_disk(str(p))) for p in sorted(parts_root.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) ))
+SAMPLES_PER_RANK=$(( (NUM_EXAMPLES + WORLD_SIZE - 1) / WORLD_SIZE ))
+LOADER_BATCHES_PER_RANK=$(( SAMPLES_PER_RANK / PER_GPU_BATCH_SIZE ))
+STEPS_PER_EPOCH=$(( (LOADER_BATCHES_PER_RANK + GRAD_ACCUM - 1) / GRAD_ACCUM ))
+SAVE_EVERY="${SAVE_EVERY:-${STEPS_PER_EPOCH}}"
+
+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=${COMPACT_VARIANT_LABEL} 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] tokenizer=${TOKENIZER_PATH}"
+echo "[launch] examples=${NUM_EXAMPLES} total_steps=${TOTAL_STEPS} world_size=${WORLD_SIZE} grad_accum=${GRAD_ACCUM} steps_per_epoch=${STEPS_PER_EPOCH} save_every=${SAVE_EVERY}"
+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 batch=${GLOBAL_BATCH_SIZE} per_gpu=${PER_GPU_BATCH_SIZE}"
+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} mask_ratio=${MIN_MASK_RATIO}->${MAX_MASK_RATIO}"
+if [[ -n "${RESUME_PATH}" ]]; then
+  echo "[launch] resume_path=${RESUME_PATH}"
+fi
+
+if [[ "${DRY_RUN}" == "1" || "${DRY_RUN}" == "true" || "${DRY_RUN}" == "True" ]]; then
+  echo "[launch] DRY_RUN=1, validated tokenizer/data/schedule; 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 1024 \
+  --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 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_fullycoupled_elftsampling_mask1_wd0p1_fp32_8gpu.sh

@@ -0,0 +1,202 @@
+#!/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:--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}"
+
+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_adamw_wd0p1_elftsampling_${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_elftsampling_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 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 \
+  --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_fullycoupled_uniform_mask1_swiglu_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:-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:--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_uniformt_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_uniform_mask1_swiglu_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_gpt2_softendpoint_mn_pilot_4gpu.sh

@@ -0,0 +1,271 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+
+export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}"
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3}"
+export OMP_NUM_THREADS="${OMP_NUM_THREADS:-1}"
+export TOKENIZERS_PARALLELISM=false
+export PYTHONUNBUFFERED=1
+export NCCL_DEBUG="${NCCL_DEBUG:-WARN}"
+export TORCH_DISTRIBUTED_TIMEOUT="${TORCH_DISTRIBUTED_TIMEOUT:-3600}"
+export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
+
+STAMP="$(date +%Y%m%d_%H%M%S)"
+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}"
+OWT_CHUNK_CACHE_DIR="${OWT_CHUNK_CACHE_DIR:-/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext_lta_cached_chunks/gpt2_len1024_train_minus_100k}"
+OWT_EXACT_REPEAT_PER_CHUNK="${OWT_EXACT_REPEAT_PER_CHUNK:-0}"
+
+NPROC_PER_NODE="${NPROC_PER_NODE:-4}"
+MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
+MASTER_PORT="${MASTER_PORT:-32251}"
+
+MAX_LEN="${MAX_LEN:-1024}"
+LOWK_CLEAN_TOKENS="${LOWK_CLEAN_TOKENS:-64,128,256}"
+CLEAN_STATE_MODE="${CLEAN_STATE_MODE:-onehot}"
+MASK_MIXTURE_LOWK_PROB="${MASK_MIXTURE_LOWK_PROB:-1.0}"
+MASK_MIXTURE_ALL_PROB="${MASK_MIXTURE_ALL_PROB:-0.0}"
+VOCAB_SIZE_OVERRIDE="${VOCAB_SIZE_OVERRIDE:-0}"
+D_MODEL="${D_MODEL:-768}"
+COND_DIM="${COND_DIM:-128}"
+N_LAYERS="${N_LAYERS:-12}"
+N_HEADS="${N_HEADS:-12}"
+DIM_FF="${DIM_FF:-3072}"
+TOTAL_STEPS="${TOTAL_STEPS:-500}"
+PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-32}"
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+NUM_WORKERS="${NUM_WORKERS:-4}"
+DATALOADER_PREFETCH_FACTOR="${DATALOADER_PREFETCH_FACTOR:-4}"
+MUON_IMPL="${MUON_IMPL:-legacy}"
+OUTPUT_WEIGHT_DECAY="${OUTPUT_WEIGHT_DECAY:--1}"
+MIN_MASK_RATIO="${MIN_MASK_RATIO:-0.1}"
+MAX_MASK_RATIO="${MAX_MASK_RATIO:-1.0}"
+MASK_RATIO_FLOOR_SCHEDULE="${MASK_RATIO_FLOOR_SCHEDULE:-none}"
+DIRICHLET_CONCENTRATION_MIN="${DIRICHLET_CONCENTRATION_MIN:-1.0}"
+DIRICHLET_CONCENTRATION_MAX="${DIRICHLET_CONCENTRATION_MAX:-1024}"
+DIRICHLET_ENDPOINT_MODE="${DIRICHLET_ENDPOINT_MODE:-categorical_dual_t}"
+DIRICHLET_SEMANTIC_T_MODE="${DIRICHLET_SEMANTIC_T_MODE:-same}"
+DIRICHLET_SEMANTIC_T_VALUE="${DIRICHLET_SEMANTIC_T_VALUE:-0.0}"
+SIMPLEX_BRIDGE_SAMPLER="${SIMPLEX_BRIDGE_SAMPLER:-dirichlet}"
+ENDPOINT_SEQUENCE_RANDOM_PROB_ALPHA="${ENDPOINT_SEQUENCE_RANDOM_PROB_ALPHA:-0.0}"
+CATEGORICAL_WRONG_BASIN_TOKEN_IDS="${CATEGORICAL_WRONG_BASIN_TOKEN_IDS:-}"
+CATEGORICAL_WRONG_BASIN_PROB="${CATEGORICAL_WRONG_BASIN_PROB:-0.0}"
+CATEGORICAL_WRONG_UNIGRAM_PROB="${CATEGORICAL_WRONG_UNIGRAM_PROB:-0.0}"
+CATEGORICAL_WRONG_UNIFORM_PROB="${CATEGORICAL_WRONG_UNIFORM_PROB:-0.0}"
+CATEGORICAL_WRONG_PROB_FLOOR="${CATEGORICAL_WRONG_PROB_FLOOR:-0.0}"
+CATEGORICAL_WRONG_BASIN_SHARED_PROB="${CATEGORICAL_WRONG_BASIN_SHARED_PROB:-0.0}"
+CATEGORICAL_WRONG_UNIGRAM_SHARED_PROB="${CATEGORICAL_WRONG_UNIGRAM_SHARED_PROB:-0.0}"
+LOGISTIC_NORMAL_SIGMA_MIN="${LOGISTIC_NORMAL_SIGMA_MIN:-0.1}"
+LOGISTIC_NORMAL_SIGMA_MAX="${LOGISTIC_NORMAL_SIGMA_MAX:-1.0}"
+LOGISTIC_NORMAL_TAU_MIN="${LOGISTIC_NORMAL_TAU_MIN:-1.0}"
+LOGISTIC_NORMAL_TAU_MAX="${LOGISTIC_NORMAL_TAU_MAX:-1.0}"
+ROLLOUT_TRAIN_PROB="${ROLLOUT_TRAIN_PROB:-0.0}"
+ROLLOUT_TRAIN_STEPS="${ROLLOUT_TRAIN_STEPS:-1}"
+ROLLOUT_TRAIN_STEPS_MIN="${ROLLOUT_TRAIN_STEPS_MIN:--1}"
+ROLLOUT_TRAIN_INFER_STEPS="${ROLLOUT_TRAIN_INFER_STEPS:-64}"
+ROLLOUT_TRAIN_TIME_MODE="${ROLLOUT_TRAIN_TIME_MODE:-fixed_steps}"
+ROLLOUT_TRAIN_S_DIST="${ROLLOUT_TRAIN_S_DIST:-uniform}"
+ROLLOUT_TRAIN_S_MIN_FRAC="${ROLLOUT_TRAIN_S_MIN_FRAC:-0.0}"
+ROLLOUT_TRAIN_S_MAX_FRAC="${ROLLOUT_TRAIN_S_MAX_FRAC:-0.125}"
+ROLLOUT_TRAIN_S_BETA_ALPHA="${ROLLOUT_TRAIN_S_BETA_ALPHA:-2.0}"
+ROLLOUT_TRAIN_S_BETA_BETA="${ROLLOUT_TRAIN_S_BETA_BETA:-6.0}"
+ROLLOUT_TRAIN_TEMP="${ROLLOUT_TRAIN_TEMP:-1.0}"
+ROLLOUT_TRAIN_MAX_GAMMA="${ROLLOUT_TRAIN_MAX_GAMMA:-1.0}"
+ROLLOUT_TRAIN_CORRUPT_ONLY="${ROLLOUT_TRAIN_CORRUPT_ONLY:-1}"
+ROLLOUT_TRAIN_SAMPLEWISE="${ROLLOUT_TRAIN_SAMPLEWISE:-0}"
+ROLLOUT_TRAIN_SELECTED_ONLY="${ROLLOUT_TRAIN_SELECTED_ONLY:-0}"
+ROLLOUT_TRAIN_COMPUTE_ALWAYS="${ROLLOUT_TRAIN_COMPUTE_ALWAYS:-0}"
+ROLLOUT_TRAIN_SYNC_T="${ROLLOUT_TRAIN_SYNC_T:-0}"
+
+TARGET_LOSS="${TARGET_LOSS:-linear_soft_kl}"
+LINEAR_SOFT_TARGET_POWER="${LINEAR_SOFT_TARGET_POWER:-1.0}"
+LINEAR_SOFT_TARGET_MIN_CONF="${LINEAR_SOFT_TARGET_MIN_CONF:-0.0}"
+LINEAR_SOFT_TARGET_MAX_CONF="${LINEAR_SOFT_TARGET_MAX_CONF:-1.0}"
+
+tag() {
+  printf "%s" "$1" | sed -e 's/,/_/g' -e 's/-/m/g' -e 's/\./p/g'
+}
+
+RUN_NAME="${RUN_NAME:-pilot_gpt2_softendpoint_mn_n${MAX_LEN}_m$(tag "${LOWK_CLEAN_TOKENS}")_${CLEAN_STATE_MODE}_${TOTAL_STEPS}_${STAMP}}"
+SAVE_DIR="${SAVE_DIR:-runs/${RUN_NAME}}"
+LOG_FILE="${LOG_FILE:-logs/softendpoint_mn_pilot_4gpu/${RUN_NAME}.log}"
+RESUME_PATH="${RESUME_PATH:-}"
+INIT_MODEL_PATH="${INIT_MODEL_PATH:-}"
+mkdir -p "$(dirname "${LOG_FILE}")" "${SAVE_DIR}"
+
+echo "[launch] gpt2 cached OWT soft-endpoint m/n pilot"
+echo "[launch] run_name=${RUN_NAME}"
+echo "[launch] save_dir=${SAVE_DIR}"
+echo "[launch] n=${MAX_LEN} m=${LOWK_CLEAN_TOKENS} clean_state_mode=${CLEAN_STATE_MODE}"
+echo "[launch] mask_mixture lowk=${MASK_MIXTURE_LOWK_PROB} all=${MASK_MIXTURE_ALL_PROB}"
+echo "[launch] model d=${D_MODEL} layers=${N_LAYERS} heads=${N_HEADS} ff=${DIM_FF} vocab_override=${VOCAB_SIZE_OVERRIDE}"
+echo "[launch] optimizer=muon muon_impl=${MUON_IMPL} weight_decay=${WEIGHT_DECAY:-0.1} output_weight_decay=${OUTPUT_WEIGHT_DECAY}"
+echo "[launch] target_loss=${TARGET_LOSS} conf=${LINEAR_SOFT_TARGET_MIN_CONF}->${LINEAR_SOFT_TARGET_MAX_CONF} power=${LINEAR_SOFT_TARGET_POWER}"
+echo "[launch] mask_ratio=${MIN_MASK_RATIO}->${MAX_MASK_RATIO}"
+echo "[launch] mask_ratio_floor_schedule=${MASK_RATIO_FLOOR_SCHEDULE}"
+echo "[launch] dirichlet C=${DIRICHLET_CONCENTRATION_MIN}->${DIRICHLET_CONCENTRATION_MAX} endpoint=${DIRICHLET_ENDPOINT_MODE} sampler=${SIMPLEX_BRIDGE_SAMPLER}"
+echo "[launch] wrong_mix seq_alpha=${ENDPOINT_SEQUENCE_RANDOM_PROB_ALPHA} wrong_floor=${CATEGORICAL_WRONG_PROB_FLOOR} unigram=${CATEGORICAL_WRONG_UNIGRAM_PROB} uniform=${CATEGORICAL_WRONG_UNIFORM_PROB} basin=${CATEGORICAL_WRONG_BASIN_PROB} basin_ids=${CATEGORICAL_WRONG_BASIN_TOKEN_IDS}"
+echo "[launch] rollout_train prob=${ROLLOUT_TRAIN_PROB} mode=${ROLLOUT_TRAIN_TIME_MODE} steps=${ROLLOUT_TRAIN_STEPS} steps_min=${ROLLOUT_TRAIN_STEPS_MIN} infer_steps=${ROLLOUT_TRAIN_INFER_STEPS} s_dist=${ROLLOUT_TRAIN_S_DIST} s_frac=${ROLLOUT_TRAIN_S_MIN_FRAC}->${ROLLOUT_TRAIN_S_MAX_FRAC} temp=${ROLLOUT_TRAIN_TEMP} corrupt_only=${ROLLOUT_TRAIN_CORRUPT_ONLY} samplewise=${ROLLOUT_TRAIN_SAMPLEWISE} selected_only=${ROLLOUT_TRAIN_SELECTED_ONLY} sync_t=${ROLLOUT_TRAIN_SYNC_T}"
+echo "[launch] cache=${OWT_CHUNK_CACHE_DIR} exact_repeat_per_chunk=${OWT_EXACT_REPEAT_PER_CHUNK}"
+if [[ -n "${RESUME_PATH}" ]]; then
+  echo "[launch] resume_path=${RESUME_PATH}"
+fi
+if [[ -n "${INIT_MODEL_PATH}" ]]; then
+  echo "[launch] init_model_path=${INIT_MODEL_PATH}"
+fi
+
+resume_args=()
+if [[ -n "${RESUME_PATH}" ]]; then
+  resume_args=(--resume_path "${RESUME_PATH}")
+fi
+init_args=()
+if [[ -n "${INIT_MODEL_PATH}" ]]; then
+  init_args=(--init_model_path "${INIT_MODEL_PATH}")
+fi
+rollout_bool_args=()
+if [[ "${ROLLOUT_TRAIN_CORRUPT_ONLY}" == "0" || "${ROLLOUT_TRAIN_CORRUPT_ONLY}" == "false" || "${ROLLOUT_TRAIN_CORRUPT_ONLY}" == "False" ]]; then
+  rollout_bool_args+=(--no-rollout_train_corrupt_only)
+else
+  rollout_bool_args+=(--rollout_train_corrupt_only)
+fi
+if [[ "${ROLLOUT_TRAIN_SAMPLEWISE}" == "1" || "${ROLLOUT_TRAIN_SAMPLEWISE}" == "true" || "${ROLLOUT_TRAIN_SAMPLEWISE}" == "True" ]]; then
+  rollout_bool_args+=(--rollout_train_samplewise)
+fi
+if [[ "${ROLLOUT_TRAIN_SELECTED_ONLY}" == "1" || "${ROLLOUT_TRAIN_SELECTED_ONLY}" == "true" || "${ROLLOUT_TRAIN_SELECTED_ONLY}" == "True" ]]; then
+  rollout_bool_args+=(--rollout_train_selected_only)
+fi
+if [[ "${ROLLOUT_TRAIN_COMPUTE_ALWAYS}" == "1" || "${ROLLOUT_TRAIN_COMPUTE_ALWAYS}" == "true" || "${ROLLOUT_TRAIN_COMPUTE_ALWAYS}" == "True" ]]; then
+  rollout_bool_args+=(--rollout_train_compute_always)
+fi
+if [[ "${ROLLOUT_TRAIN_SYNC_T}" == "1" || "${ROLLOUT_TRAIN_SYNC_T}" == "true" || "${ROLLOUT_TRAIN_SYNC_T}" == "True" ]]; then
+  rollout_bool_args+=(--rollout_train_sync_t)
+fi
+
+python -m torch.distributed.run \
+  --nnodes=1 \
+  --nproc_per_node="${NPROC_PER_NODE}" \
+  --node_rank=0 \
+  --master_addr="${MASTER_ADDR}" \
+  --master_port="${MASTER_PORT}" \
+  train.py \
+  --data_path "${DATA_PATH}" \
+  --text_column text \
+  --openwebtext_split train_minus_100k \
+  --detokenizer auto \
+  --tokenizer_path "${TOKENIZER_PATH}" \
+  --save_dir "${SAVE_DIR}" \
+  "${resume_args[@]}" \
+  "${init_args[@]}" \
+  --wrap \
+  --wrap_mode stream \
+  --wrap_record_buffer_size 200 \
+  --owt_cached_chunks \
+  --owt_chunk_cache_dir "${OWT_CHUNK_CACHE_DIR}" \
+  --owt_chunk_cache_write_batch 4096 \
+  --owt_exact_repeat_per_chunk "${OWT_EXACT_REPEAT_PER_CHUNK}" \
+  --max_len "${MAX_LEN}" \
+  --batch_size "${PER_GPU_BATCH_SIZE}" \
+  --num_workers "${NUM_WORKERS}" \
+  --dataloader_prefetch_factor "${DATALOADER_PREFETCH_FACTOR}" \
+  --global_batch_size "${GLOBAL_BATCH_SIZE}" \
+  --total_steps "${TOTAL_STEPS}" \
+  --log_every "${LOG_EVERY:-50}" \
+  --eval_every 0 \
+  --save_every "${SAVE_EVERY:-500}" \
+  --latest_every "${LATEST_EVERY:-250}" \
+  --lr "${LEARNING_RATE:-0.002}" \
+  --weight_decay "${WEIGHT_DECAY:-0.1}" \
+  --output_weight_decay "${OUTPUT_WEIGHT_DECAY}" \
+  --adam_beta1 0.9 \
+  --adam_beta2 0.95 \
+  --adam_eps 1e-8 \
+  --optimizer muon \
+  --muon_impl "${MUON_IMPL}" \
+  --muon_momentum 0.95 \
+  --muon_ns_steps 5 \
+  --muon_update_scale 1.0 \
+  --ema_decay 0.9999 \
+  --ema_start_step 0 \
+  --warmup_steps "${WARMUP_STEPS:-100}" \
+  --lr_schedule constant_warmup \
+  --min_lr 0 \
+  --grad_clip 1.0 \
+  --adamw_param_groups nanogpt \
+  --seed 123 \
+  --d_model "${D_MODEL}" \
+  --cond_dim "${COND_DIM}" \
+  --n_layers "${N_LAYERS}" \
+  --n_heads "${N_HEADS}" \
+  --dim_ff "${DIM_FF}" \
+  --vocab_size_override "${VOCAB_SIZE_OVERRIDE}" \
+  --dropout 0.0 \
+  --model_type ddit \
+  --state_format prob \
+  --bridge dirichlet \
+  --target_loss "${TARGET_LOSS}" \
+  --linear_soft_target_power "${LINEAR_SOFT_TARGET_POWER}" \
+  --linear_soft_target_min_conf "${LINEAR_SOFT_TARGET_MIN_CONF}" \
+  --linear_soft_target_max_conf "${LINEAR_SOFT_TARGET_MAX_CONF}" \
+  --rollout_train_prob "${ROLLOUT_TRAIN_PROB}" \
+  --rollout_train_steps "${ROLLOUT_TRAIN_STEPS}" \
+  --rollout_train_steps_min "${ROLLOUT_TRAIN_STEPS_MIN}" \
+  --rollout_train_infer_steps "${ROLLOUT_TRAIN_INFER_STEPS}" \
+  --rollout_train_time_mode "${ROLLOUT_TRAIN_TIME_MODE}" \
+  --rollout_train_s_dist "${ROLLOUT_TRAIN_S_DIST}" \
+  --rollout_train_s_min_frac "${ROLLOUT_TRAIN_S_MIN_FRAC}" \
+  --rollout_train_s_max_frac "${ROLLOUT_TRAIN_S_MAX_FRAC}" \
+  --rollout_train_s_beta_alpha "${ROLLOUT_TRAIN_S_BETA_ALPHA}" \
+  --rollout_train_s_beta_beta "${ROLLOUT_TRAIN_S_BETA_BETA}" \
+  --rollout_train_temp "${ROLLOUT_TRAIN_TEMP}" \
+  --rollout_train_max_gamma "${ROLLOUT_TRAIN_MAX_GAMMA}" \
+  "${rollout_bool_args[@]}" \
+  --target_prob 1.0 \
+  --min_t 0.0 \
+  --max_t 1.0 \
+  --t_sampling_mode "${T_SAMPLING_MODE:-logit_normal}" \
+  --t_sampling_logit_mean "${T_LOGIT_MEAN:--1.5}" \
+  --t_sampling_logit_std "${T_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 "${MIN_MASK_RATIO}" \
+  --max_mask_ratio "${MAX_MASK_RATIO}" \
+  --mask_ratio_floor_schedule "${MASK_RATIO_FLOOR_SCHEDULE}" \
+  --mask_mixture_lowk_prob "${MASK_MIXTURE_LOWK_PROB}" \
+  --mask_mixture_all_prob "${MASK_MIXTURE_ALL_PROB}" \
+  --mask_mixture_lowk_clean_tokens "${LOWK_CLEAN_TOKENS}" \
+  --clean_state_mode "${CLEAN_STATE_MODE}" \
+  --wrong_token_replace_prob 1.0 \
+  --wrong_token_schedule linear_t \
+  --wrong_token_exp_k 1.0 \
+  --dirichlet_concentration_min "${DIRICHLET_CONCENTRATION_MIN}" \
+  --dirichlet_concentration_max "${DIRICHLET_CONCENTRATION_MAX}" \
+  --dirichlet_endpoint_mode "${DIRICHLET_ENDPOINT_MODE}" \
+  --dirichlet_semantic_t_mode "${DIRICHLET_SEMANTIC_T_MODE}" \
+  --dirichlet_semantic_t_value "${DIRICHLET_SEMANTIC_T_VALUE}" \
+  --endpoint_sequence_random_prob_alpha "${ENDPOINT_SEQUENCE_RANDOM_PROB_ALPHA}" \
+  --categorical_wrong_basin_token_ids "${CATEGORICAL_WRONG_BASIN_TOKEN_IDS}" \
+  --categorical_wrong_basin_prob "${CATEGORICAL_WRONG_BASIN_PROB}" \
+  --categorical_wrong_unigram_prob "${CATEGORICAL_WRONG_UNIGRAM_PROB}" \
+  --categorical_wrong_uniform_prob "${CATEGORICAL_WRONG_UNIFORM_PROB}" \
+  --categorical_wrong_prob_floor "${CATEGORICAL_WRONG_PROB_FLOOR}" \
+  --categorical_wrong_basin_shared_prob "${CATEGORICAL_WRONG_BASIN_SHARED_PROB}" \
+  --categorical_wrong_unigram_shared_prob "${CATEGORICAL_WRONG_UNIGRAM_SHARED_PROB}" \
+  --categorical_wrong_from_full_vocab \
+  --simplex_bridge_sampler "${SIMPLEX_BRIDGE_SAMPLER}" \
+  --logistic_normal_sigma_min "${LOGISTIC_NORMAL_SIGMA_MIN}" \
+  --logistic_normal_sigma_max "${LOGISTIC_NORMAL_SIGMA_MAX}" \
+  --logistic_normal_tau_min "${LOGISTIC_NORMAL_TAU_MIN}" \
+  --logistic_normal_tau_max "${LOGISTIC_NORMAL_TAU_MAX}" \
+  --eps 1e-8 \
+  --noise_init logistic_normal \
+  --bridge_noise_init logistic_normal \
+  --noise_sigma -1 \
+  --allow_tf32 \
+  --ddp_gradient_as_bucket_view \
+  --bf16 \
+  2>&1 | tee -a "${LOG_FILE}"

LTA_openwebtext_dualt/scripts/launch_lta_owt_selfcond_p05_4gpu.sh

@@ -0,0 +1,65 @@
+#!/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 TOKENIZERS_PARALLELISM=false
+
+STAMP="$(date +%Y%m%d_%H%M%S)"
+export RUN_NAME="${RUN_NAME:-lta_owt_gpt2cached_len1024_selfcond_p05_rollout1_samplewise_c1024_ddit768x12_muon_ema_gbs512_4gpu_50k_${STAMP}}"
+export SAVE_DIR="${SAVE_DIR:-runs/${RUN_NAME}}"
+export LOG_FILE="${LOG_FILE:-logs/selfcond_4gpu/${RUN_NAME}.log}"
+mkdir -p "$(dirname "${LOG_FILE}")"
+
+export NPROC_PER_NODE="${NPROC_PER_NODE:-4}"
+export MASTER_PORT="${MASTER_PORT:-31994}"
+export PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-32}"
+export GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+export TOTAL_STEPS="${TOTAL_STEPS:-50000}"
+export WARMUP_STEPS="${WARMUP_STEPS:-2000}"
+export MAX_LEN="${MAX_LEN:-1024}"
+export NUM_WORKERS="${NUM_WORKERS:-4}"
+export DATALOADER_PREFETCH_FACTOR="${DATALOADER_PREFETCH_FACTOR:-4}"
+export LOG_EVERY="${LOG_EVERY:-20}"
+export SAVE_EVERY="${SAVE_EVERY:-1000}"
+export LATEST_EVERY="${LATEST_EVERY:-500}"
+export EVAL_EVERY="${EVAL_EVERY:-0}"
+
+export OWT_CACHED_CHUNKS=1
+export OWT_CHUNK_CACHE_REBUILD=0
+export OWT_CHUNK_CACHE_DIR="${OWT_CHUNK_CACHE_DIR:-/e2e-data/evad-tech-vla/wanghan58/data/small_benchmarks/langflow_2604_11748/openwebtext_lta_cached_chunks/gpt2_len1024_train_minus_100k}"
+
+export OPTIMIZER="${OPTIMIZER:-muon}"
+export LEARNING_RATE="${LEARNING_RATE:-0.002}"
+export LR_SCHEDULE="${LR_SCHEDULE:-constant_warmup}"
+export MIN_LR="${MIN_LR:-0}"
+export WEIGHT_DECAY="${WEIGHT_DECAY:-0.0}"
+export ADAM_BETA1="${ADAM_BETA1:-0.9}"
+export ADAM_BETA2="${ADAM_BETA2:-0.95}"
+export ADAM_EPS="${ADAM_EPS:-1e-8}"
+export MUON_MOMENTUM="${MUON_MOMENTUM:-0.95}"
+export MUON_NS_STEPS="${MUON_NS_STEPS:-5}"
+export MUON_UPDATE_SCALE="${MUON_UPDATE_SCALE:-1.0}"
+export EMA_DECAY="${EMA_DECAY:-0.9999}"
+export EMA_START_STEP="${EMA_START_STEP:-0}"
+
+export ALLOW_TF32="${ALLOW_TF32:-1}"
+export ACTIVATION_CHECKPOINTING="${ACTIVATION_CHECKPOINTING:-0}"
+export ACTIVATION_CHECKPOINT_INTERVAL="${ACTIVATION_CHECKPOINT_INTERVAL:-2}"
+export DDP_GRADIENT_AS_BUCKET_VIEW="${DDP_GRADIENT_AS_BUCKET_VIEW:-1}"
+export FULL_TRAIN_STATS="${FULL_TRAIN_STATS:-0}"
+export ENABLE_TORCH_COMPILE="${ENABLE_TORCH_COMPILE:-0}"
+export ALLOW_EXISTING_SAVE_DIR="${ALLOW_EXISTING_SAVE_DIR:-0}"
+
+export ROLLOUT_TRAIN_PROB="${ROLLOUT_TRAIN_PROB:-0.5}"
+export ROLLOUT_TRAIN_STEPS="${ROLLOUT_TRAIN_STEPS:-1}"
+export ROLLOUT_TRAIN_INFER_STEPS="${ROLLOUT_TRAIN_INFER_STEPS:-64}"
+export ROLLOUT_TRAIN_TEMP="${ROLLOUT_TRAIN_TEMP:-1.45}"
+export ROLLOUT_TRAIN_MAX_GAMMA="${ROLLOUT_TRAIN_MAX_GAMMA:-1.0}"
+export ROLLOUT_TRAIN_CORRUPT_ONLY="${ROLLOUT_TRAIN_CORRUPT_ONLY:-1}"
+export ROLLOUT_TRAIN_SAMPLEWISE="${ROLLOUT_TRAIN_SAMPLEWISE:-1}"
+export ROLLOUT_TRAIN_COMPUTE_ALWAYS="${ROLLOUT_TRAIN_COMPUTE_ALWAYS:-0}"
+
+bash scripts/launch_lta_owt_c1024_fullycoupled_8gpu_len1024_gpt2_cached_chunks_1m.sh

LTA_openwebtext_dualt/scripts/launch_lta_owt_t5_blockar_twostream_8gpu.sh

@@ -0,0 +1,120 @@
+#!/bin/bash
+set -e
+set -x
+set -o 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 NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+export MASTER_PORT="${MASTER_PORT:-32331}"
+export PYTHONPATH="$(pwd)${PYTHONPATH:+:$PYTHONPATH}"
+export TOKENIZERS_PARALLELISM=false
+export PYTHONUNBUFFERED=1
+export NCCL_DEBUG="${NCCL_DEBUG:-WARN}"
+
+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}"
+
+# Two-stream packs a 1024-token example into a 2048-token model sequence, so the
+# dense first version needs a smaller per-GPU batch than the regular 1024 runs.
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-128}"
+PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-4}"
+TOTAL_STEPS="${TOTAL_STEPS:-1000000}"
+WARMUP_STEPS="${WARMUP_STEPS:-2000}"
+MAX_RECORDS="${MAX_RECORDS:-0}"
+
+D_MODEL="${D_MODEL:-768}"
+N_LAYERS="${N_LAYERS:-12}"
+N_HEADS="${N_HEADS:-12}"
+DIM_FF="${DIM_FF:-3072}"
+
+RUN_NAME="${RUN_NAME:-lta_owt_t5_blockar_twostream_b8x128_mask1_uniformt_adaln_swiglu_gbs${GLOBAL_BATCH_SIZE}_${NPROC_PER_NODE}gpu_$(date +%Y%m%d_%H%M%S)}"
+
+mkdir -p logs/blockar_twostream runs/${RUN_NAME}
+
+torchrun \
+  --nnodes=1 \
+  --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}" \
+  --max_records "${MAX_RECORDS}" \
+  --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 "${WARMUP_STEPS}" \
+  --log_every 100 \
+  --eval_every 0 \
+  --save_every 1000 \
+  --latest_every 500 \
+  --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 "${D_MODEL}" \
+  --cond_dim 128 \
+  --n_layers "${N_LAYERS}" \
+  --n_heads "${N_HEADS}" \
+  --dim_ff "${DIM_FF}" \
+  --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 \
+  --block_ar_two_stream \
+  --block_ar_block_len 128 \
+  --target_prob 1.0 \
+  --min_t 0.0 \
+  --max_t 1.0 \
+  --t_sampling_mode uniform \
+  --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_all_prob 1.0 \
+  --wrong_token_replace_prob 1.0 \
+  --wrong_token_schedule linear_t \
+  --dirichlet_concentration_min 1.0 \
+  --dirichlet_concentration_max 1024 \
+  --dirichlet_endpoint_mode categorical_dual_t \
+  --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 "logs/blockar_twostream/${RUN_NAME}.log"

LTA_openwebtext_dualt/scripts/launch_lta_owt_t5_rollin_p50_randk0_4_8gpu.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:-0,1,2,3,4,5,6,7}"
+export NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+export MASTER_PORT="${MASTER_PORT:-32254}"
+
+export DATA_PATH="${DATA_PATH:-/e2e-data/evad-tech-vla/wanghan58/data/embedded-language-flows/openwebtext-t5}"
+export TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/models/hf/t5-small/tokenizer.json}"
+
+export GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+export PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-32}"
+export EPOCHS="${EPOCHS:-5}"
+export TOTAL_STEPS="${TOTAL_STEPS:-1000000}"
+export NUM_WORKERS="${NUM_WORKERS:-8}"
+export DATALOADER_PREFETCH_FACTOR="${DATALOADER_PREFETCH_FACTOR:-4}"
+export LOG_EVERY="${LOG_EVERY:-100}"
+export LATEST_EVERY="${LATEST_EVERY:-1000}"
+export ALLOW_TF32="${ALLOW_TF32:-1}"
+
+export LR="${LR:-0.002}"
+export WEIGHT_DECAY="${WEIGHT_DECAY:-0.1}"
+export OUTPUT_WEIGHT_DECAY="${OUTPUT_WEIGHT_DECAY:--1}"
+export OUTPUT_INIT_STD="${OUTPUT_INIT_STD:-0.0}"
+
+# Corrected roll-in: model-temperature rollout, synced sampled path, random K in {0,1,2,3,4}.
+export T_SAMPLING_MODE="${T_SAMPLING_MODE:-uniform}"
+export ROLLOUT_TRAIN_PROB="${ROLLOUT_TRAIN_PROB:-0.50}"
+export ROLLOUT_TRAIN_TIME_MODE="${ROLLOUT_TRAIN_TIME_MODE:-sampled_path}"
+export ROLLOUT_TRAIN_STEPS="${ROLLOUT_TRAIN_STEPS:-4}"
+export ROLLOUT_TRAIN_STEPS_MIN="${ROLLOUT_TRAIN_STEPS_MIN:-0}"
+export ROLLOUT_TRAIN_INFER_STEPS="${ROLLOUT_TRAIN_INFER_STEPS:-1}"
+export ROLLOUT_TRAIN_S_DIST="${ROLLOUT_TRAIN_S_DIST:-uniform}"
+export ROLLOUT_TRAIN_S_MIN_FRAC="${ROLLOUT_TRAIN_S_MIN_FRAC:-0.0}"
+export ROLLOUT_TRAIN_S_MAX_FRAC="${ROLLOUT_TRAIN_S_MAX_FRAC:-0.25}"
+export ROLLOUT_TRAIN_TEMP="${ROLLOUT_TRAIN_TEMP:-1.0}"
+export ROLLOUT_TRAIN_MAX_GAMMA="${ROLLOUT_TRAIN_MAX_GAMMA:-1.0}"
+export ROLLOUT_TRAIN_CORRUPT_ONLY="${ROLLOUT_TRAIN_CORRUPT_ONLY:-1}"
+export ROLLOUT_TRAIN_SAMPLEWISE="${ROLLOUT_TRAIN_SAMPLEWISE:-1}"
+export ROLLOUT_TRAIN_SELECTED_ONLY="${ROLLOUT_TRAIN_SELECTED_ONLY:-1}"
+export ROLLOUT_TRAIN_COMPUTE_ALWAYS="${ROLLOUT_TRAIN_COMPUTE_ALWAYS:-0}"
+export ROLLOUT_TRAIN_SYNC_T="${ROLLOUT_TRAIN_SYNC_T:-1}"
+
+export MIN_MASK_RATIO="${MIN_MASK_RATIO:-1.0}"
+export MAX_MASK_RATIO="${MAX_MASK_RATIO:-1.0}"
+export MASK_MIXTURE_ORIGINAL_PROB="${MASK_MIXTURE_ORIGINAL_PROB:-0.0}"
+export MASK_MIXTURE_LOWK_PROB="${MASK_MIXTURE_LOWK_PROB:-0.0}"
+export MASK_MIXTURE_LOWCORRUPT_PROB="${MASK_MIXTURE_LOWCORRUPT_PROB:-0.0}"
+export MASK_MIXTURE_BLOCK_PROB="${MASK_MIXTURE_BLOCK_PROB:-0.0}"
+export MASK_MIXTURE_ALL_PROB="${MASK_MIXTURE_ALL_PROB:-1.0}"
+
+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 RUN_NAME="${RUN_NAME:-lta_owt_t5_rollin_p50_randk0_4_uniformt_temp1_synct_mask1_gbs${GLOBAL_BATCH_SIZE}_${NPROC_PER_NODE}gpu_$(date +%Y%m%d_%H%M%S)}"
+
+bash scripts/launch_lta_owt_elfaligned_t5_logitnormal_tokenized_8gpu.sh

LTA_openwebtext_dualt/scripts/make_compact_gpt2_bpe_tokenizer.py

@@ -0,0 +1,160 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Iterable
+
+
+SPECIAL_TOKENS = ["<pad>", "</s>", "<unk>", "<|endoftext|>"]
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(
+        description=(
+            "Create compact byte-level GPT-2 BPE tokenizers by truncating the "
+            "merge table while preserving byte fallback coverage."
+        )
+    )
+    p.add_argument("--source_tokenizer", required=True)
+    p.add_argument("--output_dir", required=True)
+    p.add_argument("--vocab_size", type=int, required=True)
+    p.add_argument(
+        "--special_tokens",
+        nargs="*",
+        default=SPECIAL_TOKENS,
+        help="Special tokens to reserve first in the compact vocab.",
+    )
+    return p.parse_args()
+
+
+def normalize_merge(merge: str | list[str]) -> tuple[str, str]:
+    if isinstance(merge, str):
+        left, right = merge.split(" ", 1)
+        return left, right
+    if isinstance(merge, list) and len(merge) == 2:
+        return str(merge[0]), str(merge[1])
+    raise ValueError(f"Unsupported merge entry: {merge!r}")
+
+
+def merge_to_json(left: str, right: str, like: str | list[str]) -> str | list[str]:
+    if isinstance(like, str):
+        return f"{left} {right}"
+    return [left, right]
+
+
+def base_byte_tokens(orig_id_to_token: dict[int, str]) -> Iterable[str]:
+    # GPT-2 byte-level BPE starts with exactly 256 byte fallback tokens.
+    for idx in range(256):
+        token = orig_id_to_token.get(idx)
+        if token is None:
+            raise ValueError(f"Source tokenizer is missing byte token id={idx}")
+        yield token
+
+
+def main() -> None:
+    args = parse_args()
+    src_path = Path(args.source_tokenizer)
+    out_dir = Path(args.output_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    with src_path.open("r", encoding="utf-8") as f:
+        tokenizer = json.load(f)
+
+    model = tokenizer["model"]
+    orig_vocab: dict[str, int] = {str(k): int(v) for k, v in model["vocab"].items()}
+    orig_id_to_token = {idx: tok for tok, idx in orig_vocab.items()}
+    orig_merges = list(model.get("merges", []))
+
+    special_tokens = []
+    for tok in args.special_tokens:
+        if tok not in special_tokens:
+            special_tokens.append(tok)
+
+    if args.vocab_size <= len(special_tokens) + 256:
+        raise ValueError(
+            f"vocab_size={args.vocab_size} must exceed specials+byte alphabet "
+            f"({len(special_tokens)}+256)"
+        )
+
+    ordered_tokens: list[str] = []
+    token_set: set[str] = set()
+
+    def add_token(tok: str) -> None:
+        if tok not in token_set:
+            token_set.add(tok)
+            ordered_tokens.append(tok)
+
+    for tok in special_tokens:
+        add_token(tok)
+    for tok in base_byte_tokens(orig_id_to_token):
+        add_token(tok)
+
+    compact_merges: list[str | list[str]] = []
+    for merge in orig_merges:
+        left, right = normalize_merge(merge)
+        out = left + right
+        if left not in token_set or right not in token_set:
+            continue
+        if out not in orig_vocab:
+            continue
+        if out not in token_set:
+            if len(ordered_tokens) >= args.vocab_size:
+                break
+            add_token(out)
+        if out in token_set:
+            compact_merges.append(merge_to_json(left, right, merge))
+        if len(ordered_tokens) >= args.vocab_size:
+            break
+
+    if len(ordered_tokens) != args.vocab_size:
+        raise RuntimeError(
+            f"Could only build vocab of size {len(ordered_tokens)}; "
+            f"requested {args.vocab_size}"
+        )
+
+    compact_vocab = {tok: idx for idx, tok in enumerate(ordered_tokens)}
+
+    out = dict(tokenizer)
+    out["truncation"] = None
+    out["padding"] = None
+    out["added_tokens"] = [
+        {
+            "id": compact_vocab[tok],
+            "special": True,
+            "content": tok,
+            "single_word": False,
+            "lstrip": False,
+            "rstrip": False,
+            "normalized": True,
+        }
+        for tok in special_tokens
+    ]
+    out["post_processor"] = None
+    out_model = dict(model)
+    out_model["vocab"] = compact_vocab
+    out_model["merges"] = compact_merges
+    out_model["unk_token"] = "<unk>"
+    out["model"] = out_model
+
+    with (out_dir / "tokenizer.json").open("w", encoding="utf-8") as f:
+        json.dump(out, f, ensure_ascii=False)
+
+    meta = {
+        "source_tokenizer": str(src_path),
+        "vocab_size": args.vocab_size,
+        "actual_vocab_size": len(compact_vocab),
+        "num_merges": len(compact_merges),
+        "special_tokens": special_tokens,
+        "construction": "compact_gpt2_bpe_prefix_merges",
+    }
+    with (out_dir / "tokenizer_meta.json").open("w", encoding="utf-8") as f:
+        json.dump(meta, f, indent=2, ensure_ascii=False)
+        f.write("\n")
+
+    print(json.dumps(meta, ensure_ascii=False))
+
+
+if __name__ == "__main__":
+    main()

LTA_openwebtext_dualt/scripts/run_lta_lm1b_linear_simplex_len128_4gpu_save1k_with_watch.sh

@@ -0,0 +1,217 @@
+#!/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:-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}"
+SCORER="${SCORER:-/e2e-data/evad-tech-vla/wanghan58/models/flowtext_scorers/gpt2-large-standard}"
+
+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:-32931}"
+
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-64}"
+TOTAL_STEPS="${TOTAL_STEPS:-20000}"
+WARMUP_STEPS="${WARMUP_STEPS:-2500}"
+MAX_LEN="${MAX_LEN:-128}"
+NUM_WORKERS="${NUM_WORKERS:-0}"
+LOG_EVERY="${LOG_EVERY:-100}"
+EVAL_EVERY="${EVAL_EVERY:-0}"
+SAVE_EVERY="${SAVE_EVERY:-1000}"
+LATEST_EVERY="${LATEST_EVERY:-1000}"
+
+LR="${LR:-3e-4}"
+D_MODEL="${D_MODEL:-768}"
+COND_DIM="${COND_DIM:-128}"
+N_LAYERS="${N_LAYERS:-12}"
+N_HEADS="${N_HEADS:-12}"
+DIM_FF="${DIM_FF:-3072}"
+DROPOUT="${DROPOUT:-0.1}"
+
+LINEAR_NOISE_SIGMA="${LINEAR_NOISE_SIGMA:-3.0}"
+WRONG_TOKEN_REPLACE_PROB="${WRONG_TOKEN_REPLACE_PROB:-0.0}"
+WRONG_TOKEN_SCHEDULE="${WRONG_TOKEN_SCHEDULE:-constant}"
+ENDPOINT_TEMP_START="${ENDPOINT_TEMP_START:-1.45}"
+ENDPOINT_TEMP_END="${ENDPOINT_TEMP_END:-0.8}"
+ENDPOINT_PROJECTION="${ENDPOINT_PROJECTION:-soft}"
+FINAL_FROM="${FINAL_FROM:-model_t1}"
+
+ENABLE_WATCHER="${ENABLE_WATCHER:-1}"
+WATCH_CUDA_VISIBLE_DEVICES="${WATCH_CUDA_VISIBLE_DEVICES:-0}"
+WATCH_SAMPLES="${WATCH_SAMPLES:-64}"
+WATCH_STEPS="${WATCH_STEPS:-128}"
+WATCH_POLL_SECONDS="${WATCH_POLL_SECONDS:-30}"
+
+RESUME_PATH="${RESUME_PATH:-}"
+ALLOW_EXISTING_SAVE_DIR="${ALLOW_EXISTING_SAVE_DIR:-0}"
+
+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
+if [[ ! -e "${DATA_PATH}" ]]; then
+  echo "Missing DATA_PATH: ${DATA_PATH}" >&2
+  exit 2
+fi
+if [[ ! -f "${TOKENIZER_PATH}" ]]; then
+  echo "Missing TOKENIZER_PATH: ${TOKENIZER_PATH}" >&2
+  exit 2
+fi
+
+RUN_NAME="${RUN_NAME:-lta_lm1b_linear_simplex_len${MAX_LEN}_sigma${LINEAR_NOISE_SIGMA}_gbs${GLOBAL_BATCH_SIZE}_4gpu_20k_save1k_$(date +%Y%m%d_%H%M%S)}"
+SAVE_DIR="${SAVE_DIR:-runs/${RUN_NAME}}"
+LOG_DIR="${LOG_DIR:-logs/lm1b_linear_simplex_len128_4gpu}"
+LOG_FILE="${LOG_FILE:-${LOG_DIR}/${RUN_NAME}.log}"
+WATCH_LOG="${WATCH_LOG:-${LOG_DIR}/${RUN_NAME}.watch.log}"
+
+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
+  exit 2
+fi
+
+mkdir -p "${LOG_DIR}" runs "${SAVE_DIR}"
+
+RESUME_ARGS=()
+if [[ -n "${RESUME_PATH}" ]]; then
+  RESUME_ARGS+=(--resume_path "${RESUME_PATH}")
+fi
+
+echo "[launch] method=lm1b_linear_simplex_algebraic_len128_4gpu"
+echo "[launch] host=$(hostname) time=$(date -Iseconds)"
+echo "[launch] run_name=${RUN_NAME}"
+echo "[launch] save_dir=${SAVE_DIR}"
+echo "[launch] bridge=prob linear simplex: p_t=(1-t)*softmax(${LINEAR_NOISE_SIGMA}*N)+t*onehot(y)"
+echo "[launch] wrong_endpoint: prob=${WRONG_TOKEN_REPLACE_PROB} schedule=${WRONG_TOKEN_SCHEDULE}"
+echo "[launch] infer=linear algebraic: p_next=(1-s_next)*p0+s_next*a_theta"
+echo "[launch] watcher=${ENABLE_WATCHER} watch_gpu=${WATCH_CUDA_VISIBLE_DEVICES} watch_samples=${WATCH_SAMPLES}"
+
+run_watcher() {
+  local last_done=0
+  while true; do
+    for ckpt in "${SAVE_DIR}"/step_*.pt; do
+      [[ -f "${ckpt}" ]] || continue
+      local base step out_dir
+      base="$(basename "${ckpt}")"
+      step="${base#step_}"
+      step="${step%.pt}"
+      step=$((10#${step}))
+      if (( step <= last_done )); then
+        continue
+      fi
+      out_dir="docs/lta_samples/metrics_$(date +%Y%m%d)/${RUN_NAME}_step${step}_linear_alg_t${ENDPOINT_TEMP_START}_to_${ENDPOINT_TEMP_END}_n${WATCH_SAMPLES}"
+      echo "[watch] infer step=${step} ckpt=${ckpt} out=${out_dir}" | tee -a "${WATCH_LOG}"
+      CUDA_VISIBLE_DEVICES="${WATCH_CUDA_VISIBLE_DEVICES}" python scripts/eval_lm1b_linear_simplex_genppl.py \
+        --checkpoint "${ckpt}" \
+        --tokenizer_path "${TOKENIZER_PATH}" \
+        --scorer "${SCORER}" \
+        --out_dir "${out_dir}" \
+        --n_samples "${WATCH_SAMPLES}" \
+        --max_len "${MAX_LEN}" \
+        --steps "${WATCH_STEPS}" \
+        --batch_size 1 \
+        --score_batch 1 \
+        --noise_init logistic_normal \
+        --noise_sigma "${LINEAR_NOISE_SIGMA}" \
+        --endpoint_temp_start "${ENDPOINT_TEMP_START}" \
+        --endpoint_temp_end "${ENDPOINT_TEMP_END}" \
+        --endpoint_projection "${ENDPOINT_PROJECTION}" \
+        --endpoint_top_p 1.0 \
+        --final_from "${FINAL_FROM}" \
+        --seed "$((20260524 + step))" 2>&1 | tee -a "${WATCH_LOG}" || true
+      last_done="${step}"
+    done
+    if (( last_done >= TOTAL_STEPS )); then
+      break
+    fi
+    sleep "${WATCH_POLL_SECONDS}"
+  done
+}
+
+if [[ "${ENABLE_WATCHER}" == "1" ]]; then
+  run_watcher &
+  WATCH_PID=$!
+  echo "[launch] watcher_pid=${WATCH_PID}"
+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 all \
+  --detokenizer auto \
+  --tokenizer_path "${TOKENIZER_PATH}" \
+  --save_dir "${SAVE_DIR}" \
+  --wrap \
+  --wrap_mode stream \
+  --wrap_record_buffer_size 200 \
+  --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 "${LR}" \
+  --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 "${D_MODEL}" \
+  --cond_dim "${COND_DIM}" \
+  --n_layers "${N_LAYERS}" \
+  --n_heads "${N_HEADS}" \
+  --dim_ff "${DIM_FF}" \
+  --dropout "${DROPOUT}" \
+  --model_type ddit \
+  --state_format prob \
+  --bridge prob \
+  --target_loss hard_ce \
+  --target_prob 1.0 \
+  --min_t 0.0 \
+  --max_t 1.0 \
+  --t_sampling_mode uniform \
+  --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_all_prob 1.0 \
+  --wrong_token_replace_prob "${WRONG_TOKEN_REPLACE_PROB}" \
+  --wrong_token_schedule "${WRONG_TOKEN_SCHEDULE}" \
+  --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 "${LINEAR_NOISE_SIGMA}" \
+  "${RESUME_ARGS[@]}" \
+  --bf16 2>&1 | tee -a "${LOG_FILE}"
+
+if [[ "${ENABLE_WATCHER}" == "1" ]]; then
+  wait "${WATCH_PID}" || true
+fi

LTA_openwebtext_dualt/scripts/run_lta_owt_t5elf_absrope_time4_dirichlet_len1025_C1_to_1024_8gpu_1m_mask1_sameT_prebos_save10k.sh

@@ -0,0 +1,74 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+cd /e2e-data/evad-tech-vla/wanghan58/workspace/LTA_openwebtext_dualt
+
+# OWT through the existing ELF/T5 tokenized pipeline:
+#   [BOS] + up to 1024 existing T5 tokens.
+#
+# The ELF T5 records already end with </s>, so we only prepend BOS here.
+# T5 has no separate BOS token in this tokenizer; the repo wrapper falls back
+# to bos_id = eos_id = </s>.
+#
+# Backbone:
+#   ddit_elf = RMSNorm/SwiGLU/QK-norm + RoPE + 4 prefix time tokens.
+#   Learned absolute position embeddings are added before RoPE.
+#
+# Bridge:
+#   Dirichlet C=1->1024, mask_ratio=1.0, model t and corruption t are shared.
+
+export DATA_PATH="${DATA_PATH:-/e2e-data/evad-tech-vla/wanghan58/data/embedded-language-flows/openwebtext-t5}"
+export TOKENIZER_PATH="${TOKENIZER_PATH:-/e2e-data/evad-tech-vla/wanghan58/models/hf/t5-small/tokenizer.json}"
+export TOKENIZED_HF=1
+export TOKENIZED_PAD_TOKEN="${TOKENIZED_PAD_TOKEN:-pad}"
+export TOKENIZED_PREPEND_BOS=1
+export TOKENIZED_APPEND_EOS=0
+export TOKENIZED_STRIP_EDGE_SPECIALS=0
+
+export MAX_LEN="${MAX_LEN:-1025}"
+export VOCAB_SIZE="${VOCAB_SIZE:-32100}"
+export CMIN="${CMIN:-1}"
+export CMAX="${CMAX:-1024}"
+
+export MODEL_TYPE=ddit_elf
+export ELF_NUM_TIME_TOKENS="${ELF_NUM_TIME_TOKENS:-4}"
+export ELF_NUM_MODEL_MODE_TOKENS="${ELF_NUM_MODEL_MODE_TOKENS:-0}"
+export QK_NORM="${QK_NORM:-1}"
+export ABS_POS_EMBED=1
+export CORRUPT_T_MODE=same
+export MIN_MASK_RATIO=1.0
+export MAX_MASK_RATIO=1.0
+export CATEGORICAL_WRONG_PROB_FLOOR="${CATEGORICAL_WRONG_PROB_FLOOR:-0.0}"
+
+export CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}"
+export NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+export GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+export PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-32}"
+export TOTAL_STEPS="${TOTAL_STEPS:-1000000}"
+export WARMUP_STEPS="${WARMUP_STEPS:-2500}"
+export SAVE_EVERY="${SAVE_EVERY:-10000}"
+export LATEST_EVERY="${LATEST_EVERY:-1000}"
+export LOG_EVERY="${LOG_EVERY:-100}"
+
+export DATE_TAG="${DATE_TAG:-$(date +%Y%m%d)}"
+export RUN_NAME="${RUN_NAME:-lta_owt_t5elf_absrope_time4_dirichlet_len1025_C1_to_1024_prebos_mask1_sameT_gbs512_b32_8gpu_1m_save10k_${DATE_TAG}}"
+
+export WATCH_ENABLED="${WATCH_ENABLED:-1}"
+export WATCH_STEP_INTERVAL="${WATCH_STEP_INTERVAL:-10000}"
+export WATCH_N_SAMPLES="${WATCH_N_SAMPLES:-128}"
+export WATCH_CUDA_VISIBLE_DEVICES="${WATCH_CUDA_VISIBLE_DEVICES:-7}"
+export WATCH_DECODE_MODE="${WATCH_DECODE_MODE:-dual_line_probe}"
+export WATCH_DECODE_BATCH="${WATCH_DECODE_BATCH:-1}"
+export WATCH_DUAL_SEMANTIC_POWER="${WATCH_DUAL_SEMANTIC_POWER:-1.5}"
+export WATCH_DUAL_EARLY_TEMP="${WATCH_DUAL_EARLY_TEMP:-2.8}"
+export WATCH_DUAL_LATE_TEMP="${WATCH_DUAL_LATE_TEMP:-1.45}"
+export WATCH_DUAL_TEMP_END="${WATCH_DUAL_TEMP_END:-0.55}"
+export WATCH_DUAL_TEMP_POWER="${WATCH_DUAL_TEMP_POWER:-1.5}"
+export WATCH_ENDPOINT_TEMP="${WATCH_ENDPOINT_TEMP:-1.45}"
+export WATCH_ENDPOINT_TOP_P="${WATCH_ENDPOINT_TOP_P:-0.95}"
+export WATCH_GUMBEL_TAU_START="${WATCH_GUMBEL_TAU_START:-1.0}"
+export WATCH_GUMBEL_TAU_END="${WATCH_GUMBEL_TAU_END:-0.2}"
+export WATCH_OUT_BASE="${WATCH_OUT_BASE:-docs/lta_samples/metrics_${DATE_TAG}/owt_t5elf_absrope_time4_len1025_C1_to_1024_prebos_mask1_sameT_dualline_dirres_c${CMIN}_${CMAX}_n${WATCH_N_SAMPLES}/${RUN_NAME}}"
+export WATCH_LOG_DIR="${WATCH_LOG_DIR:-logs/owt_t5elf_absrope_time4_len1025_C1_to_1024_prebos_mask1_sameT_dualline_watch}"
+
+bash scripts/run_lta_owt_dirichlet_len1024_Cv_to_2v_8gpu_save1k_with_gumbel_watch.sh

LTA_openwebtext_dualt/scripts/run_train8_ctx1024_rollin_sweep_4gpu.sh

@@ -0,0 +1,235 @@
+#!/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}"
+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:-train8_rollin_len1024_rollin_p50_s4_i32_20260517_1855ctx1024bs128}"
+OUT_ROOT="${OUT_ROOT:-docs/lta_samples/metrics_20260517/ctx1024_rollin_sweep_bs512_ode128_${GROUP_STAMP}}"
+DRIVER_LOG="${DRIVER_LOG:-logs/ctx1024_rollin_sweep_4gpu/${GROUP_STAMP}.log}"
+CURVE_CSV="${CURVE_CSV:-${OUT_ROOT}/hit_ratio_curve.csv}"
+mkdir -p "$(dirname "${DRIVER_LOG}")" "${OUT_ROOT}"
+
+cache="${BASE_CACHE}/gpt2_len${MAX_LEN}_train8_compact_overfit"
+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,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"
+  local out_dir="${OUT_ROOT}/${config}/step_${target_step}"
+  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 flowmap \
+    --time_schedule logit_normal \
+    --time_logit_mean -1.5 \
+    --time_logit_std 0.8 \
+    --model_t_mode post \
+    --c_min 1 \
+    --c_max 512 \
+    --late_temp 1.0 \
+    --final_from state \
+    --final_decode argmax
+  python - "$out_dir" "$N_SAMPLES" "$GLOBAL_BATCH_SIZE" "$MAX_LEN" "$CURVE_CSV" "$config" "$run_name" <<'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]
+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} 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},{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
+}
+
+run_config() {
+  local config="$1"
+  local rollout_prob="$2"
+  local rollout_steps="$3"
+  local rollout_infer_steps="$4"
+  local output_wd="$5"
+  local sync_t="$6"
+  local run_name="${RUN_PREFIX:-train8_ctx1024}_${config}_${GROUP_STAMP}"
+  echo "[ctx1024-sweep] config=${config} run=${run_name} p=${rollout_prob} steps=${rollout_steps} infer=${rollout_infer_steps} 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-sweep] 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-sweep] 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_INFER_STEPS="${rollout_infer_steps}" \
+    ROLLOUT_TRAIN_TEMP=1.45 \
+    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}" \
+    RUN_NAME="${run_name}" \
+    RESUME_PATH="${resume_path}" \
+    bash scripts/launch_lta_owt_gpt2_softendpoint_mn_pilot_4gpu.sh
+    echo "[ctx1024-sweep] eval config=${config} step=${target_step}" | tee -a "${DRIVER_LOG}"
+    eval_latest "${config}" "${run_name}" "${target_step}" | tee -a "${DRIVER_LOG}"
+    if tail -n 1 "${CURVE_CSV}" | awk -F, '{exit ($7 >= 60) ? 0 : 1}'; then
+      echo "[ctx1024-sweep] early-hit config=${config}" | tee -a "${DRIVER_LOG}"
+      break
+    fi
+  done
+}
+
+if [[ -n "${WAIT_FOR_RUN}" ]]; then
+  echo "[ctx1024-sweep] 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-sweep] start stamp=${GROUP_STAMP} len=${MAX_LEN} vocab=${vocab_size} out=${OUT_ROOT}" | tee -a "${DRIVER_LOG}"
+
+CONFIGS=(
+  "p75_s4_i32_outwdm1|0.75|4|32|-1|0"
+  "p100_s4_i32_outwdm1|1.0|4|32|-1|0"
+  "p50_s8_i64_outwdm1|0.50|8|64|-1|0"
+  "p50_s4_i32_sync1_outwdm1|0.50|4|32|-1|1"
+  "p50_s8_i64_sync1_outwdm1|0.50|8|64|-1|1"
+  "p50_s4_i32_outwd0p3|0.50|4|32|0.3|0"
+  "p50_s4_i32_outwd1p0|0.50|4|32|1.0|0"
+)
+
+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 rollout_steps rollout_infer_steps output_wd sync_t <<<"${entry}"
+  run_config "${config}" "${rollout_prob}" "${rollout_steps}" "${rollout_infer_steps}" "${output_wd}" "${sync_t}"
+done
+
+echo "[ctx1024-sweep] done" | tee -a "${DRIVER_LOG}"

LTA_openwebtext_dualt/scripts/run_train8_len256_headwd_sweep_bs512_until_exact_4gpu.sh

@@ -0,0 +1,287 @@
+#!/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}"
+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}"
+STOP_EXACT_COUNT="${STOP_EXACT_COUNT:-1}"
+STOP_EXACT_REF_COUNT="${STOP_EXACT_REF_COUNT:-0}"
+INFER_STEPS="${INFER_STEPS:-128}"
+STEP_CHUNK="${STEP_CHUNK:-500}"
+MAX_TOTAL_STEPS="${MAX_TOTAL_STEPS:-100000}"
+MASTER_PORT_BASE="${MASTER_PORT_BASE:-38200}"
+PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-128}"
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+OWT_EXACT_REPEAT_PER_CHUNK="${OWT_EXACT_REPEAT_PER_CHUNK:-64}"
+LEARNING_RATE="${LEARNING_RATE:-0.002}"
+WEIGHT_DECAY="${WEIGHT_DECAY:-0.1}"
+HEAD_WD_VALUES="${HEAD_WD_VALUES:-0.1 0.3 1.0 3.0}"
+OUT_ROOT="${OUT_ROOT:-docs/lta_samples/metrics_20260517/train8_len256_headwd_sweep_bs512_until_exact_ode128}"
+DRIVER_LOG="${DRIVER_LOG:-logs/train8_len256_headwd_sweep_bs512_until_exact_4gpu/driver.log}"
+SUMMARY_JSONL="${SUMMARY_JSONL:-${OUT_ROOT}/summary_results.jsonl}"
+CURVE_CSV="${CURVE_CSV:-${OUT_ROOT}/hit_ratio_curve.csv}"
+mkdir -p "$(dirname "${DRIVER_LOG}")" "${OUT_ROOT}"
+
+tag_float() {
+  printf "%s" "$1" | sed -e 's/-/m/g' -e 's/\./p/g'
+}
+
+vocab_size_for_len() {
+  python - "$BASE_CACHE" "$MAX_LEN" <<'PY'
+import json
+import sys
+from pathlib import Path
+base = Path(sys.argv[1])
+max_len = int(sys.argv[2])
+meta = json.loads((base / f"gpt2_len{max_len}_train8_compact_overfit" / "meta.json").read_text())
+print(int(meta.get("compact_vocab_size", meta.get("vocab_size"))))
+PY
+}
+
+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
+}
+
+refresh_curve() {
+  python - "$SUMMARY_JSONL" "$CURVE_CSV" <<'PY'
+import csv
+import json
+import re
+import sys
+from pathlib import Path
+
+summary = Path(sys.argv[1])
+out = Path(sys.argv[2])
+rows = []
+seen = set()
+if summary.exists():
+    for line in summary.read_text().splitlines():
+        if not line.strip():
+            continue
+        row = json.loads(line)
+        m = re.search(
+            r"train8_n(\d+)_compactv\d+_small_headwd([^_]+)_bs512_hard_ce_(onehot|allcorrupt)",
+            row.get("run", ""),
+        )
+        if not m:
+            continue
+        max_len = int(m.group(1))
+        head_wd = m.group(2).replace("p", ".").replace("m", "-")
+        mode = m.group(3)
+        step = int(row["ckpt_step"])
+        key = (max_len, head_wd, mode, step)
+        if key in seen:
+            continue
+        seen.add(key)
+        n_gen = float(row.get("n_gen", 64))
+        rows.append({
+            "context_len": max_len,
+            "head_weight_decay": head_wd,
+            "mode": mode,
+            "step": step,
+            "train_views_seen": int(row.get("train_views_seen", step * 512)),
+            "train_tokens_seen": int(row.get("train_tokens_seen", step * 512 * max_len)),
+            "hit_ratio": float(row.get("exact_count", 0)) / n_gen,
+            "exact_count": int(row.get("exact_count", 0)),
+            "ref_coverage": float(row.get("exact_ref_count", 0)) / 8.0,
+            "exact_ref_count": int(row.get("exact_ref_count", 0)),
+            "token_acc_mean": float(row.get("token_acc_mean", 0.0)),
+        })
+rows.sort(key=lambda x: (x["context_len"], float(x["head_weight_decay"]), x["mode"], x["step"]))
+out.parent.mkdir(parents=True, exist_ok=True)
+with out.open("w", newline="", encoding="utf-8") as f:
+    fieldnames = [
+        "context_len",
+        "head_weight_decay",
+        "mode",
+        "step",
+        "train_views_seen",
+        "train_tokens_seen",
+        "hit_ratio",
+        "exact_count",
+        "ref_coverage",
+        "exact_ref_count",
+        "token_acc_mean",
+    ]
+    writer = csv.DictWriter(f, fieldnames=fieldnames)
+    writer.writeheader()
+    writer.writerows(rows)
+print(f"curve_rows={len(rows)}")
+PY
+}
+
+eval_latest() {
+  local run_name="$1"
+  local target_step="$2"
+  local head_wd="$3"
+  local cache="$4"
+  local out_dir="${OUT_ROOT}/len${MAX_LEN}/headwd_$(tag_float "${head_wd}")/${run_name}/step_${target_step}"
+  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 flowmap \
+    --time_schedule logit_normal \
+    --time_logit_mean -1.5 \
+    --time_logit_std 0.8 \
+    --model_t_mode post \
+    --c_min 1 \
+    --c_max 512 \
+    --late_temp 1.0 \
+    --final_from state \
+    --final_decode argmax
+  python - "$out_dir" "$N_SAMPLES" "$GLOBAL_BATCH_SIZE" "$MAX_LEN" "$SUMMARY_JSONL" "$STOP_EXACT_COUNT" "$STOP_EXACT_REF_COUNT" "$head_wd" <<'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])
+summary = Path(sys.argv[5])
+stop_exact_count = int(sys.argv[6])
+stop_exact_ref_count = int(sys.argv[7])
+head_wd = float(sys.argv[8])
+row = json.loads((out / "decode_token_acc.jsonl").read_text().splitlines()[-1])
+row["head_weight_decay"] = head_wd
+row["train_views_seen"] = int(row["ckpt_step"]) * global_batch
+row["train_tokens_seen"] = row["train_views_seen"] * max_len
+line = (
+    "RESULT "
+    f"len={max_len} "
+    f"head_wd={head_wd:g} "
+    f"run={row['run']} "
+    f"ckpt_step={row['ckpt_step']} "
+    f"train_views_seen={row['train_views_seen']} "
+    f"train_tokens_seen={row['train_tokens_seen']} "
+    f"token_acc_mean={row['token_acc_mean']:.4f} "
+    f"exact_count={row['exact_count']}/{n} "
+    f"exact_ref_count={row['exact_ref_count']} "
+    f"exact_refs={row['exact_ref_hits']}"
+)
+print(line, flush=True)
+summary.parent.mkdir(parents=True, exist_ok=True)
+with summary.open("a", encoding="utf-8") as f:
+    f.write(json.dumps(row, ensure_ascii=False) + "\n")
+raise SystemExit(
+    0
+    if int(row["exact_count"]) >= stop_exact_count
+    and int(row["exact_ref_count"]) >= stop_exact_ref_count
+    else 1
+)
+PY
+}
+
+cache="${BASE_CACHE}/gpt2_len${MAX_LEN}_train8_compact_overfit"
+vocab_size="$(vocab_size_for_len)"
+
+mode_specs=(
+  "onehot|1|0|1,2,4"
+  "allcorrupt|0|1|0"
+)
+
+for head_wd in ${HEAD_WD_VALUES}; do
+  head_tag="$(tag_float "${head_wd}")"
+  echo "[headwd-sweep] head_wd=${head_wd} $(date)" | tee -a "${DRIVER_LOG}"
+  for mode_spec in "${mode_specs[@]}"; do
+    IFS='|' read -r mode lowk_prob all_prob lowk_tokens <<<"${mode_spec}"
+    run_name="train8_n${MAX_LEN}_compactv${vocab_size}_small_headwd${head_tag}_bs512_hard_ce_${mode}"
+    done_flag="${OUT_ROOT}/len${MAX_LEN}/headwd_${head_tag}/${run_name}/DONE"
+    while :; do
+      if [[ -f "${done_flag}" ]]; then
+        echo "[headwd-sweep] skip done head_wd=${head_wd} mode=${mode} run=${run_name}" | tee -a "${DRIVER_LOG}"
+        break
+      fi
+
+      step_now="$(latest_step "${run_name}")"
+      if [[ "${step_now}" -ge "${MAX_TOTAL_STEPS}" ]]; then
+        echo "[headwd-sweep] hit max_total_steps head_wd=${head_wd} mode=${mode} step=${step_now}" | tee -a "${DRIVER_LOG}"
+        break
+      fi
+      target_step=$((step_now + STEP_CHUNK))
+      if [[ "${target_step}" -gt "${MAX_TOTAL_STEPS}" ]]; then
+        target_step="${MAX_TOTAL_STEPS}"
+      fi
+
+      resume_path=""
+      if [[ -f "runs/${run_name}/latest.pt" ]]; then
+        resume_path="runs/${run_name}/latest.pt"
+      fi
+
+      echo "[headwd-sweep] train head_wd=${head_wd} mode=${mode} run=${run_name} 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=$((MASTER_PORT_BASE + target_step / STEP_CHUNK)) \
+      OWT_CHUNK_CACHE_DIR="${cache}" \
+      OWT_EXACT_REPEAT_PER_CHUNK="${OWT_EXACT_REPEAT_PER_CHUNK}" \
+      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}" \
+      LOWK_CLEAN_TOKENS="${lowk_tokens}" \
+      MASK_MIXTURE_LOWK_PROB="${lowk_prob}" \
+      MASK_MIXTURE_ALL_PROB="${all_prob}" \
+      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:-50}" \
+      SAVE_EVERY="${STEP_CHUNK}" \
+      LATEST_EVERY="${STEP_CHUNK}" \
+      WARMUP_STEPS="${WARMUP_STEPS:-10}" \
+      LEARNING_RATE="${LEARNING_RATE}" \
+      WEIGHT_DECAY="${WEIGHT_DECAY}" \
+      OUTPUT_WEIGHT_DECAY="${head_wd}" \
+      MUON_IMPL="${MUON_IMPL:-optax}" \
+      TARGET_LOSS="hard_ce" \
+      CLEAN_STATE_MODE="onehot" \
+      RUN_NAME="${run_name}" \
+      RESUME_PATH="${resume_path}" \
+      bash scripts/launch_lta_owt_gpt2_softendpoint_mn_pilot_4gpu.sh
+
+      echo "[headwd-sweep] eval head_wd=${head_wd} mode=${mode} run=${run_name} step=${target_step}" | tee -a "${DRIVER_LOG}"
+      if eval_latest "${run_name}" "${target_step}" "${head_wd}" "${cache}" | tee -a "${DRIVER_LOG}"; then
+        echo "[headwd-sweep] DONE head_wd=${head_wd} mode=${mode} run=${run_name} step=${target_step}" | tee -a "${DRIVER_LOG}"
+        touch "${done_flag}"
+        refresh_curve | tee -a "${DRIVER_LOG}"
+        break
+      else
+        echo "[headwd-sweep] continue head_wd=${head_wd} mode=${mode} run=${run_name} step=${target_step}" | tee -a "${DRIVER_LOG}"
+        refresh_curve | tee -a "${DRIVER_LOG}"
+      fi
+    done
+  done
+done
+
+refresh_curve | tee -a "${DRIVER_LOG}"

LTA_openwebtext_dualt/scripts/run_train8_noise_geometry_pilots_4gpu.sh

@@ -0,0 +1,289 @@
+#!/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}"
+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}"
+STEP_CHUNK="${STEP_CHUNK:-1000}"
+MAX_TOTAL_STEPS="${MAX_TOTAL_STEPS:-30000}"
+STOP_EXACT_COUNT="${STOP_EXACT_COUNT:-64}"
+STOP_EXACT_REF_COUNT="${STOP_EXACT_REF_COUNT:-8}"
+MASTER_PORT_BASE="${MASTER_PORT_BASE:-32950}"
+PER_GPU_BATCH_SIZE="${PER_GPU_BATCH_SIZE:-128}"
+GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-512}"
+GROUP_STAMP="${GROUP_STAMP:-$(date +%Y%m%d_%H%M%S)}"
+OUT_ROOT="${OUT_ROOT:-docs/lta_samples/metrics_20260517/noise_geometry_pilots_len${MAX_LEN}_bs512_ode128_${GROUP_STAMP}}"
+DRIVER_LOG="${DRIVER_LOG:-logs/noise_geometry_pilots_4gpu/${GROUP_STAMP}.log}"
+SUMMARY_JSONL="${SUMMARY_JSONL:-${OUT_ROOT}/summary_results.jsonl}"
+CURVE_CSV="${CURVE_CSV:-${OUT_ROOT}/hit_ratio_curve.csv}"
+mkdir -p "$(dirname "${DRIVER_LOG}")" "${OUT_ROOT}"
+
+cache="${BASE_CACHE}/gpt2_len${MAX_LEN}_train8_compact_overfit"
+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,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"
+  local out_dir="${OUT_ROOT}/${config}/step_${target_step}"
+  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 flowmap \
+    --time_schedule logit_normal \
+    --time_logit_mean -1.5 \
+    --time_logit_std 0.8 \
+    --model_t_mode post \
+    --c_min 1 \
+    --c_max 512 \
+    --late_temp 1.0 \
+    --final_from state \
+    --final_decode argmax
+  python - "$config" "$out_dir" "$N_SAMPLES" "$GLOBAL_BATCH_SIZE" "$MAX_LEN" "$SUMMARY_JSONL" "$CURVE_CSV" "$STOP_EXACT_COUNT" "$STOP_EXACT_REF_COUNT" <<'PY'
+import json
+import sys
+from pathlib import Path
+
+config = sys.argv[1]
+out = Path(sys.argv[2])
+n = int(sys.argv[3])
+global_batch = int(sys.argv[4])
+max_len = int(sys.argv[5])
+summary = Path(sys.argv[6])
+curve = Path(sys.argv[7])
+stop_exact_count = int(sys.argv[8])
+stop_exact_ref_count = int(sys.argv[9])
+row = json.loads((out / "decode_token_acc.jsonl").read_text().splitlines()[-1])
+row["config"] = config
+row["train_views_seen"] = int(row["ckpt_step"]) * global_batch
+row["train_tokens_seen"] = row["train_views_seen"] * max_len
+line = (
+    "RESULT "
+    f"config={config} "
+    f"ckpt_step={row['ckpt_step']} "
+    f"views={row['train_views_seen']} "
+    f"token_acc={row['token_acc_mean']:.4f} "
+    f"exact={row['exact_count']}/{n} "
+    f"exact_refs={row['exact_ref_count']} "
+    f"hits={row['exact_ref_hits']}"
+)
+print(line, flush=True)
+summary.parent.mkdir(parents=True, exist_ok=True)
+with summary.open("a", encoding="utf-8") as f:
+    f.write(json.dumps(row, ensure_ascii=False) + "\n")
+with curve.open("a", encoding="utf-8") as f:
+    f.write(
+        f"{config},{row['ckpt_step']},{row['train_views_seen']},{row['train_tokens_seen']},"
+        f"{row['token_acc_mean']},{row['exact_count']},{row['exact_ref_count']},"
+        f"\"{row['exact_ref_hits']}\"\n"
+    )
+raise SystemExit(
+    0
+    if int(row["exact_count"]) >= stop_exact_count
+    and int(row["exact_ref_count"]) >= stop_exact_ref_count
+    else 1
+)
+PY
+}
+
+reset_noise_defaults() {
+  export MIN_MASK_RATIO=1.0
+  export MAX_MASK_RATIO=1.0
+  export MASK_MIXTURE_LOWK_PROB=0.0
+  export MASK_MIXTURE_ALL_PROB=1.0
+  export LOWK_CLEAN_TOKENS=0
+  export CLEAN_STATE_MODE=onehot
+  export TARGET_LOSS=hard_ce
+  export LINEAR_SOFT_TARGET_POWER=1.0
+  export LINEAR_SOFT_TARGET_MIN_CONF=0.0
+  export LINEAR_SOFT_TARGET_MAX_CONF=1.0
+  export DIRICHLET_CONCENTRATION_MIN=1.0
+  export DIRICHLET_CONCENTRATION_MAX=1024
+  export SIMPLEX_BRIDGE_SAMPLER=dirichlet
+  export ENDPOINT_SEQUENCE_RANDOM_PROB_ALPHA=0.0
+  export CATEGORICAL_WRONG_BASIN_TOKEN_IDS=""
+  export CATEGORICAL_WRONG_BASIN_PROB=0.0
+  export CATEGORICAL_WRONG_UNIGRAM_PROB=0.0
+  export CATEGORICAL_WRONG_UNIFORM_PROB=0.0
+  export CATEGORICAL_WRONG_BASIN_SHARED_PROB=0.0
+  export CATEGORICAL_WRONG_UNIGRAM_SHARED_PROB=0.0
+  export LOGISTIC_NORMAL_SIGMA_MIN=0.1
+  export LOGISTIC_NORMAL_SIGMA_MAX=1.0
+  export LOGISTIC_NORMAL_TAU_MIN=1.0
+  export LOGISTIC_NORMAL_TAU_MAX=1.0
+}
+
+configure_noise() {
+  local config="$1"
+  reset_noise_defaults
+  case "${config}" in
+    allcorrupt_fullvocab_dirC1_1024)
+      ;;
+    allcorrupt_seqrand1_dirC1_1024)
+      export ENDPOINT_SEQUENCE_RANDOM_PROB_ALPHA=1.0
+      ;;
+    allcorrupt_highC64_4096)
+      export DIRICHLET_CONCENTRATION_MIN=64
+      export DIRICHLET_CONCENTRATION_MAX=4096
+      ;;
+    allcorrupt_logistic_sig0p05_0p5)
+      export SIMPLEX_BRIDGE_SAMPLER=logistic_normal_linear_mean
+      export LOGISTIC_NORMAL_SIGMA_MIN=0.05
+      export LOGISTIC_NORMAL_SIGMA_MAX=0.5
+      ;;
+    allcorrupt_seqrand1_highC64_4096)
+      export ENDPOINT_SEQUENCE_RANDOM_PROB_ALPHA=1.0
+      export DIRICHLET_CONCENTRATION_MIN=64
+      export DIRICHLET_CONCENTRATION_MAX=4096
+      ;;
+    allcorrupt_unigram_shared0p5)
+      export CATEGORICAL_WRONG_UNIGRAM_PROB=1.0
+      export CATEGORICAL_WRONG_UNIGRAM_SHARED_PROB=0.5
+      ;;
+    *)
+      echo "unknown config: ${config}" >&2
+      return 2
+      ;;
+  esac
+}
+
+configs=(
+  allcorrupt_fullvocab_dirC1_1024
+  allcorrupt_seqrand1_dirC1_1024
+  allcorrupt_highC64_4096
+  allcorrupt_logistic_sig0p05_0p5
+  allcorrupt_seqrand1_highC64_4096
+  allcorrupt_unigram_shared0p5
+)
+
+echo "[noise-pilot] start stamp=${GROUP_STAMP} len=${MAX_LEN} vocab=${vocab_size} out=${OUT_ROOT}" | tee -a "${DRIVER_LOG}"
+
+round_idx=0
+while :; do
+  round_idx=$((round_idx + 1))
+  active=0
+  echo "[noise-pilot] round=${round_idx} $(date)" | tee -a "${DRIVER_LOG}"
+
+  idx=0
+  for config in "${configs[@]}"; do
+    idx=$((idx + 1))
+    run_name="train8_noisegeo_len${MAX_LEN}_${config}_${GROUP_STAMP}"
+    done_flag="${OUT_ROOT}/${config}/DONE"
+    if [[ -f "${done_flag}" ]]; then
+      echo "[noise-pilot] skip done config=${config}" | tee -a "${DRIVER_LOG}"
+      continue
+    fi
+    step_now="$(latest_step "${run_name}")"
+    if [[ "${step_now}" -ge "${MAX_TOTAL_STEPS}" ]]; then
+      echo "[noise-pilot] capped config=${config} step=${step_now}" | tee -a "${DRIVER_LOG}"
+      continue
+    fi
+    active=1
+    target_step=$((step_now + STEP_CHUNK))
+    if [[ "${target_step}" -gt "${MAX_TOTAL_STEPS}" ]]; then
+      target_step="${MAX_TOTAL_STEPS}"
+    fi
+    resume_path=""
+    if [[ -f "runs/${run_name}/latest.pt" ]]; then
+      resume_path="runs/${run_name}/latest.pt"
+    fi
+
+    configure_noise "${config}"
+    echo "[noise-pilot] train config=${config} from=${step_now} to=${target_step} C=${DIRICHLET_CONCENTRATION_MIN}->${DIRICHLET_CONCENTRATION_MAX} seq_alpha=${ENDPOINT_SEQUENCE_RANDOM_PROB_ALPHA} sampler=${SIMPLEX_BRIDGE_SAMPLER}" | 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:-0.002}" \
+    WEIGHT_DECAY="${WEIGHT_DECAY:-0.1}" \
+    MUON_IMPL="${MUON_IMPL:-legacy}" \
+    OUTPUT_WEIGHT_DECAY="${OUTPUT_WEIGHT_DECAY:--1}" \
+    RUN_NAME="${run_name}" \
+    RESUME_PATH="${resume_path}" \
+    bash scripts/launch_lta_owt_gpt2_softendpoint_mn_pilot_4gpu.sh
+
+    echo "[noise-pilot] eval config=${config} step=${target_step}" | tee -a "${DRIVER_LOG}"
+    if eval_latest "${config}" "${run_name}" "${target_step}" | tee -a "${DRIVER_LOG}"; then
+      echo "[noise-pilot] DONE config=${config} step=${target_step}" | tee -a "${DRIVER_LOG}"
+      touch "${done_flag}"
+    else
+      echo "[noise-pilot] continue config=${config} step=${target_step}" | tee -a "${DRIVER_LOG}"
+    fi
+  done
+
+  if [[ "${active}" -eq 0 ]]; then
+    echo "[noise-pilot] all configs done/capped $(date)" | tee -a "${DRIVER_LOG}"
+    break
+  fi
+done

LTA_openwebtext_dualt/scripts/sweep_categorical_c1024_diffusion_finalsample_20260506.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import importlib.util
+import sys
+from pathlib import Path
+
+
+BASE_PATH = Path(__file__).with_name("sweep_categorical_c1024_rolling_noise_20260506.py")
+spec = importlib.util.spec_from_file_location("sweep_categorical_c1024_rolling_noise_20260506", BASE_PATH)
+if spec is None or spec.loader is None:
+    raise RuntimeError(f"Could not load {BASE_PATH}")
+sweep = importlib.util.module_from_spec(spec)
+sys.modules[spec.name] = sweep
+spec.loader.exec_module(sweep)
+
+
+def finalsample_configs():
+    C = sweep.DecodeCfg
+    return [
+        C("fs_t1p3_eta1_blend_ft0p35", endpoint_temp=1.3, final_sample=True, final_temp=0.35),
+        C("fs_t1p3_eta1_blend_ft0p50", endpoint_temp=1.3, final_sample=True, final_temp=0.50),
+        C("fs_t1p3_eta1_blend_ft0p70", endpoint_temp=1.3, final_sample=True, final_temp=0.70),
+        C("fs_t1p3_eta1_blend_ft1p00", endpoint_temp=1.3, final_sample=True, final_temp=1.00),
+        C("fs_t1p3_eta1_blend_ft1p30", endpoint_temp=1.3, final_sample=True, final_temp=1.30),
+        C("fs_t1p3_eta1_blend_ft1p60", endpoint_temp=1.3, final_sample=True, final_temp=1.60),
+        C("fs_t1p2_eta1_blend_ft0p50", endpoint_temp=1.2, final_sample=True, final_temp=0.50),
+        C("fs_t1p2_eta1_blend_ft0p70", endpoint_temp=1.2, final_sample=True, final_temp=0.70),
+        C("fs_t1p2_eta1_blend_ft1p00", endpoint_temp=1.2, final_sample=True, final_temp=1.00),
+        C("fs_t1p1_eta1_blend_ft0p70", endpoint_temp=1.1, final_sample=True, final_temp=0.70),
+        C("fs_t1p3_eta1_state_ft0p35", endpoint_temp=1.3, final_from="state", final_sample=True, final_temp=0.35),
+        C("fs_t1p3_eta1_state_ft0p50", endpoint_temp=1.3, final_from="state", final_sample=True, final_temp=0.50),
+        C(
+            "fs_sp1p15_sem1p7_t1p3_eta1_blend_ft0p50",
+            endpoint_temp=1.3,
+            support_power=1.15,
+            semantic_power=1.7,
+            final_sample=True,
+            final_temp=0.50,
+        ),
+    ]
+
+
+sweep.default_configs = finalsample_configs
+
+
+if __name__ == "__main__":
+    sweep.main()