custom
code
sovereign-compute
pax-coder / docs /SOVEREIGN_NVIDIA_TRAINING_GUIDE.md
SNAPKITTYWEST's picture
chore: push pax-coder from SNAPKITTYWEST GitHub
ef6eb55 verified
|
Raw
History Blame Contribute Delete
14.8 kB

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.

# 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

# 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.

# 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:

# 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:

# 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

# 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)

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

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