[JIT archive] Add a flag to not include debug files

[gau-nernst]

🚀 The feature, motivation and pitch

While comparing TorchScript vs ONNX in terms of file size, I notice ONNX file is always smaller than TorchScript. Upon inspecting the JIT archive generated by torch.jit.trace() and torch.jit.save(), I notice that there are a lot of .debug_pkl files. It seems like they are only used for debugging purpose (#23582). When I delete them manually, I could still run the model.

For small models, especially with quantization, these debug files can be quite large compared to the model weights. For one of my models, 0.9MB (out of 6.8MB) is occupied by the .debug_pkl files. For deploying to mobile devices, any amount counts.

import torch
from torch.ao.quantization import quantize_dynamic
import timm
from pathlib import Path
import shutil

model = timm.create_model("vit_tiny_patch16_224.augreg_in21k", num_classes=0).eval()
model = quantize_dynamic(model)  # quantize to make a small model
torch.jit.save(torch.jit.trace(model, torch.randn(1, 3, 224, 224)), "test.zip")

# remove .debug_pkl files
shutil.unpack_archive("test.zip")
for debug_file in Path("test").glob("code/**/*.debug_pkl"):
    debug_file.unlink()
shutil.make_archive("test_compact", "zip", base_dir="test")

# correctness checks
model1 = torch.jit.load("test.zip")
model2 = torch.jit.load("test_compact.zip")

x = torch.randn(1, 3, 224, 224)
torch.testing.assert_close(model1(x), model2(x))

The code above produces 2 files

• test.zip: 6.7MB
• test_compact.zip: 5.6MB

I'm proposing adding a flag to torch.jit.save() to not include debug files. Example usage

torch.jit.save(model_jit, "model_path.pt", include_debug_files=False)

Alternatives

Run a script to automatically remove .debug_pkl files after JIT export, like I did above.

Additional context

No response

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel