frankenstallm/source/PLAN_hybrid_3b_fixes.md

# FRANKENSTALLM-H 3B Hybrid Model — 점검 결과 및 수정 실행 가이드

> **작성일**: 2026-03-05
> **목적**: Phase 2 검증 전, 발견된 이슈 6건을 수정하고 바로 실행 가능한 상태로 만든다.
> **다음 세션에서 이 문서를 참조하여 바로 실행할 것.**

---

## 이슈 요약 (6건)

| # | 심각도 | 이슈 | 파일 | 영향 |
|---|--------|------|------|------|
| 1 | **CRITICAL** | Mamba 블록에 FFN(channel mixer) 없음 | `model/mamba_block.py` | 37/40 레이어 capacity 부족 |
| 2 | **HIGH** | `n_groups=1` (Nemotron 표준은 8) | `configs/hybrid_3b.yaml` | B/C projection 표현력 저하 |
| 3 | **HIGH** | Hybrid 아키텍처 startup 로그 없음 | `train/pretrain.py` | 디버깅·모니터링 곤란 |
| 4 | **MEDIUM** | 체크포인트 resume 시 아키텍처 검증 없음 | `train/utils.py` | 잘못된 가중치 로드 가능 |
| 5 | **MEDIUM** | selective_scan에 NaN/Inf 감지 없음 | `model/mamba_block.py` | 수치 불안정 진단 불가 |
| 6 | **LOW** | selective_scan 입력 shape 검증 없음 | `model/mamba_block.py` | 모호한 에러 메시지 |

---

## 구현 순서 및 의존성

```
Step 1 (FFN 추가) ← 가장 먼저, 아키텍처 변경
  ├── 1a. model/config.py: mamba_d_ffn 필드 추가
  ├── 1b. model/mamba_block.py: FFN sublayer 추가
  ├── 1c. model/transformer.py: 생성자 인자 전달 + _init_weights 수정
  └── 1d. configs/hybrid_3b.yaml: mamba_d_ffn=4608 추가

Step 2 (n_groups) ← Step 1과 독립, 병렬 가능
  └── configs/hybrid_3b.yaml: n_groups=8

Step 3 (로그) ← Step 1 완료 후 (파라미터 수 정확해야)
  └── train/pretrain.py: startup 배너에 hybrid 정보 추가

Step 4 (체크포인트 검증) ← 독립
  └── train/utils.py: load_checkpoint에 config 비교 로직

Step 5-6 (NaN 감지 + shape 검증) ← 독립
  └── model/mamba_block.py: selective_scan 함수
```

**병렬 가능**: Step 1 + Step 2는 YAML만 겹침 (마지막에 합치면 됨).
Step 4, Step 5-6도 독립적으로 병렬 실행 가능.

---

## Step 1: Mamba2Block에 FFN 추가 (CRITICAL)

### 배경

- Mamba2Block은 SSM(sequence mixer)만 있고 FFN(channel mixer)이 없음
- Nemotron-H에서는 모든 Mamba 레이어 뒤에 MLP가 따라옴
- 현재 37/40 레이어에 FFN이 없어 feature mixing이 불가능
- **확정**: `mamba_d_ffn = 4608` (d_model × 1.5), 총 파라미터 ~4.5B, VRAM ~80GB/GPU

### 1a. `model/config.py` 수정

**위치**: LMConfig dataclass 내부 (line 61 이후)

**추가할 필드** (기존 `mamba_chunk_size` 뒤에):
```python
    mamba_d_ffn: Optional[int] = None  # FFN dim for Mamba blocks (None → d_ffn)
```

**`__post_init__` 추가** (line 86, hybrid validation 블록 뒤에):
```python
        # Mamba FFN dimension: default to d_ffn if not specified
        if self.mamba_d_ffn is None:
            self.mamba_d_ffn = self.d_ffn
```

**`to_dict()` 추가** (기존 mamba_chunk_size 뒤에):
```python
            "mamba_d_ffn": self.mamba_d_ffn,
```

### 1b. `model/mamba_block.py` 수정

**Import 변경** (line 19):
```python
# 변경 전:
from .layers import RMSNorm

# 변경 후:
from .layers import RMSNorm, SwiGLU
```

**`Mamba2Block.__init__` 시그니처 변경** (line 128-137):
```python
# 변경 전:
    def __init__(
        self,
        d_model: int,
        d_state: int = 128,
        head_dim: int = 64,
        expand: int = 2,
        conv_kernel: int = 4,
        n_groups: int = 1,
        chunk_size: int = 256,
    ) -> None:

# 변경 후:
    def __init__(
        self,
        d_model: int,
        d_state: int = 128,
        head_dim: int = 64,
        expand: int = 2,
        conv_kernel: int = 4,
        n_groups: int = 1,
        chunk_size: int = 256,
        d_ffn: int = 0,
        bias: bool = False,
    ) -> None:
```

