HuggingFace Hub

HuggingFace Hub

NNx ships first-class interop with the HuggingFace ecosystem via two independent mechanisms:

1. safetensors as an opt-in checkpoint format on NNCheckpoint — safe (no arbitrary-code unpickling), mmap-friendly, and readable by ComfyUI / vLLM / AutoGPTQ / transformers tools.
2. PyTorchModelHubMixin on NNModel — free save_pretrained / push_to_hub / from_pretrained for distributing models via the Hub.

Both paths require the hub extra:

pip install "thekaveh-nnx[hub]"

Without it, the rest of NNx keeps working — the integration is gated behind import-time guards. Calling any Hub method without the extra raises a clear ImportError pointing at this install line.

safetensors checkpoints

When to use safetensors

Use safetensors when any of the following is true:

• You plan to publish weights to the Hub or share them outside your machine. Pickle checkpoints can execute arbitrary code on load; safetensors cannot.
• You need to load weights into a non-Python tool (ComfyUI, vLLM, AutoGPTQ, transformers-aware loaders).
• You care about mmap-based zero-copy loads — safetensors files are laid out so tensors can be mapped directly from disk.

Pickle remains the default and is the right choice for local-only training runs where the convenience of torch.save-ing the full checkpoint (preserving OrderedDict key order and NNCheckpoint identity) outweighs the security trade-off.

Writing a safetensors checkpoint

NNCheckpoint.to_file accepts a format kwarg:

from nnx import NNCheckpoint

# Build a checkpoint as usual…
ckpt = NNCheckpoint(
    idp=...,                # NNIterationDataPoint
    model_params=model.params,
    net_params=model.net_params,
    net_state=model.net.state_dict(),
)

# …then write either format. Pickle is the default.
ckpt.to_file("checkpoint.pt")                            # legacy default
ckpt.to_file("checkpoint.safetensors", format="safetensors")

NNParams, NNModelParams, and NNIterationDataPoint are JSON-serialized into the safetensors metadata dict (the format spec limits metadata to str → str). Net tensors are detached, made contiguous, and written through safetensors' save_file. Writes are atomic — staged at <path>.tmp then os.replace-d — matching the same interrupt-safety guarantee as the pickle path.

Reading a checkpoint of either format

NNCheckpoint.from_file auto-detects the format by sniffing the first few bytes:

• Modern torch.save → ZIP container starting with b"PK\x03\x04".
• Legacy torch.save / bare pickle → \x80 PROTO opcode.
• safetensors → little-endian u64 header length followed by a JSON object (byte 8 = {).

ckpt = NNCheckpoint.from_file("checkpoint.safetensors")  # or .pt
model = NNModel.from_checkpoint(ckpt)

Publishing an NNModel to the Hub

When to use the Hub mixin

Use save_pretrained / push_to_hub / from_pretrained for distribution: shipping a trained NNModel so others can from_pretrained("you/your-model") and run it. The flat on-disk layout (model.safetensors + config.json + README.md) is what the Hub expects.

Keep using NNCheckpoint for local training state: the runs/<id>/checkpoints/ layout carries per-epoch IDPs, optimizer state sidecars, and run.id-keyed metadata that the Hub layout deliberately strips.

save_pretrained — save locally

from nnx import NNModel, NNParams, NNModelParams, Activations, Devices, Losses, Nets

model = NNModel(
    net_params=NNParams(
        input_dim=4, output_dim=2, hidden_dims=[8],
        dropout_prob=0.0, activation=Activations.RELU,
    ),
    params=NNModelParams(
        net=Nets.FEED_FWD, device=Devices.CPU, loss=Losses.CROSS_ENTROPY,
    ),
)
# …train…
model.save_pretrained("./my-model")

This writes three files into ./my-model/:

• model.safetensors — self.net.state_dict() as safetensors.
• config.json — {"net_params": <state>, "params": <state>}, using the same public state() form NNRun hashes for run.id grouping.
• README.md — auto-generated model card from the mixin.

from_pretrained — load from a local directory

from nnx import NNModel

model = NNModel.from_pretrained("./my-model")

from_pretrained reads config.json, rebuilds NNParams and NNModelParams via their public from_state constructors, then loads the safetensors weights into the freshly-built net. Bit-exact round-trip on tensors; state() form identical on the params.

push_to_hub — publish to the Hub

# One-time login (writes a token to ~/.cache/huggingface/token):
#   hf auth login

model.push_to_hub("your-user/your-model")

The mixin handles repo creation, file upload, and commit. Everything save_pretrained writes locally is pushed.

from_pretrained — load from the Hub

model = NNModel.from_pretrained("your-user/your-model")

HuggingFace's cache directory is used transparently — repeat loads hit the local cache, not the network.

What this does NOT do

• NNRun is not Hub-published. The Hub layout is per-model. To publish a full training run (idps.csv + run.yaml + every per-phase checkpoint), upload the runs/<id>/ directory directly via huggingface_hub.upload_folder.
• Optimizer state is not in the Hub config. save_pretrained writes only network weights; resuming optimizer state from a Hub-loaded model is not supported. Use NNCheckpoint for warm-resume workflows.
• The Hub mixin does not rewrite NNModel's constructor. It still takes (net_params, params) keyword args at __init__ — the mixin is purely additive.

Reference

Symbol	Kind	Notes
save_pretrained	NNModel method	Write model.safetensors + config.json + README.md to a local directory.
push_to_hub	NNModel method	Publish the same artifacts to the HuggingFace Hub.
from_pretrained	NNModel classmethod	Load from local directory or Hub repo identifier.
safetensors	Format	Via NNCheckpoint.to_file(format="safetensors"). Auto-detected on from_file.

See also

• Persistence-Runs-and-Checkpoints · Language-Modeling · GGUF-and-Ollama
• API reference · source: src/nnx/