5. HuggingFace Hub integration¶

NNx ships first-class interop with the HuggingFace ecosystem:

safetensors as an opt-in checkpoint format on NNCheckpoint — safe (no arbitrary-code unpickling), mmap-friendly, and readable by ComfyUI / vLLM / AutoGPTQ / transformers tools.
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 back at this install line.

A third, Hub-independent format — the run bundle (nnx.bundles, §3) — carries a run checkpoint's training state and calibrators as data too. It needs no optional dependency.

1. safetensors checkpoints¶

1.1. When to use it¶

Use safetensors when any of the following is true:

You plan to publish weights to the Hub or share them with anyone outside your machine. Pickle checkpoints can execute arbitrary code on load (see the security note on NNCheckpoint.from_file); 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 the underlying tensors can be mapped directly from disk.

Pickle remains the default. It is the right choice for local-only training runs where the convenience of torch.save-ing the full dataclass outweighs the security trade-off. That save preserves the OrderedDict key order and the NNCheckpoint identity.

1.2. Writing a safetensors checkpoint¶

NNCheckpoint.to_file takes 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, so a JSON wrapper is the cleanest fit. The net's tensors are detached and cloned, because safetensors refuses tensors that share storage. They are then written through safetensors' standard save_file.

Writes are atomic. The file is staged in a hidden temporary file beside <path>, flushed to disk and os.replace-d into place. This matches the same KeyboardInterrupt-safety guarantee that the pickle path provides.

1.3. Reading a checkpoint of either format¶

NNCheckpoint.from_file auto-detects which format the file was written in by sniffing the first few bytes:

