neuralfn-torch

name

neuralfn-torch

description

Build, train, and export torch-backed neural network models (GPT, Llama, MoE, Jamba, JEPA, diffusion, etc.) using the NeuralFn Python API. Use whenever the user asks to build a language model, train a transformer, use template presets, configure ModelSpec/BlockSpec, compile a torch graph, export weights, or do autoregressive inference with NeuralFn in Python code. For MCP tool operations, use neuralfn-mcp instead.

NeuralFn Torch Models

Use this skill when building, training, or exporting torch-backed models with the NeuralFn Python API. For core graph operations, see neuralfn-python-sdk. For MCP tools, see neuralfn-mcp.

For detailed torch backend and preset reference, see presets-reference.md.

Full API documentation lives in the repo at docs/ (index). For a single-file LLM-ready dump of all docs, see llms-full.txt.

End-to-end example: build, train, export

from neuralfn import build_gpt_root_graph, TorchTrainer, TorchTrainConfig
from neuralfn.config import build_llama_spec
from neuralfn.inference import export_to_pt, InferenceCache
import torch

# 1. Build model graph
spec = build_llama_spec(n_layer=4, num_heads=4, n_embd=128, num_kv_heads=2)
graph = build_gpt_root_graph(name="my_llama", model_spec=spec)

# 2. Train
trainer = TorchTrainer(graph, TorchTrainConfig(
    epochs=10, learning_rate=5e-3, batch_size=4, device="cuda"
))
losses = trainer.train(
    train_inputs=[[1,2,3,4],[2,3,4,5],[3,4,5,6]],
    train_targets=[[2,3,4,5],[3,4,5,6],[4,5,6,7]],
)

# 3. Export
export_to_pt(graph, "my_llama.pt")

# 4. Inference
cache = InferenceCache(graph, device="cuda")
prompt = torch.tensor([[1, 2, 3]], dtype=torch.long)
logits = cache.step(prompt)
next_token = logits.argmax(dim=-1)

Building graphs with presets

from neuralfn import build_gpt_root_graph
from neuralfn.config import build_nanogpt_spec

spec = build_nanogpt_spec(
    n_layer=4,      # transformer layers
    num_heads=4,    # attention heads
    n_embd=128,     # model dimension
    vocab_size=256, # auto-adjusted by trainer
)
graph = build_gpt_root_graph(name="model", model_spec=spec)

The graph has runtime="torch", training_method="torch", and a populated variant_library with attention and MLP subgraph variants.

Shipped preset catalog

Preset	Architecture	Key features
nanogpt	GPT-2 style	LayerNorm, GELU MLP, absolute position embeddings
gpt2	GPT-2	LayerNorm, GELU MLP, absolute pos, linear bias
llama	LLaMA	RMSNorm, SwiGLU, RoPE, GQA
moe / mixllama	LLaMA + MoE	RMSNorm, MoE MLP, RoPE, GQA
llama_fast	LLaMA + compile	Like llama with torch.compile
mixllama_fast	MoE + compile	Like moe with torch.compile
jamba	Jamba hybrid	Attention + Mamba interleaved, MoE
ternary_b158	BitNet b1.58	Ternary {-1, 0, 1} weights
seq2seq	Encoder-decoder	Seq2Seq objective, MoE MLP
diffusion	Discrete diffusion	Diffusion objective with denoising head
ttt_llama	TTT-Linear	Test-time training attention replacement
llm_jepa	LLM-JEPA	JEPA with EMA target encoder
dense_jepa_evo	Dense JEPA Evo	Non-semantic AR+JEPA control with dense FFNs
moe_jepa_evo	MoE JEPA Evo	Non-semantic AR+JEPA control with standard MoE routing
semantic_router_moe	Semantic Router MoE	AR-only semantic router control with shared routed MoE blocks
semantic_dense_jepa_evo	Semantic Dense JEPA Evo	Chunk-level semantic planner, dense FFNs, JEPA supervision, no route evolution
semantic_moe_jepa_evo	Semantic MoE JEPA Evo	Chunk-level semantic router, shared/semantic/free experts, JEPA supervision, route evolution
hnet_lm	H-Net	Raw byte input, byte patch embedding
universal_llama	Universal TX	ACT-based adaptive recurrence
llama_megakernel	Fused LLaMA	FusedCausalAttention, max-autotune compile
kv_pca_llama	PCA KV cache	PCA-compressed keys/values
deepseek_v3	DeepSeek-V3	MLA + auxfree-balanced MoE + shared experts
deepseek_v4	DeepSeek-V4-Pro	NSA attention + auxfree MoE + mHC residuals + QK-norm + FP8
gemma3	Gemma-2/3	Sliding-window attention + GeGLU + QK-norm + softcap
diff_transformer	Differential TX	Two-softmax differential attention + head-wise norm
qwen3_longctx	Qwen long-ctx	GQA + YaRN RoPE scaling + QK-norm
longctx_sparse_llama	Long-ctx sparse	NSA / block-sparse / sliding-window / streaming
modern_norms_llama	Modern norms	DyT + QK-norm + GeGLU
fp8_llama / mxfp4_llama	Blackwell precision	FP8 E4M3 / MXFP4 microscaled weight linears
auxfree_moe_jepa_evo, diff_semantic_moe_jepa_evo, dyt_geglu_semantic_dense_jepa_evo	NeuralFn crosses	Modern kernels × JEPA/semantic/route-evo stacks
<preset>_modern	Modernized	Any base preset + RMSNorm/QK-norm/RoPE-YaRN/GeGLU/auxfree (see MODERN_BASE_PRESETS)

