diff --git a/docs/figures/neuron/phase16a/loss_curves.png b/docs/figures/neuron/phase16a/loss_curves.png new file mode 100644 index 0000000..3e6de4a Binary files /dev/null and b/docs/figures/neuron/phase16a/loss_curves.png differ diff --git a/docs/figures/neuron/phase16a/sparsity_trace.png b/docs/figures/neuron/phase16a/sparsity_trace.png new file mode 100644 index 0000000..4053e05 Binary files /dev/null and b/docs/figures/neuron/phase16a/sparsity_trace.png differ diff --git a/notebooks/02-function-level/15-phase16a-rigl-set-regrow.ipynb b/notebooks/02-function-level/15-phase16a-rigl-set-regrow.ipynb new file mode 100644 index 0000000..3796563 --- /dev/null +++ b/notebooks/02-function-level/15-phase16a-rigl-set-regrow.ipynb @@ -0,0 +1,339 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "# 15-phase16a-rigl-set-regrow\n", + "\n", + "**neuron Phase 16a** — paradigm 의 dynamic phase 두 번째 단계. Phase 15 의 *static prune* 을 확장하여 **iterative prune + regrow** 로 학습 중 topology 가 진화하는 DST (Dynamic Sparse Training) 구현.\n", + "\n", + "핵심 가설:\n", + "1. **RigL ≥ static prune?** — 같은 50% sparsity 에서 RigL DST 가 단순 prune 보다 낮은 loss?\n", + "2. **RigL ≥ SET?** — gradient-guided regrow 가 random regrow 보다 우수?\n", + "3. **constant sparsity 유지** — 각 DST cycle 후 sparsity 일정?\n", + "4. **all-finite** — DST cycle 도중 학습 안정성?\n", + "\n", + "설계: full graph block (`hybrid_around_one_around_one`, Phase 14 최저 loss 구조) 위에서 **4 mode × 2 seed = 8 run**.\n", + "- `dense`: prune 없음 (baseline)\n", + "- `static_50`: 50% prune at step 750, regrow 없음 (Phase 15)\n", + "- `dst_set_50`: 50% prune at step 750 + random regrow (SET) every 50 steps, swap 10%\n", + "- `dst_rigl_50`: 50% prune at step 750 + gradient-based regrow (RigL) every 50 steps, swap 10%\n", + "- dst_end_step = 1300 (마지막 200 step 은 stabilize)\n", + "\n", + "데이터: TinyShakespeare (char-LM, block_size=64)\n", + "시드: [42, 123]\n", + "작성일: 2026-05-27\n", + "연관: Issue [#76](https://github.com/EinSofINTEREST/GraphLM/issues/76) (Phase 16 main [#75](https://github.com/EinSofINTEREST/GraphLM/issues/75)) / Phase 15 PR [#72](https://github.com/EinSofINTEREST/GraphLM/pull/72)" + ] + }, + { + "cell_type": "markdown", + "id": "1", + "metadata": {}, + "source": [ + "## 0. 환경 / 의존성" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "from __future__ import annotations\n", + "\n", + "import math\n", + "import statistics\n", + "from pathlib import Path\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import torch\n", + "\n", + "from graphlm.data.tinyshakespeare import (\n", + " CharTokenizer,\n", + " TinyShakespeareDataset,\n", + " load_tinyshakespeare_text,\n", + ")\n", + "from graphlm.neuron.hybrid_transformer_demo import (\n", + " HybridTransformerTrainConfig,\n", + " train_hybrid_transformer_lm,\n", + ")\n", + "from graphlm.utils import safe_perplexity\n", + "\n", + "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "print(f\"device: {device}\")\n", + "print(f\"torch: {torch.__version__}\")" + ] + }, + { + "cell_type": "markdown", + "id": "3", + "metadata": {}, + "source": [ + "## 1. Config + 데이터" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "text = load_tinyshakespeare_text()\n", + "tokenizer = CharTokenizer(text)\n", + "dataset = TinyShakespeareDataset(text, tokenizer)\n", + "vocab_size = tokenizer.vocab_size\n", + "print(f\"vocab_size = {vocab_size}, dataset size = {len(dataset)}\")\n", + "\n", + "# Phase 15 와 동일 hyperparameter (공정 비교)\n", + "HIDDEN_DIM = 128\n", + "N_HEADS = 4\n", + "FFN_DIM = 256\n", + "N_LAYERS = 4\n", + "GROUP_SIZE = 16\n", + "BLOCK_SIZE = 64\n", + "BATCH_SIZE = 32\n", + "LR = 3e-4\n", + "MAX_STEPS = 1500\n", + "PRUNE_AT_STEP = MAX_STEPS // 2 # 750\n", + "DST_PERIOD = 50 # 매 50 step 마다 DST cycle\n", + "DST_SWAP_FRACTION = 0.1 # alive 의 10% swap\n", + "DST_END_STEP = 1300 # 마지막 200 step stabilize\n", + "SEEDS = [42, 123]\n", + "ARCH = \"hybrid_around_one_around_one\"\n", + "\n", + "# 4 mode 정의\n", + "MODES = [\n", + " {\"name\": \"dense\", \"prune_fraction\": 0.0, \"prune_at_step\": None, \"regrow_method\": None},\n", + " {\n", + " \"name\": \"static_50\",\n", + " \"prune_fraction\": 0.5,\n", + " \"prune_at_step\": PRUNE_AT_STEP,\n", + " \"regrow_method\": None,\n", + " },\n", + " {\n", + " \"name\": \"dst_set_50\",\n", + " \"prune_fraction\": 0.5,\n", + " \"prune_at_step\": PRUNE_AT_STEP,\n", + " \"regrow_method\": \"random\",\n", + " },\n", + " {\n", + " \"name\": \"dst_rigl_50\",\n", + " \"prune_fraction\": 0.5,\n", + " \"prune_at_step\": PRUNE_AT_STEP,\n", + " \"regrow_method\": \"rigl\",\n", + " },\n", + "]\n", + "print(\n", + " f\"\\nDST cycle: period={DST_PERIOD}, swap_fraction={DST_SWAP_FRACTION}, end_step={DST_END_STEP}\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "5", + "metadata": {}, + "source": [ + "## 2. Sweep 실행 (4 mode × 2 seed = 8 run)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "results = {}\n", + "for mode in MODES:\n", + " for seed in SEEDS:\n", + " key = (mode[\"name\"], seed)\n", + " print(f\"\\n== mode={mode['name']} seed={seed} ==\")\n", + " cfg = HybridTransformerTrainConfig(\n", + " dataset=dataset,\n", + " vocab_size=vocab_size,\n", + " hidden_dim=HIDDEN_DIM,\n", + " n_heads=N_HEADS,\n", + " ffn_dim=FFN_DIM,\n", + " n_layers=N_LAYERS,\n", + " group_size=GROUP_SIZE,\n", + " arch=ARCH,\n", + " use_full_graph=True,\n", + " block_size=BLOCK_SIZE,\n", + " batch_size=BATCH_SIZE,\n", + " lr=LR,\n", + " max_steps=MAX_STEPS,\n", + " prune_at_step=mode[\"prune_at_step\"],\n", + " prune_fraction=mode[\"prune_fraction\"],\n", + " regrow_method=mode[\"regrow_method\"],\n", + " dst_period=DST_PERIOD if mode[\"regrow_method\"] else None,\n", + " dst_swap_fraction=DST_SWAP_FRACTION,\n", + " dst_end_step=DST_END_STEP if mode[\"regrow_method\"] else None,\n", + " seed=seed,\n", + " device=device,\n", + " )\n", + " out = train_hybrid_transformer_lm(cfg)\n", + " results[key] = out\n", + " n_cycles = len(out[\"dst_cycles\"])\n", + " print(\n", + " f\" final_loss = {out['final_loss']:.4f} (ppl = {safe_perplexity(out['final_loss']):.2f})\"\n", + " f\" sparsity = {out['final_sparsity']:.3f} dst_cycles = {n_cycles}\"\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "7", + "metadata": {}, + "source": [ + "## 3. 결과 표 + 자동 verdict" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": "print(\n f\"{'mode':>12s} {'seed':>6s} {'final_loss':>12s} {'perplexity':>12s} \"\n f\"{'sparsity':>10s} {'dst_cycles':>12s}\"\n)\nprint(\"-\" * 75)\nfor (name, seed), out in results.items():\n fl = out[\"final_loss\"]\n print(\n f\"{name:>12s} {seed:>6d} {fl:>12.4f} {safe_perplexity(fl):>12.2f} \"\n f\"{out['final_sparsity']:>10.3f} {len(out['dst_cycles']):>12d}\"\n )\n\n# mode 별 평균\nprint(\"\\n== Mode summary (mean ± σ across seeds) ==\")\nsummary = {}\nfor mode in MODES:\n name = mode[\"name\"]\n vals = [results[(name, s)][\"final_loss\"] for s in SEEDS]\n sparsities = [results[(name, s)][\"final_sparsity\"] for s in SEEDS]\n m = statistics.mean(vals)\n sd = statistics.stdev(vals) if len(vals) > 1 else 0.0\n summary[name] = (m, sd, statistics.mean(sparsities))\n print(\n f\" {name:>12s} {m:.4f} ± {sd:.4f} (ppl ≈ {safe_perplexity(m):.2f}) \"\n f\"sparsity={statistics.mean(sparsities):.3f}\"\n )\n\n# 자동 verdict\nprint(\"\\n== Verdict ==\")\ndense_loss = summary[\"dense\"][0]\nstatic_loss = summary[\"static_50\"][0]\nset_loss = summary[\"dst_set_50\"][0]\nrigl_loss = summary[\"dst_rigl_50\"][0]\n\n# 1. constant sparsity — static 은 final, DST 는 모든 cycle 별 sparsity 검증\n# (CodeRabbit #3308022917 — final 만 보면 cycle 중간 drift 미감지)\nstatic_ok = abs(summary[\"static_50\"][2] - 0.5) < 0.02\ndst_cycles_ok = all(\n abs(c[\"sparsity_after\"] - 0.5) < 0.02\n for mode_name in (\"dst_set_50\", \"dst_rigl_50\")\n for seed in SEEDS\n for c in results[(mode_name, seed)][\"dst_cycles\"]\n)\nsp_ok = static_ok and dst_cycles_ok\nverdict_1 = \"PASS\" if sp_ok else \"FAIL\"\nprint(\n f\"1. constant sparsity (static final + 모든 DST cycle, target 0.5 ±0.02): \"\n f\"{sp_ok} [{verdict_1}]\"\n)\n\n# 2. RigL ≥ static — DST RigL 이 static prune 보다 ≤ (낮거나 같음)\ndiff_rigl_static = rigl_loss - static_loss\nverdict_2 = \"PASS\" if diff_rigl_static <= 0.02 else \"FAIL\"\nprint(f\"2. RigL ≤ static + 0.02: diff = {diff_rigl_static:+.4f} [{verdict_2}]\")\n\n# 3. RigL ≥ SET — gradient-guided 가 random 보다 ≤\ndiff_rigl_set = rigl_loss - set_loss\nverdict_3 = \"PASS\" if diff_rigl_set <= 0.02 else \"FAIL\"\nprint(f\"3. RigL ≤ SET + 0.02: diff = {diff_rigl_set:+.4f} [{verdict_3}]\")\n\n# 4. all-finite\nall_finite = all(math.isfinite(out[\"final_loss\"]) for out in results.values())\nverdict_4 = \"PASS\" if all_finite else \"FAIL\"\nprint(f\"4. all-finite (DST stability): {all_finite} [{verdict_4}]\")\n\n# DST cycle 정보 (RigL seed=42 기준 샘플)\nrigl_42_cycles = results[(\"dst_rigl_50\", 42)][\"dst_cycles\"]\nprint(f\"\\n== RigL seed=42 의 DST cycles ({len(rigl_42_cycles)}건) ==\")\nfor c in rigl_42_cycles[:3]:\n print(f\" step={c['step']} swap={c['total_swap']} sparsity_after={c['sparsity_after']:.3f}\")\nif len(rigl_42_cycles) > 6:\n print(f\" ... (중략 {len(rigl_42_cycles) - 6}건) ...\")\nfor c in rigl_42_cycles[-3:]:\n print(f\" step={c['step']} swap={c['total_swap']} sparsity_after={c['sparsity_after']:.3f}\")" + }, + { + "cell_type": "markdown", + "id": "9", + "metadata": {}, + "source": [ + "## 4. Loss curve 시각화 — prune + DST cycle 표시" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(1, 1, figsize=(12, 6))\n", + "colors = {\n", + " \"dense\": \"tab:gray\",\n", + " \"static_50\": \"tab:red\",\n", + " \"dst_set_50\": \"tab:orange\",\n", + " \"dst_rigl_50\": \"tab:blue\",\n", + "}\n", + "window = 50\n", + "\n", + "for mode in MODES:\n", + " name = mode[\"name\"]\n", + " losses_per_seed = [results[(name, s)][\"losses\"] for s in SEEDS]\n", + " smoothed = []\n", + " for losses in losses_per_seed:\n", + " smoothed.append(\n", + " [\n", + " sum(losses[max(0, i - window + 1) : i + 1]) / min(i + 1, window)\n", + " for i in range(len(losses))\n", + " ]\n", + " )\n", + " arr = torch.tensor(smoothed)\n", + " mean = arr.mean(dim=0)\n", + " std = arr.std(dim=0)\n", + " steps = list(range(len(mean)))\n", + " ax.plot(steps, mean, label=name, color=colors[name], linewidth=1.5)\n", + " ax.fill_between(steps, mean - std, mean + std, color=colors[name], alpha=0.12)\n", + "\n", + "# prune step + DST end 수직선\n", + "ax.axvline(PRUNE_AT_STEP, color=\"black\", linestyle=\":\", alpha=0.5, label=f\"prune @ {PRUNE_AT_STEP}\")\n", + "ax.axvline(DST_END_STEP, color=\"gray\", linestyle=\":\", alpha=0.5, label=f\"DST end @ {DST_END_STEP}\")\n", + "\n", + "ax.set_xlabel(\"step\")\n", + "ax.set_ylabel(f\"loss (rolling mean w={window})\")\n", + "ax.set_title(\"Phase 16a — DST (RigL vs SET vs static): all 50% sparsity (mean ± σ over 2 seeds)\")\n", + "ax.legend(loc=\"upper right\")\n", + "ax.grid(alpha=0.3)\n", + "plt.tight_layout()\n", + "\n", + "out_dir = Path(\"../../runs/notebook-neuron-phase16a\")\n", + "out_dir.mkdir(parents=True, exist_ok=True)\n", + "fig.savefig(out_dir / \"loss_curves.png\", dpi=150, bbox_inches=\"tight\")\n", + "plt.show()\n", + "print(f\"saved: {out_dir / 'loss_curves.png'}\")" + ] + }, + { + "cell_type": "markdown", + "id": "11", + "metadata": {}, + "source": [ + "## 5. DST cycle sparsity 추적 (RigL)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "fig, ax = plt.subplots(1, 1, figsize=(10, 5))\n", + "for seed in SEEDS:\n", + " cycles = results[(\"dst_rigl_50\", seed)][\"dst_cycles\"]\n", + " steps = [c[\"step\"] for c in cycles]\n", + " sparsities = [c[\"sparsity_after\"] for c in cycles]\n", + " ax.plot(steps, sparsities, marker=\"o\", label=f\"RigL seed={seed}\", linewidth=1.5)\n", + "\n", + "ax.axhline(0.5, color=\"black\", linestyle=\"--\", alpha=0.5, label=\"target sparsity 0.5\")\n", + "ax.set_xlabel(\"step\")\n", + "ax.set_ylabel(\"effective sparsity (mask=0 비율)\")\n", + "ax.set_title(\"Phase 16a — DST cycle 후 sparsity (constant 유지 검증)\")\n", + "ax.legend()\n", + "ax.grid(alpha=0.3)\n", + "ax.set_ylim(0.45, 0.55)\n", + "plt.tight_layout()\n", + "fig.savefig(out_dir / \"sparsity_trace.png\", dpi=150, bbox_inches=\"tight\")\n", + "plt.show()\n", + "print(f\"saved: {out_dir / 'sparsity_trace.png'}\")" + ] + }, + { + "cell_type": "markdown", + "id": "13", + "metadata": {}, + "source": [ + "## 6. 결론 / 다음 단계\n", + "\n", + "(셀 출력 보고 사용자가 채울 영역)\n", + "\n", + "- RigL vs SET vs static 의 final loss ranking?\n", + "- DST 가 static prune 보다 의미 있는 개선?\n", + "- constant sparsity 정확히 유지되는지?\n", + "\n", + "**Phase 16b 진입 시 가설**:\n", + "- 16a 의 *constant sparsity DST* 가 paradigm 안에서 작동 입증 → 16b 의 *capacity expansion (Net2Net)* 도 비슷한 통합 메커니즘 (edge_mask) 위에서 작동 가능?\n", + "- 16a + 16b 결합 → *grow + shrink 동시 dynamic*" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "GraphLM (uv .venv)", + "language": "python", + "name": "graphlm-uv-venv" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/src/graphlm/neuron/graph_hybrid.py b/src/graphlm/neuron/graph_hybrid.py index 16f56f0..89b8499 100644 --- a/src/graphlm/neuron/graph_hybrid.py +++ b/src/graphlm/neuron/graph_hybrid.py @@ -286,6 +286,93 @@ def n_alive_edges(self) -> int: with torch.no_grad(): return int((self.edge_mask > 0).sum().item()) + # ── Phase 16a: edge regrow (DST) ──────────────────────────── + + def n_pruned_edges(self) -> int: + """현재 prune 된 edge 수 (mask=0).""" + with torch.no_grad(): + return int((self.edge_mask == 0).sum().item()) + + def regrow_random(self, n: int, *, reset_weight: bool = True) -> int: + """무작위 n 개 pruned 위치를 un-mask (SET 스타일). + + Args: + n: regrow 할 edge 수 (n_pruned_edges 이상이면 모두 살림). + reset_weight: True 면 regrown 위치의 ``weight`` 를 0 으로 초기화 — 옛 값이 학습 dynamics + 에 노이즈로 작용하지 않도록 (SET/RigL 의 일반적 권장). + + Returns: + 실제 regrow 된 edge 수. + """ + if n < 0: + raise ValueError(f"n must be >= 0, got {n}") + if n == 0: + return 0 + with torch.no_grad(): + pruned_indices = torch.nonzero(self.edge_mask == 0, as_tuple=False) # (n_pruned, 4) + n_pruned = pruned_indices.size(0) + if n_pruned == 0: + return 0 + n_to_regrow = min(n, n_pruned) + # 무작위 n 개 선택 + perm = torch.randperm(n_pruned, device=pruned_indices.device)[:n_to_regrow] + regrow_4d = pruned_indices[perm] + self._activate_edges(regrow_4d, reset_weight=reset_weight) + return n_to_regrow + + def regrow_by_score(self, n: int, scores: Tensor, *, reset_weight: bool = True) -> int: + """``scores`` 기반 상위 n 개 pruned 위치를 un-mask (RigL 스타일). + + RigL: caller 가 dense gradient magnitude 를 scores 로 전달. 가장 큰 신호가 있던 pruned + edge 를 재활성화. + + Args: + n: regrow 할 edge 수. + scores: shape == edge_mask, 각 위치의 점수 (높을수록 regrow 우선). + reset_weight: regrown 위치 ``weight=0`` 초기화 여부. + + Returns: + 실제 regrow 된 edge 수. + """ + if n < 0: + raise ValueError(f"n must be >= 0, got {n}") + if scores.shape != self.edge_mask.shape: + raise ValueError( + f"scores shape {tuple(scores.shape)} != edge_mask shape " + f"{tuple(self.edge_mask.shape)}" + ) + if n == 0: + return 0 + with torch.no_grad(): + pruned_indices = torch.nonzero(self.edge_mask == 0, as_tuple=False) + n_pruned = pruned_indices.size(0) + if n_pruned == 0: + return 0 + n_to_regrow = min(n, n_pruned) + # pruned 위치의 score 만 추출 → top-n 선택 (largest=True) + pruned_scores = scores[ + pruned_indices[:, 0], + pruned_indices[:, 1], + pruned_indices[:, 2], + pruned_indices[:, 3], + ] + _, topk_idx = torch.topk(pruned_scores, n_to_regrow, largest=True) + regrow_4d = pruned_indices[topk_idx] + self._activate_edges(regrow_4d, reset_weight=reset_weight) + return n_to_regrow + + def _activate_edges(self, indices_4d: Tensor, *, reset_weight: bool) -> None: + """주어진 4D 인덱스 위치의 mask=1 + (옵션) weight=0 초기화 (internal helper). + + ``self.weight`` 는 Parameter — `.data` 직접 수정은 autograd 우회로 비권장 + (gemini #3307952463 / Copilot #3307955940). ``with torch.no_grad()`` + in-place 갱신. + """ + i0, i1, i2, i3 = indices_4d[:, 0], indices_4d[:, 1], indices_4d[:, 2], indices_4d[:, 3] + with torch.no_grad(): + self.edge_mask[i0, i1, i2, i3] = 1.0 + if reset_weight: + self.weight[i0, i1, i2, i3] = 0.0 + def extra_repr(self) -> str: return ( f"in_features={self.in_features}, out_features={self.out_features}, " diff --git a/src/graphlm/neuron/hybrid_transformer_demo.py b/src/graphlm/neuron/hybrid_transformer_demo.py index 1fd6415..54fc80e 100644 --- a/src/graphlm/neuron/hybrid_transformer_demo.py +++ b/src/graphlm/neuron/hybrid_transformer_demo.py @@ -15,6 +15,7 @@ from collections.abc import Iterator from dataclasses import dataclass +from typing import Literal import torch import torch.nn.functional as F @@ -68,6 +69,19 @@ class HybridTransformerTrainConfig: # 살아있는 edge 중 하위 magnitude 비율 (prune_bottom_fraction 사용). 0.0 → no-op. prune_fraction: float = 0.0 + # Phase 16a: DST regrow (constant sparsity 유지) + # 동작: regrow_method != None AND dst_period != None 일 때만 DST cycle 실행. + # prune_at_step 이후 dst_period 마다 prune-and-regrow 실행 (dst_end_step 까지). + # regrow_method=None → no regrow (Phase 15 호환, static prune). + regrow_method: Literal["random", "rigl"] | None = None + # DST cycle 주기 (default None = DST cycle 미실행 = Phase 15 호환) + # (Copilot #3307955919 — 주석/동작 일치) + dst_period: int | None = None + # DST cycle 에서 매번 swap 할 alive 비율 (= prune k% + regrow k%, sparsity 보존) + dst_swap_fraction: float = 0.1 + # DST cycle 종료 step (default None = max_steps 까지). 마지막 200 step 정도는 stabilize. + dst_end_step: int | None = None + # runtime seed: int = 0 device: str = "cpu" @@ -81,6 +95,19 @@ def __post_init__(self) -> None: f"prune_at_step must be in [1, max_steps={self.max_steps}], " f"got {self.prune_at_step}" ) + # Phase 16a regrow 검증 + if self.regrow_method not in (None, "random", "rigl"): + raise ValueError( + f"regrow_method must be None / 'random' / 'rigl', got {self.regrow_method!r}" + ) + if not 0.0 <= self.dst_swap_fraction <= 1.0: + raise ValueError(f"dst_swap_fraction must be in [0, 1], got {self.dst_swap_fraction}") + if self.dst_period is not None and self.dst_period < 1: + raise ValueError(f"dst_period must be >= 1, got {self.dst_period}") + if self.dst_end_step is not None and not 1 <= self.dst_end_step <= self.max_steps: + raise ValueError( + f"dst_end_step must be in [1, max_steps={self.max_steps}], got {self.dst_end_step}" + ) class HybridGraphTransformerLM(nn.Module): @@ -203,6 +230,90 @@ def _model_sparsity(model: nn.Module) -> float: return dead / total +# ── Phase 16a: DST cycle helpers ──────────────────────────── + + +def _compute_dense_grad_scores( + model: nn.Module, x: Tensor, y: Tensor, vocab_size: int +) -> dict[str, Tensor]: + """RigL 용 dense gradient magnitude 측정. + + edge_mask 를 일시적으로 모두 1 로 설정 → forward → ``torch.autograd.grad`` 로 weight gradient + 직접 추출 → edge_mask 복원. 모델의 ``.grad`` 속성을 건드리지 않아 train loop 의 gradient + accumulation / 동결 layer 와 안전 (gemini #3307952471). + + Returns: + layer 이름 → score tensor (shape == edge_mask shape). + """ + # 1. mask 저장 + dense 설정 + saved_masks: dict[str, Tensor] = {} + target_modules: dict[str, HybridGraphLinear] = {} + for name, mod in model.named_modules(): + if isinstance(mod, HybridGraphLinear): + saved_masks[name] = mod.edge_mask.clone() + mod.edge_mask.fill_(1.0) + target_modules[name] = mod + try: + # 2. 별도 forward (현재 train step 의 gradient 와 분리) + logits = model(x) + loss = F.cross_entropy(logits.reshape(-1, vocab_size), y.reshape(-1)) + # 3. torch.autograd.grad 로 직접 gradient 추출 (.grad 미사용, side-effect 없음) + weights_to_grad = [ + mod.weight for mod in target_modules.values() if mod.weight.requires_grad + ] + grads = ( + torch.autograd.grad(loss, weights_to_grad, allow_unused=True) if weights_to_grad else [] + ) + # 4. score 추출 + scores: dict[str, Tensor] = {} + grad_iter = iter(grads) + for name, mod in target_modules.items(): + if not mod.weight.requires_grad: + # 동결 layer — 0 score (regrow 우선순위 최하) + scores[name] = torch.zeros_like(mod.weight) + continue + g = next(grad_iter) + scores[name] = ( + g.detach().abs().clone() if g is not None else torch.zeros_like(mod.weight) + ) + return scores + finally: + # 5. mask 복원 — 예외 발생해도 항상 실행 (CodeRabbit #3308022927). 미복원 시 학습이 + # all-1 mask 로 계속되어 paradigm 의 prune 효과 사라짐. + for name, mod in target_modules.items(): + mod.edge_mask.copy_(saved_masks[name]) + + +def _dst_swap_step( + model: nn.Module, + swap_fraction: float, + regrow_method: Literal["random", "rigl"], + rigl_scores: dict[str, Tensor] | None, +) -> dict[str, dict]: + """DST cycle: 각 HybridGraphLinear 의 alive 중 swap_fraction 만큼 prune + 같은 수 regrow. + + constant sparsity 유지 — prune 한 수와 regrow 한 수가 같음. + + Returns: + layer 이름 → {pruned, regrown} count + """ + summary: dict[str, dict] = {} + for name, mod in model.named_modules(): + if not isinstance(mod, HybridGraphLinear): + continue + # prune alive 의 swap_fraction + n_pruned = mod.prune_bottom_fraction(swap_fraction) + # 같은 수 regrow + if regrow_method == "rigl": + if rigl_scores is None or name not in rigl_scores: + raise RuntimeError(f"regrow_method='rigl' 인데 layer '{name}' 의 scores 없음") + n_regrown = mod.regrow_by_score(n_pruned, rigl_scores[name]) + else: # random (SET) + n_regrown = mod.regrow_random(n_pruned) + summary[name] = {"pruned": n_pruned, "regrown": n_regrown} + return summary + + def train_hybrid_transformer_lm(config: HybridTransformerTrainConfig) -> dict: """1 run 학습 — Phase 13/14/15 sweep unit. @@ -231,6 +342,7 @@ def train_hybrid_transformer_lm(config: HybridTransformerTrainConfig) -> dict: optimizer = torch.optim.AdamW(model.parameters(), lr=config.lr) losses: list[float] = [] prune_event: dict | None = None + dst_cycles: list[dict] = [] model.train() for step in range(1, config.max_steps + 1): x, y = next(data_iter) @@ -255,6 +367,33 @@ def train_hybrid_transformer_lm(config: HybridTransformerTrainConfig) -> dict: "sparsity_after": _model_sparsity(model), } + # Phase 16a: DST swap cycle (prune+regrow) after initial prune + if ( + config.regrow_method is not None + and config.dst_period is not None + and config.prune_at_step is not None + and step > config.prune_at_step + and (step - config.prune_at_step) % config.dst_period == 0 + and (config.dst_end_step is None or step <= config.dst_end_step) + and config.dst_swap_fraction > 0 + ): + rigl_scores = None + if config.regrow_method == "rigl": + rigl_scores = _compute_dense_grad_scores(model, x, y, config.vocab_size) + cycle_summary = _dst_swap_step( + model, + swap_fraction=config.dst_swap_fraction, + regrow_method=config.regrow_method, + rigl_scores=rigl_scores, + ) + dst_cycles.append( + { + "step": step, + "total_swap": sum(c["pruned"] for c in cycle_summary.values()), + "sparsity_after": _model_sparsity(model), + } + ) + n_last = min(100, len(losses)) final_loss = sum(losses[-n_last:]) / n_last if n_last > 0 else 0.0 @@ -264,6 +403,7 @@ def train_hybrid_transformer_lm(config: HybridTransformerTrainConfig) -> dict: "final_adj": _snapshot_adj(model), "final_sparsity": _model_sparsity(model), "prune_event": prune_event, + "dst_cycles": dst_cycles, } diff --git a/tests/neuron/test_graph_hybrid.py b/tests/neuron/test_graph_hybrid.py index f1d343a..ec9ca40 100644 --- a/tests/neuron/test_graph_hybrid.py +++ b/tests/neuron/test_graph_hybrid.py @@ -277,6 +277,150 @@ def test_prune_invalid_fraction_rejected(): lin.prune_bottom_fraction(1.5) +# ── Phase 16a: edge regrow (DST) ───────────────────────────── + + +def test_n_pruned_edges_consistent_with_alive(): + lin = HybridGraphLinear(16, 16, group_size=4) + total = lin.edge_mask.numel() + assert lin.n_pruned_edges() + lin.n_alive_edges() == total + + +def test_regrow_random_basic(): + """random regrow 가 정확히 n 개 살리고 alive count 증가.""" + torch.manual_seed(0) + lin = HybridGraphLinear(16, 16, group_size=4) + lin.prune_bottom_fraction(0.5) + alive_before = lin.n_alive_edges() + n_regrown = lin.regrow_random(50, reset_weight=True) + assert n_regrown == 50 + assert lin.n_alive_edges() == alive_before + 50 + + +def test_regrow_random_reset_weight_zeros(): + """reset_weight=True → 새로 살린 위치의 weight=0.""" + torch.manual_seed(0) + lin = HybridGraphLinear(16, 16, group_size=4) + lin.prune_bottom_fraction(0.5) + pruned_before = lin.edge_mask == 0 + n = lin.regrow_random(30, reset_weight=True) + # 새로 살린 위치 = pruned_before AND mask_now>0 + new_alive = pruned_before & (lin.edge_mask > 0) + assert int(new_alive.sum().item()) == n + assert torch.all(lin.weight.data[new_alive] == 0.0) + + +def test_regrow_random_no_reset_preserves_weight(): + """reset_weight=False → weight 값 보존.""" + torch.manual_seed(0) + lin = HybridGraphLinear(16, 16, group_size=4) + lin.prune_bottom_fraction(0.5) + # weight 의 prune 위치 값 기록 + pruned_before = lin.edge_mask == 0 + weight_at_pruned = lin.weight.data[pruned_before].clone() + lin.regrow_random(30, reset_weight=False) + # 일부 위치는 다시 alive 됨 — 그 위치의 weight 가 원래 값과 같은지 (변경 없음) 확인 + # 모든 pruned_before 위치의 weight 가 변경 안 됐어야 함 + assert torch.equal(lin.weight.data[pruned_before], weight_at_pruned) + + +def test_regrow_more_than_pruned_caps_at_pruned(): + """n > n_pruned 이면 모든 pruned 만큼만 regrow.""" + torch.manual_seed(0) + lin = HybridGraphLinear(16, 16, group_size=4) + lin.prune_bottom_fraction(0.3) + n_pruned = lin.n_pruned_edges() + n = lin.regrow_random(99999) + assert n == n_pruned + assert lin.n_pruned_edges() == 0 + + +def test_regrow_zero_noop(): + lin = HybridGraphLinear(16, 16, group_size=4) + lin.prune_bottom_fraction(0.5) + assert lin.regrow_random(0) == 0 + + +def test_regrow_negative_rejected(): + lin = HybridGraphLinear(16, 16, group_size=4) + with pytest.raises(ValueError, match="n must be >= 0"): + lin.regrow_random(-1) + + +def test_regrow_by_score_picks_top_n(): + """scores 기반 regrow 가 정확히 top-n 위치 선택.""" + torch.manual_seed(0) + lin = HybridGraphLinear(16, 16, group_size=4) + lin.prune_bottom_fraction(0.5) + # scores: pruned 위치마다 다른 값 (rank 가능) + scores = torch.zeros_like(lin.edge_mask) + pruned_indices = torch.nonzero(lin.edge_mask == 0, as_tuple=False) + # 첫 10 위치에 가장 큰 score + for rank, (i0, i1, i2, i3) in enumerate(pruned_indices[:10]): + scores[i0, i1, i2, i3] = 100.0 - rank + n = lin.regrow_by_score(5, scores) + assert n == 5 + # 첫 5 위치만 alive 됐는지 + for i0, i1, i2, i3 in pruned_indices[:5]: + assert lin.edge_mask[i0, i1, i2, i3] == 1.0 + # 다음 5 위치는 여전히 pruned + for i0, i1, i2, i3 in pruned_indices[5:10]: + assert lin.edge_mask[i0, i1, i2, i3] == 0.0 + + +def test_regrow_by_score_wrong_shape_rejected(): + lin = HybridGraphLinear(16, 16, group_size=4) + lin.prune_bottom_fraction(0.5) + bad_scores = torch.zeros(10) + with pytest.raises(ValueError, match="shape"): + lin.regrow_by_score(3, bad_scores) + + +def test_regrow_then_forward_gradient_flows(): + """regrow 후 forward 가 새 edge 에서 정상 grad 흐름 (mask=1 효과).""" + torch.manual_seed(0) + lin = HybridGraphLinear( + 16, + 16, + group_size=4, + adj_outer_init="uniform_around_one", + adj_inner_init="uniform_around_one", + ) + lin.prune_bottom_fraction(0.5) + # regrow 한 위치 + pruned_before = lin.edge_mask == 0 + lin.regrow_random(30, reset_weight=True) + regrown_positions = pruned_before & (lin.edge_mask > 0) + # forward + backward — regrown 위치는 weight=0 이지만 mask=1 이라 grad 는 흘러야 함 + x = torch.randn(4, 16) + lin(x).sum().backward() + # regrown 위치의 weight grad 가 nonzero (= 살아있음 증명) + grad_at_regrown = lin.weight.grad[regrown_positions] + assert (grad_at_regrown.abs().sum() > 0).item(), ( + "regrow 후 새 edge 의 weight grad 가 모두 0 — gradient 가 안 흐름" + ) + + +def test_constant_sparsity_prune_then_regrow(): + """DST cycle: 같은 수 prune + regrow → sparsity 정확히 일정 (Copilot #3307955897 강화).""" + torch.manual_seed(0) + lin = HybridGraphLinear(16, 16, group_size=4) + lin.prune_bottom_fraction(0.5) + sparsity_before_cycle = lin.effective_sparsity() + alive_before_cycle = lin.n_alive_edges() + + # DST cycle: 두 번째 prune 의 return 값을 받아 그 수만큼 정확히 regrow + n_pruned_in_cycle = lin.prune_bottom_fraction(0.2) # alive 의 20% + n_regrown = lin.regrow_random(n_pruned_in_cycle) + assert n_regrown == n_pruned_in_cycle, "DST cycle: regrow 수가 prune 수와 일치해야 함" + + # 정확히 동일 sparsity / alive 유지 (constant sparsity) + assert lin.effective_sparsity() == sparsity_before_cycle, ( + f"DST 후 sparsity 불일치: {lin.effective_sparsity()} vs {sparsity_before_cycle}" + ) + assert lin.n_alive_edges() == alive_before_cycle + + def test_edge_mask_in_state_dict(): """edge_mask 가 state_dict 에 포함되어 save/load 보존.""" lin1 = HybridGraphLinear(16, 16, group_size=4)