Modern torch.save produces a ZIP container that starts with b"PK\x03\x04".
Legacy torch.save (with _use_new_zipfile_serialization=False) and bare pickle files start with the \x80 PROTO opcode.
safetensors files start with a little-endian u64 header length followed by a JSON object — byte 8 is always {. (The u64's low byte can legitimately be 0x80, colliding with the pickle PROTO opcode. The loader therefore positively identifies safetensors via byte 8 before the pickle sniff.)

The same call works for both:

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

2. Publishing an NNModel to the Hub¶

2.1. 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 this writes (model.safetensors + config.json + README.md) is what the Hub expects, and it's what downstream tools probe for.

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

2.2. Save a model 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. A model built from a registered factory (§2.6) has no net_params; its params.net is the ModelSpec descriptor.

A recorded recipe (nnx.transforms, see Surgery §8) can change a model's topology. Such a model writes its ordered operations to "transforms", each with its id, version, targets and config. Without a recipe that list is empty. It is saved only while its topology is exactly its base plus that recipe. Unrecorded surgery is refused before anything is written, rather than failing at load.
README.md — auto-generated model card from the mixin.

2.3. Load from a local directory¶

from nnx import NNModel
model = NNModel.from_pretrained("./my-model")

from_pretrained reads config.json and rebuilds NNParams and NNModelParams via their public from_state constructors. It replays any recorded "transforms" on the fresh base: LoRA wrappers are rebuilt, and low-rank factors are allocated in their recorded shape (no SVD). It then loads the safetensors weights into the freshly-built self.net. An unknown operation or version is refused before any tensor is loaded. Tensors round-trip bit-exactly.

The params' state() round-trips too, except NNModelParams.device. It takes map_location (default "cpu"), so a GPU or MPS model reloads on the CPU unless you pass one.

The saved precision policy is resolved again on the destination device. precision= replaces it, and the legacy mixed_precision flag with it.

A model can be saved with extra top-level submodules attached to its net, such as a JEPA predictor registered as model.net._jepa_predictor. It rebuilds without them through exclude_submodules. Every other key still loads strictly (strict=False would also accept missing weights). A name the saved weights do not hold is refused. Without it, the error names the extra submodules and this option:

encoder = NNModel.from_pretrained("./jepa-model", exclude_submodules=("_jepa_predictor",))

2.4. 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 that save_pretrained writes locally is pushed.

2.5. Load from the Hub¶

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

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

2.6. Registered and runtime-only modules¶

A model built from a registered factory (ModelSpec, see Concepts §3.2) round-trips through the Hub like a built-in one. config.json stores the spec (id, version, config, seed) instead of net_params. from_pretrained rebuilds the module through the factory registry, so the same register_model_factory(...) must have run in the loading process.

An unregistered factory raises MissingModelFactoryError before anything is built. A factory whose topology no longer matches the weights raises before they load. Pass batch_adapter= to from_pretrained for a module that needs one.

A runtime-only module (NNModel(module=...), reconstructible=False) has no factory to rebuild it. save_pretrained / push_to_hub therefore reject it with MissingModelFactoryError before any directory or file is written. Register a factory and train from a ModelSpec to publish it.

3. Three artifact formats and their trust boundaries¶

NNx writes three kinds of artifact. Each has one reader, and no reader opens another's files:

Format	Written by	Holds	Read by	Trust boundary
Pickle checkpoint `runs/<id>/checkpoints/<tag>.pt` (+ `.opt.<generation>.pt` sidecar)	`NNModel.train`, `NNCheckpoint.save` / `to_file()`	weights, params, epoch record, and the training state that resumes the run	`NNCheckpoint.from_file` / `load`, `train(resume_from_run_id=...)`	Unpickles (`torch.load(weights_only=False)`): only files you produced.
safetensors checkpoint / Hub distribution	`to_file(format="safetensors")`, `save_pretrained` / `push_to_hub`	weights and params (no optimizer or RNG state)	`NNCheckpoint.from_file`, `NNModel.from_pretrained`	Data only; `from_pretrained` downloads from the Hub when given a repo id.
Run bundle `<dir>/bundle.json` + `g-<generation>/`	`nnx.bundles.export_bundle`	weights, params, epoch record, training state (a `"resume"` bundle) or none (`"inference"`), calibrator records	`inspect_bundle`, `validate_bundle`, `reconstruct_bundle`	Data only: safetensors plus schema-validated JSON, checked against SHA-256 sums before any tensor is read; never unpickles, imports code or downloads.

Export reads your own run. export_bundle(run_id, "bundle") opens the run's pickle checkpoint and publishes it as data. That checkpoint is the legacy trust boundary: files NNx wrote locally. From then on the bundle can travel. Before any tensor is read, validate_bundle checks every payload's size and SHA-256 and the generation id. It also checks that nothing is missing, unlisted, symlinked or outside the bundle.
Reconstruction takes caller-supplied registries. A registered module (ModelSpec) is rebuilt only from the factories you pass. Pass them as reconstruct_bundle(path, factories={(id, version): factory}); the process registry is the default. The bundle names none to import. A missing factory or component is reported before any model is allocated.
The formats stay distinguishable. NNCheckpoint.from_file refuses a bundle directory and a bundle's payload files. from_pretrained refuses a bundle directory. The bundle readers refuse a pickle checkpoint, a run directory and a Hub distribution (config.json). They never fall back to unpickling or downloading.
What does not fit is refused. These fail with a message, with no pickle fallback. They include module extra state that is not a tensor, a custom object in optimizer or component state, and a runtime-only module. So do a recorded topology transform NNx cannot replay from data (named by index and id) and an unknown bundle version.

See Concepts §22 and examples/run_bundle.py.

4. What this does NOT do¶

NNRun is not Hub-published. The Hub layout is per-model, not per-training-run. If you want 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.
A raw state dict or adapter-only export does not rebuild a recipe. export_state_dict() and save_lora_weights write tensors only. The recipe that produced the topology lives in checkpoints and config.json. Materialize the recipe before loading such a file.
Optimizer state is not in the Hub config. save_pretrained writes only the network weights; resuming optimizer state from a Hub-loaded model isn't supported. Use NNCheckpoint — or a run bundle (§3), which carries the training state as data — for warm-resume workflows.
The Hub mixin doesn't rewrite NNModel's constructor. It still takes (net_params, params) keyword args at __init__. It also takes the keyword-only module= / batch_adapter= for your own modules. The mixin is purely additive. Existing code keeps working unchanged.