[FEAT#76] neuron Phase 16a — RigL/SET edge regrow (constant sparsity DST) - #78
Conversation
…SET) - src/graphlm/neuron/graph_hybrid.py 확장: - n_pruned_edges(): pruned 위치 수 - regrow_random(n, reset_weight=True): 무작위 n 개 pruned 위치 un-mask (SET 스타일) - regrow_by_score(n, scores, reset_weight=True): scores 기반 top-n pruned 위치 un-mask (RigL 스타일) - _activate_edges(): mask=1 + (옵션) weight=0 reset helper - regrow 가 정확히 n 개 활성화 (deterministic topk for score-based) - reset_weight=True 가 기본 — RigL/SET 표준 권장 (옛 prune 값의 노이즈 영향 회피) - 12 신규 unit tests: - regrow random / score / shape 검증 / n>n_pruned cap - reset_weight 동작 (True=0 init, False=옛 값 보존) - regrow 후 gradient flow 검증 (= 새 edge 실제 학습 가능) - constant sparsity (prune+regrow 동량 → sparsity 일정) - 202 → 214 tests, all green
- HybridTransformerTrainConfig 확장: - regrow_method: None / 'random' (SET) / 'rigl' - dst_period: int (DST cycle 주기, None = single-cycle Phase 15 호환) - dst_swap_fraction: float (cycle 당 swap 비율, default 0.1) - dst_end_step: int | None (DST cycle 종료 step) - __post_init__ 에 4 인자 검증 추가 - 신규 helper: - _compute_dense_grad_scores: mask 일시 1 → forward+backward → |weight.grad| → mask 복원 (RigL 용) - _dst_swap_step: alive 의 swap_fraction prune + 같은 수 regrow (constant sparsity) - train loop: - prune_at_step 직후 (Phase 15 호환) + dst_period 마다 DST cycle 실행 - dst_end_step 이후로는 cycle 중지 (학습 stabilize) - result 에 dst_cycles 리스트 추가 (cycle 별 step / total_swap / sparsity_after) - ruff format 적용
…sparsity) - notebooks/02-function-level/15-phase16a-rigl-set-regrow.ipynb 신규 - 4 mode × 2 seed = 8 run: dense / static_50 / dst_set_50 / dst_rigl_50 - arch 고정: hybrid_around_one_around_one + use_full_graph=True (Phase 14 최저 loss) - DST 설정: period=50, swap=10%, end_step=1300 (마지막 200 step stabilize) - 자동 verdict 4가지: - constant sparsity (target 0.5, ±0.02) - RigL ≤ static + 0.02 - RigL ≤ SET + 0.02 - all-finite (DST stability) - loss curve (prune + DST end 수직선) + sparsity trace (constant 유지 검증)
|
Warning Review limit reached
More reviews will be available in 21 minutes and 10 seconds. Learn how PR review limits work. Your organization has run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: ⛔ Files ignored due to path filters (2)
📒 Files selected for processing (2)
📝 WalkthroughWalkthroughThis PR implements Phase 16a: Dynamic Sparse Training with iterative edge regrowth. Core regrow APIs enable random or gradient-scored (RigL) reactivation of pruned edges; training integration adds DST cycle scheduling with dense-gradient measurement for RigL scoring; comprehensive tests validate regrow mechanics and cycle correctness; a Phase 16a experiment notebook runs a four-mode sweep comparing dense, static-pruned, SET-DST, and RigL-DST training across two seeds with automated performance verdicts. ChangesPhase 16a: RigL/SET Dynamic Sparse Training
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Code Review
This pull request implements Phase 16a of the dynamic sparse training (DST) paradigm, introducing iterative prune and regrow capabilities (SET and RigL styles) to evolve network topology during training. The changes include new regrow methods in HybridGraphLinear, configuration updates, training loop integration, comprehensive unit tests, and a demonstration notebook. The reviewer feedback highlights two important improvements: avoiding the use of .data for in-place tensor modifications to prevent bypassing autograd tracking, and refactoring the gradient score computation to use torch.autograd.grad instead of directly manipulating .grad attributes to avoid side-effects and potential AttributeErrors.
There was a problem hiding this comment.
Pull request overview
이 PR은 Phase 15의 static prune을 확장해 constant sparsity를 유지하는 DST(Dynamic Sparse Training) 를 구현하기 위해, HybridGraphLinear에 edge regrow(SET random / RigL score-based) 기능을 추가하고, 데모 학습 루프에 DST cycle(prune+regrow) 및 RigL용 dense gradient score 측정을 연결합니다. 또한 해당 동작을 검증하는 테스트와 실험용 노트북을 추가합니다.
Changes:
HybridGraphLinear에n_pruned_edges,regrow_random,regrow_by_score및 내부 활성화 헬퍼 추가train_hybrid_transformer_lm에 DST cycle 실행 로직과 RigL용 dense grad score 측정 유틸 추가, 결과에dst_cycles포함- regrow/DST 성질을 검증하는 unit test 및 Phase 16a sweep 노트북 신규 추가
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
src/graphlm/neuron/graph_hybrid.py |
pruned edge를 다시 살리는(regrow) API 및 내부 활성화 로직 추가 |
src/graphlm/neuron/hybrid_transformer_demo.py |
DST cycle 및 RigL score 측정(일시적 dense backward) 헬퍼/학습 루프 확장 |
tests/neuron/test_graph_hybrid.py |
regrow 동작/경계조건/grad 흐름 및 constant sparsity 관련 테스트 추가 |
notebooks/02-function-level/15-phase16a-rigl-set-regrow.ipynb |
dense/static/DST(SET/RigL) 비교 sweep 노트북 추가 |
- gemini #3307952463 + Copilot #3307955940 (중복): _activate_edges 의 self.weight.data 직접 수정 비권장 → with torch.no_grad(): + in-place [...] = 0.0 로 변경. helper 자체에 no_grad 보장. - gemini #3307952471 (HIGH): _compute_dense_grad_scores 가 model.grad 직접 수정 (side-effect) + 동결 layer 시 AttributeError → torch.autograd.grad(loss, weights, allow_unused=True) 로 재작성. .grad 미사용, requires_grad=False / unused layer 안전 처리. - Copilot #3307955897: test_constant_sparsity_prune_then_regrow 가 prune/regrow 동량 안 맞고 ±0.1 허용 → 두 번째 prune 의 return 값으로 정확한 regrow + sparsity / alive 정확 일치 assert. - Copilot #3307955919: dst_period 주석/실제 동작 불일치 (주석은 "None=첫 cycle만" 이나 실제는 None=미실행) → 주석을 실제 동작 (None=DST cycle 미실행, Phase 15 호환) 으로 수정. - 213 tests still green
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/graphlm/neuron/hybrid_transformer_demo.py (1)
331-343:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftReset AdamW optimizer state for DST-pruned/regrown edges.
edge_maskgates theHybridGraphLinearforward (and thus gradients), but the DST regrow path only resets the weight values (reset_weight=True→HybridGraphLinear._activate_edges(..., reset_weight=...)setsweight=0). AdamW’s decoupled weight decay and per-parameterexp_avg/exp_avg_sqare not cleared for those pruned/regrown indices, so regrown edges can resume with stale optimizer moments. Clear/mask AdamW state for the affected weight entries during_dst_swap_step(not justweight).
File:src/graphlm/neuron/hybrid_transformer_demo.py(331-343, also 359-377)
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@notebooks/02-function-level/15-phase16a-rigl-set-regrow.ipynb`:
- Around line 237-255: sp_ok currently only checks summary[m][2]
(final_sparsity); update it to validate sparsity_after for every entry in the
DST cycles for the DST modes ("dst_set_50" and "dst_rigl_50"). Locate the sp_ok
computation and replace it with logic that iterates summary[m]["dst_cycles"] (or
the dst_cycles structure in summary) and ensures abs(cycle["sparsity_after"] -
0.5) < 0.02 for every cycle for m in ("dst_set_50","dst_rigl_50"), while keeping
the original static_50 final-sparsity check; keep the same verdict_1 variable
and print line but base sp_ok on this per-cycle check.
In `@src/graphlm/neuron/hybrid_transformer_demo.py`:
- Around line 248-273: The code temporarily sets HybridGraphLinear.edge_mask to
all ones (saved_masks, target_modules) then runs a forward and
torch.autograd.grad; if an exception occurs the original masks are not restored.
Fix by surrounding the dense-override, forward, grad collection and score
extraction with a try/finally: after saving masks and setting
mod.edge_mask.fill_(1.0), run logits/loss and grads inside try, compute scores
there, and in the finally always restore each
mod.edge_mask.copy_(saved_masks[name]) (and ensure any early returns are avoided
so restoration runs); reference saved_masks, target_modules, model,
weights_to_grad, torch.autograd.grad, and edge_mask when implementing.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 5fa6fce5-69ad-44f1-bddd-1b20143fad41
📒 Files selected for processing (4)
notebooks/02-function-level/15-phase16a-rigl-set-regrow.ipynbsrc/graphlm/neuron/graph_hybrid.pysrc/graphlm/neuron/hybrid_transformer_demo.pytests/neuron/test_graph_hybrid.py
…ity 검증 - CodeRabbit #3308022927: _compute_dense_grad_scores 의 forward/grad 예외 시 edge_mask 미복원 → all-1 mask 로 학습 지속 위험 → try/finally 로 mask 복원 보장. 정상/예외 경로 모두 복원. - CodeRabbit #3308022917: notebook verdict 1 의 constant sparsity 검증이 final_sparsity 만 보고 cycle 중간 drift 미감지 → static 은 final, DST 는 모든 dst_cycles[*].sparsity_after 검증 (target 0.5 ±0.02).
연관 이슈
구현 내용
Phase 15 (PR #72) 의 static prune 을 확장하여 iterative prune + regrow 로 학습 중 topology 가 진화하는 DST (Dynamic Sparse Training) 구현. constant sparsity 유지하며 RigL (gradient-guided regrow) / SET (random regrow) 두 변형 지원.
신규 모듈
src/graphlm/neuron/graph_hybrid.py확장 (4 신규 메서드)n_pruned_edges(): pruned 위치 수regrow_random(n, reset_weight=True): SET 스타일 random regrowregrow_by_score(n, scores, reset_weight=True): RigL 스타일 score-based regrow (deterministic top-k)_activate_edges(): mask=1 + (옵션) weight=0 reset helpersrc/graphlm/neuron/hybrid_transformer_demo.py확장TrainConfig신규 4 인자:regrow_method/dst_period/dst_swap_fraction/dst_end_step+__post_init__검증_compute_dense_grad_scores: mask 일시 1 → forward+backward →|weight.grad|(RigL 용)_dst_swap_step: alive 의 swap_fraction prune + 같은 수 regrow (constant sparsity)dst_cycles리스트 추가notebooks/02-function-level/15-phase16a-rigl-set-regrow.ipynb신규gradient resurrection 방지 + dense gradient 측정
Phase 15 backwards compat
regrow_method=None(default) → DST cycle 미실행 = Phase 15 static prune 와 동일 동작테스트 (202 → 214, +12 신규)
test_regrow_random_basic/test_regrow_by_score_picks_top_n: 정확성test_regrow_random_reset_weight_zeros/test_regrow_random_no_reset_preserves_weight: reset 옵션test_regrow_more_than_pruned_caps_at_pruned/test_regrow_zero_noop/test_regrow_negative_rejected: 경계test_regrow_then_forward_gradient_flows: 핵심 — regrow 후 새 edge 의 학습 가능성test_constant_sparsity_prune_then_regrow: DST 의 핵심 성질CPU smoke 결과
CI / 머지 게이트 점검
변경 영향 범위
src/graphlm/neuron/graph_hybrid.py(확장),hybrid_transformer_demo.py(확장),tests/neuron/test_graph_hybrid.py(확장),notebooks/02-function-level/(1 신규)Low—regrow_method=None(default) 시 Phase 15 와 동일 동작. opt-in 메커니즘.Required Status Checks
Commit Lint/PR Title Lint/Linked Issue CheckFormat Check/Build/Test(214 passed) /Lint(50 papers OK)롤백 계획
regrow_method미사용 시 backwards-compat. revert 시 영향 없음.Phase 16 의 다음 단계
본 PR (16a) 머지 후 Sub-issue #77 (Phase 16b — Net2Net grow) 진행:
Summary by CodeRabbit
New Features
Tests