MiniMind 설정을 줄여 모델을 만들고, 층의 결정을 숫자로 확인한다

목표

MiniMindConfig 를 128차원·4층·Q 헤드 4·KV 헤드 2·어휘 1024 로 줄여 모델을 만들고, 파라미터 수를 식과 모델에서 맞춰 본다. GQA 의 KV 캐시 절약, RMSNorm 의 출력, RoPE 의 상대 위치 성질, 학습 전 손실, MoE 의 전체·활성 파라미터를 MiniMind 코드로 직접 잰다.

왜 중요한가

설정 파일 한 장이 모델의 크기와 서빙 비용을 정한다. 파라미터를 식으로 셀 수 있으면 공유 가중치나 정규화를 빠뜨렸는지 바로 알고, KV 캐시를 토큰당 바이트로 계산할 수 있으면 한 카드에 몇 명을 받을지 계산할 수 있다. 그리고 요즘 모델이 쓰는 처방(RMSNorm·GQA·RoPE·SwiGLU·임베딩 공유)은 각각 이유가 있다. 이 실습은 설명을 믿는 대신 MiniMind 의 실제 코드로 그 성질을 숫자로 확인한다 — 평균을 빼지 않는 정규화, 절반이 된 K·V, 거리만 남는 회전.

단계

작은 설정을 /root/mm/arch/config.json 에 저장하세요(hidden_size 128, num_hidden_layers 4, vocab_size 1024, num_attention_heads 4, num_key_value_heads 2, max_position_embeddings 512).
파라미터 수를 식으로 계산하고 모델에서도 세어 /root/mm/arch/count.json 에 intermediate_size·per_layer·formula·model 로 적으세요.
첫 층 어텐션의 q_proj·k_proj 출력 크기와, fp32 로 토큰 하나를 캐시할 때의 바이트(GQA 그대로 / KV 헤드가 Q 헤드 수와 같다면)를 /root/mm/arch/gqa.json 에 적으세요.
torch.manual_seed(0) 뒤 randn(4, 128)*5+3 을 첫 층의 input_layernorm 에 넣어, 입력·출력의 RMS 와 출력의 평균을 /root/mm/arch/rmsnorm.json 에 적으세요.
torch.manual_seed(0) 뒤 q·k 를 하나씩 뽑아 MiniMind 의 apply_rotary_pos_emb 로 (3,7)·(103,107)·(3,50) 자리에 두고 내적을 /root/mm/arch/rope.json 에 적으세요.
mmkit.seed_all(0) 으로 만든 학습 전 모델의 손실을 검증 자료 앞 16×128 토큰으로 재어 /root/mm/arch/init_loss.json 에 loss·ln_vocab 로 적으세요.
같은 설정에 use_moe=True 를 더한 모델의 전체·활성 파라미터를 /root/mm/arch/moe.json 에 적으세요.
/root/mm/arch/report.md 에 ## 파라미터는 어디에 ## GQA 와 KV 캐시 ## RoPE 와 RMSNorm 세 절을 쓰고, 2단계의 파라미터 수와 3단계의 GQA 토큰당 캐시 바이트를 넣으세요.

참고

모델은 import mmkit; mmkit.new_model(설정 dict) 로 만듭니다. 안에서 /opt/minimind/model/model_minimind.py 의 MiniMindForCausalLM(MiniMindConfig(**설정)) 을 부를 뿐입니다 — 코드를 직접 열어 보세요.
층의 부품 이름: model.model.layers[i].self_attn.q_proj, .input_layernorm, .mlp.gate_proj, model.model.freqs_cos.
흔한 실수: 공유된 lm_head 를 한 번 더 세는 것, q_norm·k_norm(head_dim 크기)을 빼먹는 것, KV 캐시를 복사(repeat_kv) 뒤의 크기로 세는 것.
원문: model_minimind.py · RMSNorm · RoFormer(RoPE) · GQA · GLU Variants

작은 설정 한 장

/root/mm/arch/config.json 에 hidden_size 128, num_hidden_layers 4, vocab_size 1024, num_attention_heads 4, num_key_value_heads 2, max_position_embeddings 512 를 적으세요. 나머지 값은 MiniMindConfig 의 기본값을 씁니다.

MiniMind-3 은 768·8·6400·8·4 입니다. 층 수와 차원을 줄이면 파라미터가 차원의 제곱으로 줄어 CPU 로도 몇 분 안에 학습됩니다. 채점기는 이 파일로 실제 모델을 만들어 봅니다.

파라미터를 식으로 센다

파라미터 수를 식으로 계산하고(formula), 모델에서도 세어(model) /root/mm/arch/count.json 에 intermediate_size·per_layer·formula·model 로 적으세요. 두 값이 같아야 합니다.

