Spaces:
Sleeping
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. 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.
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.
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
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.