[FEAT#67] neuron Phase 13 — HybridGraphLinear 의 Transformer FFN 통합 + RMSNorm - #68
Conversation
- src/graphlm/neuron/rms_norm.py 신규 — LLaMA / Mistral 표준 norm - mean centering 생략 + bias 없음 (LayerNorm 대비 파라미터 ↓, 안정성 동등) - weight=1 초기화로 function preservation 보장 - mixed precision 안전성 위해 float32 cast 후 RMS 계산 - 8 unit tests (shape / scaling / gradient / 입력 검증)
- src/graphlm/neuron/hybrid_transformer.py 신규 - HybridGraphFFN: HybridGraphLinear 로 fc1 + fc2 (GELU) - HybridGraphTransformerBlock: pre-norm (RMSNorm + attn + RMSNorm + FFN) - PlainFFN / PlainTransformerBlock: 공정 비교용 baseline (norm 은 RMSNorm 통일) - make_block: 4 arch dispatch (plain / hybrid_full_full / full_around_one / around_one_around_one) - identity outer 는 FFN 의 rectangular 구조상 미지원 (ValueError 명시) - 14 unit tests — FFN 동치성 / block 동치성 / gradient flow / arch dispatch - __init__.py exports 갱신
… helper - src/graphlm/neuron/hybrid_transformer_demo.py 신규 - HybridGraphTransformerLM: token+pos emb → N blocks → RMSNorm → lm_head - train_hybrid_transformer_lm: 1 run unit (4 arch dispatch) - count_parameters: arch 간 파라미터 수 reporting - _snapshot_adj: hybrid arch 의 block 별 FFN fc1/fc2 outer/inner snapshot (Phase 12 hierarchy 호환) - ruff format / unused import 정리 - 120 tests all green
…_around_one / around_one_around_one)
- notebooks/02-function-level/12-phase13-hybrid-transformer.ipynb 신규
- 4 arch × 2 seed = 8 run, max_steps=1500, TinyShakespeare char-LM
- model: hidden=128, ffn=256, n_layers=4, n_heads=4, block_size=64
- 자동 verdict 4 가지: function preservation / scale-corrected ≤ plain / full around_one 안정성 / RMSNorm NaN 없음
- loss curve (mean ± σ rolling) + hybrid adj heatmap (outer / inner)
- output figure → runs/notebook-neuron-phase13/{loss_curves,hybrid_adj}.png
|
Warning Review limit reached
More reviews will be available in 38 minutes and 9 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 selected for processing (4)
📝 WalkthroughWalkthroughPhase 13 introduces RMSNorm and new transformer blocks that use HybridGraphLinear in the FFN path, plain baselines, a small HybridGraphTransformerLM and training harness, tests verifying function preservation and gradients, updated public exports, and a notebook that sweeps four architecture variants on TinyShakespeare reporting losses, verdicts, and adjacency visualizations. ChangesPhase 13 Hybrid Transformer
Sequence Diagram(s)sequenceDiagram
participant Trainer
participant Model as HybridGraphTransformerLM
participant Blocks as TransformerBlocks
participant FFN as HybridGraphFFN
participant Optimizer as AdamW
Trainer->>Model: train_hybrid_transformer_lm(config)
Model->>Model: set seed, build model, move to device
loop for each batch (max_steps)
Trainer->>Model: Forward pass (tokens)
Model->>Blocks: token/pos embed -> block stack
Blocks->>Blocks: RMSNorm -> CausalSelfAttention -> residual add
Blocks->>FFN: RMSNorm -> HybridGraphFFN -> residual add
FFN->>FFN: HybridGraphLinear(fc1 adj_outer/adj_inner) -> GELU -> HybridGraphLinear(fc2)
Blocks-->>Model: block outputs
Model->>Model: final RMSNorm -> lm_head -> logits
Trainer->>Optimizer: loss -> backward -> step
Optimizer->>FFN: update weight, adj_outer, adj_inner
Trainer->>Model: record per-step loss
end
Trainer->>Model: snapshot_adj() (hybrid only)
Model-->>Trainer: per-block adj_outer/adj_inner tensors
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
🚥 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 13 of the project, integrating the HybridGraphLinear layer into the standard Transformer FFN position and introducing RMSNorm as a modern normalization layer. It includes the core implementations of HybridGraphFFN, HybridGraphTransformerBlock, PlainTransformerBlock, and RMSNorm, along with a demo notebook, training helpers, and comprehensive unit tests. The review feedback highlights two important issues: a potential mixed-precision dtype mismatch in RMSNorm.forward when multiplying by self.weight, and a potential OverflowError in the notebook when calculating perplexity using math.exp on high loss values.
There was a problem hiding this comment.
Pull request overview
이 PR은 Phase 12에서 확립된 HybridGraphLinear를 Transformer FFN 위치에 통합하고, norm을 RMSNorm으로 교체한 Phase 13 실험용 모듈/데모/테스트/노트북을 추가합니다. plain 대비 hybrid_* 아키텍처를 동일한 블록 구조에서 공정 비교할 수 있게 구성되어 neuron 실험 라인을 Transformer 백본으로 확장하는 변경입니다.
Changes:
RMSNorm신규 구현 및 단위 테스트 추가HybridGraphFFN/HybridGraphTransformerBlock및make_block(4-arch dispatch) 신규 추가 + 동치성/그라디언트 테스트 추가- Phase 13 sweep 실행을 위한 demo LM/train helper 및 실험 노트북 추가,
neuron.__init__export 갱신
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
src/graphlm/neuron/rms_norm.py |
Phase 13용 RMSNorm 구현 추가 |
src/graphlm/neuron/hybrid_transformer.py |
Hybrid FFN + pre-norm Transformer block + arch dispatch 추가 |
src/graphlm/neuron/hybrid_transformer_demo.py |
char-LM 데모 모델/학습 헬퍼 및 adj 스냅샷 유틸 추가 |
src/graphlm/neuron/__init__.py |
Phase 13 신규 심볼 export 추가 |
tests/neuron/test_rms_norm.py |
RMSNorm 기능/검증 테스트 추가 |
tests/neuron/test_hybrid_transformer.py |
FFN/블록 동치성, 그라디언트, dispatch 테스트 추가 |
notebooks/02-function-level/12-phase13-hybrid-transformer.ipynb |
4 arch × 2 seed sweep 및 시각화/자동 verdict 노트북 추가 |
- src/graphlm/neuron/rms_norm.py — RMSNorm.forward 에서 self.weight 를 x_dtype 으로 cast (gemini #3303153077): mixed precision (FP16/BF16) 입력 시 weight 가 float32 로 남아 있어 출력 dtype 이 강제로 float32 promotion 되던 문제 — residual connection dtype mismatch 회피 - notebooks/02-function-level/12-phase13-hybrid-transformer.ipynb — safe_perplexity 헬퍼 도입 (gemini #3303153101): 학습 초반 큰 loss 에서 math.exp OverflowError 회피 (cap=20.0)
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@src/graphlm/neuron/hybrid_transformer_demo.py`:
- Around line 113-130: Replace the long parameter list on
train_hybrid_transformer_lm with a single dataclass (e.g.,
HybridTransformerTrainConfig) that contains the current scalars: dataset:
TinyShakespeareDataset, vocab_size: int, seed: int, arch: Arch, hidden_dim: int,
n_heads: int, ffn_dim: int, n_layers: int, group_size: int, block_size: int,
batch_size: int, lr: float, max_steps: int, device: str = "cpu", dropout: float
= 0.0; update the function signature to accept that config object
(train_hybrid_transformer_lm(config: HybridTransformerTrainConfig)) and change
all internal references from individual parameters (hidden_dim, n_heads, lr,
etc.) to config.hidden_dim, config.n_heads, config.lr, etc.; add the dataclass
import and the new type to annotations and update any call sites/tests to
construct and pass the new config instance when calling
train_hybrid_transformer_lm.
🪄 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: 351e0665-b8a9-4bc2-8c1b-f623db17cc70
📒 Files selected for processing (7)
notebooks/02-function-level/12-phase13-hybrid-transformer.ipynbsrc/graphlm/neuron/__init__.pysrc/graphlm/neuron/hybrid_transformer.pysrc/graphlm/neuron/hybrid_transformer_demo.pysrc/graphlm/neuron/rms_norm.pytests/neuron/test_hybrid_transformer.pytests/neuron/test_rms_norm.py
- src/graphlm/neuron/rms_norm.py (Copilot #3303168589): floating-point 입력 검증 - int 등 비-floating 입력 silent cast 회피 — TypeError 명시 (LayerNorm 정책과 동일) - src/graphlm/neuron/hybrid_transformer.py (Copilot #3303168649 / #3303168686): - HybridGraphFFN / PlainFFN 둘 다 dropout 인자 추가 — backbone.FFN 의 fc2-뒤-dropout 패턴 일관 - HybridGraphTransformerBlock / PlainTransformerBlock 의 dropout 인자가 attention 만이 아니라 FFN 까지 전달 - 7 신규 unit tests — non-floating 거부 / dtype roundtrip / FFN dropout / block dropout 전파 - 152 tests all green
- CodeRabbit #3303186824 — project rule "5+ args → dataclass" 적용 - src/graphlm/neuron/hybrid_transformer_demo.py: - HybridTransformerTrainConfig (frozen dataclass) 신규 — data/model/train/runtime 4 그룹 - train_hybrid_transformer_lm(config) 단일 인자로 단순화 (기존 14 kwargs) - notebooks/02-function-level/12-phase13-hybrid-transformer.ipynb §2 — config 객체 구성 후 train 호출 - 152 tests all green
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
src/graphlm/neuron/hybrid_transformer_demo.py (1)
179-179: ⚡ Quick winExtract magic number to a named constant.
The hardcoded
100violates the guideline "Replace magic numbers with named constants" and makes the averaging window less discoverable.As per coding guidelines, magic numbers should be replaced with named constants.
♻️ Proposed fix
Add a module-level constant near the top of the file (after imports):
# 학습 마지막 N step 의 loss 평균 계산 window _FINAL_LOSS_AVERAGING_WINDOW = 100Then update line 179:
- n_last = min(100, len(losses)) + n_last = min(_FINAL_LOSS_AVERAGING_WINDOW, len(losses))🤖 Prompt for 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. In `@src/graphlm/neuron/hybrid_transformer_demo.py` at line 179, Replace the magic number 100 used when computing n_last (n_last = min(100, len(losses))) with a module-level named constant (e.g., _FINAL_LOSS_AVERAGING_WINDOW) declared near the top of src/graphlm/neuron/hybrid_transformer_demo.py (after imports); then change the assignment to n_last = min(_FINAL_LOSS_AVERAGING_WINDOW, len(losses)) so the averaging window is discoverable and configurable while keeping the existing behavior of using the smaller of the window and len(losses).
🤖 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/12-phase13-hybrid-transformer.ipynb`:
- Line 43: Extract the safe_perplexity function into a reusable module under the
graphlm package (e.g., create graphlm.utils or graphlm.metrics) keeping the same
signature (safe_perplexity(loss: float, cap: float = 20.0) -> float) and
tests/docstring, export it, then remove the in-notebook definition and replace
it with an import like `from graphlm.utils import safe_perplexity` in the
notebook; ensure any references to safe_perplexity (e.g., calls in plotting or
evaluation code) continue to work unchanged and update package __init__ if
needed to expose the symbol.
- Line 131: The stability checks currently use math.isnan which treats ±inf as
valid; update the verdict logic to use math.isfinite for numeric-finiteness
checks: change the full stability check (variable verdict_3 that tests aa_loss
with not math.isnan(aa_loss) and aa_loss < plain_loss + 0.5) to require
math.isfinite(aa_loss) and the final RMSNorm stability aggregator (all_finite
computed with not math.isnan(out["final_loss"])) to use
math.isfinite(out["final_loss"]) so infinities fail the stability verdicts.
---
Nitpick comments:
In `@src/graphlm/neuron/hybrid_transformer_demo.py`:
- Line 179: Replace the magic number 100 used when computing n_last (n_last =
min(100, len(losses))) with a module-level named constant (e.g.,
_FINAL_LOSS_AVERAGING_WINDOW) declared near the top of
src/graphlm/neuron/hybrid_transformer_demo.py (after imports); then change the
assignment to n_last = min(_FINAL_LOSS_AVERAGING_WINDOW, len(losses)) so the
averaging window is discoverable and configurable while keeping the existing
behavior of using the smaller of the window and len(losses).
🪄 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: 04e28d5a-75f8-49d4-a6b3-55104e41e5e1
⛔ Files ignored due to path filters (2)
docs/figures/neuron/phase13/hybrid_adj.pngis excluded by!**/*.pngdocs/figures/neuron/phase13/loss_curves.pngis excluded by!**/*.png
📒 Files selected for processing (6)
notebooks/02-function-level/12-phase13-hybrid-transformer.ipynbsrc/graphlm/neuron/hybrid_transformer.pysrc/graphlm/neuron/hybrid_transformer_demo.pysrc/graphlm/neuron/rms_norm.pytests/neuron/test_hybrid_transformer.pytests/neuron/test_rms_norm.py
- CodeRabbit #3304306120 — safe_perplexity 를 노트북 → src/graphlm/utils/metrics.py 로 이전 (project rule: 노트북은 analysis 만, 로직은 src/) - graphlm.utils 에서 export, 8 unit tests 추가 - 노트북은 `from graphlm.utils import safe_perplexity` 만 사용 - CodeRabbit #3304306127 — verdict 안정성 검사를 isnan → isfinite 로 교체 (verdict 3, 4) - isnan 만 쓰면 ±inf 가 silent PASS — 학습 발산 미감지 위험 - 152 → 160 tests all green
연관 이슈
구현 내용
Phase 12 에서
HybridGraphLinear(outer group + inner channel dual routing) 를 paradigm 의 ultimate 표현으로 확립. Phase 13 은 이를 실제 Transformer 의 FFN 위치 에 통합하고 RMSNorm 을 도입.신규 모듈
src/graphlm/neuron/rms_norm.py—RMSNorm(LLaMA / Mistral 표준 norm)src/graphlm/neuron/hybrid_transformer.py— Transformer block 통합HybridGraphFFN: HybridGraphLinear fc1 + fc2 + GELUHybridGraphTransformerBlock: pre-norm (RMSNorm + CausalSelfAttention + RMSNorm + FFN + residual)PlainFFN/PlainTransformerBlock: 공정 비교용 baseline (norm 은 RMSNorm 통일)make_block: 4 arch dispatchsrc/graphlm/neuron/hybrid_transformer_demo.py— LM + 학습 helperHybridGraphTransformerLM: small char-LM (token+pos emb → N blocks → RMSNorm → lm_head)train_hybrid_transformer_lm: 1 run sweep unit_snapshot_adj: block 별 fc1/fc2 outer/inner snapshot (Phase 12 hierarchy 호환)notebooks/02-function-level/12-phase13-hybrid-transformer.ipynb— 4 arch × 2 seed sweep + 자동 verdict + 시각화4 가지 arch
plain— RMSNorm + 표준 nn.Linear FFN (baseline)hybrid_full_full— function preserving baselinehybrid_full_around_one— Phase 11 channel-level scale-correctedhybrid_around_one_around_one— outer/inner 둘 다 scale-correctedidentity outer 미지원 — FFN 은 rectangular (hidden ≠ ffn) 구조라 정방 identity 정의 불가. ValueError 명시.
테스트 (122 → 145, +23 신규)
test_rms_norm.py— 8 tests (shape / RMS unit / weight learnable / 입력 검증)test_hybrid_transformer.py— 15 tests (FFN/block function preservation, gradient flow, 4 arch dispatch)0-init 금지 + magnitude rule 일관 적용
HybridGraphFFN도 underlying HybridGraphLinear 의 ValueError 그대로 상속RMSNorm의 weight 는 1 로 시작 (0-init 회피, multiplier 위치 magnitude rule)CI / 머지 게이트 점검
변경 영향 범위
src/graphlm/neuron/(3 신규 파일),notebooks/02-function-level/(1 신규 노트북),tests/neuron/(2 신규 테스트 파일)Low— 신규 모듈 추가만, 기존 모듈 변경 없음.__init__.pyexports 만 갱신.Required Status Checks
Commit LintPR Title LintLinked Issue CheckFormat Check(로컬 통과)Build(로컬 통과)Test(로컬 145 passed)Lint(로컬 통과)롤백 계획
다음 단계 (Phase 14 후보)
Summary by CodeRabbit
New Features
Tests