Common config keys

Key	Default	Description
n_layer / num_layers	4	Transformer layers
n_head / num_heads	4	Attention heads
n_embd / model_dim	128	Model dimension
vocab_size	256	Vocabulary (auto-adjusted by trainer)
num_kv_heads	2	GQA key/value heads
mlp_multiplier	8/3 (llama) or 4 (gpt2)	MLP hidden multiplier
multiple_of	256	Round MLP width to multiple
experts	8	MoE: number of experts
top_k	2	MoE: experts per token
dropout_p	0.0 or 0.1	Dropout rate
tie_embeddings	varies	Tie embedding/LM head weights
logit_softcap	0.0	Tanh softcap (>0 enables)
ttt_hidden_dim	32	TTT hidden dimension
byte_patch_size	4	H-Net byte patch size
max_recurrence_steps	4	Universal TX max steps

Programmatic spec building

from neuralfn.config import build_llama_spec, ModelSpec
from neuralfn.torch_templates import build_model_stage_graph, build_gpt_template_payload

# Build a ModelSpec directly
spec = build_llama_spec(n_layer=6, n_embd=256, num_heads=8, num_kv_heads=4)

# Build just the model stage subgraph
stage_graph = build_model_stage_graph("model_stage", spec)

# Build a full payload for server/editor template APIs
payload = build_gpt_template_payload("my_model", {"preset": "llama", "n_layer": 6, "n_embd": 256})

Spec builders: build_nanogpt_spec, build_nanogpt_megakernel_spec, build_gpt2_spec, build_gpt2_megakernel_spec, build_llama_spec, build_mixllama_spec, build_llama_fast_spec, build_llama_fast_megakernel_spec, build_mixllama_fast_spec, build_mixllama_fast_megakernel_spec, build_jamba_hybrid_spec, build_ternary_b158_spec, build_decoder2encoder_moe_spec, build_diffllama_spec, build_ttt_llama_spec, build_llm_jepa_spec, build_dense_jepa_evo_spec, build_moe_jepa_evo_spec, build_semantic_router_moe_spec, build_semantic_router_moe_megakernel_spec, build_jepa_semantic_hybrid_spec, build_jepa_semantic_hybrid_megakernel_spec, build_semantic_dense_jepa_evo_spec, build_semantic_moe_jepa_evo_spec, build_hnet_lm_spec, build_universal_llama_spec, build_llama_megakernel_spec, build_kv_pca_llama_spec, and build_composed_lm_spec. Frontier builders: build_deepseek_v3_spec, build_deepseek_v4_spec, build_gemma3_spec, build_diff_transformer_spec, build_qwen3_longctx_spec, build_longctx_sparse_llama_spec, build_modern_norms_llama_spec, build_fp8_llama_spec, build_mxfp4_llama_spec, build_auxfree_moe_jepa_evo_spec, build_diff_semantic_moe_jepa_evo_spec, build_dyt_geglu_semantic_dense_jepa_evo_spec. Modernized variants are generated as <preset>_modern (dispatch strips the suffix and applies _apply_modern_profile).

TorchTrainConfig

Field	Default	Description
learning_rate	3e-4	Adam learning rate
epochs	50	Training epochs
batch_size	8	Batch size
weight_decay	0.01	AdamW weight decay
device	"cuda"	Device ("cuda", "cpu")
amp_dtype	"float32"	AMP dtype; float32 disables autocast
compile	False	Use torch.compile
activation_checkpointing	False	Gradient checkpointing
fsdp2_enabled	False	FSDP2 sharding
max_steps	None	Step cap (None = epoch-based)

Training with datasets

# With inline data
losses = trainer.train(
    train_inputs=[[1,2,3,4],[2,3,4,5]],
    train_targets=[[2,3,4,5],[3,4,5,6]],
)

# With HuggingFace dataset name (must be downloaded first via server API)
losses = trainer.train(dataset_names=["HuggingFaceFW__fineweb"], seq_len=64)

Dataset roles by objective:

• AR / H-Net / Universal: tokens, targets
• Seq2Seq: enc_tokens, dec_tokens, targets
• Diffusion / JEPA: tokens
• Semantic routing presets: tokens, targets, plus semantic_data_source -> sem_targets

CompiledTorchGraph

from neuralfn.torch_backend import CompiledTorchGraph

compiled = CompiledTorchGraph(graph)  # compiles NeuronGraph to nn.Module
compiled.to("cuda")

# Forward pass
outputs = compiled(token_ids, targets)

