Blame for README.md

---
title: Tensor Pruning API
emoji: ✂️
colorFrom: blue
colorTo: purple
sdk: docker
app_port: 7860
---

# Tensor Pruning API

A Docker Space for [the Tensor Pruning API](https://github.com/AngshulMajumdar/Tensor-Pruning-API). It selects query-weight masks for dense Llama-layout checkpoints with a single retention count across all layers. The CPU selector uses the nested weighted \((2,4,8)\) residual objective from the API. A job publishes one Boolean NumPy mask per query projection, with axes `[head, channel, input]`, plus the tensor names and retained counts in its status response.

## Run on Hugging Face

Unzip the release and run the following from the directory containing `publish.py`. Log in through the Hugging Face CLI using an account with permission to create Spaces. Omit the model flag if you will attach the model later.

```sh
python -m pip install -U huggingface_hub
hf auth login
python publish.py --repo Angshul/Tensor-Pruning-API --model-repo OWNER/MODEL
```

The helper creates a private Docker Space with the chosen model mounted read-only at `/models`. If the repository already exists, review its contents before uploading the deployment files with the Hugging Face UI or CLI. Use a model that your account can access. If `config.json` and the weight files are in a mounted subdirectory, send that relative directory as `checkpoint`.

In Space Settings, set a secret named `TENSOR_PRUNING_API_TOKEN` to an operator-chosen random token. Configure `PRUNING_THREADS` to match the assigned CPUs if necessary. Keep `PRUNING_WORK_ROOT` on local storage that supports memory mapping and file preallocation. The default is `/home/user/work`. A model mount supplies the checkpoint. The job's query staging files, index scratch, and masks occupy local disk.

## Submit and retrieve a job

Open the Space's direct app URL from its Hugging Face page. The root opens the API documentation at `/docs`. The health route is `/health`. The example uses a private Space, so it sends a Hugging Face access token for the proxy and a separate pruning token for the app.

```sh
curl -fsS "$APP_URL/v1/checkpoint-jobs" \
  -H "Authorization: Bearer $HF_TOKEN" \
  -H "X-Tensor-Pruning-Token: $PRUNING_TOKEN" \
  -H 'Content-Type: application/json' \
  -d '{"checkpoint":".","keep":1610612736,"seed":0}'

curl -fsS "$APP_URL/v1/jobs/$JOB_ID" \
  -H "Authorization: Bearer $HF_TOKEN" \
  -H "X-Tensor-Pruning-Token: $PRUNING_TOKEN"

curl -fsS "$APP_URL/v1/jobs/$JOB_ID/masks/0" \
  -H "Authorization: Bearer $HF_TOKEN" \
  -H "X-Tensor-Pruning-Token: $PRUNING_TOKEN" \
  -o mask-0000.npy
```

The example retention count is 30% of the Llama 3.1 70B query slots. The service runs one job at a time and bounds the waiting queue. Status and mask routes also serve the original staged NumPy-manifest API at `POST /v1/jobs`. Download completed masks before a Space restart, which can clear local storage. Clear obsolete job directories when the service is stopped to recover disk space. Interrupted jobs can be resubmitted.

The checkpoint endpoint accepts the standard `model.safetensors` file or `model.safetensors.index.json` and its shards. It reads one `model.layers.N.self_attn.q_proj.weight` tensor per layer. It supports dense FP16, BF16, FP32, and FP64 weights. The worker leaves model weights unchanged. Mask application is a downstream step.

The default objective uses unit input moments. To use activation weighting, supply `moments_manifest` as a relative path to a JSON object mapping each eligible tensor name to a corresponding NumPy vector file. Stage the JSON and vectors under `PRUNING_INPUT_ROOT`, which defaults to `/home/user/work/inputs`, or mount a separate read-only volume there. Dense-model calibration is performed externally.

## Size a Llama 3.1 70B job

The standard layout has 80 layers and \(8192\times8192\) query matrices, yielding 5,368,709,120 eligible slots. For FP16 or BF16 and an interior retention budget:

| Local file use | Bytes per query slot | Approximate decimal GB |
| --- | ---: | ---: |
| Staged query weights | 2 | 10.74 |
| Channel-local partition indices | 4 | 21.47 |
| Output Boolean masks | 1 | 5.37 |
| Total | 7 | 37.58 |

The API checks available disk before staging and requires roughly 38.66 GB free for this job, including its reserve. Older output masks, the image, and other local files consume additional disk space. The model repository must be mounted separately for the sizing above. The worker reads only the query byte ranges from the mounted shards and memory-maps local working files. The figures are disk usage estimates. RAM and elapsed time depend on local I/O, page caching, and CPU allocation. The 70B layout was checked with synthetic metadata. A full 70B selection run remains to be measured.

## Run locally

```sh
docker build -t tensor-pruning-space .
docker run --rm -p 7860:7860 \
  -v "$CHECKPOINT_DIR:/models:ro" \
  -e TENSOR_PRUNING_API_TOKEN \
  -e PRUNING_THREADS=4 \
  tensor-pruning-space
```

Use `http://localhost:7860/docs` with the pruning token. The local container needs sufficient free disk or a writable local bind mount for the work root. The Python integration tests can be run with `python -m pip install '.[space,test]'` followed by `python -m unittest discover -s tests -v`.