#!/usr/bin/env python3 """AGILLM4.1 mainline single-file trainer/inference runtime. AGILLM4.1 is the promoted AGILLM4 mainline evolved from the AGILLM3.5 prototype, and it is larger than AGILLM3/AGILLM3.5. Resumed checkpoints are the source of truth for the exact architecture, with AGILLM4 presets available for fresh starts. This file is mechanically folded from AGILLM4 plus compatibility patches: - DeepSeek-V4-Pro tokenizer/checkpoint support by default - DeepSeek-V3.2 legacy compatibility support through the agillm35 shim - AR + SAT checkpoint schema compatibility; NAT can be disabled with --agillm3_compat - DiffusionBlock training support and optional async side-update ingestion """ from __future__ import annotations # Single-file module alias: helper code still imports the historical module names. import sys as _agillm41_sys _agillm41_sys.modules.setdefault("nB300_agillm4", _agillm41_sys.modules[__name__]) _agillm41_sys.modules.setdefault("agillm35", _agillm41_sys.modules[__name__]) _agillm41_sys.modules.setdefault("agillm41", _agillm41_sys.modules[__name__]) _agillm41_sys.modules.setdefault("dblocks_train", _agillm41_sys.modules[__name__]) _agillm41_sys.modules.setdefault("fused_ce", _agillm41_sys.modules[__name__]) _agillm41_sys.modules.setdefault("anchor_memory", _agillm41_sys.modules[__name__]) import types as _agillm41_types # ===== BEGIN agillm_checkpoint_provenance.py (folded) ===== _AGILLM_CHECKPOINT_PROVENANCE_SOURCE = '"""agillm_checkpoint_provenance.py — git-style lineage tracking for checkpoints.\n\nEvery full checkpoint (.pt) carries a `provenance` dict that records:\n - warmstart source & its provenance (chained like git commits)\n - training step, tokens seen, loss (total + per-head)\n - training script name + SHA256, full argv\n - creation time, hostname, PID, GPU metrics\n - inference samples (3 short generations from the model)\n - dataset provenance snapshot\n\nCLI usage:\n python3 agillm_checkpoint_provenance.py show \n python3 agillm_checkpoint_provenance.py lineage \n python3 agillm_checkpoint_provenance.py compare \n"""\n\nfrom __future__ import annotations\n\nimport argparse\nimport hashlib\nimport json\nimport os\nimport platform\nimport re\nimport subprocess\nimport sys\nimport time\nimport pathlib\nfrom typing import Any, Dict, List, Optional, Tuple\n\n# ---------------------------------------------------------------------------\n# Schema key\n# ---------------------------------------------------------------------------\nPROVENANCE_KEY = "agillm43_provenance"\nPROVENANCE_SCHEMA_VERSION = 1\n\n# ---------------------------------------------------------------------------\n# Provenance dict shape\n# ---------------------------------------------------------------------------\n"""\nprovenance = {\n "schema_version": 1,\n "checkpoint_type": "full" | "delta",\n\n # Identity\n "created_at_iso": "2026-06-23T03:14:00Z",\n "created_at_unix": 1750000000.0,\n "hostname": "agillm43-boxa",\n "pid": 1372905,\n "lane": "a0",\n\n # Training state\n "step": 13886,\n "seen_tok": 850000000,\n "loss": 2.345,\n "loss_ar": 2.1,\n "loss_sat": 0.15,\n "loss_nat": 0.095,\n "batch_size": 56,\n "block_size": 1536,\n\n # Source\n "train_script": "agillm41.py",\n "train_script_sha256": "abc123...",\n "train_argv": "--warmstart_from /workspace/... --preset agillm4_floor ...",\n\n # Warmstart chain (like git parent)\n "warmstart_source_path": "/workspace/agillm4_v100_master_ckpts/pretrain_step02182564.pt",\n "warmstart_source_provenance": { ... } or None,\n\n # Config snapshot\n "cfg_keys": ["dmodel", "layers", "heads", ...],\n\n # Inference samples (3 short generations)\n "inference_samples": [\n {"prompt": "The meaning of life is", "generation": " to find", "tokens": 5},\n ...\n ],\n\n # GPU state at save time\n "gpu": {\n "allocated_gb": 30.5,\n "reserved_gb": 31.2,\n "peak_allocated_gb": 32.0,\n },\n\n # Dataset provenance fragment\n "dataset_provenance": { ... },\n\n # Tokenizer info\n "tokenizer_id": "...",\n}\n"""\n\n\n# ---------------------------------------------------------------------------\n# Utilities\n# ---------------------------------------------------------------------------\n\ndef _iso_now() -> str:\n return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())\n\n\ndef _sha256_file(path: pathlib.Path) -> str:\n h = hashlib.sha256()\n with open(path, "rb") as f:\n while True:\n chunk = f.read(1 << 20)\n if not chunk:\n break\n h.update(chunk)\n return h.hexdigest()\n\n\ndef _sha256_bytes(data: bytes) -> str:\n return hashlib.sha256(data).hexdigest()\n\n\ndef _gpu_metrics() -> dict:\n """Collect GPU memory usage if CUDA is available."""\n try:\n import torch\n if not torch.cuda.is_available():\n return {}\n return {\n "allocated_gb": round(torch.cuda.memory_allocated() / (1024**3), 2),\n "reserved_gb": round(torch.cuda.memory_reserved() / (1024**3), 2),\n "peak_allocated_gb": round(torch.cuda.max_memory_allocated() / (1024**3), 2),\n }\n except Exception:\n return {}\n\n\ndef _script_sha256() -> Tuple[str, str]:\n """SHA256 of the running training script. Returns (basename, hexdigest)."""\n try:\n main = sys.modules.get("__main__")\n if main and hasattr(main, "__file__") and main.__file__:\n p = pathlib.Path(main.__file__).resolve()\n return p.name, _sha256_file(p)\n except Exception:\n pass\n return ("", "")\n\n\ndef _script_argv() -> str:\n return " ".join(sys.argv)\n\n\ndef _read_proc_cmdline(pid: str = "self") -> str:\n try:\n raw = pathlib.Path("/proc") / str(pid) / "cmdline"\n data = raw.read_bytes()\n return " ".join(part.decode("utf-8", "replace") for part in data.split(b"\\0") if part)\n except Exception:\n return ""\n\n\ndef _safe_env_snapshot() -> dict:\n """Capture useful launch env without leaking tokens or credentials."""\n prefixes = ("AGILLM", "CUDA_", "HF_HUB_", "HF_DATASETS_", "PYTORCH_", "OMP_", "MKL_")\n allow = {\n "CUDA_VISIBLE_DEVICES",\n "HF_HUB_DISABLE_XET",\n "HF_DATASETS_TRUST_REMOTE_CODE",\n "PYTORCH_CUDA_ALLOC_CONF",\n "OMP_NUM_THREADS",\n "MKL_NUM_THREADS",\n }\n secret_fragments = ("TOKEN", "SECRET", "PASSWORD", "PASSWD", "KEY", "CREDENTIAL", "AUTH", "COOKIE")\n out = {}\n for key, value in sorted(os.environ.items()):\n if not (key in allow or key.startswith(prefixes)):\n continue\n if any(fragment in key.upper() for fragment in secret_fragments):\n out[key] = ""\n else:\n out[key] = str(value)[:2048]\n return out\n\n\ndef _redact_text(text: str) -> str:\n secret_fragments = ("TOKEN", "SECRET", "PASSWORD", "PASSWD", "API_KEY", "AUTH", "COOKIE")\n lines = []\n for line in str(text).splitlines()[:240]:\n upper = line.upper()\n if any(fragment in upper for fragment in secret_fragments):\n lines.append("")\n else:\n lines.append(line[:4096])\n return "\\n".join(lines)\n\n\ndef _launch_metadata() -> dict:\n meta = {\n "schema": "agillm.launch.v1",\n "argv": list(sys.argv),\n "argv_string": _script_argv(),\n "cwd": "",\n "pid": os.getpid(),\n "ppid": os.getppid(),\n "proc_cmdline": _read_proc_cmdline("self"),\n "parent_proc_cmdline": _read_proc_cmdline(str(os.getppid())),\n "env": _safe_env_snapshot(),\n }\n try:\n meta["cwd"] = str(pathlib.Path.cwd())\n except Exception:\n pass\n launch_script = os.environ.get("AGILLM43_LAUNCH_SCRIPT") or os.environ.get("AGILLM_LAUNCH_SCRIPT") or ""\n launch_command = os.environ.get("AGILLM43_LAUNCH_COMMAND") or os.environ.get("AGILLM_LAUNCH_COMMAND") or ""\n if launch_command:\n meta["launch_command"] = _redact_text(launch_command)\n if launch_script:\n sp = pathlib.Path(launch_script)\n info = {"path": str(sp)}\n try:\n if sp.exists() and sp.is_file():\n info["size_bytes"] = sp.stat().st_size\n info["sha256"] = _sha256_file(sp)\n info["preview_redacted"] = _redact_text(sp.read_text(errors="replace"))\n except Exception as exc:\n info["error"] = str(exc)\n meta["launch_script"] = info\n return meta\n\n\ndef _infer_samples(core, ar_h, sat_h, tok, device: str, prompt_texts: List[str],\n max_new: int = 32, temperature: float = 0.5, top_k: int = 20) -> List[dict]:\n """Generate a few short inference samples from the model.\n\n This is called at save time with gradients off (torch.no_grad).\n If anything fails, returns an empty list — never crashes a save.\n """\n samples = []\n try:\n import torch\n core.eval()\n ar_h.eval()\n if sat_h is not None:\n sat_h.eval()\n\n for prompt in prompt_texts:\n try:\n input_ids = tok.encode(prompt, return_tensors="pt").to(device)\n if input_ids.numel() == 0:\n continue\n generated = input_ids.clone()\n for _ in range(max_new):\n with torch.no_grad():\n h = core(generated, None)\n logits = ar_h(h[:, -1:])\n probs = torch.softmax(logits[:, -1] / max(temperature, 1e-8), dim=-1)\n if top_k > 0:\n vals, idxs = torch.topk(probs, min(top_k, probs.size(-1)))\n probs = torch.zeros_like(probs).scatter_(-1, idxs, vals)\n next_id = torch.multinomial(probs, 1)\n generated = torch.cat([generated, next_id], dim=1)\n if next_id.item() == 0: # EOS\n break\n text = tok.decode(generated[0].tolist(), skip_special_tokens=True)\n new_tokens = generated.size(1) - input_ids.size(1)\n samples.append({\n "prompt": prompt,\n "generation": text[len(prompt):] if text.startswith(prompt) else text,\n "tokens": new_tokens,\n })\n except Exception:\n samples.append({"prompt": prompt, "generation": "", "tokens": 0})\n except Exception:\n pass\n return samples\n\n\n# ---------------------------------------------------------------------------\n# Core provenance construction\n# ---------------------------------------------------------------------------\n\ndef _step_from_text(text: Optional[str]) -> Optional[int]:\n m = re.search(r"step(\\d+)", str(text or ""))\n return int(m.group(1)) if m else None\n\n\ndef _origin_step_from_provenance(prov: Optional[dict]) -> int:\n if not isinstance(prov, dict):\n return 0\n for key in ("global_origin_step", "warmstart_base_step"):\n try:\n value = int(prov.get(key) or 0)\n except Exception:\n value = 0\n if value > 0:\n return value\n parent = prov.get("warmstart_source_path") or prov.get("source_path") or ""\n parent_step = _step_from_text(parent)\n if parent_step and parent_step > 0: # AGILLM-LINEAGE-FIX 20260702\n return int(parent_step)\n return 0\n\n\ndef _origin_seen_tok_from_provenance(prov: Optional[dict]) -> int:\n if not isinstance(prov, dict):\n return 0\n for key in ("global_origin_seen_tok", "warmstart_base_seen_tok"):\n try:\n value = int(prov.get(key) or 0)\n except Exception:\n value = 0\n if value > 0:\n return value\n return 0\n\n\ndef collect(args, *, step: int, seen_tok: int, loss: float,\n loss_ar: Optional[float] = None, loss_sat: Optional[float] = None,\n loss_nat: Optional[float] = None,\n batch_size: int = 0, block_size: int = 0,\n warmstart_source_path: Optional[str] = None,\n warmstart_source_provenance: Optional[dict] = None,\n dataset_provenance: Optional[dict] = None,\n lane: str = "",\n inference_samples: Optional[list] = None,\n checkpoint_type: str = "full",\n _sample_core=None, _sample_ar=None, _sample_sat=None,\n _sample_tok=None, _sample_device: str = "",\n _sample_prompts: Optional[List[str]] = None) -> dict:\n """Build a provenance dict to embed in the checkpoint."""\n\n script_name, script_sha = _script_sha256()\n\n prov: dict = {\n "schema_version": PROVENANCE_SCHEMA_VERSION,\n "checkpoint_type": checkpoint_type,\n "created_at_iso": _iso_now(),\n "created_at_unix": time.time(),\n "hostname": platform.node(),\n "pid": os.getpid(),\n "lane": lane or "",\n "step": int(step),\n "seen_tok": int(seen_tok),\n "loss": float(loss),\n "batch_size": int(batch_size),\n "block_size": int(block_size),\n "train_script": script_name,\n "train_argv": _script_argv(),\n "launch": _launch_metadata(),\n "gpu": _gpu_metrics(),\n }\n\n if script_sha:\n prov["train_script_sha256"] = script_sha\n\n if loss_ar is not None:\n prov["loss_ar"] = float(loss_ar)\n if loss_sat is not None:\n prov["loss_sat"] = float(loss_sat)\n if loss_nat is not None:\n prov["loss_nat"] = float(loss_nat)\n\n source_step = _step_from_text(warmstart_source_path)\n origin_step = _origin_step_from_provenance(warmstart_source_provenance)\n origin_seen_tok = _origin_seen_tok_from_provenance(warmstart_source_provenance)\n if not origin_step and source_step and source_step > 0: # AGILLM-LINEAGE-FIX 20260702\n origin_step = int(source_step)\n\n prov["local_step"] = int(step)\n if source_step is not None:\n prov["warmstart_source_step"] = int(source_step)\n prov["global_origin_step"] = int(origin_step or 0)\n prov["warmstart_base_step"] = int(origin_step or 0)\n prov["effective_global_step"] = int((origin_step + int(step)) if origin_step else int(step))\n prov["global_origin_seen_tok"] = int(origin_seen_tok or 0)\n prov["warmstart_base_seen_tok"] = int(origin_seen_tok or 0)\n prov["effective_seen_tok"] = int(int(origin_seen_tok or 0) + int(seen_tok))\n\n if warmstart_source_path:\n prov["warmstart_source_path"] = str(warmstart_source_path)\n if warmstart_source_provenance:\n prov["warmstart_source_provenance"] = warmstart_source_provenance\n\n if dataset_provenance:\n prov["dataset_provenance"] = dataset_provenance\n\n if inference_samples is not None:\n prov["inference_samples"] = inference_samples\n elif _sample_core is not None and _sample_ar is not None and _sample_tok is not None:\n try:\n prompts = _sample_prompts or ["The meaning of", "def hello():", "2 + 2 ="]\n prov["inference_samples"] = _infer_samples(\n _sample_core, _sample_ar, _sample_sat,\n _sample_tok, _sample_device or "cpu", prompts, max_new=12)\n except Exception:\n prov["inference_samples"] = []\n\n return prov\n\n\ndef embed(state_dict: dict, provenance: dict) -> dict:\n """Embed provenance into the checkpoint state dict (mutates + returns)."""\n state_dict[PROVENANCE_KEY] = provenance\n return state_dict\n\n\n# ---------------------------------------------------------------------------\n# Extraction (lightweight — only reads provenance from .pt wrapper)\n# ---------------------------------------------------------------------------\n\ndef extract(path: pathlib.Path) -> Optional[dict]:\n """Extract the provenance dict from a saved .pt checkpoint.\n\n This reads only the top-level wrapper, not the full model weights.\n For zstd-wrapped checkpoints, it only decompresses enough to find the\n provenance key.\n\n Returns None if no provenance is found.\n """\n try:\n import torch\n # The checkpoint may be zstd-wrapped. Load the wrapper first.\n wrapper = torch.load(str(path), map_location="cpu", weights_only=False)\n if not isinstance(wrapper, dict):\n return None\n\n # If zstd-wrapped, decompress and get inner dict\n inner = wrapper\n if wrapper.get("__agillm43_payload_codec__") == "agillm43_zstd_torch_v1":\n import zstandard as zstd\n raw = zstd.ZstdDecompressor().decompress(bytes(wrapper["payload"].tolist()))\n import io\n inner = torch.load(io.BytesIO(raw), map_location="cpu", weights_only=False)\n\n if not isinstance(inner, dict):\n return None\n\n provenance = inner.get(PROVENANCE_KEY)\n if provenance is not None:\n return provenance\n\n # Fallback: check for sidecar\n sidecar = path.with_suffix(".provenance.json")\n if sidecar.exists():\n return json.loads(sidecar.read_text())\n\n return None\n except Exception:\n return None\n\n\ndef extract_provenance_sidecar(ckpt_path: pathlib.Path) -> Optional[dict]:\n """Read the .provenance.json sidecar without touching the .pt at all."""\n sidecar = ckpt_path.with_suffix(".provenance.json")\n if sidecar.exists():\n try:\n return json.loads(sidecar.read_text())\n except Exception:\n pass\n return None\n\n\ndef write_sidecar(ckpt_path: pathlib.Path, provenance: dict) -> None:\n """Write .provenance.json sidecar beside the checkpoint."""\n sidecar = ckpt_path.with_suffix(".provenance.json")\n tmp = sidecar.with_suffix(".provenance.json.tmp")\n try:\n tmp.write_text(json.dumps(provenance, indent=2, sort_keys=True) + "\\n")\n tmp.replace(sidecar)\n except Exception as exc:\n print(f"[provenance] WARNING: failed to write sidecar {sidecar}: {exc}")\n\n\n# ---------------------------------------------------------------------------\n# Display / CLI\n# ---------------------------------------------------------------------------\n\ndef format_provenance(prov: dict, indent: int = 0) -> str:\n """Format a provenance dict as a readable block."""\n pad = " " * indent\n lines = [f"{pad}┌── Checkpoint Provenance ──"]\n if not prov:\n return f"{pad}└── (no provenance)"\n\n def kv(k, v, default="—"):\n val = v if v is not None else default\n return f"{pad} {k}: {val}"\n\n lines.append(kv("Schema version", prov.get("schema_version")))\n lines.append(kv("Type", prov.get("checkpoint_type")))\n lines.append(kv("Step", prov.get("step")))\n lines.append(kv("Tokens seen", f"{prov.get(\'seen_tok\', 0):,}"))\n lines.append(kv("Loss", prov.get("loss")))\n if prov.get("loss_ar") is not None:\n lines.append(kv(" ├ AR loss", prov["loss_ar"]))\n if prov.get("loss_sat") is not None:\n lines.append(kv(" ├ SAT loss", prov["loss_sat"]))\n if prov.get("loss_nat") is not None:\n lines.append(kv(" └ NAT loss", prov["loss_nat"]))\n lines.append(kv("Batch / Block", f"{prov.get(\'batch_size\')} / {prov.get(\'block_size\')}"))\n lines.append(kv("Created (ISO)", prov.get("created_at_iso")))\n lines.append(kv("Hostname", prov.get("hostname")))\n lines.append(kv("PID", prov.get("pid")))\n lines.append(kv("Lane", prov.get("lane", "—")))\n lines.append(kv("Train script", prov.get("train_script")))\n if prov.get("train_script_sha256"):\n lines.append(kv(" └ SHA256", prov["train_script_sha256"][:16] + "..."))\n gpu = prov.get("gpu", {})\n if gpu:\n lines.append(kv("GPU alloc/resrv/peak",\n f"{gpu.get(\'allocated_gb\', \'?\')}G / {gpu.get(\'reserved_gb\', \'?\')}G / {gpu.get(\'peak_allocated_gb\', \'?\')}G"))\n\n ws = prov.get("warmstart_source_path")\n if ws:\n lines.append(kv("Warmstart source", ws))\n wprov = prov.get("warmstart_source_provenance")\n if wprov:\n lines.append(f"{pad} └ step={wprov.get(\'step\', \'?\')} loss={wprov.get(\'loss\', \'?\')}")\n\n samples = prov.get("inference_samples", [])\n if samples:\n lines.append(f"{pad}Inference samples ({len(samples)}):")\n for i, s in enumerate(samples):\n gen = s.get("generation", "")\n if len(gen) > 60:\n gen = gen[:60] + "..."\n lines.append(f"{pad} [{i}] prompt={s.get(\'prompt\',\'\')!r}")\n lines.append(f"{pad} → {gen!r} ({s.get(\'tokens\', 0)} tokens)")\n\n lines.append(f"{pad}└──")\n return "\\n".join(lines)\n\n\ndef show_lineage(path: pathlib.Path, max_depth: int = 32) -> List[dict]:\n """Walk the provenance chain (like git log) and return ordered list [oldest..newest]."""\n chain: List[dict] = []\n seen = set()\n current = path.resolve() if path.exists() else path\n\n for _ in range(max_depth):\n prov = extract(current)\n if prov is None:\n break\n\n key = str(current)\n if key in seen:\n break\n seen.add(key)\n\n entry = prov.copy()\n entry["_checkpoint_path"] = str(current)\n chain.append(entry)\n\n # Walk to warmstart parent\n ws = prov.get("warmstart_source_path")\n if not ws:\n break\n wprov = prov.get("warmstart_source_provenance")\n if not wprov:\n break\n current = pathlib.Path(ws)\n # Avoid infinite loop if parent points to itself\n if str(current) == key:\n break\n else:\n chain.append({"_checkpoint_path": f"(truncated at {max_depth} hops)"})\n\n chain.reverse() # oldest first\n return chain\n\n\ndef format_lineage(chain: List[dict]) -> str:\n """Format a lineage chain as a readable tree."""\n lines = ["Checkpoint Lineage (oldest → newest):", ""]\n for i, entry in enumerate(chain):\n path = entry.get("_checkpoint_path", "?")\n step = entry.get("step", "?")\n loss = entry.get("loss", "?")\n iso = entry.get("created_at_iso", "?")\n ws = entry.get("warmstart_source_path", "")\n marker = "●" if i == len(chain) - 1 else "│" if i < len(chain) - 1 else "○"\n lines.append(f" {marker} step={step} loss={loss} {iso}")\n lines.append(f" │ {path}")\n if ws and i < len(chain) - 1:\n lines.append(f" │ warmstart ← {pathlib.Path(ws).name}")\n lines.append("")\n return "\\n".join(lines)\n\n\n# ---------------------------------------------------------------------------\n# CLI\n# ---------------------------------------------------------------------------\n\ndef _cmd_show(args_cli):\n path = pathlib.Path(args_cli.checkpoint)\n if not path.exists():\n print(f"ERROR: {path} not found")\n sys.exit(1)\n prov = extract(path)\n if prov is None:\n prov = extract_provenance_sidecar(path)\n if prov is None:\n print(f"No provenance found in {path}")\n sys.exit(1)\n print(format_provenance(prov))\n if args_cli.verbose:\n print("\\nFull provenance JSON:")\n print(json.dumps(prov, indent=2, sort_keys=True))\n\n\ndef _cmd_lineage(args_cli):\n path = pathlib.Path(args_cli.checkpoint)\n if not path.exists():\n print(f"ERROR: {path} not found")\n sys.exit(1)\n chain = show_lineage(path, max_depth=args_cli.max_depth)\n print(format_lineage(chain))\n\n\ndef _cmd_compare(args_cli):\n a = pathlib.Path(args_cli.checkpoint_a)\n b = pathlib.Path(args_cli.checkpoint_b)\n for p, label in [(a, "A"), (b, "B")]:\n if not p.exists():\n print(f"ERROR: {label}={p} not found")\n sys.exit(1)\n\n pa = extract(a) or {}\n pb = extract(b) or {}\n\n def safe(key, d, default="—"):\n return d.get(key, default)\n\n print(f"Compare: {a.name} vs {b.name}")\n print()\n keys = ["step", "seen_tok", "loss", "loss_ar", "loss_sat", "loss_nat",\n "batch_size", "block_size", "created_at_iso", "hostname", "lane"]\n for k in keys:\n va = safe(k, pa)\n vb = safe(k, pb)\n changed = " ←" if str(va) != str(vb) else ""\n print(f" {k:20s} {str(va):>20s} {str(vb):>20s}{changed}")\n\n sa = pa.get("inference_samples", [])\n sb = pb.get("inference_samples", [])\n if sa or sb:\n print()\n print(f" Inference samples: A={len(sa)} B={len(sb)}")\n\n\ndef main():\n parser = argparse.ArgumentParser(\n description="agillm checkpoint provenance — git for checkpoints",\n formatter_class=argparse.RawDescriptionHelpFormatter,\n epilog=__doc__,\n )\n sub = parser.add_subparsers(dest="command")\n\n p_show = sub.add_parser("show", help="Show provenance for a checkpoint")\n p_show.add_argument("checkpoint", type=str, help="Path to .pt checkpoint")\n p_show.add_argument("-v", "--verbose", action="store_true", help="Also dump full JSON")\n\n p_lineage = sub.add_parser("lineage", help="Show full warmstart chain (git log)")\n p_lineage.add_argument("checkpoint", type=str, help="Path to .pt checkpoint")\n p_lineage.add_argument("--max-depth", type=int, default=32, help="Max hops to follow")\n\n p_cmp = sub.add_parser("compare", help="Compare two checkpoints")\n p_cmp.add_argument("checkpoint_a", type=str)\n p_cmp.add_argument("checkpoint_b", type=str)\n\n args_cli = parser.parse_args()\n if args_cli.command == "show":\n _cmd_show(args_cli)\n elif args_cli.command == "lineage":\n _cmd_lineage(args_cli)\n elif args_cli.command == "compare":\n _cmd_compare(args_cli)\n else:\n parser.print_help()\n sys.exit(1)\n\n\nif __name__ == "__main__":\n main()\n' _agillm_provenance = _agillm41_types.ModuleType("agillm_checkpoint_provenance") _agillm_provenance.__file__ = __file__ + "#agillm_checkpoint_provenance" exec(compile(_AGILLM_CHECKPOINT_PROVENANCE_SOURCE, _agillm_provenance.__file__, "exec"), _agillm_provenance.__dict__) _agillm41_sys.modules.setdefault("agillm_checkpoint_provenance", _agillm_provenance) # ===== END agillm_checkpoint_provenance.py (folded) ===== # ===== BEGIN anchor_memory.py ===== #!/usr/bin/env python3 from dataclasses import dataclass import torch import torch.nn as nn import torch.nn.functional as F @dataclass class AnchorMemoryConfig: d_model: int heads: int anchor_stride: int = 256 max_anchors: int = 2048 dropout: float = 0.0 class AnchorCompressor(nn.Module): """Compress local token spans into trainable anchor vectors.""" def __init__(self, d_model: int, anchor_stride: int): super().__init__() self.anchor_stride = anchor_stride self.score = nn.Linear(d_model, 1) self.mix = nn.Sequential( nn.LayerNorm(d_model), nn.Linear(d_model, 4 * d_model), nn.GELU(), nn.Linear(4 * d_model, d_model), ) def forward(self, x: torch.Tensor) -> torch.Tensor: bsz, seq, dim = x.shape pad = (-seq) % self.anchor_stride if pad: x = F.pad(x, (0, 0, 0, pad)) chunks = x.view(bsz, -1, self.anchor_stride, dim) weights = self.score(chunks).softmax(dim=2) pooled = (chunks * weights).sum(dim=2) return pooled + self.mix(pooled) class AnchorMemoryLayer(nn.Module): """Local-token stream reads from a bounded bank of learned anchors.""" def __init__(self, cfg: AnchorMemoryConfig): super().__init__() self.cfg = cfg self.compress = AnchorCompressor(cfg.d_model, cfg.anchor_stride) self.q_ln = nn.LayerNorm(cfg.d_model) self.mem_ln = nn.LayerNorm(cfg.d_model) self.read = nn.MultiheadAttention( cfg.d_model, cfg.heads, dropout=cfg.dropout, batch_first=True, ) self.gate = nn.Sequential(nn.Linear(2 * cfg.d_model, cfg.d_model), nn.Sigmoid()) self.out_ln = nn.LayerNorm(cfg.d_model) def forward( self, x: torch.Tensor, memory: torch.Tensor | None = None, *, detach_memory: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: new_anchors = self.compress(x) if detach_memory: new_anchors = new_anchors.detach() if memory is None: bank = new_anchors else: bank = torch.cat([memory, new_anchors], dim=1) if bank.size(1) > self.cfg.max_anchors: bank = bank[:, -self.cfg.max_anchors :] recalled, _ = self.read(self.q_ln(x), self.mem_ln(bank), self.mem_ln(bank), need_weights=False) gate = self.gate(torch.cat([x, recalled], dim=-1)) mixed = x + gate * recalled return self.out_ln(mixed), bank def smoke_test() -> None: cfg = AnchorMemoryConfig(d_model=128, heads=8, anchor_stride=32, max_anchors=64) layer = AnchorMemoryLayer(cfg) x = torch.randn(2, 256, 128) y, memory = layer(x) assert y.shape == x.shape assert memory.shape == (2, 8, 128) y2, memory2 = layer(x, memory) assert y2.shape == x.shape assert memory2.shape == (2, 16, 128) print("anchor_memory smoke OK", y.shape, memory2.shape) # ===== END anchor_memory.py ===== # ===== BEGIN fused_ce.py ===== """Fused cross-entropy: streams over the VOCAB dimension (online-softmax) so the [N x V] logit matrix is NEVER materialized -- only [N x vchunk]. Custom backward recomputes softmax per vocab-chunk (grad = softmax - onehot). This is the DiffusionBlocks 'process in chunks, don't hold the whole thing' idea applied to the output head instead of network depth.""" import torch class FusedCE(torch.autograd.Function): @staticmethod def forward(ctx, h, W, tgt, vchunk=16384): with torch.cuda.amp.autocast(enabled=True): hf = h.float() Wf = W.float() N, d = h.shape V = W.shape[0] m = torch.full((N,), -1e30, device=h.device, dtype=torch.float32) s = torch.zeros(N, device=h.device, dtype=torch.float32) zt = torch.zeros(N, device=h.device, dtype=torch.float32) for c in range(0, V, vchunk): lg = hf @ Wf[c:c+vchunk].T # [N,vchunk] transient only cm = lg.max(1).values nm = torch.maximum(m, cm) s = s * torch.exp(m - nm) + torch.exp(lg - nm[:, None]).sum(1) m = nm ic = (tgt >= c) & (tgt < c+vchunk) if ic.any(): zt[ic] = lg[ic, tgt[ic] - c].float() lse = m + torch.log(s) ctx.save_for_backward(h, W, tgt, lse) ctx.vchunk = vchunk return (lse - zt).mean() @staticmethod def backward(ctx, go): h, W, tgt, lse = ctx.saved_tensors vc = ctx.vchunk N, d = h.shape V = W.shape[0] with torch.cuda.amp.autocast(enabled=True): hf = h.float() Wc_all = W.float() gh = torch.zeros_like(hf) gW = torch.zeros(W.shape, device=W.device, dtype=torch.float32) sc = float(go) / N for c in range(0, V, vc): Wc = Wc_all[c:c+vc] p = torch.exp(hf @ Wc.T - lse[:, None]) # softmax chunk [N,vchunk] ic = (tgt >= c) & (tgt < c+vc) if ic.any(): p[ic, tgt[ic] - c] -= 1.0 p *= sc gh += p @ Wc gW[c:c+vc] += p.T @ hf return gh.to(h.dtype), gW.to(W.dtype), None, None def fused_ce(h, W, tgt, vchunk=16384): return FusedCE.apply(h.reshape(-1, h.size(-1)), W, tgt.reshape(-1), vchunk) # ===== END fused_ce.py ===== # ===== BEGIN dblocks_train.py ===== """DiffusionBlocks training mode folded into AGILLM-4 (gated by --dblock). Block-wise EDM denoising on the real Encoder blocks, supervising AR + SAT(fixed+var) + NAT each step on ONE block, with grad-checkpointed layers and fused vocab-streaming CE. Reuses the live data stream / optimizer / checkpointing of nB300_agillm4. Lazy-imports nB300 inside functions to avoid a circular import. """ import math import random import time from collections import defaultdict import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import torch.utils.checkpoint as _ck # Optional CuPy hook for future AGILLM agents. # Keep the main trainer on PyTorch CUDA: autograd, AMP, SDPA, MoE, and DBlock # losses are already torch-native. This helper is deliberately lazy and disabled # by default so importing the trainer never depends on CuPy or CUDA toolkit # headers. Use it only for side/offline NumPy-heavy, non-autograd helpers such as # checkpoint/delta diagnostics, custom array probes, or preprocessing experiments. _CUPY_DISABLED = object() _OPTIONAL_CUPY = _CUPY_DISABLED def _optional_cupy_backend(reason=""): """Return cupy when AGILLM_ENABLE_CUPY=1, otherwise None. CuPy is useful for large NumPy-style array work on CUDA/ROCm hosts, but it is not a replacement for torch in the AGILLM4.3 training hot path. Callers must keep data on the GPU and avoid CPU<->GPU ping-pong. On Vast CUDA images, CuPy may need CUDA_PATH=/usr/local/cuda so elementwise kernels can find headers. """ global _OPTIONAL_CUPY import os as _os if _os.environ.get("AGILLM_ENABLE_CUPY", "0") != "1": return None if _OPTIONAL_CUPY is _CUPY_DISABLED: if not _os.environ.get("CUDA_PATH") and _os.path.exists("/usr/local/cuda"): _os.environ["CUDA_PATH"] = "/usr/local/cuda" try: import cupy as _cp # type: ignore _OPTIONAL_CUPY = _cp label = f" for {reason}" if reason else "" print(f"[cupy] optional backend enabled{label}: cupy={_cp.__version__}", flush=True) except Exception as exc: _OPTIONAL_CUPY = None print(f"[cupy] optional backend unavailable: {type(exc).__name__}: {exc}", flush=True) return _OPTIONAL_CUPY SD = 0.5 def _profile_active(state, args): limit = int(getattr(args, "profile_steps", 0) or 0) return limit > 0 and int(state.get("profile_n", 0)) < limit def _profile_add(state, name, seconds): if seconds is None: return prof = state.setdefault("profile_times", defaultdict(float)) prof[name] += float(seconds) def _profile_tic(enabled): if not enabled: return None if torch.cuda.is_available(): torch.cuda.synchronize() return time.perf_counter() def _profile_toc(state, name, start): if start is None: return if torch.cuda.is_available(): torch.cuda.synchronize() _profile_add(state, name, time.perf_counter() - start) def _profile_step_done(state, args): limit = int(getattr(args, "profile_steps", 0) or 0) if limit <= 0: return n_prev = int(state.get("profile_n", 0)) if n_prev >= limit: return state["profile_n"] = n_prev + 1 n = int(state["profile_n"]) log_every = max(1, int(getattr(args, "profile_log_every", 25) or 25)) if n % log_every != 0 and n != limit: return times = state.get("profile_times", {}) keys = [ "data_stream", "tensor", "setup", "ar_forward", "ar_ce", "ar_backward", "sat_forward", "sat_ce", "sat_backward", "nat_forward", "nat_ce", "nat_backward", "opt_step", "step_total", ] parts = [] for key in keys: val = float(times.get(key, 0.0)) * 1000.0 / max(1, n) if val > 0.01: parts.append(f"{key}={val:.2f}ms") print(f"[profile] n={n}/{limit} avg " + " ".join(parts), flush=True) def _cdf(x): return 0.5 * (1 + math.erf(x / math.sqrt(2))) def _ppf(p): return float(torch.erfinv(torch.tensor(2 * p - 1.0)) * math.sqrt(2)) def _dblock_sigma_config(args=None): smin = float(getattr(args, "dblock_sigma_min", 0.002) if args is not None else 0.002) smax = float(getattr(args, "dblock_sigma_max", 80.0) if args is not None else 80.0) pm = float(getattr(args, "dblock_sigma_pmean", -1.2) if args is not None else -1.2) ps = float(getattr(args, "dblock_sigma_pstd", 1.2) if args is not None else 1.2) smin = max(smin, 1e-6) smax = max(smax, smin * 1.0001) ps = max(ps, 1e-6) return smin, smax, pm, ps def _block_sigmas(B, smin=0.002, smax=80.0, pm=-1.2, ps=1.2): smin = max(float(smin), 1e-6) smax = max(float(smax), smin * 1.0001) ps = max(float(ps), 1e-6) a, b = _cdf((math.log(smin) - pm) / ps), _cdf((math.log(smax) - pm) / ps) return [float(np.exp(pm + ps * _ppf(a + (b - a) * (i / B)))) for i in range(B + 1)] def _edm_pre(s): s = s[:, None, None] return SD**2 / (s**2 + SD**2), s * SD / (s**2 + SD**2) ** 0.5, 1 / (s**2 + SD**2) ** 0.5 def _edm_w(s, wmax=5.0): return float(((s**2 + SD**2) / (s * SD) ** 2).clamp(max=wmax).mean()) _DBLOCK_ROUTER_EVENT_FEATURES = 10 _DBLOCK_ROUTER_HISTORY = 32 class _DblockLearnedRouter(nn.Module): # Transformer DBlock router conditioned on the network's running representation # plus a bounded route/outcome memory. Sequence = [CTX] + B block tokens + H # recent outcome tokens, so routing can learn from what the model is seeing now # and what the previous routing choices actually did to loss. def __init__(self, ctx_dim, d_model=64, heads=4, layers=2, feat_dim=6, n_blocks_max=64, history=_DBLOCK_ROUTER_HISTORY, event_dim=_DBLOCK_ROUTER_EVENT_FEATURES): super().__init__() d_model = max(16, int(d_model)) heads = max(1, int(heads)) if d_model % heads != 0: heads = 1 self.ctx_dim = int(ctx_dim) self.feat_dim = int(feat_dim) self.history = max(0, int(history)) self.event_dim = int(event_dim) self.block_emb = nn.Embedding(int(n_blocks_max), d_model) self.feat_proj = nn.Linear(int(feat_dim), d_model) self.ctx_proj = nn.Linear(int(ctx_dim), d_model) self.event_proj = nn.Linear(self.event_dim, d_model) self.kind_emb = nn.Embedding(3, d_model) self.event_pos = nn.Embedding(max(1, self.history), d_model) self.cls = nn.Parameter(torch.zeros(1, 1, d_model)) enc = nn.TransformerEncoderLayer( d_model=d_model, nhead=heads, dim_feedforward=max(32, d_model * 4), dropout=0.0, activation="gelu", batch_first=True, norm_first=True, ) self.encoder = nn.TransformerEncoder(enc, num_layers=max(1, int(layers))) self.ln = nn.LayerNorm(d_model) self.value = nn.Sequential( nn.LayerNorm(d_model * 2), nn.Linear(d_model * 2, d_model), nn.GELU(), nn.Linear(d_model, 1), ) nn.init.normal_(self.cls, std=0.02) @staticmethod def _fit_last_dim(x, dim): if x.size(-1) == dim: return x if x.size(-1) > dim: return x[..., :dim] return F.pad(x, (0, dim - x.size(-1))) def forward(self, block_ids, feats, ctx, history=None): feats = self._fit_last_dim(feats.float(), self.feat_dim) ctx = self._fit_last_dim(ctx.float(), self.ctx_dim) B = feats.size(1) bt = self.block_emb(block_ids.clamp(min=0, max=self.block_emb.num_embeddings - 1)) + self.feat_proj(feats) bt = bt + self.kind_emb(torch.ones(B, dtype=torch.long, device=feats.device)).unsqueeze(0) ctx_tok = self.cls + self.ctx_proj(ctx).unsqueeze(1) ctx_tok = ctx_tok + self.kind_emb(torch.zeros(1, dtype=torch.long, device=feats.device)).view(1, 1, -1) tokens = [ctx_tok, bt] if history is not None and self.history > 0: if not torch.is_tensor(history): history = torch.tensor(history, dtype=feats.dtype, device=feats.device) else: history = history.to(device=feats.device, dtype=feats.dtype) if history.dim() == 2: history = history.unsqueeze(0) if history.dim() == 3 and history.numel() > 0: if history.size(0) == 1 and feats.size(0) > 1: history = history.expand(feats.size(0), -1, -1) elif history.size(0) != feats.size(0): history = history[:1].expand(feats.size(0), -1, -1) if history.size(1) > self.history: history = history[:, -self.history :, :] history = self._fit_last_dim(history, self.event_dim) H = history.size(1) if H > 0: pos = torch.arange(H, dtype=torch.long, device=feats.device).clamp(max=max(0, self.history - 1)) kind = torch.full((H,), 2, dtype=torch.long, device=feats.device) ht = self.event_proj(history) + self.event_pos(pos).unsqueeze(0) + self.kind_emb(kind).unsqueeze(0) tokens.append(ht) h = self.ln(self.encoder(torch.cat(tokens, dim=1))) ctx_h = h[:, 0:1, :].expand(-1, B, -1) block_h = h[:, 1 : 1 + B, :] return self.value(torch.cat([block_h, ctx_h], dim=-1)).squeeze(-1) def _dblock_router_mode(args): return str(getattr(args, "dblock_router", "heuristic") or "heuristic").lower() def _dblock_router_enabled(args): return _dblock_router_mode(args) in {"transformer", "learned", "neural"} def _dblock_router_boot(state, args, ctx_dim=None): if not _dblock_router_enabled(args): return hidden = int(getattr(args, "dblock_router_hidden", 64) or 64) heads = int(getattr(args, "dblock_router_heads", 4) or 4) layers = int(getattr(args, "dblock_router_layers", 2) or 2) lr = float(getattr(args, "dblock_router_lr", 0.002) or 0.002) history = max(8, min(128, int(getattr(args, "dblock_router_history", _DBLOCK_ROUTER_HISTORY) or _DBLOCK_ROUTER_HISTORY))) cdim = int(ctx_dim or state.get("router_ctx_dim", 0) or 64) state["router_ctx_dim"] = cdim router = _DblockLearnedRouter(ctx_dim=cdim, d_model=hidden, heads=heads, layers=layers, history=history).to("cpu") state["router"] = router state["router_opt"] = torch.optim.AdamW(router.parameters(), lr=lr, weight_decay=1e-3) state["router_target_ema"] = None state["router_target_abs_ema"] = None state["router_train_loss"] = None state["router_last"] = None state["router_history"] = [] state["router_history_limit"] = history print( f"[dblock] learned_router=ctx_seq_transformer hidden={hidden} heads={heads} layers={layers} ctx_dim={cdim} history={history} lr={lr:g} " f"blend={float(getattr(args, 'dblock_router_blend', 0.35)):.2f} " f"ramp_steps={int(getattr(args, 'dblock_router_ramp_steps', 256) or 0)}", flush=True, ) def _dblock_router_features(state, args): B = int(state["B"]) step = int(state.get("step", 0)) counts = list(state.get("counts", [0 for _ in range(B)])) if len(counts) != B: counts = [0 for _ in range(B)] emas = list(state.get("loss_ema", [None for _ in range(B)])) if len(emas) != B: emas = [None for _ in range(B)] last_seen = list(state.get("last_seen", [-1 for _ in range(B)])) if len(last_seen) != B: last_seen = [-1 for _ in range(B)] bsig = list(state.get("bsig", _block_sigmas(B, *_dblock_sigma_config(args)))) max_count = max(1, max(counts) if counts else 1) known = [float(x) for x in emas if x is not None and math.isfinite(float(x))] center = sum(known) / len(known) if known else 0.0 scale = (sum((x - center) ** 2 for x in known) / len(known)) ** 0.5 if len(known) > 1 else max(1.0, abs(center) * 0.05) scale = max(1e-3, scale) stale = [step - last_seen[i] if last_seen[i] >= 0 else step + 1 for i in range(B)] max_stale = int(getattr(args, "dblock_max_stale_steps", 64) or 0) stale_denom = float(max(1, max_stale if max_stale > 0 else max(stale) if stale else 1)) logs = [math.log(max(1e-9, float(x))) for x in bsig] log_min = min(logs) if logs else 0.0 log_span = max(1e-6, (max(logs) - log_min) if logs else 1.0) feats = [] for i in range(B): ema = emas[i] known_flag = 1.0 if ema is not None and math.isfinite(float(ema)) else 0.0 loss_z = 0.0 if not known_flag else max(-5.0, min(5.0, (float(ema) - center) / scale)) lo = logs[min(i, len(logs) - 1)] if logs else 0.0 hi = logs[min(i + 1, len(logs) - 1)] if logs else lo sig_mid = ((0.5 * (lo + hi)) - log_min) / log_span feats.append([ loss_z, known_flag, float(counts[i]) / float(max_count), max(0.0, float(max_count - counts[i]) / float(max_count)), min(1.0, max(0.0, float(stale[i]) / stale_denom)), float(sig_mid), ]) block_ids = torch.arange(B, dtype=torch.long).unsqueeze(0) ft = torch.tensor([feats], dtype=torch.float32) cdim = int(state.get("router_ctx_dim", 0) or 0) ctx = state.get("router_ctx") if torch.is_tensor(ctx) and cdim > 0 and ctx.numel() == cdim: cv = ctx.detach().reshape(1, cdim).float() else: cv = torch.zeros(1, max(1, cdim)) return block_ids, ft, cv def _dblock_router_clip(x, lo=-5.0, hi=5.0): try: x = float(x) except Exception: return 0.0 if not math.isfinite(x): return 0.0 return max(lo, min(hi, x)) def _dblock_router_history_features(state, args): limit = int(state.get("router_history_limit", getattr(args, "dblock_router_history", _DBLOCK_ROUTER_HISTORY)) or 0) limit = max(0, min(128, limit)) if limit <= 0: return torch.zeros((1, 0, _DBLOCK_ROUTER_EVENT_FEATURES), dtype=torch.float32) hist = list(state.get("router_history", []))[-limit:] if not hist: return torch.zeros((1, 0, _DBLOCK_ROUTER_EVENT_FEATURES), dtype=torch.float32) B = int(state["B"]) step = int(state.get("step", 0)) losses = [] for rec in hist: try: loss = float(rec.get("loss", 0.0)) except Exception: loss = 0.0 if math.isfinite(loss): losses.append(loss) center = sum(losses) / len(losses) if losses else 0.0 scale = (sum((x - center) ** 2 for x in losses) / len(losses)) ** 0.5 if len(losses) > 1 else max(1.0, abs(center) * 0.05) scale = max(1e-3, scale) rows = [] for rec in hist: rec_step = int(rec.get("step", -1)) block = max(0, min(B - 1, int(rec.get("block", 0)))) age = max(0, step - rec_step) try: rec_loss = float(rec.get("loss", center)) except Exception: rec_loss = center loss = _dblock_router_clip((rec_loss - center) / scale) rows.append([ float(block) / float(max(1, B - 1)), _dblock_router_clip(rec.get("target", 0.0)), loss, max(0.0, min(1.0, float(rec.get("count_norm", 0.0)))), max(0.0, min(1.0, float(rec.get("stale_norm", 0.0)))), min(1.0, math.log1p(age) / math.log1p(max(2, limit))), min(1.0, math.log1p(max(0, rec_step)) / math.log1p(10000.0)), 1.0 if float(rec.get("router_choice", 0.0)) > 0.0 else 0.0, max(0.0, min(1.0, float(rec.get("blend", 0.0)))), 1.0, ]) return torch.tensor([rows], dtype=torch.float32) def _dblock_router_append_history(state, args, bi, loss_float, target_val): limit = int(state.get("router_history_limit", getattr(args, "dblock_router_history", _DBLOCK_ROUTER_HISTORY)) or _DBLOCK_ROUTER_HISTORY) limit = max(0, min(128, limit)) if limit <= 0: return B = int(state["B"]) step = int(state.get("step", 0)) counts = list(state.get("counts", [0 for _ in range(B)])) if len(counts) != B: counts = [0 for _ in range(B)] last_seen = list(state.get("last_seen", [-1 for _ in range(B)])) if len(last_seen) != B: last_seen = [-1 for _ in range(B)] max_count = max(1, max(counts) if counts else 1) stale = step - last_seen[int(bi)] if 0 <= int(bi) < len(last_seen) and last_seen[int(bi)] >= 0 else step + 1 max_stale = int(getattr(args, "dblock_max_stale_steps", 64) or 0) stale_denom = float(max(1, max_stale if max_stale > 0 else stale)) route = state.get("router_last") router_choice = 0.0 blend = 0.0 if isinstance(route, dict): router_choice = 1.0 if int(route.get("choice", -1)) == int(bi) else 0.0 blend = float(route.get("blend", 0.0)) hist = state.setdefault("router_history", []) hist.append({ "step": int(step), "block": int(bi), "loss": float(loss_float), "target": float(target_val), "count_norm": float(counts[int(bi)]) / float(max_count) if 0 <= int(bi) < len(counts) else 0.0, "stale_norm": min(1.0, max(0.0, float(stale) / stale_denom)), "router_choice": router_choice, "blend": blend, }) if len(hist) > limit: del hist[:-limit] def _dblock_router_norm(xs): vals = [0.0 if not math.isfinite(float(x)) else float(x) for x in xs] if not vals: return vals mean = sum(vals) / len(vals) scale = max(1e-6, (sum((x - mean) ** 2 for x in vals) / len(vals)) ** 0.5) return [(x - mean) / scale for x in vals] def _dblock_fleet_lane_keys(args): keys = [] for env_key in ("AGILLM_FLEET_LANE", "AGILLM_WORKER_ID", "AGILLM_LANE_ID"): val = os.environ.get(env_key, "") if val: keys.append(str(val)) save_dir = str(getattr(args, "save_dir", "") or "") if save_dir: keys.append(os.path.basename(save_dir.rstrip("/"))) keys.append(save_dir) return [k for i, k in enumerate(keys) if k and k not in keys[:i]] def _dblock_fleet_router_scores(state, args, base_scores): state["fleet_router_last"] = None if not base_scores: return None try: cfg = get_hot_config() except Exception: return None spec = cfg.get("dblock_fleet_router") or cfg.get("dblock_fleet_route") if not isinstance(spec, dict): return None if str(spec.get("enabled", True)).lower() in {"0", "false", "off", "no"}: return None lanes = spec.get("lanes") if isinstance(spec.get("lanes"), dict) else {} lane_key = None lane = None for key in _dblock_fleet_lane_keys(args): cand = lanes.get(key) if isinstance(cand, dict): lane_key, lane = key, cand break if lane is None: return None bias = lane.get("bias", lane.get("block_bias", lane.get("biases"))) if not isinstance(bias, (list, tuple)): return None B = int(state.get("B", len(base_scores)) or len(base_scores)) if len(bias) != B or len(base_scores) != B: return None vals = [] for x in bias: try: fx = float(x) except Exception: fx = 0.0 vals.append(0.0 if not math.isfinite(fx) else max(-3.0, min(3.0, fx))) if not any(abs(x) > 1e-9 for x in vals): return None strength = float(lane.get("strength", spec.get("strength", 0.20)) or 0.0) strength = max(0.0, min(1.0, strength)) if strength <= 1e-9: return None base = [float(x) if math.isfinite(float(x)) else 0.0 for x in base_scores] mean = sum(base) / len(base) scale = max(1e-3, (sum((x - mean) ** 2 for x in base) / len(base)) ** 0.5) adjusted = [base[i] + strength * scale * vals[i] for i in range(B)] state["fleet_router_last"] = { "lane": str(lane_key), "role": str(lane.get("role", "")), "strength": float(strength), "bias": [float(x) for x in vals], "updated_at": spec.get("updated_at", ""), } return adjusted def _dblock_router_choose(state, args, heuristic_scores): state["router_last"] = None if not _dblock_router_enabled(args): return None router = state.get("router") if router is None: return None B = int(state["B"]) step = int(state.get("step", 0)) warmup = int(getattr(args, "dblock_warmup_steps", max(8, B * 2))) ramp_steps = int(getattr(args, "dblock_router_ramp_steps", 256) or 0) blend_base = max(0.0, min(1.0, float(getattr(args, "dblock_router_blend", 0.35) or 0.0))) if step < warmup or blend_base <= 0.0: return None ramp = 1.0 if ramp_steps <= 0 else min(1.0, max(0.0, float(step - warmup) / float(ramp_steps))) blend = blend_base * ramp if blend <= 1e-6: return None history_features = _dblock_router_history_features(state, args) with torch.no_grad(): router.eval() pred = router(*_dblock_router_features(state, args), history=history_features)[0].detach().cpu().tolist() h = _dblock_router_norm(heuristic_scores) q = _dblock_router_norm(pred) if len(h) != B or len(q) != B: return None counts = state.get("counts", [0 for _ in range(B)]) combined = [(1.0 - blend) * h[i] + blend * q[i] for i in range(B)] choice = max(range(B), key=lambda i: (combined[i], -counts[i], -i)) state["router_last"] = { "mode": "ctx_seq_transformer", "choice": int(choice), "blend": float(blend), "history": int(history_features.size(1)), "pred": [float(x) for x in pred], } return choice def _dblock_router_update(state, args, bi, loss_value): if not _dblock_router_enabled(args): return router, opt = state.get("router"), state.get("router_opt") if router is None or opt is None: return try: loss_float = float(loss_value) except Exception: return if not math.isfinite(loss_float): return baseline = state.get("router_target_ema") scale = state.get("router_target_abs_ema") if baseline is None or not math.isfinite(float(baseline)): baseline = loss_float if scale is None or not math.isfinite(float(scale)) or float(scale) < 1e-3: scale = max(1.0, abs(loss_float) * 0.05) target_val = max(-5.0, min(5.0, (loss_float - float(baseline)) / max(1e-3, float(scale)))) router.train() pred = router(*_dblock_router_features(state, args), history=_dblock_router_history_features(state, args))[0, int(bi)] fit_loss = F.smooth_l1_loss(pred, pred.detach().new_tensor(target_val)) opt.zero_grad(set_to_none=True) fit_loss.backward() nn.utils.clip_grad_norm_(router.parameters(), 1.0) opt.step() diff = abs(loss_float - float(baseline)) state["router_target_ema"] = 0.98 * float(baseline) + 0.02 * loss_float state["router_target_abs_ema"] = 0.98 * float(scale) + 0.02 * max(1e-3, diff) state["router_train_loss"] = float(fit_loss.detach().cpu()) _dblock_router_append_history(state, args, bi, loss_float, target_val) def _dblock_get_candidates(L): c = [] # 1. Uniform candidates for b in [2, 3, 4, 6] for b in [2, 3, 4, 6]: per = max(1, L // b) asg = [list(range(i * per, (i + 1) * per)) for i in range(b)] asg[-1] = list(range((b - 1) * per, L)) c.append((b, asg, f"Uniform-{b}")) # 2. Non-uniform candidates for B=3 # Middle-heavy (e.g. 25%, 50%, 25%) m_h = [max(1, L // 4), max(1, L // 2)] m_h.append(L - sum(m_h)) asg = [] curr = 0 for size in m_h: asg.append(list(range(curr, curr + size))) curr += size c.append((3, asg, "Middle-Heavy-3")) # End-heavy (e.g. 20%, 35%, 45%) e_h = [max(1, int(L * 0.20)), max(1, int(L * 0.35))] e_h.append(L - sum(e_h)) asg = [] curr = 0 for size in e_h: asg.append(list(range(curr, curr + size))) curr += size c.append((3, asg, "End-Heavy-3")) # Start-heavy (e.g. 45%, 35%, 20%) s_h = [max(1, int(L * 0.45)), max(1, int(L * 0.35))] s_h.append(L - sum(s_h)) asg = [] curr = 0 for size in s_h: asg.append(list(range(curr, curr + size))) curr += size c.append((3, asg, "Start-Heavy-3")) # 3. Non-uniform candidates for B=4 # Middle-heavy (e.g. 20%, 30%, 30%, 20%) m_h4 = [max(1, int(L * 0.20)), max(1, int(L * 0.30)), max(1, int(L * 0.30))] m_h4.append(L - sum(m_h4)) asg = [] curr = 0 for size in m_h4: asg.append(list(range(curr, curr + size))) curr += size c.append((4, asg, "Middle-Heavy-4")) # End-heavy (e.g. 15%, 25%, 30%, 30%) e_h4 = [max(1, int(L * 0.15)), max(1, int(L * 0.25)), max(1, int(L * 0.30))] e_h4.append(L - sum(e_h4)) asg = [] curr = 0 for size in e_h4: asg.append(list(range(curr, curr + size))) curr += size c.append((4, asg, "End-Heavy-4")) return c def _dblock_init(core, args): L = len(core.blocks) auto_search = getattr(args, "auto_dblock_search", False) if auto_search: candidates = _dblock_get_candidates(L) print(f"[dblock] Auto Search enabled with {len(candidates)} candidates.") B, asg, name = candidates[0] state = { "auto_search": True, "candidates": candidates, "candidate_idx": 0, "search_step": 0, "search_interval": 20, "scores": [], } else: B = int(getattr(args, "dblock_blocks", 4)) sp = max(1, L // B) asg = [list(range(i * sp, (i + 1) * sp)) for i in range(B)] asg[-1] = list(range((B - 1) * sp, L)) state = {"auto_search": False} bsig = _block_sigmas(B, *_dblock_sigma_config(args)) schedule = getattr(args, "dblock_schedule", "loss_balanced") print(f"[dblock] DiffusionBlocks mode: {L} layers -> {B} blocks {asg}") print(f"[dblock] schedule={schedule} sigma boundaries: {[round(x, 3) for x in bsig]}") state.update({ "B": B, "assign": asg, "bsig": bsig, "step": 0, "counts": [0 for _ in range(B)], "loss_ema": [None for _ in range(B)], "last_seen": [-1 for _ in range(B)], }) if bool(getattr(args, "dblock_looped", False)): loop_layers = int(getattr(args, "dblock_loop_layers", 0) or 0) if loop_layers <= 0: loop_layers = max(1, L // max(1, B)) loop_layers = max(1, min(loop_layers, L)) loop_start = max(0, min(int(getattr(args, "dblock_loop_start", 0) or 0), L - loop_layers)) loop_group = list(range(loop_start, loop_start + loop_layers)) if not hasattr(core, "dblock_loop_embed"): d = int(getattr(core.emb, "embedding_dim", 0)) core.dblock_loop_embed = nn.Embedding(B, d).to(core.emb.weight.device) nn.init.normal_(core.dblock_loop_embed.weight, mean=0.0, std=0.02) state.update({ "looped": True, "loop_group": loop_group, "loop_layers": loop_layers, "loop_start": loop_start, }) print( f"[dblock-looped] enabled: shared_layers={loop_group} bands={B} " f"unrolled_depth={loop_layers * B} one-band-per-step no_bptt=True", flush=True, ) _dblock_router_boot(state, args, ctx_dim=int(getattr(core.emb, "embedding_dim", 0)) or None) return state def _choose_block(state, args): if not state.get("auto_search", False) and state.get("step", 0) % 100 == 0: try: cfg = get_hot_config() if "dblock_blocks" in cfg: new_B = int(cfg["dblock_blocks"]) if new_B != state.get("B"): L = sum(len(x) for x in state["assign"]) if "assign" in state else 28 new_sp = max(1, L // new_B) new_asg = [list(range(i * new_sp, (i + 1) * new_sp)) for i in range(new_B)] new_asg[-1] = list(range((new_B - 1) * new_sp, L)) print(f"[dblock] Dynamically adjusting block configuration from hot_config: B={state['B']} -> {new_B}, assign={new_asg}", flush=True) state["B"] = new_B state["assign"] = new_asg state["bsig"] = _block_sigmas(new_B, *_dblock_sigma_config(args)) state["counts"] = [0] * new_B state["loss_ema"] = [None] * new_B state["last_seen"] = [-1] * new_B except Exception as e: print(f"[dblock] Error reloading hot_config in _choose_block: {e}", flush=True) if state.get("auto_search", False) and state["candidate_idx"] < len(state["candidates"]): state["search_step"] += 1 if "search_start_time" not in state: state["search_start_time"] = time.perf_counter() state["search_tokens"] = 0 if state["search_step"] >= state["search_interval"]: valid_emas = [e for e in state["loss_ema"] if e is not None] avg_loss = sum(valid_emas) / max(1, len(valid_emas)) if valid_emas else float('inf') elapsed = time.perf_counter() - state["search_start_time"] tokens = state.get("search_tokens", 0) tokps = tokens / max(1e-9, elapsed) cand = state["candidates"][state["candidate_idx"]] cand_name = cand[2] if len(cand) > 2 else f"Candidate-{state['candidate_idx']}" state["scores"].append({ "idx": state["candidate_idx"], "B": state["B"], "assign": state["assign"], "name": cand_name, "loss": avg_loss, "tokps": tokps }) print(f"[dblock] Candidate {state['candidate_idx']} ({cand_name}) complete: loss={avg_loss:.4f} speed={tokps:.1f} tok/s", flush=True) state["candidate_idx"] += 1 state["search_step"] = 0 if "search_start_time" in state: del state["search_start_time"] state["search_tokens"] = 0 if state["candidate_idx"] < len(state["candidates"]): B, asg, cand_name = state["candidates"][state["candidate_idx"]] state["B"] = B state["assign"] = asg state["bsig"] = _block_sigmas(B, *_dblock_sigma_config(args)) state["counts"] = [0] * B state["loss_ema"] = [None] * B state["last_seen"] = [-1] * B print(f"[dblock] Switched to candidate {state['candidate_idx']} ({cand_name}): {B} blocks {asg}", flush=True) else: # Select the candidate with highest speed/loss utility best_cand = None best_utility = -1.0 for score_entry in state["scores"]: loss = score_entry["loss"] tokps = score_entry["tokps"] utility = tokps / max(1e-3, loss) score_entry["utility"] = utility if utility > best_utility: best_utility = utility best_cand = score_entry B = best_cand["B"] asg = best_cand["assign"] state["B"] = B state["assign"] = asg state["bsig"] = _block_sigmas(B, *_dblock_sigma_config(args)) state["auto_search"] = False print(f"[dblock] Search complete. Locked in best candidate {best_cand['name']} (Utility={best_utility:.2f}, Loss={best_cand['loss']:.4f}, Speed={best_cand['tokps']:.1f} tok/s): {B} blocks {asg}", flush=True) B = state["B"] schedule = str(getattr(args, "dblock_schedule", "loss_balanced") or "loss_balanced").lower() step = int(state.get("step", 0)) counts = state.setdefault("counts", [0 for _ in range(B)]) if len(counts) != B: counts[:] = [0 for _ in range(B)] emas = state.setdefault("loss_ema", [None for _ in range(B)]) if len(emas) != B: emas[:] = [None for _ in range(B)] last_seen = state.setdefault("last_seen", [-1 for _ in range(B)]) if len(last_seen) != B: last_seen[:] = [-1 for _ in range(B)] state["router_last"] = None state["fleet_router_last"] = None if schedule == "random": return random.randrange(B) if schedule == "roundrobin": return step % B explore = max(0.0, min(1.0, float(getattr(args, "dblock_explore", 0.05)))) warmup = int(getattr(args, "dblock_warmup_steps", max(8, B * 2))) def least_trained(): return min(range(B), key=lambda i: (counts[i], last_seen[i], i)) if step < warmup or any(c == 0 for c in counts): return least_trained() max_stale = int(getattr(args, "dblock_max_stale_steps", 64) or 0) stale = [step - last_seen[i] if last_seen[i] >= 0 else step + 1 for i in range(B)] if max_stale > 0 and max(stale) >= max_stale: return max(range(B), key=lambda i: (stale[i], -counts[i], -i)) max_count = max(counts) if counts else 0 min_count = min(counts) if counts else 0 max_skew = float(getattr(args, "dblock_max_count_skew", 1.35) or 0.0) if max_skew > 1.0 and min_count > 0 and (max_count / max(1, min_count)) > max_skew: return least_trained() if explore > 0.0 and random.random() < explore: return least_trained() stale_bonus = float(getattr(args, "dblock_stale_bonus", 0.35) or 0.0) undertrain_bonus = float(getattr(args, "dblock_undertrain_bonus", 0.25) or 0.0) stale_denom = float(max(1, max_stale if max_stale > 0 else max(stale) if stale else 1)) count_denom = float(max(1, max_count)) def score(i): loss_score = -1.0 if emas[i] is None else float(emas[i]) stale_score = stale_bonus * min(1.0, max(0.0, stale[i] / stale_denom)) undertrain_score = undertrain_bonus * max(0.0, (max_count - counts[i]) / count_denom) return (loss_score + stale_score + undertrain_score, -counts[i], stale[i], -i) base_scores = [float(score(i)[0]) for i in range(B)] route_scores = _dblock_fleet_router_scores(state, args, base_scores) or base_scores if route_scores is base_scores: heuristic_choice = max(range(B), key=score) else: heuristic_choice = max(range(B), key=lambda i: (route_scores[i], -counts[i], stale[i], -i)) learned_choice = _dblock_router_choose(state, args, route_scores) return heuristic_choice if learned_choice is None else learned_choice def _sample_sigma(ids, lo, hi, args, state): cur_step = int(state.get("step", 0)) curriculum = int(getattr(args, "dblock_sigma_curriculum_steps", 0)) if curriculum > 0: frac = min(1.0, max(0.05, (cur_step + 1) / float(curriculum))) hi = lo * ((hi / max(lo, 1e-8)) ** frac) mode = str(getattr(args, "dblock_sigma_sampling", "lognormal") or "lognormal").lower() if mode in {"lognormal", "truncated_lognormal", "edm"}: _, _, pm, ps = _dblock_sigma_config(args) qa = _cdf((math.log(max(lo, 1e-6)) - pm) / ps) qb = _cdf((math.log(max(hi, lo * 1.0001)) - pm) / ps) qa = min(max(qa, 1e-7), 1.0 - 1e-7) qb = min(max(qb, qa + 1e-7), 1.0 - 1e-7) n = int(ids.size(0)) if bool(getattr(args, "dblock_sigma_stratified", True)) and n > 1: # Beyond the DBT paper: randomized quantile strata reduce Monte Carlo # variance of the conditional p_noise integral for each block. u = (torch.arange(n, device=ids.device, dtype=torch.float32) + torch.rand((), device=ids.device)) / float(n) u = u.index_select(0, torch.randperm(n, device=ids.device)) else: u = torch.rand(n, device=ids.device, dtype=torch.float32) q = qa + (qb - qa) * u q = q.clamp(1e-7, 1.0 - 1e-7) z = torch.erfinv(2.0 * q - 1.0) * math.sqrt(2.0) return torch.exp(torch.tensor(pm, device=ids.device, dtype=torch.float32) + float(ps) * z) sig_np = np.exp( np.random.uniform( math.log(max(lo, 1e-4)), math.log(max(hi, lo + 1e-4)), ids.size(0), ).astype("float32") ) return torch.from_numpy(sig_np).to(ids.device) def _maybe_log( state, args, bi, layers, ar_val, sat_val, nat_val, total_val, peak_alloc, peak_reserved, objective=None, raw_avg=None, raw_total=None, edm_weight=None, ): log_every = int(getattr(args, "dblock_log_every", 50)) step = int(state.get("step", 0)) if log_every <= 0 or step % log_every != 0: return counts_list = state.get("counts", []) last_seen = state.get("last_seen", [-1 for _ in counts_list]) counts = ",".join(str(x) for x in counts_list) emas = ",".join("nan" if x is None else f"{x:.2f}" for x in state.get("loss_ema", [])) stale = ",".join(str(max(0, step - int(last_seen[i]))) for i in range(min(len(counts_list), len(last_seen)))) mem = "" if peak_alloc is not None: mem = f" peak_alloc={peak_alloc:.2f}GB peak_reserved={peak_reserved:.2f}GB" display = float(raw_avg) if raw_avg is not None and math.isfinite(float(raw_avg)) else float(total_val) raw_part = "" if raw_total is not None: raw_part += f" raw_sum={float(raw_total):.3f}" if edm_weight is not None: raw_part += f" edm_w={float(edm_weight):.3f}" route = state.get("router_last") if isinstance(route, dict): pred = ",".join(f"{float(x):.2f}" for x in route.get("pred", [])) hist = route.get("history") hist_part = "" if hist is None else f" hist={int(hist)}" raw_part += f" router={route.get('mode', 'none')} blend={float(route.get('blend', 0.0)):.2f}{hist_part} pred=[{pred}]" rloss = state.get("router_train_loss") if rloss is not None: raw_part += f" router_fit={float(rloss):.3f}" fleet = state.get("fleet_router_last") if isinstance(fleet, dict): fbias = fleet.get("bias", []) top = [] try: top = sorted(range(len(fbias)), key=lambda j: abs(float(fbias[j])), reverse=True)[:3] except Exception: top = [] top_part = ",".join(f"{j}:{float(fbias[j]):+.2f}" for j in top) raw_part += f" fleet={fleet.get('lane', '')} role={fleet.get('role', '')} strength={float(fleet.get('strength', 0.0)):.2f}" if top_part: raw_part += f" fleet_bias=[{top_part}]" print( f"[dblock] step={step} block={bi} obj={objective or 'mixed'} layers={layers} " f"loss={display:.3f} weighted={total_val:.3f} ar={ar_val:.3f} sat={sat_val:.3f} nat={nat_val:.3f}" f"{raw_part} counts=[{counts}] ema=[{emas}] stale=[{stale}]{mem}", flush=True, ) def _update_stats(state, bi, loss_value, args=None): if args is not None: _dblock_router_update(state, args, bi, loss_value) B = state["B"] counts = state.setdefault("counts", [0 for _ in range(B)]) emas = state.setdefault("loss_ema", [None for _ in range(B)]) last_seen = state.setdefault("last_seen", [-1 for _ in range(B)]) if len(last_seen) != B: last_seen[:] = [-1 for _ in range(B)] counts[bi] += 1 last_seen[bi] = int(state.get("step", 0)) prev = emas[bi] beta = 0.96 emas[bi] = float(loss_value) if prev is None else beta * float(prev) + (1.0 - beta) * float(loss_value) state["step"] = int(state.get("step", 0)) + 1 def _activation_offload_enabled(args): return bool(getattr(args, "dblock_activation_offload", False)) and torch.cuda.is_available() def _activation_offload_hooks(args): min_bytes = int(float(getattr(args, "dblock_activation_offload_min_mb", 1.0) or 1.0) * 1024 * 1024) def pack(t): if not torch.is_tensor(t) or not t.is_cuda or not t.is_floating_point() or t.numel() * t.element_size() < min_bytes: return t return ("cpu_offload", t.device, t.detach().to("cpu", non_blocking=True)) def unpack(x): if isinstance(x, tuple) and len(x) == 3 and x[0] == "cpu_offload": _, dev, cpu_t = x return cpu_t.to(dev, non_blocking=True) return x return torch.autograd.graph.saved_tensors_hooks(pack, unpack) def _dblock_sublayer_base_mode(args): mode = str(getattr(args, "dblock_sublayer_mode", "off") or "off").strip().lower().replace("-", "_") if mode in {"none", "disabled"}: return "off" return mode def _dblock_sublayer_mode_for_layer(args, state, block_idx, layer_pos): mode = _dblock_sublayer_base_mode(args) if mode == "split_alt": step = int((state or {}).get("step", 0)) return "attn_only" if ((step + int(block_idx) + int(layer_pos)) % 2 == 0) else "ffn_only" if mode == "cycle": step = int((state or {}).get("step", 0)) return ("full", "ffn_only", "attn_only")[(step + int(block_idx) + int(layer_pos)) % 3] return mode def _run_block_forward(block, x, mask, sublayer_mode="off"): mode = str(sublayer_mode or "off").strip().lower().replace("-", "_") if mode in {"off", "full"}: return block(x, mask) if mode == "attn_only": n = x.size(1) return x + block.mha(block.ln1(x), mask, rel_bias_tokens=n) if mode == "ffn_only": return x + block.ff(block.ln2(x)) raise ValueError(f"unknown DBlock sublayer mode: {sublayer_mode}") def _run_block(block, x, mask, use_checkpoint, args=None, sublayer_mode="off"): if use_checkpoint: return _ck.checkpoint(lambda y, block=block, mode=sublayer_mode: _run_block_forward(block, y, mask, mode), x, use_reentrant=False) if args is not None and _activation_offload_enabled(args): with _activation_offload_hooks(args): return _run_block_forward(block, x, mask, sublayer_mode) return _run_block_forward(block, x, mask, sublayer_mode) def _dblock_checkpoint_this_layer(args, base_enabled, layer_pos, layer_count=None): if not base_enabled: return False pos = int(layer_pos) count = int(layer_count or 0) skip_tail = max(0, int(getattr(args, "dblock_checkpoint_skip_tail", 0) or 0)) if skip_tail > 0 and count > 0 and pos >= max(0, count - skip_tail): return False stride = int(getattr(args, "dblock_checkpoint_stride", 1) or 1) if stride <= 0: return False if stride == 1: return True return (pos % stride) == 0 def _dblock_loop_condition(core, h, block_idx, args): emb = getattr(core, "dblock_loop_embed", None) if emb is None: return h idx = torch.tensor([int(block_idx)], device=h.device, dtype=torch.long) cond = emb(idx).to(dtype=h.dtype).view(1, 1, -1) return h + float(getattr(args, "dblock_loop_cond_scale", 1.0) or 0.0) * cond def _maybe_register_looped_infer(core, sd, args): """Looped checkpoints carry 'dblock_loop_embed.weight' in their core state. Recreate the matching embedding on the inference core (so the strict core load accepts it) and flip args into looped mode so the EDM block-chain decodes through the single shared looped group with loop-index conditioning.""" core_sd = sd.get("core") if isinstance(sd, dict) else None if not isinstance(core_sd, dict): return w = core_sd.get("dblock_loop_embed.weight") if w is None: return bands = int(w.shape[0]) d = int(getattr(core.emb, "embedding_dim", 0)) or int(w.shape[1]) if not hasattr(core, "dblock_loop_embed"): core.dblock_loop_embed = nn.Embedding(bands, d).to(core.emb.weight.device) try: setattr(args, "dblock_looped", True) setattr(args, "dblock_blocks", bands) except Exception: pass print("[dblock-looped] inference: shared looped group, bands=%d" % bands, flush=True) def _sample_token_loss_inputs(hidden, targets, max_tokens): max_tokens = int(max_tokens or 0) if max_tokens <= 0: return hidden.contiguous(), targets.contiguous(), int(targets.numel()), int(targets.numel()) flat_targets = targets.reshape(-1) total = int(flat_targets.numel()) if total <= max_tokens: return hidden.contiguous(), targets.contiguous(), total, total # With-replacement sampling avoids building a full randperm each step; the sampled # mean remains an unbiased estimator of the dense token CE mean. idx = torch.randint(total, (max_tokens,), device=targets.device) flat_hidden = hidden.reshape(total, hidden.size(-1)) return flat_hidden.index_select(0, idx).contiguous(), flat_targets.index_select(0, idx).contiguous(), int(max_tokens), total def _choose_objectives(state, args, ar_weight, sat_weight, nat_weight, do_sat_periodic, do_nat_periodic): mode = str(getattr(args, "dblock_objective_mode", "periodic") or "periodic").lower() if mode != "stochastic": return ar_weight > 0.0, sat_weight > 0.0 and do_sat_periodic, nat_weight > 0.0 and do_nat_periodic, "periodic" choices = [] probs = [] if ar_weight > 0.0: choices.append("ar") probs.append(max(0.0, _dblock_hot_float(args, "dblock_ar_prob", 0.80, min_value=0.0))) if sat_weight > 0.0 and not getattr(args, "ar_only", False): choices.append("sat") probs.append(max(0.0, _dblock_hot_float(args, "dblock_sat_prob", 0.10, min_value=0.0))) if nat_weight > 0.0 and not getattr(args, "ar_only", False): choices.append("nat") probs.append(max(0.0, _dblock_hot_float(args, "dblock_nat_prob", 0.10, min_value=0.0))) if not choices: return False, False, False, "none" total = sum(probs) if total <= 0.0: probs = [1.0 / len(choices) for _ in choices] else: probs = [p / total for p in probs] picked = random.choices(choices, weights=probs, k=1)[0] return picked == "ar", picked == "sat", picked == "nat", picked def _dblock_step(core, ar_h, sat_h, nat_h, opt, scaler, args, ids, state): M = _agillm41_sys.modules[__name__] if state is not None and state.get("auto_search", False): state["search_tokens"] = state.get("search_tokens", 0) + ids.numel() prof = _profile_active(state, args) _step_t = _profile_tic(prof) if torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() _setup_t = _profile_tic(prof) B = state["B"] asg = state["assign"] bs = state["bsig"] T = ids.size(1) use_layer_checkpoint = bool(getattr(args, "grad_checkpoint", False)) if _dblock_router_enabled(args): with torch.no_grad(): _rc_emb = core.emb(ids) state["router_ctx"] = _rc_emb.mean(dim=(0, 1)).detach().float().to("cpu") del _rc_emb bi = _choose_block(state, args) lo, hi = sorted([bs[bi], bs[bi + 1]]) layers = asg[bi] if state.get("looped", False): layers = state.get("loop_group") or layers sig = _sample_sigma(ids, lo, hi, args, state) cs, co, ci = _edm_pre(sig) w = _edm_w(sig, float(getattr(args, "dblock_edm_wmax", 5.0))) SATB = M.SAT_BLOCK ar_weight = _dblock_hot_float(args, "dblock_ar_weight", 1.0, min_value=0.0) sat_weight = _dblock_hot_float(args, "dblock_sat_weight", 1.0, min_value=0.0) nat_weight = ( _dblock_hot_float(args, "dblock_nat_weight", 1.0, min_value=0.0) * _dblock_hot_float(args, "nat_loss_weight", 1.0, names=["nat_loss_weight", "dblock_nat_loss_weight"], min_value=0.0) ) do_sat_periodic = (not getattr(args, "ar_only", False)) and ( int(getattr(args, "sat_every", 1)) <= 1 or ((int(state.get("step", 0)) + 1) % int(getattr(args, "sat_every", 1)) == 0) ) do_nat_periodic = ( nat_h is not None and (not getattr(args, "ar_only", False)) and int(getattr(args, "nat_every", 1)) > 0 and ( int(getattr(args, "nat_every", 1)) <= 1 or ((int(state.get("step", 0)) + 1) % int(getattr(args, "nat_every", 1)) == 0) ) ) run_ar, run_sat, run_nat, objective = _choose_objectives( state, args, ar_weight, sat_weight, nat_weight, do_sat_periodic, do_nat_periodic ) _profile_toc(state, "setup", _setup_t) ar_val = 0.0 sat_val = 0.0 nat_val = 0.0 ar_raw_val = 0.0 sat_raw_val = 0.0 nat_raw_val = 0.0 if run_ar: causal = M.causal_mask(T, structured=M.use_structured_masks(args)) _t = _profile_tic(prof) with M.amp(args.amp): emb = core.emb(ids) zt = emb + sig[:, None, None] * torch.randn_like(emb) h = _dblock_loop_condition(core, ci * zt, bi, args) if state.get("looped", False) else ci * zt for lpos, li in enumerate(layers): mode = _dblock_sublayer_mode_for_layer(args, state, bi, lpos) h = _run_block(core.blocks[li], h, causal, _dblock_checkpoint_this_layer(args, use_layer_checkpoint, lpos, len(layers)), args, mode) Dn = core.ln(cs * zt + co * h) _profile_toc(state, "ar_forward", _t) _t = _profile_tic(prof) ar_hidden, ar_targets, ar_used, ar_total = _sample_token_loss_inputs( Dn[:, :-1], ids[:, 1:], _dblock_loss_token_cap(args, "ar") ) ar_raw = fused_ce(ar_hidden, ar_h.proj.weight, ar_targets) ar_raw_val = float(ar_raw.detach()) ar = ar_weight * w * ar_raw ar_val = float(ar.detach()) _profile_toc(state, "ar_ce", _t) _t = _profile_tic(prof) _aux = _collect_moe_aux(core, getattr(args,'moe_aux_coef',0.0), getattr(args,'moe_z_coef',0.0)) if torch.is_tensor(_aux): ar = ar + _aux.to(ar.dtype) scaler.scale(ar).backward() _profile_toc(state, "ar_backward", _t) del causal, emb, zt, h, Dn, ar_hidden, ar_targets, ar_raw, ar, ar_used, ar_total if run_sat: smask = M.sat_mask(T, structured=M.use_structured_masks(args)) _t = _profile_tic(prof) with M.amp(args.amp): emb2 = core.emb(ids) zt2 = emb2 + sig[:, None, None] * torch.randn_like(emb2) h2 = _dblock_loop_condition(core, ci * zt2, bi, args) if state.get("looped", False) else ci * zt2 for lpos, li in enumerate(layers): mode = _dblock_sublayer_mode_for_layer(args, state, bi, lpos) h2 = _run_block(core.blocks[li], h2, smask, _dblock_checkpoint_this_layer(args, use_layer_checkpoint, lpos, len(layers)), args, mode) Ds = core.ln(cs * zt2 + co * h2) _profile_toc(state, "sat_forward", _t) _t = _profile_tic(prof) # SAT decode uses the latest SAT_BLOCK hidden states to emit the next # SAT_BLOCK tokens. Train that contract densely across the context. sat_ctx = Ds[:, :-SATB] sat_tgt = ids[:, SATB:] if sat_ctx.size(1) == 0 or sat_ctx.size(1) != sat_tgt.size(1): sat_ctx = Ds[:, :-1] sat_tgt = ids[:, 1:] sat_hidden, sat_targets, sat_used, sat_total = _sample_token_loss_inputs( sat_ctx, sat_tgt, _dblock_loss_token_cap(args, "sat") ) sat_gate_ctx = sat_ctx[:, ::SATB] with M.amp(args.amp): satf = fused_ce(sat_hidden, sat_h.proj.weight, sat_targets) satv = ( M.EMIT_LAMBDA * F.cross_entropy( sat_h.gate(sat_gate_ctx.reshape(-1, sat_gate_ctx.size(-1)).float()), torch.ones(sat_gate_ctx.numel() // sat_gate_ctx.size(-1), dtype=torch.long, device=ids.device), ) if sat_h.gate is not None and sat_gate_ctx.size(1) > 0 else 0.0 ) sat_raw = satf + satv sat_raw_val = float(sat_raw.detach()) sat = sat_weight * w * sat_raw _profile_toc(state, "sat_ce", _t) sat_val = float(sat.detach()) _t = _profile_tic(prof) _aux = _collect_moe_aux(core, getattr(args,'moe_aux_coef',0.0), getattr(args,'moe_z_coef',0.0)) if torch.is_tensor(_aux): sat = sat + _aux.to(sat.dtype) scaler.scale(sat).backward() _profile_toc(state, "sat_backward", _t) del smask, emb2, zt2, h2, Ds, sat_hidden, sat_targets, sat_gate_ctx, satf, satv, sat_raw, sat if run_nat: ratio = min(max(float(getattr(args, "nat_mask_ratio", 0.5)), 0.05), 0.95) nat_mode = str(getattr(args, "dblock_nat_embed_noise_mode", "off") or "off").strip().lower() nat_noise_scale = max(0.0, float(getattr(args, "dblock_nat_embed_noise_scale", 1.0) or 1.0)) nat_ids = M._nat_ids_for_training(ids, int(getattr(args, "nat_max_tokens", 0))) _t = _profile_tic(prof) with M.amp(args.amp): nat_in = nat_ids.clone() m = M._nat_corruption_mask(nat_ids, ratio, args) if nat_mode in {"visible", "mask_plus_noise"}: clean_hn = core.emb(nat_ids) if nat_mode == "mask_plus_noise": nat_in[m] = M.BLANK hn = core.emb(nat_in) else: hn = clean_hn.clone() nat_noise = sig[:, None, None].to(clean_hn.dtype) * nat_noise_scale * torch.randn_like(clean_hn) hn = hn.clone() # mask_plus_noise must not leak the clean target embedding at masked # positions. The old code used clean_hn + noise, so training saw the # answer token while inference only has BLANK slots. noise_base = clean_hn if nat_mode == "visible" else hn hn[m] = (noise_base + nat_noise)[m] else: nat_in[m] = M.BLANK hn = core.emb(nat_in) if state.get("looped", False): hn = _dblock_loop_condition(core, hn, bi, args) for lpos, li in enumerate(layers): mode = _dblock_sublayer_mode_for_layer(args, state, bi, lpos) hn = _run_block(core.blocks[li], hn, None, _dblock_checkpoint_this_layer(args, use_layer_checkpoint, lpos, len(layers)), args, mode) Dnat = core.ln(hn) _profile_toc(state, "nat_forward", _t) _t = _profile_tic(prof) nat_hidden = Dnat[m] nat_targets = nat_ids[m] nat_hidden, nat_targets, nat_used, nat_total = _sample_token_loss_inputs( nat_hidden.unsqueeze(0), nat_targets.unsqueeze(0), _dblock_loss_token_cap(args, "nat") ) nat_raw = fused_ce(nat_hidden, nat_h.proj.weight, nat_targets) nat_raw_val = float(nat_raw.detach()) nat = nat_weight * w * nat_raw nat_val = float(nat.detach()) _profile_toc(state, "nat_ce", _t) _t = _profile_tic(prof) _aux = _collect_moe_aux(core, getattr(args,'moe_aux_coef',0.0), getattr(args,'moe_z_coef',0.0)) if torch.is_tensor(_aux): nat = nat + _aux.to(nat.dtype) scaler.scale(nat).backward() _profile_toc(state, "nat_backward", _t) del nat_ids, nat_in, m, hn, Dnat, nat_hidden, nat_targets, nat_raw, nat, nat_used, nat_total total_val = ar_val + sat_val + nat_val raw_total_val = ar_raw_val + sat_raw_val + nat_raw_val raw_count = int(bool(run_ar)) + int(bool(run_sat)) + int(bool(run_nat)) raw_avg_val = raw_total_val / max(1, raw_count) if not math.isfinite(total_val): opt.zero_grad(set_to_none=True) if torch.cuda.is_available(): torch.cuda.empty_cache() print(f"[dblock] non-finite loss {total_val}; skipped optimizer step", flush=True) _profile_toc(state, "step_total", _step_t) _profile_step_done(state, args) _update_stats(state, bi, total_val, args) return total_val _spike_k = float(getattr(args, "loss_spike_skip", 0.0)) if _spike_k > 0.0: _ema = state.get("spike_ema") if _ema is not None and _ema <= 0.0: _ema = None; state.pop("spike_ema", None) # reset degenerate zero-EMA if _ema is not None and math.isfinite(_ema) and math.isfinite(raw_avg_val) and raw_avg_val > _spike_k * _ema: opt.zero_grad(set_to_none=True) if torch.cuda.is_available(): torch.cuda.empty_cache() print(f"[dblock] loss spike raw_avg={raw_avg_val:.2f} > {_spike_k}x EMA={_ema:.2f}; skipped optimizer step", flush=True) _profile_toc(state, "step_total", _step_t) _profile_step_done(state, args) _update_stats(state, bi, total_val, args) return total_val if math.isfinite(raw_avg_val) and raw_avg_val > 1e-3: # skip near-zero state["spike_ema"] = raw_avg_val if _ema is None else (0.98 * _ema + 0.02 * raw_avg_val) _t = _profile_tic(prof) scaler.unscale_(opt) nn.utils.clip_grad_norm_([p for g in opt.param_groups for p in g["params"]], 1.0) scaler.step(opt) scaler.update() opt.zero_grad(set_to_none=True) _profile_toc(state, "opt_step", _t) peak_alloc = None peak_reserved = None if torch.cuda.is_available(): peak_alloc = torch.cuda.max_memory_allocated() / (1024**3) peak_reserved = torch.cuda.max_memory_reserved() / (1024**3) _profile_toc(state, "step_total", _step_t) _profile_step_done(state, args) _update_stats(state, bi, total_val, args) _maybe_log( state, args, bi, layers, ar_val, sat_val, nat_val, total_val, peak_alloc, peak_reserved, objective=objective, raw_avg=raw_avg_val, raw_total=raw_total_val, edm_weight=w, ) return raw_avg_val # ===== END dblocks_train.py ===== # ===== BEGIN nB300_agillm4.py ===== #!/usr/bin/env python3 # n.py - Joint AR+SAT+NAT Trainer with Expansion Ratio Testing # Enhanced inference: checkpoint name, tok/s, UK time import argparse, copy, json, math, pathlib, random, time, os, sys, threading, hashlib, re, subprocess from pathlib import Path from contextlib import nullcontext from typing import Dict, Any, List, Optional, Tuple from datetime import datetime, timezone _ASCII_LOG_TRANSLATION = str.maketrans({ "\u2018": "'", "\u2019": "'", "\u201a": "'", "\u201b": "'", "\u201c": '"', "\u201d": '"', "\u201e": '"', "\u201f": '"', "\u2013": "-", "\u2014": "-", "\u2212": "-", "\u2026": "...", "\u00a0": " ", }) def _ascii_log_text(text: str) -> str: return str(text).translate(_ASCII_LOG_TRANSLATION).encode("ascii", "replace").decode("ascii") class _AsciiLogStream: def __init__(self, wrapped): self._wrapped = wrapped def write(self, text): return self._wrapped.write(_ascii_log_text(text)) def flush(self): return self._wrapped.flush() def isatty(self): return self._wrapped.isatty() def fileno(self): return self._wrapped.fileno() @property def encoding(self): return "ascii" def __getattr__(self, name): return getattr(self._wrapped, name) if ( not sys.stdout.isatty() and os.environ.get("NB300_RAW_UNICODE_LOGS", "").lower() not in {"1", "true", "yes"} ): sys.stdout = _AsciiLogStream(sys.stdout) sys.stderr = _AsciiLogStream(sys.stderr) STATUS_SCRIPT_PATH = Path(__file__).resolve() STATUS_DEFAULT_LOG = STATUS_SCRIPT_PATH.parent / "train.log" STATUS_DEFAULT_SAVE_DIR = STATUS_SCRIPT_PATH.parent / "ckpts_expansion" _STATUS_PROGRESS_RE = re.compile( r"^\[(?P\d+(?:\.\d+)?)%\]\s+" r"(?P[\d,]+)/(?P[\d,]+)\s+tok\s+\|\s+" r"(?P[\d.]+)\s+tok/s\s+\|\s+" r"loss=(?P-?[\d.]+)\s+B=(?P\d+)\s+L=(?P\d+)" r"(?:\s+step=(?P\d+))?" r"(?:\s+eta=(?P\S+))?" r"(?:\s+elapsed=(?P\S+))?" r"\s*$" ) _STATUS_DELTA_RE = re.compile(r"\[delta\]\s+saved\s+(?P\S+?\.pt)\s+\((?P[0-9a-f]+)\.\.\.\)") _STATUS_STEP_RE = re.compile(r"step(?P\d+)") def _status_iso(ts: Optional[float]) -> Optional[str]: if ts is None: return None return datetime.fromtimestamp(ts, tz=timezone.utc).astimezone().isoformat(timespec="seconds") def _status_human_duration(seconds: Optional[float]) -> Optional[str]: if seconds is None: return None total = max(0, int(seconds)) days, rem = divmod(total, 86400) hours, rem = divmod(rem, 3600) minutes, secs = divmod(rem, 60) parts = [] if days: parts.append(f"{days}d") if hours or parts: parts.append(f"{hours}h") if minutes or parts: parts.append(f"{minutes}m") parts.append(f"{secs}s") return " ".join(parts) def _status_compact_duration(seconds: Optional[float]) -> str: if seconds is None: return "unknown" try: if not math.isfinite(float(seconds)): return "unknown" except Exception: return "unknown" total = max(0, int(seconds)) years, rem = divmod(total, 365 * 86400) days, rem = divmod(rem, 86400) hours, rem = divmod(rem, 3600) minutes, secs = divmod(rem, 60) if years: return f"{years}y{days}d{hours}h" if days: return f"{days}d{hours}h{minutes}m" if hours: return f"{hours}h{minutes}m{secs}s" if minutes: return f"{minutes}m{secs}s" return f"{secs}s" def _status_format_int(value: Optional[int]) -> str: return "?" if value is None else f"{value:,}" def _status_parse_step(text: str) -> Optional[int]: match = _STATUS_STEP_RE.search(str(text or "")) return int(match.group("step")) if match else None def _agillm43_lineage_info(source_path: Optional[str], source_provenance: Optional[dict], save_dir: str = "") -> Dict[str, Any]: source_path = str(source_path or "") try: source_abs = os.path.abspath(source_path) if source_path else "" except Exception: source_abs = source_path try: save_abs = os.path.abspath(str(save_dir or "")) if save_dir else "" except Exception: save_abs = str(save_dir or "") master_marker = f"{os.sep}agillm4_v100_master_ckpts{os.sep}" if not source_path: warmstart_kind = "from_scratch" elif master_marker in source_abs: warmstart_kind = "warmstarted_from_master" elif save_abs and source_abs.startswith(save_abs + os.sep): warmstart_kind = "warmstarted_from_lane_checkpoint" else: warmstart_kind = "warmstarted_from_non_master_checkpoint" source_step = _status_parse_step(source_path) origin_step = 0 origin_seen_tok = 0 if isinstance(source_provenance, dict): for key in ("global_origin_step", "warmstart_base_step"): try: value = int(source_provenance.get(key) or 0) except Exception: value = 0 if value > 0: origin_step = value break if origin_step <= 0: parent = source_provenance.get("warmstart_source_path") or source_provenance.get("source_path") or "" parent_step = _status_parse_step(parent) if parent_step and parent_step > 0: # AGILLM-LINEAGE-FIX 20260702: was >= 1_000_000 (broke warmstart from <1M-step recovery ckpts) origin_step = parent_step for key in ("global_origin_seen_tok", "warmstart_base_seen_tok"): try: value = int(source_provenance.get(key) or 0) except Exception: value = 0 if value > 0: origin_seen_tok = value break if origin_step <= 0 and source_step and source_step > 0: # AGILLM-LINEAGE-FIX 20260702: was 'master or >= 1_000_000' (broke <1M-step recovery warmstarts) origin_step = int(source_step) return { "source_path": source_path, "source_step": int(source_step or 0), "warmstart_kind": warmstart_kind, "created_from_scratch": warmstart_kind == "from_scratch", "source_is_master_checkpoint": warmstart_kind == "warmstarted_from_master", "source_is_lane_checkpoint": warmstart_kind == "warmstarted_from_lane_checkpoint", "source_is_non_master_checkpoint": warmstart_kind == "warmstarted_from_non_master_checkpoint", "warmstart_base_step": int(origin_step or 0), "global_origin_step": int(origin_step or 0), "warmstart_base_seen_tok": int(origin_seen_tok or 0), "global_origin_seen_tok": int(origin_seen_tok or 0), } def _status_resolve_ckpt_path(raw_path: str, base_dir: Path) -> Path: ckpt_path = Path(raw_path) return ckpt_path if ckpt_path.is_absolute() else (base_dir / ckpt_path).resolve() def _status_read_cmdline(proc_dir: Path) -> Optional[List[str]]: try: data = (proc_dir / "cmdline").read_bytes().split(b"\0") return [item.decode("utf-8", errors="ignore") for item in data if item] except Exception: return None def _status_get_arg_value(args: List[str], flag: str) -> Optional[str]: for idx, arg in enumerate(args): if arg == flag and idx + 1 < len(args): return args[idx + 1] prefix = flag + "=" if arg.startswith(prefix): return arg[len(prefix):] return None def _status_resolve_proc_arg(proc_dir: Path, raw_arg: str) -> Optional[Path]: try: arg_path = Path(raw_arg) if arg_path.is_absolute(): return arg_path.resolve() cwd = Path(os.readlink(proc_dir / "cwd")) return (cwd / arg_path).resolve() except Exception: return None def _status_proc_uptime(proc_dir: Path) -> Optional[float]: try: proc_uptime = float((Path("/proc") / "uptime").read_text().split()[0]) stat_text = (proc_dir / "stat").read_text() after = stat_text[stat_text.rfind(")") + 2:].split() start_ticks = float(after[19]) clock_ticks = os.sysconf(os.sysconf_names["SC_CLK_TCK"]) return max(0.0, proc_uptime - (start_ticks / clock_ticks)) except Exception: return None def _status_find_trainers(script_path: Path) -> List[Dict[str, Any]]: matches: List[Dict[str, Any]] = [] for proc_dir in Path("/proc").iterdir(): if not proc_dir.name.isdigit(): continue args = _status_read_cmdline(proc_dir) if not args or "train" not in args: continue if Path(args[0]).name in {"bash", "dash", "sh"} and "-c" in args[:3]: # Launch wrappers carry the trainer argv for exit logging, but are not trainers. continue resolved_script = None for arg in args: if Path(arg).name != script_path.name: continue candidate = _status_resolve_proc_arg(proc_dir, arg) if candidate == script_path: resolved_script = candidate break if resolved_script is None: continue uptime_seconds = _status_proc_uptime(proc_dir) try: cwd = str(Path(os.readlink(proc_dir / "cwd"))) except Exception: cwd = None save_dir_arg = _status_get_arg_value(args, "--save_dir") save_dir_resolved = _status_resolve_proc_arg(proc_dir, save_dir_arg) if save_dir_arg else None matches.append({ "pid": int(proc_dir.name), "cmdline": " ".join(args), "args": args, "cwd": cwd, "save_dir_arg": save_dir_arg, "save_dir_resolved": str(save_dir_resolved) if save_dir_resolved is not None else None, "uptime_seconds": round(uptime_seconds, 3) if uptime_seconds is not None else None, "uptime_human": _status_human_duration(uptime_seconds), }) return sorted(matches, key=lambda item: item["pid"]) def _status_parse_progress_line(line: str) -> Optional[Dict[str, Any]]: match = _STATUS_PROGRESS_RE.match(line.strip()) if not match: return None tok_per_sec = float(match.group("tok_s")) loss = float(match.group("loss")) return { "raw_line": line.strip(), "percent": float(match.group("percent")), "seen_tokens": int(match.group("seen").replace(",", "")), "target_tokens": int(match.group("target").replace(",", "")), "tok_per_sec": int(tok_per_sec) if tok_per_sec.is_integer() else tok_per_sec, "loss": loss, "batch": int(match.group("batch")), "block": int(match.group("block")), "step": int(match.group("step")) if match.group("step") else None, "eta": match.group("eta"), "elapsed": match.group("elapsed"), } def _status_parse_delta_line(line: str) -> Optional[Dict[str, Any]]: match = _STATUS_DELTA_RE.search(line) if not match: return None name = match.group("name") return { "raw_line": line.strip(), "name": name, "step": _status_parse_step(name), "sha_prefix": match.group("sha"), "source": "log", } def _status_scan_log(log_path: Path) -> tuple[Dict[str, Any], Optional[Dict[str, Any]], Optional[Dict[str, Any]], List[str]]: now = time.time() info: Dict[str, Any] = { "path": str(log_path), "exists": log_path.exists(), "mtime": None, "mtime_iso": None, "age_seconds": None, "age_human": None, "size_bytes": None, } warnings: List[str] = [] if not log_path.exists(): warnings.append(f"train log missing: {log_path}") return info, None, None, warnings try: st = log_path.stat() info["mtime"] = st.st_mtime info["mtime_iso"] = _status_iso(st.st_mtime) info["age_seconds"] = round(max(0.0, now - st.st_mtime), 3) info["age_human"] = _status_human_duration(info["age_seconds"]) info["size_bytes"] = st.st_size except Exception as exc: warnings.append(f"failed to stat train log: {exc}") last_progress = None last_delta = None try: with log_path.open("r", encoding="utf-8", errors="ignore") as handle: for raw_line in handle: line = raw_line.rstrip("\n") progress = _status_parse_progress_line(line) if progress is not None: last_progress = progress delta = _status_parse_delta_line(line) if delta is not None: last_delta = delta except Exception as exc: warnings.append(f"failed to read train log: {exc}") return info, last_progress, last_delta, warnings def _status_latest_full_checkpoint(save_dir: Path, base_dir: Path) -> tuple[Dict[str, Any], List[str]]: latest_path = save_dir / "latest.json" info: Dict[str, Any] = { "metadata_path": str(latest_path), "exists": latest_path.exists(), "raw_path": None, "checkpoint_path": None, "checkpoint_name": None, "checkpoint_exists": None, "step": None, "checkpoint_mtime": None, "checkpoint_mtime_iso": None, } warnings: List[str] = [] if not latest_path.exists(): warnings.append(f"latest.json missing: {latest_path}") return info, warnings try: payload = json.loads(latest_path.read_text(encoding="utf-8")) except Exception as exc: warnings.append(f"failed to parse latest.json: {exc}") return info, warnings raw_path = payload.get("path") info["raw_path"] = raw_path info["step"] = payload.get("step") for key in ( "warmstart_kind", "warmstart_source_path", "checkpoint_summary", "effective_global_step", "global_origin_step", "warmstart_base_step", "effective_seen_tok", "global_origin_seen_tok", "warmstart_base_seen_tok", ): if key in payload: info[key] = payload.get(key) provenance = payload.get("agillm43_provenance") or {} if isinstance(provenance, dict): info["agillm43_provenance"] = provenance for key in ( "effective_global_step", "global_origin_step", "warmstart_base_step", "effective_seen_tok", "global_origin_seen_tok", "warmstart_base_seen_tok", ): if key not in info and key in provenance: info[key] = provenance.get(key) if raw_path: ckpt_path = _status_resolve_ckpt_path(raw_path, base_dir) info["checkpoint_path"] = str(ckpt_path) info["checkpoint_name"] = ckpt_path.name info["checkpoint_exists"] = ckpt_path.exists() if ckpt_path.exists(): try: st = ckpt_path.stat() info["checkpoint_mtime"] = st.st_mtime info["checkpoint_mtime_iso"] = _status_iso(st.st_mtime) except Exception as exc: warnings.append(f"failed to stat full checkpoint: {exc}") else: warnings.append(f"latest.json points to missing checkpoint: {ckpt_path}") return info, warnings def _status_newest_delta(save_dir: Path) -> tuple[Optional[Dict[str, Any]], List[str]]: warnings: List[str] = [] if not save_dir.exists(): warnings.append(f"save dir missing: {save_dir}") return None, warnings try: candidates = [item for item in save_dir.glob("*_delta_step*.pt") if item.is_file()] except Exception as exc: warnings.append(f"failed to list delta checkpoints: {exc}") return None, warnings if not candidates: warnings.append(f"no delta checkpoints found in {save_dir}") return None, warnings newest = max(candidates, key=lambda item: item.stat().st_mtime) st = newest.stat() info = { "path": str(newest), "name": newest.name, "step": _status_parse_step(newest.name), "mtime": st.st_mtime, "mtime_iso": _status_iso(st.st_mtime), "size_bytes": st.st_size, "source": "disk", } sidecar = newest.with_suffix(".provenance.json") info["provenance_sidecar_path"] = str(sidecar) info["provenance_sidecar_exists"] = sidecar.exists() if sidecar.exists(): try: provenance = json.loads(sidecar.read_text(encoding="utf-8")) info["agillm43_provenance"] = provenance for key in ( "warmstart_kind", "warmstart_source_path", "local_step", "effective_global_step", "global_origin_step", "warmstart_base_step", "effective_seen_tok", "global_origin_seen_tok", "warmstart_base_seen_tok", ): if key in provenance: info[key] = provenance.get(key) except Exception as exc: warnings.append(f"failed to parse delta provenance sidecar {sidecar}: {exc}") return info, warnings def _status_gpu_info() -> tuple[Optional[Dict[str, Any]], List[str]]: warnings: List[str] = [] try: result = subprocess.run( [ "nvidia-smi", "--query-gpu=name,utilization.gpu,memory.used,memory.total,temperature.gpu,power.draw", "--format=csv,noheader,nounits", ], capture_output=True, text=True, timeout=5, check=False, ) except FileNotFoundError: return None, warnings except Exception as exc: warnings.append(f"failed to query GPU status: {exc}") return None, warnings if result.returncode != 0: warnings.append(result.stderr.strip() or "nvidia-smi returned non-zero exit status") return None, warnings lines = [line.strip() for line in result.stdout.splitlines() if line.strip()] if not lines: return None, warnings if len(lines) > 1: warnings.append("multiple GPUs detected; reporting the first GPU only") parts = [part.strip() for part in lines[0].split(",")] if len(parts) != 6: warnings.append(f"unexpected nvidia-smi format: {lines[0]}") return None, warnings def _parse_int(raw: str) -> Optional[int]: try: return int(float(raw)) except Exception: return None def _parse_float(raw: str) -> Optional[float]: try: return float(raw) except Exception: return None return { "name": parts[0], "utilization_gpu": _parse_int(parts[1]), "memory_used_mib": _parse_int(parts[2]), "memory_total_mib": _parse_int(parts[3]), "temperature_c": _parse_int(parts[4]), "power_draw_w": _parse_float(parts[5]), }, warnings def _status_choose_delta(from_log: Optional[Dict[str, Any]], from_disk: Optional[Dict[str, Any]], warnings: List[str]) -> Optional[Dict[str, Any]]: if from_log and from_disk: log_step = from_log.get("step") disk_step = from_disk.get("step") if log_step is not None and disk_step is not None: if log_step != disk_step: warnings.append( f"log delta step {log_step} and newest on-disk delta step {disk_step} differ; using the newer step" ) if disk_step >= log_step: merged = dict(from_disk) merged["source"] = "disk+log" if disk_step == log_step else "disk" if disk_step == log_step: merged["sha_prefix"] = from_log.get("sha_prefix") return merged return dict(from_log) return dict(from_disk) if from_disk: return dict(from_disk) if from_log: return dict(from_log) return None def _collect_status(log_path: Path, save_dir: Path) -> tuple[Dict[str, Any], int]: checked_at = time.time() requested_save_dir = save_dir.expanduser() log_path = log_path.expanduser() status: Dict[str, Any] = { "checked_at": checked_at, "checked_at_iso": _status_iso(checked_at), "running": False, "process": None, "progress": None, "delta_checkpoint": None, "delta_from_log": None, "delta_on_disk": None, "latest_full_checkpoint": None, "log": None, "gpu": None, "save_dir": { "requested_path": str(requested_save_dir), "path": str(requested_save_dir), "exists": requested_save_dir.exists(), "source": "requested", }, "warnings": [], } warnings = status["warnings"] matches = _status_find_trainers(STATUS_SCRIPT_PATH) requested_resolved = requested_save_dir.resolve() save_dir_matches = [ item for item in matches if item.get("save_dir_resolved") and Path(item["save_dir_resolved"]).resolve() == requested_resolved ] if save_dir_matches: matches = save_dir_matches elif len(matches) > 1: warnings.append(f"no active trainer command line matched requested save_dir exactly: {requested_resolved}") if len(matches) > 1: status["error"] = f"multiple active {STATUS_SCRIPT_PATH.name} train processes found" status["processes"] = matches return status, 1 if matches: status["running"] = True status["process"] = matches[0] save_dir = requested_save_dir if status["process"] and status["process"].get("cwd"): proc_cwd = Path(status["process"]["cwd"]) alt_save_dir = (proc_cwd / requested_save_dir.name).resolve() if alt_save_dir != requested_save_dir and alt_save_dir.exists(): requested_delta, _ = _status_newest_delta(requested_save_dir) requested_full, _ = _status_latest_full_checkpoint(requested_save_dir, STATUS_SCRIPT_PATH.parent) alt_delta, _ = _status_newest_delta(alt_save_dir) alt_full, _ = _status_latest_full_checkpoint(alt_save_dir, proc_cwd) requested_score = int(requested_delta is not None) + int(bool(requested_full.get("checkpoint_exists"))) alt_score = int(alt_delta is not None) + int(bool(alt_full.get("checkpoint_exists"))) if alt_score > requested_score: save_dir = alt_save_dir status["save_dir"] = { "requested_path": str(requested_save_dir), "path": str(save_dir), "exists": save_dir.exists(), "source": "process_cwd_fallback", } warnings.append( f"using process cwd save dir fallback: {save_dir} (requested {requested_save_dir})" ) log_info, progress, delta_from_log, log_warnings = _status_scan_log(log_path) warnings.extend(log_warnings) status["log"] = log_info status["progress"] = progress status["delta_from_log"] = delta_from_log latest_base_dir = STATUS_SCRIPT_PATH.parent if status["save_dir"].get("source") == "process_cwd_fallback" and status["process"] and status["process"].get("cwd"): latest_base_dir = Path(status["process"]["cwd"]) latest_full, latest_warnings = _status_latest_full_checkpoint(save_dir, latest_base_dir) warnings.extend(latest_warnings) status["latest_full_checkpoint"] = latest_full delta_on_disk, delta_warnings = _status_newest_delta(save_dir) warnings.extend(delta_warnings) status["delta_on_disk"] = delta_on_disk status["delta_checkpoint"] = _status_choose_delta(delta_from_log, delta_on_disk, warnings) gpu, gpu_warnings = _status_gpu_info() warnings.extend(gpu_warnings) status["gpu"] = gpu if status["running"] and log_info.get("age_seconds") is not None and log_info["age_seconds"] > 600: warnings.append(f"train log appears stale while trainer is running ({log_info['age_human']} old)") if log_info.get("exists") and progress is None: warnings.append("no parseable progress line found in train log") latest_step = latest_full.get("step") if latest_full else None delta_step = status["delta_checkpoint"].get("step") if status["delta_checkpoint"] else None if latest_step is not None and delta_step is not None and latest_step < delta_step: warnings.append(f"latest.json step {latest_step} lags newest delta step {delta_step}") if not status["running"] and progress is None: warnings.append("no active trainer process found") return status, 0 def _format_status_text(status: Dict[str, Any]) -> str: lines = [f"AGILLM status @ {status.get('checked_at_iso')}"] if status.get("error"): lines.append(f"Error: {status['error']}") for proc in status.get("processes", []): lines.append(f"- pid {proc.get('pid')}: {proc.get('cmdline')}") return "\n".join(lines) process = status.get("process") if status.get("running") and process: lines.append(f"Process: RUNNING | pid {process.get('pid')} | uptime {process.get('uptime_human') or 'unknown'}") lines.append(f"Cmd: {process.get('cmdline')}") else: lines.append("Process: NOT RUNNING") progress = status.get("progress") if progress: eta = progress.get("eta") if not eta and progress.get("tok_per_sec"): remaining = max(0, progress["target_tokens"] - progress["seen_tokens"]) eta = _status_compact_duration(remaining / float(progress["tok_per_sec"])) lines.append( "Progress: " f"{progress['percent']:.1f}% | " f"{_status_format_int(progress['seen_tokens'])}/{_status_format_int(progress['target_tokens'])} tok | " f"{progress['tok_per_sec']} tok/s | loss {progress['loss']:.3f} | " f"B={progress['batch']} L={progress['block']}" + (f" | step {progress['step']}" if progress.get("step") else "") + (f" | ETA {eta}" if eta else "") ) else: lines.append("Progress: unavailable") log_info = status.get("log") or {} if log_info.get("exists"): lines.append( f"Log: {log_info.get('path')} | updated {log_info.get('age_human') or 'unknown'} ago | " f"mtime {log_info.get('mtime_iso')}" ) else: lines.append(f"Log: missing ({log_info.get('path')})") delta = status.get("delta_checkpoint") if delta: line = f"Delta: {delta.get('name')} | step {delta.get('step')} | source {delta.get('source')}" if delta.get("path"): line += f" | {delta['path']}" lines.append(line) else: lines.append("Delta: unavailable") latest_full = status.get("latest_full_checkpoint") or {} if latest_full.get("exists"): lines.append( f"Latest full: step {latest_full.get('step')} | {latest_full.get('checkpoint_path') or latest_full.get('raw_path')}" ) else: lines.append(f"Latest full: unavailable ({latest_full.get('metadata_path')})") gpu = status.get("gpu") if gpu: lines.append( f"GPU: {gpu.get('name')} | {gpu.get('utilization_gpu')}% | " f"{gpu.get('memory_used_mib')}/{gpu.get('memory_total_mib')} MiB | " f"{gpu.get('temperature_c')}C | {gpu.get('power_draw_w')} W" ) warnings = status.get("warnings") or [] if warnings: lines.append("Warnings:") lines.extend(f"- {warning}" for warning in warnings) return "\n".join(lines) def _emit_status(log_path: Path, save_dir: Path, as_json: bool) -> int: status, exit_code = _collect_status(log_path, save_dir) if as_json: print(json.dumps(status, indent=2, sort_keys=True)) else: print(_format_status_text(status)) return exit_code def _run_status_command(argv: List[str]) -> int: parser = argparse.ArgumentParser(prog=f"{STATUS_SCRIPT_PATH.name} status", description="Read-only training status") parser.add_argument("--json", dest="json_output", action="store_true", help="Emit machine-readable JSON") parser.add_argument("--log", type=Path, default=STATUS_DEFAULT_LOG, help="Path to the training log") parser.add_argument("--save_dir", type=Path, default=STATUS_DEFAULT_SAVE_DIR, help="Checkpoint directory") args = parser.parse_args(argv) return _emit_status(args.log, args.save_dir, args.json_output) def _maybe_handle_status_fastpath() -> None: if len(sys.argv) > 1 and sys.argv[1] == "status": raise SystemExit(_run_status_command(sys.argv[2:])) _maybe_handle_status_fastpath() import torch import torch.utils.checkpoint as torch_checkpoint # SafeProgress - Claude-safe progress (discrete lines, not single growing line) class SafeProgress: def __init__(self, total, initial=0, unit="tok", print_every=100, print_every_sec=60, initial_step=0): self.total, self.n, self.unit = total, initial, unit self.initial = initial self.last_print, self.postfix = initial, {} self.print_every = max(1, int(print_every)) self.print_every_sec = max(1, int(print_every_sec)) self.step = int(initial_step or 0) self.last_print_step = self.step self.start_time = __import__('time').time() self.last_print_time = self.start_time def update(self, n=1): self.n += n self.step += 1 now = __import__('time').time() if ( self.step == 1 or (self.step - self.last_print_step) >= self.print_every or (now - self.last_print_time) >= self.print_every_sec ): self._print(now) self.last_print = self.n self.last_print_step = self.step self.last_print_time = now def set_postfix(self, **kwargs): self.postfix = kwargs def _print(self, now=None): now = now or __import__('time').time() elapsed = now - self.start_time rate = (self.n - self.initial) / elapsed if elapsed > 0 else 0 pct = 100 * self.n / self.total if self.total > 0 else 0 pf = ' '.join(f"{k}={v}" for k,v in self.postfix.items()) remaining = max(0, self.total - self.n) eta = _status_compact_duration(remaining / rate) if rate > 0 else "unknown" elapsed_s = _status_compact_duration(elapsed) print( f"[{pct:.4f}%] {self.n:,}/{self.total:,} {self.unit} | " f"{rate:.2f} tok/s | {pf} step={self.step} eta={eta} elapsed={elapsed_s}", flush=True, ) def close(self): self._print(); print("Done.", flush=True) import torch.nn as nn import torch.nn.functional as F import signal import os from datasets import load_dataset, DownloadConfig from transformers import AutoTokenizer, logging as hf_log # from tqdm.auto import tqdm # DISABLED - kills Claude context # ─────────────────────────────── HOT DATASET LOADING ─────────────────────────────── HOT_CONFIG_PATH = Path(os.environ.get("AGILLM_HOT_CONFIG") or os.environ.get("AGILLM_HOT_CONFIG_PATH") or "/workspace/hot_config.json") DEFAULT_LANGUAGE_PRETRAIN_SOURCES = os.environ.get( "AGILLM_DEFAULT_LANGUAGE_PRETRAIN_SOURCES", "HuggingFaceFW/fineweb,HuggingFaceFW/fineweb-edu:sample-10BT,wikimedia/wikipedia:20231101.en,allenai/c4:en,Skylion007/openwebtext,tiiuae/falcon-refinedweb,EleutherAI/proof-pile-2,allenai/dolma:v1_6-sample", ) _hot_config_cache = {"mtime": 0, "data": {}} def get_hot_config() -> dict: """Load hot_config.json with caching, return empty dict if missing""" try: if HOT_CONFIG_PATH.exists(): mtime = HOT_CONFIG_PATH.stat().st_mtime if mtime > _hot_config_cache["mtime"]: with open(HOT_CONFIG_PATH) as f: _hot_config_cache["data"] = json.load(f) _hot_config_cache["mtime"] = mtime return _hot_config_cache["data"] except Exception as e: print(f"[hot_config] Error loading: {e}") return {} def _hot_int_from_config(cfg: dict, names: list[str], default: int) -> int: """Read a non-negative int from hot_config, accepting top-level or dblock-nested keys.""" if not isinstance(cfg, dict): return int(default) candidates = [] for name in names: candidates.append(cfg.get(name)) nested = cfg.get("dblock") if isinstance(nested, dict): for name in names: candidates.append(nested.get(name)) for value in candidates: if value is None or value == "": continue try: return max(0, int(value)) except Exception: continue return int(default) def _hot_float_from_config(cfg: dict, names: list[str], default: float, min_value=None, max_value=None) -> float: """Read a float from hot_config, accepting top-level or dblock-nested keys.""" if not isinstance(cfg, dict): return float(default) candidates = [] for name in names: candidates.append(cfg.get(name)) nested = cfg.get("dblock") if isinstance(nested, dict): for name in names: candidates.append(nested.get(name)) for value in candidates: if value is None or value == "": continue try: out = float(value) except Exception: continue if min_value is not None: out = max(float(min_value), out) if max_value is not None: out = min(float(max_value), out) return out return float(default) _hot_dblock_loss_tokens_seen = {} _hot_dblock_float_seen = {} def _dblock_hot_float(args, attr: str, default: float, names=None, min_value=None, max_value=None) -> float: cli_default = float(getattr(args, attr, default) or default) keys = list(names or [attr]) if attr not in keys: keys.insert(0, attr) try: cfg = get_hot_config() except Exception: return cli_default value = _hot_float_from_config(cfg, keys, cli_default, min_value=min_value, max_value=max_value) key = (id(args), attr) if _hot_dblock_float_seen.get(key) != value: _hot_dblock_float_seen[key] = value if value != cli_default: print(f"[hot_config] {attr}={value} (cli_default={cli_default})", flush=True) return value def _dblock_loss_token_cap(args, objective: str) -> int: """Hot-reload AR/SAT/NAT sampled CE token caps from hot_config.json. Supported hot_config keys: dblock_loss_tokens: shared default for AR/SAT/NAT dblock_ar_loss_tokens / dblock_sat_loss_tokens / dblock_nat_loss_tokens dblock: {loss_tokens, ar_loss_tokens, sat_loss_tokens, nat_loss_tokens} """ obj = str(objective or "").strip().lower() attr = f"dblock_{obj}_loss_tokens" default = int(getattr(args, attr, 0) or 0) try: cfg = get_hot_config() except Exception: return default value = _hot_int_from_config( cfg, [attr, f"{obj}_loss_tokens", "dblock_loss_tokens", "loss_tokens"], default, ) key = (id(args), attr) if _hot_dblock_loss_tokens_seen.get(key) != value: _hot_dblock_loss_tokens_seen[key] = value if value != default: print(f"[hot_config] {attr}={value} (cli_default={default})", flush=True) return value def _dataset_config_to_csv(value) -> str: if isinstance(value, list): return ",".join(str(item).strip() for item in value if str(item).strip()) return str(value or "").strip() def _dataset_specs_csv(value: str) -> list[str]: return [part.strip() for part in str(value or "").split(",") if part.strip()] def _dataset_spec_without_weight(spec: str) -> str: head, sep, tail = str(spec or "").strip().rpartition("|") if sep: try: float(tail) return head.strip() except Exception: pass return str(spec or "").strip() def _dataset_merge_csv(*groups: str) -> str: """Merge dataset CSV groups; later duplicate specs update weight/config but never remove defaults.""" ordered = [] by_key = {} for group in groups: for spec in _dataset_specs_csv(group): key = _dataset_spec_without_weight(spec) if not key: continue if key not in by_key: ordered.append(key) by_key[key] = spec return ",".join(by_key[key] for key in ordered if key in by_key) def _dataset_remove_csv(sources: str, removals: str) -> str: """Remove dataset specs by unweighted key or substring.""" specs = _dataset_specs_csv(sources) remove_specs = _dataset_specs_csv(removals) if not specs or not remove_specs: return str(sources or "").strip() remove_keys = {_dataset_spec_without_weight(spec) for spec in remove_specs if spec} remove_needles = {key.lower() for key in remove_keys if key} kept = [] for spec in specs: key = _dataset_spec_without_weight(spec) key_l = key.lower() if key in remove_keys or any(needle and needle in key_l for needle in remove_needles): continue kept.append(spec) return ",".join(kept) def _looks_like_numeracy_source(spec: str) -> bool: base = _dataset_spec_without_weight(spec).lower() return "agillm_math_numeracy" in base or "math_numeracy_synth" in base def _looks_numeracy_only_sources(sources: str) -> bool: specs = _dataset_specs_csv(sources) return bool(specs) and all(_looks_like_numeracy_source(spec) for spec in specs) def _language_pretrain_fallback_sources() -> str: return str( os.environ.get("AGILLM_LANGUAGE_PRETRAIN_SOURCES") or globals().get("DEFAULT_LANGUAGE_PRETRAIN_SOURCES", "") or globals().get("DEFAULT_PRETRAIN_SOURCES", "") ).strip() def _augment_numeracy_only_sources(default_sources: str) -> str: default_sources = str(default_sources or "").strip() disabled = str(os.environ.get("AGILLM_DISABLE_LANGUAGE_FALLBACK", "")).strip().lower() in {"1", "true", "yes", "on"} if disabled or not _looks_numeracy_only_sources(default_sources): return default_sources language_sources = _language_pretrain_fallback_sources() if not language_sources: return default_sources print( "[dataset-policy] numeracy-only pretrain source replaced with built-in language pretrain mix; " "numeracy_weight=0", flush=True, ) return language_sources def get_hot_datasets(default_sources: str) -> str: """Merge hot_config datasets into the safe default mix instead of replacing it.""" cfg = get_hot_config() sources = _augment_numeracy_only_sources(default_sources) hot_ds = _dataset_config_to_csv(cfg.get("datasets")) if hot_ds: sources = _dataset_merge_csv(sources, hot_ds) print(f"[hot_config] Merged datasets into default mix: {hot_ds}", flush=True) append_ds = _dataset_config_to_csv(cfg.get("datasets_append") or cfg.get("extra_datasets")) if append_ds: sources = _dataset_merge_csv(sources, append_ds) print(f"[hot_config] Appended datasets: {append_ds}", flush=True) remove_ds = _dataset_config_to_csv(cfg.get("datasets_remove") or cfg.get("remove_datasets")) if remove_ds: before = sources sources = _dataset_remove_csv(sources, remove_ds) if sources != before: print(f"[hot_config] Removed datasets from mix: {remove_ds}", flush=True) return sources def _dataset_source_summary(sources: str) -> dict: specs = _dataset_specs_csv(sources) return { "count": len(specs), "specs": specs, "has_language_mix": any(("fineweb" in s.lower()) or ("wikipedia" in s.lower()) or ("c4" in s.lower()) or ("proof-pile" in s.lower()) or ("txt360" in s.lower()) for s in specs), "has_numeracy": any(_looks_like_numeracy_source(s) for s in specs), } def _dataset_provenance(phase_name: str, requested_source: str, effective_source: str, args, *, use_hot_config: bool = True, val_requested: str = "", val_effective: str = "") -> dict: cfg = get_hot_config() if use_hot_config else {} hot_mtime = None try: hot_mtime = HOT_CONFIG_PATH.stat().st_mtime if HOT_CONFIG_PATH.exists() else None except Exception: hot_mtime = None summary = _dataset_source_summary(effective_source) return { "schema": "agillm.dataset_provenance.v1", "phase": str(phase_name), "source_requested": str(requested_source or ""), "source_effective": str(effective_source or ""), "source_count": int(summary["count"]), "source_specs": list(summary["specs"]), "has_language_mix": bool(summary["has_language_mix"]), "has_numeracy": bool(summary["has_numeracy"]), "hot_config_path": str(HOT_CONFIG_PATH), "hot_config_mtime": hot_mtime, "hot_config_used": bool(use_hot_config), "hot_config_has_datasets": bool(cfg.get("datasets")), "hot_config_has_append": bool(cfg.get("datasets_append") or cfg.get("extra_datasets")), "val_source_requested": str(val_requested or ""), "val_source_effective": str(val_effective or ""), "dataset_field_text": str(getattr(args, "dataset_field_text", "text")), "chat": bool(getattr(args, "chat", False)), } # DISABLED: # Auto-rotating log to prevent context-window suicide # DISABLED: try: # DISABLED: from rotating_log import install_rotating_log # DISABLED: install_rotating_log() # DISABLED: except ImportError: # pass # Running without rotation # ───────────────────────── ASCII Sanitizer ───────────────────────── def _ascii_safe(s): if not isinstance(s, str): return s return (s .replace('\u2019', "'").replace('\u2018', "'") .replace('\u201C', '"').replace('\u201D', '"') .replace('\u2014', '-').replace('\u2013', '-') .replace('\u2026', '...') .replace('\u00A0', ' ')) # ───────────────────────── ANSI Colors ───────────────────────── class Colors: RESET = "\033[0m" BOLD = "\033[1m" PROMPT = "\033[36m" GEN = "\033[0m" INFO = "\033[90m" WARN = "\033[93m" # ───────────────────────── Globals ───────────────────────── hf_log.set_verbosity_error() DEV = torch.device("cuda" if torch.cuda.is_available() else "cpu") torch.backends.cuda.matmul.allow_tf32 = True try: torch.set_float32_matmul_precision("high") except Exception: pass TOKENIZER_ID = os.environ.get("TOKENIZER_ID", "deepseek-ai/DeepSeek-V4-Pro") SYNTHETIC_TOKENIZER = os.environ.get("AGILLM_SYNTHETIC_TOKENIZER", "").lower() in {"1", "true", "yes"} class _SyntheticTokenizer: pad_token = "<|pad|>" pad_token_id = 0 eos_token_id = 1 sep_token_id = 1 def __init__(self, vocab_size: int): self.vocab_size = vocab_size self.backend_tokenizer = self def add_special_tokens(self, _tokens): return 0 def get_vocab(self): return {f"tok_{i}": i for i in range(self.vocab_size)} def encode(self, text): return [2 + (ord(ch) % max(1, self.vocab_size - 2)) for ch in str(text)] def decode(self, ids, skip_special_tokens=True): return " ".join(f"tok{int(i)}" for i in ids if not skip_special_tokens or int(i) > 1) def to_str(self): return json.dumps({"type": "synthetic", "vocab_size": self.vocab_size}) if SYNTHETIC_TOKENIZER: tok = _SyntheticTokenizer(int(os.environ.get("AGILLM_SYNTHETIC_VOCAB", "8192"))) print(f"[tokenizer] synthetic tokenizer enabled vocab={tok.vocab_size}") else: _tok_src = os.environ.get("TOKENIZER_DIR", "/workspace/tokenizers/deepseek-v4-pro") if not os.path.isdir(_tok_src): _tok_src = TOKENIZER_ID try: tok = AutoTokenizer.from_pretrained(_tok_src, use_fast=True, trust_remote_code=True, local_files_only=True) except Exception as _tok_exc: print(f"[tokenizer] offline load from {_tok_src} failed ({_tok_exc}); network fallback {TOKENIZER_ID}", flush=True) tok = AutoTokenizer.from_pretrained(TOKENIZER_ID, use_fast=True, trust_remote_code=True) if tok.pad_token is None: tok.add_special_tokens({"pad_token": "<|pad|>"}) # ─── Fix tokenizer Ġ/▁ mismatch ─── # Some DeepSeek tokenizer releases use Ġ (U+0120) for space-prefixed tokens, # but some transformers versions set the Metaspace pre-tokenizer to use # ▁ (U+2581) instead, causing encode/decode to lose all spaces. def _set_backend_tokenizer(tokenizer, backend) -> None: """Swap a fast tokenizer backing tokenizers.Tokenizer across transformers versions. Modern transformers expose backend_tokenizer as a READ-ONLY property backed by _tokenizer; older versions allow direct assignment. Setting _tokenizer is what makes the checkpoint tokenizer-restore actually take effect (it was failing silently).""" try: tokenizer._tokenizer = backend return except Exception: pass tokenizer.backend_tokenizer = backend def _tokenizer_payload() -> dict: """Embed enough tokenizer state for checkpoints/deltas to be self-contained. tokenizer_json is the exact fast-tokenizer backend. tokenizer_bundle stores the small save_pretrained() files as text for environments that need config/special token metadata too. This is intentionally best-effort so a tokenizer hiccup never aborts a model save. """ out = {"tokenizer_payload_schema": 2} try: out["tokenizer_id"] = TOKENIZER_ID except Exception: pass try: out["tokenizer_json"] = tok.backend_tokenizer.to_str() except Exception as e: print(f"[tokenizer] WARNING: could not embed tokenizer_json in checkpoint: {e}") try: out["tokenizer_special"] = { "pad_token": getattr(tok, "pad_token", None), "pad_token_id": getattr(tok, "pad_token_id", None), "eos_token": getattr(tok, "eos_token", None), "eos_token_id": getattr(tok, "eos_token_id", None), "sep_token": getattr(tok, "sep_token", None), "sep_token_id": getattr(tok, "sep_token_id", None), "vocab_size": len(tok.get_vocab()) if hasattr(tok, "get_vocab") else None, } except Exception: pass try: import tempfile bundle = {} with tempfile.TemporaryDirectory(prefix="agillm_tok_") as td: tok.save_pretrained(td) for item in Path(td).iterdir(): if item.is_file() and item.stat().st_size <= 64 * 1024 * 1024: try: bundle[item.name] = item.read_text(encoding="utf-8") except UnicodeDecodeError: import base64 bundle[item.name] = {"base64": base64.b64encode(item.read_bytes()).decode("ascii")} if bundle: out["tokenizer_bundle"] = bundle except Exception as e: print(f"[tokenizer] WARNING: could not embed tokenizer bundle in checkpoint: {e}") return out def _tokenizer_sidecar_paths(path): try: p = Path(path) except Exception: return [] return [ Path(str(p) + ".tokenizer.json"), p.with_suffix(p.suffix + ".tokenizer.json"), p.parent / (p.name + ".tokenizer.json"), ] def _read_tokenizer_sidecar(path): import json as _json if not path: return {} for sidecar in _tokenizer_sidecar_paths(path): try: if sidecar.exists(): obj = _json.loads(sidecar.read_text(encoding="utf-8")) if isinstance(obj, dict): obj.setdefault("tokenizer_sidecar", str(sidecar)) return obj except Exception as exc: print(f"[tokenizer] WARNING: could not read tokenizer sidecar {sidecar}: {exc}") return {} def _write_tokenizer_sidecar(path, payload) -> None: """Write tokenizer metadata beside a full checkpoint and as latest.tokenizer.json.""" try: p = Path(path) data = dict(payload or {}) if data.get("tokenizer_json") and not data.get("tokenizer_payload_schema"): data["tokenizer_payload_schema"] = 2 data.setdefault("tokenizer_payload_schema", 2) data["checkpoint_name"] = p.name data["checkpoint_path"] = str(p) for out in (Path(str(p) + ".tokenizer.json"), p.parent / "latest.tokenizer.json"): tmp = Path(str(out) + ".tmp") tmp.write_text(json.dumps(data, ensure_ascii=False, sort_keys=True), encoding="utf-8") tmp.replace(out) except Exception as exc: print(f"[tokenizer] WARNING: could not write tokenizer sidecar for {path}: {exc}") def _apply_tokenizer_special(payload) -> None: try: spec = payload.get("tokenizer_special") if hasattr(payload, "get") else None if not isinstance(spec, dict): return if spec.get("pad_token") is not None: tok.pad_token = spec.get("pad_token") if spec.get("eos_token") is not None: tok.eos_token = spec.get("eos_token") if spec.get("sep_token") is not None: tok.sep_token = spec.get("sep_token") except Exception as exc: print(f"[tokenizer] WARNING: special-token restore skipped: {exc}") def _restore_tokenizer_from_ckpt(d, ckpt_path=None) -> None: """Make tok match what a checkpoint/delta was trained with. Embedded tokenizer_json is exact and preferred. A sidecar produced for older checkpoints is next. Runtime TOKENIZER_ID is last-resort compatibility only. Never raises: a tokenizer issue must not abort load/infer. """ try: payload = d if hasattr(d, "get") else {} if ckpt_path: sidecar = _read_tokenizer_sidecar(ckpt_path) if sidecar: merged = dict(sidecar) # Embedded checkpoint fields win, but sidecars can fill schema, # special-token metadata, or bundle files missing from old saves. merged.update({k: v for k, v in payload.items() if str(k).startswith("tokenizer_") and v is not None}) payload = merged tj = payload.get("tokenizer_json") if hasattr(payload, "get") else None if tj: from tokenizers import Tokenizer as _Tokenizer _set_backend_tokenizer(tok, _Tokenizer.from_str(tj)) _apply_tokenizer_special(payload) source = payload.get("tokenizer_sidecar") or "checkpoint" print(f"[tokenizer] Restored from {source}") return tid = payload.get("tokenizer_id") if hasattr(payload, "get") else None if tid and tid != TOKENIZER_ID: print(f"[tokenizer] WARNING: checkpoint trained with tokenizer_id={tid} but runtime TOKENIZER_ID={TOKENIZER_ID}; set TOKENIZER_ID to match") elif tid: print(f"[tokenizer] checkpoint tokenizer_id={tid} matches runtime (no embedded json)") else: print("[tokenizer] no tokenizer embedded in checkpoint; using runtime default") except Exception as e: print(f"[tokenizer] WARNING: tokenizer restore skipped: {e}") def _fix_tokenizer_space_mismatch(tokenizer): try: import json as _json from tokenizers import Tokenizer as _Tokenizer bt = tokenizer.backend_tokenizer tj = _json.loads(bt.to_str()) pre = tj.get("pre_tokenizer", {}) needs_fix = (pre.get("type") == "Metaspace" and pre.get("replacement") == "\u2581") if not needs_fix: return # Check if vocab actually uses Ġ (U+0120) for spaces vocab = tj.get("model", {}).get("vocab", {}) has_gpt2_space = any(k.startswith("\u0120") for k in list(vocab.keys())[:500]) if not has_gpt2_space: return # Patch pre_tokenizer: ▁ -> Ġ tj["pre_tokenizer"]["replacement"] = "\u0120" # Patch decoder: ▁ -> Ġ in Replace step for step in tj.get("decoder", {}).get("decoders", []): if step.get("type") == "Replace": pat = step.get("pattern", {}) if pat.get("String") == "\u2581": pat["String"] = "\u0120" # Rebuild backend tokenizer fixed = _Tokenizer.from_str(_json.dumps(tj)) _set_backend_tokenizer(tokenizer, fixed) # Verify fix test_ids = tokenizer.encode("hello world") test_dec = tokenizer.decode(test_ids, skip_special_tokens=True) if "hello world" in test_dec: print("[tokenizer] Fixed Ġ/▁ space mismatch") else: print(f"[tokenizer] WARNING: fix applied but decode test failed: {repr(test_dec)}") except Exception as e: print(f"[tokenizer] Could not fix space mismatch: {e}") if not SYNTHETIC_TOKENIZER: _fix_tokenizer_space_mismatch(tok) # ─── Tokenizer startup health check ─── # Abort early if tokenizer can't roundtrip spaces — prevents silent data corruption def _tokenizer_health_check(tokenizer): import transformers as _tf ver = _tf.__version__ print(f"[tokenizer] transformers={ver}, tokenizers={__import__('tokenizers').__version__}") # Warn on known-bad versions try: from packaging.version import Version if Version(ver) >= Version('5.0.0'): print(f'[tokenizer] WARNING: transformers {ver} may have Metaspace bug — verify carefully') except ImportError: pass # Roundtrip tests — must preserve spaces tests = [ 'Water boils at one hundred degrees', 'The quick brown fox jumps over the lazy dog', 'Hello world! This is a test sentence with spaces.', ] for text in tests: ids = tokenizer.encode(text) decoded = tokenizer.decode(ids, skip_special_tokens=True) if ' ' not in decoded: print(f'[tokenizer] FATAL: Roundtrip lost all spaces!') print(f' Input: {repr(text)}') print(f' Encoded: {ids[:20]}...') print(f' Decoded: {repr(decoded)}') print(f'[tokenizer] ABORTING — fix tokenizer before training!') sys.exit(1) # Check decoded is reasonably close to input if text.lower().split()[:3] != decoded.lower().split()[:3]: print(f'[tokenizer] WARNING: Roundtrip diverged:') print(f' Input: {repr(text[:60])}') print(f' Decoded: {repr(decoded[:60])}') print(f'[tokenizer] Health check PASSED — spaces preserved in roundtrip') if not SYNTHETIC_TOKENIZER: _tokenizer_health_check(tok) VOCAB, BLANK, EOS = ( max(tok.get_vocab().values()) + 1, int(getattr(tok, "pad_token_id", 0) or 0), tok.eos_token_id if tok.eos_token_id is not None else tok.sep_token_id ) # ───────────────────────── PRESETS ───────────────────────── PRESETS: Dict[str, Dict[str, int]] = { "femto_1x": dict(d=16, layers=1, heads=1, rank=16), "femto_12x": dict(d=16, layers=1, heads=1, rank=192), "femto_24x": dict(d=16, layers=1, heads=1, rank=384), "pico_1x": dict(d=32, layers=1, heads=2, rank=16), "pico_3x": dict(d=32, layers=1, heads=2, rank=48), "pico_6x": dict(d=32, layers=1, heads=2, rank=96), "pico_12x": dict(d=32, layers=1, heads=2, rank=192), "pico_24x": dict(d=32, layers=1, heads=2, rank=384), "pico_48x": dict(d=32, layers=1, heads=2, rank=768), "nano_1x": dict(d=64, layers=2, heads=4, rank=16), "nano_3x": dict(d=64, layers=2, heads=4, rank=48), "nano_6x": dict(d=64, layers=2, heads=4, rank=96), "nano_12x": dict(d=64, layers=2, heads=4, rank=192), "nano_24x": dict(d=64, layers=2, heads=4, rank=384), "nano_48x": dict(d=64, layers=2, heads=4, rank=768), "nano_96x": dict(d=64, layers=2, heads=4, rank=1536), "micro_3x": dict(d=128, layers=4, heads=8, rank=48), "micro_6x": dict(d=128, layers=4, heads=8, rank=96), "micro_12x": dict(d=128, layers=4, heads=8, rank=192), "micro_24x": dict(d=128, layers=4, heads=8, rank=384), "small": dict(d=512, layers=8, heads=16, rank=64), "smallx2": dict(d=512, layers=16, heads=16, rank=64), "base": dict(d=768, layers=12, heads=24, rank=96), "base18": dict(d=768, layers=18, heads=24, rank=96), "large": dict(d=1024, layers=24, heads=16, rank=128), # AGILLM-4 tiers. These are intentionally above the ~700M AGILLM-3 size. # Approx dense parameter count with the current untied embedding+AR+SAT+NAT heads: # agillm4_floor ~= 1.21B, agillm4_main ~= 1.70B, agillm4_big ~= 2.40B. "agillm4_floor": dict(d=1280, layers=28, heads=20, rank=160), "agillm4_main": dict(d=1536, layers=32, heads=24, rank=192), "agillm4_big": dict(d=1792, layers=36, heads=28, rank=224), } DEFAULT_BLOCK = 1122 DEFAULT_BATCH = 4 SAT_BLOCK = 2 LR_CORE, LR_HEAD = 5e-5, 2e-4 EMIT_LAMBDA = 0.1 DEFAULT_SAVE_SEC = 24 * 3600 DEFAULT_DELTA_STEPS = 0 # step-triggered delta saves disabled; use DEFAULT_DELTA_SEC DEFAULT_DELTA_SEC = int(os.environ.get("AGILLM43_DELTA_EVERY_SEC", "3600")) # lightweight weight-only save every N seconds DEFAULT_MAX_DELTAS = 5 # keep last N deltas (older pruned after full save) CKDIR = pathlib.Path("ckpts_expansion") DEFAULT_PRETRAIN_SOURCES = "LLM360/TxT360,OpenTransformer/goddess-crawl,OpenTransformer/agillm-crawl-data,OpenTransformer/web-crawl-2026,OpenTransformer/web-crawl-clean-v2,OpenTransformer/scraped-web-data,OpenTransformer/turbo-crawl,OpenTransformer/sft-data-clean,OpenTransformer/web-crawl-v1,HuggingFaceFW/fineweb,wikimedia/wikipedia:20231101.en,allenai/c4:en,EleutherAI/proof-pile-2" DEFAULT_AFTER_SFT_SOURCES = "mlabonne/opc-sft-stage2-chat,HuggingFaceH4/ultrachat_200k@train_sft" DEFAULT_AFTER_SFT_BLOCK = 768 DEFAULT_ATTN_BACKEND = os.environ.get("AGILLM_ATTN_BACKEND", "manual") def _env_int(name: str, default: int) -> int: try: return int(os.environ.get(name, default)) except (TypeError, ValueError): return default DEFAULT_SUBLINEAR_WINDOW = _env_int("AGILLM_SUBLINEAR_WINDOW", 256) DEFAULT_SUBLINEAR_STRIDE = _env_int("AGILLM_SUBLINEAR_STRIDE", 64) DEFAULT_SUBLINEAR_MAX_ANCHORS = _env_int("AGILLM_SUBLINEAR_MAX_ANCHORS", 256) DEFAULT_SUBLINEAR_CHUNK = _env_int("AGILLM_SUBLINEAR_CHUNK", 128) DEFAULT_SUBLINEAR_SINKS = _env_int("AGILLM_SUBLINEAR_SINKS", 4) DEFAULT_SUBLINEAR_RECENT_ANCHORS = _env_int("AGILLM_SUBLINEAR_RECENT_ANCHORS", -1) # -1 = half of max anchors DEFAULT_SUBLINEAR_POOLED_LANDMARKS = bool(_env_int("AGILLM_SUBLINEAR_POOLED_LANDMARKS", 0)) DEFAULT_ANCHOR_MEMORY = bool(_env_int("AGILLM_ANCHOR_MEMORY", 0)) DEFAULT_ANCHOR_STRIDE = _env_int("AGILLM_ANCHOR_STRIDE", 256) DEFAULT_ANCHOR_MAX = _env_int("AGILLM_ANCHOR_MAX", 2048) DEFAULT_ANCHOR_POSITION = _env_int("AGILLM_ANCHOR_POSITION", -1) # -1 = stack middle DEFAULT_KV_BUFFER = bool(_env_int("AGILLM_KV_BUFFER", 0)) DEFAULT_MOE_FFN = bool(_env_int("AGILLM_MOE_FFN", 0)) DEFAULT_MOE_EXPERTS = _env_int("AGILLM_MOE_EXPERTS", 4) DEFAULT_MOE_TOP_K = _env_int("AGILLM_MOE_TOP_K", 1) DEFAULT_MOE_MLP_MULT = _env_int("AGILLM_MOE_MLP_MULT", 4) AGILLM4_TOKEN_PARAM_RATIO = 100.0 # ───────────────────────── UK Time Helper ───────────────────────── def get_uk_time() -> str: utc_now = datetime.now(timezone.utc) year = utc_now.year march_last = datetime(year, 3, 31, 1, 0, tzinfo=timezone.utc) while march_last.weekday() != 6: march_last = march_last.replace(day=march_last.day - 1) oct_last = datetime(year, 10, 31, 1, 0, tzinfo=timezone.utc) while oct_last.weekday() != 6: oct_last = oct_last.replace(day=oct_last.day - 1) if march_last <= utc_now < oct_last: uk_offset = 1 tz_name = "BST" else: uk_offset = 0 tz_name = "GMT" from datetime import timedelta uk_time = utc_now + timedelta(hours=uk_offset) return uk_time.strftime(f'%Y-%m-%d %H:%M:%S {tz_name}') # ───────────────────────── Utilities ───────────────────────── def rng_state(): if DEV.type == "cuda": try: return torch.cuda.get_rng_state(DEV) except TypeError: return torch.cuda.get_rng_state() return torch.get_rng_state() def _is_probably_ckpt(path: pathlib.Path) -> bool: try: return path.is_file() and path.suffix == ".pt" and not path.name.endswith(".pt.tmp") and path.stat().st_size > (1<<20) except Exception: return False def _resolve_ckpt(path: pathlib.Path) -> pathlib.Path | None: try: if path.is_dir(): cands = sorted([p for p in path.glob("*.pt") if _is_probably_ckpt(p)], key=lambda p: p.stat().st_mtime, reverse=True) return cands[0] if cands else None if path.suffix == ".tmp": solid = path.with_suffix("") return solid if _is_probably_ckpt(solid) else _resolve_ckpt(path.parent) return path if _is_probably_ckpt(path) else _resolve_ckpt(path.parent) except Exception: return None def _try_load(path: pathlib.Path, map_location="cpu"): try: return _agillm43_load_pt(path, map_location=map_location, weights_only=False) except Exception as e: print(f"[ckpt-skip] {path} not usable: {e}") return None def _prune_checkpoints(save_dir: pathlib.Path, phase_name: str, max_ckpts: int): if max_ckpts is None or max_ckpts <= 0: return try: pattern = f"{phase_name}_step*.pt" pinned = _pinned_basenames(save_dir) if '_pinned_basenames' in globals() else set() ckpts = sorted( [p for p in save_dir.glob(pattern) if _is_probably_ckpt(p) and not p.name.endswith('.resume_delta.pt') and p.name not in pinned], key=lambda p: p.stat().st_mtime ) excess = len(ckpts) - max_ckpts if excess > 0: for p in ckpts[:excess]: try: for sidecar in ( p, p.with_name(p.name + ".tokenizer.json"), p.with_suffix(".provenance.json"), ): try: if sidecar.exists(): sidecar.unlink() except Exception: pass print(f" [prune] deleted old {p.name} (+ sidecars)") except Exception: pass except Exception as e: print(f"[ckpt-prune] error: {e}") def print_expansion_info(cfg: dict, tie_weights: bool = False, plain: bool = False): d_k = cfg["d"] // cfg["heads"] rank = cfg["rank"] ratio = rank / d_k regime = "COMPRESSION" if ratio < 1 else ("IDENTITY" if ratio == 1 else "EXPANSION") tie_str = "YES" if tie_weights else "NO" if plain: print("[attention_config]") print(f"d_model={cfg['d']} heads={cfg['heads']} d_k={d_k}") print(f"layers={cfg['layers']} tie_weights={tie_str}") print(f"rank={rank} ratio={ratio:.1f}x regime={regime}") return print(f"┌─────────────────────────────────────────┐") print(f"│ TUNEABLE ATTENTION CONFIG │") print(f"├─────────────────────────────────────────┤") print(f"│ d_model: {cfg['d']:4d} heads: {cfg['heads']:2d} d_k: {d_k:3d} │") print(f"│ layers: {cfg['layers']:4d} tie_weights: {tie_str:3s} │") print(f"│ rank: {rank:4d} ratio: {ratio:.1f}x [{regime:11s}] │") print(f"└─────────────────────────────────────────┘") # ───────────────────────── AMP helper ───────────────────────── try: from torch.amp import autocast as _ac, GradScaler except ImportError: from torch.cuda.amp import autocast as _ac, GradScaler def _auto_amp_dtype(): if DEV.type == "cuda": try: if torch.cuda.is_bf16_supported(): return torch.bfloat16 return torch.float16 except Exception: return torch.float16 return torch.float32 def amp(enabled: bool): if not enabled or DEV.type != "cuda": return nullcontext() dtype = _auto_amp_dtype() if hasattr(torch, "amp") and hasattr(torch.amp, "autocast"): try: return torch.amp.autocast("cuda", dtype=dtype) except TypeError: try: return torch.amp.autocast(device_type="cuda", dtype=dtype) except TypeError: pass return torch.cuda.amp.autocast(dtype=dtype) def _needs_grad_scaler() -> bool: return bool(DEV.type == "cuda" and _auto_amp_dtype() == torch.float16) # ───────────────────────── Chat & Data Stream ───────────────────────── def _coerce_role(r: str) -> str: r = (r or "").lower() if r in {"user", "human", "customer"}: return "user" if r in {"assistant", "gpt", "bot"}: return "assistant" if r in {"system", "context"}: return "system" return r or "user" def _chat_content(m: dict) -> str: content = m.get("content", m.get("text", m.get("value", ""))) return content if isinstance(content, str) else "" def _chat_role(m: dict) -> str: return _coerce_role(m.get("role", m.get("from", m.get("speaker", "")))) def _fallback_chat_template(messages: list[dict], add_generation_prompt: bool) -> str: parts = [] for m in messages: role = _chat_role(m) content = _chat_content(m).strip() if not content: continue if role == "system": parts.append(f"System: {content}") elif role == "assistant": parts.append(f"Assistant: {content}") else: parts.append(f"User: {content}") if add_generation_prompt and (not parts or not parts[-1].startswith("Assistant:")): parts.append("Assistant:") return "\n".join(parts) def _render_chat_text_from_ex(ex: dict, messages_key: str, add_generation_prompt: bool) -> Optional[str]: msgs = ex.get(messages_key) if msgs is None: for alt in ("conversations", "dialog", "turns"): if isinstance(ex.get(alt), list): msgs = ex[alt]; break if isinstance(msgs, list) and msgs and isinstance(msgs[0], dict): norm = [] for m in msgs: content = _chat_content(m) if not isinstance(content, str) or not content: continue norm.append({"role": _chat_role(m), "content": content}) if not norm: return None try: return tok.apply_chat_template(norm, tokenize=False, add_generation_prompt=add_generation_prompt) except Exception: return _fallback_chat_template(norm, add_generation_prompt) for a, b in (("prompt", "response"), ("instruction", "output"), ("question", "answer")): if isinstance(ex.get(a), str) and isinstance(ex.get(b), str): return f"User: {ex[a]}\nAssistant: {ex[b]}" return None def _parse_dataset_ref(ds_name: str): split = "train" ref = ds_name if "@" in ref: ref, split = ref.rsplit("@", 1) split = split or "train" if ":" in ref: base, config = ref.split(":", 1) else: base, config = ref, None return base, config, split _DATASET_COMPAT_RULES = [ # Keep dataset-specific scars in one place. These are name-pattern fixes for # repos whose HF auto-builder, schema, or default config is known to bite # streaming pretraining. (re.compile(r"^EleutherAI/proof-pile-2$"), {"loader": "proof_pile_direct"}), (re.compile(r"^allenai/dolma$"), {"loader": "dolma_url_manifest", "default_config": "v1_6-sample"}), (re.compile(r"^tiiuae/falcon-refinedweb$"), {"text_fields": ("content", "text")}), (re.compile(r"^HuggingFaceFW/fineweb-edu$"), {"default_config": "sample-10BT"}), (re.compile(r"^Salesforce/wikitext$"), {"default_config": "wikitext-103-raw-v1"}), ] def _dataset_compat(base: str) -> dict: for pattern, rule in _DATASET_COMPAT_RULES: try: if pattern.match(base or ""): return rule except Exception: continue return {} def _dataset_text_fields_for_source(ds_name: str, preferred: str = "text") -> List[str]: base, _config, _split = _parse_dataset_ref(ds_name) compat = _dataset_compat(base) fields = [] def add(field): if isinstance(field, str) and field and field not in fields: fields.append(field) add(preferred) for field in compat.get("text_fields", ()): add(field) for field in ("text", "content", "raw_content", "document", "body"): add(field) return fields _PROOF_PILE_REPO = "EleutherAI/proof-pile-2" _PROOF_PILE_URL_BASE = f"https://huggingface.co/datasets/{_PROOF_PILE_REPO}/resolve/main/" _PROOF_PILE_FILE_CACHE = {} _DOLMA_REPO = "allenai/dolma" _DOLMA_FILE_CACHE = {} def _dolma_data_files(config: Optional[str], split: str) -> List[str]: # The Dolma HF builder can hit UnicodeDecodeError by treating compressed # payload bytes as text. Its repo exposes URL manifests; feed those URLs # to the JSON builder directly instead. subset_ref = (config or os.environ.get("AGILLM_DOLMA_SUBSET", "") or "v1_6-sample").strip() split_ref = (split or "train").strip() or "train" cache_key = (subset_ref, split_ref) cached = _DOLMA_FILE_CACHE.get(cache_key) if cached: return cached if split_ref != "train": raise FileNotFoundError(f"{_DOLMA_REPO} manifest loader only supports train split, got {split_ref!r}") manifest = subset_ref if subset_ref.startswith("urls/") else f"urls/{subset_ref}.txt" try: from huggingface_hub import hf_hub_download manifest_path = hf_hub_download(_DOLMA_REPO, manifest, repo_type="dataset") urls = [line.strip() for line in Path(manifest_path).read_text().splitlines() if line.strip() and not line.startswith("#")] except Exception as exc: raise RuntimeError(f"could not resolve {_DOLMA_REPO} manifest {manifest}: {exc}") from exc if not urls: raise FileNotFoundError(f"empty {_DOLMA_REPO} manifest {manifest}") _DOLMA_FILE_CACHE[cache_key] = urls return urls def _proof_pile_data_files(config: Optional[str], split: str) -> List[str]: # The HF auto-builder for proof-pile-2 can try to UTF-8 decode compressed # .jsonl.zst bytes. Loading the repo's shards explicitly through the JSON # builder keeps this language source usable while preserving one logical # interleave source. subset_ref = (config or os.environ.get("AGILLM_PROOF_PILE_SUBSET", "") or "all").strip() split_ref = (split or "train").strip() or "train" cache_key = (subset_ref, split_ref) cached = _PROOF_PILE_FILE_CACHE.get(cache_key) if cached: return cached if subset_ref.lower() in {"", "all", "default", "full"}: subsets = ["algebraic-stack", "arxiv", "open-web-math"] else: subsets = [s.strip() for s in re.split(r"[+;]", subset_ref) if s.strip()] try: from huggingface_hub import list_repo_files repo_files = list_repo_files(_PROOF_PILE_REPO, repo_type="dataset") except Exception as exc: raise RuntimeError(f"could not list {_PROOF_PILE_REPO} shards: {exc}") from exc prefixes = tuple(f"{subset}/{split_ref}/" for subset in subsets) shard_paths = sorted( f for f in repo_files if f.endswith(".jsonl.zst") and f.startswith(prefixes) ) if not shard_paths: raise FileNotFoundError( f"no {_PROOF_PILE_REPO} .jsonl.zst shards for subset={subset_ref!r} split={split_ref!r}" ) urls = [_PROOF_PILE_URL_BASE + f for f in shard_paths] _PROOF_PILE_FILE_CACHE[cache_key] = urls return urls def _open_stream_one(ds_name: str, seed: int, streaming: bool = True): dc = DownloadConfig(max_retries=5, use_etag=True, resume_download=True) base, config, split = _parse_dataset_ref(ds_name) compat = _dataset_compat(base) if config is None and compat.get("default_config"): config = str(compat["default_config"]) print(f"[dataset-policy] {base} default_config={config}", flush=True) if not streaming: print(f"[download] Downloading {ds_name} (non-streaming)...") if base == "json": data_files = {"train": config} ds = load_dataset("json", data_files=data_files, split=split, streaming=streaming, download_config=dc) elif compat.get("loader") == "proof_pile_direct": urls = _proof_pile_data_files(config, split) data_files = {split: urls} subset_ref = config or os.environ.get("AGILLM_PROOF_PILE_SUBSET", "") or "all" print( f"[dataset-policy] proof-pile direct jsonl.zst loader subset={subset_ref} split={split} shards={len(urls)}", flush=True, ) ds = load_dataset("json", data_files=data_files, split=split, streaming=streaming, download_config=dc) elif compat.get("loader") == "dolma_url_manifest": urls = _dolma_data_files(config, split) data_files = {split: urls} subset_ref = config or os.environ.get("AGILLM_DOLMA_SUBSET", "") or "v1_6-sample" print( f"[dataset-policy] dolma direct json.gz loader subset={subset_ref} split={split} shards={len(urls)}", flush=True, ) ds = load_dataset("json", data_files=data_files, split=split, streaming=streaming, download_config=dc) else: ds = load_dataset(base, config, split=split, streaming=streaming, download_config=dc) if config else \ load_dataset(base, split=split, streaming=streaming, download_config=dc) if streaming: return iter(ds.shuffle(buffer_size=200, seed=seed)) # AGILLM-OOM-FIX 20260702: was 1000, OOM-killed at step 1 on 31GB RAM else: print(f"[download] Got {len(ds):,} examples. Shuffling...") ds = ds.shuffle(seed=seed) return iter(ds) def token_stream(ds_names: str, target: int, seed: int = 42, chat: bool = False, chat_messages_key: str = "messages", sft_add_generation_prompt: bool = False, dataset_field_text: str = "text", streaming: bool = True, use_hot_config: bool = True): if use_hot_config: ds_names = get_hot_datasets(ds_names) # HOT LOAD raw = [s.strip() for s in ds_names.split(",") if s.strip()] if not raw: return # Weighted interleave across sources, with an online quality router on top. # Base weights express policy; the router learns which sources yield bounded, # clean, useful examples instead of rewarding giant records for token volume. sources, weights = [], [] for s in raw: w = 1.0 head, sep, tail = s.rpartition("|") if sep: try: w = float(tail); s = head except ValueError: pass sources.append(s); weights.append(max(w, 0.0)) if sum(weights) <= 0: weights = [1.0] * len(sources) try: max_example_tokens = int(os.environ.get("AGILLM_MAX_EXAMPLE_TOKENS", "4096") or 0) except Exception: max_example_tokens = 4096 max_example_tokens = max(0, max_example_tokens) _rng = random.Random(seed) its = [None] * len(sources) emitted = 0 fail_counts = [0] * len(sources) disabled_until = [0.0] * len(sources) last_retry_log = [0.0] * len(sources) backoff_base = 2.0 max_cooldown = float(os.environ.get("AGILLM_STREAM_SOURCE_MAX_COOLDOWN_SEC", "300") or 300) fatal_cooldown = float(os.environ.get("AGILLM_STREAM_SOURCE_FATAL_COOLDOWN_SEC", "1800") or 1800) fatal_errors = {"DataFilesNotFoundError", "ArrowInvalid", "CastError", "FileNotFoundError"} router_enabled = str(os.environ.get("AGILLM_DATASET_NN_ROUTER", "1")).lower() not in {"0", "false", "off", "no"} router_state_path = Path(os.environ.get("AGILLM_DATASET_ROUTER_STATE", "/workspace/agillm_dataset_router_state.json")) router_explore = max(0.0, min(float(os.environ.get("AGILLM_DATASET_ROUTER_EXPLORE", "0.03") or 0.03), 0.50)) router_lr = max(0.0, min(float(os.environ.get("AGILLM_DATASET_ROUTER_LR", "0.03") or 0.03), 0.20)) router_min_score = max(0.01, min(float(os.environ.get("AGILLM_DATASET_ROUTER_MIN_SCORE", "0.05") or 0.05), 1.0)) router_sharpness = max(1.0, min(float(os.environ.get("AGILLM_DATASET_ROUTER_SHARPNESS", "3.0") or 3.0), 8.0)) router_log_sec = max(30.0, float(os.environ.get("AGILLM_DATASET_ROUTER_LOG_SEC", "300") or 300)) router_save_sec = max(10.0, float(os.environ.get("AGILLM_DATASET_ROUTER_SAVE_SEC", "60") or 60)) router_target_tokens = max(64.0, float(os.environ.get("AGILLM_DATASET_ROUTER_TARGET_TOKENS", str(max(512, min(max_example_tokens or 4096, 2048)))) or 2048)) router_min_quality = max(0.0, min(1.0, float(os.environ.get("AGILLM_DATASET_ROUTER_MIN_QUALITY", "0.45") or 0.45))) router_last_log = 0.0 router_last_save = 0.0 def _env_bool(name, default=False): return str(os.environ.get(name, "1" if default else "0")).strip().lower() not in {"", "0", "false", "off", "no"} def _env_float(name, default, lo=None, hi=None): try: val = float(os.environ.get(name, str(default)) or default) except Exception: val = float(default) if lo is not None: val = max(float(lo), val) if hi is not None: val = min(float(hi), val) return val agent_enabled = _env_bool("AGILLM_DATASET_AGENT_ROUTER", False) agent_timeout = _env_float("AGILLM_DATASET_AGENT_TIMEOUT_SEC", 8.0, 1.0, 60.0) agent_min_interval = _env_float("AGILLM_DATASET_AGENT_MIN_INTERVAL_SEC", 600.0, 30.0, 86400.0) agent_source_interval = _env_float("AGILLM_DATASET_AGENT_SOURCE_INTERVAL_SEC", 900.0, 30.0, 86400.0) agent_fail_threshold = int(_env_float("AGILLM_DATASET_AGENT_FAILS", 2.0, 1.0, 50.0)) agent_min_pulls = int(_env_float("AGILLM_DATASET_AGENT_MIN_PULLS", 4.0, 1.0, 1000.0)) agent_err_threshold = _env_float("AGILLM_DATASET_AGENT_ERR_EMA", 0.18, 0.01, 1.0) agent_empty_threshold = _env_float("AGILLM_DATASET_AGENT_EMPTY_EMA", 0.20, 0.01, 1.0) agent_latency_threshold = _env_float("AGILLM_DATASET_AGENT_LATENCY_SEC", 20.0, 1.0, 600.0) agent_min_conf = _env_float("AGILLM_DATASET_AGENT_MIN_CONF", 0.25, 0.0, 1.0) agent_default_penalty = _env_float("AGILLM_DATASET_AGENT_PENALTY", 0.35, 0.01, 1.0) agent_default_cooldown = _env_float("AGILLM_DATASET_AGENT_COOLDOWN_SEC", 900.0, 30.0, 86400.0) agent_disable_sec = _env_float("AGILLM_DATASET_AGENT_DISABLE_SEC", 21600.0, 60.0, 604800.0) agent_last_call = 0.0 def _sigmoid(x): if x < -40.0: return 0.0 if x > 40.0: return 1.0 return 1.0 / (1.0 + math.exp(-x)) def _load_router_state(): default_weights = [-0.15, 0.85, 1.40, -2.00, -0.25, 0.90, -2.50, 2.40, -3.00, -2.80, -1.60, -0.80] default = { "schema": "agillm.dataset_router.v2", "updated_utc": "", "weights": list(default_weights), "sources": {}, "agent": {}, } try: if router_state_path.exists(): loaded = json.loads(router_state_path.read_text()) if isinstance(loaded, dict): default.update({k: loaded.get(k, default[k]) for k in default}) if not isinstance(default.get("sources"), dict): default["sources"] = {} if default.get("schema") != "agillm.dataset_router.v2": default["schema"] = "agillm.dataset_router.v2" default["weights"] = list(default_weights) if not isinstance(default.get("weights"), list) or len(default["weights"]) != len(default_weights): default["weights"] = list(default_weights) except Exception as exc: print(f"[dataset-router] warning: could not load {router_state_path}: {exc}", flush=True) return default router = _load_router_state() router.setdefault("agent", {}) try: agent_last_call = float(router["agent"].get("last_call", 0.0) or 0.0) except Exception: agent_last_call = 0.0 def _source_state(src): st = router.setdefault("sources", {}).setdefault(src, {}) st.setdefault("ok_ema", 0.55) st.setdefault("err_ema", 0.05) st.setdefault("lat_ema", 1.0) st.setdefault("tok_ema", 256.0) st.setdefault("token_fit_ema", 0.50) st.setdefault("quality_ema", 0.65) st.setdefault("replacement_ema", 0.0) st.setdefault("control_ema", 0.0) st.setdefault("repeat_ema", 0.0) st.setdefault("short_ema", 0.05) st.setdefault("empty_ema", 0.05) st.setdefault("pulls", 0) st.setdefault("tokens", 0) st.setdefault("errors", 0) st.setdefault("empty", 0) st.setdefault("last_ok", 0.0) st.setdefault("last_error", "") st.setdefault("last_score", 0.5) st.setdefault("last_quality", 0.65) st.setdefault("agent_score_mult", 1.0) st.setdefault("agent_penalty_until", 0.0) st.setdefault("agent_last_check", 0.0) st.setdefault("agent_last_action", "") st.setdefault("agent_last_reason", "") st.setdefault("agent_last_error", "") return st for src in sources: _source_state(src) source_text_fields = [_dataset_text_fields_for_source(src, dataset_field_text) for src in sources] def _router_features(i, now): total_w = max(sum(weights), 1e-9) base = max(weights[i], 0.0) / total_w st = _source_state(sources[i]) return [ 1.0, min(1.0, base * len(weights)), float(st.get("ok_ema", 0.55)), float(st.get("err_ema", 0.05)), min(1.0, float(st.get("lat_ema", 1.0)) / 15.0), float(st.get("token_fit_ema", 0.50)), float(st.get("empty_ema", 0.05)), float(st.get("quality_ema", 0.65)), float(st.get("replacement_ema", 0.0)), float(st.get("control_ema", 0.0)), float(st.get("repeat_ema", 0.0)), float(st.get("short_ema", 0.05)), ] def _router_score(i, now): if not router_enabled: return 1.0 ws = router.get("weights") or [] feats = _router_features(i, now) z = sum(float(w) * float(f) for w, f in zip(ws, feats)) score = max(router_min_score, min(1.0, _sigmoid(z))) st = _source_state(sources[i]) try: until = float(st.get("agent_penalty_until", 0.0) or 0.0) mult = max(0.01, min(2.0, float(st.get("agent_score_mult", 1.0) or 1.0))) except Exception: until, mult = 0.0, 1.0 if until > now: score = max(router_min_score, min(1.0, score * mult)) elif until or mult != 1.0: st["agent_score_mult"] = 1.0 st["agent_penalty_until"] = 0.0 st["last_score"] = score return score def _save_router_state(force=False): nonlocal router_last_save now = time.time() if not force and now - router_last_save < router_save_sec: return router_last_save = now try: router["updated_utc"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime(now)) tmp = router_state_path.with_suffix(router_state_path.suffix + f".{os.getpid()}.tmp") tmp.parent.mkdir(parents=True, exist_ok=True) tmp.write_text(json.dumps(router, indent=2, sort_keys=True) + "\n") tmp.replace(router_state_path) except Exception as exc: print(f"[dataset-router] warning: could not save {router_state_path}: {exc}", flush=True) def _agent_read_secret(env_names, paths): for name in env_names: val = os.environ.get(name, "") if val.strip(): return val.strip() for raw_path in paths: try: p = Path(raw_path).expanduser() if p.exists(): val = p.read_text(errors="ignore").strip() if val: return val except Exception: pass return "" def _agent_provider_key_model(): pref = str(os.environ.get("AGILLM_DATASET_AGENT_PROVIDER", "auto") or "auto").strip().lower() deepseek_key = _agent_read_secret( ("DEEPSEEK_API_KEY", "AGILLM_DEEPSEEK_API_KEY"), ( "/root/.config/agillm/deepseek_api_key", "/workspace/private/deepseek_api_key", "/workspace/agillm_private/deepseek_api_key", ), ) openrouter_key = _agent_read_secret( ("OPENROUTER_API_KEY", "AGILLM_OPENROUTER_API_KEY"), ( "/root/.config/agillm/openrouter_api_key", "/workspace/private/openrouter_api_key", "/workspace/agillm_private/openrouter_api_key", ), ) deepseek_model = os.environ.get("AGILLM_DATASET_AGENT_DEEPSEEK_MODEL", "deepseek-chat") openrouter_model = os.environ.get("AGILLM_DATASET_AGENT_OPENROUTER_MODEL", "deepseek/deepseek-chat-v3-0324") if pref == "deepseek": return "deepseek", deepseek_key, deepseek_model, "configured" if deepseek_key else "missing-key" if pref == "openrouter": return "openrouter", openrouter_key, openrouter_model, "configured" if openrouter_key else "missing-key" if deepseek_key: return "deepseek", deepseek_key, deepseek_model, "configured" if openrouter_key: return "openrouter", openrouter_key, openrouter_model, "configured" return "auto", "", "", "missing-key" def _agent_extract_json(text): text = str(text or "").strip() if not text: return {} try: obj = json.loads(text) return obj if isinstance(obj, dict) else {} except Exception: pass start, end = text.find("{"), text.rfind("}") if start >= 0 and end > start: try: obj = json.loads(text[start:end + 1]) return obj if isinstance(obj, dict) else {} except Exception: return {} return {} def _agent_call(provider, key, model, payload): import urllib.error import urllib.request if provider == "deepseek": url = "https://api.deepseek.com/chat/completions" headers = {"Authorization": "Bearer " + key, "Content-Type": "application/json"} elif provider == "openrouter": url = "https://openrouter.ai/api/v1/chat/completions" headers = { "Authorization": "Bearer " + key, "Content-Type": "application/json", "HTTP-Referer": "https://join.opentransformers.online", "X-Title": "AGILLM dataset router", } else: return False, "unknown_provider" system = ( "You are a dataset routing policy agent for an active neural-network training run. " "Return compact JSON only. You may advise rerouting, cooldown, penalizing, disabling, keeping, or recovering a dataset source. " "Never create, rewrite, summarize, or transform training samples. " "Allowed actions: keep, penalize, cooldown, disable, recover. " "Use score_multiplier between 0.01 and 2.0 and cooldown_sec as seconds." ) body = { "model": model, "messages": [ {"role": "system", "content": system}, {"role": "user", "content": json.dumps(payload, sort_keys=True)}, ], "temperature": 0, "max_tokens": 180, } data = json.dumps(body).encode("utf-8") req = urllib.request.Request(url, data=data, headers=headers, method="POST") try: with urllib.request.urlopen(req, timeout=agent_timeout) as resp: raw = resp.read(32768).decode("utf-8", errors="replace") parsed = json.loads(raw) content = (((parsed.get("choices") or [{}])[0].get("message") or {}).get("content") or "") if not content and isinstance(parsed.get("output"), str): content = parsed["output"] return True, content except urllib.error.HTTPError as exc: return False, f"HTTP{getattr(exc, 'code', 'error')}" except Exception as exc: return False, type(exc).__name__ def _agent_maybe_advise(i, event): nonlocal agent_last_call if not agent_enabled or i is None: return now = time.time() st = _source_state(sources[i]) pulls = int(st.get("pulls", 0)) errors = int(st.get("errors", 0)) if pulls < agent_min_pulls and errors < agent_fail_threshold: return bad_enough = ( fail_counts[i] >= agent_fail_threshold or errors >= agent_fail_threshold or float(st.get("err_ema", 0.0)) >= agent_err_threshold or float(st.get("empty_ema", 0.0)) >= agent_empty_threshold or float(st.get("lat_ema", 0.0)) >= agent_latency_threshold ) if not bad_enough: return if now - agent_last_call < agent_min_interval: return if now - float(st.get("agent_last_check", 0.0) or 0.0) < agent_source_interval: return provider, key, model, status = _agent_provider_key_model() if not key: router.setdefault("agent", {})["last_status"] = status st["agent_last_check"] = now st["agent_last_error"] = status _save_router_state(force=True) return st["agent_last_check"] = now router.setdefault("agent", {})["last_call"] = now router["agent"]["last_provider"] = provider router["agent"]["last_model"] = model agent_last_call = now payload = { "source_index": i, "source": sources[i], "event": str(event or "failure")[:120], "policy": "reroute/cooldown only; never generate or modify data", "stats": { "pulls": pulls, "errors": errors, "empty": int(st.get("empty", 0)), "fail_count": int(fail_counts[i]), "ok_ema": float(st.get("ok_ema", 0.0)), "err_ema": float(st.get("err_ema", 0.0)), "empty_ema": float(st.get("empty_ema", 0.0)), "lat_ema": float(st.get("lat_ema", 0.0)), "tok_ema": float(st.get("tok_ema", 0.0)), "token_fit_ema": float(st.get("token_fit_ema", 0.0)), "quality_ema": float(st.get("quality_ema", 0.0)), "replacement_ema": float(st.get("replacement_ema", 0.0)), "control_ema": float(st.get("control_ema", 0.0)), "repeat_ema": float(st.get("repeat_ema", 0.0)), "router_score": float(st.get("last_score", 0.5)), "disabled_for_sec": max(0.0, float(disabled_until[i]) - now), "agent_score_mult": float(st.get("agent_score_mult", 1.0) or 1.0), }, "return_schema": { "action": "keep|penalize|cooldown|disable|recover", "score_multiplier": 0.35, "cooldown_sec": 900, "confidence": 0.5, "reason": "short reason", }, } ok, content = _agent_call(provider, key, model, payload) if not ok: st["agent_last_error"] = str(content)[:120] print(f"[dataset-agent] provider={provider} model={model} src={i}:{sources[i][:42]} error={content}", flush=True) _save_router_state(force=True) return advice = _agent_extract_json(content) action = str(advice.get("action", "keep") or "keep").strip().lower() if action not in {"keep", "penalize", "cooldown", "disable", "recover"}: action = "keep" try: confidence = max(0.0, min(1.0, float(advice.get("confidence", 0.0) or 0.0))) except Exception: confidence = 0.0 if confidence < agent_min_conf: action = "keep" try: mult = max(0.01, min(2.0, float(advice.get("score_multiplier", agent_default_penalty) or agent_default_penalty))) except Exception: mult = agent_default_penalty try: cooldown_sec = max(0.0, float(advice.get("cooldown_sec", agent_default_cooldown) or agent_default_cooldown)) except Exception: cooldown_sec = agent_default_cooldown reason = str(advice.get("reason", "") or "")[:180] if action == "recover": st["agent_score_mult"] = 1.0 st["agent_penalty_until"] = 0.0 disabled_until[i] = 0.0 elif action == "penalize": st["agent_score_mult"] = min(float(st.get("agent_score_mult", 1.0) or 1.0), mult) st["agent_penalty_until"] = max(float(st.get("agent_penalty_until", 0.0) or 0.0), now + max(cooldown_sec, agent_default_cooldown)) elif action == "cooldown": st["agent_score_mult"] = min(float(st.get("agent_score_mult", 1.0) or 1.0), mult) until = now + max(cooldown_sec, agent_default_cooldown) st["agent_penalty_until"] = max(float(st.get("agent_penalty_until", 0.0) or 0.0), until) disabled_until[i] = max(disabled_until[i], until) elif action == "disable": st["agent_score_mult"] = min(float(st.get("agent_score_mult", 1.0) or 1.0), min(mult, agent_default_penalty)) until = now + max(cooldown_sec, agent_disable_sec) st["agent_penalty_until"] = max(float(st.get("agent_penalty_until", 0.0) or 0.0), until) disabled_until[i] = max(disabled_until[i], until) st["agent_last_action"] = action st["agent_last_reason"] = reason st["agent_last_error"] = "" router.setdefault("agent", {})["last_status"] = "ok" _save_router_state(force=True) print( f"[dataset-agent] provider={provider} model={model} src={i}:{sources[i][:42]} " f"event={str(event)[:40]} action={action} mult={mult:.2f} cooldown={cooldown_sec:.0f}s conf={confidence:.2f} reason={reason}", flush=True, ) def _score_text_sample(text, token_count): preview = str(text or "")[:65536] n = max(1, len(preview)) repl = preview.count("\ufffd") / n control = sum(1 for ch in preview if ord(ch) < 32 and ch not in "\n\r\t") / n long_runs = 0 run = 1 prev = "" for ch in preview: if ch == prev: run += 1 else: if run >= 12: long_runs += run prev = ch run = 1 if run >= 12: long_runs += run repeat = long_runs / n whitespace = sum(1 for ch in preview if ch.isspace()) / n alpha = sum(1 for ch in preview if ch.isalpha()) / n digit = sum(1 for ch in preview if ch.isdigit()) / n tok = max(0.0, float(token_count or 0.0)) token_fit = max(0.0, min(1.0, 1.0 - abs(tok - router_target_tokens) / max(router_target_tokens, 1.0))) short = 1.0 if tok < min(128.0, router_target_tokens * 0.25) else 0.0 quality = 1.0 quality -= min(0.55, repl * 18.0) quality -= min(0.40, control * 28.0) quality -= min(0.35, repeat * 7.0) if whitespace < 0.04 or whitespace > 0.55: quality -= 0.12 if alpha < 0.18 and digit > 0.35: quality -= 0.16 if tok < 32: quality -= 0.35 elif tok < 128: quality -= 0.12 quality = max(0.0, min(1.0, quality)) return quality, token_fit, repl, control, repeat, short def _router_update(i, label, feat, token_count=0, latency=0.0, err="", empty=False, quality=None, token_fit=None, replacement_rate=0.0, control_rate=0.0, repeat_rate=0.0, short=0.0): if i is None: return st = _source_state(sources[i]) try: label = max(0.0, min(1.0, float(label))) except Exception: label = 0.0 alpha = 0.04 q = float(st.get("quality_ema", 0.65) if quality is None else max(0.0, min(1.0, float(quality)))) fit = float(st.get("token_fit_ema", 0.50) if token_fit is None else max(0.0, min(1.0, float(token_fit)))) replacement_rate = max(0.0, min(1.0, float(replacement_rate or 0.0))) control_rate = max(0.0, min(1.0, float(control_rate or 0.0))) repeat_rate = max(0.0, min(1.0, float(repeat_rate or 0.0))) short = max(0.0, min(1.0, float(short or 0.0))) st["pulls"] = int(st.get("pulls", 0)) + 1 st["ok_ema"] = (1.0 - alpha) * float(st.get("ok_ema", 0.55)) + alpha * label st["err_ema"] = (1.0 - alpha) * float(st.get("err_ema", 0.05)) + alpha * (1.0 - label) st["lat_ema"] = (1.0 - alpha) * float(st.get("lat_ema", 1.0)) + alpha * max(float(latency or 0.0), 0.0) st["tok_ema"] = (1.0 - alpha) * float(st.get("tok_ema", 256.0)) + alpha * max(float(token_count or 0.0), 0.0) st["token_fit_ema"] = (1.0 - alpha) * float(st.get("token_fit_ema", 0.50)) + alpha * fit st["quality_ema"] = (1.0 - alpha) * float(st.get("quality_ema", 0.65)) + alpha * q st["replacement_ema"] = (1.0 - alpha) * float(st.get("replacement_ema", 0.0)) + alpha * replacement_rate st["control_ema"] = (1.0 - alpha) * float(st.get("control_ema", 0.0)) + alpha * control_rate st["repeat_ema"] = (1.0 - alpha) * float(st.get("repeat_ema", 0.0)) + alpha * repeat_rate st["short_ema"] = (1.0 - alpha) * float(st.get("short_ema", 0.05)) + alpha * short st["empty_ema"] = (1.0 - alpha) * float(st.get("empty_ema", 0.05)) + alpha * (1.0 if empty else 0.0) st["last_quality"] = q if label >= 0.5: st["tokens"] = int(st.get("tokens", 0)) + int(token_count or 0) st["last_ok"] = time.time() st["last_error"] = "" else: st["errors"] = int(st.get("errors", 0)) + 1 st["last_error"] = str(err or "bad_sample")[:120] if empty: st["empty"] = int(st.get("empty", 0)) + 1 if router_enabled and feat and router_lr > 0: pred = _sigmoid(sum(float(w) * float(f) for w, f in zip(router["weights"], feat))) grad = label - pred router["weights"] = [max(-8.0, min(8.0, float(w) + router_lr * grad * float(f))) for w, f in zip(router["weights"], feat)] _save_router_state(force=(label < 0.5 or int(st.get("pulls", 0)) <= 3 or (int(st.get("pulls", 0)) % 25 == 0))) def _choose_source(available, now): if not router_enabled or _rng.random() < router_explore: return _rng.choices(available, weights=[weights[i] for i in available])[0] eff = [] for i in available: score = _router_score(i, now) eff.append(max(1e-9, weights[i] * (score ** router_sharpness))) if sum(eff) <= 0: eff = [weights[i] for i in available] return _rng.choices(available, weights=eff)[0] agent_provider, agent_key, agent_model, agent_status = _agent_provider_key_model() if not agent_enabled: agent_desc = "off" elif agent_key: agent_desc = f"{agent_provider}:{agent_model}" else: agent_desc = f"{agent_provider}:missing-key" print( f"[dataset-router] nn={'on' if router_enabled else 'off'} explore={router_explore:.3f} " f"agent={agent_desc} state={router_state_path} sources={len(sources)}", flush=True, ) while emitted < target: now = time.time() available = [i for i, w in enumerate(weights) if w > 0.0 and disabled_until[i] <= now] if not available: next_ready = min(disabled_until) if disabled_until else now + 1.0 sleep_s = max(1.0, min(30.0, next_ready - now)) print(f"[stream-retry] all sources cooling down, sleeping {sleep_s:.1f}s", flush=True) time.sleep(sleep_s) continue if router_enabled and now - router_last_log >= router_log_sec: rows = [] for i in range(len(sources)): st = _source_state(sources[i]) rows.append((float(st.get("last_score", _router_score(i, now))), i, st)) rows.sort(reverse=True) msg = "; ".join( f"{i}:{sources[i][:36]} score={score:.2f} q={st.get('quality_ema', 0):.2f} fit={st.get('token_fit_ema', 0):.2f} ok={st.get('ok_ema', 0):.2f} err={st.get('err_ema', 0):.2f} tok={st.get('tok_ema', 0):.0f}" for score, i, st in rows[:5] ) print(f"[dataset-router] {msg}", flush=True) router_last_log = now src_idx = _choose_source(available, now) feat = _router_features(src_idx, now) t0 = time.perf_counter() try: if its[src_idx] is None: its[src_idx] = _open_stream_one(sources[src_idx], seed + src_idx, streaming=streaming) ex = next(its[src_idx]) text = None if isinstance(ex, dict): if chat: text = _render_chat_text_from_ex(ex, chat_messages_key, sft_add_generation_prompt) if text is None: for field in source_text_fields[src_idx]: if isinstance(ex.get(field), str): text = ex[field] break if not isinstance(text, str) or not text.strip(): _router_update(src_idx, 0, feat, latency=time.perf_counter() - t0, err="empty_or_missing_text", empty=True) _agent_maybe_advise(src_idx, "empty_or_missing_text") continue if fail_counts[src_idx]: print(f"[stream-recover] {sources[src_idx]} recovered after {fail_counts[src_idx]} failures", flush=True) fail_counts[src_idx] = 0 disabled_until[src_idx] = 0.0 max_example_chars = int(os.environ.get("AGILLM_MAX_EXAMPLE_CHARS", str(max(8192, (max_example_tokens or 4096) * 8))) or 0) if max_example_chars and len(text) > max_example_chars: span_chars = max(1, len(text) - max_example_chars + 1) start_chars = _rng.randrange(span_chars) text = text[start_chars:start_chars + max_example_chars] enc = tok.encode(text) if EOS is not None and (len(enc) == 0 or enc[-1] != EOS): enc = enc + [EOS] if max_example_tokens and len(enc) > max_example_tokens: span = max(1, len(enc) - max_example_tokens + 1) start = _rng.randrange(span) enc = enc[start:start + max_example_tokens] if not enc: _router_update(src_idx, 0, feat, latency=time.perf_counter() - t0, err="empty_tokens", empty=True) _agent_maybe_advise(src_idx, "empty_tokens") continue quality, token_fit, replacement_rate, control_rate, repeat_rate, short = _score_text_sample(text, len(enc)) label = quality if quality >= router_min_quality else max(0.0, quality * 0.5) _router_update(src_idx, label, feat, token_count=len(enc), latency=time.perf_counter() - t0, quality=quality, token_fit=token_fit, replacement_rate=replacement_rate, control_rate=control_rate, repeat_rate=repeat_rate, short=short) for t in enc: yield t emitted += 1 if emitted >= target: _save_router_state(force=True) return except StopIteration: its[src_idx] = None # exhausted: reopen on next pick (stream cycles) except Exception as e: its[src_idx] = None fail_counts[src_idx] += 1 err = type(e).__name__ _router_update(src_idx, 0, feat, latency=time.perf_counter() - t0, err=err) cooldown = min(max_cooldown, backoff_base ** min(fail_counts[src_idx], 8)) if err in fatal_errors: cooldown = max(cooldown, fatal_cooldown) disabled_until[src_idx] = time.time() + cooldown _agent_maybe_advise(src_idx, err) if time.time() - last_retry_log[src_idx] > 15.0 or fail_counts[src_idx] <= 2: print( f"[stream-retry] {sources[src_idx]} error: {err}, " f"cooling {cooldown:.1f}s failures={fail_counts[src_idx]}", flush=True, ) last_retry_log[src_idx] = time.time() # ───────────────────────── ALiBi ───────────────────────── def _alibi_slopes(n_heads: int): def pow2slopes(n): start = 2 ** (-2 ** -(math.log2(n) - 3)) ratio = start return [start * (ratio ** i) for i in range(n)] if math.log2(n_heads).is_integer(): vals = pow2slopes(n_heads) else: closest = 2 ** math.floor(math.log2(n_heads)) vals = pow2slopes(closest) extra = pow2slopes(2 * closest) vals += extra[0::2][: n_heads - closest] return torch.tensor(vals, device=DEV).view(1, n_heads, 1, 1) def alibi_bias(n_heads: int, n_tokens: int): i = torch.arange(n_tokens, device=DEV).view(1, 1, n_tokens, 1) j = torch.arange(n_tokens, device=DEV).view(1, 1, 1, n_tokens) dist = (j - i).clamp_min(0) return -_alibi_slopes(n_heads) * dist class StructuredAttentionMask: """Symbolic attention rules for sublinear attention. Dense masks are O(T^2). This object carries the rule so sublinear attention can apply it only to the gathered local/anchor candidate keys: O(T * candidates). """ __slots__ = ("kind", "q_len", "k_len", "query_base", "block") def __init__(self, kind: str, q_len: int, k_len: int = None, query_base: int = 0, block: int = 1): self.kind = (kind or "none").lower() self.q_len = int(q_len) self.k_len = int(k_len if k_len is not None else q_len) self.query_base = int(query_base) self.block = max(1, int(block)) def to_dense(self, device=None, dtype=torch.float32): device = device or DEV if self.kind in {"none", "nat", "bidirectional", "unrestricted"}: return None q_pos = torch.arange(self.query_base, self.query_base + self.q_len, device=device, dtype=torch.long).view(self.q_len, 1) k_pos = torch.arange(self.k_len, device=device, dtype=torch.long).view(1, self.k_len) if self.kind == "causal": allow = k_pos <= q_pos elif self.kind in {"sat", "block_causal", "block-causal"}: allow = (k_pos // self.block) <= (q_pos // self.block) else: raise ValueError(f"unknown structured attention mask kind: {self.kind}") zeros = torch.zeros((self.q_len, self.k_len), device=device, dtype=dtype) neg = torch.full_like(zeros, float("-inf")) return torch.where(allow, zeros, neg).unsqueeze(0).unsqueeze(0) def _is_structured_attention_mask(mask) -> bool: return isinstance(mask, StructuredAttentionMask) def use_structured_masks(args=None, backend: str = None) -> bool: backend = (backend or getattr(args, "attn_backend", "") or "").lower() return backend == "sublinear" and not bool(getattr(args, "no_structured_masks", False)) # ───────────────────────── Model components ───────────────────────── class KVBuffer: """Preallocated K/V cache for decode. Replaces torch.cat-based growth. Layout matches MHA-internal head-major shape [B, H, T, d_k]. Caller sizes once; each ``append`` writes ``length:length+n`` slots in place and grows ``length``. ``view()`` returns slices of the live region so attention sees only filled positions. """ __slots__ = ("k", "v", "length", "capacity") def __init__( self, batch: int, heads: int, capacity: int, d_k: int, device, dtype, ): self.k = torch.empty(batch, heads, capacity, d_k, device=device, dtype=dtype) self.v = torch.empty(batch, heads, capacity, d_k, device=device, dtype=dtype) self.length = 0 self.capacity = capacity def append(self, k_new: torch.Tensor, v_new: torch.Tensor): n = k_new.size(2) end = self.length + n if end > self.capacity: raise RuntimeError( f"KVBuffer overflow: length={self.length} + n={n} > capacity={self.capacity}" ) self.k[:, :, self.length:end].copy_(k_new) self.v[:, :, self.length:end].copy_(v_new) self.length = end def view(self): return self.k[:, :, :self.length], self.v[:, :, :self.length] class TuneableAttentionMHA(nn.Module): def __init__( self, d: int, h: int, r: int, use_relpos: bool = True, attn_backend: str = DEFAULT_ATTN_BACKEND, sublinear_window: int = DEFAULT_SUBLINEAR_WINDOW, sublinear_stride: int = DEFAULT_SUBLINEAR_STRIDE, sublinear_max_anchors: int = DEFAULT_SUBLINEAR_MAX_ANCHORS, sublinear_chunk: int = DEFAULT_SUBLINEAR_CHUNK, sublinear_sinks: int = DEFAULT_SUBLINEAR_SINKS, sublinear_recent_anchors: int = DEFAULT_SUBLINEAR_RECENT_ANCHORS, sublinear_pooled_landmarks: bool = DEFAULT_SUBLINEAR_POOLED_LANDMARKS, tie_kv: bool = False, ): super().__init__() assert d % h == 0 self.h, self.dk, self.r = h, d // h, r self.use_relpos = use_relpos self.attn_backend = (attn_backend or "manual").lower() self.sublinear_window = max(1, int(sublinear_window)) self.sublinear_stride = max(0, int(sublinear_stride)) self.sublinear_max_anchors = max(0, int(sublinear_max_anchors)) self.sublinear_chunk = max(1, int(sublinear_chunk)) self.sublinear_sinks = max(0, int(sublinear_sinks)) recent = int(sublinear_recent_anchors) if recent < 0: recent = self.sublinear_max_anchors // 2 self.sublinear_recent_anchors = min(max(0, recent), self.sublinear_max_anchors) self.sublinear_pooled_landmarks = bool(sublinear_pooled_landmarks) # Exact n1 harvest: one fused QKV projection is mathematically the same # as three independent bias-free Linear(d, d) projections with their # weights stacked along out_features. # Q-K=V (arXiv 2606.04032): tie Key & Value into one shared projection. # For r>dk, reshape_heads==reshape_v so k_new IS v_new (exact) -> clean 50% KV-cache cut # and -33% qkv params. Gated; default off preserves the 3*d checkpoint layout. self.tie_kv = bool(tie_kv) self.qkv = nn.Linear(d, (2 if self.tie_kv else 3) * d, bias=False) self.U = nn.Parameter(torch.randn(self.dk, r)) nn.init.orthogonal_(self.U) self.proj = nn.Linear(h * self.dk, d, bias=False) self.drop = nn.Dropout(0.1) # Exact n1 harvest: for expansion ranks, (q @ U) @ (k @ U).T is # q @ (U @ U.T) @ k.T. This keeps score/cache width at d_k with no # quality change. Inference caches the metric and training recomputes # it so gradients through U are unchanged. self._metric_cache: Optional[torch.Tensor] = None self._metric_cache_ver: int = -1 self._metric_cache_param_id: int = -1 self._metric_cache_data_ptr: int = -1 self._metric_cache_shape: Tuple[int, int] = (-1, -1) def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): qkv_key = prefix + "qkv.weight" if qkv_key not in state_dict: qk = prefix + "q.weight" kk = prefix + "k.weight" vk = prefix + "v.weight" if qk in state_dict and kk in state_dict and vk in state_dict: fused = _cat_legacy_weight_blocks([state_dict[qk], state_dict[kk], state_dict[vk]]) if fused is not None: state_dict[qkv_key] = fused state_dict.pop(qk) state_dict.pop(kk) state_dict.pop(vk) return super()._load_from_state_dict( state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs, ) def _proj_qk(self, x): B, N, _ = x.shape return (x.view(B, N, self.h, self.dk).transpose(1, 2) @ self.U) def _reshape_v(self, x): B, N, _ = x.shape return x.view(B, N, self.h, self.dk).transpose(1, 2) def _reshape_heads(self, x): B, N, _ = x.shape return x.view(B, N, self.h, self.dk).transpose(1, 2) def _get_metric(self) -> torch.Tensor: if torch.is_grad_enabled(): return self.U @ self.U.T cur_ver = self.U._version cur_param_id = id(self.U) cur_data_ptr = int(self.U.data_ptr()) cur_shape = tuple(self.U.shape) cache = self._metric_cache if ( cache is None or cache.dtype != self.U.dtype or cache.device != self.U.device or self._metric_cache_ver != cur_ver or self._metric_cache_param_id != cur_param_id or self._metric_cache_data_ptr != cur_data_ptr or self._metric_cache_shape != cur_shape ): cache = (self.U @ self.U.T).detach() self._metric_cache = cache self._metric_cache_ver = cur_ver self._metric_cache_param_id = cur_param_id self._metric_cache_data_ptr = cur_data_ptr self._metric_cache_shape = cur_shape return cache def train(self, mode: bool = True): if mode: self._metric_cache = None self._metric_cache_ver = -1 self._metric_cache_param_id = -1 self._metric_cache_data_ptr = -1 self._metric_cache_shape = (-1, -1) return super().train(mode) def _structured_valid(self, attn_mask, q_pos, idx): if not _is_structured_attention_mask(attn_mask): return None kind = attn_mask.kind if kind in {"none", "nat", "bidirectional", "unrestricted"}: return torch.ones_like(idx, dtype=torch.bool) if kind == "causal": return idx <= q_pos[:, None] if kind in {"sat", "block_causal", "block-causal"}: block = max(1, int(attn_mask.block)) return (idx // block) <= (q_pos[:, None] // block) raise ValueError(f"unknown structured attention mask kind: {kind}") def _sublinear_anchor_positions(self, k_len: int, device): anchor_start = self.sublinear_stride - 1 if self.sublinear_stride <= 0 or self.sublinear_max_anchors <= 0 or anchor_start >= k_len: anchors = torch.empty(0, device=device, dtype=torch.long) else: all_anchors = torch.arange(anchor_start, k_len, self.sublinear_stride, device=device, dtype=torch.long) if all_anchors.numel() <= self.sublinear_max_anchors: anchors = all_anchors else: recent_budget = min(self.sublinear_recent_anchors, self.sublinear_max_anchors) span_budget = max(0, self.sublinear_max_anchors - recent_budget) parts = [] if span_budget > 0: span_sel = torch.linspace(0, all_anchors.numel() - 1, span_budget, device=device).round().long().unique() parts.append(all_anchors[span_sel]) if recent_budget > 0: parts.append(all_anchors[-recent_budget:]) anchors = torch.cat(parts).unique() if parts else torch.empty(0, device=device, dtype=torch.long) if self.sublinear_sinks > 0 and k_len > 0: sinks = torch.arange(min(self.sublinear_sinks, k_len), device=device, dtype=torch.long) anchors = torch.cat([sinks, anchors]).unique() if anchors.numel() else sinks return anchors def _sublinear_attention(self, q, k, v, attn_mask=None, rel_bias_tokens=None): """Local-window + landmark attention: O(N * (window + N/stride)).""" bsz, heads, q_len, _ = q.shape k_len = k.size(2) device = q.device query_base = max(0, k_len - q_len) outputs = [] scale = 1.0 / math.sqrt(self.dk) slopes = None if self.use_relpos and rel_bias_tokens is not None: slopes = _alibi_slopes(self.h).to(device=device, dtype=torch.float32) anchors = self._sublinear_anchor_positions(k_len, device) anchor_k = anchor_v = None if anchors.numel() and self.sublinear_pooled_landmarks and self.sublinear_stride > 1: # Optional pooled landmarks: each global anchor summarizes its stride segment. # This is off by default because it adds cumsum work; enable after benchmarking. ends = anchors + 1 starts = (ends - self.sublinear_stride).clamp_min(0) zero_k = k.new_zeros(k.size(0), k.size(1), 1, k.size(3)) zero_v = v.new_zeros(v.size(0), v.size(1), 1, v.size(3)) prefix_k = torch.cat([zero_k, k.cumsum(dim=2)], dim=2) prefix_v = torch.cat([zero_v, v.cumsum(dim=2)], dim=2) denom = (ends - starts).to(dtype=k.dtype).view(1, 1, -1, 1).clamp_min(1) anchor_k = (prefix_k[:, :, ends, :] - prefix_k[:, :, starts, :]) / denom anchor_v = (prefix_v[:, :, ends, :] - prefix_v[:, :, starts, :]) / denom offsets = torch.arange( -self.sublinear_window, self.sublinear_window + 1, device=device, dtype=torch.long, ) for q_start in range(0, q_len, self.sublinear_chunk): q_end = min(q_len, q_start + self.sublinear_chunk) cur = q_end - q_start q_pos = torch.arange(query_base + q_start, query_base + q_end, device=device, dtype=torch.long) local_raw = q_pos[:, None] + offsets[None, :] local_valid = (local_raw >= 0) & (local_raw < k_len) local_idx = local_raw.clamp(0, max(0, k_len - 1)) k_local = k[:, :, local_idx, :] v_local = v[:, :, local_idx, :] if anchors.numel(): anchor_idx = anchors.view(1, -1).expand(cur, -1) local_lo = (q_pos - self.sublinear_window).clamp_min(0).view(-1, 1) local_hi = (q_pos + self.sublinear_window).clamp_max(max(0, k_len - 1)).view(-1, 1) # Drop anchor copies already present in the local window; duplicates bias softmax mass. anchor_valid = (anchor_idx < local_lo) | (anchor_idx > local_hi) idx = torch.cat([local_idx, anchor_idx], dim=1) valid = torch.cat([local_valid, anchor_valid], dim=1) if anchor_k is not None and anchor_v is not None: k_anchor = anchor_k.unsqueeze(2).expand(-1, -1, cur, -1, -1) v_anchor = anchor_v.unsqueeze(2).expand(-1, -1, cur, -1, -1) else: k_anchor = k[:, :, anchor_idx, :] v_anchor = v[:, :, anchor_idx, :] k_sel = torch.cat([k_local, k_anchor], dim=-2) v_sel = torch.cat([v_local, v_anchor], dim=-2) else: idx = local_idx valid = local_valid k_sel = k_local v_sel = v_local structured_valid = self._structured_valid(attn_mask, q_pos, idx) if structured_valid is not None: valid = valid & structured_valid scores = (q[:, :, q_start:q_end, :].unsqueeze(-2) * k_sel).sum(dim=-1) * scale if slopes is not None: dist = (q_pos.view(1, 1, cur, 1) - idx.view(1, 1, cur, -1)).abs().to(torch.float32) scores = scores + (-slopes * dist).to(scores.dtype) if torch.is_tensor(attn_mask) and attn_mask.size(-1) == k_len and attn_mask.size(-2) >= q_end: mask_q = attn_mask[..., q_start:q_end, :] gather_idx = idx.view(1, 1, cur, -1).expand(mask_q.size(0), mask_q.size(1), cur, idx.size(1)) scores = scores + torch.gather(mask_q, -1, gather_idx) scores = scores.masked_fill(~valid.view(1, 1, cur, -1), float("-inf")) weights = torch.softmax(scores.float(), dim=-1).to(v.dtype) outputs.append((weights.unsqueeze(-1) * v_sel).sum(dim=-2)) return torch.cat(outputs, dim=2) def forward(self, x, mask=None, rel_bias_tokens=None, kv_cache=None, use_cache=False): if self.tie_kv: q_lin, kv_lin = self.qkv(x).chunk(2, dim=-1) k_lin = v_lin = kv_lin else: q_lin, k_lin, v_lin = self.qkv(x).chunk(3, dim=-1) if self.r > self.dk: q = self._reshape_heads(q_lin) @ self._get_metric() k_new = self._reshape_heads(k_lin) v_new = k_new if self.tie_kv else self._reshape_v(v_lin) else: q = self._proj_qk(q_lin) k_new = self._proj_qk(k_lin) v_new = self._reshape_v(v_lin) if kv_cache is None: k, v = k_new, v_new elif isinstance(kv_cache, KVBuffer): if use_cache: kv_cache.append(k_new, v_new) k, v = kv_cache.view() else: k, v = k_new, v_new else: k_cached, v_cached = kv_cache if use_cache: k = torch.cat([k_cached, k_new], dim=2) v = torch.cat([v_cached, v_new], dim=2) else: k, v = k_new, v_new attn_mask = mask if self.attn_backend != "sublinear" and _is_structured_attention_mask(attn_mask): attn_mask = attn_mask.to_dense(device=q.device, dtype=q.dtype) if self.attn_backend != "sublinear" and self.use_relpos and rel_bias_tokens is not None: rel = alibi_bias(self.h, rel_bias_tokens)[:, :, -q.size(2):, :].to(device=q.device, dtype=q.dtype) attn_mask = rel if attn_mask is None else attn_mask + rel if self.attn_backend == "sdpa" and attn_mask is not None and attn_mask.dtype != torch.bool and attn_mask.dtype != q.dtype: attn_mask = attn_mask.to(dtype=q.dtype) if self.attn_backend == "sdpa": try: z = F.scaled_dot_product_attention( q, k, v, attn_mask=attn_mask, dropout_p=0.0, scale=1.0 / math.sqrt(self.dk), ) except TypeError: # Older torch lacks the scale kwarg. Rescale q so SDPA's default sqrt(r) # denominator matches the historical AGILLM sqrt(d_k) denominator. q_scaled = q * math.sqrt(q.size(-1) / self.dk) z = F.scaled_dot_product_attention(q_scaled, k, v, attn_mask=attn_mask, dropout_p=0.0) elif self.attn_backend == "sublinear": z = self._sublinear_attention(q, k, v, attn_mask=attn_mask, rel_bias_tokens=rel_bias_tokens) else: att = (q @ k.transpose(-1, -2)) / math.sqrt(self.dk) if attn_mask is not None: att = att + attn_mask z = att.softmax(-1).to(v.dtype) @ v z = z.transpose(1, 2).reshape(x.size(0), x.size(1), -1) out = self.drop(self.proj(z)) if not use_cache: return out new_kv = kv_cache if isinstance(kv_cache, KVBuffer) else (k, v) return out, new_kv class MoEFFN(nn.Module): def __init__(self, d: int, mlp_mult: int = 4, experts: int = 4, top_k: int = 1, shared_experts: int = 0, shared_mlp_mult: int = 0): super().__init__() self.d = int(d) self.mlp_mult = max(1, int(mlp_mult)) self.num_experts = max(1, int(experts)) self.top_k = min(max(1, int(top_k)), self.num_experts) hidden = self.mlp_mult * self.d self.router = nn.Linear(self.d, self.num_experts, bias=False) self.experts = nn.ModuleList([ nn.Sequential(nn.Linear(self.d, hidden), nn.ReLU(), nn.Linear(hidden, self.d)) for _ in range(self.num_experts) ]) # Shared experts (DeepSeek/ST-MoE style): always-on FFN added to the routed # output, giving every token a consistent fallback representation -> lower # routing variance, smoother optimization. Output layer is ZERO-INITIALISED so # the shared path is a no-op at step 0, making it mergeable into an existing # checkpoint without disruption (it then learns to contribute). self.num_shared = max(0, int(shared_experts)) if self.num_shared > 0: shidden = max(1, int(shared_mlp_mult) or self.mlp_mult) * self.d self.shared = nn.ModuleList([ nn.Sequential(nn.Linear(self.d, shidden), nn.ReLU(), nn.Linear(shidden, self.d)) for _ in range(self.num_shared) ]) for blk in self.shared: nn.init.zeros_(blk[2].weight); nn.init.zeros_(blk[2].bias) else: self.shared = None # Detached FFN input stashed each training forward; the router aux loss is # recomputed OUTSIDE the gradient-checkpoint boundary by _collect_moe_aux(). self.last_router_input = None # Inference-only expert streaming: block-stream can keep only router/shared # paths resident and page selected routed experts on demand. self.expert_stream = False self.expert_stream_empty_cache = True self.expert_stream_stats = {"loads": 0, "tokens": 0} def set_expert_stream(self, enabled: bool, empty_cache: bool = True): self.expert_stream = bool(enabled) self.expert_stream_empty_cache = bool(empty_cache) return self def _run_expert(self, expert, rows): if self.expert_stream and torch.is_tensor(rows) and rows.is_cuda: expert.to(rows.device) try: out = expert(rows) finally: expert.to("cpu") self.expert_stream_stats["loads"] = int(self.expert_stream_stats.get("loads", 0)) + 1 self.expert_stream_stats["tokens"] = int(self.expert_stream_stats.get("tokens", 0)) + int(rows.size(0)) if self.expert_stream_empty_cache and torch.cuda.is_available(): torch.cuda.empty_cache() return out return expert(rows) def _shared_out(self, flat): if self.shared is None: return 0.0 s = self.shared[0](flat) for blk in self.shared[1:]: s = s + blk(flat) return s def forward(self, x): orig_shape = x.shape flat = x.reshape(-1, orig_shape[-1]) if self.training: # Stash the detached input (no autograd graph) so the load-balance loss # can be recomputed after the block forward. Computing it here would run # without grad (checkpoint's no-grad first pass) or pin block activations # across the checkpoint boundary and blow up VRAM. self.last_router_input = flat.detach() router_in = flat.to(self.router.weight.dtype) if flat.dtype != self.router.weight.dtype else flat scores = self.router(router_in).float() if self.top_k == 1: probs = scores.softmax(dim=-1) chosen = probs.argmax(dim=-1) out = torch.zeros_like(flat) for expert_id, expert in enumerate(self.experts): mask = chosen == expert_id if not bool(mask.any()): continue gate = probs[mask, expert_id].to(flat.dtype).clamp_min(1e-6) # Keep the forward value equal to the selected expert while # sending a straight-through gradient into the top-1 router. gate_st = (gate / gate.detach()).unsqueeze(-1) out[mask] = self._run_expert(expert, flat[mask]) * gate_st if self.shared is not None: out = out + self._shared_out(flat) return out.reshape(orig_shape) vals, idx = torch.topk(scores, k=self.top_k, dim=-1) weights = vals.softmax(dim=-1).to(flat.dtype) out = torch.zeros_like(flat) for rank in range(self.top_k): chosen = idx[:, rank] weight = weights[:, rank].unsqueeze(-1) for expert_id, expert in enumerate(self.experts): rows = (chosen == expert_id).nonzero(as_tuple=False).flatten() if rows.numel() == 0: continue out.index_add_(0, rows, self._run_expert(expert, flat.index_select(0, rows)) * weight.index_select(0, rows)) if self.shared is not None: out = out + self._shared_out(flat) return out.reshape(orig_shape) def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs): legacy = { "0.weight": "0.weight", "0.bias": "0.bias", "2.weight": "2.weight", "2.bias": "2.bias", } seeded = False for expert_idx, expert in enumerate(self.experts): expert_state = expert.state_dict() for legacy_suffix, expert_suffix in legacy.items(): src_key = prefix + legacy_suffix dst_key = prefix + f"experts.{expert_idx}." + expert_suffix src = state_dict.get(src_key) tgt = expert_state.get(expert_suffix) if dst_key not in state_dict and torch.is_tensor(src) and torch.is_tensor(tgt) and tuple(src.shape) == tuple(tgt.shape): state_dict[dst_key] = src seeded = True if seeded and prefix + "router.weight" not in state_dict: state_dict[prefix + "router.weight"] = self.router.weight.detach().clone() if seeded: for suffix in legacy: state_dict.pop(prefix + suffix, None) return super()._load_from_state_dict( state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs, ) def _collect_moe_aux(model, aux_coef=0.0, z_coef=0.0): """Sum and clear the MoE load-balance / router-z losses. Recomputes the router on the detached FFN input stashed during the forward, so it works with gradient checkpointing (router logits are available WITH grad here, outside the checkpointed region) and pins no block activations (the input is detached, so only router.weight receives gradient). Returns a scalar tensor to add to the loss before backward(), or 0.0 when disabled / nothing stashed. Verified on a 4090 (28L/d1280, AMP+grad_checkpoint): peak VRAM delta ~1MB. """ total = None for m in model.modules(): if isinstance(m, MoEFFN): inp = m.last_router_input m.last_router_input = None if inp is None or (aux_coef <= 0 and z_coef <= 0): continue router_in = inp.to(m.router.weight.dtype) if inp.dtype != m.router.weight.dtype else inp scores = m.router(router_in).float() probs = scores.softmax(dim=-1) importance = probs.mean(dim=0) top1 = probs.argmax(dim=-1) load = torch.bincount(top1, minlength=m.num_experts).to(importance.dtype) / max(1, top1.numel()) if aux_coef > 0: lb = aux_coef * m.num_experts * (load.detach() * importance).sum() total = lb if total is None else total + lb if z_coef > 0: zl = z_coef * (torch.logsumexp(scores, dim=-1) ** 2).mean() total = zl if total is None else total + zl return total if total is not None else 0.0 class Block(nn.Module): def __init__( self, d: int, h: int, r: int, attn_backend: str = DEFAULT_ATTN_BACKEND, sublinear_window: int = DEFAULT_SUBLINEAR_WINDOW, sublinear_stride: int = DEFAULT_SUBLINEAR_STRIDE, sublinear_max_anchors: int = DEFAULT_SUBLINEAR_MAX_ANCHORS, sublinear_chunk: int = DEFAULT_SUBLINEAR_CHUNK, sublinear_sinks: int = DEFAULT_SUBLINEAR_SINKS, sublinear_recent_anchors: int = DEFAULT_SUBLINEAR_RECENT_ANCHORS, sublinear_pooled_landmarks: bool = DEFAULT_SUBLINEAR_POOLED_LANDMARKS, moe_ffn: bool = DEFAULT_MOE_FFN, moe_experts: int = DEFAULT_MOE_EXPERTS, moe_top_k: int = DEFAULT_MOE_TOP_K, moe_mlp_mult: int = DEFAULT_MOE_MLP_MULT, moe_shared_experts: int = 0, moe_shared_mlp_mult: int = 0, tie_kv: bool = False, ): super().__init__() self.ln1, self.ln2 = nn.LayerNorm(d), nn.LayerNorm(d) self.mha = TuneableAttentionMHA( d, h, r, attn_backend=attn_backend, sublinear_window=sublinear_window, sublinear_stride=sublinear_stride, sublinear_max_anchors=sublinear_max_anchors, sublinear_chunk=sublinear_chunk, sublinear_sinks=sublinear_sinks, sublinear_recent_anchors=sublinear_recent_anchors, sublinear_pooled_landmarks=sublinear_pooled_landmarks, tie_kv=tie_kv, ) self.ff = ( MoEFFN(d, mlp_mult=moe_mlp_mult, experts=moe_experts, top_k=moe_top_k, shared_experts=moe_shared_experts, shared_mlp_mult=moe_shared_mlp_mult) if moe_ffn else nn.Sequential(nn.Linear(d, 4 * d), nn.ReLU(), nn.Linear(4 * d, d)) ) def forward(self, x, mask, kv=None, use_cache=False, total_seq_len=None): if use_cache: y, new_kv = self.mha(self.ln1(x), mask, rel_bias_tokens=total_seq_len, kv_cache=kv, use_cache=True) x = x + y + self.ff(self.ln2(x + y)) return x, new_kv else: n = x.size(1) x = x + self.mha(self.ln1(x), mask, rel_bias_tokens=n) return x + self.ff(self.ln2(x)) class Encoder(nn.Module): def __init__( self, cfg, tie_weights: bool = False, attn_backend: str = DEFAULT_ATTN_BACKEND, grad_checkpoint: bool = False, sublinear_window: int = DEFAULT_SUBLINEAR_WINDOW, sublinear_stride: int = DEFAULT_SUBLINEAR_STRIDE, sublinear_max_anchors: int = DEFAULT_SUBLINEAR_MAX_ANCHORS, sublinear_chunk: int = DEFAULT_SUBLINEAR_CHUNK, sublinear_sinks: int = DEFAULT_SUBLINEAR_SINKS, sublinear_recent_anchors: int = DEFAULT_SUBLINEAR_RECENT_ANCHORS, sublinear_pooled_landmarks: bool = DEFAULT_SUBLINEAR_POOLED_LANDMARKS, anchor_memory: bool = DEFAULT_ANCHOR_MEMORY, anchor_stride: int = DEFAULT_ANCHOR_STRIDE, anchor_max: int = DEFAULT_ANCHOR_MAX, anchor_position: int = DEFAULT_ANCHOR_POSITION, moe_ffn: Optional[bool] = None, moe_experts: Optional[int] = None, moe_top_k: Optional[int] = None, moe_mlp_mult: Optional[int] = None, moe_shared_experts: Optional[int] = None, moe_shared_mlp_mult: Optional[int] = None, tie_kv: Optional[bool] = None, ): super().__init__() d, l, h, r = cfg["d"], cfg["layers"], cfg["heads"], cfg["rank"] if tie_kv is None: tie_kv = bool(cfg.get("tie_kv", False)) if moe_ffn is None: moe_ffn = bool(cfg.get("moe_ffn", DEFAULT_MOE_FFN)) if moe_experts is None: moe_experts = int(cfg.get("moe_experts", DEFAULT_MOE_EXPERTS)) if moe_top_k is None: moe_top_k = int(cfg.get("moe_top_k", DEFAULT_MOE_TOP_K)) if moe_mlp_mult is None: moe_mlp_mult = int(cfg.get("moe_mlp_mult", DEFAULT_MOE_MLP_MULT)) moe_experts = max(1, int(moe_experts)) moe_top_k = min(max(1, int(moe_top_k)), moe_experts) moe_mlp_mult = max(1, int(moe_mlp_mult)) if moe_shared_experts is None: moe_shared_experts = int(cfg.get("moe_shared_experts", 0)) if moe_shared_mlp_mult is None: moe_shared_mlp_mult = int(cfg.get("moe_shared_mlp_mult", 0)) moe_shared_experts = max(0, int(moe_shared_experts)) self.emb = nn.Embedding(VOCAB, d) self.blocks = nn.ModuleList([ Block( d, h, r, attn_backend=attn_backend, sublinear_window=sublinear_window, sublinear_stride=sublinear_stride, sublinear_max_anchors=sublinear_max_anchors, sublinear_chunk=sublinear_chunk, sublinear_sinks=sublinear_sinks, sublinear_recent_anchors=sublinear_recent_anchors, sublinear_pooled_landmarks=sublinear_pooled_landmarks, moe_ffn=bool(moe_ffn), moe_experts=moe_experts, moe_top_k=moe_top_k, moe_mlp_mult=moe_mlp_mult, moe_shared_experts=moe_shared_experts, moe_shared_mlp_mult=moe_shared_mlp_mult, tie_kv=bool(tie_kv), ) for _ in range(l) ]) self.ln = nn.LayerNorm(d) self.tie_weights = tie_weights self.attn_backend = attn_backend self.grad_checkpoint = grad_checkpoint self.sublinear_window = sublinear_window self.sublinear_stride = sublinear_stride self.sublinear_max_anchors = sublinear_max_anchors self.sublinear_chunk = sublinear_chunk self.sublinear_sinks = sublinear_sinks self.sublinear_recent_anchors = sublinear_recent_anchors self.sublinear_pooled_landmarks = bool(sublinear_pooled_landmarks) self.moe_ffn = bool(moe_ffn) self.moe_experts = moe_experts self.moe_top_k = moe_top_k self.moe_mlp_mult = moe_mlp_mult self.moe_shared_experts = moe_shared_experts self.anchor_memory_enabled = bool(anchor_memory) self.anchor_stride = int(anchor_stride) self.anchor_max = int(anchor_max) n_layers = int(cfg["layers"]) if int(anchor_position) < 0: self.anchor_position = n_layers // 2 else: self.anchor_position = min(int(anchor_position), n_layers - 1) if self.anchor_memory_enabled: am_cfg = AnchorMemoryConfig( d_model=int(cfg["d"]), heads=int(cfg["heads"]), anchor_stride=self.anchor_stride, max_anchors=self.anchor_max, ) self.anchor = AnchorMemoryLayer(am_cfg) else: self.anchor = None def forward(self, ids, mask, kv_caches=None, use_cache=False, total_seq_len=None, inputs_embeds=None): # SwiReasoning: latent steps inject a continuous thought vector instead of a # discrete token embedding. inputs_embeds is [B, T, d]. x = self.emb(ids) if inputs_embeds is None else inputs_embeds if not use_cache: for i, blk in enumerate(self.blocks): if self.grad_checkpoint and self.training: x = torch_checkpoint.checkpoint(lambda y, block=blk: block(y, mask), x, use_reentrant=False) else: x = blk(x, mask) if self.anchor is not None and i == self.anchor_position: if self.grad_checkpoint and self.training: x, _ = torch_checkpoint.checkpoint(self.anchor, x, use_reentrant=False) else: x, _ = self.anchor(x) return self.ln(x) new_kvs = [] for i, blk in enumerate(self.blocks): kv = kv_caches[i] if kv_caches else None x, kv_out = blk(x, mask, kv, use_cache=True, total_seq_len=total_seq_len) new_kvs.append(kv_out) if self.anchor is not None and i == self.anchor_position: x, _ = self.anchor(x) return self.ln(x), new_kvs class ARHead(nn.Module): def __init__(self, d, tie_weights: bool = False, embedding_weight: nn.Parameter = None): super().__init__() self.tie_weights = tie_weights if tie_weights and embedding_weight is not None: self.proj = nn.Linear(d, VOCAB, bias=False) self.proj.weight = embedding_weight else: self.proj = nn.Linear(d, VOCAB) def forward(self, h): return self.proj(h) class NATHead(nn.Module): def __init__(self, d, tie_weights: bool = False, embedding_weight: nn.Parameter = None): super().__init__() self.tie_weights = tie_weights if tie_weights and embedding_weight is not None: self.proj = nn.Linear(d, VOCAB, bias=False) self.proj.weight = embedding_weight else: self.proj = nn.Linear(d, VOCAB) def forward(self, h): return self.proj(h) class SATHead(nn.Module): def __init__(self, d, mode="var", tie_weights: bool = False, embedding_weight: nn.Parameter = None, mlp: bool = False): super().__init__() self.tie_weights = tie_weights self.mlp = bool(mlp) if self.mlp: self.proj = nn.Sequential( nn.Linear(d, d), nn.GELU(), nn.Linear(d, VOCAB), ) elif tie_weights and embedding_weight is not None: self.proj = nn.Linear(d, VOCAB, bias=False) self.proj.weight = embedding_weight else: self.proj = nn.Linear(d, VOCAB) self.gate = nn.Linear(d, 2) if mode == "var" else None def forward(self, h_last): return self.proj(h_last), (self.gate(h_last[:, 0]) if self.gate else None) # ───────────────────────── Masks ───────────────────────── def causal_mask(n, structured: bool = False): if structured: return StructuredAttentionMask("causal", q_len=n, k_len=n, query_base=0) return torch.triu(torch.full((1, 1, n, n), float("-inf"), device=DEV), 1) def sat_mask(n, block=SAT_BLOCK, structured: bool = False): if structured: return StructuredAttentionMask("sat", q_len=n, k_len=n, query_base=0, block=block) idx = torch.arange(n, device=DEV) grp = idx.unsqueeze(0) // block allow = (grp.T == grp) | (grp.T > grp) return torch.where(allow, 0.0, float("-inf")).unsqueeze(0).unsqueeze(0) def sat_mask_cached(new_len: int, cached_len: int, block=SAT_BLOCK, structured: bool = False): total_len = cached_len + new_len if structured: return StructuredAttentionMask("sat", q_len=new_len, k_len=total_len, query_base=cached_len, block=block) q_idx = torch.arange(cached_len, total_len, device=DEV).unsqueeze(1) k_idx = torch.arange(total_len, device=DEV).unsqueeze(0) q_grp = q_idx // block k_grp = k_idx // block allow = q_grp >= k_grp return torch.where(allow, 0.0, float("-inf")).unsqueeze(0).unsqueeze(0) # ───────────────────────── Checkpoint helpers ───────────────────────── # ───────────────────────── Delta Checkpoints (weight-only, async) ───────────────────────── _delta_lock = threading.Lock() _delta_thread: Optional[threading.Thread] = None def _sha256_file(path: pathlib.Path) -> str: """Compute SHA256 of a file for integrity verification.""" h = hashlib.sha256() with open(path, "rb") as f: for chunk in iter(lambda: f.read(1 << 20), b""): h.update(chunk) return h.hexdigest() _AGILLM43_TENSOR_CODEC_MAGIC = "__agillm43_tensor_state_codec__" _AGILLM43_PAYLOAD_CODEC_MAGIC = "__agillm43_payload_codec__" _AGILLM43_TENSOR_CODEC_VERSION = "agillm43_tensor_state_v3_rowq8c" def _agillm43_dtype_name(dtype) -> str: return str(dtype).replace("torch.", "") def _agillm43_dtype_from_name(name: str): return getattr(torch, str(name).replace("torch.", "")) def _agillm43_zstd_compress(data: bytes, level: int = 1) -> bytes: try: import zstandard as zstd return zstd.ZstdCompressor(level=int(level)).compress(data) except Exception: import zlib return b"ZLIB" + zlib.compress(data, max(1, min(9, int(level)))) def _agillm43_payload_bytes(data) -> bytes: if torch.is_tensor(data): return data.detach().cpu().contiguous().numpy().tobytes() return bytes(data) def _agillm43_byte_tensor(data: bytes) -> torch.Tensor: import warnings with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) return torch.frombuffer(memoryview(data), dtype=torch.uint8).clone() def _agillm43_zstd_decompress(data: bytes) -> bytes: data = _agillm43_payload_bytes(data) if data.startswith(b"ZLIB"): import zlib return zlib.decompress(data[4:]) import zstandard as zstd return zstd.ZstdDecompressor().decompress(data) def _agillm43_tensor_bytes(t: torch.Tensor) -> bytes: tc = t.detach().cpu().contiguous() return tc.view(torch.uint8).numpy().tobytes() def _agillm43_tensor_from_bytes(raw: bytes, dtype_name: str, shape): dtype = _agillm43_dtype_from_name(dtype_name) import warnings with warnings.catch_warnings(): warnings.simplefilter("ignore", UserWarning) return torch.frombuffer(memoryview(raw), dtype=dtype).clone().reshape(tuple(int(x) for x in shape)) def _agillm43_zstd_level_from_codec(codec: str, default: int = 1) -> int: text = str(codec or "").strip().lower() level = int(default or 1) try: import re as _re m = _re.search(r"zstd(?:[-_]?level)?[-_]?([0-9]{1,2})", text) if m: level = int(m.group(1)) else: env_level = os.environ.get("AGILLM43_ZSTD_LEVEL") if env_level: level = int(env_level) except Exception: level = int(default or 1) return max(1, min(22, int(level))) def _agillm43_pack_aux_tensor(tensor: torch.Tensor, zstd_level: int = 1): raw = _agillm43_tensor_bytes(tensor) compressed = _agillm43_zstd_compress(raw, zstd_level) if len(compressed) < len(raw): return _agillm43_byte_tensor(compressed), "zstd", len(compressed) return _agillm43_byte_tensor(raw), "raw", len(raw) def _agillm43_unpack_aux_tensor(data, codec: str, dtype_name: str, shape): raw = _agillm43_zstd_decompress(data) if codec == "zstd" else _agillm43_payload_bytes(data) return _agillm43_tensor_from_bytes(raw, dtype_name, shape) def _agillm43_encode_tensor_state(state, mode: str = "adaptive-zstd", zstd_level: int = 1): """Problem-specific tensor codec for DBlock lease/update payloads. Modes: - off/raw/none: return the input unchanged. - zstd/lossless-zstd: lossless per-tensor zstd bytes. - fp16-zstd: cast floating tensors to fp16 before zstd. - int8-zstd/q8-zstd: symmetric per-tensor int8 + zstd. - rowq8-zstd/int8-rowwise-zstd: last-axis row-wise int8 + zstd, optimized for AGILLM4.3 projection/embedding matrices with outlier rows. - adaptive-zstd/auto: choose global int8 when it passes the AGILLM4.3 side-update error budget, otherwise row-wise int8 for matrix-like tensors when that passes, otherwise fp16. This is the production default for DBlock federation traffic because it is usually smaller and faster to decompress than fp16-zstd on AGILLM4.3 block weights. """ if not isinstance(state, dict): return state mode = str(mode or "off").strip().lower() if mode in {"", "off", "none", "raw", "false", "0"}: return state if mode in {"auto", "adaptive", "agillm-auto", "agillm43-auto"}: mode = "adaptive-zstd" q8_rms_max = float(os.environ.get("AGILLM43_CODEC_Q8_RMS_MAX", "0.0060") or 0.0060) q8_max_abs = float(os.environ.get("AGILLM43_CODEC_Q8_MAX_ABS", "0.020") or 0.020) adaptive_exact = str(os.environ.get("AGILLM43_CODEC_ADAPTIVE_EXACT", "0")).lower() in {"1", "true", "yes", "on"} rowq8_scale_dtype = str(os.environ.get("AGILLM43_CODEC_ROWQ8_SCALE_DTYPE", "float16") or "float16").lower() tensors = {} plain = {} source_total = 0 raw_total = 0 packed_total = 0 tensor_count = 0 pack_counts = defaultdict(int) def make_int8_candidate(src: torch.Tensor): f = src.float() maxabs = float(f.abs().max().item()) if f.numel() else 0.0 scale = max(maxabs / 127.0, 1.0e-12) q = torch.clamp(torch.round(f / scale), -127, 127).to(torch.int8).contiguous() if mode.startswith("adaptive"): # Fast bound/estimate for uniform symmetric quantization. Exact scans are # available for lab runs, but the federation hot path needs encode speed. rms = float(scale / math.sqrt(12.0)) maxerr = float(scale * 0.5) if adaptive_exact: err = q.float().mul(scale).sub(f) rms = float(err.pow(2).mean().sqrt().item()) if err.numel() else 0.0 maxerr = float(err.abs().max().item()) if err.numel() else 0.0 else: rms = 0.0 maxerr = 0.0 return q, scale, rms, maxerr def make_rowwise_int8_candidate(src: torch.Tensor): rowq8_min_cols = int(os.environ.get("AGILLM43_CODEC_ROWQ8_MIN_COLS", "64") or 64) if src.ndim < 2 or int(src.shape[-1]) < rowq8_min_cols or src.numel() == 0: return None f = src.float() cols = int(f.shape[-1]) rows = int(f.numel() // cols) flat = f.reshape(rows, cols) scales = flat.abs().amax(dim=1).div(127.0).clamp_min(1.0e-12).to(torch.float32).contiguous() q = torch.clamp(torch.round(flat / scales[:, None]), -127, 127).to(torch.int8).contiguous() recon_scales = scales.to(torch.float16).float() if rowq8_scale_dtype not in {"fp32", "float32"} else scales if mode.startswith("adaptive"): # Same hot-path bound as global int8, but per row. Include the tiny # fp16-scale storage error used by the production rowq8c payload. # Near-threshold candidates get one exact refinement pass; that keeps # adaptive from falling back to fp16 on AGILLM projection rows whose # conservative bound is pessimistic but actual error is inside budget. scale_err = recon_scales.sub(scales).abs() rms = float(torch.sqrt(torch.mean((recon_scales / math.sqrt(12.0)).pow(2))).item()) if scale_err.numel(): rms += float(torch.sqrt(torch.mean((scale_err * 64.0).pow(2))).item()) maxerr = float((recon_scales.abs().max() * 0.5 + scale_err.max() * 127.0).item()) if scale_err.numel() else float((recon_scales.abs().max() * 0.5).item()) refine_margin = float(os.environ.get("AGILLM43_CODEC_ROWQ8_REFINE_MARGIN", "1.50") or 1.50) if adaptive_exact or (rms <= q8_rms_max * refine_margin and maxerr <= q8_max_abs * refine_margin): err = q.float().mul(recon_scales[:, None]).sub(flat) rms = float(err.pow(2).mean().sqrt().item()) if err.numel() else 0.0 maxerr = float(err.abs().max().item()) if err.numel() else 0.0 else: rms = 0.0 maxerr = 0.0 return q.reshape(tuple(src.shape)), scales, rows, cols, rms, maxerr def pack_rowwise_scales(scales: torch.Tensor): if rowq8_scale_dtype in {"fp32", "float32"}: stored = scales.to(torch.float32).contiguous() else: stored = scales.to(torch.float16).contiguous() data, codec, nbytes = _agillm43_pack_aux_tensor(stored, zstd_level) return data, codec, _agillm43_dtype_name(stored.dtype), int(nbytes) for key, value in state.items(): if not torch.is_tensor(value): plain[key] = value continue src = value.detach().cpu().contiguous() source_total += int(src.numel() * src.element_size()) orig_dtype = _agillm43_dtype_name(src.dtype) pack_kind = "lossless" scale = None scales = None scales_data = None scales_codec = None scales_dtype = None rows = None cols = None scale_nbytes = 0 err_rms = None err_max_abs = None rowwise_mode = mode.startswith("rowq8") or mode.startswith("int8-row") or mode.startswith("q8-row") if src.is_floating_point() and rowwise_mode: rowq = make_rowwise_int8_candidate(src) if rowq is None: packed_tensor, scale, err_rms, err_max_abs = make_int8_candidate(src) pack_kind = "int8_symmetric" else: packed_tensor, scales, rows, cols, err_rms, err_max_abs = rowq scales = scales.to(torch.float32).contiguous() scales_data, scales_codec, scales_dtype, scale_nbytes = pack_rowwise_scales(scales) scales = None pack_kind = "int8_rowwise" elif src.is_floating_point() and (mode.startswith("int8") or mode.startswith("q8")): packed_tensor, scale, err_rms, err_max_abs = make_int8_candidate(src) pack_kind = "int8_symmetric" elif src.is_floating_point() and mode.startswith("adaptive") and src.dtype != torch.float16: q, q_scale, q_rms, q_max = make_int8_candidate(src) if q_rms <= q8_rms_max and q_max <= q8_max_abs: packed_tensor = q scale = q_scale err_rms = q_rms err_max_abs = q_max pack_kind = "int8_symmetric" else: rowq = make_rowwise_int8_candidate(src) if rowq is not None: rq, rq_scales, rq_rows, rq_cols, rq_rms, rq_max = rowq if rq_rms <= q8_rms_max and rq_max <= q8_max_abs: packed_tensor = rq scales = rq_scales.to(torch.float32).contiguous() scales_data, scales_codec, scales_dtype, scale_nbytes = pack_rowwise_scales(scales) scales = None rows = rq_rows cols = rq_cols err_rms = rq_rms err_max_abs = rq_max pack_kind = "int8_rowwise" else: packed_tensor = src.to(torch.float16).contiguous() pack_kind = "fp16" err = packed_tensor.float().sub(src.float()) err_rms = float(err.pow(2).mean().sqrt().item()) if err.numel() else 0.0 err_max_abs = float(err.abs().max().item()) if err.numel() else 0.0 else: packed_tensor = src.to(torch.float16).contiguous() pack_kind = "fp16" err = packed_tensor.float().sub(src.float()) err_rms = float(err.pow(2).mean().sqrt().item()) if err.numel() else 0.0 err_max_abs = float(err.abs().max().item()) if err.numel() else 0.0 elif src.is_floating_point() and mode.startswith("fp16") and src.dtype != torch.float16: packed_tensor = src.to(torch.float16).contiguous() pack_kind = "fp16" else: packed_tensor = src raw = _agillm43_tensor_bytes(packed_tensor) raw_total += len(raw) compressed = _agillm43_zstd_compress(raw, zstd_level) if len(compressed) < len(raw): data_bytes = compressed codec = "zstd" else: data_bytes = raw codec = "raw" packed_total += len(data_bytes) + scale_nbytes data = _agillm43_byte_tensor(data_bytes) pack_counts[pack_kind] += 1 item = { "shape": list(src.shape), "orig_dtype": orig_dtype, "packed_dtype": _agillm43_dtype_name(packed_tensor.dtype), "pack_kind": pack_kind, "scale": scale, "scales": scales, "scales_data": scales_data, "scales_codec": scales_codec, "scales_dtype": scales_dtype, "rows": rows, "cols": cols, "scale_nbytes": scale_nbytes, "codec": codec, "raw_nbytes": len(raw), "packed_nbytes": len(data_bytes), "data": data, } if err_rms is not None: item["err_rms"] = float(err_rms) if err_max_abs is not None: item["err_max_abs"] = float(err_max_abs) tensors[key] = item tensor_count += 1 return { _AGILLM43_TENSOR_CODEC_MAGIC: _AGILLM43_TENSOR_CODEC_VERSION, "mode": mode, "zstd_level": int(zstd_level), "q8_rms_max": float(q8_rms_max), "q8_max_abs": float(q8_max_abs), "adaptive_exact": bool(adaptive_exact), "tensor_count": tensor_count, "pack_counts": dict(pack_counts), "source_nbytes": int(source_total), "raw_nbytes": int(raw_total), "packed_nbytes": int(packed_total), "plain": plain, "tensors": tensors, } def _agillm43_decode_tensor_state(state): if not (isinstance(state, dict) and str(state.get(_AGILLM43_TENSOR_CODEC_MAGIC, "")).startswith("agillm43_tensor_state_v")): return state out = dict(state.get("plain") or {}) for key, item in (state.get("tensors") or {}).items(): data = item.get("data", b"") raw = _agillm43_zstd_decompress(data) if item.get("codec") == "zstd" else _agillm43_payload_bytes(data) packed = _agillm43_tensor_from_bytes(raw, item.get("packed_dtype"), item.get("shape")) if item.get("pack_kind") == "int8_symmetric": scale = float(item.get("scale") or 1.0) value = packed.float().mul_(scale) elif item.get("pack_kind") == "int8_rowwise": rows = int(item.get("rows") or 0) scales_data = item.get("scales_data") if scales_data is not None: rows = rows or int(item.get("scale_rows") or 0) scales = _agillm43_unpack_aux_tensor(scales_data, item.get("scales_codec"), item.get("scales_dtype") or "float16", [rows]).float() else: scales = item.get("scales") if not torch.is_tensor(scales): raise ValueError(f"rowwise tensor codec missing scales for {key}") scales = scales.float() rows = rows or int(scales.numel()) value = packed.float().reshape(rows, -1).mul_(scales.reshape(rows, 1)).reshape(tuple(int(x) for x in item.get("shape"))) else: value = packed out[key] = value return out def _agillm43_tensor_state_summary(state) -> dict: if isinstance(state, dict) and str(state.get(_AGILLM43_TENSOR_CODEC_MAGIC, "")).startswith("agillm43_tensor_state_v"): source = int(state.get("source_nbytes") or state.get("raw_nbytes") or 0) raw = int(state.get("raw_nbytes") or 0) packed = int(state.get("packed_nbytes") or 0) return { "codec": state.get(_AGILLM43_TENSOR_CODEC_MAGIC), "mode": state.get("mode"), "tensors": int(state.get("tensor_count") or 0), "pack_counts": dict(state.get("pack_counts") or {}), "source_nbytes": source, "raw_nbytes": raw, "packed_nbytes": packed, "ratio": (float(source) / float(packed)) if packed > 0 else 0.0, "post_transform_ratio": (float(raw) / float(packed)) if packed > 0 else 0.0, } return {"codec": "raw"} _AGILLM43_SHARDED_CODEC_MAGIC = "agillm43_block_sharded_torch_v1" def _agillm43_is_sharded_codec(codec: str) -> bool: mode = str(codec or "").strip().lower().replace("_", "-") return mode in {"sharded", "sharded-zstd", "block-sharded", "block-sharded-zstd", "blocks", "blocks-zstd"} def _agillm43_shard_dir_for_path(path) -> pathlib.Path: return pathlib.Path(str(path) + ".shards") def _agillm43_shard_safe_name(name: str) -> str: return re.sub(r"[^A-Za-z0-9_.-]+", "_", str(name)).strip("._") or "shard" def _agillm43_save_sharded_pt(obj, path, zstd_level: int = 1): """Save a checkpoint as a manifest plus independently loadable block/head shards.""" import shutil path = pathlib.Path(path) shard_dir = _agillm43_shard_dir_for_path(path) if shard_dir.exists(): shutil.rmtree(shard_dir) shard_dir.mkdir(parents=True, exist_ok=True) shard_codec = "zstd" if int(zstd_level or 0) > 0 else "off" entries = [] def _save_entry(name, payload, target, mode="set"): shard_name = _agillm43_shard_safe_name(name) + ".pt" shard_path = shard_dir / shard_name info = _agillm43_save_pt(payload, shard_path, codec=shard_codec, zstd_level=zstd_level) entries.append({ "name": str(name), "target": list(target), "mode": str(mode), "shard": shard_name, "codec": info.get("codec", "raw"), "nbytes": int(shard_path.stat().st_size), "sha256": _sha256_file(shard_path), }) if not isinstance(obj, dict): skeleton = {"payload": None} _save_entry("payload", obj, ["payload"], "set") else: skeleton = dict(obj) core_sd = skeleton.get("core") if isinstance(core_sd, dict): core_base = {} block_groups = {} for key, value in core_sd.items(): m = re.match(r"blocks\.(\d+)\.", str(key)) if m: block_groups.setdefault(int(m.group(1)), {})[key] = value else: core_base[key] = value skeleton["core"] = core_base for idx in sorted(block_groups): _save_entry(f"core_block_{idx:04d}", block_groups[idx], ["core"], "dict_update") for top_key in ("ar", "sat", "nat", "opt", "scaler"): if top_key in skeleton: payload = skeleton.pop(top_key) _save_entry(top_key, payload, [top_key], "set") manifest = { _AGILLM43_SHARDED_CODEC_MAGIC: _AGILLM43_SHARDED_CODEC_MAGIC, "format_version": 1, "codec": "block-sharded", "shard_codec": shard_codec, "shard_dir": shard_dir.name, "skeleton": skeleton, "entries": entries, } torch.save(manifest, path, _use_new_zipfile_serialization=False) return {"codec": "block-sharded", "shards": len(entries), "shard_dir": str(shard_dir), "shard_codec": shard_codec} def _agillm43_load_sharded_pt(path, manifest: dict, map_location="cpu", weights_only=False, skip_keys=None): path = pathlib.Path(path) skip = {str(k) for k in (skip_keys or set())} shard_dir = path.parent / str(manifest.get("shard_dir") or (path.name + ".shards")) state = dict(manifest.get("skeleton") or {}) for entry in manifest.get("entries") or []: target = list(entry.get("target") or []) if not target: continue top_key = str(target[0]) if top_key in skip: continue shard_path = shard_dir / str(entry.get("shard")) payload = _agillm43_load_pt(shard_path, map_location=map_location, weights_only=weights_only) if entry.get("mode") == "dict_update": base = state.setdefault(top_key, {}) if not isinstance(base, dict): raise ValueError(f"sharded checkpoint target {top_key!r} is not a dict") base.update(payload) else: state[top_key] = payload if set(state.keys()) == {"payload"}: return state["payload"] return state def _agillm43_finalize_pt_save(tmp, path, info: dict): """Atomically publish a .pt and, for sharded packages, its shard directory.""" tmp = pathlib.Path(tmp) path = pathlib.Path(path) if info.get("codec") != "block-sharded": tmp.replace(path) return import shutil tmp_shards = pathlib.Path(info.get("shard_dir") or _agillm43_shard_dir_for_path(tmp)) final_shards = _agillm43_shard_dir_for_path(path) if final_shards.exists(): shutil.rmtree(final_shards) tmp.replace(path) if tmp_shards.exists(): tmp_shards.replace(final_shards) manifest = torch.load(path, map_location="cpu", weights_only=False) if isinstance(manifest, dict) and manifest.get(_AGILLM43_SHARDED_CODEC_MAGIC) == _AGILLM43_SHARDED_CODEC_MAGIC: manifest["shard_dir"] = final_shards.name torch.save(manifest, path, _use_new_zipfile_serialization=False) info["shard_dir"] = str(final_shards) def _agillm43_save_pt(obj, path, codec: str = "off", zstd_level: int = 1): codec = str(codec or "off").strip().lower() zstd_level = _agillm43_zstd_level_from_codec(codec, zstd_level) if _agillm43_is_sharded_codec(codec): return _agillm43_save_sharded_pt(obj, path, zstd_level=zstd_level) if codec in {"", "off", "none", "raw", "false", "0"}: torch.save(obj, path, _use_new_zipfile_serialization=False) return {"codec": "raw"} import io buf = io.BytesIO() torch.save(obj, buf, _use_new_zipfile_serialization=False) raw = buf.getvalue() packed = _agillm43_zstd_compress(raw, zstd_level) if len(packed) >= len(raw): torch.save(obj, path, _use_new_zipfile_serialization=False) return {"codec": "raw", "raw_nbytes": len(raw), "packed_nbytes": len(packed), "zstd_level": int(zstd_level)} wrapper = { _AGILLM43_PAYLOAD_CODEC_MAGIC: "agillm43_zstd_torch_v1", "codec": "zstd", "zstd_level": int(zstd_level), "requested_codec": codec, "raw_nbytes": len(raw), "packed_nbytes": len(packed), "payload": _agillm43_byte_tensor(packed), } torch.save(wrapper, path, _use_new_zipfile_serialization=False) return {"codec": "zstd", "raw_nbytes": len(raw), "packed_nbytes": len(packed), "zstd_level": int(zstd_level), "ratio": float(len(raw)) / max(1.0, float(len(packed)))} def _agillm43_decompress_cache_enabled() -> bool: text = str(os.environ.get("AGILLM43_DECOMPRESS_CACHE", "1") or "1").strip().lower() return text not in {"0", "false", "no", "off", "disable", "disabled"} _AGILLM43_ZSTD_FRAME_MAGIC = b"\x28\xb5\x2f\xfd" def _agillm43_source_sha256_sidecar(path: pathlib.Path) -> str: candidates = [ path.with_suffix(path.suffix + ".sha256"), path.with_suffix(".sha256"), ] for sidecar in candidates: try: if not sidecar.exists(): continue text = sidecar.read_text(errors="ignore").strip() m = re.search(r"\b([0-9a-fA-F]{64})\b", text) if m: return m.group(1).lower() except Exception: pass return "" def _agillm43_decompress_cache_info(path) -> dict: path = pathlib.Path(path) st = path.stat() source_sha256 = _agillm43_source_sha256_sidecar(path) source_id = source_sha256 or f"{int(st.st_size):x}-{int(st.st_mtime_ns):x}" source_id = re.sub(r"[^0-9a-fA-F._-]", "_", source_id)[:20] cache_dir = path.parent / ".agillm43_decompressed_cache" cache_path = cache_dir / f"{path.name}.{source_id}.raw.pt" return { "source_path": str(path.resolve()), "source_name": path.name, "source_size": int(st.st_size), "source_mtime_ns": int(st.st_mtime_ns), "source_sha256": source_sha256, "source_id": source_id, "cache_dir": cache_dir, "cache_path": cache_path, "manifest_path": cache_path.with_suffix(cache_path.suffix + ".json"), } def _agillm43_manifest_matches(info: dict, manifest: dict) -> bool: if int(manifest.get("schema_version") or 0) != 1: return False if str(manifest.get("source_path") or "") != str(info.get("source_path") or ""): return False if int(manifest.get("source_size") or -1) != int(info.get("source_size") or -2): return False cached_sha = str(manifest.get("source_sha256") or "") source_sha = str(info.get("source_sha256") or "") if source_sha or cached_sha: return cached_sha == source_sha return int(manifest.get("source_mtime_ns") or -1) == int(info.get("source_mtime_ns") or -2) def _agillm43_find_decompressed_cache(path): if not _agillm43_decompress_cache_enabled(): return None try: info = _agillm43_decompress_cache_info(path) cache_path = info["cache_path"] manifest_path = info["manifest_path"] if not (cache_path.exists() and manifest_path.exists()): return None if cache_path.stat().st_size <= 0: return None manifest = json.loads(manifest_path.read_text()) if not _agillm43_manifest_matches(info, manifest): return None return cache_path except Exception: return None def _agillm43_file_looks_like_zstd_wrapper(path: pathlib.Path) -> bool: try: with open(path, "rb") as f: head = f.read(1 << 20) return ( _AGILLM43_PAYLOAD_CODEC_MAGIC.encode("utf-8") in head and b"agillm43_zstd_torch_v1" in head ) except Exception: return False def _agillm43_find_embedded_zstd_frame_offset(path: pathlib.Path, max_scan: int = 256 << 20): try: magic = _AGILLM43_ZSTD_FRAME_MAGIC chunk_size = 16 << 20 scanned = 0 prev = b"" with open(path, "rb") as f: while scanned < int(max_scan): chunk = f.read(chunk_size) if not chunk: return None hay = prev + chunk idx = hay.find(magic) if idx >= 0: return int(scanned - len(prev) + idx) prev = hay[-(len(magic) - 1):] scanned += len(chunk) except Exception: return None return None def _agillm43_zstd_frame_content_size(path: pathlib.Path, offset: int) -> int: try: import zstandard as zstd with open(path, "rb") as f: f.seek(int(offset)) head = f.read(32) params = zstd.get_frame_parameters(head) size = int(getattr(params, "content_size", 0) or 0) return max(0, size) except Exception: return 0 def _agillm43_stream_decompress_file_frame_to_file(path: pathlib.Path, offset: int, out_file) -> int: import zstandard as zstd counter = {"n": 0} class _CountingWriter: def write(self, chunk): out_file.write(chunk) counter["n"] += len(chunk) return len(chunk) with open(path, "rb") as f: f.seek(int(offset)) zstd.ZstdDecompressor().copy_stream(f, _CountingWriter()) return int(counter["n"]) def _agillm43_write_decompressed_cache_from_file_frame(path) -> pathlib.Path | None: path = pathlib.Path(path) if not _agillm43_file_looks_like_zstd_wrapper(path): return None offset = _agillm43_find_embedded_zstd_frame_offset(path) if offset is None: return None info = _agillm43_decompress_cache_info(path) cache_path = info["cache_path"] manifest_path = info["manifest_path"] if cache_path.exists() and manifest_path.exists() and cache_path.stat().st_size > 0: try: manifest = json.loads(manifest_path.read_text()) if _agillm43_manifest_matches(info, manifest): return cache_path except Exception: pass cache_path.parent.mkdir(parents=True, exist_ok=True) raw_nbytes = _agillm43_zstd_frame_content_size(path, offset) if raw_nbytes > 0: try: import shutil free = int(shutil.disk_usage(str(cache_path.parent)).free) reserve = 512 * 1024 * 1024 if free < raw_nbytes + reserve: raise RuntimeError( f"not enough free disk for decompressed checkpoint cache: " f"need about {(raw_nbytes + reserve) / (1024 ** 3):.2f}GB, " f"free {free / (1024 ** 3):.2f}GB" ) except RuntimeError: raise except Exception: pass tmp_path = cache_path.with_name(f".{cache_path.name}.{os.getpid()}.tmp") tmp_manifest = manifest_path.with_name(f".{manifest_path.name}.{os.getpid()}.tmp") try: if tmp_path.exists(): tmp_path.unlink() print(f"[ckpt-cache] streaming embedded zstd once to {cache_path.name}", flush=True) with open(tmp_path, "wb") as f: written = _agillm43_stream_decompress_file_frame_to_file(path, offset, f) if raw_nbytes > 0 and int(written) != raw_nbytes: raise RuntimeError(f"decompressed cache size mismatch: wrote {written}, expected {raw_nbytes}") os.replace(str(tmp_path), str(cache_path)) manifest = { "schema_version": 1, "cache_kind": "agillm43_decompressed_pt", "source_path": info["source_path"], "source_name": info["source_name"], "source_size": info["source_size"], "source_mtime_ns": info["source_mtime_ns"], "source_sha256": info["source_sha256"], "source_id": info["source_id"], "zstd_frame_offset": int(offset), "raw_nbytes": int(written), "created_at_unix": time.time(), } tmp_manifest.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") os.replace(str(tmp_manifest), str(manifest_path)) return cache_path finally: for stale in (tmp_path, tmp_manifest): try: if stale.exists(): stale.unlink() except Exception: pass class _Agillm43PayloadReader: def __init__(self, data, offset: int = 0, chunk_size: int = 8 << 20): self._owner = None if torch.is_tensor(data): tensor = data.detach().cpu().contiguous() self._owner = tensor.numpy() view = memoryview(self._owner) else: if isinstance(data, memoryview): self._owner = data view = data else: self._owner = data view = memoryview(data if isinstance(data, (bytes, bytearray)) else bytes(data)) self._view = view.cast("B") self._pos = max(0, int(offset)) self._chunk_size = max(1 << 20, int(chunk_size or (8 << 20))) def read(self, n: int = -1) -> bytes: if self._pos >= len(self._view): return b"" if n is None or n < 0: n = self._chunk_size n = min(int(n), self._chunk_size) end = min(len(self._view), self._pos + n) out = self._view[self._pos:end].tobytes() self._pos = end return out def prefix(self, n: int) -> bytes: end = min(len(self._view), int(n)) return self._view[:end].tobytes() def _agillm43_stream_decompress_payload_to_file(data, out_file) -> int: reader = _Agillm43PayloadReader(data) if reader.prefix(4) == b"ZLIB": import zlib reader = _Agillm43PayloadReader(data, offset=4) dec = zlib.decompressobj() total = 0 while True: chunk = reader.read() if not chunk: break raw = dec.decompress(chunk) if raw: out_file.write(raw) total += len(raw) tail = dec.flush() if tail: out_file.write(tail) total += len(tail) return total import zstandard as zstd counter = {"n": 0} class _CountingWriter: def write(self, chunk): out_file.write(chunk) counter["n"] += len(chunk) return len(chunk) zstd.ZstdDecompressor().copy_stream(reader, _CountingWriter()) return int(counter["n"]) def _agillm43_write_decompressed_cache(path, wrapper: dict) -> pathlib.Path: info = _agillm43_decompress_cache_info(path) cache_path = info["cache_path"] manifest_path = info["manifest_path"] if cache_path.exists() and manifest_path.exists() and cache_path.stat().st_size > 0: try: manifest = json.loads(manifest_path.read_text()) if _agillm43_manifest_matches(info, manifest): return cache_path except Exception: pass cache_path.parent.mkdir(parents=True, exist_ok=True) raw_nbytes = int(wrapper.get("raw_nbytes") or 0) if raw_nbytes > 0: try: import shutil free = int(shutil.disk_usage(str(cache_path.parent)).free) reserve = 512 * 1024 * 1024 if free < raw_nbytes + reserve: raise RuntimeError( f"not enough free disk for decompressed checkpoint cache: " f"need about {(raw_nbytes + reserve) / (1024 ** 3):.2f}GB, " f"free {free / (1024 ** 3):.2f}GB" ) except RuntimeError: raise except Exception: pass tmp_path = cache_path.with_name(f".{cache_path.name}.{os.getpid()}.tmp") tmp_manifest = manifest_path.with_name(f".{manifest_path.name}.{os.getpid()}.tmp") try: if tmp_path.exists(): tmp_path.unlink() print(f"[ckpt-cache] decompressing once to {cache_path.name}", flush=True) with open(tmp_path, "wb") as f: written = _agillm43_stream_decompress_payload_to_file(wrapper["payload"], f) if raw_nbytes > 0 and int(written) != raw_nbytes: raise RuntimeError(f"decompressed cache size mismatch: wrote {written}, expected {raw_nbytes}") os.replace(str(tmp_path), str(cache_path)) manifest = { "schema_version": 1, "cache_kind": "agillm43_decompressed_pt", "source_path": info["source_path"], "source_name": info["source_name"], "source_size": info["source_size"], "source_mtime_ns": info["source_mtime_ns"], "source_sha256": info["source_sha256"], "source_id": info["source_id"], "raw_nbytes": int(written), "created_at_unix": time.time(), } tmp_manifest.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n") os.replace(str(tmp_manifest), str(manifest_path)) return cache_path finally: for stale in (tmp_path, tmp_manifest): try: if stale.exists(): stale.unlink() except Exception: pass def _agillm43_load_pt(path, map_location="cpu", weights_only=False, skip_keys=None): path = pathlib.Path(path) cached = _agillm43_find_decompressed_cache(path) if cached is not None: try: print(f"[ckpt-cache] using decompressed cache {cached.name}", flush=True) return torch.load(cached, map_location=map_location, weights_only=weights_only) except Exception as exc: print(f"[ckpt-cache] ignoring invalid cache {cached.name}: {exc}", flush=True) try: cached.unlink() cached.with_suffix(cached.suffix + ".json").unlink(missing_ok=True) except Exception: pass if _agillm43_decompress_cache_enabled(): cached = _agillm43_write_decompressed_cache_from_file_frame(path) if cached is not None: print(f"[ckpt-cache] using decompressed cache {cached.name}", flush=True) return torch.load(cached, map_location=map_location, weights_only=weights_only) obj = torch.load(path, map_location=map_location, weights_only=weights_only) if isinstance(obj, dict) and obj.get(_AGILLM43_SHARDED_CODEC_MAGIC) == _AGILLM43_SHARDED_CODEC_MAGIC: return _agillm43_load_sharded_pt(path, obj, map_location=map_location, weights_only=weights_only, skip_keys=skip_keys) if isinstance(obj, dict) and obj.get(_AGILLM43_PAYLOAD_CODEC_MAGIC) == "agillm43_zstd_torch_v1": if _agillm43_decompress_cache_enabled(): import gc as _gc cached = _agillm43_write_decompressed_cache(path, obj) del obj _gc.collect() print(f"[ckpt-cache] using decompressed cache {cached.name}", flush=True) return torch.load(cached, map_location=map_location, weights_only=weights_only) import io raw = _agillm43_zstd_decompress(obj["payload"]) return torch.load(io.BytesIO(raw), map_location=map_location, weights_only=weights_only) return obj def _do_delta_save(tensors: dict, path: pathlib.Path, meta: dict, codec: str = "zstd"): """Background worker: write weight-only checkpoint + checksum.""" try: path.parent.mkdir(exist_ok=True, parents=True) tmp = path.with_suffix(path.suffix + ".dtmp") payload = {"weights": tensors, **meta} info = _agillm43_save_pt(payload, tmp, codec=codec, zstd_level=1) _agillm43_finalize_pt_save(tmp, path, info) digest = _sha256_file(path) # Write sidecar checksum path.with_suffix(".sha256").write_text(f"{digest} {path.name}\n") if info.get("codec") == "zstd": print(f" [delta] saved {path.name} ({digest[:12]}...) codec=zstd ratio={info.get('ratio', 0.0):.2f}x") elif info.get("codec") == "block-sharded": print(f" [delta] saved {path.name} ({digest[:12]}...) codec=block-sharded shards={info.get('shards', 0)}") else: print(f" [delta] saved {path.name} ({digest[:12]}...) codec=raw") except Exception as e: print(f" [delta] FAILED {path.name}: {e}") def _delete_delta_artifacts(path: pathlib.Path): for sidecar in ( path, path.with_suffix(".sha256"), path.with_suffix(path.suffix + ".upload.sha256"), path.with_suffix(path.suffix + ".dtmp"), ): try: if sidecar.exists(): sidecar.unlink() except Exception: pass def _unwrap_compiled_module(module: nn.Module) -> nn.Module: """Return the original module when torch.compile wrapped it.""" return getattr(module, "_orig_mod", module) def _checkpoint_state_dict(module: nn.Module) -> dict: """State dict with stable keys, even when module is torch.compile'd.""" return _unwrap_compiled_module(module).state_dict() def _strip_orig_mod_prefix(state: dict) -> dict: """Accept older deltas accidentally saved from compiled modules.""" if not isinstance(state, dict): return state prefix = "_orig_mod." if not any(isinstance(k, str) and k.startswith(prefix) for k in state): return state return { (k[len(prefix):] if isinstance(k, str) and k.startswith(prefix) else k): v for k, v in state.items() } def _cat_legacy_weight_blocks(blocks: list) -> Optional[torch.Tensor]: if not blocks or not all(torch.is_tensor(t) for t in blocks): return None first = blocks[0] tail_shape = tuple(first.shape[1:]) if any(t.dtype != first.dtype or t.device != first.device for t in blocks): return None if any(t.ndim != first.ndim or tuple(t.shape[1:]) != tail_shape for t in blocks): return None return torch.cat(blocks, dim=0).contiguous() def _fuse_qkv_in_state_dict(sd: dict) -> dict: """Fold legacy q/k/v.weight triples into qkv.weight before loading/filtering.""" if not isinstance(sd, dict): return sd prefixes = set() for key in list(sd.keys()): for suffix in (".q.weight", ".k.weight", ".v.weight"): if isinstance(key, str) and key.endswith(suffix): prefixes.add(key[: -len(suffix)]) for prefix in prefixes: qk, kk, vk = prefix + ".q.weight", prefix + ".k.weight", prefix + ".v.weight" fk = prefix + ".qkv.weight" if qk in sd and kk in sd and vk in sd and fk not in sd: fused = _cat_legacy_weight_blocks([sd[qk], sd[kk], sd[vk]]) if fused is not None: sd[fk] = fused sd.pop(qk) sd.pop(kk) sd.pop(vk) return sd def _expand_dense_ffn_to_moe_state_dict(sd: dict, target_sd: dict) -> dict: if not isinstance(sd, dict) or not isinstance(target_sd, dict): return sd out = dict(sd) seeded_prefixes: set[str] = set() for target_key, target in target_sd.items(): if not isinstance(target_key, str) or ".ff.experts." not in target_key: continue match = re.match(r"(blocks\.\d+\.ff\.)experts\.\d+\.(0|2)\.(weight|bias)$", target_key) if not match: continue prefix = match.group(1) legacy_key = f"{prefix}{match.group(2)}.{match.group(3)}" src = out.get(legacy_key) if target_key not in out and torch.is_tensor(src) and torch.is_tensor(target) and tuple(src.shape) == tuple(target.shape): out[target_key] = src seeded_prefixes.add(prefix) for prefix in seeded_prefixes: router_key = prefix + "router.weight" router_target = target_sd.get(router_key) if router_key not in out and torch.is_tensor(router_target): out[router_key] = router_target.detach().clone() for legacy_suffix in ("0.weight", "0.bias", "2.weight", "2.bias"): out.pop(prefix + legacy_suffix, None) return out def _reconcile_shared_expert_keys(sd: dict, target_sd: dict) -> dict: """Warm-start compat between shared-expert (4.3) and shared-less (4.2) checkpoints. - Shared-less checkpoint into a model WITH shared experts: fill the missing `.ff.shared.` keys from the freshly initialised module values. The shared output layer is zero-initialised, so the warm-started model is numerically identical to the source checkpoint at step 0 (it then learns to contribute). - Shared-expert checkpoint into a model WITHOUT them: drop the `.ff.shared.` keys (everything transferable is kept; only the shared path is shed). """ if not isinstance(sd, dict) or not isinstance(target_sd, dict): return sd out = dict(sd) filled = 0 dropped = 0 for key, target in target_sd.items(): if isinstance(key, str) and ".ff.shared." in key and key not in out and torch.is_tensor(target): out[key] = target.detach().clone() filled += 1 for key in list(out.keys()): if isinstance(key, str) and ".ff.shared." in key and key not in target_sd: out.pop(key) dropped += 1 if filled: print(f"[warm-start] shared experts: {filled} keys init fresh (zero-init no-op)", flush=True) if dropped: print(f"[warm-start] shared experts: {dropped} checkpoint keys dropped (model has none)", flush=True) return out def _prepare_core_state_dict_for_load(core: nn.Module, sd: dict) -> dict: sd = _strip_orig_mod_prefix(sd) sd = _fuse_qkv_in_state_dict(dict(sd)) if isinstance(sd, dict) else sd if isinstance(sd, dict): sd = _expand_dense_ffn_to_moe_state_dict(sd, core.state_dict()) sd = _reconcile_shared_expert_keys(sd, core.state_dict()) return sd def _split_qkv_in_state_dict_for_test(sd: dict) -> dict: out = dict(sd) for key in list(out.keys()): if not isinstance(key, str) or not key.endswith(".qkv.weight"): continue base = key[: -len(".qkv.weight")] q, k, v = out.pop(key).chunk(3, dim=0) out[base + ".q.weight"] = q.clone() out[base + ".k.weight"] = k.clone() out[base + ".v.weight"] = v.clone() return out def _clone_opt_value(value): if torch.is_tensor(value): return value.detach().clone() return copy.deepcopy(value) def _optimizer_param_name_lookup(core, ar_h, sat_h, nat_h=None) -> dict[int, str]: out = {} for prefix, module in (("core", core), ("ar", ar_h), ("sat", sat_h), ("nat", nat_h)): if module is None: continue for name, param in module.named_parameters(): out.setdefault(id(param), f"{prefix}.{name}") return out def _optimizer_group_param_names(opt, core, ar_h, sat_h, nat_h=None) -> List[List[str]]: lookup = _optimizer_param_name_lookup(core, ar_h, sat_h, nat_h) return [ [lookup.get(id(param), f"") for param in group["params"]] for group in opt.param_groups ] def _legacy_names_for_current_param(name: str) -> List[str]: if name.endswith(".qkv.weight"): base = name[: -len(".qkv.weight")] return [base + ".q.weight", base + ".k.weight", base + ".v.weight"] return [name] def _fuse_legacy_optimizer_param_state(states: List[dict]) -> Optional[dict]: if len(states) < 2 or any(not isinstance(state, dict) for state in states): return None common = set(states[0]) for state in states[1:]: common &= set(state) out = {} for key in common: vals = [state[key] for state in states] if all(torch.is_tensor(v) for v in vals): shape = vals[0].shape if vals[0].ndim > 0 and all(v.shape == shape for v in vals[1:]): out[key] = torch.cat([v.detach().clone() for v in vals], dim=0).contiguous() else: out[key] = vals[0].detach().clone() else: out[key] = copy.deepcopy(vals[0]) return out def _fuse_legacy_qkv_optimizer_state(opt_state: dict, opt, core, ar_h, sat_h, nat_h=None) -> Optional[dict]: """Remap pre-QKV-fusion AdamW state to the current fused parameter layout.""" if not isinstance(opt_state, dict) or "state" not in opt_state or "param_groups" not in opt_state: return None current_sd = opt.state_dict() current_names = _optimizer_group_param_names(opt, core, ar_h, sat_h, nat_h) legacy_names = [ [legacy for name in group_names for legacy in _legacy_names_for_current_param(name)] for group_names in current_names ] if len(legacy_names) != len(opt_state.get("param_groups", [])): return None legacy_name_to_pid = {} for group_idx, names in enumerate(legacy_names): old_params = list(opt_state["param_groups"][group_idx].get("params", [])) if len(names) != len(old_params): return None for name, pid in zip(names, old_params): legacy_name_to_pid[name] = pid new_groups = [] for group_idx, current_group in enumerate(current_sd["param_groups"]): new_group = copy.deepcopy(opt_state["param_groups"][group_idx]) new_group["params"] = list(current_group["params"]) if "param_names" in new_group: new_group["param_names"] = list(current_names[group_idx]) new_groups.append(new_group) old_states = opt_state.get("state", {}) new_states = {} for group_names, current_group in zip(current_names, current_sd["param_groups"]): for name, new_pid in zip(group_names, current_group["params"]): legacy_set = _legacy_names_for_current_param(name) if len(legacy_set) > 1: old_pids = [legacy_name_to_pid.get(legacy) for legacy in legacy_set] if all(pid in old_states for pid in old_pids): fused = _fuse_legacy_optimizer_param_state([old_states[pid] for pid in old_pids]) if fused is not None: new_states[new_pid] = fused continue old_pid = legacy_name_to_pid.get(name) if old_pid in old_states: new_states[new_pid] = {key: _clone_opt_value(value) for key, value in old_states[old_pid].items()} return {"state": new_states, "param_groups": new_groups} def _optimizer_state_compatibility_reason(opt_state: dict, opt) -> tuple[bool, str]: """Return whether a checkpoint optimizer state is safe to load into opt. torch Optimizer.load_state_dict can accept mismatched param-group option dictionaries and only fail later at step time. A diagnostic PowerStep checkpoint uses momentum/beta groups; AdamW-family optimizers need betas/eps. Cross-family states are treated as weight-only resumes. """ if not isinstance(opt_state, dict): return False, "checkpoint has no optimizer state" groups = opt_state.get("param_groups") if not isinstance(groups, list) or not groups: return False, "checkpoint optimizer has no param_groups" saved_keys = set() for group in groups: if isinstance(group, dict): saved_keys.update(group.keys()) saved_keys.discard("params") saved_keys.discard("param_names") cls = opt.__class__.__name__.lower() saved_powerstep = bool({"momentum", "beta"} & saved_keys) saved_adam = bool({"betas", "eps"} & saved_keys) wants_powerstep = "powerstep" in cls wants_adam = "adam" in cls if wants_powerstep and saved_adam and not saved_powerstep: return False, f"Adam-style checkpoint optimizer keys {sorted(saved_keys)} do not match {opt.__class__.__name__}" if wants_adam: if saved_powerstep and not saved_adam: return False, f"PowerStep checkpoint optimizer keys {sorted(saved_keys)} do not match {opt.__class__.__name__}" if not saved_adam: return False, f"checkpoint optimizer keys {sorted(saved_keys)} are missing Adam betas/eps for {opt.__class__.__name__}" return True, "compatible" def _agillm43_release_loaded_checkpoint(ck): """Drop large checkpoint payloads promptly after resume.""" try: if isinstance(ck, dict): ck.clear() except Exception: pass try: import gc as _gc _gc.collect() except Exception: pass try: import ctypes as _ctypes _ctypes.CDLL("libc.so.6").malloc_trim(0) except Exception: pass def save_delta(core, ar_h, sat_h, nat_h, step: int, seen_tok: int, save_dir: pathlib.Path, phase_name: str, delta_codec: str = "zstd3", provenance=None, origin_tag: str = "", dt_tag: str = "", role_tag: str = ""): """Save weight-only delta in background thread. Non-blocking.""" global _delta_thread # Wait for any previous delta write to finish if _delta_thread is not None and _delta_thread.is_alive(): _delta_thread.join(timeout=60) # Snapshot weights to CPU (detach from GPU graph) with _delta_lock: tensors = { "core": {k: v.detach().cpu() for k, v in _checkpoint_state_dict(core).items()}, "ar": {k: v.detach().cpu() for k, v in _checkpoint_state_dict(ar_h).items()}, "sat": {k: v.detach().cpu() for k, v in _checkpoint_state_dict(sat_h).items()}, } if nat_h is not None: tensors["nat"] = {k: v.detach().cpu() for k, v in _checkpoint_state_dict(nat_h).items()} meta = {"step": step, "seen_tok": seen_tok, "wall_time": time.time(), "delta": True, "agillm43_delta_codec": str(delta_codec or "off"), **_tokenizer_payload()} # Add provenance to delta checkpoints so hourly durable artifacts carry lineage. try: if provenance is not None: _agillm_provenance.embed(meta, dict(provenance)) else: _agillm_provenance.embed(meta, _agillm_provenance.collect(None, step=step, seen_tok=seen_tok, loss=0.0, batch_size=0, block_size=0, checkpoint_type="delta")) except Exception: pass path = save_dir / f"{phase_name}_delta_step{step:08d}{origin_tag}{dt_tag}{role_tag}.pt" _delta_thread = threading.Thread(target=_do_delta_save, args=(tensors, path, meta, delta_codec), daemon=True) _delta_thread.start() def _prune_delta_files_to_count(save_dir: pathlib.Path, phase_name: str, keep_count: int): """Keep only the newest keep_count complete delta files.""" try: pattern = f"{phase_name}_delta_step*.pt" deltas = sorted( [p for p in save_dir.glob(pattern) if p.stat().st_size > 0], key=lambda p: p.stat().st_mtime ) excess = len(deltas) - max(0, keep_count) if excess > 0: for p in deltas[:excess]: _delete_delta_artifacts(p) print(f" [delta-prune] deleted {p.name}") except Exception as e: print(f" [delta-prune] error: {e}") def _prune_deltas(save_dir: pathlib.Path, phase_name: str, max_deltas: int): """Keep only the most recent max_deltas delta files.""" if max_deltas is None or max_deltas <= 0: return _prune_delta_files_to_count(save_dir, phase_name, max_deltas) def _pinned_basenames(save_dir: pathlib.Path) -> set: try: txt = (save_dir / ".pinned").read_text() return {ln.strip().split("/")[-1] for ln in txt.splitlines() if ln.strip() and not ln.strip().startswith("#")} except Exception: return set() def _disk_hygiene(save_dir, phase_name: str, args, reason: str = ""): """In-file disk auto-prune so the training disk never wedges (a full disk makes Python unable to even start -> watchdog crash-loop). All AGILLM-4.2 disk pruning lives here in the single file rather than an external janitor that can silently die. Conservative: removes orphan *.tmp partial writes, full checkpoints beyond --max_ckpts, deltas beyond --delta_max_keep, stale side-cycle rounds and applied async-update artifacts, and escalates under --disk_free_floor_gb. NEVER deletes the newest full checkpoint, the resume/seed deltas, files younger than 2 min, or anything listed in /.pinned. Best-effort: never raises into the training loop.""" import shutil, glob as _glob try: save_dir = pathlib.Path(save_dir) ws = save_dir.parent pinned = _pinned_basenames(save_dir) floor = float(getattr(args, "disk_free_floor_gb", 0.0) or 0.0) now = time.time() def free_gb(): try: return shutil.disk_usage(str(save_dir)).free / (1024 ** 3) except Exception: return 1e9 def young(p, secs=120): try: return (now - p.stat().st_mtime) < secs except Exception: return True def rm(p): try: if p.name in pinned: return False if p.is_dir(): shutil.rmtree(p, ignore_errors=True) else: p.unlink() print(f" [disk] pruned {p.name}", flush=True) return True except Exception: return False def newest_first(paths): return sorted(paths, key=lambda p: p.stat().st_mtime, reverse=True) # 1) orphan partial writes (a live save's *.tmp is younger than 2 min) for t in save_dir.glob("*.tmp"): if not young(t): rm(t) # 2) full checkpoints beyond --max_ckpts (keep newest) keep_full = max(1, int(getattr(args, "max_ckpts", 2) or 2)) fulls = newest_first([p for p in save_dir.glob(f"{phase_name}_step*.pt") if not p.name.endswith(".resume_delta.pt")]) for p in fulls[keep_full:]: if not young(p): rm(p) # 3) deltas beyond --delta_max_keep keep_delta = max(1, int(getattr(args, "delta_max_keep", 1) or 1)) deltas = newest_first(list(save_dir.glob(f"{phase_name}_delta_step*.pt"))) for p in deltas[keep_delta:]: if not young(p): rm(p) # 4) transient side artifacts (side-cycle rounds, applied async updates) rounds = ws / "agillm41_side_rounds" rdirs = newest_first([d for d in rounds.glob("side_cycle_*") if d.is_dir()]) if rounds.exists() else [] for p in rdirs[2:]: rm(p) su = ws / "agillm41_side_updates" inc = su / "incoming" if inc.exists(): for p in newest_first(list(inc.glob("*.pt")))[4:]: if not young(p): rm(p) for sub in ("accepted", "rejected"): d = su / sub if d.exists(): for p in d.glob("*"): if not young(p, 600): rm(p) # 4b) V100 federation-cutover artifacts (fed14_* round/results/cache staging and # per-GPU side_updates_g*). These are named differently from the legacy # side_rounds / side_updates layout swept in section 4, so the original glob # never matched them and they accumulated (root cause of the 2026-06 disk creep). # Keep the newest round + results dir (an in-flight round is recent => young()); # applied side-updates already live bounded in agillm41_side_updates/incoming. try: fed_round = newest_first([d for d in ws.glob("agillm_v100_fed14_round_*") if d.is_dir()]) fed_res = newest_first([d for d in ws.glob("agillm_v100_fed14_results_*") if d.is_dir()]) for p in fed_round[1:] + fed_res[1:]: if not young(p, 1800): rm(p) for hb in ws.glob("agillm_v100_fed14_round_*.heartbeat.jsonl"): if not young(hb, 1800): rm(hb) for c in ws.glob("agillm_v100_fed14_cache"): if c.is_dir() and not young(c, 1800): rm(c) for gd in ws.glob("agillm41_side_updates_g*"): inc_g = gd / "incoming" if inc_g.exists(): for p in newest_first(list(inc_g.glob("*.pt")))[4:]: if not young(p): rm(p) for sub in ("accepted", "rejected"): d = gd / sub if d.exists(): for p in d.glob("*"): if not young(p, 600): rm(p) except Exception: pass # 5) escalate under the free-space floor (transient + extra ckpts only) if floor > 0 and free_gb() < floor: print(f" [disk] below floor {floor:.0f}GB (free {free_gb():.1f}GB){(' ' + reason) if reason else ''}; escalating", flush=True) for p in rdirs[1:]: rm(p) for p in newest_first(list(save_dir.glob(f"{phase_name}_delta_step*.pt")))[1:]: if not young(p): rm(p) for p in newest_first([p for p in save_dir.glob(f"{phase_name}_step*.pt") if not p.name.endswith(".resume_delta.pt")])[1:]: if not young(p): rm(p) print(f" [disk] after escalation: {free_gb():.1f}GB free", flush=True) except Exception as e: print(f"[disk-hygiene] error: {e}", flush=True) def _build_val_set(source, chat_cfg, args, block): """Capture a fixed held-out token sample (val_seed stream) as (1, block+1) CPU batches. A fixed sample re-evaluated periodically gives a comparable loss curve over training.""" n = int(getattr(args, "val_tokens", 0) or 0) if n <= 0: return [] want = max(1, n // (block + 1)) * (block + 1) val_source_requested = str(getattr(args, "val_source", "") or "").strip() val_source = val_source_requested if val_source and _looks_numeracy_only_sources(val_source) and not _looks_numeracy_only_sources(source): print( "[dataset-policy] val_source is numeracy-only; using effective language pretrain mix for validation", flush=True, ) val_source = source use_hot_config = False else: use_hot_config = not bool(val_source) val_source = val_source or source print( f"[val] building held-out set from {val_source} " f"(hot_config={'on' if use_hot_config else 'off'}, seed {getattr(args, 'val_seed', 1337)})", flush=True, ) toks = [] try: for t in token_stream( val_source, want, seed=int(getattr(args, "val_seed", 1337)), chat=chat_cfg.get("chat", False), chat_messages_key=chat_cfg.get("key", "messages"), sft_add_generation_prompt=chat_cfg.get("gen_prompt", False), dataset_field_text=chat_cfg.get("text_field", "text"), streaming=True, use_hot_config=use_hot_config, ): toks.append(int(t)) if len(toks) >= want: break except Exception as e: print(f"[val] failed to build val set ({type(e).__name__}: {e}); validation disabled", flush=True) return [] batches = [torch.tensor(toks[i:i + block + 1], dtype=torch.long).unsqueeze(0) for i in range(0, len(toks) - block, block + 1)] print(f"[val] held-out set ready: {len(batches)} batches x {block + 1} tokens (seed {getattr(args, 'val_seed', 1337)})", flush=True) return batches def _run_validation(core, ar_h, val_batches, args, step): """Full-stack AR cross-entropy on the fixed held-out batches (no_grad, eval mode).""" if not val_batches: return None was_training = core.training core.eval(); ar_h.eval() tot_ce, tot_tok = 0.0, 0 try: with torch.no_grad(): for ids_cpu in val_batches: ids = ids_cpu.to(DEV) with amp(args.amp): h = core(ids, causal_mask(ids.size(1), structured=use_structured_masks(args))) ce = fused_ce(h[:, :-1], ar_h.proj.weight, ids[:, 1:]) ntok = ids.size(1) - 1 tot_ce += float(ce.detach()) * ntok tot_tok += ntok except Exception as e: print(f"[val] eval error ({type(e).__name__}: {e}); skipping this round", flush=True) if was_training: core.train(); ar_h.train() return None if was_training: core.train(); ar_h.train() ce = tot_ce / max(1, tot_tok) ppl = math.exp(min(20.0, ce)) print(f"[val] step={step} tokens={tot_tok} ce={ce:.4f} ppl={ppl:.2f}", flush=True) return ce def _load_module_state_compatible(module: nn.Module, state: dict, label: str = "module") -> int: """Load matching tensors only; skip obsolete untied vocab matrices for tied heads.""" if not isinstance(state, dict): return 0 state = _strip_orig_mod_prefix(state) tgt_sd = module.state_dict() tied = bool(getattr(module, "tie_weights", False)) filt = {} skipped = [] for k, v in state.items(): if tied and k == "proj.weight": skipped.append(k) continue if k in tgt_sd and hasattr(v, "shape") and v.shape == tgt_sd[k].shape: filt[k] = v else: skipped.append(k) if filt: module.load_state_dict(filt, strict=False) if tied and skipped: print(f"[ckpt] {label}: tied head active; skipped old untied tensors: {', '.join(skipped[:4])}{'...' if len(skipped)>4 else ''}") return len(filt) class _skip_param_init: """Suppress torch.nn.init.* tensor fills while constructing inference models. Every parameter is overwritten from the checkpoint immediately after construction, so constructor random init is pure startup cost. Params the checkpoint cannot supply are re-initialized afterwards. """ _FILLS = ( "uniform_", "normal_", "trunc_normal_", "constant_", "ones_", "zeros_", "eye_", "dirac_", "xavier_uniform_", "xavier_normal_", "kaiming_uniform_", "kaiming_normal_", "orthogonal_", "sparse_", ) def __enter__(self): import torch.nn.init as _init self._saved = {} for name in self._FILLS: fn = getattr(_init, name, None) if fn is None: continue self._saved[name] = fn def _noop(tensor, *args, **kwargs): return tensor setattr(_init, name, _noop) return self def __exit__(self, *exc): import torch.nn.init as _init for name, fn in self._saved.items(): setattr(_init, name, fn) return False def _reinit_params_missing_from_state(core: nn.Module, sd_core: dict): if not isinstance(sd_core, dict): return present = set(_strip_orig_mod_prefix(sd_core).keys()) missing_mods = {} for name, _ in core.named_parameters(): if name in present: continue mod_name = name.rsplit(".", 1)[0] if "." in name else "" missing_mods.setdefault(mod_name, name) reinit = 0 for mod_name, param_name in missing_mods.items(): try: mod = core.get_submodule(mod_name) if mod_name else core except AttributeError: mod = None if mod is not None and hasattr(mod, "reset_parameters"): mod.reset_parameters() reinit += 1 else: print(f"[infer] WARNING: param {param_name} absent from checkpoint and module has no reset_parameters; it may be uninitialized", flush=True) if reinit: print(f"[infer] reinitialized {reinit} module(s) for checkpoint-missing parameters", flush=True) def load_delta(path: pathlib.Path, core, ar_h, sat_h, nat_h=None): """Load weight-only delta. Returns (step, seen_tok) or raises.""" # Verify checksum if sidecar exists sha_path = path.with_suffix(".sha256") if sha_path.exists(): expected = sha_path.read_text().split()[0] actual = _sha256_file(path) if expected != actual: raise ValueError(f"Checksum mismatch for {path.name}: expected {expected[:12]}... got {actual[:12]}...") print(f" [delta] checksum OK for {path.name}") ck = _agillm43_load_pt(path, map_location="cpu", weights_only=False) if not ck.get("delta"): raise ValueError(f"{path.name} is not a delta checkpoint") core.load_state_dict(_prepare_core_state_dict_for_load(core, ck["weights"]["core"])) _load_module_state_compatible(ar_h, ck["weights"].get("ar", {}), "ar") _load_module_state_compatible(sat_h, ck["weights"].get("sat", {}), "sat") if nat_h is not None: nat_sd = ck["weights"].get("nat") if nat_sd is not None: _load_module_state_compatible(nat_h, nat_sd, "nat") else: print("[nat] Delta has no NAT head; keeping fresh NAT initialization") _restore_tokenizer_from_ckpt(ck, path) step = ck.get("step", 0) seen_tok = ck.get("seen_tok", 0) _agillm43_release_loaded_checkpoint(ck) return step, seen_tok def _flush_delta(): """Wait for any in-flight delta save to complete.""" global _delta_thread if _delta_thread is not None and _delta_thread.is_alive(): print(" [delta] flushing in-flight write...") _delta_thread.join(timeout=120) def save_ckpt(path: pathlib.Path, core, ar_h, sat_h, nat_h, opt, scaler, meta, codec: str = "zstd", provenance=None): path.parent.mkdir(exist_ok=True, parents=True) tmp = path.with_suffix(path.suffix + ".tmp") tokenizer_payload = _tokenizer_payload() tokenizer_payload.setdefault("tokenizer_payload_schema", 2) state = { "core": _checkpoint_state_dict(core), "ar": _checkpoint_state_dict(ar_h), "sat": _checkpoint_state_dict(sat_h), "opt": opt.state_dict(), "scaler": scaler.state_dict(), "cfg": meta.get("cfg"), **tokenizer_payload, "transformers_version": __import__("transformers").__version__, "tokenizers_version": __import__("tokenizers").__version__, "tie_weights": meta.get("tie_weights", False), **{k: v for k, v in meta.items() if k not in ("cfg", "tie_weights")} } if nat_h is not None: state["nat"] = _checkpoint_state_dict(nat_h) ckpt_codec = str(codec or "off") state["agillm43_ckpt_codec"] = ckpt_codec if provenance is not None: try: provenance = dict(provenance) except Exception: provenance = {"raw_provenance_repr": repr(provenance)} source_path = str(provenance.get("warmstart_source_path") or "") try: save_root = str(path.parent.resolve()) except Exception: save_root = str(path.parent) if not source_path: warmstart_kind = "from_scratch" else: source_abs = os.path.abspath(source_path) save_abs = os.path.abspath(save_root) master_marker = f"{os.sep}agillm4_v100_master_ckpts{os.sep}" if master_marker in source_abs: warmstart_kind = "warmstarted_from_master" elif source_abs.startswith(save_abs + os.sep): warmstart_kind = "warmstarted_from_lane_checkpoint" else: warmstart_kind = "warmstarted_from_non_master_checkpoint" provenance["checkpoint_path"] = str(path) provenance["warmstart_kind"] = warmstart_kind provenance["created_from_scratch"] = warmstart_kind == "from_scratch" provenance["source_is_master_checkpoint"] = warmstart_kind == "warmstarted_from_master" provenance["source_is_lane_checkpoint"] = warmstart_kind == "warmstarted_from_lane_checkpoint" provenance["source_is_non_master_checkpoint"] = warmstart_kind == "warmstarted_from_non_master_checkpoint" state["agillm43_provenance"] = provenance state["agillm43_warmstart_kind"] = warmstart_kind state["agillm43_warmstart_source_path"] = source_path state["agillm43_checkpoint_summary"] = f"{warmstart_kind}; source={source_path or 'none'}; path={path}" info = _agillm43_save_pt(state, tmp, codec=ckpt_codec, zstd_level=1) _agillm43_finalize_pt_save(tmp, path, info) _write_tokenizer_sidecar(path, {k: state.get(k) for k in ("tokenizer_payload_schema", "tokenizer_id", "tokenizer_json", "tokenizer_bundle", "tokenizer_special", "transformers_version", "tokenizers_version") if state.get(k) is not None}) if provenance is not None: try: globals().get("_agillm_provenance").write_sidecar(path, provenance) except Exception as exc: print(f"[provenance] WARNING: failed to write sidecar for {path}: {exc}") latest_payload = {"path": str(path), "step": meta["step"]} if provenance is not None: latest_payload["agillm43_provenance"] = provenance latest_payload["warmstart_kind"] = provenance.get("warmstart_kind") latest_payload["warmstart_source_path"] = provenance.get("warmstart_source_path", "") latest_payload["checkpoint_summary"] = state.get("agillm43_checkpoint_summary") if meta.get("dataset_provenance"): latest_payload["dataset_provenance"] = meta.get("dataset_provenance") latest_payload["source_effective"] = meta.get("dataset_provenance", {}).get("source_effective", "") (path.parent / "latest.json").write_text(json.dumps(latest_payload)) if info.get("codec") == "zstd": print(f"\n✓ saved checkpoint {path.name} codec=zstd ratio={info.get('ratio', 0.0):.2f}x") elif info.get("codec") == "block-sharded": print(f"\n✓ saved checkpoint {path.name} codec=block-sharded shards={info.get('shards', 0)} dir={path.name}.shards") else: print(f"\n✓ saved checkpoint {path.name} codec=raw") def load_ckpt(path, core, ar_h, sat_h, opt, scaler, nat_h=None): p = _resolve_ckpt(path) or path ck = _try_load(p, map_location="cpu") if ck is None: raise FileNotFoundError(f"No valid checkpoint at {p}") core.load_state_dict(_prepare_core_state_dict_for_load(core, ck["core"])) _load_module_state_compatible(ar_h, ck.get("ar", {}), "ar") _load_module_state_compatible(sat_h, ck.get("sat", {}), "sat") if nat_h is not None: if "nat" in ck: _load_module_state_compatible(nat_h, ck["nat"], "nat") else: print("[nat] Checkpoint has no NAT head; keeping fresh NAT initialization") opt_state_loaded = False if opt.__class__.__name__ == "PowerStep": print("[ckpt] PowerStep optimizer selected; resetting checkpoint optimizer state") else: opt_state = ck.get("opt") compatible, reason = _optimizer_state_compatibility_reason(opt_state, opt) if not compatible: print(f"[ckpt] WARNING: optimizer state incompatible; resetting optimizer ({reason})") else: try: opt.load_state_dict(opt_state) opt_state_loaded = True except Exception as exc: fused_opt = _fuse_legacy_qkv_optimizer_state(opt_state, opt, core, ar_h, sat_h, nat_h) if fused_opt is not None: fused_compatible, fused_reason = _optimizer_state_compatibility_reason(fused_opt, opt) if fused_compatible: try: opt.load_state_dict(fused_opt) opt_state_loaded = True print("[ckpt] Converted legacy q/k/v optimizer state to fused qkv layout") except Exception as exc2: print(f"[ckpt] WARNING: optimizer state incompatible; resetting optimizer ({type(exc).__name__}: {exc}; qkv remap failed: {type(exc2).__name__}: {exc2})") else: print(f"[ckpt] WARNING: fused optimizer state incompatible; resetting optimizer ({fused_reason})") else: print(f"[ckpt] WARNING: optimizer state incompatible; resetting optimizer ({type(exc).__name__}: {exc})") if opt_state_loaded: try: scaler.load_state_dict(ck["scaler"]) except Exception as exc: print(f"[ckpt] WARNING: scaler state incompatible; resetting scaler ({type(exc).__name__}: {exc})") else: print("[ckpt] scaler state reset with optimizer state") # Restore tokenizer from checkpoint (embedded json preferred; never raises) _restore_tokenizer_from_ckpt(ck, p) # Warn if transformers version changed since checkpoint was saved if "transformers_version" in ck: import transformers as _tf if ck["transformers_version"] != _tf.__version__: print(f"[tokenizer] WARNING: checkpoint saved with transformers={ck['transformers_version']}, now running {_tf.__version__}") step = ck.get("step", 0) seen_tok = ck.get("seen_tok", 0) wall_time = ck.get("wall_time", time.time()) _agillm43_release_loaded_checkpoint(ck) return step, seen_tok, wall_time def _safe_load_any(path: pathlib.Path, tgt: nn.Module, key: str | None = None): p = _resolve_ckpt(path) or path if not p.exists(): return 0 ck = _try_load(p, map_location="cpu") if ck is None: return 0 sd = ck.get(key, ck) if key else ck if isinstance(sd, dict) and "state_dict" in sd: sd = sd["state_dict"] if isinstance(tgt, Encoder) or key == "core": sd = _prepare_core_state_dict_for_load(tgt, sd) else: sd = _strip_orig_mod_prefix(sd) sd = _fuse_qkv_in_state_dict(dict(sd)) if isinstance(sd, dict) else sd if not isinstance(sd, dict): return 0 tgt_sd = tgt.state_dict() filt = {k: v for k, v in sd.items() if k in tgt_sd and hasattr(v, "shape") and v.shape == tgt_sd[k].shape} if filt: tgt.load_state_dict(filt, strict=False) return len(filt) def infer_cfg_from_ckpt(path: pathlib.Path): p = _resolve_ckpt(path) or path if not p.exists(): return None sd = _try_load(p, map_location="cpu") if sd is None: return None if "cfg" in sd: return dict(sd["cfg"]) return None def _infer_cfg_from_delta_checkpoint(sd: dict) -> tuple[dict, bool, str]: """Infer model config for weight-only delta checkpoints. Delta checkpoints intentionally omit optimizer/scaler and can omit cfg. Native inference still needs the original architecture. Recover it from provenance when possible, then validate/fill from tensor shapes. """ weights = sd.get("weights") or {} core = weights.get("core") or {} ar = weights.get("ar") or {} emb = core.get("emb.weight") if not torch.is_tensor(emb) or emb.ndim != 2: raise ValueError("delta checkpoint missing core emb.weight; cannot infer cfg") d = int(emb.shape[1]) layer_ids = [] for key in core.keys(): if not key.startswith("blocks."): continue parts = key.split(".") if len(parts) > 2 and parts[1].isdigit(): layer_ids.append(int(parts[1])) if not layer_ids: raise ValueError("delta checkpoint has no block tensors; cannot infer layer count") layers = max(layer_ids) + 1 u = core.get("blocks.0.mha.U") if not torch.is_tensor(u) or u.ndim != 2: raise ValueError("delta checkpoint missing blocks.0.mha.U; cannot infer attention rank") dk = int(u.shape[0]) rank = int(u.shape[1]) if dk <= 0 or d % dk != 0: raise ValueError(f"delta checkpoint incompatible d/dk: d={d} dk={dk}") heads = d // dk prov = sd.get("agillm43_provenance") or {} train_argv = str(prov.get("train_argv") or "") if isinstance(prov, dict) else "" tokens = train_argv.split() preset_name = "" if "--preset" in tokens: idx = tokens.index("--preset") if idx + 1 < len(tokens): preset_name = tokens[idx + 1] if not preset_name: for name in PRESETS.keys(): if ("--preset " + name) in train_argv: preset_name = name break cfg = None source = "shapes" if preset_name in PRESETS: cand = dict(PRESETS[preset_name]) if int(cand.get("d", -1)) == d and int(cand.get("layers", -1)) == layers and int(cand.get("rank", -1)) == rank: cfg = cand source = "provenance:" + preset_name if cfg is None: matches = [] for name, cand in PRESETS.items(): if int(cand.get("d", -1)) == d and int(cand.get("layers", -1)) == layers and int(cand.get("rank", -1)) == rank: matches.append((name, dict(cand))) if matches: source = "preset:" + matches[0][0] cfg = matches[0][1] else: cfg = {"d": d, "layers": layers, "heads": heads, "rank": rank} cfg["d"] = d cfg["layers"] = layers cfg["heads"] = heads cfg["rank"] = rank qkv = core.get("blocks.0.mha.qkv.weight") tie_kv = bool(torch.is_tensor(qkv) and int(qkv.shape[0]) == 2 * d) cfg["tie_kv"] = tie_kv router = core.get("blocks.0.ff.router.weight") moe_ffn = torch.is_tensor(router) cfg["moe_ffn"] = bool(moe_ffn) if moe_ffn: cfg["moe_experts"] = int(router.shape[0]) cfg["moe_top_k"] = int(cfg.get("moe_top_k", 1) or 1) exp0 = core.get("blocks.0.ff.experts.0.0.weight") if torch.is_tensor(exp0) and exp0.ndim == 2: cfg["moe_mlp_mult"] = max(1, int(exp0.shape[0]) // d) shared_ids = set() for key in core.keys(): if not key.startswith("blocks.0.ff.shared."): continue parts = key.split(".") if len(parts) > 4 and parts[4].isdigit(): shared_ids.add(int(parts[4])) cfg["moe_shared_experts"] = len(shared_ids) shared0 = core.get("blocks.0.ff.shared.0.0.weight") if torch.is_tensor(shared0) and shared0.ndim == 2: cfg["moe_shared_mlp_mult"] = max(1, int(shared0.shape[0]) // d) else: cfg["moe_shared_mlp_mult"] = int(cfg.get("moe_shared_mlp_mult", 0) or 0) ar_weight = ar.get("proj.weight") if isinstance(ar, dict) else None ar_bias = ar.get("proj.bias") if isinstance(ar, dict) else None tie_weights = bool(sd.get("tie_weights", False)) if not tie_weights: tie_weights = "--tie_weights" in train_argv if not tie_weights and torch.is_tensor(ar_weight) and tuple(ar_weight.shape) == tuple(emb.shape) and ar_bias is None: tie_weights = True return cfg, tie_weights, source # ───────────────────────── Training Logic ───────────────────────── def _load_infer_head_state(module: nn.Module, state: dict, name: str): """Load inference heads across small checkpoint/schema drifts. Some older AGILLM-4 full checkpoints were saved before the current SAT/NAT head bias fields existed. For inference, preserve the old behavior by explicitly zero-filling missing bias tensors, while still failing on missing non-bias weights or shape mismatches. """ if not isinstance(state, dict): module.load_state_dict(state) return module_state = module.state_dict() patched = dict(state) zero_filled = [] shape_mismatch = [] for key, target in module_state.items(): if key not in patched and key.endswith('.bias') and torch.is_tensor(target): patched[key] = torch.zeros_like(target) zero_filled.append(key) for key, value in list(patched.items()): target = module_state.get(key) if target is None or not torch.is_tensor(value) or not torch.is_tensor(target): continue if tuple(value.shape) != tuple(target.shape): shape_mismatch.append(f"{key}: ckpt={tuple(value.shape)} model={tuple(target.shape)}") patched.pop(key) if shape_mismatch: raise RuntimeError(f"{name} checkpoint shape mismatch: " + "; ".join(shape_mismatch[:6])) loaded = module.load_state_dict(patched, strict=False) missing = [key for key in loaded.missing_keys if key not in zero_filled] if missing: raise RuntimeError(f"{name} checkpoint missing required keys: " + ", ".join(missing[:12])) notes = [] if zero_filled: notes.append("zero-filled " + ", ".join(zero_filled[:6])) if loaded.unexpected_keys: notes.append("ignored unexpected " + ", ".join(loaded.unexpected_keys[:6])) if notes: print(f"[infer-compat] {name}: " + "; ".join(notes), flush=True) def _sat_head_mlp_from_state(sd: dict) -> bool: sat_sd = sd.get("sat", {}) if sd.get("delta") and "weights" in sd: sat_sd = sd["weights"].get("sat", sat_sd) return any(str(key).startswith("proj.2.") for key in sat_sd) def _parse_grow_plan(s: str) -> List[int]: return sorted(set([int(x.strip()) for x in s.split(",") if x.strip() and int(x.strip()) >= 128])) def _count_enabled_params(*modules) -> int: seen_data_ptrs = set() total = 0 for m in modules: if m is None: continue for p in m.parameters(): if p.data_ptr() not in seen_data_ptrs: seen_data_ptrs.add(p.data_ptr()) total += p.numel() return total def _target_token_ratio(args) -> float: if getattr(args, "token_param_ratio", 0.0) and args.token_param_ratio > 0: return float(args.token_param_ratio) if str(getattr(args, "preset", "")).startswith("agillm4_"): return AGILLM4_TOKEN_PARAM_RATIO return 51.2 if args.chilla_max_double else 25.0 def _phase_freeze(core: nn.Module, *, freeze_core: bool, unfreeze_ln: bool, train_emb: bool): for p in core.parameters(): p.requires_grad = not freeze_core if freeze_core: if unfreeze_ln: for blk in core.blocks: for p in blk.ln1.parameters(): p.requires_grad = True for p in blk.ln2.parameters(): p.requires_grad = True for p in core.ln.parameters(): p.requires_grad = True if train_emb: for p in core.emb.parameters(): p.requires_grad = True def _side_update_unique_path(directory: pathlib.Path, name: str) -> pathlib.Path: directory.mkdir(parents=True, exist_ok=True) dest = directory / name if not dest.exists(): return dest stem, suffix = dest.stem, dest.suffix stamp = time.strftime("%Y%m%d-%H%M%S", time.gmtime()) for idx in range(1000): candidate = directory / f"{stem}.{stamp}.{idx}{suffix}" if not candidate.exists(): return candidate return directory / f"{stem}.{stamp}.{os.getpid()}{suffix}" def _side_update_move(path: pathlib.Path, directory: pathlib.Path) -> pathlib.Path: dest = _side_update_unique_path(directory, path.name) try: path.replace(dest) except OSError: import shutil shutil.move(str(path), str(dest)) return dest def _apply_async_side_updates(core: nn.Module, cfg: dict, args, step: int) -> tuple[list[dict], list[dict]]: update_dir_s = str(getattr(args, "async_update_dir", "") or "").strip() alpha = float(getattr(args, "async_update_alpha", 1.0) or 0.0) if not update_dir_s or alpha <= 0.0: return [], [] update_dir = pathlib.Path(update_dir_s) if not update_dir.exists(): return [], [] max_updates = max(1, int(getattr(args, "async_update_max_per_check", 1) or 1)) max_age = float(getattr(args, "async_update_max_age_sec", 0.0) or 0.0) accepted_dir = pathlib.Path(getattr(args, "async_update_accepted_dir", "") or (update_dir.parent / "accepted")) rejected_dir = pathlib.Path(getattr(args, "async_update_rejected_dir", "") or (update_dir.parent / "rejected")) param_map = dict(core.named_parameters()) buffer_map = dict(core.named_buffers()) now = time.time() applied: list[dict] = [] rejected: list[dict] = [] candidates = sorted( [p for p in update_dir.glob("*.pt") if p.is_file() and not p.name.endswith(".tmp")], key=lambda p: p.stat().st_mtime, ) for path in candidates[:max_updates]: reject_reason = "" try: if max_age > 0 and now - path.stat().st_mtime > max_age: reject_reason = f"stale update older than {max_age:g}s" raise ValueError(reject_reason) upd = _agillm43_load_pt(path, map_location="cpu", weights_only=False) kind = upd.get("kind") if kind not in {"agillm35_dblock_slice_update", "agillm4_dblock_slice_update", "agillm41_dblock_slice_update"}: raise ValueError(f"bad update kind {kind!r}") if dict(upd.get("cfg", {})) != dict(cfg): raise ValueError("cfg mismatch") update_mode = "state_lerp" block_state = upd.get("block_state") block_delta_state = upd.get("block_delta_state") if block_delta_state is not None: update_mode = "delta_add" block_codec = _agillm43_tensor_state_summary(block_delta_state) block_state = _agillm43_decode_tensor_state(block_delta_state) else: block_codec = _agillm43_tensor_state_summary(block_state) block_state = _agillm43_decode_tensor_state(block_state) if not isinstance(block_state, dict) or not block_state: raise ValueError("missing block_state or block_delta_state") changed = 0 with torch.no_grad(): for key, value in block_state.items(): target = param_map.get(key) if target is None: target = buffer_map.get(key) if target is None: raise KeyError(f"unknown core key {key}") if tuple(value.shape) != tuple(target.shape): raise ValueError(f"{key} shape mismatch update={tuple(value.shape)} target={tuple(target.shape)}") src = value.to(device=target.device, dtype=target.dtype, non_blocking=True) if update_mode == "delta_add": if not target.is_floating_point(): raise ValueError(f"{key} delta update targets non-floating tensor") target.add_(src, alpha=alpha) elif alpha >= 1.0: target.copy_(src) else: target.lerp_(src, alpha) changed += 1 del src dest = _side_update_move(path, accepted_dir) rec = { "path": str(dest), "worker_id": upd.get("worker_id"), "block_id": upd.get("block_id"), "layers": upd.get("layers"), "tokens": int(upd.get("tokens") or 0), "tok_per_sec": float(upd.get("tok_per_sec") or 0.0), "alpha": alpha, "keys": changed, "block_codec": block_codec, "update_mode": update_mode, } applied.append(rec) print(json.dumps({"event": "async_side_update_applied", "step": step, **rec}), flush=True) except Exception as exc: try: dest = _side_update_move(path, rejected_dir) except Exception: dest = path err = reject_reason or str(exc) print( json.dumps( { "event": "async_side_update_rejected", "step": step, "path": str(dest), "error": err, } ), flush=True, ) try: upd_partial = _agillm43_load_pt(dest, map_location="cpu", weights_only=False) if dest.exists() else {} except Exception: upd_partial = {} rejected.append({ "path": str(dest), "worker_id": upd_partial.get("worker_id"), "block_id": upd_partial.get("block_id"), "layers": upd_partial.get("layers"), "error": err, }) return applied, rejected # ── HF federation dataset logging ───────────────────────────────────────────── _HF_FED_UPDATES_REPO = "OpenTransformer/AGILLM-4.3-fed-updates" _HF_FED_ROUNDS_REPO = "OpenTransformer/AGILLM-4.3-fed-rounds" def _hf_fed_log_rows_bg(repo_id: str, rows: list, step: int) -> None: """Append JSONL rows to an HF dataset repo in a fire-and-forget background thread.""" if not rows: return token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") if not token: return try: from huggingface_hub import HfApi except ImportError: return def _upload(): try: api = HfApi(token=token) ts = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime()) fname = f"data/{step:08d}-{ts}-{os.getpid()}.jsonl" content = "\n".join(json.dumps(r, separators=(",", ":")) for r in rows) + "\n" api.upload_file( path_or_fileobj=content.encode(), path_in_repo=fname, repo_id=repo_id, repo_type="dataset", commit_message=f"fed log step {step}", ) except Exception as exc: print(f"[hf-fed-log] {repo_id} upload failed: {exc}", flush=True) threading.Thread(target=_upload, daemon=True).start() def _hf_fed_log_side_updates(applied: list, rejected: list, step: int) -> None: """Log accepted/rejected side-updates to HF AGILLM-4.3-fed-updates.""" ts = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) rows = [] for rec in applied: rows.append({ "ts_utc": ts, "step": step, "status": "accepted", "worker_id": rec.get("worker_id"), "block_id": rec.get("block_id"), "layers": rec.get("layers"), "tokens": rec.get("tokens"), "tok_per_sec": rec.get("tok_per_sec"), "alpha": rec.get("alpha"), "keys": rec.get("keys"), "update_mode": rec.get("update_mode"), "block_codec": rec.get("block_codec"), }) for rec in rejected: rows.append({ "ts_utc": ts, "step": step, "status": "rejected", "worker_id": rec.get("worker_id"), "block_id": rec.get("block_id"), "layers": rec.get("layers"), "tokens": None, "tok_per_sec": None, "alpha": None, "keys": None, "update_mode": None, "block_codec": None, "error": rec.get("error"), }) _hf_fed_log_rows_bg(_HF_FED_UPDATES_REPO, rows, step) def _hf_fed_log_round(step: int, seen_tok: int, loss: float, role_tag: str, origin_tag: str) -> None: """Log a delta-save event (federation round boundary) to HF AGILLM-4.3-fed-rounds.""" row = { "ts_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "step": step, "seen_tok": int(seen_tok), "loss": round(float(loss), 6), "role_tag": role_tag, "origin_tag": origin_tag, } _hf_fed_log_rows_bg(_HF_FED_ROUNDS_REPO, [row], step) # ── end HF federation dataset logging ───────────────────────────────────────── def _optimizer_param_groups(core, ar_h, sat_h, lr_core: float, lr_head: float, nat_h=None): # Shared/tied vocab projections must appear in only one optimizer group. # VRAM-first AGILLM-4 uses one embedding/projection tensor for AR/SAT/NAT. seen: set[int] = set() groups = [] def add(params, lr): unique = [] for p in params: if not p.requires_grad: continue key = id(p) if key in seen: continue seen.add(key) unique.append(p) if unique: groups.append({"params": unique, "lr": lr}) add(core.parameters(), lr_core) add(ar_h.parameters(), lr_head) add(sat_h.parameters(), lr_head) if nat_h is not None: add(nat_h.parameters(), lr_head) return groups class PowerStep(torch.optim.Optimizer): """Memory-efficient optimizer (arXiv:2605.10335): heavy-ball momentum + signed power transform, a SINGLE buffer (no Adam second moment). Update: m_t = gamma*m_{t-1} + g_t ; theta -= lr * (sign(m)*|m|^beta + wd*theta) beta in (0,1) gives Adam-like coordinate adaptivity; beta=1 -> SGD-momentum, beta=0 -> signSGD-momentum. Half the optimizer state of Adam. Faithful AGILLM-4.2 dblock-step benchmark (small model, real EDM objective, bf16): converged faster and to a LOWER loss than AdamW/paged_adamw8bit (EMA 6.6 vs 8.7-9.5). Note: its update scale differs from Adam, so it needs its own LR (~1e-3 vs Adam's 3e-4). The fp32 momentum buffer here lives in VRAM (~+3GB at 1B params); for the 24GB 4090 a paged or int8-quantized buffer (per the paper) is the deployment path.""" def __init__(self, params, lr=1e-3, momentum=0.9, beta=0.1, weight_decay=0.0, int8=False, paged=False): if not 0.0 <= beta <= 1.0: raise ValueError(f"beta must be in [0,1], got {beta}") if int8 and paged: raise ValueError("choose at most one of PowerStep int8 / paged") # Memory modes for the single momentum buffer (VRAM is the constraint; RAM is cheap): # default -> fp32 buffer in VRAM (fastest). # int8=True -> blockwise-int8 buffer in VRAM (paper's headline; ~1/4 VRAM). # paged=True -> fp32 buffer in pinned CPU RAM (~0 persistent VRAM; spends RAM+PCIe). self._int8 = bool(int8); self._paged = bool(paged) if self._int8: import bitsandbytes.functional as _bnbF self._bnbF = _bnbF super().__init__(params, dict(lr=lr, momentum=momentum, beta=beta, weight_decay=weight_decay)) @torch.no_grad() def step(self, closure=None): loss = None if closure is not None: with torch.enable_grad(): loss = closure() EPS = 1e-12 for group in self.param_groups: lr = group["lr"]; gamma = group["momentum"]; beta = group["beta"]; wd = group["weight_decay"] if self._int8 or self._paged: # Per-tensor path (blockwise-int8 in VRAM, or fp32 buffer in CPU RAM). for p in group["params"]: if p.grad is None: continue g = p.grad st = self.state[p] if self._int8: m = (torch.zeros_like(p, dtype=torch.float32) if "mq" not in st else self._bnbF.dequantize_blockwise(st["mq"], st["mstate"])) m.mul_(gamma).add_(g.float()) u = (m * (m.abs() + EPS).pow(beta - 1.0)).to(p.dtype) st["mq"], st["mstate"] = self._bnbF.quantize_blockwise(m) else: if "m" not in st: st["m"] = torch.zeros(p.shape, dtype=torch.float32, pin_memory=torch.cuda.is_available()) m = st["m"].to(p.device, non_blocking=True) m.mul_(gamma).add_(g.float()) u = (m * (m.abs() + EPS).pow(beta - 1.0)).to(p.dtype) st["m"].copy_(m, non_blocking=True) if wd != 0: p.mul_(1.0 - lr * wd) p.add_(u, alpha=-lr) continue # Fast multi-tensor (foreach) path for the default in-VRAM fp32 buffer: # batches the elementwise update across all params -> few kernel launches, # matching fused optimizers instead of one launch set per parameter. params, grads, ms = [], [], [] for p in group["params"]: if p.grad is None: continue st = self.state[p] if "m" not in st: st["m"] = torch.zeros_like(p, memory_format=torch.preserve_format) params.append(p); grads.append(p.grad); ms.append(st["m"]) if not params: continue # m = gamma*m + g torch._foreach_mul_(ms, gamma) torch._foreach_add_(ms, grads) # u = sign(m)*|m|^beta = m * (|m|+eps)^(beta-1) (avoids a separate sign pass) absm = torch._foreach_abs(ms) torch._foreach_add_(absm, EPS) torch._foreach_pow_(absm, beta - 1.0) us = torch._foreach_mul(ms, absm) if wd != 0: torch._foreach_mul_(params, 1.0 - lr * wd) torch._foreach_add_(params, us, alpha=-lr) return loss def make_optimizer(args, core, ar_h, sat_h, lr_core: float, lr_head: float, nat_h=None): groups = _optimizer_param_groups(core, ar_h, sat_h, lr_core, lr_head, nat_h) opt_name = getattr(args, "optimizer", "adamw") if opt_name == "adamw": return torch.optim.AdamW(groups) if opt_name == "powerstep": return PowerStep(groups, momentum=float(getattr(args, "powerstep_momentum", 0.9)), beta=float(getattr(args, "powerstep_beta", 0.1)), weight_decay=float(getattr(args, "weight_decay", 0.0) or 0.0), int8=bool(getattr(args, "powerstep_int8", False)), paged=bool(getattr(args, "powerstep_paged", False))) if opt_name in {"adamw8bit", "paged_adamw8bit"}: try: import bitsandbytes as bnb except Exception as exc: raise RuntimeError( f"--optimizer {opt_name} requires bitsandbytes. Install it in the training env first." ) from exc if opt_name == "paged_adamw8bit": return bnb.optim.PagedAdamW8bit(groups) return bnb.optim.AdamW8bit(groups) raise ValueError(f"unknown optimizer: {opt_name}") def _oom_backoff_state_path(args) -> pathlib.Path: configured = str(getattr(args, "oom_memory_path", "") or "").strip() if configured: return pathlib.Path(configured).expanduser() return pathlib.Path(args.save_dir) / "oom_backoff_state.json" def _oom_backoff_now() -> str: return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) def _oom_backoff_cuda_info() -> Dict[str, Any]: info: Dict[str, Any] = {"device": str(DEV), "gpu_name": "cpu", "gpu_total_gb": 0.0} if DEV.type == "cuda": try: prop = torch.cuda.get_device_properties(DEV) info["gpu_name"] = str(prop.name) info["gpu_total_gb"] = round(float(prop.total_memory) / (1024 ** 3), 3) except Exception: pass return info def _oom_backoff_signature(args, block: int) -> Dict[str, Any]: gpu = _oom_backoff_cuda_info() return { "preset": str(getattr(args, "preset", "")), "block": int(block), "amp": bool(getattr(args, "amp", False)), "optimizer": str(getattr(args, "optimizer", "")), "attn_backend": str(getattr(args, "attn_backend", "")), "grad_checkpoint": bool(getattr(args, "grad_checkpoint", False)), "dblock": bool(getattr(args, "dblock", False)), "dblock_blocks": int(getattr(args, "dblock_blocks", 0) or 0), "dblock_ar_prob": float(getattr(args, "dblock_ar_prob", 0.0) or 0.0), "dblock_sat_prob": float(getattr(args, "dblock_sat_prob", 0.0) or 0.0), "dblock_nat_prob": float(getattr(args, "dblock_nat_prob", 0.0) or 0.0), "sat_every": int(getattr(args, "sat_every", 0) or 0), "nat_every": int(getattr(args, "nat_every", 0) or 0), "oom_auto_backoff": bool(getattr(args, "oom_auto_backoff", False)), "ckpt_codec": str(getattr(args, "ckpt_codec", "") or ""), "delta_codec": str(getattr(args, "delta_codec", "") or ""), "dblock_blocks": int(getattr(args, "dblock_blocks", 0) or 0), "dblock_checkpoint_stride": int(getattr(args, "dblock_checkpoint_stride", 1) or 0), "dblock_checkpoint_skip_tail": int(getattr(args, "dblock_checkpoint_skip_tail", 0) or 0), "dblock_activation_offload": bool(getattr(args, "dblock_activation_offload", False)), "dblock_objective_mode": str(getattr(args, "dblock_objective_mode", "")), "ar_only": bool(getattr(args, "ar_only", False)), "sat_every": int(getattr(args, "sat_every", 0) or 0), "nat_every": int(getattr(args, "nat_every", 0) or 0), "nat_max_tokens": int(getattr(args, "nat_max_tokens", 0) or 0), "moe_ffn": bool(getattr(args, "moe_ffn", False)), "moe_experts": int(getattr(args, "moe_experts", 0) or 0), "moe_top_k": int(getattr(args, "moe_top_k", 0) or 0), "gpu_name": gpu.get("gpu_name", "unknown"), "gpu_total_gb": gpu.get("gpu_total_gb", 0.0), } def _oom_backoff_key(signature: Dict[str, Any]) -> str: raw = json.dumps(signature, sort_keys=True, separators=(",", ":")).encode("utf-8") return hashlib.sha256(raw).hexdigest()[:20] def _oom_backoff_load(path: pathlib.Path) -> Dict[str, Any]: try: if path.exists(): data = json.loads(path.read_text()) if isinstance(data, dict): data.setdefault("schema", "agillm.oom_backoff.v1") data.setdefault("entries", {}) return data except Exception as exc: print(f"[oom-backoff] warning: failed to read {path}: {exc}", flush=True) return {"schema": "agillm.oom_backoff.v1", "entries": {}} def _oom_backoff_save(path: pathlib.Path, state: Dict[str, Any]) -> None: try: path.parent.mkdir(parents=True, exist_ok=True) state["updated_utc"] = _oom_backoff_now() tmp = path.with_suffix(path.suffix + ".tmp") tmp.write_text(json.dumps(state, indent=2, sort_keys=True) + "\n") tmp.replace(path) except Exception as exc: print(f"[oom-backoff] warning: failed to write {path}: {exc}", flush=True) def _oom_backoff_entry(state: Dict[str, Any], key: str, signature: Dict[str, Any]) -> Dict[str, Any]: entries = state.setdefault("entries", {}) entry = entries.get(key) if not isinstance(entry, dict): entry = {} entries[key] = entry entry["signature"] = signature entry.setdefault("successes", 0) entry.setdefault("ooms", 0) entry.setdefault("events", []) return entry def _oom_backoff_features(signature: Dict[str, Any], batch: int, block: int) -> List[float]: total_gb = float(signature.get("gpu_total_gb", 0.0) or 0.0) return [ min(2.0, max(0.0, float(batch) / 128.0)), min(2.0, max(0.0, float(block) / 4096.0)), min(2.0, max(0.0, total_gb / 80.0)), 1.0 if signature.get("dblock") else 0.0, min(2.0, max(0.0, float(signature.get("dblock_blocks", 0) or 0) / 32.0)), min(2.0, max(0.0, float(signature.get("dblock_checkpoint_stride", 1) or 0) / 8.0)), 1.0 if signature.get("amp") else 0.0, 1.0 if "8bit" in str(signature.get("optimizer", "")) else 0.0, 1.0 / max(1.0, float(signature.get("sat_every", 1) or 1)), 1.0 / max(1.0, float(signature.get("nat_every", 1) or 1)), ] def _oom_mlp_init(entry: Dict[str, Any], key: str, n_features: int) -> Dict[str, Any]: mlp = entry.get("mlp") if isinstance(mlp, dict) and len(mlp.get("w1", [])) == 8: return mlp seed = int(hashlib.sha256(("oom-mlp:" + key).encode("utf-8")).hexdigest()[:8], 16) rng = random.Random(seed) hidden = 8 mlp = { "w1": [[rng.uniform(-0.05, 0.05) for _ in range(n_features)] for _ in range(hidden)], "b1": [0.0 for _ in range(hidden)], "w2": [rng.uniform(-0.05, 0.05) for _ in range(hidden)], "b2": 0.0, "seen": 0, } entry["mlp"] = mlp return mlp def _oom_mlp_forward(mlp: Dict[str, Any], features: List[float]) -> Tuple[float, List[float]]: hidden: List[float] = [] for row, bias in zip(mlp.get("w1", []), mlp.get("b1", [])): z = float(bias) + sum(float(w) * float(x) for w, x in zip(row, features)) hidden.append(math.tanh(z)) logit = float(mlp.get("b2", 0.0)) + sum(float(w) * h for w, h in zip(mlp.get("w2", []), hidden)) logit = max(-30.0, min(30.0, logit)) prob = 1.0 / (1.0 + math.exp(-logit)) return prob, hidden def _oom_mlp_update(entry: Dict[str, Any], key: str, signature: Dict[str, Any], batch: int, block: int, label: int) -> float: features = _oom_backoff_features(signature, batch, block) mlp = _oom_mlp_init(entry, key, len(features)) prob, hidden = _oom_mlp_forward(mlp, features) lr = 0.04 dlogit = prob - float(label) old_w2 = [float(w) for w in mlp["w2"]] for j, h in enumerate(hidden): mlp["w2"][j] = float(mlp["w2"][j]) - lr * dlogit * h mlp["b2"] = float(mlp.get("b2", 0.0)) - lr * dlogit for j, h in enumerate(hidden): dh = dlogit * old_w2[j] * (1.0 - h * h) for i, x in enumerate(features): mlp["w1"][j][i] = float(mlp["w1"][j][i]) - lr * dh * float(x) mlp["b1"][j] = float(mlp["b1"][j]) - lr * dh mlp["seen"] = int(mlp.get("seen", 0) or 0) + 1 return prob def _oom_backoff_peak_gb() -> float: if DEV.type != "cuda": return 0.0 try: return round(float(torch.cuda.max_memory_allocated()) / (1024 ** 3), 4) except Exception: return 0.0 def _oom_backoff_start(args, phase_name: str, block: int, requested_batch: int) -> Tuple[int, Dict[str, Any], pathlib.Path, str, Dict[str, Any]]: path = _oom_backoff_state_path(args) state = _oom_backoff_load(path) signature = _oom_backoff_signature(args, block) key = _oom_backoff_key(signature) entry = _oom_backoff_entry(state, key, signature) batch = int(requested_batch) reasons: List[str] = [] safe = int(entry.get("safe_batch", 0) or 0) oom = int(entry.get("oom_batch", 0) or 0) if oom > 0 and batch >= oom: cap = max(1, int(math.floor(oom * float(getattr(args, "oom_backoff_safety", 0.92) or 0.92)))) if safe > 0 and safe < oom: cap = min(cap, safe) batch = min(batch, cap) reasons.append(f"known OOM at B={oom}") try: threshold = float(getattr(args, "oom_predict_threshold", 0.70) or 0.70) mlp = _oom_mlp_init(entry, key, len(_oom_backoff_features(signature, batch, block))) if int(mlp.get("seen", 0) or 0) >= 6: while batch > 1: prob, _hidden = _oom_mlp_forward(mlp, _oom_backoff_features(signature, batch, block)) if prob < threshold: break nb = max(1, int(math.floor(batch * float(getattr(args, "oom_backoff_safety", 0.92) or 0.92)))) if nb >= batch: nb = batch - 1 reasons.append(f"MLP p_oom={prob:.2f} at B={batch}") batch = nb except Exception as exc: print(f"[oom-backoff] predictor warning: {exc}", flush=True) if batch != requested_batch: print( f"[oom-backoff] {phase_name}: startup cap Batch {requested_batch} -> {batch} " f"({'; '.join(reasons) or 'persistent memory'}) state={path}", flush=True, ) _oom_backoff_save(path, state) return int(batch), state, path, key, signature def _oom_backoff_next_batch(args, entry: Dict[str, Any], current_batch: int) -> int: safe = int(entry.get("safe_batch", 0) or 0) factor = float(getattr(args, "oom_backoff_safety", 0.92) or 0.92) candidate = max(1, int(math.floor(current_batch * factor))) if candidate >= current_batch: candidate = current_batch - 1 if safe > 0 and safe < current_batch: candidate = min(candidate, safe) return max(1, int(candidate)) def _oom_backoff_record( args, state: Dict[str, Any], path: pathlib.Path, key: str, signature: Dict[str, Any], *, outcome: str, batch: int, block: int, step: int, phase_name: str, peak_gb: float = 0.0, ) -> Dict[str, Any]: entry = _oom_backoff_entry(state, key, signature) label = 1 if outcome == "oom" else 0 prob = _oom_mlp_update(entry, key, signature, int(batch), int(block), label) event = { "utc": _oom_backoff_now(), "outcome": outcome, "batch": int(batch), "block": int(block), "step": int(step), "phase": phase_name, "peak_gb": float(peak_gb or 0.0), "mlp_p_oom_before": round(float(prob), 4), } events = entry.setdefault("events", []) events.append(event) del events[:-64] if outcome == "oom": entry["ooms"] = int(entry.get("ooms", 0) or 0) + 1 prior = int(entry.get("oom_batch", 0) or 0) entry["oom_batch"] = int(batch) if prior <= 0 else min(prior, int(batch)) entry["last_oom_utc"] = event["utc"] entry["last_oom_peak_gb"] = float(peak_gb or 0.0) else: entry["successes"] = int(entry.get("successes", 0) or 0) + 1 prior = int(entry.get("safe_batch", 0) or 0) entry["safe_batch"] = max(prior, int(batch)) entry["last_safe_utc"] = event["utc"] entry["last_safe_peak_gb"] = float(peak_gb or 0.0) _oom_backoff_save(path, state) return entry def _oom_backoff_enabled(args) -> bool: return bool(getattr(args, "oom_auto_backoff", True)) def _nat_ids_for_training(ids: torch.Tensor, max_tokens: int) -> torch.Tensor: if max_tokens and max_tokens > 0 and ids.size(1) > max_tokens: return ids[:, -max_tokens:] return ids def _nat_span_len(T: int, ratio: float, max_tokens: int = 0) -> int: target = max(1, min(T, int(round(T * max(0.01, min(0.95, float(ratio))))))) hi = min(T, max(target, target * 2)) if max_tokens and max_tokens > 0: hi = min(hi, int(max_tokens)) lo = max(1, min(hi, target // 2 if target > 1 else 1)) return random.randint(lo, hi) def _nat_corruption_mask(ids: torch.Tensor, ratio: float, args) -> torch.Tensor: """Mask schedule for NAT CMLM training. Random single-token holes are easy because nearby clean target tokens leak most of the answer. Inference asks NAT to fill contiguous/all-BLANK future spans. Mix random, contiguous, and right-suffix spans so training matches that use. """ B, T = ids.shape ratio = max(0.05, min(0.95, float(ratio))) span_prob = max(0.0, min(1.0, float(getattr(args, "nat_span_mask_prob", 0.35) or 0.0))) suffix_prob = max(0.0, min(1.0, float(getattr(args, "nat_suffix_mask_prob", 0.20) or 0.0))) max_span = int(getattr(args, "nat_span_max_tokens", 0) or 0) mask = torch.empty((B, T), device=ids.device, dtype=torch.bool) for b in range(B): r = random.random() if r < suffix_prob: row = torch.zeros((T,), device=ids.device, dtype=torch.bool) n = _nat_span_len(T, ratio, max_span) row[-n:] = True elif r < suffix_prob + span_prob: row = torch.zeros((T,), device=ids.device, dtype=torch.bool) n = _nat_span_len(T, ratio, max_span) start = random.randint(0, max(0, T - n)) row[start:start + n] = True else: row = torch.rand((T,), device=ids.device) < ratio if not bool(row.any()): row[random.randrange(T)] = True mask[b] = row return mask def _train_phase( args, phase_name: str, core, ar_h, sat_h, nat_h, opt, scaler, start_step, seen_tok, resume_wall_time, cfg, source, steps, block_size, batch_size, chat_cfg: dict, max_ckpts: int, target_tokens_override: Optional[int] = None, tie_weights: bool = False, streaming: bool = True, lineage: Optional[Dict[str, Any]] = None, provenance_cache: Optional[Dict[str, Any]] = None ): BLOCK = block_size BATCH_REQUESTED = int(batch_size) BATCH = BATCH_REQUESTED oom_state: Dict[str, Any] = {} oom_state_path = pathlib.Path(args.save_dir) / "oom_backoff_state.json" oom_key = "" oom_signature: Dict[str, Any] = {} oom_good_steps = 0 if _oom_backoff_enabled(args): BATCH, oom_state, oom_state_path, oom_key, oom_signature = _oom_backoff_start(args, phase_name, BLOCK, BATCH) if lineage is None: lineage = {} if target_tokens_override is not None: target_tokens = target_tokens_override else: ratio = _target_token_ratio(args) param_count = _count_enabled_params(core, ar_h, sat_h, nat_h) target_tokens = int(ratio * param_count) print(f"[{phase_name}] token_param_ratio={ratio:g} param_count={param_count:,} target_tokens={target_tokens:,}") if steps: phase_target_tokens = steps * BLOCK * BATCH total_tokens_needed = seen_tok + phase_target_tokens else: total_tokens_needed = target_tokens if total_tokens_needed <= seen_tok: print(f"[{phase_name}] target {total_tokens_needed} already reached.") return start_step, seen_tok, resume_wall_time data_seed = int(getattr(args, "data_seed", 42)) if data_seed < 0: # Streaming restarts from the dataset head with a fixed shuffle seed, so every # restart re-trains the same early data. Derive a per-resume seed instead: # deterministic for a given checkpoint, different across restarts. data_seed = 42 + int(start_step) print(f"[data] per-restart shuffle seed {data_seed} (derived from resume step)", flush=True) effective_source = get_hot_datasets(source) val_requested = str(getattr(args, "val_source", "") or "").strip() if val_requested and _looks_numeracy_only_sources(val_requested) and not _looks_numeracy_only_sources(effective_source): val_effective = effective_source else: val_effective = val_requested or effective_source dataset_meta = _dataset_provenance( phase_name, source, effective_source, args, use_hot_config=True, val_requested=val_requested, val_effective=val_effective, ) print( f"[dataset-policy] phase={phase_name} sources={dataset_meta['source_count']} " f"language_mix={dataset_meta['has_language_mix']} numeracy={dataset_meta['has_numeracy']}", flush=True, ) val_batches = _build_val_set(effective_source, chat_cfg, args, BLOCK) last_val_mono = time.monotonic() stream = token_stream( effective_source, total_tokens_needed, seed=data_seed, chat=chat_cfg.get("chat", False), chat_messages_key=chat_cfg.get("key", "messages"), sft_add_generation_prompt=chat_cfg.get("gen_prompt", False), dataset_field_text=chat_cfg.get("text_field", "text"), streaming=streaming, use_hot_config=False, ) ce_tok = nn.CrossEntropyLoss(label_smoothing=0.1) ce_gate = nn.CrossEntropyLoss() ctc = nn.CTCLoss(blank=BLANK, zero_infinity=True) pbar = SafeProgress(total=total_tokens_needed, initial=seen_tok, unit="tok", initial_step=start_step) if start_step or seen_tok: print(f"[{phase_name}] resume counters: step={int(start_step)} seen_tok={int(seen_tok)} current_B={int(BATCH)} current_L={int(BLOCK)}", flush=True) grow_plan = _parse_grow_plan(args.grow_plan) if args.auto_grow else [] buf: list[int] = [] batch_accum: list[list[int]] = [] step = start_step steps_since_last_grow = 0 oom_retries = 0 MAX_OOM_RETRIES = int(getattr(args, "oom_retries_before_backoff", 0) or 0) now_wall = time.time() last_save_mono = time.monotonic() - (now_wall - (resume_wall_time or now_wall)) last_delta_step = start_step last_delta_mono = last_save_mono last_heartbeat_mono = time.monotonic() _disk_hygiene(pathlib.Path(args.save_dir), phase_name, args, reason="startup") # Derive origin tag from warmstart path for checkpoint naming _ws_path = getattr(args, "warmstart_from", None) or getattr(args, "resume", None) or "" _ws_m = re.search(r"step(\d+)", pathlib.Path(_ws_path).name) if _ws_path else None _origin_tag = f"_from{int(_ws_m.group(1)):08d}" if _ws_m else "" _role_tag = f"_{getattr(args, 'ckpt_role', '').strip()}" if getattr(args, "ckpt_role", "").strip() else "" if val_batches: _run_validation(core, ar_h, val_batches, args, step) print(f"[{phase_name}] Starting. Goal: {total_tokens_needed:,} tokens. Batch={BATCH}, Block={BLOCK}") print( f"[{phase_name}] AR_ONLY={args.ar_only}, SAT_EVERY={args.sat_every}, " f"NAT_EVERY={args.nat_every}, TIE_WEIGHTS={tie_weights}, STREAMING={streaming}" ) _flush_flag = [False] _terminate_after_flush = [False] def _signal_name(signum): try: return signal.Signals(signum).name except Exception: return str(signum) def _on_flush_signal(signum, frame): _flush_flag[0] = True print(f"\n[{phase_name}] flush signal received ({_signal_name(signum)}); will checkpoint at next step") def _on_terminate_signal(signum, frame): _flush_flag[0] = True _terminate_after_flush[0] = True print(f"\n[{phase_name}] {_signal_name(signum)} received; will checkpoint at next step and exit cleanly") try: signal.signal(signal.SIGUSR1, _on_flush_signal) for _term_sig in (getattr(signal, "SIGTERM", None), getattr(signal, "SIGINT", None), getattr(signal, "SIGHUP", None)): if _term_sig is not None: signal.signal(_term_sig, _on_terminate_signal) print(f"[{phase_name}] on-demand flush ready: kill -USR1 {os.getpid()} or touch {pathlib.Path(args.save_dir) / 'FLUSH_NOW'}") print(f"[{phase_name}] graceful termination flush ready: SIGTERM/SIGINT/SIGHUP will save a checkpoint then exit") except (ValueError, OSError): pass _DBS = _dblock_init(core, args) if getattr(args,'dblock',False) else None if DEV.type == "cuda": try: torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() print( f"[vram] training-start cache cleared: " f"alloc={torch.cuda.memory_allocated() / (1024**3):.2f}GB " f"reserved={torch.cuda.memory_reserved() / (1024**3):.2f}GB " f"structured_masks={use_structured_masks(args)}", flush=True, ) except Exception: pass while seen_tok < total_tokens_needed: _profile_batch = _DBS is not None and int(getattr(args, "profile_steps", 0) or 0) > 0 and int(_DBS.get("profile_n", 0)) < int(getattr(args, "profile_steps", 0) or 0) _data_t = time.perf_counter() if _profile_batch else None try: while len(buf) < BLOCK: buf.append(next(stream)) except StopIteration: break if _profile_batch: try: _db_prof = _agillm41_sys.modules[__name__] _db_prof._profile_add(_DBS, "data_stream", time.perf_counter() - _data_t) except Exception: pass seq = buf[:BLOCK] buf = buf[BLOCK:] batch_accum.append(seq) if len(batch_accum) < BATCH: continue _tensor_t = time.perf_counter() if _profile_batch else None ids = torch.tensor(batch_accum, device=DEV) if _profile_batch: if DEV.type == "cuda": try: torch.cuda.synchronize() except Exception: pass try: _db_prof = _agillm41_sys.modules[__name__] _db_prof._profile_add(_DBS, "tensor", time.perf_counter() - _tensor_t) except Exception: pass batch_accum = [] tgt_ar = ids.clone() try: if getattr(args, "dblock", False): loss_value = _dblock_step(core, ar_h, sat_h, nat_h, opt, scaler, args, ids, _DBS) _prov_loss = float(loss_value) else: with amp(args.amp): h_ar = core(ids, causal_mask(ids.size(1), structured=use_structured_masks(args))) logits_ar = ar_h(h_ar)[:, :-1] loss_ar = ce_tok(logits_ar.reshape(-1, VOCAB), tgt_ar[:, 1:].reshape(-1)) loss_value = float(loss_ar.detach().item()) _aux = _collect_moe_aux(core, getattr(args,'moe_aux_coef',0.0), getattr(args,'moe_z_coef',0.0)) if torch.is_tensor(_aux): loss_ar = loss_ar + _aux.to(loss_ar.dtype) scaler.scale(loss_ar).backward() del h_ar, logits_ar, loss_ar do_sat = (not args.ar_only) and (args.sat_every <= 1 or ((step + 1) % args.sat_every == 0)) if do_sat: # Same AR+SAT objective as a summed loss, but sequential backward keeps # only one core-forward activation graph live at a time on 24GB cards. with amp(args.amp): h_sat = core(ids, sat_mask(ids.size(1), structured=use_structured_masks(args))) sat_ctx = h_sat[:, :-SAT_BLOCK] tgt_sat = ids[:, SAT_BLOCK:] if sat_ctx.size(1) == 0 or sat_ctx.size(1) != tgt_sat.size(1): sat_ctx = h_sat[:, :-1] tgt_sat = ids[:, 1:] logits_sat = sat_h.proj(sat_ctx) loss_sat = ce_tok(logits_sat.reshape(-1, VOCAB), tgt_sat.reshape(-1)) if sat_h.gate is not None: sat_gate_ctx = sat_ctx[:, ::SAT_BLOCK] gate_targets = torch.ones( sat_gate_ctx.numel() // sat_gate_ctx.size(-1), device=DEV, dtype=torch.long ) loss_sat += EMIT_LAMBDA * ce_gate( sat_h.gate(sat_gate_ctx.reshape(-1, sat_gate_ctx.size(-1))), gate_targets ) loss_value += float(loss_sat.detach().item()) _aux = _collect_moe_aux(core, getattr(args,'moe_aux_coef',0.0), getattr(args,'moe_z_coef',0.0)) if torch.is_tensor(_aux): loss_sat = loss_sat + _aux.to(loss_sat.dtype) scaler.scale(loss_sat).backward() del h_sat, logits_sat, loss_sat do_nat = ( nat_h is not None and (not args.ar_only) and args.nat_every > 0 and (args.nat_every <= 1 or ((step + 1) % args.nat_every == 0)) ) if do_nat: nat_ids = _nat_ids_for_training(ids, args.nat_max_tokens) with amp(args.amp): # Mask-predict (CMLM) objective: corrupt a fraction of positions # with BLANK and reconstruct them from surrounding context. The # old CTC objective fed the clean target as input, so the head # only learned to copy and collapsed at inference on all-BLANK # input. This conditions on real context and cannot collapse. nat_in = nat_ids.clone() ratio = min(max(float(args.nat_mask_ratio), 0.05), 0.95) mask = _nat_corruption_mask(nat_ids, ratio, args) nat_in[mask] = BLANK h_nat = core(nat_in, None) logits_nat = nat_h(h_nat) loss_nat = F.cross_entropy(logits_nat[mask].float(), nat_ids[mask]) loss_nat = float(args.nat_loss_weight) * loss_nat loss_value += float(loss_nat.detach().item()) _aux = _collect_moe_aux(core, getattr(args,'moe_aux_coef',0.0), getattr(args,'moe_z_coef',0.0)) if torch.is_tensor(_aux): loss_nat = loss_nat + _aux.to(loss_nat.dtype) scaler.scale(loss_nat).backward() del nat_ids, nat_in, mask, h_nat, logits_nat, loss_nat _prov_loss = float(loss_value) scaler.unscale_(opt) nn.utils.clip_grad_norm_([p for group in opt.param_groups for p in group["params"]], 1.0) scaler.step(opt) scaler.update() opt.zero_grad(set_to_none=True) except RuntimeError as e: msg = str(e).lower() if "out of memory" in msg or "cuda error" in msg: batch_accum = [] try: del ids, tgt_ar except Exception: pass opt.zero_grad(set_to_none=True) scaler = GradScaler(enabled=(args.amp and _needs_grad_scaler())) peak_gb = _oom_backoff_peak_gb() if DEV.type == "cuda": try: torch.cuda.empty_cache() except Exception: pass try: torch.cuda.synchronize() except Exception: pass if _oom_backoff_enabled(args): _oom_backoff_record(args, oom_state, oom_state_path, oom_key, oom_signature, outcome="oom", batch=BATCH, block=BLOCK, step=step, phase_name=phase_name, peak_gb=peak_gb) oom_retries += 1 if oom_retries <= MAX_OOM_RETRIES: print(f"\n[{phase_name} OOM] Retry {oom_retries}/{MAX_OOM_RETRIES} at Batch={BATCH}, clearing VRAM...") time.sleep(2) continue oom_retries = 0 if BATCH > 1: entry = _oom_backoff_entry(oom_state, oom_key, oom_signature) if _oom_backoff_enabled(args) else {} _nb = _oom_backoff_next_batch(args, entry, BATCH) if _oom_backoff_enabled(args) else max(1, int(BATCH * 0.85)) if _nb >= BATCH: _nb = BATCH - 1 print(f"\n[{phase_name} OOM] Reducing Batch: {BATCH} -> {_nb} (persistent learned backoff, state={oom_state_path})") BATCH = _nb oom_good_steps = 0 time.sleep(2) else: new_block = max(128, int(BLOCK * 0.8)) new_block = max(128, (new_block // 128) * 128) if new_block >= BLOCK: new_block = max(128, BLOCK - 128) print(f"\n[{phase_name} OOM] Reducing Block: {BLOCK} -> {new_block}") BLOCK = new_block oom_good_steps = 0 if _oom_backoff_enabled(args): BATCH, oom_state, oom_state_path, oom_key, oom_signature = _oom_backoff_start(args, phase_name, BLOCK, BATCH) time.sleep(2) steps_since_last_grow = 0 continue raise step += 1 # Periodic tokenizer spot-check: verify training data has spaces if step % 1000 == 0: try: sample_text = tok.decode(ids[0][:50].tolist(), skip_special_tokens=True) if len(sample_text) > 20 and " " not in sample_text: print(f"\n[tokenizer] ALERT step {step}: decoded batch has NO SPACES!") print(f" Sample: {repr(sample_text[:80])}") print(" Check transformers version!") except Exception: pass oom_retries = 0 if _oom_backoff_enabled(args): oom_good_steps += 1 good_every = max(1, int(getattr(args, "oom_warmup_good_steps", 16) or 16)) if oom_good_steps in (1, good_every) or (oom_good_steps % max(1, good_every * 4) == 0): _oom_backoff_record(args, oom_state, oom_state_path, oom_key, oom_signature, outcome="success", batch=BATCH, block=BLOCK, step=step, phase_name=phase_name, peak_gb=_oom_backoff_peak_gb()) toks_processed = BLOCK * BATCH seen_tok += toks_processed pbar.set_postfix(loss=f"{loss_value:.3f}", B=BATCH, L=BLOCK) pbar.update(toks_processed) async_every = int(getattr(args, "async_update_every_steps", 0) or 0) if async_every > 0 and (step % async_every) == 0: _hf_fed_log_side_updates(*_apply_async_side_updates(core, cfg, args, step), step) empty_cache_every = int(getattr(args, "empty_cache_every_steps", 0) or 0) if DEV.type == "cuda" and empty_cache_every > 0 and (step % empty_cache_every) == 0: try: torch.cuda.empty_cache() except Exception: pass heartbeat_every = int(getattr(args, "heartbeat_every_sec", 300) or 0) now_mono = time.monotonic() if heartbeat_every > 0 and now_mono - last_heartbeat_mono >= heartbeat_every: mem = "" if DEV.type == "cuda": try: mem = ( f" gpu_alloc={torch.cuda.memory_allocated() / (1024**3):.2f}GB" f" gpu_reserved={torch.cuda.memory_reserved() / (1024**3):.2f}GB" f" gpu_peak={torch.cuda.max_memory_allocated() / (1024**3):.2f}GB" ) except Exception: mem = "" try: heartbeat_payload = { "schema": "agillm.run_state.v1", "model": "AGILLM4.3", "phase": "training", "trainer_phase": phase_name, "pid": int(os.getpid()), "step": int(step), "seen_tok": int(seen_tok), "loss": float(loss_value), "batch_size": int(BATCH), "requested_batch_size": int(BATCH_REQUESTED), "block": int(BLOCK), "oom_backoff": { "enabled": bool(_oom_backoff_enabled(args)), "state_path": str(oom_state_path), "key": str(oom_key), }, "dblock": bool(getattr(args, "dblock", False)), "dblock_blocks": int(getattr(args, "dblock_blocks", 0) or 0), "dblock_ar_prob": float(getattr(args, "dblock_ar_prob", 0.0) or 0.0), "dblock_sat_prob": float(getattr(args, "dblock_sat_prob", 0.0) or 0.0), "dblock_nat_prob": float(getattr(args, "dblock_nat_prob", 0.0) or 0.0), "sat_every": int(getattr(args, "sat_every", 0) or 0), "nat_every": int(getattr(args, "nat_every", 0) or 0), "oom_auto_backoff": bool(getattr(args, "oom_auto_backoff", False)), "ckpt_codec": str(getattr(args, "ckpt_codec", "") or ""), "delta_codec": str(getattr(args, "delta_codec", "") or ""), "structured_masks": bool(use_structured_masks(args)), "device": str(DEV), "save_dir": str(args.save_dir), "dataset_provenance": dataset_meta, "warmstart": lineage, "warmstart_source_path": lineage.get("source_path", ""), "warmstart_kind": lineage.get("warmstart_kind", ""), "warmstart_base_step": int(lineage.get("warmstart_base_step", 0) or 0), "global_origin_step": int(lineage.get("global_origin_step", 0) or 0), "effective_global_step": int((int(lineage.get("global_origin_step", 0) or 0) + int(step)) if int(lineage.get("global_origin_step", 0) or 0) > 0 else int(step)), "warmstart_base_seen_tok": int(lineage.get("warmstart_base_seen_tok", 0) or 0), "global_origin_seen_tok": int(lineage.get("global_origin_seen_tok", 0) or 0), "effective_seen_tok": int(int(lineage.get("global_origin_seen_tok", 0) or 0) + int(seen_tok)), "updated_utc": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), } if DEV.type == "cuda": try: heartbeat_payload["gpu"] = { "allocated_gb": round(torch.cuda.memory_allocated() / (1024**3), 4), "reserved_gb": round(torch.cuda.memory_reserved() / (1024**3), 4), "peak_allocated_gb": round(torch.cuda.max_memory_allocated() / (1024**3), 4), } except Exception: pass hb_path = pathlib.Path(args.save_dir) / "run_state.json" hb_tmp = hb_path.with_suffix(".json.tmp") hb_tmp.write_text(json.dumps(heartbeat_payload, sort_keys=True) + "\n") hb_tmp.replace(hb_path) top_path = pathlib.Path(args.save_dir).parent / "agillm43_run_state.json" merged = {} if top_path.exists(): try: merged = json.loads(top_path.read_text()) except Exception: merged = {} if isinstance(merged, dict): merged.update(heartbeat_payload) merged["phase"] = "training" merged["destructive_actions_allowed"] = False top_tmp = top_path.with_suffix(".json.tmp") top_tmp.write_text(json.dumps(merged, indent=2, sort_keys=True) + "\n") top_tmp.replace(top_path) except Exception as exc: print(f"[heartbeat-json] warning: {exc}", flush=True) print( f"[heartbeat] phase={phase_name} pid={os.getpid()} step={step} " f"seen_tok={seen_tok} loss={loss_value:.3f} B={BATCH} L={BLOCK} " f"dblock={bool(getattr(args, 'dblock', False))} structured_masks={use_structured_masks(args)}{mem}", flush=True, ) last_heartbeat_mono = now_mono if val_batches and int(getattr(args, "val_every_sec", 0) or 0) > 0 and \ (time.monotonic() - last_val_mono) >= int(args.val_every_sec): _run_validation(core, ar_h, val_batches, args, step) last_val_mono = time.monotonic() _flush_sentinel = pathlib.Path(args.save_dir) / "FLUSH_NOW" if _flush_flag[0] or _flush_sentinel.exists(): _flush_flag[0] = False try: _flush_sentinel.unlink() except FileNotFoundError: pass _ck_name = f"{phase_name}_step{step:08d}{_origin_tag}{time.strftime('_%Y%m%dT%H%MZ', time.gmtime())}{_role_tag}.pt" _flush_delta() _disk_hygiene(pathlib.Path(args.save_dir), phase_name, args, reason="pre-flush-save") _prune_checkpoints(pathlib.Path(args.save_dir), phase_name, max_ckpts) _prov = _agillm_provenance.collect(args, step=step, seen_tok=seen_tok, loss=_prov_loss, batch_size=BATCH_REQUESTED, block_size=BLOCK, warmstart_source_path=getattr(args, 'warmstart_from', None) or getattr(args, 'resume', None), warmstart_source_provenance=provenance_cache, dataset_provenance=dataset_meta, lane=phase_name or "", _sample_core=core, _sample_ar=ar_h, _sample_sat=sat_h, _sample_tok=tok, _sample_device=DEV) save_ckpt(pathlib.Path(args.save_dir) / _ck_name, core, ar_h, sat_h, nat_h, opt, scaler, meta={"cfg": cfg, "step": step, "seen_tok": seen_tok, "wall_time": time.time(), "tie_weights": tie_weights, "dataset_provenance": dataset_meta}, codec=getattr(args, "ckpt_codec", "zstd3"), provenance=_prov) _prune_checkpoints(pathlib.Path(args.save_dir), phase_name, max_ckpts) last_save_mono = time.monotonic() _prune_deltas(pathlib.Path(args.save_dir), phase_name, args.delta_max_keep) last_delta_step = step last_delta_mono = time.monotonic() print(f"[{phase_name}] ON-DEMAND flush saved {_ck_name} at step {step}") if _terminate_after_flush[0]: setattr(args, "_agillm_terminate_after_flush", True) print(f"[{phase_name}] termination checkpoint complete; stopping training loop cleanly", flush=True) return step, seen_tok, time.time() _save_sec = get_hot_config().get("save_every_sec", args.save_every_sec) try: _save_sec = float(_save_sec) except Exception: _save_sec = args.save_every_sec if _save_sec > 0: now_mono = time.monotonic() if now_mono - last_save_mono >= _save_sec: ck_name = f"{phase_name}_step{step:08d}{_origin_tag}{time.strftime('_%Y%m%dT%H%MZ', time.gmtime())}{_role_tag}.pt" _flush_delta() # wait for any in-flight delta before full save _disk_hygiene(pathlib.Path(args.save_dir), phase_name, args, reason="pre-save") _prune_checkpoints(pathlib.Path(args.save_dir), phase_name, max_ckpts) _prov = _agillm_provenance.collect(args, step=step, seen_tok=seen_tok, loss=_prov_loss, batch_size=BATCH_REQUESTED, block_size=BLOCK, warmstart_source_path=getattr(args, 'warmstart_from', None) or getattr(args, 'resume', None), warmstart_source_provenance=provenance_cache, dataset_provenance=dataset_meta, lane=phase_name or "", _sample_core=core, _sample_ar=ar_h, _sample_sat=sat_h, _sample_tok=tok, _sample_device=DEV) save_ckpt(pathlib.Path(args.save_dir) / ck_name, core, ar_h, sat_h, nat_h, opt, scaler, meta={"cfg": cfg, "step": step, "seen_tok": seen_tok, "wall_time": time.time(), "tie_weights": tie_weights, "dataset_provenance": dataset_meta}, codec=getattr(args, "ckpt_codec", "zstd3"), provenance=_prov) _prune_checkpoints(pathlib.Path(args.save_dir), phase_name, max_ckpts) last_save_mono = now_mono # Prune old deltas after a full save (they're superseded) _prune_deltas(pathlib.Path(args.save_dir), phase_name, args.delta_max_keep) last_delta_step = step # reset delta counter after full save last_delta_mono = now_mono # ── Delta checkpoint (time-based preferred, optional step fallback, weight-only, async) ── hot_cfg = get_hot_config() _delta_steps = hot_cfg.get("delta_every_steps", args.delta_every_steps) try: _delta_steps = int(_delta_steps) except Exception: _delta_steps = args.delta_every_steps _delta_sec = hot_cfg.get("delta_every_sec", args.delta_every_sec) try: _delta_sec = float(_delta_sec) except Exception: _delta_sec = args.delta_every_sec now_mono = time.monotonic() _delta_due_by_steps = _delta_steps > 0 and (step - last_delta_step) >= _delta_steps _delta_due_by_time = _delta_sec > 0 and (now_mono - last_delta_mono) >= _delta_sec if _delta_due_by_steps or _delta_due_by_time: save_root = pathlib.Path(args.save_dir) # AGILLM4 production runs on small rented disks. When keep=1, prune # old deltas before the async writer creates the next multi-GB file. if args.delta_max_keep and args.delta_max_keep > 0: _flush_delta() _prune_delta_files_to_count(save_root, phase_name, args.delta_max_keep - 1) _delta_prov = _agillm_provenance.collect(args, step=step, seen_tok=seen_tok, loss=_prov_loss, batch_size=BATCH_REQUESTED, block_size=BLOCK, warmstart_source_path=getattr(args, 'warmstart_from', None) or getattr(args, 'resume', None), warmstart_source_provenance=provenance_cache, dataset_provenance=dataset_meta, lane=phase_name or "", checkpoint_type="delta") save_delta(core, ar_h, sat_h, nat_h, step, seen_tok, save_root, phase_name, getattr(args, "delta_codec", "zstd3"), provenance=_delta_prov, origin_tag=_origin_tag, dt_tag=time.strftime("_%Y%m%dT%H%MZ", time.gmtime()), role_tag=_role_tag) last_delta_step = step last_delta_mono = now_mono _hf_fed_log_round(step, seen_tok, loss_value, _role_tag, _origin_tag) if args.auto_grow: steps_since_last_grow += 1 if steps_since_last_grow >= args.grow_every_steps: steps_since_last_grow = 0 try: idx = grow_plan.index(BLOCK) if idx + 1 < len(grow_plan): BLOCK = grow_plan[idx + 1] print(f"[{phase_name} Grow] Block -> {BLOCK}") if DEV.type == "cuda": torch.cuda.empty_cache() except ValueError: grow_plan = sorted(set(grow_plan + [BLOCK])) pbar.close() _flush_delta() # ensure any in-flight delta completes before final save if phase_name != "sft": _prov = _agillm_provenance.collect(args, step=step, seen_tok=seen_tok, loss=_prov_loss, batch_size=BATCH_REQUESTED, block_size=BLOCK, warmstart_source_path=getattr(args, 'warmstart_from', None) or getattr(args, 'resume', None), warmstart_source_provenance=provenance_cache, dataset_provenance=dataset_meta, lane=phase_name or "", _sample_core=core, _sample_ar=ar_h, _sample_sat=sat_h, _sample_tok=tok, _sample_device=DEV) save_ckpt(pathlib.Path(args.save_dir) / f"{phase_name}_final.pt", core, ar_h, sat_h, nat_h, opt, scaler, meta={"cfg": cfg, "step": step, "seen_tok": seen_tok, "wall_time": time.time(), "tie_weights": tie_weights, "dataset_provenance": dataset_meta}, codec=getattr(args, "ckpt_codec", "zstd3"), provenance=_prov) else: print("[sft] Skipping duplicate sft_final.pt; final.pt will contain the SFT result.") return step, seen_tok, time.time() # ───────────────────────── Main Orchestrator ───────────────────────── def train(args): if getattr(args, "agillm3_compat", False): args.no_nat_head = True args.nat_every = 0 args.dblock_nat_weight = 0.0 args.dblock_nat_prob = 0.0 args.reinit_nat = False args.seed_nat_from_ar = False print(f"[agillm4.1] legacy compatibility mode: tokenizer={TOKENIZER_ID}, AR+SAT checkpoint schema, NAT disabled") cfg = PRESETS[args.preset].copy() tie_weights = args.tie_weights print_expansion_info(cfg, tie_weights) if not args.fresh: if args.warmstart_from: src_probe = pathlib.Path(args.warmstart_from) elif args.resume: src_probe = pathlib.Path(args.resume) else: src_probe = pathlib.Path(args.save_dir) / "final.pt" prev_cfg = infer_cfg_from_ckpt(src_probe) else: prev_cfg = None if prev_cfg: cfg.update({k: v for k, v in prev_cfg.items() if k in cfg}) if args.x2 and prev_cfg.get("layers"): cfg["layers"] = max(cfg["layers"], prev_cfg["layers"] * 2) if args.rank: cfg["rank"] = args.rank if args.x2 and not prev_cfg: cfg["layers"] *= 2 prev_moe = prev_cfg if isinstance(prev_cfg, dict) else {} if bool(getattr(args, "tie_kv", False)): cfg["tie_kv"] = True requested_moe = bool(getattr(args, "moe_ffn", DEFAULT_MOE_FFN)) if requested_moe or bool(prev_moe.get("moe_ffn", False)): cfg["moe_ffn"] = True cfg["moe_experts"] = int(getattr(args, "moe_experts", DEFAULT_MOE_EXPERTS) if requested_moe else prev_moe.get("moe_experts", DEFAULT_MOE_EXPERTS)) cfg["moe_top_k"] = int(getattr(args, "moe_top_k", DEFAULT_MOE_TOP_K) if requested_moe else prev_moe.get("moe_top_k", DEFAULT_MOE_TOP_K)) cfg["moe_mlp_mult"] = int(getattr(args, "moe_mlp_mult", DEFAULT_MOE_MLP_MULT) if requested_moe else prev_moe.get("moe_mlp_mult", DEFAULT_MOE_MLP_MULT)) cfg["moe_shared_experts"] = int(getattr(args, "moe_shared_experts", 0) if requested_moe else prev_moe.get("moe_shared_experts", 0)) cfg["moe_shared_mlp_mult"] = int(getattr(args, "moe_shared_mlp_mult", 0) if requested_moe else prev_moe.get("moe_shared_mlp_mult", 0)) else: cfg["moe_ffn"] = False use_nat_head = not bool(getattr(args, "no_nat_head", False)) if not use_nat_head: cfg["nat_head"] = False args.nat_every = 0 args.dblock_nat_weight = 0.0 args.dblock_nat_prob = 0.0 print(f"Config: {cfg}") print( "AGILLM4.1 single-file runtime: " f"attn_backend={args.attn_backend} grad_checkpoint={args.grad_checkpoint} " f"sublinear_window={args.sublinear_window} sublinear_stride={args.sublinear_stride} " f"sublinear_max_anchors={args.sublinear_max_anchors} sublinear_chunk={args.sublinear_chunk} " f"sublinear_sinks={args.sublinear_sinks} sublinear_recent_anchors={args.sublinear_recent_anchors} " f"sublinear_pooled_landmarks={args.sublinear_pooled_landmarks} " f"moe_ffn={cfg.get('moe_ffn', False)} moe_experts={cfg.get('moe_experts', 0)} " f"moe_top_k={cfg.get('moe_top_k', 0)} moe_mlp_mult={cfg.get('moe_mlp_mult', 0)}" ) core = Encoder( cfg, tie_weights=tie_weights, attn_backend=args.attn_backend, grad_checkpoint=args.grad_checkpoint, sublinear_window=args.sublinear_window, sublinear_stride=args.sublinear_stride, sublinear_max_anchors=args.sublinear_max_anchors, sublinear_chunk=args.sublinear_chunk, sublinear_sinks=args.sublinear_sinks, sublinear_recent_anchors=args.sublinear_recent_anchors, sublinear_pooled_landmarks=args.sublinear_pooled_landmarks, anchor_memory=getattr(args, "anchor_memory", DEFAULT_ANCHOR_MEMORY), anchor_stride=getattr(args, "anchor_stride", DEFAULT_ANCHOR_STRIDE), anchor_max=getattr(args, "anchor_max", DEFAULT_ANCHOR_MAX), anchor_position=getattr(args, "anchor_position", DEFAULT_ANCHOR_POSITION), ).to(DEV) ar_h = ARHead(cfg["d"], tie_weights=tie_weights, embedding_weight=core.emb.weight if tie_weights else None).to(DEV) sat_h = SATHead(cfg["d"], mode="var", tie_weights=tie_weights, embedding_weight=core.emb.weight if tie_weights else None).to(DEV) nat_h = NATHead(cfg["d"], tie_weights=tie_weights, embedding_weight=core.emb.weight if tie_weights else None).to(DEV) if use_nat_head else None if bool(getattr(args, "dblock_looped", False)): loop_bands = max(1, int(getattr(args, "dblock_blocks", 4) or 4)) core.dblock_loop_embed = nn.Embedding(loop_bands, int(cfg["d"])).to(DEV) nn.init.normal_(core.dblock_loop_embed.weight, mean=0.0, std=0.02) print(f"[dblock-looped] registered loop-index embedding: bands={loop_bands} dim={int(cfg['d'])}", flush=True) total_params = _count_enabled_params(core, ar_h, sat_h, nat_h) print(f"Total parameters: {total_params:,}") if tie_weights: head_names = "AR/SAT/NAT" if nat_h is not None else "AR/SAT" print(f"{Colors.WARN}[weight-tying] Embedding and {head_names} vocab projections share one tensor (VRAM-first){Colors.RESET}") _agillm_provenance_cache = None _agillm_loaded_source_path = "" resume_source_requested = bool(getattr(args, "resume_delta", None) or getattr(args, "resume", None)) # Full resume and resume-delta paths load their exact source below. Avoid an # extra best-guess warm-start here; it double-loads multi-GB checkpoints and # can trip the 32GB Vast container memory limit before training starts. if not args.fresh and (getattr(args, "warmstart_from", None) or not resume_source_requested): src = pathlib.Path(args.warmstart_from) if args.warmstart_from else pathlib.Path(args.save_dir) / "final.pt" src = _resolve_ckpt(src) if src: loaded = _safe_load_any(src, core, key="core") _safe_load_any(src, ar_h, key="ar") _safe_load_any(src, sat_h, key="sat") nat_loaded = _safe_load_any(src, nat_h, key="nat") if nat_h is not None else 0 if nat_h is not None and not nat_loaded: print("[nat] Warm-start source has no NAT head; NAT head initialized fresh") if loaded: print(f"Warm-start loaded from {src}") _agillm_loaded_source_path = str(src) _agillm_provenance_cache = _agillm_provenance.extract(src) else: _agillm_provenance_cache = None if not _agillm_loaded_source_path and (getattr(args, "warmstart_from", None) or getattr(args, "resume", None) or getattr(args, "resume_delta", None)): _agillm_loaded_source_path = str(getattr(args, "warmstart_from", None) or getattr(args, "resume", None) or getattr(args, "resume_delta", None)) _agillm_lineage = _agillm43_lineage_info(_agillm_loaded_source_path, _agillm_provenance_cache, args.save_dir) print( f"[lineage] warmstart_kind={_agillm_lineage.get('warmstart_kind')} " f"source={_agillm_lineage.get('source_path') or 'none'} " f"origin_step={_agillm_lineage.get('global_origin_step', 0)}", flush=True, ) _phase_freeze(core, freeze_core=args.freeze_core, unfreeze_ln=args.unfreeze_ln, train_emb=args.train_emb) opt = make_optimizer(args, core, ar_h, sat_h, args.lr_core, args.lr_head, nat_h) scaler = GradScaler(enabled=(args.amp and _needs_grad_scaler())) start_step, seen_tok, last_wall = 0, 0, None if args.resume_delta and not args.fresh: delta_step, delta_tok = load_delta(pathlib.Path(args.resume_delta), core, ar_h, sat_h, nat_h) start_step, seen_tok, last_wall = delta_step, delta_tok, None print(f"Resumed from DELTA at step {start_step} (optimizer state reset — momentum rebuilds in ~100 steps)") elif args.resume and not args.fresh: start_step, seen_tok, last_wall = load_ckpt(pathlib.Path(args.resume), core, ar_h, sat_h, opt, scaler, nat_h) print(f"Resumed from step {start_step}") if getattr(args, "seed_nat_from_ar", False) and nat_h is not None and ar_h is not None: # Seed the non-autoregressive (NAT) head from the trained AR head ("father"). # Same hidden->vocab projection shape, so NAT starts knowing the token # distribution instead of from random/blank -> faster, no collapse. with torch.no_grad(): nat_h.proj.weight.copy_(ar_h.proj.weight) if nat_h.proj.bias is not None: if getattr(ar_h.proj, "bias", None) is not None: nat_h.proj.bias.copy_(ar_h.proj.bias) else: nat_h.proj.bias.zero_() print("[nat] Seeded NAT head from the AR head ('father') for the mask-predict objective") elif getattr(args, "reinit_nat", False) and nat_h is not None: for _m in nat_h.modules(): if isinstance(_m, nn.Linear): nn.init.normal_(_m.weight, mean=0.0, std=0.02) if _m.bias is not None: nn.init.zeros_(_m.bias) print("[nat] Reinitialized NAT head weights (random) for the mask-predict objective") # torch.compile AFTER loading checkpoint (key names differ) if args.compile: print("[torch.compile] Compiling model...") core = torch.compile(core, mode="reduce-overhead") ar_h = torch.compile(ar_h, mode="reduce-overhead") sat_h = torch.compile(sat_h, mode="reduce-overhead") if nat_h is not None: nat_h = torch.compile(nat_h, mode="reduce-overhead") print("[torch.compile] Done.") step, seen_tok, last_wall = _train_phase( args, "pretrain", core, ar_h, sat_h, nat_h, opt, scaler, start_step, seen_tok, last_wall, cfg, args.source, args.steps, args.block or DEFAULT_BLOCK, args.batch_size or DEFAULT_BATCH, chat_cfg={"chat": args.chat, "key": args.chat_messages_key, "gen_prompt": args.sft_add_generation_prompt, "text_field": args.dataset_field_text}, max_ckpts=args.max_ckpts, target_tokens_override=args.target_tokens, tie_weights=tie_weights, lineage=_agillm_lineage, provenance_cache=_agillm_provenance_cache ) if getattr(args, "_agillm_terminate_after_flush", False): print("[train] graceful termination after checkpoint; skipping final.pt duplicate", flush=True) return if (not args.after_sft_source) and (args.after_sft_steps and args.after_sft_steps > 0): args.after_sft_source = DEFAULT_AFTER_SFT_SOURCES args.after_sft_chat = True if args.after_sft_add_generation_prompt is None: args.after_sft_add_generation_prompt = True if not args.after_sft_block: args.after_sft_block = DEFAULT_AFTER_SFT_BLOCK if args.after_sft_source and args.after_sft_steps and args.after_sft_steps > 0: print("\n[Orchestrator] Starting Post-Pretraining SFT Phase...") _phase_freeze(core, freeze_core=args.after_sft_freeze_core, unfreeze_ln=args.after_sft_unfreeze_ln, train_emb=args.after_sft_train_emb) opt = make_optimizer( args, core, ar_h, sat_h, args.after_sft_lr_core or args.lr_core, args.after_sft_lr_head or args.lr_head, nat_h, ) step, seen_tok, last_wall = _train_phase( args, "sft", core, ar_h, sat_h, nat_h, opt, scaler, step, seen_tok, last_wall, cfg, args.after_sft_source, args.after_sft_steps, args.after_sft_block or DEFAULT_AFTER_SFT_BLOCK, args.batch_size or DEFAULT_BATCH, chat_cfg={ "chat": args.after_sft_chat, "key": args.after_sft_chat_messages_key, "gen_prompt": args.after_sft_add_generation_prompt if args.after_sft_add_generation_prompt is not None else args.sft_add_generation_prompt, "text_field": args.after_sft_dataset_field_text }, max_ckpts=args.max_ckpts, target_tokens_override=None, tie_weights=tie_weights, streaming=True, lineage=_agillm_lineage, provenance_cache=_agillm_provenance_cache ) if getattr(args, "_agillm_terminate_after_flush", False): print("[train] graceful termination after checkpoint; skipping final.pt duplicate", flush=True) return final_effective_source = get_hot_datasets(args.source) final_dataset_meta = _dataset_provenance("final", args.source, final_effective_source, args) _prov = _agillm_provenance.collect(args, step=step, seen_tok=seen_tok, loss=0.0, batch_size=int(args.batch_size or DEFAULT_BATCH), block_size=int(args.block or DEFAULT_BLOCK), warmstart_source_path=getattr(args, 'warmstart_from', None) or getattr(args, 'resume', None), warmstart_source_provenance=_agillm_provenance_cache, dataset_provenance=final_dataset_meta, lane="final", _sample_core=core, _sample_ar=ar_h, _sample_sat=sat_h, _sample_tok=tok, _sample_device=DEV) save_ckpt(pathlib.Path(args.save_dir) / "final.pt", core, ar_h, sat_h, nat_h, opt, scaler, meta={"cfg": cfg, "step": step, "seen_tok": seen_tok, "wall_time": time.time(), "tie_weights": tie_weights, "dataset_provenance": final_dataset_meta}, codec=getattr(args, "ckpt_codec", "zstd3"), provenance=_prov) print("🎉 All Training Complete") # ───────────────────────── Sampling ───────────────────────── def _apply_penalties(logits, ids, n, rep_p, pres_p, freq_p): if ids.numel() == 0: return logits hist = ids[0, -n:].long() if n > 0 else ids[0].long() uniq, counts = torch.unique(hist, return_counts=True) if pres_p or freq_p: logits[..., uniq] -= (pres_p + freq_p * counts.float()) if rep_p != 1.0: sel = logits[..., uniq] logits[..., uniq] = torch.where(sel > 0, sel / rep_p, sel * rep_p) return logits def _suppress_eos(logits, args, force=False): if (force or getattr(args, "ignore_eos", False)) and EOS is not None: logits = logits.clone() logits[..., int(EOS)] = -1e9 return logits def _sample(logits, T, top_k, top_p, min_p, greedy): if greedy: return logits.argmax(-1, keepdim=True) probs = (logits / max(T, 1e-8)).softmax(-1) if top_k: v, i = torch.topk(probs, min(top_k, probs.size(-1))) probs = torch.zeros_like(probs).scatter_(-1, i, v) if top_p < 1.0: s_probs, s_idx = torch.sort(probs, descending=True, dim=-1) probs = torch.zeros_like(probs).scatter_(-1, s_idx, s_probs * (torch.cumsum(s_probs, -1) <= top_p).float()) if min_p > 0: probs[probs < min_p] = 0 if probs.sum() == 0: return logits.argmax(-1, keepdim=True) return probs.div_(probs.sum()).multinomial(1) def _swi_entropy(probs): """Shannon entropy (nats) of a [B, V] distribution, averaged over batch.""" p = probs.clamp_min(1e-12) return float(-(p * p.log()).sum(-1).mean()) def _swi_soft_embed(core, probs, top_k): """Continuous 'thought' = probability-weighted average of token embeddings. The model's next-token belief stays in superposition in hidden space rather than collapsing to one discrete token. Restricting to top-k mass keeps it sharp. """ E = core.emb.weight # [V, d] if top_k and 0 < top_k < probs.size(-1): v, i = torch.topk(probs, top_k, dim=-1) # [B, k] v = v / v.sum(-1, keepdim=True).clamp_min(1e-12) thought = (v.unsqueeze(-1) * E[i]).sum(1) # [B, d] else: thought = probs.to(E.dtype) @ E # [B, d] return thought.unsqueeze(1).to(E.dtype) # [B, 1, d] def _swireasoning_decode(core, ar_h, ids, args, min_new): """Training-free SwiReasoning decode for the AR path. Alternates between two reasoning regimes, gated by next-token entropy: EXPLICIT — sample a real token (model thinks out loud). LATENT — inject a continuous thought embedding and emit NO token; model reasons silently in hidden space (token-efficient). Policy: diffuse / rising entropy → drop into latent to explore in superposition; low / sharply-falling entropy → switch back to explicit to consolidate. --swi_max_switches and --swi_think_budget cap overthinking. """ use_struct = use_structured_masks(args) seq_len = ids.size(1) h, kvs = core(ids, causal_mask(seq_len, structured=use_struct), use_cache=True, total_seq_len=seq_len) mode = "latent" if getattr(args, "swi_start_latent", False) else "explicit" switches = latent_run = think_steps = emitted = 0 prev_H = None n_latent = n_explicit = 0 while emitted < args.max_new and think_steps < args.swi_max_steps: logits_last = ar_h(h)[:, -1].float() probs_raw = (logits_last / max(args.temperature, 1e-8)).softmax(-1) H = _swi_entropy(probs_raw) dH = 0.0 if prev_H is None else (H - prev_H) prev_H = H thinking = think_steps < args.swi_think_budget if thinking and switches < args.swi_max_switches: if mode == "latent": if (H < args.swi_explicit_thresh or dH < -args.swi_eps or latent_run >= args.swi_max_latent): mode, switches, latent_run = "explicit", switches + 1, 0 else: if H > args.swi_latent_thresh and dH > args.swi_eps: mode, switches = "latent", switches + 1 else: mode = "explicit" if mode == "latent": thought = _swi_soft_embed(core, probs_raw, args.swi_topk) seq_len += 1; think_steps += 1; latent_run += 1; n_latent += 1 h, kvs = core(None, None, kv_caches=kvs, use_cache=True, total_seq_len=seq_len, inputs_embeds=thought) continue logits = _apply_penalties(logits_last, ids, args.penalty_last_n, args.repetition_penalty, args.presence_penalty, args.frequency_penalty) logits = _suppress_eos(logits, args, emitted < min_new) nxt = _sample(logits, args.temperature, args.top_k, args.top_p, args.min_p, args.greedy) ids = torch.cat([ids, nxt], 1) emitted += 1; think_steps += 1; n_explicit += 1 if EOS is not None and not getattr(args, "ignore_eos", False) and int(nxt.item()) == int(EOS): break seq_len += 1 h, kvs = core(nxt, None, kv_caches=kvs, use_cache=True, total_seq_len=seq_len) saved = (n_latent / max(1, n_latent + n_explicit)) * 100.0 print(f"[swi] explicit={n_explicit} latent={n_latent} switches={switches} " f"({saved:.0f}% of reasoning steps emitted no token)") return ids def _dblock_block_layers(core, dblock_blocks): L = len(core.blocks) B = max(1, int(dblock_blocks)) per = max(1, L // B) groups = [] for b in range(B): lo = b * per hi = L if b == B - 1 else (b + 1) * per groups.append(list(range(lo, hi))) return groups def _dblock_select_block(sigma, bsig): for b in range(len(bsig) - 1): if bsig[b] <= sigma <= bsig[b + 1]: return b return 0 if sigma < bsig[0] else len(bsig) - 2 def _block_stream_enabled(args) -> bool: return bool(getattr(args, "block_stream", False)) def _block_stream_compute_device(args=None): return DEV def _moe_expert_stream_enabled(args) -> bool: return bool(getattr(args, "moe_expert_stream", False)) def _dtype_from_arg(args, attr: str, flag: str): name = str(getattr(args, attr, "fp32") or "fp32").lower() if name in {"fp32", "float32", "none"}: return None if name in {"fp16", "float16", "half"}: return torch.float16 if name in {"bf16", "bfloat16"}: return torch.bfloat16 raise ValueError(f"unsupported {flag} {name!r}") def _block_stream_dtype(args): return _dtype_from_arg(args, "block_stream_dtype", "--block_stream_dtype") def _infer_dtype(args): return _dtype_from_arg(args, "infer_dtype", "--infer_dtype") def _modules_have_dtype(modules, dtype) -> bool: if dtype is None: return True for module in modules: if module is None: continue for tensor in module.parameters(recurse=True): if tensor.dtype != dtype: return False for tensor in module.buffers(recurse=True): if tensor.dtype != dtype: return False return True def _cast_modules_dtype(modules, dtype) -> bool: if dtype is None or _modules_have_dtype(modules, dtype): return False for module in modules: if module is not None: module.to(dtype=dtype) return True def _block_stream_empty_cache(args) -> bool: return bool(getattr(args, "block_stream_empty_cache", True)) and torch.cuda.is_available() def _block_stream_kv_cache_enabled(args) -> bool: return bool(getattr(args, "block_stream_kv_cache", True)) def _block_stream_cache_pages_mode(args): explicit = getattr(args, "block_stream_cache_pages", None) if explicit is None: return "auto" return "on" if bool(explicit) else "off" def _block_stream_cache_pages_enabled(args) -> bool: effective = getattr(args, "_block_stream_cache_pages_effective", None) if effective is not None: return bool(effective) return _block_stream_cache_pages_mode(args) == "on" def _module_tensor_bytes(mod) -> int: total = 0 for t in list(mod.parameters(recurse=True)) + list(mod.buffers(recurse=True)): total += int(t.numel()) * int(t.element_size()) return total def _configure_block_stream_page_cache(args, core): mode = _block_stream_cache_pages_mode(args) if mode == "off": args._block_stream_cache_pages_effective = False args._block_stream_cache_pages_reason = "explicit-off" return if mode == "on": args._block_stream_cache_pages_effective = True args._block_stream_cache_pages_reason = "explicit-on" return if not torch.cuda.is_available() or DEV.type != "cuda": args._block_stream_cache_pages_effective = False args._block_stream_cache_pages_reason = "auto-no-cuda" return try: device_index = DEV.index if getattr(DEV, "index", None) is not None else torch.cuda.current_device() free, total = torch.cuda.mem_get_info(device_index) except (TypeError, ValueError): free, total = torch.cuda.mem_get_info() page_bytes = sum(_module_tensor_bytes(blk) for blk in core.blocks) allocated = torch.cuda.memory_allocated() reserved = torch.cuda.memory_reserved() reusable = max(0, int(reserved) - int(allocated)) usable = int(free) + int(reusable) # This is an incremental fit check, not total model size. At this point the # embedding, heads, CUDA context, and allocator slabs are already resident; # measured page-cache peak is lower than raw block parameter bytes + safety. effective_page_bytes = int(page_bytes * 0.75) safety = max(128 * 1024 * 1024, int(total * 0.005)) effective_need = effective_page_bytes + int(safety) enabled = int(usable) > int(effective_need) args._block_stream_cache_pages_effective = bool(enabled) args._block_stream_cache_pages_reason = ( f"auto usable={usable/1e9:.2f}GB free={free/1e9:.2f}GB " f"reuse={reusable/1e9:.2f}GB need={effective_need/1e9:.2f}GB raw={page_bytes/1e9:.2f}GB" ) # AGILLM-PARTIAL-PAGE-CACHE 20260703: when the full page set does not fit # (small cards like the 2GB 840M), keep as many leading blocks resident as # the budget allows and stream only the rest. args._block_stream_cached_layers = None if not enabled: n_blocks = max(1, len(core.blocks)) per_block_eff = max(1, int((page_bytes / n_blocks) * 0.75)) budget = int(usable) - int(safety) n_cached = min(n_blocks, max(0, budget // per_block_eff)) if n_cached > 0: args._block_stream_cached_layers = frozenset(range(int(n_cached))) args._block_stream_cache_pages_reason += f" partial={int(n_cached)}/{n_blocks}" def _block_stream_kv_store_device(args): name = str(getattr(args, "block_stream_kv_device", "cuda") or "cuda").lower() if name in {"cuda", "gpu"} and torch.cuda.is_available(): return DEV return torch.device("cpu") def _block_stream_kv_to_device(kv, device): if kv is None or isinstance(kv, KVBuffer): return kv k, v = kv if k.device == device and v.device == device: return kv return (k.to(device, non_blocking=True), v.to(device, non_blocking=True)) def _block_stream_kv_to_store(kv, device): if kv is None or isinstance(kv, KVBuffer): return kv k, v = kv if device.type == "cpu": return (k.detach().to("cpu", non_blocking=True), v.detach().to("cpu", non_blocking=True)) return (k.detach(), v.detach()) def _block_stream_layer_pages(core, args): page_layers = int(getattr(args, "block_stream_page_layers", 1) or 0) if page_layers <= 0: return _dblock_block_layers(core, int(getattr(args, "dblock_blocks", 4) or 4)) page_layers = max(1, page_layers) return [list(range(i, min(i + page_layers, len(core.blocks)))) for i in range(0, len(core.blocks), page_layers)] def _block_stream_release(mod, args): mod.to("cpu") if _block_stream_empty_cache(args): torch.cuda.empty_cache() def _block_stream_load_block(block, device, args): if _moe_expert_stream_enabled(args) and isinstance(getattr(block, "ff", None), MoEFFN): block.ln1.to(device) block.ln2.to(device) block.mha.to(device) block.ff.router.to(device) if block.ff.shared is not None: block.ff.shared.to(device) for expert in block.ff.experts: expert.to("cpu") block.ff.set_expert_stream(True, bool(getattr(args, "moe_expert_stream_empty_cache", True))) return block return block.to(device) def _block_stream_release_block(block, args, layer_idx=None): if _block_stream_cache_pages_enabled(args): return cached = getattr(args, "_block_stream_cached_layers", None) if cached is not None and layer_idx is not None and layer_idx in cached: return if isinstance(getattr(block, "ff", None), MoEFFN): block.ff.set_expert_stream(False, bool(getattr(args, "moe_expert_stream_empty_cache", True))) block.to("cpu") if _block_stream_empty_cache(args): torch.cuda.empty_cache() def _moe_expert_stream_stats(core): loads = 0 tokens = 0 for mod in core.modules(): if isinstance(mod, MoEFFN): st = getattr(mod, "expert_stream_stats", None) or {} loads += int(st.get("loads", 0)) tokens += int(st.get("tokens", 0)) return loads, tokens def _moe_expert_stream_reset_stats(core): for mod in core.modules(): if isinstance(mod, MoEFFN): mod.expert_stream_stats = {"loads": 0, "tokens": 0} def _block_stream_maybe_anchor(core, layer_idx, x, args): if core.anchor is None or layer_idx != core.anchor_position: return x device = _block_stream_compute_device(args) core.anchor.to(device) x, _ = core.anchor(x) _block_stream_release(core.anchor, args) return x @torch.no_grad() def _block_stream_forward(core, ids, mask, args): """Run Encoder.forward while paging blocks through the compute device.""" device = _block_stream_compute_device(args) core.emb.to(device) core.ln.to(device) ids = ids.to(device) x = core.emb(ids) for page in _block_stream_layer_pages(core, args): resident = [_block_stream_load_block(core.blocks[li], device, args) for li in page] try: for li, blk in zip(page, resident): x = _run_block(blk, x, mask, False, args) x = _block_stream_maybe_anchor(core, li, x, args) finally: for li, blk in zip(page, resident): _block_stream_release_block(blk, args, layer_idx=li) return core.ln(x) @torch.no_grad() def _block_stream_forward_cached(core, ids, mask, kv_caches, total_seq_len, args): """Block-stream AR/SAT decode with KV cache. We still page layer weights through the compute device, but avoid recomputing the full prefix for every emitted token. KV tensors can stay on CUDA for speed or be stored on CPU for the lowest resident VRAM. """ device = _block_stream_compute_device(args) kv_store_device = _block_stream_kv_store_device(args) core.emb.to(device) core.ln.to(device) ids = ids.to(device) x = core.emb(ids) new_kvs = [None] * len(core.blocks) for page in _block_stream_layer_pages(core, args): resident = [_block_stream_load_block(core.blocks[li], device, args) for li in page] try: for li, blk in zip(page, resident): kv = kv_caches[li] if kv_caches else None kv = _block_stream_kv_to_device(kv, device) x, kv_out = blk(x, mask, kv, use_cache=True, total_seq_len=total_seq_len) x = _block_stream_maybe_anchor(core, li, x, args) new_kvs[li] = _block_stream_kv_to_store(kv_out, kv_store_device) finally: for li, blk in zip(page, resident): _block_stream_release_block(blk, args, layer_idx=li) return core.ln(x), new_kvs def _edm_denoise_block(core, layers, z, sigma_t, mask, args, block_idx=None): cs, co, ci = _edm_pre(sigma_t) h = ci * z if block_idx is not None and getattr(core, "dblock_loop_embed", None) is not None: h = _dblock_loop_condition(core, h, block_idx, args) if _block_stream_enabled(args): device = _block_stream_compute_device(args) for li in layers: blk = _block_stream_load_block(core.blocks[li], device, args) try: h = _run_block(blk, h, mask, False, args) h = _block_stream_maybe_anchor(core, li, h, args) finally: _block_stream_release_block(blk, args, layer_idx=li) else: for li in layers: h = _run_block(core.blocks[li], h, mask, False, args) return cs * z + co * h @torch.no_grad() def _dblock_euler_hidden(core, ids, args): """DiffusionBlocks EDM Euler block-chain hidden state (faithful reverse ODE), adapted to agillm4.1's causal AR head. --euler_start_sigma tunes context conditioning (SDEdit-style); returns LayerNorm'd hidden [B,T,d].""" import numpy as _np dblock_blocks = int(getattr(args, "dblock_blocks", 4) or 4) steps = max(dblock_blocks, int(getattr(args, "euler_steps", 0) or (dblock_blocks * 2))) bsig = _block_sigmas(dblock_blocks, *_dblock_sigma_config(args)) looped = bool(getattr(args, "dblock_looped", False)) and getattr(core, "dblock_loop_embed", None) is not None if looped: _ll = int(getattr(args, "dblock_loop_layers", 0) or 0) or max(1, len(core.blocks) // max(1, dblock_blocks)) _ll = max(1, min(_ll, len(core.blocks))) _ls = max(0, min(int(getattr(args, "dblock_loop_start", 0) or 0), len(core.blocks) - _ll)) _loop_group = list(range(_ls, _ls + _ll)) groups = [_loop_group for _ in range(dblock_blocks)] else: groups = _dblock_block_layers(core, dblock_blocks) sigma_min = float(bsig[0]) start = float(getattr(args, "euler_start_sigma", 0.0) or 0.0) if start <= 0.0: start = float(bsig[-1]) start = max(start, sigma_min * 2) mask = causal_mask(ids.size(1), structured=use_structured_masks(args)) e = core.emb(ids) lo, hi = math.log(sigma_min), math.log(start) sched = [float(_np.exp(hi + (lo - hi) * (i / steps))) for i in range(steps + 1)] z = e + sched[0] * torch.randn_like(e) with amp(getattr(args, "amp", False)): for i in range(steps): s_cur, s_next = sched[i], sched[i + 1] b = _dblock_select_block(s_cur, bsig) sig_t = torch.full((ids.size(0),), s_cur, device=ids.device, dtype=z.dtype) D = _edm_denoise_block(core, groups[b], z, sig_t, mask, args, block_idx=(b if looped else None)) z = z + ((s_next - s_cur) / s_cur) * (z - D) sig0 = torch.full((ids.size(0),), sigma_min, device=ids.device, dtype=z.dtype) D0 = _edm_denoise_block(core, groups[0], z, sig0, mask, args, block_idx=(0 if looped else None)) return core.ln(D0) @torch.no_grad() def _agillm43_prepare_infer_instance(args): global DEV _requested_device = getattr(args, "device", "auto") _effective_device = _requested_device if _effective_device == "auto": _effective_device = "cuda" if torch.cuda.is_available() else "cpu" if _effective_device == "cuda" and not torch.cuda.is_available(): raise RuntimeError("--device cuda requested but CUDA is not available") DEV = torch.device(_effective_device) if DEV.type == "cpu" and bool(getattr(args, "block_stream", False)): print("[infer] --block_stream requested with --device cpu; disabling block_stream", flush=True) args.block_stream = False print(f"[infer] device={DEV} requested={_requested_device} cuda_available={torch.cuda.is_available()}", flush=True) if DEV.type == "cpu": _cpu_threads = int(getattr(args, "cpu_threads", 0) or 0) if _cpu_threads <= 0: _cpu_threads = max(1, min(16, int(os.cpu_count() or 1))) try: torch.set_num_threads(_cpu_threads) print(f"[infer] cpu_threads={_cpu_threads}", flush=True) except Exception as exc: print(f"[infer] warning: could not set cpu_threads={_cpu_threads}: {exc}", flush=True) _cpu_interop_threads = int(getattr(args, "cpu_interop_threads", 0) or 0) if _cpu_interop_threads > 0: try: torch.set_num_interop_threads(_cpu_interop_threads) print(f"[infer] cpu_interop_threads={_cpu_interop_threads}", flush=True) except Exception as exc: print(f"[infer] warning: could not set cpu_interop_threads={_cpu_interop_threads}: {exc}", flush=True) if args.mode == "ar": if args.temperature is None: args.temperature = 0.7 if args.top_k is None: args.top_k = 0 if args.repetition_penalty is None: args.repetition_penalty = 1.3 if args.presence_penalty is None: args.presence_penalty = 0.0 if args.frequency_penalty is None: args.frequency_penalty = 0.3 if args.penalty_last_n is None: args.penalty_last_n = 128 if args.var is None: args.var = False elif args.mode == "sat": if args.temperature is None: args.temperature = 0.5 if args.top_k is None: args.top_k = 30 if args.repetition_penalty is None: args.repetition_penalty = 2.0 if args.presence_penalty is None: args.presence_penalty = 0.6 if args.frequency_penalty is None: args.frequency_penalty = 1.0 if args.penalty_last_n is None: args.penalty_last_n = 200 if args.var is None: args.var = True else: if args.temperature is None: args.temperature = 0.25 if args.top_k is None: args.top_k = 0 if args.repetition_penalty is None: args.repetition_penalty = 2.0 if args.presence_penalty is None: args.presence_penalty = 0.8 if args.frequency_penalty is None: args.frequency_penalty = 1.2 if args.penalty_last_n is None: args.penalty_last_n = 512 if args.var is None: args.var = False min_new = int(getattr(args, "min_new", 0) or 0) if args.mode == "sat": min_new = max(min_new, SAT_BLOCK) path = _resolve_ckpt(pathlib.Path(args.ckpt)) or pathlib.Path(args.ckpt) _t_stage = time.perf_counter() sd = _agillm43_load_pt(path, map_location="cpu", weights_only=False, skip_keys={"opt", "scaler"}) print(f"[load-profile] checkpoint_load={time.perf_counter() - _t_stage:.1f}s", flush=True) _t_stage = time.perf_counter() # Inference never needs optimizer/scaler state. Drop it before model construction # so block-stream runs keep CPU RAM pressure lower after checkpoint load. if isinstance(sd, dict): sd.pop("opt", None) sd.pop("scaler", None) import gc as _gc _gc.collect() # Restore tokenizer from checkpoint (embedded json preferred; never raises) _restore_tokenizer_from_ckpt(sd, path) print(f"[load-profile] tokenizer_restore={time.perf_counter() - _t_stage:.1f}s", flush=True) # Warn if transformers version changed since checkpoint was saved if "transformers_version" in sd: import transformers as _tf if sd["transformers_version"] != _tf.__version__: print(f"[tokenizer] WARNING: checkpoint saved with transformers={sd['transformers_version']}, now running {_tf.__version__}") # Handle delta checkpoints (weight-only, often no cfg) if sd.get("delta"): cfg, tie_weights, cfg_source = _infer_cfg_from_delta_checkpoint(sd) print("[infer] Delta checkpoint detected, cfg_source=%s d=%s layers=%s heads=%s rank=%s tie_kv=%s moe_ffn=%s tie_weights=%s" % ( cfg_source, cfg.get("d"), cfg.get("layers"), cfg.get("heads"), cfg.get("rank"), bool(cfg.get("tie_kv", False)), bool(cfg.get("moe_ffn", False)), bool(tie_weights), ), flush=True) weights = sd.get("weights") or {} # Remap: delta stores under sd["weights"]["core"/"ar"/"sat"/"nat"] sd["core"] = weights["core"] sd["ar"] = weights["ar"] sd["sat"] = weights["sat"] if "nat" in weights: sd["nat"] = weights["nat"] else: cfg = sd["cfg"] tie_weights = sd.get("tie_weights", False) plain_output = ( bool(getattr(args, "plain_output", False)) or bool(getattr(args, "claude_friendly", False)) or not sys.stdout.isatty() ) uk_time = get_uk_time() ckpt_name = path.name if plain_output: print(f"[infer] inference_time={uk_time}") print(f"[infer] checkpoint={ckpt_name}") else: print(f"┌─────────────────────────────────────────────────┐") print(f"│ INFERENCE @ {uk_time:<35s} │") print(f"├─────────────────────────────────────────────────┤") print(f"│ Checkpoint: {ckpt_name:<35s} │") print(f"└─────────────────────────────────────────────────┘") print_expansion_info(cfg, tie_weights, plain=plain_output) _t_stage = time.perf_counter() block_stream = _block_stream_enabled(args) infer_dtype = None if block_stream else _infer_dtype(args) preload_dtype = _block_stream_dtype(args) if block_stream else infer_dtype resident_dtype = (infer_dtype is not None and not block_stream) core_device = torch.device("cpu") if (block_stream or resident_dtype) else DEV old_default_dtype = torch.get_default_dtype() if preload_dtype is not None: torch.set_default_dtype(preload_dtype) try: with _skip_param_init(): core = Encoder( cfg, tie_weights=tie_weights, attn_backend=args.attn_backend, sublinear_window=args.sublinear_window, sublinear_stride=args.sublinear_stride, sublinear_max_anchors=args.sublinear_max_anchors, sublinear_chunk=args.sublinear_chunk, sublinear_sinks=args.sublinear_sinks, sublinear_recent_anchors=args.sublinear_recent_anchors, sublinear_pooled_landmarks=args.sublinear_pooled_landmarks, anchor_memory=getattr(args, "anchor_memory", DEFAULT_ANCHOR_MEMORY), anchor_stride=getattr(args, "anchor_stride", DEFAULT_ANCHOR_STRIDE), anchor_max=getattr(args, "anchor_max", DEFAULT_ANCHOR_MAX), anchor_position=getattr(args, "anchor_position", DEFAULT_ANCHOR_POSITION), ).to(core_device) print(f"[load-profile] encoder_construct={time.perf_counter() - _t_stage:.1f}s", flush=True) _t_stage = time.perf_counter() head_device = torch.device("cpu") if resident_dtype else DEV ar_h = ARHead(cfg["d"], tie_weights=tie_weights, embedding_weight=core.emb.weight if tie_weights else None).to(head_device) sat_head_mlp = bool(sd.get("sat_head_mlp", False) or _sat_head_mlp_from_state(sd)) sat_h = SATHead(cfg["d"], mlp=sat_head_mlp, tie_weights=tie_weights, embedding_weight=core.emb.weight if tie_weights else None).to(head_device) nat_h = NATHead(cfg["d"], tie_weights=tie_weights, embedding_weight=core.emb.weight if tie_weights else None).to(head_device) if ("nat" in sd or args.mode == "nat") else None finally: torch.set_default_dtype(old_default_dtype) _reinit_params_missing_from_state(core, sd["core"] if isinstance(sd.get("core"), dict) else {}) _maybe_register_looped_infer(core, sd, args) print(f"[load-profile] heads_construct={time.perf_counter() - _t_stage:.1f}s", flush=True) _t_stage = time.perf_counter() if preload_dtype is not None: if not _cast_modules_dtype((core, ar_h, sat_h, nat_h), preload_dtype): print(f"[infer] preload_dtype already={str(preload_dtype).replace('torch.', '')}", flush=True) print(f"[load-profile] preload_dtype_cast={time.perf_counter() - _t_stage:.1f}s", flush=True) _t_stage = time.perf_counter() core.load_state_dict(_prepare_core_state_dict_for_load(core, sd["core"])) print(f"[load-profile] core_state_load={time.perf_counter() - _t_stage:.1f}s", flush=True) _t_stage = time.perf_counter() ar_h.load_state_dict(sd["ar"]) _load_infer_head_state(sat_h, sd["sat"], "SATHead") if nat_h is not None: if "nat" not in sd: raise ValueError("NAT inference requested, but this checkpoint has no NAT head") _load_infer_head_state(nat_h, sd["nat"], "NATHead") print(f"[load-profile] head_state_load={time.perf_counter() - _t_stage:.1f}s", flush=True) _t_stage = time.perf_counter() core.eval() ar_h.eval() sat_h.eval() if nat_h is not None: nat_h.eval() if resident_dtype: _cast_modules_dtype((core, ar_h, sat_h, nat_h), infer_dtype) core.to(DEV) ar_h.to(DEV) sat_h.to(DEV) if nat_h is not None: nat_h.to(DEV) print(f"[infer] infer_dtype={str(infer_dtype).replace('torch.', '')} resident=True device={DEV}") if block_stream: stream_dtype = _block_stream_dtype(args) if stream_dtype is not None: _cast_modules_dtype((core, ar_h, sat_h, nat_h), stream_dtype) print(f"[infer] block_stream_dtype={str(stream_dtype).replace('torch.', '')}") core.emb.to(DEV) core.ln.to(DEV) if core.anchor is not None: core.anchor.to("cpu") for blk in core.blocks: blk.to("cpu") if _block_stream_empty_cache(args): torch.cuda.empty_cache() _configure_block_stream_page_cache(args, core) page_desc = "dblock" if int(getattr(args, "block_stream_page_layers", 1) or 0) <= 0 else f"{int(getattr(args, 'block_stream_page_layers', 1))} layer(s)" moe_desc = " moe_expert_stream=True" if _moe_expert_stream_enabled(args) else "" page_cache_reason = getattr(args, "_block_stream_cache_pages_reason", "") page_cache_desc = f" page_cache={_block_stream_cache_pages_enabled(args)}" if page_cache_reason: page_cache_desc += f" ({page_cache_reason})" if _block_stream_kv_cache_enabled(args): kv_desc = f" KV cache=True kv_device={_block_stream_kv_store_device(args)}" else: kv_desc = " KV cache=False full-prefix recompute=True" print(f"[infer] block_stream=True device={DEV} page={page_desc}{moe_desc};{page_cache_desc}{kv_desc}") if _moe_expert_stream_enabled(args): _moe_expert_stream_reset_stats(core) print(f"[load-profile] device_placement={time.perf_counter() - _t_stage:.1f}s", flush=True) total_params = _count_enabled_params(core, ar_h, sat_h, nat_h) if total_params >= 1_000_000_000: param_str = f"{total_params / 1_000_000_000:.2f}B" elif total_params >= 1_000_000: param_str = f"{total_params / 1_000_000:.2f}M" elif total_params >= 1_000: param_str = f"{total_params / 1_000:.2f}K" else: param_str = f"{total_params}" print(f"Model size: {param_str} parameters ({total_params:,})") try: del sd import gc as _gc _gc.collect() except Exception: pass return { "path": path, "cfg": cfg, "tie_weights": tie_weights, "plain_output": plain_output, "block_stream": block_stream, "resident_dtype": resident_dtype, "core": core, "ar_h": ar_h, "sat_h": sat_h, "nat_h": nat_h, } @torch.no_grad() def _agillm43_generate_from_instance(inst, args): global DEV core = inst["core"] ar_h = inst["ar_h"] sat_h = inst["sat_h"] nat_h = inst["nat_h"] block_stream = bool(inst.get("block_stream", False)) resident_dtype = bool(inst.get("resident_dtype", False)) plain_output = ( bool(getattr(args, "plain_output", False)) or bool(getattr(args, "claude_friendly", False)) or not sys.stdout.isatty() ) min_new = int(getattr(args, "min_new", 0) or 0) if args.mode == "sat": min_new = max(min_new, SAT_BLOCK) # AGILLM-STREAM 20260703: machine-readable per-commit markers (plain only). stream = bool(getattr(args, "stream", False)) and plain_output prompt_tokens = tok.encode(args.prompt) prompt_len = len(prompt_tokens) ids = torch.tensor([prompt_tokens], device=DEV) if ids.size(1) == 0: ids = torch.tensor([[EOS]], device=DEV) prompt_len = 1 mode_str = args.mode if args.mode == "sat": mode_str = f"sat-{'var' if args.var else 'fixed'}" if plain_output: print(f"Generating ({mode_str})...") else: print(f"{Colors.INFO}Generating ({mode_str})...{Colors.RESET}") if (block_stream or resident_dtype) and torch.cuda.is_available(): torch.cuda.reset_peak_memory_stats() start = time.time() if args.mode == "ar" and getattr(args, "swi_reasoning", False): if getattr(args, "block_stream", False) or getattr(args, "sampler", "ar") == "euler": print("[swi] --swi_reasoning needs plain KV decode " "(no --block_stream / --sampler euler); falling back to standard AR.") args.swi_reasoning = False if args.mode == "ar" and getattr(args, "swi_reasoning", False): ids = _swireasoning_decode(core, ar_h, ids, args, min_new) elif args.mode == "ar": _euler = getattr(args, "sampler", "ar") == "euler" if stream: print("[STREAM_BEGIN] " + json.dumps({"mode": "ar", "slots": int(args.max_new)}), flush=True) block_stream_kv = block_stream and _block_stream_kv_cache_enabled(args) kvs = None if not _euler and block_stream_kv: h, kvs = _block_stream_forward_cached( core, ids, causal_mask(ids.size(1), structured=use_structured_masks(args)), None, ids.size(1), args, ) elif not _euler and not block_stream: h, kvs = core(ids, causal_mask(ids.size(1), structured=use_structured_masks(args)), use_cache=True, total_seq_len=ids.size(1)) for _ in range(args.max_new): if _euler: h = _dblock_euler_hidden(core, ids, args) elif block_stream and not block_stream_kv: h = _block_stream_forward(core, ids, causal_mask(ids.size(1), structured=use_structured_masks(args)), args) logits = ar_h(h)[:, -1].float() logits = _apply_penalties(logits, ids, args.penalty_last_n, args.repetition_penalty, args.presence_penalty, args.frequency_penalty) logits = _suppress_eos(logits, args) nxt = _sample(logits, args.temperature, args.top_k, args.top_p, args.min_p, args.greedy) ids = torch.cat([ids, nxt], 1) if stream: print("[STREAM_AR] " + json.dumps({"i": int(ids.size(1) - prompt_len) - 1, "text": tok.decode([int(nxt.item())], skip_special_tokens=True)}), flush=True) if EOS is not None and not getattr(args, "ignore_eos", False) and int(nxt.item()) == int(EOS): break if not _euler: if block_stream_kv: h, kvs = _block_stream_forward_cached(core, ids[:, -1:], None, kvs, ids.size(1), args) elif not block_stream: h, kvs = core(ids[:, -1:], None, kv_caches=kvs, use_cache=True, total_seq_len=ids.size(1)) elif args.mode == "nat": # Iterative mask-predict decode (CMLM): keep the prompt fixed and fill the # BLANK slots, committing confident predictions each pass. Unlike the # original straight argmax path, this applies the same anti-repetition # penalties and sampler used by AR/SAT at each committed position. n_fill = max(1, int(args.max_new)) ids = torch.tensor([prompt_tokens + [BLANK] * n_fill], device=DEV) remaining = set(range(prompt_len, prompt_len + n_fill)) passes = max(1, int(args.nat_passes)) def _nat_history(current_ids: torch.Tensor): keep = current_ids[0] != BLANK if bool(keep.any()): return current_ids[:, keep] return current_ids[:, :max(1, prompt_len)] def _nat_pick(logits_pos: torch.Tensor, current_ids: torch.Tensor): logits_pos = logits_pos.clone() logits_pos[..., BLANK] = -1e9 logits_pos = _apply_penalties( logits_pos, _nat_history(current_ids), args.penalty_last_n, args.repetition_penalty, args.presence_penalty, args.frequency_penalty, ) logits_pos = _suppress_eos(logits_pos, args) nat_greedy = bool(getattr(args, "nat_greedy", True)) return _sample(logits_pos, args.temperature, args.top_k, args.top_p, args.min_p, args.greedy or nat_greedy) if stream: print("[STREAM_BEGIN] " + json.dumps({"mode": "nat", "slots": n_fill}), flush=True) for p in range(passes): if not remaining: break h = _block_stream_forward(core, ids, None, args) if block_stream else core(ids, None) logits = nat_h(h).float() logits[..., BLANK] = -1e9 conf = logits.softmax(-1).amax(-1) k_min = max(1, -(-len(remaining) // (passes - p))) conf_threshold = getattr(args, "nat_conf_threshold", 0.9) confident_positions = [q for q in remaining if float(conf[0, q]) > conf_threshold] if len(confident_positions) > k_min: ordered = sorted(confident_positions, key=lambda q: float(conf[0, q]), reverse=True) else: ordered = sorted(remaining, key=lambda q: float(conf[0, q]), reverse=True)[:k_min] for pos in ordered: nxt = _nat_pick(logits[:, pos, :], ids) ids[0, pos] = int(nxt.reshape(-1)[0]) remaining.discard(pos) if stream: print("[STREAM_NAT] " + json.dumps({"pos": int(pos) - prompt_len, "text": tok.decode([int(ids[0, pos])], skip_special_tokens=True)}), flush=True) if remaining: h = _block_stream_forward(core, ids, None, args) if block_stream else core(ids, None) logits = nat_h(h).float() logits[..., BLANK] = -1e9 for pos in sorted(remaining): nxt = _nat_pick(logits[:, pos, :], ids) ids[0, pos] = int(nxt.reshape(-1)[0]) if stream: print("[STREAM_NAT] " + json.dumps({"pos": int(pos) - prompt_len, "text": tok.decode([int(ids[0, pos])], skip_special_tokens=True)}), flush=True) else: if stream: print("[STREAM_BEGIN] " + json.dumps({"mode": "sat", "slots": int(args.max_new)}), flush=True) cached_len = ids.size(1) block_stream_kv = block_stream and _block_stream_kv_cache_enabled(args) if block_stream_kv: h, kvs = _block_stream_forward_cached( core, ids, sat_mask(ids.size(1), structured=use_structured_masks(args)), None, cached_len, args, ) elif block_stream: h = _block_stream_forward(core, ids, sat_mask(ids.size(1), structured=use_structured_masks(args)), args) kvs = None else: h, kvs = core(ids, sat_mask(ids.size(1), structured=use_structured_masks(args)), use_cache=True, total_seq_len=cached_len) h_buffer = h[:, -SAT_BLOCK:] added = 0 stop = False # Align to a SAT block boundary with AR tokens before block emission. while ids.size(1) % SAT_BLOCK != 0 and added < args.max_new: logits = ar_h(h)[:, -1].float() logits = _apply_penalties(logits, ids, args.penalty_last_n, args.repetition_penalty, args.presence_penalty, args.frequency_penalty) logits = _suppress_eos(logits, args, added < min_new) nxt = _sample(logits, args.temperature, args.top_k, args.top_p, args.min_p, args.greedy) ids = torch.cat([ids, nxt], 1) added += 1 if stream: print("[STREAM_SAT] " + json.dumps({"i": int(added) - 1, "text": tok.decode([int(nxt.item())], skip_special_tokens=True)}), flush=True) if EOS is not None and not getattr(args, "ignore_eos", False) and int(nxt.item()) == int(EOS): stop = True break if block_stream: if block_stream_kv: h, kvs = _block_stream_forward_cached(core, nxt, None, kvs, ids.size(1), args) cached_len = ids.size(1) h_buffer = torch.cat([h_buffer, h], dim=1)[:, -SAT_BLOCK:] else: h = _block_stream_forward(core, ids, sat_mask(ids.size(1), structured=use_structured_masks(args)), args) h_buffer = h[:, -SAT_BLOCK:] else: h, kvs = core(nxt, None, kv_caches=kvs, use_cache=True, total_seq_len=ids.size(1)) cached_len = ids.size(1) h_buffer = torch.cat([h_buffer, h], dim=1)[:, -SAT_BLOCK:] while added < args.max_new and not stop: logits_all, gate = sat_h(h_buffer) logits_all = logits_all.float() if gate is not None: gate = gate.float() stride = SAT_BLOCK if (not args.var or gate is None) else (gate.softmax(-1).multinomial(1).item() + 1) stride = min(int(stride), logits_all.size(1)) new_tokens = [] for i in range(int(stride)): logits = logits_all[:, i].clone() # BLANK is the SAT/NAT mask-filler token; with this tokenizer it is # ALSO the EOS id (pad==eos==1), so an unbanned SAT head "ends" on # every filler prediction while NAT (which bans BLANK) keeps going. # Ban it here exactly like the NAT path does. logits[..., BLANK] = -1e9 logits = _apply_penalties(logits, ids, args.penalty_last_n, args.repetition_penalty, args.presence_penalty, args.frequency_penalty) logits = _suppress_eos(logits, args, added < min_new) nxt = _sample(logits, args.temperature, args.top_k, args.top_p, args.min_p, args.greedy) new_tokens.append(nxt) ids = torch.cat([ids, nxt], 1) added += 1 if stream: print("[STREAM_SAT] " + json.dumps({"i": int(added) - 1, "text": tok.decode([int(nxt.item())], skip_special_tokens=True)}), flush=True) if EOS is not None and not getattr(args, "ignore_eos", False) and int(nxt.item()) == int(EOS): stop = True break if added >= args.max_new: break if stop or added >= args.max_new: break new_ids = torch.cat(new_tokens, dim=1) if block_stream: if block_stream_kv: mask = sat_mask_cached(new_ids.size(1), cached_len, structured=use_structured_masks(args)) h, kvs = _block_stream_forward_cached(core, new_ids, mask, kvs, ids.size(1), args) cached_len = ids.size(1) h_buffer = torch.cat([h_buffer, h], dim=1)[:, -SAT_BLOCK:] else: h = _block_stream_forward(core, ids, sat_mask(ids.size(1), structured=use_structured_masks(args)), args) h_buffer = h[:, -SAT_BLOCK:] else: mask = sat_mask_cached(new_ids.size(1), cached_len, structured=use_structured_masks(args)) h, kvs = core(new_ids, mask, kv_caches=kvs, use_cache=True, total_seq_len=ids.size(1)) cached_len = ids.size(1) h_buffer = torch.cat([h_buffer, h], dim=1)[:, -SAT_BLOCK:] elapsed = time.time() - start gen_tokens = len(ids[0]) - prompt_len tok_per_sec = gen_tokens / elapsed if elapsed > 0 else 0 if (block_stream or resident_dtype) and torch.cuda.is_available(): peak_alloc_gb = torch.cuda.max_memory_allocated() / 1e9 peak_reserved_gb = torch.cuda.max_memory_reserved() / 1e9 label = "block_stream" if block_stream else "resident" print(f"[infer] {label}_cuda_peak_alloc={peak_alloc_gb:.2f}GB peak_reserved={peak_reserved_gb:.2f}GB") if block_stream and _moe_expert_stream_enabled(args): loads, tokens = _moe_expert_stream_stats(core) print(f"[infer] moe_expert_stream_loads={loads} routed_tokens={tokens}") all_tokens = ids[0].tolist() prompt_text = tok.decode(all_tokens[:prompt_len], skip_special_tokens=True) gen_text = tok.decode(all_tokens[prompt_len:], skip_special_tokens=True) safe_prompt = _ascii_safe(prompt_text) if plain_output else prompt_text safe_gen = _ascii_safe(gen_text) if plain_output else gen_text if plain_output: print(f"{safe_prompt}{safe_gen}") print(f"[{elapsed:.2f}s | {gen_tokens} tokens | {tok_per_sec:.1f} tok/s]") else: print(f"{Colors.PROMPT}{safe_prompt}{Colors.RESET}{safe_gen}") print(f"{Colors.INFO}[{elapsed:.2f}s | {gen_tokens} tokens | {tok_per_sec:.1f} tok/s]{Colors.RESET}") if getattr(args, "claude_friendly", False): claude_prompt = _ascii_safe(prompt_text) claude_gen = _ascii_safe(gen_text) print("[CLAUDE_FRIENDLY_START]") print(f"[mode={mode_str}]") print("[prompt_input]") print(claude_prompt) print("[completion]") print(claude_gen) print("[prompt_plus_completion]") print(f"{claude_prompt}{claude_gen}") print(f"[stats] {elapsed:.2f}s | {gen_tokens} tokens | {tok_per_sec:.1f} tok/s") print("[CLAUDE_FRIENDLY_END]") @torch.no_grad() def infer(args): inst = _agillm43_prepare_infer_instance(args) return _agillm43_generate_from_instance(inst, args) @torch.no_grad() def infer_server(args): # AGILLM-WARM-SERVER-PORT 20260703: keep one loaded checkpoint/model alive for many stdin JSON prompts. args.plain_output = True inst = _agillm43_prepare_infer_instance(args) print("[INFER_SERVER_READY]", flush=True) request_fields = { "prompt", "mode", "max_new", "min_new", "temperature", "top_k", "top_p", "min_p", "greedy", "ignore_eos", "nat_passes", "nat_greedy", "nat_conf_threshold", "var", "repetition_penalty", "presence_penalty", "frequency_penalty", "penalty_last_n", "stream", "claude_friendly", "sampler", "euler_steps", "euler_start_sigma", "dblock_blocks", "swi_reasoning", "swi_latent_thresh", "swi_explicit_thresh", "swi_eps", "swi_max_switches", "swi_max_latent", "swi_think_budget", "swi_max_steps", "swi_topk", "swi_start_latent", "block_stream", "block_stream_kv_cache", "block_stream_kv_device", "block_stream_cache_pages", "moe_expert_stream", } for raw_line in sys.stdin: raw_line = raw_line.strip() if not raw_line: continue try: req = json.loads(raw_line) if str(req.get("cmd", "infer")).lower() in {"quit", "exit", "stop"}: print("[INFER_SERVER_STOPPING]", flush=True) break run_args = copy.copy(args) for key, value in req.items(): if key in request_fields: setattr(run_args, key, value) run_args.plain_output = True print("[INFER_SERVER_RESULT_START]", flush=True) _agillm43_generate_from_instance(inst, run_args) print("[INFER_SERVER_RESULT_END]", flush=True) except Exception as exc: print(f"[INFER_SERVER_ERROR] {type(exc).__name__}: {exc}", flush=True) print("[INFER_SERVER_RESULT_END]", flush=True) # ───────────────────────── CLI ───────────────────────── # ------------------------- AGILLM4.3 native supervisor ------------------------- def _agillm43_now_iso(): import time return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()) def _agillm43_log_json(log_path, event, **fields): import json from pathlib import Path payload = {"event": event, "at": _agillm43_now_iso()} payload.update(fields) line = json.dumps(payload, separators=(",", ":")) print(line, flush=True) try: lp = Path(log_path) lp.parent.mkdir(parents=True, exist_ok=True) with lp.open("a", encoding="utf-8") as f: f.write(line + "\n") except Exception: pass def _agillm43_cmdline(pid): from pathlib import Path try: raw = Path(f"/proc/{int(pid)}/cmdline").read_bytes() return [x.decode("utf-8", "ignore") for x in raw.split(b"\0") if x] except Exception: return [] def _agillm43_matching_pids(kind): import os from pathlib import Path me = os.getpid() pids = [] for proc in Path("/proc").glob("[0-9]*"): try: pid = int(proc.name) except ValueError: continue if pid == me: continue cmd = _agillm43_cmdline(pid) if not cmd: continue exe = Path(cmd[0]).name.lower() if "python" not in exe: continue joined = " ".join(cmd) if "agillm41.py" not in joined: continue if kind == "train" and " train " in f" {joined} ": pids.append(pid) elif kind == "supervise" and " supervise " in f" {joined} ": pids.append(pid) return sorted(set(pids)) def _agillm43_gpu_pids(): import subprocess pids = [] try: out = subprocess.check_output( ["nvidia-smi", "--query-compute-apps=pid", "--format=csv,noheader,nounits"], text=True, stderr=subprocess.DEVNULL, timeout=10, ) for line in out.splitlines(): line = line.strip().split(",", 1)[0].strip() if line.isdigit(): pids.append(int(line)) except Exception: pass return pids def _agillm43_latest_step(save_dir): import json from pathlib import Path try: return int(json.loads((Path(save_dir) / "latest.json").read_text()).get("step", 0)) except Exception: return 0 def _agillm43_kill(pid, sig): import os try: os.kill(int(pid), sig) return True except Exception: return False def _agillm43_prepare_env(save_dir, side_dir): import os from pathlib import Path env = os.environ.copy() env.setdefault("TOKENIZERS_PARALLELISM", "false") env.setdefault("TOKENIZER_ID", "deepseek-ai/DeepSeek-V4-Pro") env.setdefault("AGILLM_ATTN_BACKEND", "sublinear") env.pop("PYTORCH_CUDA_ALLOC_CONF", None) shm = Path("/dev/shm") if shm.is_dir() and os.access(shm, os.W_OK): tmp = shm / "agillm_tmp" tmp.mkdir(parents=True, exist_ok=True) env.update({"TMPDIR": str(tmp), "TMP": str(tmp), "TEMP": str(tmp)}) hf_token_path = Path("/root/.cache/huggingface/token") if hf_token_path.exists(): token = hf_token_path.read_text(errors="ignore").strip() if token: env["HF_TOKEN"] = token env["HUGGING_FACE_HUB_TOKEN"] = token def _agillm43_load_secret_file(env_name, paths): if env.get(env_name, "").strip(): return True for raw_path in paths: try: p = Path(raw_path) if p.exists(): val = p.read_text(errors="ignore").strip() if val: env[env_name] = val return True except Exception: pass return False have_deepseek = _agillm43_load_secret_file( "DEEPSEEK_API_KEY", ( "/root/.config/agillm/deepseek_api_key", "/workspace/private/deepseek_api_key", "/workspace/agillm_private/deepseek_api_key", ), ) have_openrouter = _agillm43_load_secret_file( "OPENROUTER_API_KEY", ( "/root/.config/agillm/openrouter_api_key", "/workspace/private/openrouter_api_key", "/workspace/agillm_private/openrouter_api_key", ), ) env.setdefault("AGILLM_MAX_EXAMPLE_TOKENS", "4096") env.setdefault("AGILLM_MAX_EXAMPLE_CHARS", "32768") env.setdefault("AGILLM_DATASET_NN_ROUTER", "1") env.setdefault("AGILLM_DATASET_ROUTER_EXPLORE", "0.08") env.setdefault("AGILLM_DATASET_ROUTER_MIN_SCORE", "0.12") env.setdefault("AGILLM_DATASET_ROUTER_SHARPNESS", "2.0") env.setdefault("AGILLM_DATASET_ROUTER_TARGET_TOKENS", "2048") if have_deepseek or have_openrouter: env.setdefault("AGILLM_DATASET_AGENT_ROUTER", "0") env.setdefault("AGILLM_DATASET_AGENT_PROVIDER", "auto") Path(save_dir).mkdir(parents=True, exist_ok=True) for name in ("incoming", "accepted", "rejected"): (Path(side_dir) / name).mkdir(parents=True, exist_ok=True) return env def _agillm43_prune_save_dir(save_dir): import os from pathlib import Path d = Path(save_dir) for tmp in d.glob("*.tmp"): try: tmp.unlink() except Exception: pass ckpts = sorted([p for p in d.glob("pretrain_step*.pt") if not p.name.endswith(".resume_delta.pt")], key=lambda x: x.stat().st_mtime, reverse=True) for old in ckpts[1:]: try: old.unlink() except Exception: pass def _agillm43_latest_checkpoint_path(save_dir): import glob import json import os from pathlib import Path save = Path(save_dir) src = "" try: src = json.loads((save / "latest.json").read_text()).get("path", "") except Exception: src = "" if src and Path(src).exists(): return str(Path(src)) candidates = sorted([p for p in glob.glob(str(save / "pretrain_step*.pt")) if not str(p).endswith(".resume_delta.pt")], key=os.path.getmtime) return candidates[-1] if candidates else "" def _agillm43_convert_resume_delta(save_dir, log_path): import os import re from pathlib import Path import torch save = Path(save_dir) shm = Path(os.environ.get("SHM_DIR", "/dev/shm")) if not (shm.is_dir() and os.access(shm, os.W_OK)): shm = save out = shm / "agillm43_resume.delta.pt" mark = out.parent / ".agillm43_resume.step" src = _agillm43_latest_checkpoint_path(save) if not src: seed = save / "agillm42_tiekv_seed.delta.pt" _agillm43_log_json(log_path, "native_supervisor_resume_seed", path=str(seed)) return str(seed) src_path = Path(src) m = re.search(r"step0*([0-9]+)", src_path.name) fstep = m.group(1) if m else "" try: st = src_path.stat() src_meta = { "path": str(src_path.resolve()), "name": src_path.name, "size": int(st.st_size), "mtime_ns": int(st.st_mtime_ns), "step": int(fstep) if fstep else None, } except Exception: src_meta = { "path": str(src_path), "name": src_path.name, "step": int(fstep) if fstep else None, } def _resume_delta_mark_matches(): if not (out.exists() and mark.exists()): return False try: payload = json.loads(mark.read_text().strip() or "{}") except Exception: # Old marker files only stored the step number. Rebuild once so a # stale delta from a failed probe cannot replay over a good full ckpt. return False if not isinstance(payload, dict): return False return all(payload.get(k) == v for k, v in src_meta.items()) if _resume_delta_mark_matches(): _agillm43_log_json(log_path, "native_supervisor_resume_delta_current", source=src_meta, path=str(out)) return str(out) ck = _agillm43_load_pt(src_path, map_location="cpu", weights_only=False) tok_keys = ("tokenizer_payload_schema", "tokenizer_id", "tokenizer_json", "tokenizer_bundle", "tokenizer_special", "transformers_version", "tokenizers_version") tok_payload = {} sidecar_payload = _read_tokenizer_sidecar(src_path) tok_payload.update({k: v for k, v in sidecar_payload.items() if k in tok_keys and v is not None}) tok_payload.update({k: ck.get(k) for k in tok_keys if isinstance(ck, dict) and ck.get(k) is not None}) if not tok_payload.get("tokenizer_json") or not tok_payload.get("tokenizer_bundle") or not tok_payload.get("tokenizer_special"): runtime_payload = _tokenizer_payload() tok_payload = {**runtime_payload, **tok_payload} tok_payload.setdefault("tokenizer_payload_schema", 2) src_meta["tokenizer_payload_schema"] = int(tok_payload.get("tokenizer_payload_schema", 2) or 2) delta = { "delta": True, "weights": {k: ck[k] for k in ("core", "ar", "sat", "nat") if k in ck}, "step": ck.get("step", 0), "seen_tok": ck.get("seen_tok", 0), "cfg": ck.get("cfg"), "source_checkpoint": src_meta, **tok_payload, } tmp = str(out) + ".tmp" _agillm43_save_pt(delta, tmp, codec=os.environ.get("AGILLM43_DELTA_CODEC", "zstd3")) os.replace(tmp, out) mark.write_text(json.dumps(src_meta, sort_keys=True)) try: Path(str(out) + ".sha256").unlink() except FileNotFoundError: pass _agillm43_log_json(log_path, "native_supervisor_resume_delta_converted", src=str(src_path), source=src_meta, path=str(out), step=int(delta.get("step", 0))) return str(out) AGILLM43_PROFILE_CHOICES = ("normal", "ar_repair", "full_ar_repair", "sat_repair", "sat_probe", "nat_repair") def _agillm43_profile_config(profile): profile = str(profile or "normal").lower() profiles = { "normal": { "ar_prob": "0.60", "sat_prob": "0.25", "nat_prob": "0.15", "ar_loss_tokens": os.environ.get("AGILLM43_DBLOCK_AR_LOSS_TOKENS", os.environ.get("AGILLM43_DBLOCK_LOSS_TOKENS", "2048")), "sat_loss_tokens": os.environ.get("AGILLM43_DBLOCK_SAT_LOSS_TOKENS", os.environ.get("AGILLM43_DBLOCK_LOSS_TOKENS", "2048")), "nat_loss_tokens": os.environ.get("AGILLM43_DBLOCK_NAT_LOSS_TOKENS", os.environ.get("AGILLM43_DBLOCK_LOSS_TOKENS", "2048")), "sat_every": "1", "nat_every": "4", }, "ar_repair": { # Hybrid-safe recovery mode. Keep AR emphasis for text quality, but # never disable SAT/NAT; AGILLM-4.3 is meant to recover as a hybrid. "ar_prob": "0.55", "sat_prob": "0.30", "nat_prob": "0.15", "ar_loss_tokens": "768", "sat_loss_tokens": "768", "nat_loss_tokens": "512", "sat_every": "1", "nat_every": "4", }, "full_ar_repair": { # Historical profile name retained, but AGILLM4.3 remains a hybrid: # DBLOCK + AR + SAT + NAT all stay live during repair. "ar_prob": "0.60", "sat_prob": "0.25", "nat_prob": "0.15", "ar_loss_tokens": "1024", "sat_loss_tokens": "768", "nat_loss_tokens": "512", "sat_every": "1", "nat_every": "4", "batch_size": "2", "block": "768", "steps": "500", "lr_core": "1e-5", "lr_head": "5e-5", "save_every_sec": "900", }, "sat_repair": { "ar_prob": "0.45", "sat_prob": "0.40", "nat_prob": "0.15", "ar_loss_tokens": "512", "sat_loss_tokens": "1024", "nat_loss_tokens": "512", "sat_every": "1", "nat_every": "4", }, "sat_probe": { "ar_prob": "0.05", "sat_prob": "0.90", "nat_prob": "0.05", "ar_loss_tokens": "256", "sat_loss_tokens": "2048", "nat_loss_tokens": "256", "sat_every": "1", "nat_every": "4", }, "nat_repair": { # Recovery profile for post-fix NAT training: keep AR/SAT alive, but # give NAT enough objective incidence and dense token coverage to catch up. "ar_prob": "0.45", "sat_prob": "0.25", "nat_prob": "0.30", "ar_loss_tokens": "512", "sat_loss_tokens": "1024", "nat_loss_tokens": "4096", "sat_every": "1", "nat_every": "1", "nat_loss_weight": "1.0", "nat_span_mask_prob": "0.45", "nat_suffix_mask_prob": "0.35", }, } if profile not in profiles: raise ValueError(f"unknown AGILLM4.3 profile {profile!r}; choose one of {', '.join(AGILLM43_PROFILE_CHOICES)}") cfg = profiles[profile].copy() cfg["name"] = profile return cfg def _agillm43_train_argv(save_dir, side_dir, resume_delta, profile="normal", warmstart_from=None): import sys from pathlib import Path script = str(Path(__file__).resolve()) incoming = str(Path(side_dir) / "incoming") accepted = str(Path(side_dir) / "accepted") rejected = str(Path(side_dir) / "rejected") prof = _agillm43_profile_config(profile) return [ sys.executable, "-u", script, "train", "--preset", "agillm4_floor", "--tie_kv", "--resume_delta", resume_delta, *(["--warmstart_from", str(warmstart_from)] if warmstart_from else []), "--dblock", "--dblock_blocks", os.environ.get("AGILLM43_DBLOCK_BLOCKS", "14"), "--dblock_schedule", "loss_balanced", "--dblock_router", "transformer", "--dblock_router_blend", "0.35", "--dblock_router_ramp_steps", "256", "--dblock_warmup_steps", "16", "--dblock_sigma_curriculum_steps", "2000", "--dblock_sigma_sampling", "lognormal", "--dblock_sigma_stratified", "--dblock_log_every", "25", "--dblock_objective_mode", "stochastic", "--dblock_ar_prob", prof["ar_prob"], "--dblock_sat_prob", prof["sat_prob"], "--dblock_nat_prob", prof["nat_prob"], "--nat_loss_weight", prof.get("nat_loss_weight", "1.0"), "--dblock_ar_loss_tokens", prof["ar_loss_tokens"], "--dblock_sat_loss_tokens", prof["sat_loss_tokens"], "--dblock_nat_loss_tokens", prof["nat_loss_tokens"], "--moe_ffn", "--moe_experts", "2", "--moe_top_k", "1", "--moe_mlp_mult", "4", "--moe_shared_experts", "1", "--moe_shared_mlp_mult", "2", "--moe_aux_coef", "0.01", "--moe_z_coef", "0.001", "--tie_weights", "--batch_size", prof.get("batch_size", os.environ.get("AGILLM43_BATCH_SIZE", "22")), "--block", prof.get("block", os.environ.get("AGILLM43_BLOCK", "1536")), *(["--steps", prof["steps"]] if "steps" in prof else []), "--amp", "--attn_backend", os.environ.get("AGILLM43_ATTN_BACKEND", "sdpa"), "--sublinear_window", "128", "--sublinear_stride", "128", "--sublinear_max_anchors", "128", "--sublinear_chunk", "128", "--sublinear_sinks", "4", "--sublinear_recent_anchors", "64", "--no-sublinear_pooled_landmarks", "--dblock_checkpoint_stride", "1", "--optimizer", "adamw8bit", "--loss_spike_skip", "3.0", "--sat_every", prof["sat_every"], "--nat_every", prof["nat_every"], *(["--lr_core", prof["lr_core"], "--lr_head", prof["lr_head"]] if "lr_core" in prof and "lr_head" in prof else []), "--nat_max_tokens", "768", "--nat_mask_ratio", "0.5", "--nat_span_mask_prob", prof.get("nat_span_mask_prob", "0.35"), "--nat_suffix_mask_prob", prof.get("nat_suffix_mask_prob", "0.20"), "--token_param_ratio", "55", "--val_tokens", "32768", "--val_every_sec", "3600", "--val_source", "json:/workspace/agillm_math_numeracy_synth/train.jsonl", "--data_seed", "-1", "--save_dir", str(save_dir), "--save_every_sec", prof.get("save_every_sec", "14400"), "--heartbeat_every_sec", "300", "--empty_cache_every_steps", "0", "--delta_every_steps", "0", "--delta_every_sec", str(DEFAULT_DELTA_SEC), "--delta_max_keep", "1", "--max_ckpts", "1", "--async_update_dir", incoming, "--async_update_every_steps", os.environ.get("AGILLM43_ASYNC_UPDATE_EVERY_STEPS", "50"), "--async_update_alpha", os.environ.get("AGILLM43_ASYNC_UPDATE_ALPHA", "0.10"), "--async_update_max_per_check", "2", "--async_update_max_age_sec", "86400", "--async_update_accepted_dir", accepted, "--async_update_rejected_dir", rejected, ] def _agillm43_dedupe_trainers(log_path, keep_pid=None): import signal pids = _agillm43_matching_pids("train") if len(pids) <= 1: return pids gpu = [p for p in _agillm43_gpu_pids() if p in pids] keep = int(keep_pid) if keep_pid in pids else (gpu[0] if gpu else pids[0]) for pid in pids: if pid == keep: continue _agillm43_log_json(log_path, "native_supervisor_kill_duplicate", pid=pid, keep=keep) _agillm43_kill(pid, signal.SIGTERM) return [keep] def supervise_agillm43(args): import os import subprocess import time from pathlib import Path log_path = args.log save_dir = args.save_dir side_dir = args.side_dir pause_file = Path(args.pause_file) script_dir = Path(__file__).resolve().parent os.chdir(script_dir) env = _agillm43_prepare_env(save_dir, side_dir) profile = str(getattr(args, "profile", None) or os.environ.get("AGILLM43_PROFILE", "normal")) _agillm43_profile_config(profile) _agillm43_log_json(log_path, "native_supervisor_start", pid=os.getpid(), save_dir=str(save_dir), side_dir=str(side_dir), profile=profile) while True: while pause_file.exists(): _agillm43_log_json(log_path, "native_supervisor_paused", pause=str(pause_file)) time.sleep(5) if args.dedupe: _agillm43_dedupe_trainers(log_path) live = _agillm43_matching_pids("train") if live: if args.once: _agillm43_log_json(log_path, "native_supervisor_existing_trainer", pids=live) return 0 time.sleep(max(1, args.sleep_sec)) continue _agillm43_prune_save_dir(save_dir) resume_src = _agillm43_latest_checkpoint_path(save_dir) resume_delta = _agillm43_convert_resume_delta(save_dir, log_path) argv = _agillm43_train_argv(save_dir, side_dir, resume_delta, profile=profile, warmstart_from=resume_src) _agillm43_log_json(log_path, "native_supervisor_launch", profile=profile, warmstart_from=resume_src, argv=" ".join(argv)) with open(log_path, "a", encoding="utf-8", buffering=1) as lf: child = subprocess.Popen(argv, cwd=str(script_dir), env=env, stdout=lf, stderr=subprocess.STDOUT) if args.once: _agillm43_log_json(log_path, "native_supervisor_launched_once", pid=child.pid) return 0 while child.poll() is None: if args.dedupe: _agillm43_dedupe_trainers(log_path, keep_pid=child.pid) time.sleep(max(1, args.sleep_sec)) _agillm43_log_json(log_path, "native_supervisor_trainer_exit", pid=child.pid, rc=child.returncode) time.sleep(max(1, args.sleep_sec)) def hotpatch_agillm43(args): import os import signal import subprocess import time from pathlib import Path log_path = args.log save_dir = Path(args.save_dir) pause_file = Path(args.pause_file) pause_file.touch() _agillm43_log_json(log_path, "native_hotpatch_pause", pause=str(pause_file)) try: pids = _agillm43_dedupe_trainers(log_path) pids = _agillm43_matching_pids("train") if pids: gpu = [p for p in _agillm43_gpu_pids() if p in pids] keep = gpu[0] if gpu else pids[0] before = _agillm43_latest_step(save_dir) _agillm43_log_json(log_path, "native_hotpatch_flush_requested", pid=keep, before_step=before) (save_dir / "FLUSH_NOW").touch() _agillm43_kill(keep, signal.SIGUSR1) deadline = time.time() + args.wait_flush_sec while time.time() < deadline: cur = _agillm43_latest_step(save_dir) if cur > before: _agillm43_log_json(log_path, "native_hotpatch_flush_done", latest_step=cur) break time.sleep(5) else: cur = _agillm43_latest_step(save_dir) _agillm43_log_json(log_path, "native_hotpatch_flush_timeout", latest_step=cur, before_step=before) if not args.force: return 2 else: _agillm43_log_json(log_path, "native_hotpatch_no_trainer") for spid in _agillm43_matching_pids("supervise"): if spid == os.getpid(): continue _agillm43_log_json(log_path, "native_hotpatch_stop_supervisor", pid=spid) _agillm43_kill(spid, signal.SIGTERM) if args.kill_tmux: subprocess.run(["tmux", "kill-session", "-t", args.tmux_session], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) time.sleep(2) for pid in _agillm43_matching_pids("train"): _agillm43_log_json(log_path, "native_hotpatch_stop_trainer", pid=pid) _agillm43_kill(pid, signal.SIGTERM) deadline = time.time() + 120 while time.time() < deadline and _agillm43_matching_pids("train"): time.sleep(2) for pid in _agillm43_matching_pids("train"): _agillm43_log_json(log_path, "native_hotpatch_kill_stubborn", pid=pid) _agillm43_kill(pid, signal.SIGKILL) pause_file.unlink(missing_ok=True) cmd = [ "python3", "-u", str(Path(__file__).resolve()), "supervise", "--save_dir", str(save_dir), "--side_dir", args.side_dir, "--log", log_path, "--pause_file", str(pause_file), "--sleep_sec", str(args.sleep_sec), "--profile", str(args.profile), ] if args.tmux: import shlex quoted = " ".join(shlex.quote(part) for part in cmd) subprocess.run(["tmux", "new-session", "-d", "-s", args.tmux_session, quoted], check=False) if not _agillm43_matching_pids("supervise"): with open(args.nohup_log, "a", encoding="utf-8") as lf: subprocess.Popen(cmd, cwd=str(Path(__file__).resolve().parent), stdout=lf, stderr=subprocess.STDOUT, start_new_session=True) _agillm43_log_json(log_path, "native_hotpatch_start_supervisor_nohup_fallback", log=args.nohup_log) else: _agillm43_log_json(log_path, "native_hotpatch_start_supervisor_tmux", session=args.tmux_session) else: with open(args.nohup_log, "a", encoding="utf-8") as lf: subprocess.Popen(cmd, cwd=str(Path(__file__).resolve().parent), stdout=lf, stderr=subprocess.STDOUT, start_new_session=True) _agillm43_log_json(log_path, "native_hotpatch_start_supervisor_nohup", log=args.nohup_log) deadline = time.time() + args.wait_start_sec while time.time() < deadline: live = _agillm43_matching_pids("train") if len(live) == 1: _agillm43_log_json(log_path, "native_hotpatch_restart_done", pid=live[0], latest_step=_agillm43_latest_step(save_dir)) return 0 if len(live) > 1: _agillm43_dedupe_trainers(log_path) time.sleep(3) _agillm43_log_json(log_path, "native_hotpatch_restart_timeout", trainer_count=len(_agillm43_matching_pids("train"))) return 3 finally: try: pause_file.unlink() except FileNotFoundError: pass def main(): ap = argparse.ArgumentParser(description="AGILLM Expansion Ratio Testing") sub = ap.add_subparsers(dest="cmd", required=True) tr = sub.add_parser("train") tr.add_argument("--preset", choices=PRESETS.keys(), default="large") tr.add_argument("--rank", type=int) tr.add_argument("--block", type=int, default=DEFAULT_BLOCK) tr.add_argument("--batch_size", type=int, default=DEFAULT_BATCH) tr.add_argument("--source", default=DEFAULT_PRETRAIN_SOURCES) tr.add_argument("--target_tokens", type=int) tr.add_argument("--token_param_ratio", type=float, default=0.0, help="If --target_tokens is omitted, train to this tokens:param ratio. AGILLM-4 presets default to 100.") tr.add_argument("--steps", type=int) tr.add_argument("--amp", action="store_true") tr.add_argument("--compile", action="store_true", help="Use torch.compile for speedup") tr.add_argument("--attn_backend", choices=["manual", "sdpa", "sublinear"], default=DEFAULT_ATTN_BACKEND, help="AGILLM-4 attention backend. sublinear uses local-window plus landmark candidates.") tr.add_argument("--grad_checkpoint", action="store_true", help="Recompute transformer blocks during backward to trade speed for longer context.") tr.add_argument("--sublinear_window", type=int, default=DEFAULT_SUBLINEAR_WINDOW, help="For --attn_backend sublinear, attend to this many local tokens on each side.") tr.add_argument("--sublinear_stride", type=int, default=DEFAULT_SUBLINEAR_STRIDE, help="For --attn_backend sublinear, use every Nth token as a landmark candidate.") tr.add_argument("--sublinear_max_anchors", type=int, default=DEFAULT_SUBLINEAR_MAX_ANCHORS, help="For --attn_backend sublinear, cap landmark candidates per query chunk.") tr.add_argument("--sublinear_chunk", type=int, default=DEFAULT_SUBLINEAR_CHUNK, help="For --attn_backend sublinear, query chunk size controlling peak gather memory.") tr.add_argument("--sublinear_sinks", type=int, default=DEFAULT_SUBLINEAR_SINKS, help="For sublinear attention, always include this many first-token attention sinks.") tr.add_argument("--sublinear_recent_anchors", type=int, default=DEFAULT_SUBLINEAR_RECENT_ANCHORS, help="For capped sublinear anchors, reserve this many anchors for the recent tail; -1 uses half.") tr.add_argument("--sublinear_pooled_landmarks", action=argparse.BooleanOptionalAction, default=DEFAULT_SUBLINEAR_POOLED_LANDMARKS, help="Use stride-segment pooled K/V summaries for sublinear landmark anchors.") tr.add_argument("--no_structured_masks", action="store_true", help="Disable structured causal/SAT masks for sublinear attention and fall back to dense masks.") tr.add_argument("--anchor_memory", action="store_true", help="Enable anchor-memory long-context augmentation (one AnchorMemoryLayer at mid-stack).") tr.add_argument("--anchor_stride", type=int, default=DEFAULT_ANCHOR_STRIDE, help="Token span compressed into one anchor (default 256).") tr.add_argument("--anchor_max", type=int, default=DEFAULT_ANCHOR_MAX, help="Max anchors retained in the rolling memory bank.") tr.add_argument("--anchor_position", type=int, default=DEFAULT_ANCHOR_POSITION, help="Block index after which to insert anchor memory (-1 = stack middle).") tr.add_argument("--kv_buffer", action="store_true", help="Use preallocated KV buffer instead of torch.cat-based cache growth.") tr.add_argument("--optimizer", choices=["adamw", "adamw8bit", "paged_adamw8bit", "powerstep"], default="adamw", help="Optimizer backend. 8-bit options reduce VRAM on 24GB production runs. 'powerstep' (arXiv:2605.10335) uses a single momentum buffer; in a faithful dblock-step benchmark it converged below Adam, but needs its own LR (~1e-3) and an int8/paged buffer to fit at B=6.") tr.add_argument("--powerstep_beta", type=float, default=0.1, help="PowerStep signed-power exponent beta in (0,1); 0.1 is the paper's recommended value.") tr.add_argument("--powerstep_momentum", type=float, default=0.9, help="PowerStep heavy-ball momentum coefficient gamma.") tr.add_argument("--powerstep_int8", action="store_true", help="PowerStep: store the momentum buffer as blockwise int8 in VRAM (~1/4 VRAM; needs bitsandbytes).") tr.add_argument("--powerstep_paged", action="store_true", help="PowerStep: keep the momentum buffer in pinned CPU RAM (~0 persistent VRAM, spends RAM+PCIe).") tr.add_argument("--save_every_sec", type=int, default=DEFAULT_SAVE_SEC) tr.add_argument("--disk_free_floor_gb", type=float, default=12.0, help="In-file disk auto-prune: when free space drops below this, escalate pruning of transient artifacts and old checkpoints. 0 disables the floor (routine keep-count pruning still runs).") tr.add_argument("--val_tokens", type=int, default=0, help="Held-out validation set size in tokens (sampled once from --val_seed stream at startup). 0 disables validation.") tr.add_argument("--val_every_sec", type=int, default=3600, help="Run held-out validation every N seconds (requires --val_tokens > 0).") tr.add_argument("--val_seed", type=int, default=1337, help="Shuffle seed for the held-out validation stream (distinct from the training data seed).") tr.add_argument("--val_source", default="", help="Optional validation-only dataset source. When set, bypasses hot_config so health probes are comparable across restarts.") tr.add_argument("--data_seed", type=int, default=42, help="Training stream shuffle seed. -1 derives a per-restart seed from the resume step so restarts do not re-train identical early data.") tr.add_argument("--heartbeat_every_sec", type=int, default=300, help="Print lightweight trainer heartbeat/status lines every N seconds; 0 disables.") tr.add_argument("--oom_auto_backoff", action=argparse.BooleanOptionalAction, default=True, help="Persist learned CUDA OOM batch/block limits and cap future launches before they OOM.") tr.add_argument("--oom_memory_path", default="", help="Optional JSON path for persistent OOM backoff memory. Defaults to /oom_backoff_state.json.") tr.add_argument("--oom_backoff_safety", type=float, default=0.92, help="Safety multiplier used after a known OOM or high OOM prediction.") tr.add_argument("--oom_predict_threshold", type=float, default=0.70, help="Tiny online MLP OOM probability above which startup batch is capped.") tr.add_argument("--oom_warmup_good_steps", type=int, default=16, help="Steps at one batch size before it is re-recorded as a stable safe batch.") tr.add_argument("--oom_retries_before_backoff", type=int, default=0, help="OOM retries at the same batch before reducing. 0 immediately backs off and remembers.") tr.add_argument("--empty_cache_every_steps", type=int, default=0, help="Call torch.cuda.empty_cache() every N train steps; useful for VRAM-first runs where lower reserved VRAM matters more than speed.") tr.add_argument("--profile_steps", type=int, default=0, help="Profile the first N DBlock training steps with in-process CUDA timers; 0 disables.") tr.add_argument("--profile_log_every", type=int, default=25, help="Print averaged profiler timings every N profiled steps.") tr.add_argument("--delta_every_steps", type=int, default=DEFAULT_DELTA_STEPS, help="Weight-only delta save every N steps (0=off; production should prefer --delta_every_sec)") tr.add_argument("--delta_every_sec", type=int, default=DEFAULT_DELTA_SEC, help="Weight-only delta save every N seconds (0=off)") tr.add_argument("--delta_max_keep", type=int, default=DEFAULT_MAX_DELTAS, help="Max delta checkpoints to keep") tr.add_argument("--delta_codec", default=os.environ.get("AGILLM43_DELTA_CODEC", "zstd3"), help="Delta checkpoint payload codec: off/raw, zstd/zstdN, or block-sharded-zstd. Sharded mode writes a manifest plus independently loadable shards.") tr.add_argument("--ckpt_codec", default=os.environ.get("AGILLM43_CKPT_CODEC", "zstd3"), help="Full checkpoint payload codec: off/raw, zstd/zstdN, or block-sharded-zstd. Sharded mode writes a manifest plus per-block/head shards and is accepted by load_ckpt, infer, and resume-delta conversion.") tr.add_argument("--resume_delta", type=str, help="Resume from a delta (weight-only, no optimizer state)") tr.add_argument("--async_update_dir", default="", help="Optional incoming directory for verified DBlock side updates. Empty disables async side updates.") tr.add_argument("--async_update_every_steps", type=int, default=0, help="Poll --async_update_dir every N master steps. Side workers never block master progress.") tr.add_argument("--async_update_alpha", type=float, default=1.0, help="Blend factor for accepted side updates: 1.0 copies side block weights; lower values lerp into live weights.") tr.add_argument("--async_update_max_per_check", type=int, default=1, help="Maximum side-update files to apply per poll.") tr.add_argument("--async_update_max_age_sec", type=float, default=0.0, help="Reject incoming side updates older than this many seconds. 0 disables age rejection.") tr.add_argument("--async_update_accepted_dir", default="", help="Directory for applied side-update files. Defaults to a sibling accepted/ directory.") tr.add_argument("--async_update_rejected_dir", default="", help="Directory for rejected side-update files. Defaults to a sibling rejected/ directory.") tr.add_argument("--save_dir", default=str(CKDIR)) tr.add_argument("--resume", type=str) tr.add_argument("--x2", action="store_true") tr.add_argument("--warmstart_from", type=str) tr.add_argument("--ckpt_role", type=str, default="", help="Federation role tag embedded in checkpoint filenames (e.g. master, lease, coordinator). Empty = no tag.") tr.add_argument("--fresh", action="store_true") tr.add_argument("--max_ckpts", type=int, default=2) tr.add_argument("--chilla_max_double", action="store_true") tr.add_argument("--tie_weights", action="store_true") tr.add_argument("--ar_only", action="store_true") tr.add_argument("--agillm3_compat", action="store_true", help="Legacy AGILLM3/3.5 checkpoint mode. Use TOKENIZER_ID=deepseek-ai/DeepSeek-V3.2 or the agillm35.py shim for the old tokenizer contract.") tr.add_argument("--no_nat_head", action="store_true", help="Do not instantiate/save a NAT head. Keeps AGILLM3 AR+SAT checkpoint schema and reduces params/RAM.") tr.add_argument("--sat_every", type=int, default=1, help="Train SAT every N steps. Default 1 keeps AR+SAT every step.") tr.add_argument("--nat_every", type=int, default=1, help="Train NAT every N steps with a mask-predict objective. Default 1 keeps AR+SAT+NAT every step.") tr.add_argument("--nat_loss_weight", type=float, default=1.0) tr.add_argument("--nat_expand", type=int, default=2, help="Legacy NAT expansion factor; retained for checkpoint/script compatibility.") tr.add_argument("--nat_max_tokens", type=int, default=0, help="Optional cap for NAT target tokens per batch; 0 uses the whole block.") tr.add_argument("--nat_span_mask_prob", type=float, default=0.35, help="NAT CMLM probability of replacing random holes with one contiguous masked span.") tr.add_argument("--nat_suffix_mask_prob", type=float, default=0.20, help="NAT CMLM probability of training on a right-suffix masked span, matching generation.") tr.add_argument("--nat_span_max_tokens", type=int, default=0, help="Maximum NAT contiguous/suffix span length; 0 derives it from --nat_mask_ratio.") tr.add_argument("--dblock_nat_embed_noise_mode", choices=["off", "visible", "mask_plus_noise"], default="mask_plus_noise", help="NAT embedding noise mode. off=standard BLANK masking. visible=add noise to clean embeddings. mask_plus_noise=BLANK mask + noise on masked positions.") tr.add_argument("--dblock_nat_embed_noise_scale", type=float, default=1.0, help="Scale factor for embedding noise in NAT hybrid modes.") tr.add_argument("--nat_mask_ratio", type=float, default=0.5, help="Fraction of positions masked to BLANK for the NAT mask-predict (CMLM) objective.") tr.add_argument("--tie_kv", action=argparse.BooleanOptionalAction, default=False, help="Q-K=V: tie Key & Value into one projection (~50%% KV cache, -33%% qkv params). Trained-in only; not loadable into a 3-proj checkpoint.") tr.add_argument("--moe_ffn", action=argparse.BooleanOptionalAction, default=DEFAULT_MOE_FFN, help="Use Mixture-of-Experts feed-forward layers inside the transformer blocks.") tr.add_argument("--moe_experts", type=int, default=DEFAULT_MOE_EXPERTS, help="Number of FFN experts per transformer block when --moe_ffn is enabled.") tr.add_argument("--moe_top_k", type=int, default=DEFAULT_MOE_TOP_K, help="Router top-k experts per token when --moe_ffn is enabled.") tr.add_argument("--moe_mlp_mult", type=int, default=DEFAULT_MOE_MLP_MULT, help="Expert hidden-size multiplier; 4 preserves dense FFN checkpoint shape for seeding.") tr.add_argument("--moe_shared_experts", type=int, default=0, help="Always-on shared experts added to the routed output (DeepSeek/ST-MoE style). 0 disables. Output is zero-init so it merges into an existing checkpoint as a no-op then learns to contribute.") tr.add_argument("--moe_shared_mlp_mult", type=int, default=0, help="Hidden-size multiplier for shared experts (0 = same as --moe_mlp_mult). Use a smaller value (1-2) to limit added VRAM.") tr.add_argument("--moe_aux_coef", type=float, default=0.0, help="Weight for the MoE load-balance (Switch) aux loss. 0 disables (legacy). ~0.01 keeps both experts utilised under top-1 routing. Checkpoint-safe (router recomputed outside the checkpoint).") tr.add_argument("--moe_z_coef", type=float, default=0.0, help="Weight for the MoE router z-loss (router-logit magnitude regularizer). 0 disables. ~0.001 stabilizes routing.") tr.add_argument("--loss_spike_skip", type=float, default=0.0, help="Skip the optimizer step when the mean raw CE exceeds this multiple of its EMA (dblock path). 0 disables. ~3.0 drops pathological noisy-batch spikes.") tr.add_argument("--dblock", action="store_true", help="DiffusionBlocks block-wise denoising training (low VRAM).") tr.add_argument("--dblock_looped", action="store_true", help="Experimental opt-in recurrent-depth DBlock mode: reuse one shared physical layer group across all sigma bands with a learned loop-index embedding. Single sampled band per step, no BPTT. Default off.") tr.add_argument("--dblock_loop_layers", type=int, default=0, help="Number of physical layers in the shared looped DBlock group. 0 chooses layers/dblock_blocks.") tr.add_argument("--dblock_loop_start", type=int, default=0, help="First physical layer index for the shared looped DBlock group.") tr.add_argument("--dblock_loop_cond_scale", type=float, default=1.0, help="Scale for the learned loop-index embedding added at shared block entry.") tr.add_argument("--auto_dblock_search", action="store_true", help="Auto-search block configs") tr.add_argument("--dblock_blocks", type=int, default=4, help="Partition layers into this many DiffusionBlocks blocks.") tr.add_argument("--dblock_schedule", choices=["random", "roundrobin", "loss_balanced"], default="loss_balanced", help="How --dblock chooses the next layer block. loss_balanced focuses blocks whose EMA loss is highest after warmup.") tr.add_argument("--dblock_router", choices=["heuristic", "transformer"], default="heuristic", help="Optional learned sequence-Transformer scheduler for DBlock layer-band selection; coverage guards still enforce fairness.") tr.add_argument("--dblock_router_hidden", type=int, default=64, help="Hidden width for the context/history sequence-Transformer DBlock router.") tr.add_argument("--dblock_router_heads", type=int, default=4, help="Attention heads for the context/history sequence-Transformer DBlock router.") tr.add_argument("--dblock_router_layers", type=int, default=2, help="Transformer encoder layers for the context/history sequence-Transformer DBlock router.") tr.add_argument("--dblock_router_lr", type=float, default=0.002, help="Online learning rate for the context/history sequence-Transformer DBlock router.") tr.add_argument("--dblock_router_blend", type=float, default=0.35, help="Max blend of learned-router score into heuristic DBlock score after ramp-up.") tr.add_argument("--dblock_router_ramp_steps", type=int, default=256, help="DBlock steps over which the learned router ramps from 0 to --dblock_router_blend.") tr.add_argument("--dblock_warmup_steps", type=int, default=16, help="Initial DBlock steps spent covering every block before loss-balanced scheduling.") tr.add_argument("--dblock_explore", type=float, default=0.08, help="Exploration rate for loss-balanced DBlock scheduling.") tr.add_argument("--dblock_max_stale_steps", type=int, default=64, help="Force the stalest DBlock after this many unselected DBlock steps; 0 disables.") tr.add_argument("--dblock_max_count_skew", type=float, default=1.35, help="Force least-trained DBlock when max/min sampled block counts exceed this ratio; <=1 disables.") tr.add_argument("--dblock_stale_bonus", type=float, default=0.35, help="Loss-score bonus for stale DBlocks before the hard stale guard triggers.") tr.add_argument("--dblock_undertrain_bonus", type=float, default=0.25, help="Loss-score bonus for under-sampled DBlocks before the hard count-skew guard triggers.") tr.add_argument("--dblock_log_every", type=int, default=25, help="Print DBlock block/loss/VRAM diagnostics every N DBlock steps; 0 disables.") tr.add_argument("--dblock_sublayer_mode", choices=["off", "full", "attn_only", "ffn_only", "split_alt", "cycle"], default="off", help="Experimental dormant knob: train only transformer sublayers inside selected DiffusionBlocks. off/full keeps normal Block.forward; attn_only trains LN1+attention residual; ffn_only trains LN2+FFN/MoE residual; split_alt alternates attention/FFN by step; cycle rotates full/FFN/attention.") tr.add_argument("--dblock_checkpoint_stride", type=int, default=1, help="With --grad_checkpoint in --dblock mode, checkpoint one layer every N selected block layers; 1=all layers, 2=alternate, 0=off.") tr.add_argument("--dblock_checkpoint_skip_tail", type=int, default=0, help="Experimental DBlock speed knob: do not checkpoint this many final layers in the selected block, reducing backward recompute at higher VRAM cost.") tr.add_argument("--dblock_activation_offload", action="store_true", help="Experimental DBlock speed knob: for non-checkpointed block layers, offload saved backward tensors to CPU RAM instead of recomputing.") tr.add_argument("--dblock_activation_offload_min_mb", type=float, default=1.0, help="Minimum CUDA tensor size in MB to offload under --dblock_activation_offload.") tr.add_argument("--dblock_sigma_curriculum_steps", type=int, default=2000, help="Warm sigma ranges from easy to full span over this many DBlock steps; 0 disables.") tr.add_argument("--dblock_sigma_sampling", choices=["lognormal", "truncated_lognormal", "edm", "log_uniform"], default="lognormal", help="Sigma sampling inside each DBlock interval. lognormal/truncated_lognormal follows the DBT/EDM p_noise conditional; log_uniform is the legacy sampler.") tr.add_argument("--dblock_sigma_stratified", action=argparse.BooleanOptionalAction, default=True, help="Use randomized quantile strata for log-normal DBlock sigma sampling; reduces per-step sigma Monte Carlo variance.") tr.add_argument("--dblock_sigma_min", type=float, default=0.002, help="Minimum sigma for DBlock equi-probability partitioning.") tr.add_argument("--dblock_sigma_max", type=float, default=80.0, help="Maximum sigma for DBlock equi-probability partitioning.") tr.add_argument("--dblock_sigma_pmean", type=float, default=-1.2, help="Mean of log(sigma) for DBlock log-normal p_noise.") tr.add_argument("--dblock_sigma_pstd", type=float, default=1.2, help="Stddev of log(sigma) for DBlock log-normal p_noise.") tr.add_argument("--dblock_edm_wmax", type=float, default=5.0, help="Cap for EDM loss weighting in DBlock mode.") tr.add_argument("--dblock_ar_weight", type=float, default=1.0) tr.add_argument("--dblock_sat_weight", type=float, default=1.0) tr.add_argument("--dblock_nat_weight", type=float, default=1.0) tr.add_argument("--dblock_objective_mode", choices=["periodic", "stochastic"], default="periodic", help="DBlock objective scheduler. stochastic samples one objective per step to reduce redundant AR/SAT/NAT forwards.") tr.add_argument("--dblock_ar_prob", type=float, default=0.80, help="Stochastic DBlock probability for AR objective.") tr.add_argument("--dblock_sat_prob", type=float, default=0.10, help="Stochastic DBlock probability for SAT objective.") tr.add_argument("--dblock_nat_prob", type=float, default=0.10, help="Stochastic DBlock probability for NAT objective.") tr.add_argument("--dblock_ar_loss_tokens", type=int, default=0, help="If >0, uniformly sample this many AR target positions per DBlock step for stochastic token-level CE. Hot-configurable via dblock_ar_loss_tokens or dblock.loss_tokens.") tr.add_argument("--dblock_sat_loss_tokens", type=int, default=0, help="If >0, uniformly sample this many SAT target positions per DBlock step. Hot-configurable via dblock_sat_loss_tokens or dblock.loss_tokens.") tr.add_argument("--dblock_nat_loss_tokens", type=int, default=0, help="If >0, uniformly sample this many NAT target positions per DBlock step. Hot-configurable via dblock_nat_loss_tokens or dblock.loss_tokens.") tr.add_argument("--reinit_nat", action="store_true", help="Reinitialize NAT head weights after load (use once when switching to mask-predict).") tr.add_argument("--seed_nat_from_ar", action="store_true", help="Seed the NAT head from the trained AR head ('father') after load instead of random init.") tr.add_argument("--freeze_core", action="store_true") tr.add_argument("--unfreeze_ln", action="store_true") tr.add_argument("--train_emb", action="store_true") tr.add_argument("--lr_core", type=float, default=LR_CORE) tr.add_argument("--lr_head", type=float, default=LR_HEAD) tr.add_argument("--chat", action="store_true") tr.add_argument("--chat_messages_key", default="messages") tr.add_argument("--dataset_field_text", default="text") tr.add_argument("--sft_add_generation_prompt", action="store_true") tr.add_argument("--auto_grow", action="store_true") tr.add_argument("--grow_plan", default="576,640,768,896,1024,1122") tr.add_argument("--grow_every_steps", type=int, default=50000) tr.add_argument("--after_sft_source", default="") tr.add_argument("--after_sft_steps", type=int, default=0) tr.add_argument("--after_sft_chat", action="store_true") tr.add_argument("--after_sft_chat_messages_key", default="messages") tr.add_argument("--after_sft_dataset_field_text", default="text") tr.add_argument("--after_sft_add_generation_prompt", type=bool, default=None) tr.add_argument("--after_sft_block", type=int, default=0) tr.add_argument("--after_sft_freeze_core", action="store_true") tr.add_argument("--after_sft_unfreeze_ln", action="store_true") tr.add_argument("--after_sft_train_emb", action="store_true") tr.add_argument("--after_sft_lr_core", type=float, default=0.0) tr.add_argument("--after_sft_lr_head", type=float, default=0.0) inf = sub.add_parser("infer") inf.add_argument("--mode", choices=["ar", "sat", "nat"], required=True) inf.add_argument("--sampler", choices=["ar", "euler"], default="ar", help="ar=KV decode; euler=DiffusionBlocks EDM Euler sampler.") inf.add_argument("--euler_steps", type=int, default=0, help="Euler ODE steps (0=2x dblock_blocks).") inf.add_argument("--euler_start_sigma", type=float, default=0.0, help="Euler start noise (0=sigma_max; lower=stronger context conditioning).") inf.add_argument("--dblock_blocks", type=int, default=4, help="Number of DiffusionBlocks for the Euler sampler.") inf.add_argument("--ckpt", required=True) inf.add_argument("--device", choices=["auto", "cuda", "cpu"], default="auto", help="Inference compute device. auto uses CUDA when available; cpu forces CPU-only inference.") inf.add_argument("--cpu_threads", type=int, default=0, help="CPU inference intra-op threads. 0=auto, capped at 16; only used when --device resolves to cpu.") inf.add_argument("--cpu_interop_threads", type=int, default=0, help="CPU inference inter-op threads. 0=PyTorch default; only used when --device resolves to cpu.") inf.add_argument("--prompt", default="", help="Prompt text for single-shot inference; optional when --server is set.") inf.add_argument("--max_new", type=int, default=120) inf.add_argument("--min_new", type=int, default=0, help="Minimum generated tokens before EOS can stop decoding. SAT enforces at least one block.") inf.add_argument("--temperature", type=float, default=None) inf.add_argument("--greedy", action="store_true") inf.add_argument("--top_k", type=int, default=None) inf.add_argument("--top_p", type=float, default=0.9) inf.add_argument("--min_p", type=float, default=0.0) inf.add_argument("--repetition_penalty", type=float, default=None) inf.add_argument("--presence_penalty", type=float, default=None) inf.add_argument("--frequency_penalty", type=float, default=None) inf.add_argument("--penalty_last_n", type=int, default=None) inf.add_argument("--var", action="store_true", default=None) inf.add_argument("--no-var", dest="var", action="store_false") inf.add_argument("--claude-friendly", action="store_true", help="Also print an artifact-free prompt/completion block for downstream JSON consumers") inf.add_argument("--plain-output", "--no-color", dest="plain_output", action="store_true", help="Use plain ASCII/no ANSI output for redirected inference logs") inf.add_argument("--attn_backend", choices=["manual", "sdpa", "sublinear"], default=DEFAULT_ATTN_BACKEND) inf.add_argument("--sublinear_window", type=int, default=DEFAULT_SUBLINEAR_WINDOW) inf.add_argument("--sublinear_stride", type=int, default=DEFAULT_SUBLINEAR_STRIDE) inf.add_argument("--sublinear_max_anchors", type=int, default=DEFAULT_SUBLINEAR_MAX_ANCHORS) inf.add_argument("--sublinear_chunk", type=int, default=DEFAULT_SUBLINEAR_CHUNK) inf.add_argument("--sublinear_sinks", type=int, default=DEFAULT_SUBLINEAR_SINKS) inf.add_argument("--sublinear_recent_anchors", type=int, default=DEFAULT_SUBLINEAR_RECENT_ANCHORS) inf.add_argument("--sublinear_pooled_landmarks", action=argparse.BooleanOptionalAction, default=DEFAULT_SUBLINEAR_POOLED_LANDMARKS) inf.add_argument("--no_structured_masks", action="store_true") inf.add_argument("--nat_expand", type=int, default=2) inf.add_argument("--nat_passes", type=int, default=4) inf.add_argument("--nat_greedy", action=argparse.BooleanOptionalAction, default=True, help="Use greedy token picks inside NAT mask-predict refinement by default.") inf.add_argument("--ignore_eos", action="store_true", help="Never stop on (or sample) EOS: suppress its logit and emit exactly max_new tokens. For base-model / SAT-head testing.") # AGILLM-STREAM 20260703: emit [STREAM_BEGIN]/[STREAM_NAT]/[STREAM_AR] # marker lines per committed token so UIs can render the decode live # (NAT renders as a canvas filling in confidence order). inf.add_argument("--stream", action="store_true", help="Emit [STREAM_*] marker lines per committed token for live UI rendering (plain-output only).") # ── SwiReasoning: entropy-gated explicit/latent AR decode ────────────────── inf.add_argument("--swi_reasoning", action="store_true", help="Enable SwiReasoning: alternate between explicit token CoT and silent latent reasoning, gated by next-token entropy. AR + plain KV decode only.") inf.add_argument("--swi_latent_thresh", type=float, default=2.5, help="Entropy (nats) above which an explicit step switches to latent (low confidence -> think silently).") inf.add_argument("--swi_explicit_thresh", type=float, default=1.0, help="Entropy (nats) below which a latent step switches back to explicit (high confidence -> consolidate out loud).") inf.add_argument("--swi_eps", type=float, default=0.05, help="Min entropy delta (nats) to count as a confidence trend when deciding to switch.") inf.add_argument("--swi_max_switches", type=int, default=8, help="Max latent<->explicit switches during thinking phase. After budget is spent decoder stays explicit.") inf.add_argument("--swi_max_latent", type=int, default=16, help="Max consecutive latent steps before forcing back to explicit.") inf.add_argument("--swi_think_budget", type=int, default=256, help="Total reasoning steps (latent+explicit) allowed to switch; after this stays explicit to finish.") inf.add_argument("--swi_max_steps", type=int, default=4096, help="Hard cap on total think_steps (latent+explicit) before stopping.") inf.add_argument("--swi_topk", type=int, default=20, help="Top-k mass to use for the soft thought embedding in latent steps.") inf.add_argument("--swi_start_latent", action="store_true", help="Begin in latent mode instead of explicit (starts silent).") # AGILLM-INFER-SPEED-PORT 20260703: checkpoint cache, skip-init, dtype-cast guard, and load-profile timings. inf.add_argument("--infer_dtype", choices=["fp32", "fp16", "bf16"], default="fp32", help="Resident inference dtype. fp16/bf16 load on CPU, convert, then move the model to CUDA to avoid fp32 VRAM spikes.") inf.add_argument("--block_stream", action="store_true", help="VRAM-saving inference: keep heads/embeddings resident and page Encoder blocks through the compute device.") inf.add_argument("--block_stream_page_layers", type=int, default=1, help="Layers per resident page for --block_stream. 1=lowest VRAM; 0=use --dblock_blocks pages.") inf.add_argument("--block_stream_empty_cache", action=argparse.BooleanOptionalAction, default=True, help="Call torch.cuda.empty_cache() after each streamed page unload.") inf.add_argument("--block_stream_dtype", choices=["fp32", "fp16", "bf16"], default="fp32", help="Weight/activation dtype for --block_stream. fp16 halves CPU->GPU transfer bytes on CUDA-capable cards.") inf.add_argument("--block_stream_kv_cache", action=argparse.BooleanOptionalAction, default=True, help="Use KV cache for AR/SAT --block_stream decode instead of recomputing the full prefix each token.") inf.add_argument("--block_stream_kv_device", choices=["cuda", "cpu"], default="cuda", help="Where --block_stream keeps KV cache tensors. cuda is faster; cpu minimizes resident VRAM.") inf.add_argument("--block_stream_cache_pages", action=argparse.BooleanOptionalAction, default=None, help="Auto by default: keep streamed layer pages resident when VRAM allows. Use --no-block_stream_cache_pages for strict low-VRAM streaming.") inf.add_argument("--moe_expert_stream", action="store_true", help="With --block_stream, keep routed MoE experts on CPU and page only selected experts through the compute device.") inf.add_argument("--moe_expert_stream_empty_cache", action=argparse.BooleanOptionalAction, default=True, help="Call torch.cuda.empty_cache() after unloading each streamed MoE expert.") inf.add_argument("--server", action="store_true", help="Keep one loaded inference instance alive and accept JSON requests on stdin.") sup = sub.add_parser("supervise", help="Native AGILLM4.3 trainer supervisor") sup.add_argument("--save_dir", default="/workspace/agillm4_4090_ckpts") sup.add_argument("--side_dir", default="/workspace/agillm41_side_updates") sup.add_argument("--log", default="/workspace/agillm41_master_train.log") sup.add_argument("--pause_file", default="/tmp/agillm43_master_watchdog.pause") sup.add_argument("--sleep_sec", type=int, default=15) sup.add_argument("--dedupe", action=argparse.BooleanOptionalAction, default=True) sup.add_argument("--once", action="store_true") sup.add_argument("--profile", choices=AGILLM43_PROFILE_CHOICES, default="normal", help="Training launch profile: normal, ar_repair, full_ar_repair, sat_repair, sat_probe, or nat_repair.") hp = sub.add_parser("hotpatch", help="Flush checkpoint and restart under native AGILLM4.3 supervisor") hp.add_argument("--save_dir", default="/workspace/agillm4_4090_ckpts") hp.add_argument("--side_dir", default="/workspace/agillm41_side_updates") hp.add_argument("--log", default="/workspace/agillm41_master_train.log") hp.add_argument("--pause_file", default="/tmp/agillm43_master_watchdog.pause") hp.add_argument("--wait_flush_sec", type=int, default=900) hp.add_argument("--wait_start_sec", type=int, default=300) hp.add_argument("--sleep_sec", type=int, default=15) hp.add_argument("--profile", choices=AGILLM43_PROFILE_CHOICES, default="normal", help="Training launch profile used by the restarted supervisor.") hp.add_argument("--force", action="store_true") hp.add_argument("--tmux", action=argparse.BooleanOptionalAction, default=True) hp.add_argument("--tmux_session", default="master_wd") hp.add_argument("--kill_tmux", action=argparse.BooleanOptionalAction, default=True) hp.add_argument("--nohup_log", default="/workspace/agillm41_native_supervisor.nohup") st = sub.add_parser("status", help="Read-only training status") st.add_argument("--json", dest="json_output", action="store_true") st.add_argument("--log", type=str, default=str(STATUS_DEFAULT_LOG)) st.add_argument("--save_dir", type=str, default=str(STATUS_DEFAULT_SAVE_DIR)) args = ap.parse_args() if args.cmd == "train": train(args) elif args.cmd == "infer": if not getattr(args, "server", False) and not getattr(args, "prompt", ""): ap.error("infer requires --prompt unless --server is set") infer_server(args) if getattr(args, "server", False) else infer(args) elif args.cmd == "supervise": raise SystemExit(supervise_agillm43(args)) elif args.cmd == "hotpatch": raise SystemExit(hotpatch_agillm43(args)) elif args.cmd == "status": raise SystemExit(_emit_status(Path(args.log), Path(args.save_dir), args.json_output)) else: raise SystemExit(f"unknown command: {args.cmd}") if __name__ == "__main__": main() # ===== END nB300_agillm4.py =====