Skip to main content
The trtllm-prune command removes weights from TensorRT-LLM checkpoints, useful for reducing checkpoint size during testing or creating lightweight checkpoint skeletons.

Overview

Pruning creates a checkpoint with empty weight tensors while preserving the checkpoint structure and metadata. This is useful for:
  • Testing checkpoint loading without full model weights
  • Reducing storage space for checkpoint distribution
  • Creating lightweight checkpoint templates

Command Syntax

Required Arguments

string
required
Path to the input TensorRT-LLM checkpoint directory containing config.json and rank*.safetensors files
string
required
Output directory for the pruned checkpoint

Optional Arguments

boolean
default:"false"
Remove all weights in the checkpoint. If not set, only prunable weights (attention QKV, projection, MLP weights) are removed.

Prunable Weights

By default, the following weight types are pruned:
  • attention.qkv.weight - Query, Key, Value projection weights
  • attention.proj.weight - Attention output projection
  • mlp.fc.weight - MLP fully connected layer
  • mlp.proj.weight - MLP projection layer
  • mlp.gate.weight - MLP gate weights (for gated activations)
Non-prunable weights (embeddings, normalization layers, etc.) are preserved unless --prune_all is specified.

Examples

Prune Specific Weights

This creates a pruned checkpoint with attention and MLP weights removed but embedding and normalization weights preserved.

Prune All Weights

This removes all weights, creating a minimal checkpoint skeleton.

Output Structure

The pruned checkpoint maintains the same structure as the input:
The config.json file is updated with "is_pruned": true to indicate the checkpoint has been pruned.

Use Cases

Create lightweight checkpoints to test model loading logic without requiring full model weights:
Distribute checkpoint structure without weights for testing or validation purposes.
Reduce storage requirements for intermediate checkpoints during development.

trtllm-build

Build TensorRT engines from checkpoints

trtllm-refit

Update engine weights from checkpoints