EXP-28 Phase 0: Feasibility — Profile Gemma 4 31B + Adapt Sheared LLaMA

[CalebisGross]

Parent: #386 (Project Bespoke epic)

Goal

Validate that structured pruning of Gemma 4 31B is feasible and prepare all tooling for the MI300X pruning session.

Tasks

Research & Code Review

• Clone/examine Sheared LLaMA codebase
• Map Sheared LLaMA's pruning mask code to Gemma 4 31B architecture
• Identify required adaptations: 5:1 sliding/full attention pattern, KV head sharing, dual RoPE
• Verify Gemma 4 31B layer structure: which layers are sliding window, which are full attention

Architecture Analysis

• Load Gemma 4 31B (quantized Q4/Q8 with CPU offload, or on MI300X) and profile:
  • Per-layer activation magnitude on mnemonic encoding tasks
  • Attention entropy per head
  • FFN neuron activation sparsity
  • Gate/residual stream contribution per layer
• Identify candidate layers for removal (lowest importance)
• Design target architectures: map 31B (60 layers) → 8B, 4B, 2B, 1.5B shapes

MI300X Preparation

• Write MI300X setup script (install dependencies: Sheared LLaMA, transformers, Gemma 4)
• Prepare data transfer: v7 encoding dataset → droplet
• Estimate training time and cost per phase
• VRAM budget: verify 31B fits for full fine-tune with gradient checkpointing + Muon optimizer

Sliding/Full Attention Constraint (Carmack flag)

• Document which of the 60 layers are sliding vs full attention
• Define pruning constraint: pruned model must preserve at least 1 full-attention layer per N sliding layers for global context
• Test: what happens if you naively remove all full-attention layers? (quality collapse expected)

VRAM Budget Concern

31B bf16 = 62GB. With Adam: params (62GB) + grads (62GB) + optimizer (124GB) = 248GB > 192GB MI300X.

Options:

1. Muon optimizer (no momentum for matrices, ~40% less state) → ~180GB, fits tight
2. Gradient checkpointing + activation offload → reduces activation memory
3. Mixed precision: bf16 params + fp32 optimizer partitioned → standard DeepSpeed approach
4. 8-bit Adam (bitsandbytes) → 62 + 62 + 62 = 186GB, fits

Resolve before Phase 1.

Definition of Done

• Sheared LLaMA adaptation code written and tested on a small model (Gemma 4 E2B as proxy)
• Gemma 4 31B layer importance profile completed
• Target architectures defined for each pruning level
• MI300X scripts ready to run
• VRAM budget resolved

[CalebisGross]

Research Findings — Honest Assessment

Gemma 4 31B Architecture (Confirmed)

60 layers with 5:1 sliding/full attention pattern:

• Full attention layers (10): 5, 11, 17, 23, 29, 35, 41, 47, 53, 59
• Sliding attention layers (50): all others

Critical structural differences between layer types:

Property	Sliding (50 layers)	Full (10 layers)
Q heads	32	32
KV heads	16 (GQA 2:1)	4 (GQA 8:1)
Head dim	256	512
v_proj	Yes	No (K=V)
RoPE	Standard, theta=10K	Proportional, theta=1M, 25% of dims
Window	1024 tokens	Full context

FFN is uniform: gated MLP, 21504 intermediate, gelu_tanh. No PLE, no MoE, no KV sharing on 31B. Tied embeddings (vocab=262144). HF class: Gemma4ForCausalLM.

Sheared LLaMA Codebase — Adaptation Difficulty: HIGH

The code (github.com/princeton-nlp/LLM-Shearing) is deeply tied to:

1. MosaicML Composer — not standard PyTorch or HuggingFace. Would need to learn/maintain a separate training framework.
2. LLaMA-specific model code — all attention/FFN pruning assumes uniform layers, no GQA, no sliding/full distinction.
3. Uniform Q/K/V dimensions — Gemma 4's full layers have different head dim (512 vs 256), no v_proj, and different GQA ratio. The mask code would need fundamental changes.

Estimated adaptation effort: 2-4 weeks of engineering to write composer_gemma4.py, extend L0Module for heterogeneous layers, handle GQA pruning, and test.

VRAM Budget — PROBLEM

31B bf16 + Adam full fine-tune:

• Params: 62GB
• Gradients: 62GB
• Adam states (momentum + variance, fp32): 124GB
• Total: ~248GB > 192GB MI300X

Must resolve before Phase 1. Options:

1. Muon optimizer (no momentum for matrices): ~180GB — tight fit
2. 8-bit Adam (bitsandbytes): ~186GB — tight fit
3. LoRA fine-tune instead of full (only adapter grads): ~80GB — fits easily
4. DeepSpeed ZeRO-3 with CPU offload: fits anything but 10x slower
5. Gradient checkpointing reduces activations but doesn't help with optimizer state

Three Approaches — Honest Comparison

Approach A: Full Sheared LLaMA Port

• Pros: State-of-the-art joint optimization, best quality at target size
• Cons: 2-4 weeks engineering, Composer dependency, novel GQA/heterogeneous-layer adaptation
• Risk: High — we'd be doing novel research on pruning heterogeneous architectures