# Trace (returns dict of node_id -> tensor stats)
trace = compiled.trace(token_ids, targets)

# Sync weights back to graph JSON
compiled.sync_state_back(graph)

Weight export/import

from neuralfn.inference import export_to_pt, import_from_pt, export_quantized_pt, import_quantized_pt

export_to_pt(graph, "model.pt")
import_from_pt(graph, "model.pt")

export_quantized_pt(graph, "model_q.pt", scheme="int8")   # or "ternary"
import_quantized_pt(graph, "model_q.pt")

InferenceCache (autoregressive generation)

from neuralfn.inference import InferenceCache
import torch

cache = InferenceCache(graph, device="cuda")
prompt = torch.tensor([[1, 2, 3, 4]], dtype=torch.long)
logits = cache.step(prompt)         # full prompt
next_tok = logits.argmax(dim=-1)
logits2 = cache.step(next_tok.unsqueeze(1))  # single token step
cache.reset()                       # clear for new sequence

Works with graphs that have kv_cache_read / kv_cache_write nodes. For training graphs (2 inputs), dummy targets are supplied automatically.

Experimental Presets

semantic_router_moe [Experimental]

• Preset: semantic_router_moe [Experimental]
• Load in Python: from neuralfn.config import build_semantic_router_moe_spec; then spec = build_semantic_router_moe_spec(**kwargs) and build_gpt_root_graph(name=..., model_spec=spec).
• Load via MCP / server: load_gpt_template(name=..., preset="semantic_router_moe", config={...}) [Experimental].
• What it does [Experimental]: AR-only MixLLaMA/MoE control preset that computes a vocab-grounded semantic route once from the pre-block hidden state, hashes it, teacher-forces/auto-selects one expert per semantic vocabulary dimension, broadcasts that route across the whole sequence, and applies it to every MoE block. Trains next-token CE plus semantic-alignment loss, with no JEPA encoder/EMA path.
• Disclaimer [Experimental]: Research-control preset only. It exists to isolate the router hypothesis before adding JEPA complexity.

semantic_moe_jepa_evo [Experimental]

• Preset: semantic_moe_jepa_evo [Experimental]
• Load in Python: from neuralfn.config import build_semantic_moe_jepa_evo_spec; then spec = build_semantic_moe_jepa_evo_spec(**kwargs) and build_gpt_root_graph(name=..., model_spec=spec).
• Load via MCP / server: load_gpt_template(name=..., preset="semantic_moe_jepa_evo", config={...}) [Experimental].
• What it does [Experimental]: Full Semantic MoE JEPA Evo architecture. Dense causal attention stays on the AR path; a prefix-safe chunk planner predicts semantic latents and route distributions; routes combine always-on shared experts, semantic-vocabulary experts, and free learned experts; and the trainer can periodically evolve route bias/table state.
• Config rules [Experimental]: experts must equal semantic_shared_experts + NUM_VOCAB_DIMS + semantic_free_experts. Defaults are route_chunk_size=32, semantic_shared_experts=2, semantic_free_experts=8, route_evo_fraction=0.10, and route_evo_population=8.
• Disclaimer [Experimental]: Research prototype only. Graph shape, loss mix, and route-evolution behavior may change.

semantic_dense_jepa_evo [Experimental]

• Preset: semantic_dense_jepa_evo [Experimental]
• Load in Python: from neuralfn.config import build_semantic_dense_jepa_evo_spec; then spec = build_semantic_dense_jepa_evo_spec(**kwargs) and build_gpt_root_graph(name=..., model_spec=spec).
• Load via MCP / server: load_gpt_template(name=..., preset="semantic_dense_jepa_evo", config={...}) [Experimental].
• What it does [Experimental]: Dense control for the Semantic JEPA Evo architecture. It keeps the prefix-safe chunk planner, JEPA target supervision, AR CE, JEPA latent alignment, and semantic-alignment losses, but uses dense LLaMA FFNs with no expert dispatch, route losses, or route-evolution loop.
• Config rules [Experimental]: route_chunk_size controls planner chunk boundaries. Expert-count and route-evolution fields are ignored by the dense decoder path.
• Disclaimer [Experimental]: Dense comparison/control preset only. Graph layout and tuning knobs may change.

jepa_semantic_hybrid [Experimental]

• Preset: jepa_semantic_hybrid [Experimental]
• Load in Python: from neuralfn.config import build_jepa_semantic_hybrid_spec; then spec = build_jepa_semantic_hybrid_spec(**kwargs) and build_gpt_root_graph(name=..., model_spec=spec).
• Load via MCP / server: load_gpt_template(name=..., preset="jepa_semantic_hybrid", config={...}) [Experimental].
• What it does [Experimental]: Joint Embedding Predictive Architecture (JEPA) combined with a vocab-grounded semantic state, LSH bucketing, a fixed dimension-to-expert semantic router, and routed full-sequence attention experts. sem_targets are categorical topic IDs with ignore sentinels, not quantized semantic vectors.
• Disclaimer [Experimental]: Research prototype only—APIs, graph shape, and training behavior may change without notice.