docs/SOVEREIGN_NVIDIA_TRAINING_GUIDE.md

# Sovereign NVIDIA Training Guide

**How SnapKitty trains NVIDIA's own model on NVIDIA's own hardware to produce verified NVIDIA kernels.**

---

## Why Nemotron + NVIDIA Megatron

Most AI code generators are trained on GitHub scrapes and hope for the best. We took a different approach:

| Decision | Why |
|----------|-----|
| **NVIDIA Nemotron** as base model | Nemotron was built by NVIDIA. Its internal weights already encode CUDA semantics, PTX instruction behavior, tensor core data paths, and memory hierarchy. We don't teach it NVIDIA — it already *is* NVIDIA. |
| **NVIDIA Megatron** as training framework | Megatron-LM is NVIDIA's own distributed training framework. Tensor parallelism, pipeline parallelism, sequence parallelism — all designed for NVIDIA hardware by NVIDIA engineers. |
| **RTX 3080 / RTX 4090** as target hardware | We generate kernels for the same GPUs we train on. The model writes PTX for the machine it runs on. |
| **PAX formal verification** as training signal | Every training example is a proven-correct kernel. The model learns what correct GPU code looks like because it only ever sees correct GPU code. |

The result: a model that writes NVIDIA GPU kernels with mathematical correctness proofs attached, trained by NVIDIA's framework on NVIDIA's hardware using NVIDIA's model.

**No external dependencies. No cloud APIs. Sovereign compute.**

---

## The Stack (All NVIDIA, All the Way Down)

```
┌────────────────────────────────────────────────────────────────────┐
│                    SNAPKITTY SOVEREIGN COMPUTE                      │
├────────────────────────────────────────────────────────────────────┤
│                                                                    │
│   Model:      NVIDIA Nemotron 70B                                  │
│   Framework:  NVIDIA Megatron-LM (tensor + pipeline parallelism)   │
│   Hardware:   NVIDIA RTX 3080 (sm_86) / RTX 4090 (sm_89)          │
│   ISA:        NVIDIA PTX (mma.sync, cp.async, ldmatrix, TMA)       │
│   Proofs:     Lean 4 (verified against NVIDIA hardware model)      │
│   Inference:  Deterministic (temperature=0.0, top_k=1)             │
│                                                                    │
│   Every layer is NVIDIA.                                           │
│   Every kernel is proven.                                          │
│   Every output is deterministic.                                   │
│                                                                    │
└────────────────────────────────────────────────────────────────────┘
```

---

## Hardware Targets

SnapKitty maintains verified kernel libraries for multiple NVIDIA architectures:

### RTX 3080 — Ampere (sm_86)

```
Architecture:  Ampere
Compute:       sm_86
VRAM:          10 GB GDDR6X (760 GB/s)
Tensor Cores:  3rd gen
Key PTX:       mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32
Async Copy:    cp.async.ca.shared.global + commit/wait groups
Pipeline:      3-stage (proven overlap bound: 1 - 1/3 = 66.7% utilization floor)
```

**Verified kernels:**
- `rtx_gemm_wmma.cu` — WMMA reference (128x128x32 CTA tiles)
- `rtx_gemm_ptx.cu` — Raw PTX mma.sync with ldmatrix
- `rtx_gemm_pipeline.cu` — 3-stage async pipeline with proven throughput
- `rtx_gemm_epilogue.cu` — Fused Bias+GeLU / Residual+GeLU epilogues

### RTX 4090 — Ada Lovelace (sm_89)

```
Architecture:  Ada Lovelace
Compute:       sm_89
VRAM:          24 GB GDDR6X (1008 GB/s)
Tensor Cores:  4th gen
Key PTX:       cp.async.bulk.tensor (TMA — Tensor Memory Accelerator)
Cluster:       __cluster_dims__ + barrier.cluster.* + multicast TMA
Pipeline:      4+ stage (TMA enables deeper overlap)
```

**Verified kernels:**
- `rtx_gemm_tma.cu` — TMA cluster algebra with multicast
- Cluster coherence invariant: `forall c in cluster. TMA_load(c) -> visible(c') within 1 cycle`
- Multicast law: `TMA_multicast(mask, T) = XOR_{c in mask} TMA_unicast(c, T)`

### What This Means for You

You tell PAX-Coder which GPU you have. It generates a kernel targeting exactly that architecture — not a generic CUDA kernel that might work, but a PTX-level implementation proven correct for your specific hardware.

```bash
# RTX 3080 (sm_86) — cp.async pipeline, no TMA
ollama run pax-coder "Write a verified GEMM for sm_86 with 3-stage pipeline"

# RTX 4090 (sm_89) — TMA cluster, deep pipeline
ollama run pax-coder "Write a verified GEMM for sm_89 with TMA multicast"
```

---

## Training Nemotron with Megatron-LM

### Why This Combination Works

Nemotron 70B already understands:
- CUDA memory hierarchy (global → L2 → shared → registers)
- PTX instruction semantics (what `mma.sync` actually computes)
- Tensor core data layouts (row-major A, column-major B, m16n8k8 fragments)
- Warp-level primitives (`shfl.sync`, `vote.sync`, `match.sync`)