Approach B: Layer-Drop + Width-Reduce (Simple)

• Profile layer importance → surgically remove least-important layers → reduce hidden dim via SliceGPT (PCA, no training) → fine-tune result with LoRA/spokes
• Pros: Implementable in days, standard HuggingFace tooling, no Composer
• Cons: Sub-optimal cuts (no joint optimization), may need more continued pretraining to recover
• Risk: Medium — simpler engineering, but quality might be lower than Approach A

Approach C: Knowledge Distillation

• Train a fresh 2B model from scratch using 31B as teacher
• Pros: Total architectural freedom, no pruning complexity, proven technique
• Cons: Need massive data (50-100B tokens), expensive, slow
• Risk: Medium-high — the cost of training from scratch is significant

Recommendation

Start with Approach B (layer-drop + width-reduce) as Phase 0 validation. It answers the key question — "does a 2B-class model derived from 31B beat native 2B on encoding?" — in days, not weeks. If the answer is yes, invest in Approach A for optimal quality. If no, we saved weeks of wasted engineering.

Specific Phase 0 plan:

1. Load Gemma 4 31B (Q4 with CPU offload, or on MI300X)
2. Run encoding tasks, collect per-layer activation statistics
3. Identify the 20 most important layers (preserving ≥2 full-attention layers)
4. Surgically extract those 20 layers into a new model config
5. Apply SliceGPT width reduction (5376 → 2048 hidden) if needed
6. Fine-tune with LoRA on encoding data
7. Evaluate: does it beat Qwen 2B + spokes?

This tells us if the vision is viable before committing to heavy engineering.

[CalebisGross]

E2B Profiling Results — Methodology Validated

Ran profile_layer_importance.py on Gemma 4 E2B (35 layers) with 30 encoding inputs. Script committed at f6c1a62.

Per-Layer Importance (E2B, 35 layers)

Top 5 most important:

Layer	Type	Importance
34 (last)	full_attention	1.785
0 (first)	sliding	1.564
4	full_attention	1.480
23	full_attention	1.015
3	sliding	0.774

Bottom 5 (pruning candidates):

Layer	Type	Importance
30	sliding	0.081
31	sliding	0.099
32	sliding	0.138
29	full_attention	0.182
33	sliding	0.233

Key Observations

1. Classic U-curve: First and last layers are most important. Middle-late layers contribute least.
2. Full attention layers matter more on average (0.68 vs 0.57 sliding) but the least-important full layer (29, importance 0.18) is less important than many sliding layers. Layer type alone doesn't determine importance.
3. Late sliding layers are nearly dead: Layers 30-33 have importance <0.24 — barely modify the residual stream.
4. The profile suggests aggressive pruning is viable: 10+ layers could be removed with minimal quality impact.

What's Needed for 31B

31B (62GB bf16) exceeds local hardware (48GB GPU+RAM). Options:

1. MI300X droplet — load bf16, profile all 60 layers (~20 min)
2. GGUF Q4 via llama-cpp-python — need custom hook code, less clean

MI300X is the right path. Script is ready:

python training/scripts/profile_layer_importance.py \
    --model google/gemma-4-31B-it \
    --inputs training/data/v7_inputs/all_inputs_clean.jsonl \
    --num-examples 50 \
    --max-tokens 512 \
    --device cuda \
    --output gemma4_31b_layer_profile.json

Prediction for 31B

Based on E2B pattern (same architecture family), expect:

• Layers 0-5 and 55-59: highest importance (U-curve endpoints)
• Layers 35-54: lowest importance (late-middle region)
• Full attention layers (5,11,17,23,29,35,41,47,53,59) with mixed importance — some critical, some prunable
• Target 20-layer extraction should preserve layers from both ends + key full-attention layers

[CalebisGross]

Tooling Complete — Ready for MI300X

Scripts committed (branch feat/exp25-faithfulness-probe)

1. profile_layer_importance.py — Per-layer importance profiling via forward-pass hooks. Metrics: residual contribution, cosine drift, composite importance. Validated on E2B (35 layers).

2. extract_pruned_model.py — Surgical layer extraction with config updates. Handles Gemma 4 sliding/full attention pattern, disables KV sharing post-extraction. Dry-run mode for planning.

3. mi300x_bespoke.sh — One-shot MI300X script:

   • Installs deps, downloads Gemma 4 31B (~62GB)
   • Profiles all 60 layers on 50 encoding inputs
   • Extracts at 3 target sizes: 40L (~8B), 20L (~4B), 12L (~2B)
   • Runs encoding smoke test on each
   • ~2-3 hours total including download

Known Issue: E2B != 31B for extraction testing

E2B extraction fails due to PLE + KV sharing + double-wide MLP — features the 31B doesn't have. The 31B is structurally simpler (no PLE, no KV sharing, uniform FFN). Extraction code is correct for 31B but untested until we run on MI300X.

Checklist update

• Profile script written and tested on E2B
• Extraction script written
• MI300X setup script ready
• E2B profiling shows clear importance signal (U-curve, full attn layers matter more)
• BLOCKED: Run on Gemma 4 31B (needs MI300X)
• Analyze 31B profile, design target architecture
• Extract and evaluate pruned models