nemotron-extraction 1.0.0__tar.gz

nemotron_extraction-1.0.0/CITATION.bib

@@ -0,0 +1,8 @@
+@misc{lobo_santos2026nemotron_extraction,
+  author       = {Lobo-Santos, Antonio},
+  title        = {{Nemotron Extraction}},
+  year         = {2026},
+  howpublished = {Software, version 1.0.0},
+  institution  = {Deutsches Forschungszentrum f\"ur K\"unstliche Intelligenz (DFKI)},
+  note         = {Version 1.0.0}
+}

nemotron_extraction-1.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 nemotron-extraction contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

nemotron_extraction-1.0.0/MANIFEST.in

@@ -0,0 +1,6 @@
+include LICENSE NOTICE CITATION.bib
+include uv.lock
+include report.tex report.pdf
+recursive-include benchmarks *.py
+prune outputs
+prune internal_docs

nemotron_extraction-1.0.0/NOTICE

@@ -0,0 +1,5 @@
+The worker hooks, residual reconstruction, attention projection resolution,
+and token-row assembly are adapted from the MIT-licensed nemotron-sae
+sae_pipeline extraction and hook implementation (2026 contributors).
+This package uses vLLM's worker-extension API directly; it does not vendor
+or require IBM vLLM-Hook. Model weights and datasets retain their own licenses.

nemotron_extraction-1.0.0/PKG-INFO

@@ -0,0 +1,127 @@
+Metadata-Version: 2.4
+Name: nemotron-extraction
+Version: 1.0.0
+Summary: Small, token-aligned activation hooks for Nemotron in PyTorch and vLLM
+Author: Antonio Lobo-Santos
+License-Expression: MIT
+Project-URL: Organization, https://www.dfki.de/
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+License-File: NOTICE
+Requires-Dist: torch>=2.7
+Requires-Dist: safetensors>=0.4.5
+Provides-Extra: hf
+Requires-Dist: transformers<6,>=4.57.3; extra == "hf"
+Requires-Dist: accelerate>=1.0; extra == "hf"
+Provides-Extra: vllm
+Requires-Dist: vllm<0.30,>=0.20.1; sys_platform == "linux" and extra == "vllm"
+Provides-Extra: benchmark
+Requires-Dist: datasets<6,>=3; extra == "benchmark"
+Requires-Dist: huggingface-hub<2,>=0.25; extra == "benchmark"
+Dynamic: license-file
+
+# Nemotron extraction 1.0
+
+Activations describe a model's internal state as it processes tokens.
+This package captures them from the official
+[`nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16`](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16)
+checkpoint through two backends: Hugging Face Transformers, running with PyTorch,
+and vLLM, an inference engine.
+
+A hook captures activations at a chosen location (site).
+In vLLM, hooks run inside worker processes on the graphics processor (GPU).
+The driver installs hooks through `collective_rpc` and submits token IDs.
+Workers copy captures into pinned host buffers on a separate CUDA stream and
+save Safetensors files, allowing transfers and writes to overlap inference.
+Request IDs preserve sequence order, including repeats.
+
+Hooks reconstruct residuals and clone tensors before later operations overwrite them.
+
+## Install and use
+
+Python 3.10+; vLLM requires Linux and an NVIDIA GPU.
+
+```bash
+uv venv
+uv pip install 'nemotron-extraction[hf,vllm,benchmark]==1.0.0'
+```
+
+```python
+import torch
+from nemotron_extraction import HookSite, VLLMExtractor
+
+sites = [HookSite(12, "resid_post"), HookSite(12, "attn_out_prelinear")]
+with VLLMExtractor(
+    "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", sites,
+) as extractor:
+    batch = extractor.extract(torch.tensor([[1, 2, 3]], dtype=torch.long))
+    print(batch.activations["L12_resid_post"].shape)  # (3, hidden_width)
+```
+
+Both support `await extract_async(token_ids)` and
+`stream_to_disk(token_batches, output_dir)`, yielding saved-file metadata.
+At most two batches await saving. Layers are zero-based. Sites:
+`resid_pre`, `resid_post`, `mamba_out`, `moe_out`, `attn_out`,
+`attn_out_prelinear`; mixer-specific sites require a compatible layer.
+Inputs are rectangular int64 batches without padding; outputs in host memory have shape
+`(batch × sequence_length, site_width)` in sequence-major order.
+Use separate calls for different lengths.
+
+Extraction requires eager execution, tensor parallelism of one, and disabled
+prefix caching and chunked prefill. Only prompt activations are captured.
+
+## Validation
+
+One script freezes inputs, compares both backends sequentially for English,
+German and Spanish, and saves JSON results and Markdown reports.
+uv installs the package and dependencies automatically:
+
+```bash
+uv run benchmarks/validate.py all
+```
+
+Hugging Face Transformers and vLLM produce close activations, but only 3/18 sites
+passed all agreement checks per language. Interchangeability for sparse autoencoder training
+remains unvalidated.
+
+The table summarises a batch-8 validation sample with 16,384 tokens/language from
+FineWeb-Edu (English) and FineWeb2 (German/Spanish). Cosine measures direction
+(1 means aligned). Normalised mean squared error (MSE) divides squared error by
+the reference's squared norm. Both backends share the Hugging Face Transformers
+reference root mean square as their scale.
+It retains magnitude differences; 0 means identical. Values are mean +/- population standard deviation (SD)
+across 18 equally weighted sites.
+
+| Language | Mean cosine +/- SD | Normalised MSE +/- SD |
+|---|---:|---:|
+| English | 0.999002 +/- 0.001144 | 0.001879 +/- 0.002171 |
+| German | 0.998623 +/- 0.001719 | 0.002650 +/- 0.003230 |
+| Spanish | 0.998701 +/- 0.001619 | 0.002435 +/- 0.002989 |
+
+The discrepancy remains unresolved. The default benchmark uses exactly
+1,000,000 tokens per language.
+See the [report](report.pdf) and [source](report.tex).
+
+Runs used Python 3.12, Hugging Face Transformers 4.57.6/PyTorch 2.8.0a0, and
+vLLM 0.27.1/PyTorch 2.14.0a0.
+
+## Contributing
+
+Include a reproducer and preserve hook semantics, token ordering and attribution.
+Run `uv sync --group dev`, `uv run pytest` and `uv build`.
+Compile the report with `pdflatex report.tex`.
+
+## Citation
+
+Antonio Lobo-Santos, DFKI. Download [BibTeX](CITATION.bib):
+
+```bibtex
+@misc{lobo_santos2026nemotron_extraction,
+  author = {Lobo-Santos, Antonio},
+  title = {{Nemotron Extraction}},
+  year = {2026},
+  howpublished = {Software, version 1.0.0},
+  note = {DFKI}
+}
+```

nemotron_extraction-1.0.0/README.md

@@ -0,0 +1,104 @@
+# Nemotron extraction 1.0
+
+Activations describe a model's internal state as it processes tokens.
+This package captures them from the official
+[`nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16`](https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16)
+checkpoint through two backends: Hugging Face Transformers, running with PyTorch,
+and vLLM, an inference engine.
+
+A hook captures activations at a chosen location (site).
+In vLLM, hooks run inside worker processes on the graphics processor (GPU).
+The driver installs hooks through `collective_rpc` and submits token IDs.
+Workers copy captures into pinned host buffers on a separate CUDA stream and
+save Safetensors files, allowing transfers and writes to overlap inference.
+Request IDs preserve sequence order, including repeats.
+
+Hooks reconstruct residuals and clone tensors before later operations overwrite them.
+
+## Install and use
+
+Python 3.10+; vLLM requires Linux and an NVIDIA GPU.
+
+```bash
+uv venv
+uv pip install 'nemotron-extraction[hf,vllm,benchmark]==1.0.0'
+```
+
+```python
+import torch
+from nemotron_extraction import HookSite, VLLMExtractor
+
+sites = [HookSite(12, "resid_post"), HookSite(12, "attn_out_prelinear")]
+with VLLMExtractor(
+    "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16", sites,
+) as extractor:
+    batch = extractor.extract(torch.tensor([[1, 2, 3]], dtype=torch.long))
+    print(batch.activations["L12_resid_post"].shape)  # (3, hidden_width)
+```
+
+Both support `await extract_async(token_ids)` and
+`stream_to_disk(token_batches, output_dir)`, yielding saved-file metadata.
+At most two batches await saving. Layers are zero-based. Sites:
+`resid_pre`, `resid_post`, `mamba_out`, `moe_out`, `attn_out`,
+`attn_out_prelinear`; mixer-specific sites require a compatible layer.
+Inputs are rectangular int64 batches without padding; outputs in host memory have shape
+`(batch × sequence_length, site_width)` in sequence-major order.
+Use separate calls for different lengths.
+
+Extraction requires eager execution, tensor parallelism of one, and disabled
+prefix caching and chunked prefill. Only prompt activations are captured.
+
+## Validation
+
+One script freezes inputs, compares both backends sequentially for English,
+German and Spanish, and saves JSON results and Markdown reports.
+uv installs the package and dependencies automatically:
+
+```bash
+uv run benchmarks/validate.py all
+```
+
+Hugging Face Transformers and vLLM produce close activations, but only 3/18 sites
+passed all agreement checks per language. Interchangeability for sparse autoencoder training
+remains unvalidated.
+
+The table summarises a batch-8 validation sample with 16,384 tokens/language from
+FineWeb-Edu (English) and FineWeb2 (German/Spanish). Cosine measures direction
+(1 means aligned). Normalised mean squared error (MSE) divides squared error by
+the reference's squared norm. Both backends share the Hugging Face Transformers
+reference root mean square as their scale.
+It retains magnitude differences; 0 means identical. Values are mean +/- population standard deviation (SD)
+across 18 equally weighted sites.
+
+| Language | Mean cosine +/- SD | Normalised MSE +/- SD |
+|---|---:|---:|
+| English | 0.999002 +/- 0.001144 | 0.001879 +/- 0.002171 |
+| German | 0.998623 +/- 0.001719 | 0.002650 +/- 0.003230 |
+| Spanish | 0.998701 +/- 0.001619 | 0.002435 +/- 0.002989 |
+
+The discrepancy remains unresolved. The default benchmark uses exactly
+1,000,000 tokens per language.
+See the [report](report.pdf) and [source](report.tex).
+
+Runs used Python 3.12, Hugging Face Transformers 4.57.6/PyTorch 2.8.0a0, and
+vLLM 0.27.1/PyTorch 2.14.0a0.
+
+## Contributing
+
+Include a reproducer and preserve hook semantics, token ordering and attribution.
+Run `uv sync --group dev`, `uv run pytest` and `uv build`.
+Compile the report with `pdflatex report.tex`.
+
+## Citation
+
+Antonio Lobo-Santos, DFKI. Download [BibTeX](CITATION.bib):
+
+```bibtex
+@misc{lobo_santos2026nemotron_extraction,
+  author = {Lobo-Santos, Antonio},
+  title = {{Nemotron Extraction}},
+  year = {2026},
+  howpublished = {Software, version 1.0.0},
+  note = {DFKI}
+}
+```

nemotron_extraction-1.0.0/benchmarks/aggregate.py

@@ -0,0 +1,69 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = ["nemotron-extraction==1.0.0"]
+# ///
+"""Aggregate aligned range receipts into exactly 1M validated tokens/language."""
+import argparse
+import importlib.util
+import json
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+
+spec = importlib.util.spec_from_file_location("validate", Path(__file__).with_name("validate.py"))
+validate = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(validate)
+
+
+def aggregate(paths, output, expected=1_000_000):
+    receipts = []
+    for path in paths:
+        hf = json.loads((path / "hf.json").read_text())
+        vllm = json.loads((path / "vllm.json").read_text())
+        validate.report(argparse.Namespace(output=path))
+        receipts.append((hf["start_token"], path, hf, vllm))
+    receipts.sort()
+    first = receipts[0]
+    end = 0
+    combined = {}
+    for start, path, hf, vllm in receipts:
+        if start != end:
+            raise ValueError("Ranges must be contiguous, non-overlapping and begin at zero")
+        for engine, record, baseline in (("hf", hf, first[2]), ("vllm", vllm, first[3])):
+            for key in ("language", "manifest_sha256", "batch_size", "sequence_length", "topology", "software", "script_sha256", "source", "gpu", "gpu_memory_utilization", "kv_cache_memory_bytes", "cuda_graphs_enabled"):
+                if record[key] != baseline[key]:
+                    raise ValueError(f"Ranges disagree on {engine} {key}")
+        end += hf["tokens"]
+        cosines = load_file(str(path / "cosines.safetensors"))
+        for key, metric in vllm["agreement"].items():
+            if cosines[key].numel() != hf["tokens"]:
+                raise ValueError("Cosine rows do not match counted tokens")
+            target = combined.setdefault(key, validate.Agreement())
+            target.cosines.append(cosines[key])
+            target.error2 += metric["error_squared_sum"]
+            target.reference2 += metric["reference_squared_sum"]
+            target.max_error = max(target.max_error, metric["maximum_absolute_error"])
+    if end != expected:
+        raise ValueError(f"Expected exactly {expected:,} tokens, got {end:,}")
+    output.mkdir(parents=True, exist_ok=True)
+    for engine, column in (("hf", 2), ("vllm", 3)):
+        result = dict(first[column])
+        result.update(tokens=end, start_token=0, ranges=len(receipts))
+        for key in ("model_load_seconds", "hook_install_seconds", "startup_seconds", "warmup_seconds", "inference_seconds", "capture_copy_seconds", "driver_read_seconds", "extraction_seconds", "comparison_seconds", "persistence_seconds", "wall_seconds"):
+            result[key] = sum(row[column][key] for row in receipts)
+        result["tokens_per_second"] = end / result["extraction_seconds"]
+        result["peak_gpu_memory_bytes"] = max(row[column]["peak_gpu_memory_bytes"] for row in receipts)
+        if engine == "vllm":
+            result["agreement"] = {key: value.result() for key, value in combined.items()}
+            result["passed"] = all(value["passed"] for value in result["agreement"].values())
+        validate.write_json(output / f"{engine}.json", result)
+    validate.report(argparse.Namespace(output=output))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--outputs", nargs="+", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    args = parser.parse_args()
+    aggregate(args.outputs, args.output)

nemotron_extraction-1.0.0/benchmarks/check_install.py

@@ -0,0 +1,25 @@
+"""Verify installed worker discovery and engine interfaces outside the checkout."""
+import importlib.metadata
+import inspect
+from pathlib import Path
+import nemotron_extraction
+from nemotron_extraction.worker import ActivationHookExtension
+
+print("package:", Path(nemotron_extraction.__file__).resolve())
+print("worker:", ActivationHookExtension.__module__)
+from vllm import AsyncEngineArgs
+from vllm.v1.engine.async_llm import AsyncLLM
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
+from vllm.v1.engine.input_processor import InputProcessor
+
+print("vllm:", importlib.metadata.version("vllm"))
+for name in ("collective_rpc", "generate", "shutdown", "abort"):
+    print(name, inspect.signature(getattr(AsyncLLM, name)))
+for name in ("max_num_seqs", "kv_cache_memory_bytes", "worker_extension_cls", "enforce_eager"):
+    assert name in inspect.signature(AsyncEngineArgs).parameters, name
+source = inspect.getsource(GPUModelRunner)
+for fragment in ("self.input_batch", "self.query_start_loc", "num_computed_tokens_cpu"):
+    assert fragment in source, f"Unsupported runner: missing {fragment}"
+print("request-ID assignment:")
+print(inspect.getsource(InputProcessor.assign_request_id))
+print("Installed package and worker interfaces verified")

nemotron_extraction-1.0.0/benchmarks/schedule.py

@@ -0,0 +1,52 @@
+"""Estimate tight job bounds from the three completed pilot receipts.
+
+Prints a sequential job schedule; it does not submit jobs.
+uv run --no-project benchmarks/schedule.py --pilot PATH --batch-size 4
+"""
+import argparse
+import json
+import math
+from pathlib import Path
+
+
+def schedule(pilot, batch_size):
+    jobs = []
+    records = {}
+    for language in ("en", "de", "es"):
+        path = pilot / f"{language}-0"
+        hf = json.loads((path / "hf.json").read_text())
+        vllm = json.loads((path / "vllm.json").read_text())
+        if hf["tokens"] != 16384 or vllm["tokens"] != 16384:
+            raise ValueError("Need complete 16,384-token pilot receipts")
+        if hf["batch_size"] != batch_size or vllm["batch_size"] != batch_size:
+            raise ValueError("Both pilot backends must use the selected batch size")
+        records[language] = (hf, vllm)
+    # A new allocation can have cold weights even if later pilot languages were warm.
+    startup = max(pair[0]["startup_seconds"] for pair in records.values())
+    startup += max(pair[1]["startup_seconds"] for pair in records.values())
+    for language, (hf, vllm) in records.items():
+        # Wall minus startup includes warmup, cleanup, persistence and comparison.
+        fixed = max(0, hf["wall_seconds"] - hf["startup_seconds"] - hf["extraction_seconds"] - hf["persistence_seconds"])
+        fixed += max(0, vllm["wall_seconds"] - vllm["startup_seconds"] - vllm["extraction_seconds"] - vllm["comparison_seconds"])
+        per_token = (hf["extraction_seconds"] + hf["persistence_seconds"] +
+                     vllm["extraction_seconds"] + vllm["comparison_seconds"]) / 16384
+        available = (3600 - 120) / 1.25 - startup - fixed
+        maximum = int(available / per_token / 2048) * 2048
+        if maximum < 2048:
+            raise ValueError("Model startup leaves no useful work in a one-hour job")
+        start = 0
+        while start < 1_000_000:
+            count = min(maximum, 1_000_000 - start)
+            minutes = math.ceil((1.25 * (startup + fixed + count * per_token) + 120) / 300) * 5
+            jobs.append({"language": language, "start": start, "tokens": count,
+                         "minutes": min(60, minutes), "batch_size": batch_size})
+            start += count
+    return jobs
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--pilot", type=Path, required=True)
+    parser.add_argument("--batch-size", type=int, choices=(1, 2, 4, 8, 16, 32, 64, 128), default=4)
+    args = parser.parse_args()
+    print(json.dumps(schedule(args.pilot, args.batch_size), indent=2))