We're not teaching a generic language model what CUDA is. We're taking NVIDIA's own model — which already has CUDA baked into its weights — and fine-tuning it to produce **formally verified** NVIDIA code.

The fine-tuning signal is the PAX corpus: ~2,400 verified triples of `(Lean 4 proof, PTX kernel, Futhark spec)`. After training, the model doesn't just write CUDA — it writes proven-correct CUDA.

### Training Configuration

```yaml
# Megatron-LM config for Nemotron PAX fine-tuning
model:
  name: nvidia/nemotron-70b-instruct
  tensor_parallel_size: 4
  pipeline_parallel_size: 2
  sequence_length: 4096
  
training:
  micro_batch_size: 1
  global_batch_size: 64
  learning_rate: 1.5e-5
  min_learning_rate: 1.0e-6
  lr_warmup_steps: 100
  lr_decay_style: cosine
  weight_decay: 0.01
  clip_grad: 1.0
  bf16: true
  
data:
  dataset: pax-verified-corpus
  format: nemotron_chat_template
  categories:
    - fp16_rounding      # IEEE-754 binary16 proofs
    - gemm_kernels       # mma.sync implementations
    - pipeline_overlap   # cp.async throughput bounds
    - epilogue_fusion    # Bias+GeLU algebraic laws
    - warp_primitives    # shfl.sync reductions
    - architecture       # PAX axiom mappings

loss:
  type: po_weighted_cross_entropy
  weights:
    lean4_proof: 2.0     # Proof correctness is highest priority
    ptx_kernel: 1.5      # Implementation correctness
    futhark_spec: 1.0    # Spec adherence
    certificate: 0.5     # PO tagging

inference:
  temperature: 0.0       # Deterministic — proofs don't have "creative" answers
  top_k: 1
  repetition_penalty: 1.0
```

### The PO-Weighted Loss Function

Standard cross-entropy treats every token equally. We weight proof tokens higher than comment tokens:

```
L = -sum_i w(category_i) * log P(token_i | context)

Where:
  w(lean4_proof)  = 2.0  — getting a theorem wrong is unacceptable
  w(ptx_kernel)   = 1.5  — implementation must match the proof
  w(futhark_spec) = 1.0  — spec is the reference
  w(certificate)  = 0.5  — tagging is secondary
```

This produces a model that prioritizes correctness over style.

### Deterministic Generation

PAX-Coder runs at temperature 0.0 with top_k=1. There is no sampling, no creativity, no stochastic variation.

Why: A proof is either correct or it isn't. `2 + 2 = 4` every time. A model generating formal proofs must be deterministic.

```python
# Inference — zero entropy
output = model.generate(
    input_ids,
    temperature=0.0,
    top_k=1,
    top_p=1.0,
    repetition_penalty=1.0,
    do_sample=False,
    max_new_tokens=2048
)
```

Same input → same kernel → same proof. Every time.

---

## Training on RTX 3080 (Single GPU)

For the public PAX-Coder (7B, based on DeepSeek-Coder), single-GPU training fits on the RTX 3080:

```bash
# VRAM budget — RTX 3080 10GB:
#   Base model (4-bit QLoRA)  ~4.2 GB
#   LoRA adapters (r=32)     ~0.1 GB
#   Gradients (8-bit paged)  ~1.5 GB
#   Activations (GC)         ~1.8 GB
#   Dataset buffer           ~0.5 GB
#   Total                    ~8.1 GB (1.9 GB headroom)

# One command:
./run_training.sh

# What it does:
#   1. Checks free VRAM (needs ~8GB)
#   2. Extracts training data from PAX corpus → JSONL
#   3. Fine-tunes DeepSeek-Coder-7B with QLoRA
#   4. Exports to GGUF for Ollama
#   5. ~4-6 hours on RTX 3080
```

### Full Nemotron 70B (Multi-GPU)

The full sovereign Nemotron model requires distributed training via Megatron-LM:

```bash
# Multi-node launch (4× A100 80GB or 8× RTX 4090 24GB)
torchrun \
  --nproc_per_node=4 \
  --nnodes=1 \
  --master_port=29500 \
  pretrain_gpt.py \
  --tensor-model-parallel-size 4 \
  --pipeline-model-parallel-size 1 \
  --num-layers 80 \
  --hidden-size 8192 \
  --num-attention-heads 64 \
  --seq-length 4096 \
  --micro-batch-size 1 \
  --global-batch-size 64 \
  --lr 1.5e-5 \
  --train-iters 2000 \
  --bf16 \
  --data-path pax-verified-corpus \
  --save checkpoints/nemotron-pax \
  --load nvidia/nemotron-70b-instruct
```

---

## What Makes This Sovereign