**FFN 서브레이어 추가** (line 192, `self.out_proj` 뒤에):
```python
        # --- FFN sublayer (channel mixer) ---
        if d_ffn > 0:
            self.ffn_norm = RMSNorm(d_model)
            self.ffn = SwiGLU(d_model, d_ffn, bias=bias)
        else:
            self.ffn_norm = None
            self.ffn = None
```

**`forward()` 수정** (line 280):
```python
# 변경 전:
        return residual + self.out_proj(y)

# 변경 후:
        x = residual + self.out_proj(y)
        # FFN sublayer (channel mixer)
        if self.ffn is not None:
            x = x + self.ffn(self.ffn_norm(x))
        return x
```

### 1c. `model/transformer.py` 수정

**Mamba2Block 생성자 호출 변경** (line 124-132):
```python
# 변경 전:
                    layers.append(Mamba2Block(
                        d_model=config.d_model,
                        d_state=config.mamba_d_state,
                        head_dim=config.mamba_head_dim,
                        expand=config.mamba_expand,
                        conv_kernel=config.mamba_conv_kernel,
                        n_groups=config.mamba_n_groups,
                        chunk_size=config.mamba_chunk_size,
                    ))

# 변경 후:
                    layers.append(Mamba2Block(
                        d_model=config.d_model,
                        d_state=config.mamba_d_state,
                        head_dim=config.mamba_head_dim,
                        expand=config.mamba_expand,
                        conv_kernel=config.mamba_conv_kernel,
                        n_groups=config.mamba_n_groups,
                        chunk_size=config.mamba_chunk_size,
                        d_ffn=config.mamba_d_ffn,
                        bias=config.bias,
                    ))
```

**`_init_weights` 수정** (line 180-182):
```python
# 변경 전:
        # Mamba2Block handles its own parameter init (A_log, D, dt_bias, etc.)
        if isinstance(module, Mamba2Block):
            return

# 변경 후 (이 3줄을 삭제):
# 삭제 이유: FFN 추가 후 내부 SwiGLU의 nn.Linear가 init 필요.
# A_log, D, dt_bias는 nn.Parameter이므로 isinstance(nn.Linear) 체크에 걸리지 않아
# 자동으로 스킵됨 (Mamba2Block.__init__에서 직접 초기화됨).
```

### 1d. `configs/hybrid_3b.yaml` 수정

```yaml
# mamba_chunk_size: 256 뒤에 추가:
  mamba_d_ffn: 4608
```

### Step 1 검증

```bash
cd /PROJECT/0325120031_A/ghong/taketimes/llm-bang
CUDA_VISIBLE_DEVICES=0 python -c "
import torch, sys
sys.path.insert(0, '.')
from model import LLM, LMConfig

config = LMConfig.from_yaml('configs/hybrid_3b.yaml')
print(f'mamba_d_ffn = {config.mamba_d_ffn}')

model = LLM(config)
total = sum(p.numel() for p in model.parameters())
print(f'Total params: {total:,} ({total/1e9:.2f}B)')

# Forward test
x = torch.randint(0, 64000, (1, 128))
logits, loss = model(x, targets=x)
print(f'Forward OK: logits shape={logits.shape}, loss={loss.item():.4f}')

# Backward test
loss.backward()
grads_ok = all(p.grad is not None for p in model.parameters() if p.requires_grad)
print(f'Backward OK: all grads exist = {grads_ok}')
"
# 예상 출력: Total params ~4.5B, Forward/Backward OK
```

---

## Step 2: n_groups 수정

### `configs/hybrid_3b.yaml`

```yaml
# 변경 전:
  mamba_n_groups: 1

# 변경 후:
  mamba_n_groups: 8
```

### 검증

n_heads(= d_inner / head_dim = 6144 / 64 = 96) % 8 == 0 ✓
Step 1 검증 스크립트에서 함께 확인됨 (assertion이 `__init__`에 있음).

---

## Step 3: 하이브리드 아키텍처 startup 로그 추가

### `train/pretrain.py` 수정

**위치**: line 296-297 (모델 파라미터 출력 부분) 뒤에 추가

