Publish NanoJev Core ML conversion source without trained weights

source/STATUS.md

@@ -1,6 +1,7 @@
 # NanoJev status
 
 - Pinned current checkpoint and audited native source; full checkpoint SHA256 matches upstream and contains 596,250,498 trained parameters across 322 tensors. Historical tracker revision unresolved.
-- Pinned full checkpoint downloaded. Native Qwen3 encoder plus set-head exporter matched the trained model on a 4-candidate Choice request: same argmax, maximum logit error `1.07e-6`.
-- Core ML export and runtime parity for Choice, Boolean, and Score remain pending.
+- Pinned full checkpoint downloaded. Native Qwen3 encoder plus decision-head exporter matched the trained model on Choice, Boolean, and Score requests: same argmax for all three, maximum logit error `1.67e-6`.
+- FP16 Core ML L128/K4 encoder and full decision-head packages exported in 22.7 seconds. Encoder package is about 1.1 GiB; head package is about 420 KiB.
+- Core ML runtime parity for Choice, Boolean, and Score remains pending.
 - Trained-weight redistribution license unresolved. Public HF model artifact is not yet authorized by the source metadata.

source/convert-coreml.py

@@ -46,6 +46,7 @@ def main() -> None:
     parser.add_argument("--candidates", type=int, default=4)
     parser.add_argument("--output-dir", type=Path, default=ROOT / "build")
     parser.add_argument("--parity-only", action="store_true")
+    parser.add_argument("--trace-only", action="store_true")
     args = parser.parse_args()
     if args.candidates < 2:
         raise ValueError("candidate bucket must be at least 2 for Boolean output")
@@ -62,6 +63,11 @@ def main() -> None:
     args.output_dir.mkdir(parents=True, exist_ok=True)
     started = time.perf_counter()
     encoder_trace = torch.jit.trace(encoder, tuple(torch.from_numpy(value) for value in inputs.values()))
+    if args.trace_only:
+        for index, node in enumerate(encoder_trace.inlined_graph.nodes()):
+            if node.kind() == "aten::Int":
+                print(f"trace node {index}: {node} scope={node.scopeName()}", flush=True)
+        return
     encoder_coreml = ct.convert(
         encoder_trace,
         convert_to="mlprogram",

source/export_model.py

@@ -16,8 +16,7 @@ def rotate_half(value: torch.Tensor) -> torch.Tensor:
     return torch.cat((-second, first), dim=-1)
 
 
-def rms_norm(hidden: torch.Tensor, norm: nn.Module) -> torch.Tensor:
-    width = norm.weight.numel()
+def rms_norm(hidden: torch.Tensor, norm: nn.Module, width: int) -> torch.Tensor:
     doubled = torch.cat((hidden, -hidden), dim=-1)
     normalized = F.layer_norm(doubled, (width * 2,), eps=norm.variance_epsilon)
     return normalized[..., :width] * norm.weight
@@ -51,7 +50,7 @@ class NanoEncoder(nn.Module):
 
     def _layer(self, layer: nn.Module, hidden: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
         attention = layer.self_attn
-        normalized = rms_norm(hidden, layer.input_layernorm)
+        normalized = rms_norm(hidden, layer.input_layernorm, self.hidden)
         q = attention.q_proj(normalized).view(self.candidates, self.length, self.heads, self.head_dim)
         k = attention.k_proj(normalized).view(self.candidates, self.length, self.kv_heads, self.head_dim)
         v = attention.v_proj(normalized).view(self.candidates, self.length, self.kv_heads, self.head_dim)
@@ -68,7 +67,7 @@ class NanoEncoder(nn.Module):
             self.candidates, self.length, self.heads * self.head_dim
         )
         hidden = hidden + attention.o_proj(attended)
-        return hidden + layer.mlp(rms_norm(hidden, layer.post_attention_layernorm))
+        return hidden + layer.mlp(rms_norm(hidden, layer.post_attention_layernorm, self.hidden))
 
     def forward(
         self, input_ids: torch.Tensor, attention_mask: torch.Tensor, eos_map: torch.Tensor
@@ -78,7 +77,7 @@ class NanoEncoder(nn.Module):
         hidden = self.embed_tokens(input_ids.long())
         for layer in self.layers:
             hidden = self._layer(layer, hidden, mask)
-        hidden = rms_norm(hidden, self.norm).float()
+        hidden = rms_norm(hidden, self.norm, self.hidden).float()
         return torch.matmul(eos_map, hidden).transpose(0, 1)