| Property | What it means |
|----------|--------------|
| **No cloud dependency** | Runs on local NVIDIA hardware. No API keys, no rate limits, no vendor lock-in. |
| **No trust dependency** | Every output is machine-checked. You don't trust the model — you verify its proofs. |
| **No data dependency** | Training corpus is self-generated from the PAX codebase. Not GitHub scrapes. |
| **Deterministic** | Same prompt → same output. Auditable, reproducible, provable. |
| **Hardware-native** | Model writes for the GPU it runs on. No abstraction layers. Raw PTX. |

This is what sovereign compute means: you own the hardware, you own the model, you own the training data, and you can verify every output.

---

## Custom NVIDIA Builds

SnapKitty offers custom kernel builds targeting your specific NVIDIA hardware:

### What You Get

```
┌──────────────────────────────────────────────────────┐
│  YOUR GPU → PAX-Coder generates:                     │
│                                                      │
│  1. Lean 4 correctness proof (zero sorry)            │
│  2. PTX kernel targeting YOUR sm_XX arch             │
│  3. Futhark functional spec (ground truth)           │
│  4. PAX certificate (which POs are satisfied)        │
│  5. NCU-ready binary (compile + profile)             │
│                                                      │
│  Not generic CUDA. YOUR hardware. PROVEN correct.    │
└──────────────────────────────────────────────────────┘
```

### Supported Architectures

| GPU | Architecture | Compute | Key Feature | Status |
|-----|-------------|---------|-------------|--------|
| RTX 3080 | Ampere | sm_86 | cp.async 3-stage pipeline | **Verified** |
| RTX 3090 | Ampere | sm_86 | Same as 3080 + 24GB VRAM | **Verified** |
| RTX 4090 | Ada Lovelace | sm_89 | TMA + cluster multicast | **Verified** |
| A100 | Ampere | sm_80 | Async copy + large shared | **Verified** |
| H100 | Hopper | sm_90 | TMA + warp specialization | **Spec complete** |

### How to Order

```bash
# 1. Tell us your GPU
echo "RTX 4090" | pax-coder --target sm_89

# 2. Tell us your computation
echo "128x128 GEMM with Residual+GeLU epilogue, FP16 in, FP32 accum"

# 3. Get back:
#    - verified_gemm_sm89.ptx    (your kernel)
#    - verified_gemm_sm89.lean   (your proof)
#    - verified_gemm_sm89.fut    (your spec)
#    - CERTIFICATE.json          (PO1-PO8 status)
```

---

## The Competitive Advantage

| | Generic Code Models | cuBLAS | PAX-Coder |
|--|---|---|---|
| **Correctness** | Hope-based | Tested, not proven | Machine-checked proof |
| **Hardware targeting** | Generic CUDA | Black box | Architecture-specific PTX |
| **Reproducibility** | Temperature sampling | Deterministic | Deterministic + auditable |
| **Verification** | None | Benchmarks | Lean 4 formal proof |
| **Customization** | Prompt engineering | Library calls | Fine-tuned for YOUR arch |
| **Sovereignty** | Cloud API | Proprietary | Runs on YOUR GPU |

---

## Repo Structure (PAX-Coder)

```
pax-coder/
├── PAX/                        Lean 4 formal proofs (training source)
├── src/
│   ├── rtx_gemm_ptx.cu        sm_86 GEMM — raw mma.sync
│   ├── rtx_gemm_pipeline.cu   sm_86 3-stage async — proven overlap
│   ├── rtx_gemm_epilogue.cu   sm_86 Bias+GeLU fusion — proven bounds
│   └── pax_kernel.fut         Futhark functional spec
├── demo/
│   ├── demo.py                Live inference demo
│   └── showcase_examples.jsonl Example prompts + outputs
├── docs/
│   ├── PAX_ARCHITECTURE.md    5 axioms → 8 proof obligations
│   └── SOVEREIGN_NVIDIA_TRAINING_GUIDE.md   (this file)
├── train.py                   QLoRA fine-tuning (RTX 3080 single-GPU)
├── export_training_data.py    PAX corpus → JSONL extraction
├── run_training.sh            One-command training launcher
├── Modelfile                  Ollama deployment
├── LICENSE.tri                BSL-1.1 / AGPL-3.0 / MPL-2.0
└── SOVEREIGN_NODE_KEY.md      Production access
```

---

## Getting Started

### Option 1: Use PAX-Coder directly (pre-trained, public)

```bash
ollama pull Snapkitty/pax-coder
ollama run pax-coder "Write a verified GEMM for my RTX 3080"
```

### Option 2: Train your own PAX model on your NVIDIA GPU

```bash
git clone https://github.com/SNAPKITTYWEST/pax-coder
cd pax-coder
./run_training.sh  # ~4-6h on RTX 3080
```

### Option 3: Custom sovereign build (enterprise)

Contact `licensing@snapkittywest.dev` for:
- Architecture-specific kernel libraries
- On-premises Nemotron deployment
- Formal verification consulting
- Custom PO audits

---

## License

Tri-licensed: BSL-1.1 / AGPL-3.0 / MPL-2.0

Copyright (C) 2026 Bel Esprit D'Accord Irrevocable Trust
SnapKitty Collective Limited

Authors: Ahmad Ali Parr — Jessica Westerhoff