한 층 = q·k·v·o 투영 + q_norm·k_norm(각 head_dim) + gate·up·down(각 hidden×intermediate) + RMSNorm 둘(각 hidden). 여기에 임베딩(어휘×hidden, 출력층과 공유라 한 번)과 마지막 RMSNorm 을 더합니다. intermediate_size 는 ceil(hidden·π/64)·64 입니다.

GQA 가 줄이는 것

첫 층 어텐션의 q_proj.out_features·k_proj.out_features·n_rep 와, fp32 로 토큰 하나를 캐시하는 바이트(K·V 둘 × 층 × KV 헤드 × head_dim × 4) — GQA 그대로(kv_bytes_per_token_gqa)와 KV 헤드가 Q 헤드 수와 같을 때(kv_bytes_per_token_mha) — 를 /root/mm/arch/gqa.json 에 적으세요.

MiniMind 의 Attention 은 K·V 를 num_key_value_heads 개만 만들고, 계산할 때 repeat_kv 로 n_rep 번 복사합니다. 캐시(past_kv)에는 복사 전의 K·V 가 들어갑니다.

RMSNorm 은 평균을 빼지 않는다

torch.manual_seed(0) 뒤 x = torch.randn(4, 128) * 5 + 3 을 만들어 첫 층의 input_layernorm 에 넣고, 입력과 출력의 RMS(행마다 sqrt(mean(x²)) 의 평균)와 출력 전체의 평균을 /root/mm/arch/rmsnorm.json 에 rms_before·rms_after·mean_after 로 적으세요.

MiniMind 의 RMSNorm 은 x * rsqrt(mean(x²) + eps) 에 가중치(처음엔 1)를 곱합니다. RMS 는 1 로 돌아오지만 평균을 빼지 않으므로 입력의 +3 이 출력에 흔적으로 남습니다 — LayerNorm 이라면 평균이 0 이 됩니다.

RoPE 는 거리만 남긴다

torch.manual_seed(0) 뒤 q = torch.randn(1,1,1,head_dim), k = torch.randn(1,1,1,head_dim) 을 뽑고, 모델의 freqs_cos·freqs_sin 과 apply_rotary_pos_emb 로 q 는 첫 위치, k 는 둘째 위치로 돌려 내적한 값을 (3,7)·(103,107)·(3,50) 에 대해 /root/mm/arch/rope.json 에 dot_3_7·dot_103_107·dot_3_50 로 적으세요.

apply_rotary_pos_emb(q, k, cos[p:p+1], sin[p:p+1]) 는 q·k 를 같은 위치로 돌리므로, q 와 k 를 따로 불러 각자의 위치로 돌리세요. 거리가 4 로 같은 두 쌍의 내적은 소수점 다섯째 자리까지 같아야 합니다.

학습 전 손실은 ln(어휘)

mmkit.seed_all(0) 으로 시드를 고정한 뒤 모델을 새로 만들어, /opt/mm/ref/val.npy 앞 16×128 토큰(view(16, 128))의 손실(model(x, labels=x).loss)을 /root/mm/arch/init_loss.json 에 loss·ln_vocab 로 적으세요.

무작위로 초기화된 모델은 모든 토큰에 거의 같은 확률(1/1024)을 줍니다. 그 교차 엔트로피가 ln 1024 ≈ 6.93 입니다. 이 값에서 멀면 초기화나 라벨 처리가 잘못된 것입니다.

MoE 는 크지만 한 토큰은 적게 쓴다

1단계 설정에 use_moe=True 를 더한 모델(기본 전문가 4개·토큰당 1개)의 전체 파라미터와, 토큰 하나가 실제로 쓰는 파라미터(전체 − 전문가 하나의 크기×전문가 수 + 전문가 하나의 크기×토큰당 전문가 수)를 /root/mm/arch/moe.json 에 num_experts·top_k·total_params·active_params 로 적으세요.

전문가 하나의 크기는 이름에 mlp.experts.0. 이 든 파라미터를 더하면 됩니다(층마다 하나씩이므로 모든 층의 0번 전문가 합). MiniMind 의 trainer_utils.get_model_params 가 같은 셈으로 '198M-A64M' 을 찍습니다.

구성을 설명하는 보고

/root/mm/arch/report.md 에 ## 파라미터는 어디에 ## GQA 와 KV 캐시 ## RoPE 와 RMSNorm 세 절을 쓰고, 2단계의 파라미터 수(model)와 3단계의 kv_bytes_per_token_gqa 를 숫자로 넣으세요.

첫 절에는 FFN·어텐션·임베딩이 각각 몇 %인지, 둘째 절에는 GQA 가 캐시를 몇 배 줄였는지, 셋째 절에는 4·5단계에서 본 성질을 한 줄씩 적으면 됩니다.