diff --git a/baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md b/baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md new file mode 100644 index 00000000..942ee0d7 --- /dev/null +++ b/baseline/nanogpt_one_head/ATTENTION_RESIDUALS.md @@ -0,0 +1,176 @@ +# Attention Residuals baseline + +This experiment adds **Full Attention Residuals (AttnRes)** as a controlled architectural baseline inside the existing one-head nanoGPT suite. + +## What changes + +Standard residual routing accumulates sublayer outputs through a fixed additive path. Full AttnRes replaces that accumulation: before each attention or MLP sublayer, it forms the sublayer input by content-dependent attention over the embedding and all preceding sublayer outputs. Each routing point has a learned pseudo-query vector; prior outputs are RMS-normalized to form routing keys, while the original outputs remain the values. The selected mixture is passed through the normal pre-norm sublayer, and only that sublayer output is appended as the next depth value. + +The implementation provides one router before attention and one router before the MLP. The existing causal self-attention computation itself is unchanged. + +## What does not change + +Every matched comparison keeps the existing one-head nanoGPT protocol fixed except for residual routing: + +- FineWeb-Edu sample and pinned revision +- GPT-2 tokenizer +- 80M training tokens, 1M validation tokens, 1M test tokens +- context length 256 +- one transformer block +- one attention head +- embedding width 128 +- zero dropout, bias disabled, tied token/output embeddings +- batch size 4 with 8 gradient-accumulation steps +- fixed evaluation probes and BLEU protocol +- WeightWatcher ERG/randomization settings +- separated `W_Q`, `W_K`, `W_V`, `W_O`, `W_MLP_IN`, and `W_MLP_OUT` matrices + +AttnRes pseudo-query vectors are 1-D parameter tensors of length `n_embd`. They are deliberately excluded from `transformer_matrix_items()`, so WeightWatcher continues to analyze the same six matrices per block. Under Muon, the six 2-D hidden matrices remain in the Muon group while the AttnRes queries enter the auxiliary AdamW group. + +## Command-line setup + +From `baseline/nanogpt_one_head`: + +```bash +python -m pip install -e . +export PYTORCH_ENABLE_MPS_FALLBACK=1 +``` + +Prepare the pinned corpus once and reuse it for every standard/AttnRes comparison: + +```bash +rg-onehead-prepare --config configs/reference.yaml +``` + +All scientific runs use the same command-line trainer: + +```text +rg-onehead-train +``` + +This preserves the existing checkpointing, restart, evaluation, WeightWatcher, MPS/CUDA/CPU/TPU device selection, and result layout. + +## Short run: exact one-epoch reference comparison + +This is the quickest clean comparison and uses all three canonical seeds by default. + +Standard residual control: + +```bash +rg-onehead-train \ + --config configs/reference.yaml \ + --optimizer muon \ + --data-root /tmp/rg-nanogpt-one-head/data \ + --results-root /tmp/rg-nanogpt-short-standard/results \ + --device auto \ + --no-resume +``` + +AttnRes, with the same data, model dimensions, seeds, token budget, Muon hyperparameters, schedule, and probes: + +```bash +rg-onehead-train \ + --config configs/attnres_reference.yaml \ + --optimizer muon \ + --data-root /tmp/rg-nanogpt-one-head/data \ + --results-root /tmp/rg-nanogpt-short-attnres/results \ + --device auto \ + --no-resume +``` + +For a fast pilot before committing to all three seeds, append `--seeds 1337` to both commands. + +## Long run A: exact matched historical 10-epoch comparison + +This preserves the existing validated one-epoch Muon warmup/cosine horizon and then holds the LR floors through epoch 10. It is the correct historical comparison because it changes only residual routing. + +Standard residual control: + +```bash +rg-onehead-train \ + --config configs/muon_10epochs.yaml \ + --optimizer muon \ + --seeds 1337 \ + --data-root /tmp/rg-nanogpt-one-head/data \ + --results-root /tmp/rg-nanogpt-long-muon-standard/results \ + --device auto \ + --no-resume +``` + +Matched AttnRes: + +```bash +rg-onehead-train \ + --config configs/attnres_muon_10epochs.yaml \ + --optimizer muon \ + --seeds 1337 \ + --data-root /tmp/rg-nanogpt-one-head/data \ + --results-root /tmp/rg-nanogpt-long-muon-attnres/results \ + --device auto \ + --no-resume +``` + +## Long run B: convergence-oriented full-horizon cosine pair + +The historical 10-epoch protocol spends epochs 1--10 at the LR floor. For a long run intended to continue optimizing rather than primarily observe late-horizon behavior, the `*_longcosine.yaml` pair uses a conservative learning-rate schedule across all ten epochs: + +```text +training horizon: 10 epochs +LR schedule horizon: 10 epochs +Muon matrix peak LR: 0.0100 +Muon matrix floor LR: 0.0005 +auxiliary peak LR: 0.0003 +auxiliary floor LR: 0.00001 +warmup: 1% of the 10-epoch schedule +``` + +The lower Muon peak relative to the one-epoch reference reduces long-horizon instability risk, while the nonzero floor prevents the optimization from becoming effectively frozen. This is a convergence-oriented candidate rather than a claim that its hyperparameters are globally optimal; actual convergence must be established from validation trajectories. + +Standard residual control: + +```bash +rg-onehead-train \ + --config configs/muon_10epochs_longcosine.yaml \ + --optimizer muon \ + --seeds 1337 \ + --data-root /tmp/rg-nanogpt-one-head/data \ + --results-root /tmp/rg-nanogpt-longcosine-standard/results \ + --device auto \ + --no-resume +``` + +Matched AttnRes: + +```bash +rg-onehead-train \ + --config configs/attnres_muon_10epochs_longcosine.yaml \ + --optimizer muon \ + --seeds 1337 \ + --data-root /tmp/rg-nanogpt-one-head/data \ + --results-root /tmp/rg-nanogpt-longcosine-attnres/results \ + --device auto \ + --no-resume +``` + +Do not attribute a gain to AttnRes unless it appears against the matching standard-residual config at the same token/step budget. + +## Why keep one block first? + +Full AttnRes becomes more expressive as depth increases. In a one-block transformer, the attention router initially has only the embedding output available, and the MLP router can select between the embedding and the attention-sublayer output. This is intentionally conservative: it preserves the architecture, data exposure, token budget, and all six transformer matrices of the current baseline and therefore isolates the effect of residual routing as cleanly as possible. The only added trainable parameters are two length-128 pseudo-query vectors. + +The model implementation itself is depth-ready and keeps one head at every depth. A later multi-block study can test the larger routing advantage without changing the AttnRes implementation. + +## Primary comparison metrics + +For convergence speed, compare at matched optimizer steps and matched tokens: + +- train, validation, and test loss +- validation and test next-token accuracy +- perplexity +- BLEU probe +- step/token count to fixed validation-loss thresholds +- best validation loss and the step at which it occurs + +For the RG/WeightWatcher analysis, retain the existing per-matrix metrics for all six matrices and compare trajectories of alpha, randomization distance, ERG quantities, and spectral diagnostics. AttnRes routing weights can additionally be inspected with `model.attention_residual_weights()` without altering the WeightWatcher matrix set. + +The key causal question is not whether the AttnRes run ends with a better number after a fixed horizon, but whether it reaches the same validation-loss or accuracy threshold in fewer matched tokens while preserving or improving out-of-sample behavior. diff --git a/baseline/nanogpt_one_head/configs/attnres_muon_10epochs.yaml b/baseline/nanogpt_one_head/configs/attnres_muon_10epochs.yaml new file mode 100644 index 00000000..6d95a371 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/attnres_muon_10epochs.yaml @@ -0,0 +1,106 @@ +protocol: + name: rg_nanogpt_one_head_attnres_muon_10epochs + version: 5 + description: Full Attention Residuals with ordinary Muon for ten corpus-equivalent epochs, exactly matched to the existing long-Muon baseline except for residual routing. + +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 80000000 + val_tokens: 1000000 + test_tokens: 1000000 + +model: + vocab_size: 50257 + block_size: 256 + n_layer: 1 + n_head: 1 + n_embd: 128 + dropout: 0.0 + bias: false + tie_weights: true + residual_mode: full_attnres + +training: + seeds: [1337] + batch_size: 4 + grad_accum_steps: 8 + target_epochs: 10.0 + epoch_interval: 0.25 + eval_interval_steps: 250 + eval_batches: 64 + checkpoint_interval_steps: 250 + grad_clip: 1.0 + +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.005 + warmup_fraction: 0.10 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + momentum: 0.90 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + + adamw: + display_name: AdamW + family: adamw + learning_rate: 0.0006 + min_learning_rate: 0.00006 + warmup_fraction: 0.01 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + weight_decay: 0.10 + + muon: + display_name: Muon + auxiliary AdamW + family: muon + matrix_learning_rate: 0.02 + matrix_min_learning_rate: 0.002 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 0.00003 + warmup_fraction: 0.05 + lr_schedule_epochs: 1.0 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + matrix_weight_decay: 0.01 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + aux_weight_decay: 0.01 + +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 + +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + +runtime: + matmul_precision: high + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true diff --git a/baseline/nanogpt_one_head/configs/attnres_muon_10epochs_longcosine.yaml b/baseline/nanogpt_one_head/configs/attnres_muon_10epochs_longcosine.yaml new file mode 100644 index 00000000..f96bad39 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/attnres_muon_10epochs_longcosine.yaml @@ -0,0 +1,106 @@ +protocol: + name: rg_nanogpt_one_head_attnres_muon_10epochs_longcosine + version: 5 + description: Full Attention Residuals one-head nanoGPT using the same conservative ten-epoch Muon cosine schedule as its standard-residual control. + +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 80000000 + val_tokens: 1000000 + test_tokens: 1000000 + +model: + vocab_size: 50257 + block_size: 256 + n_layer: 1 + n_head: 1 + n_embd: 128 + dropout: 0.0 + bias: false + tie_weights: true + residual_mode: full_attnres + +training: + seeds: [1337] + batch_size: 4 + grad_accum_steps: 8 + target_epochs: 10.0 + epoch_interval: 0.25 + eval_interval_steps: 250 + eval_batches: 64 + checkpoint_interval_steps: 250 + grad_clip: 1.0 + +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.005 + warmup_fraction: 0.02 + lr_schedule_epochs: 10.0 + schedule: warmup_cosine + momentum: 0.90 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + + adamw: + display_name: AdamW + family: adamw + learning_rate: 0.0004 + min_learning_rate: 0.00002 + warmup_fraction: 0.01 + lr_schedule_epochs: 10.0 + schedule: warmup_cosine + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + weight_decay: 0.10 + + muon: + display_name: Muon + auxiliary AdamW, conservative long cosine + family: muon + matrix_learning_rate: 0.01 + matrix_min_learning_rate: 0.0005 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 0.00001 + warmup_fraction: 0.01 + lr_schedule_epochs: 10.0 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + matrix_weight_decay: 0.01 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + aux_weight_decay: 0.01 + +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 + +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + +runtime: + matmul_precision: high + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true diff --git a/baseline/nanogpt_one_head/configs/attnres_reference.yaml b/baseline/nanogpt_one_head/configs/attnres_reference.yaml new file mode 100644 index 00000000..6ba45cf3 --- /dev/null +++ b/baseline/nanogpt_one_head/configs/attnres_reference.yaml @@ -0,0 +1,103 @@ +protocol: + name: rg_nanogpt_one_head_attnres_reference + version: 5 + description: Full Attention Residuals one-block, one-head nanoGPT matched exactly to the one-epoch FineWeb-Edu reference protocol except for residual routing. + +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 80000000 + val_tokens: 1000000 + test_tokens: 1000000 + +model: + vocab_size: 50257 + block_size: 256 + n_layer: 1 + n_head: 1 + n_embd: 128 + dropout: 0.0 + bias: false + tie_weights: true + residual_mode: full_attnres + +training: + seeds: [1337, 2027, 4099] + batch_size: 4 + grad_accum_steps: 8 + target_epochs: 1.0 + epoch_interval: 0.125 + eval_interval_steps: 250 + eval_batches: 64 + checkpoint_interval_steps: 250 + grad_clip: 1.0 + +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.005 + warmup_fraction: 0.10 + schedule: warmup_cosine + momentum: 0.90 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + + adamw: + display_name: AdamW + family: adamw + learning_rate: 0.0006 + min_learning_rate: 0.00006 + warmup_fraction: 0.01 + schedule: warmup_cosine + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + weight_decay: 0.10 + + muon: + display_name: Muon + auxiliary AdamW + family: muon + matrix_learning_rate: 0.02 + matrix_min_learning_rate: 0.002 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 0.00003 + warmup_fraction: 0.05 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + matrix_weight_decay: 0.01 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + aux_weight_decay: 0.01 + +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 + +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + +runtime: + matmul_precision: high + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true diff --git a/baseline/nanogpt_one_head/configs/muon_10epochs_longcosine.yaml b/baseline/nanogpt_one_head/configs/muon_10epochs_longcosine.yaml new file mode 100644 index 00000000..6abed7fc --- /dev/null +++ b/baseline/nanogpt_one_head/configs/muon_10epochs_longcosine.yaml @@ -0,0 +1,106 @@ +protocol: + name: rg_nanogpt_one_head_muon_10epochs_longcosine + version: 5 + description: Standard residual one-head nanoGPT control using a conservative ten-epoch Muon cosine schedule for long-horizon convergence studies. + +dataset: + name: HuggingFaceFW/fineweb-edu + config: sample-10BT + split: train + revision: 593b3a867298afb8ce42625a270ef20ddcad28f9 + tokenizer: gpt2 + train_tokens: 80000000 + val_tokens: 1000000 + test_tokens: 1000000 + +model: + vocab_size: 50257 + block_size: 256 + n_layer: 1 + n_head: 1 + n_embd: 128 + dropout: 0.0 + bias: false + tie_weights: true + residual_mode: standard + +training: + seeds: [1337] + batch_size: 4 + grad_accum_steps: 8 + target_epochs: 10.0 + epoch_interval: 0.25 + eval_interval_steps: 250 + eval_batches: 64 + checkpoint_interval_steps: 250 + grad_clip: 1.0 + +optimizer_profiles: + sgd_momentum: + display_name: SGD + Nesterov momentum + family: sgd + learning_rate: 0.05 + min_learning_rate: 0.005 + warmup_fraction: 0.02 + lr_schedule_epochs: 10.0 + schedule: warmup_cosine + momentum: 0.90 + dampening: 0.0 + nesterov: true + weight_decay: 0.01 + + adamw: + display_name: AdamW + family: adamw + learning_rate: 0.0004 + min_learning_rate: 0.00002 + warmup_fraction: 0.01 + lr_schedule_epochs: 10.0 + schedule: warmup_cosine + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + weight_decay: 0.10 + + muon: + display_name: Muon + auxiliary AdamW, conservative long cosine + family: muon + matrix_learning_rate: 0.01 + matrix_min_learning_rate: 0.0005 + aux_learning_rate: 0.0003 + aux_min_learning_rate: 0.00001 + warmup_fraction: 0.01 + lr_schedule_epochs: 10.0 + schedule: warmup_cosine + momentum: 0.95 + nesterov: true + newton_schulz_steps: 5 + muon_epsilon: 1.0e-7 + matrix_weight_decay: 0.01 + beta1: 0.90 + beta2: 0.95 + epsilon: 1.0e-8 + aux_weight_decay: 0.01 + +evaluation: + train_probe_seed: 21001 + validation_probe_seed: 22001 + test_probe_seed: 23001 + bleu_probe_seed: 24001 + bleu_examples: 64 + bleu_prompt_tokens: 64 + bleu_continuation_tokens: 32 + bleu_batch_size: 4 + +weightwatcher: + enabled: true + ERG: true + randomize: true + strict: true + min_evals: 20 + +runtime: + matmul_precision: high + mps_fallback: true + deterministic_algorithms: false + empty_mps_cache_after_weightwatcher: true diff --git a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py index af1c3447..824eb895 100644 --- a/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py +++ b/baseline/nanogpt_one_head/src/rg_nanogpt_one_head/model.py @@ -18,12 +18,13 @@ class GPTConfig: dropout: float = 0.0 bias: bool = False tie_weights: bool = True + residual_mode: str = "standard" def __post_init__(self) -> None: - if self.n_layer != 1 or self.n_head != 1: + if self.n_layer < 1 or self.n_head != 1: raise ValueError( - "the reference architecture is fixed to one block and one " - "attention head" + "the reference architecture requires at least one block and " + "exactly one attention head" ) if self.n_embd % self.n_head != 0: raise ValueError("n_embd must be divisible by n_head") @@ -31,6 +32,10 @@ def __post_init__(self) -> None: raise ValueError("invalid GPT configuration") if not 0.0 <= self.dropout < 1.0: raise ValueError("dropout must be in [0, 1)") + if self.residual_mode not in {"standard", "full_attnres"}: + raise ValueError( + "residual_mode must be 'standard' or 'full_attnres'" + ) class LayerNorm(nn.Module): @@ -49,6 +54,44 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ) +class RMSNorm(nn.Module): + """Parameter-free RMS normalization used only for AttnRes routing keys.""" + + def __init__(self, eps: float = 1e-6) -> None: + super().__init__() + self.eps = float(eps) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + scale = torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps) + return x * scale.to(dtype=x.dtype) + + +class DepthAttentionRouter(nn.Module): + """Full Attention Residuals router over preceding sublayer outputs. + + Each routing point owns one learned pseudo-query vector. Previous outputs + are RMS-normalized to form keys while the unnormalized outputs remain the + values. Softmax is taken over depth, independently for every token. + """ + + def __init__(self, width: int) -> None: + super().__init__() + self.query = nn.Parameter(torch.zeros(width)) + self.norm = RMSNorm() + self.last_mean_weights: torch.Tensor | None = None + + def forward(self, states: list[torch.Tensor]) -> torch.Tensor: + if not states: + raise ValueError("AttnRes requires at least one preceding output") + values = torch.stack(states, dim=0) # [depth, batch, time, width] + keys = self.norm(values) + logits = torch.einsum("d,nbtd->nbt", self.query, keys) + weights = logits.softmax(dim=0) + if not torch.jit.is_scripting(): + self.last_mean_weights = weights.detach().mean(dim=(1, 2)).cpu() + return torch.einsum("nbt,nbtd->btd", weights, values) + + class CausalSelfAttention(nn.Module): def __init__(self, cfg: GPTConfig) -> None: super().__init__() @@ -110,9 +153,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ).transpose(1, 2) dropout_p = self.dropout if self.training else 0.0 if q.device.type == "xla": - # Use core matmul/mask/softmax operations on TPU. This avoids - # depending on accelerator-specific SDPA kernel registration while - # preserving the same causal attention equation. y = self._xla_math_attention( q, k, @@ -150,15 +190,48 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: class Block(nn.Module): def __init__(self, cfg: GPTConfig) -> None: super().__init__() + self.residual_mode = cfg.residual_mode self.ln1 = LayerNorm(cfg.n_embd, cfg.bias) self.attn = CausalSelfAttention(cfg) self.ln2 = LayerNorm(cfg.n_embd, cfg.bias) self.mlp = MLP(cfg) + self.attn_res_router = ( + DepthAttentionRouter(cfg.n_embd) + if cfg.residual_mode == "full_attnres" + else None + ) + self.mlp_res_router = ( + DepthAttentionRouter(cfg.n_embd) + if cfg.residual_mode == "full_attnres" + else None + ) def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.residual_mode != "standard": + states = self.forward_attnres([x]) + return states[-1] x = x + self.attn(self.ln1(x)) return x + self.mlp(self.ln2(x)) + def forward_attnres( + self, + states: list[torch.Tensor], + ) -> list[torch.Tensor]: + if self.attn_res_router is None or self.mlp_res_router is None: + raise RuntimeError("forward_attnres called for a standard block") + + # Full AttnRes replaces additive residual accumulation. Each sublayer + # attends over the embedding and all preceding sublayer outputs to form + # its input h_l, and only f_l(Norm(h_l)) is appended as the next value. + attn_input = self.attn_res_router(states) + attn_output = self.attn(self.ln1(attn_input)) + states.append(attn_output) + + mlp_input = self.mlp_res_router(states) + mlp_output = self.mlp(self.ln2(mlp_input)) + states.append(mlp_output) + return states + class GPT(nn.Module): def __init__(self, cfg: GPTConfig) -> None: @@ -202,8 +275,14 @@ def hidden_states(self, idx: torch.Tensor) -> torch.Tensor: x = self.drop( self.token_embedding(idx) + self.position_embedding(positions) ) - for block in self.blocks: - x = block(x) + if self.cfg.residual_mode == "full_attnres": + layer_outputs = [x] + for block in self.blocks: + layer_outputs = block.forward_attnres(layer_outputs) + x = layer_outputs[-1] + else: + for block in self.blocks: + x = block(x) return self.ln_f(x) def forward( @@ -221,7 +300,6 @@ def forward( return logits, loss def next_token_logits(self, idx: torch.Tensor) -> torch.Tensor: - # Apply the expensive vocabulary projection only to the final position. hidden = self.hidden_states(idx)[:, -1:, :] return self.lm_head(hidden) @@ -249,11 +327,31 @@ def generate_greedy( def parameter_count(self) -> int: return sum(parameter.numel() for parameter in self.parameters()) + def attention_residual_weights(self) -> dict[str, list[float]]: + """Return the latest mean depth-routing weights for diagnostics.""" + result: dict[str, list[float]] = {} + for block_index, block in enumerate(self.blocks): + for name, router in ( + ("ATTN", block.attn_res_router), + ("MLP", block.mlp_res_router), + ): + if router is None or router.last_mean_weights is None: + continue + result[f"L{block_index:02d}_{name}"] = [ + float(value) for value in router.last_mean_weights.tolist() + ] + return result + def transformer_matrix_items( model: GPT, ) -> list[tuple[str, str, int, torch.Tensor]]: - """Return the six transformer matrices used by WeightWatcher and Muon.""" + """Return Q/K/V/O and MLP matrices used by WeightWatcher and Muon. + + AttnRes pseudo-query vectors are intentionally excluded: they are 1-D + routing parameters, not transformer weight matrices, and therefore remain + in the auxiliary optimizer group rather than being folded into Muon/WW. + """ items: list[tuple[str, str, int, torch.Tensor]] = [] for block_index, block in enumerate(model.blocks): matrices = ( diff --git a/baseline/nanogpt_one_head/tests/test_attention_residuals.py b/baseline/nanogpt_one_head/tests/test_attention_residuals.py new file mode 100644 index 00000000..33c4e3fb --- /dev/null +++ b/baseline/nanogpt_one_head/tests/test_attention_residuals.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +import math + +import torch + +from rg_nanogpt_one_head.model import GPT, GPTConfig, transformer_matrix_items +from rg_nanogpt_one_head.optimizers import make_optimizer_handles + + +def _tiny_config(*, residual_mode: str, n_layer: int = 1) -> GPTConfig: + return GPTConfig( + vocab_size=97, + block_size=16, + n_layer=n_layer, + n_head=1, + n_embd=32, + dropout=0.0, + bias=False, + tie_weights=False, + residual_mode=residual_mode, + ) + + +def _muon_profile() -> dict: + return { + "family": "muon", + "matrix_learning_rate": 0.01, + "matrix_min_learning_rate": 0.0005, + "aux_learning_rate": 0.0003, + "aux_min_learning_rate": 0.00001, + "momentum": 0.95, + "nesterov": True, + "newton_schulz_steps": 3, + "muon_epsilon": 1.0e-7, + "matrix_weight_decay": 0.01, + "beta1": 0.90, + "beta2": 0.95, + "epsilon": 1.0e-8, + "aux_weight_decay": 0.01, + } + + +def test_full_attnres_forward_backward_is_finite() -> None: + torch.manual_seed(7) + model = GPT(_tiny_config(residual_mode="full_attnres")) + idx = torch.randint(0, model.cfg.vocab_size, (3, 12)) + targets = torch.randint(0, model.cfg.vocab_size, (3, 12)) + + logits, loss = model(idx, targets) + + assert logits.shape == (3, 12, model.cfg.vocab_size) + assert loss is not None and torch.isfinite(loss) + loss.backward() + + block = model.blocks[0] + assert block.attn_res_router is not None + assert block.mlp_res_router is not None + assert block.attn_res_router.query.grad is not None + assert block.mlp_res_router.query.grad is not None + assert torch.isfinite(block.attn_res_router.query.grad).all() + assert torch.isfinite(block.mlp_res_router.query.grad).all() + + +def test_one_block_routing_weights_have_expected_depth() -> None: + torch.manual_seed(11) + model = GPT(_tiny_config(residual_mode="full_attnres")) + idx = torch.randint(0, model.cfg.vocab_size, (2, 10)) + model(idx) + + weights = model.attention_residual_weights() + assert set(weights) == {"L00_ATTN", "L00_MLP"} + assert len(weights["L00_ATTN"]) == 1 + assert len(weights["L00_MLP"]) == 2 + assert math.isclose(weights["L00_ATTN"][0], 1.0, abs_tol=1e-6) + assert math.isclose(sum(weights["L00_MLP"]), 1.0, abs_tol=1e-6) + + +def test_transformer_matrices_remain_separate_from_attnres_queries() -> None: + model = GPT(_tiny_config(residual_mode="full_attnres")) + items = transformer_matrix_items(model) + + assert [item[1] for item in items] == [ + "W_Q", + "W_K", + "W_V", + "W_O", + "W_MLP_IN", + "W_MLP_OUT", + ] + assert len(items) == 6 + matrix_ids = {id(item[3]) for item in items} + assert id(model.blocks[0].attn_res_router.query) not in matrix_ids + assert id(model.blocks[0].mlp_res_router.query) not in matrix_ids + + +def test_muon_keeps_attnres_queries_in_auxiliary_optimizer() -> None: + model = GPT(_tiny_config(residual_mode="full_attnres")) + handles = make_optimizer_handles(model, _muon_profile()) + + assert [handle.role for handle in handles] == ["primary", "auxiliary"] + primary_ids = { + id(parameter) + for group in handles[0].optimizer.param_groups + for parameter in group["params"] + } + auxiliary_ids = { + id(parameter) + for group in handles[1].optimizer.param_groups + for parameter in group["params"] + } + + for _, _, _, matrix in transformer_matrix_items(model): + assert id(matrix) in primary_ids + for block in model.blocks: + assert id(block.attn_res_router.query) not in primary_ids + assert id(block.mlp_res_router.query) not in primary_ids + assert id(block.attn_res_router.query) in auxiliary_ids + assert id(block.mlp_res_router.query) in auxiliary_ids + + +def test_full_attnres_is_depth_ready_while_preserving_one_head() -> None: + torch.manual_seed(13) + model = GPT(_tiny_config(residual_mode="full_attnres", n_layer=3)) + idx = torch.randint(0, model.cfg.vocab_size, (2, 8)) + logits, _ = model(idx) + + assert logits.shape == (2, 8, model.cfg.vocab_size) + assert len(transformer_matrix_items(model)) == 18 + weights = model.attention_residual_weights() + assert len(weights["L00_ATTN"]) == 1 + assert len(weights["L00_MLP"]) == 2 + assert len(weights["L01_ATTN"]) == 3 + assert len(weights["L01_MLP"]) == 4 + assert len(weights["L02_ATTN"]) == 5 + assert len(weights["L02_MLP"]) == 6