```python
    if is_main_process():
        total_params = sum(p.numel() for p in model.parameters())
        print(f"Model parameters: {total_params:,}")
        print(f"LMConfig: {lm_config}")

        # --- 여기부터 추가 ---
        if lm_config.use_hybrid:
            pattern = lm_config.hybrid_pattern.split()
            m_count = sum(1 for p in pattern if p == 'M')
            a_count = sum(1 for p in pattern if p == 'A')
            mamba_params = sum(
                p.numel() for n, p in model.named_parameters()
                if 'layers.' in n and pattern[int(n.split('.')[1])] == 'M'
            )
            attn_params = sum(
                p.numel() for n, p in model.named_parameters()
                if 'layers.' in n and pattern[int(n.split('.')[1])] == 'A'
            )
            other_params = total_params - mamba_params - attn_params
            print(
                f"  arch     : Hybrid Mamba-Transformer\n"
                f"  layers   : {m_count} Mamba + {a_count} Attention = {len(pattern)} total\n"
                f"  params   : Mamba {mamba_params/1e6:.0f}M + "
                f"Attn {attn_params/1e6:.0f}M + Other {other_params/1e6:.0f}M\n"
                f"  mamba cfg: d_state={lm_config.mamba_d_state}, "
                f"head_dim={lm_config.mamba_head_dim}, "
                f"expand={lm_config.mamba_expand}, "
                f"n_groups={lm_config.mamba_n_groups}, "
                f"d_ffn={lm_config.mamba_d_ffn}"
            )
        # --- 추가 끝 ---
```

### 검증

Step 1 검증 실행 시 로그에 hybrid 정보가 출력되는지 확인.

---

## Step 4: 체크포인트 resume 아키텍처 검증

### `train/utils.py` — `load_checkpoint()` 수정

**위치**: line 179 (`raw_model.load_state_dict(...)`) 직전에 추가

```python
    # --- Architecture validation ---
    config_path = ckpt_dir / "config.yaml"
    if config_path.exists() and hasattr(raw_model, "config"):
        with open(config_path, "r", encoding="utf-8") as f:
            saved_cfg = yaml.safe_load(f)
        current_cfg = raw_model.config.to_dict()
        critical_keys = [
            "d_model", "n_layers", "n_heads", "n_kv_heads", "vocab_size",
            "use_hybrid", "hybrid_pattern",
        ]
        mismatches = []
        for key in critical_keys:
            saved_val = saved_cfg.get(key)
            current_val = current_cfg.get(key)
            if saved_val is not None and saved_val != current_val:
                mismatches.append(
                    f"  {key}: checkpoint={saved_val} vs current={current_val}"
                )
        if mismatches:
            raise ValueError(
                f"Checkpoint architecture mismatch!\n"
                f"Checkpoint dir: {ckpt_dir}\n"
                + "\n".join(mismatches)
                + "\nUse --config matching the checkpoint, or start fresh."
            )
    # --- End architecture validation ---
```

**참고**: `yaml`은 이미 `train/utils.py` line 23에서 import 되어 있음.

### 검증

```bash
# 의도적으로 다른 config로 resume 시도
CUDA_VISIBLE_DEVICES=0 python train/pretrain.py \
    --config configs/small.yaml \
    --train_data data/3b_train.bin \
    --resume checkpoints/hybrid_3b_run1/checkpoint-0001000
# 예상: ValueError "Checkpoint architecture mismatch!" 출력
```

---

## Step 5: selective_scan NaN/Inf 감지

### `model/mamba_block.py` — `selective_scan()` 수정

**위치**: line 94 (`y[:, t, :, :] = y_t.to(x.dtype)`) 뒤에 추가

```python
        # Periodic NaN/Inf check (every 512 steps, < 1% overhead)
        if t % 512 == 511:
            if not torch.isfinite(h).all():
                raise RuntimeError(
                    f"NaN/Inf in Mamba SSM state at timestep {t}/{seq_len}. "
                    f"h stats: min={h.min().item():.4e}, max={h.max().item():.4e}, "
                    f"A_log range=[{A_log.min().item():.4f}, {A_log.max().item():.4f}]"
                )
```

### 검증

```bash
CUDA_VISIBLE_DEVICES=0 python -c "
import torch, sys
sys.path.insert(0, '.')
from model.mamba_block import Mamba2Block

block = Mamba2Block(d_model=256, d_state=64, head_dim=32, d_ffn=384)
x = torch.randn(1, 1024, 256)

# 정상 케이스
y = block(x)
print(f'Normal: output shape={y.shape}, finite={torch.isfinite(y).all()}')

# NaN 주입 테스트
block.A_log.data.fill_(100.0)  # 매우 큰 값 → exp(100) overflow
try:
    y = block(x)
    print('WARNING: NaN not detected!')
except RuntimeError as e:
    print(f'NaN correctly detected: {e}')
"
```

---

## Step 6: selective_scan 입력 shape 검증

### `model/mamba_block.py` — `selective_scan()` 수정

**위치**: line 49 (`batch, seq_len, n_heads, head_dim = x.shape`) 직전에 추가

```python
    # Input shape validation
    assert x.ndim == 4, f"x expected 4D (B,L,n_heads,head_dim), got {x.shape}"
    assert dt.ndim == 3, f"dt expected 3D (B,L,n_heads), got {dt.shape}"
    assert B.ndim == 4, f"B expected 4D (B,L,n_groups,d_state), got {B.shape}"
    assert C.ndim == 4, f"C expected 4D (B,L,n_groups,d_state), got {C.shape}"
```

---

## 최종 검증 절차 (모든 Step 완료 후)

### 1. 모델 생성 + Forward/Backward (단일 GPU)

```bash
cd /PROJECT/0325120031_A/ghong/taketimes/llm-bang
CUDA_VISIBLE_DEVICES=0 python -c "
import torch, sys
sys.path.insert(0, '.')
from model import LLM, LMConfig

config = LMConfig.from_yaml('configs/hybrid_3b.yaml')
model = LLM(config).cuda()

total = sum(p.numel() for p in model.parameters())
print(f'Total params: {total:,} ({total/1e9:.2f}B)')
assert 4.0e9 < total < 5.0e9, f'Expected ~4.5B params, got {total/1e9:.2f}B'

# Forward
x = torch.randint(0, 64000, (2, 512)).cuda()
logits, loss = model(x, targets=x)
print(f'Forward: logits={logits.shape}, loss={loss.item():.4f}')

# Backward
loss.backward()
no_grad = [n for n, p in model.named_parameters() if p.requires_grad and p.grad is None]
assert len(no_grad) == 0, f'Missing gradients: {no_grad}'
print(f'Backward: all {sum(1 for p in model.parameters() if p.requires_grad)} params have grad')

# VRAM
print(f'VRAM: {torch.cuda.memory_allocated()/1e9:.1f}GB allocated')
"
```

### 2. DDP 8-GPU 테스트 (10 steps)

```bash
cd /PROJECT/0325120031_A/ghong/taketimes/llm-bang
torchrun --nproc_per_node=8 --master_port=29501 train/pretrain.py \
    --config configs/hybrid_3b.yaml \
    --train_data data/3b_train.bin \
    --batch_size 2 \
    --lr 1e-4 \
    --warmup_steps 5 \
    --grad_accum 1 \
    --max_steps 10 \
    --checkpoint_dir /tmp/hybrid_test_ckpt \
    --use_fp8
# 예상: 10 steps 완료, 체크포인트 저장, startup 배너에 hybrid 정보 출력
```

### 3. 체크포인트 Resume 테스트

```bash
# Step 2 체크포인트에서 resume
torchrun --nproc_per_node=8 --master_port=29501 train/pretrain.py \
    --config configs/hybrid_3b.yaml \
    --train_data data/3b_train.bin \
    --batch_size 2 \
    --lr 1e-4 \
    --warmup_steps 5 \
    --grad_accum 1 \
    --max_steps 20 \
    --checkpoint_dir /tmp/hybrid_test_ckpt \
    --resume /tmp/hybrid_test_ckpt/checkpoint-0000010 \
    --use_fp8
# 예상: step 10에서 이어서 step 20까지 학습
```

---

## 수정하지 않는 것들 (의도적 제외)

- **sequential scan 성능**: Python for-loop는 느리지만 구조 변경이 큼. 별도 태스크로 chunked SSD 구현
- **FP8 + Mamba 혼합**: 현재 설계(Mamba=bf16, Attention=FP8)가 올바름. te.fp8_autocast는 te 모듈만 영향
- **DDP 설정**: find_unused_parameters=False, gradient_as_bucket_view=True 모두 정상
- **pure Transformer 모드**: use_hybrid=False면 기존 동작 유지 (하위 호환)

---

## 수정 대상 파일 요약

| 파일 | Step | 변경 내용 |
|------|------|----------|
| `model/config.py` | 1a | `mamba_d_ffn` 필드 + `__post_init__` + `to_dict()` |
| `model/mamba_block.py` | 1b, 5, 6 | SwiGLU import, FFN sublayer, NaN 감지, shape 검증 |
| `model/transformer.py` | 1c | Mamba2Block 생성자에 d_ffn/bias 전달, `_init_weights` 수정 |
| `configs/hybrid_3b.yaml` | 1d, 2 | `mamba_d_ffn: 4608`, `mamba_n_groups: 8` |
| `train/pretrain.py` | 3 | Hybrid startup 로그 |
| `train/utils.py` | 4 | `load_checkpoint()` 아키텍처 검증 |

---

## 실행 지시 (다음 세션용)

이 문서를 참조하여 다음 명령을 내리면 됩니다:

> "이 문서(hashed-drifting-harp.md)의 Step 1~6을 순서대로 실행해 줘.
> Step 1+2는 병렬로, Step 3~6은 독립적으로 진행.
> 각 Step 완료 후 해당 검증을 실행하고,
> 전체 완료 후 최종 검증 절차 3단계를 모두 실행해 줘."