feat: resumable model-free PTQ with meta-tensor skip - #2993
Conversation
Add resume/skip pathway to process_file(): when the output shard already exists, load its tensor metadata via safe_open on meta device instead of re-computing quantization. This enables interrupted runs to resume from where they left off. Also adds os.makedirs in process_file_microscale_scheme() to ensure the output directory exists before saving, and hardcodes format to float-quantized in save_utils for correct config generation. Companion: vllm-project/compressed-tensors#810 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
|
Important Review skippedAuto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review. Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed. |
|
The quality checks have failed. Please run |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces optimization and safety improvements to the model-free entrypoint, including skipping processing if the output file already exists and ensuring the output directory is created before saving. Feedback on these changes highlights a critical issue where using safe_open with device="meta" and calling get_tensor will cause a runtime error, suggesting instead to parse the safetensors header directly. Additionally, hardcoding the quantization format to "float-quantized" is flagged as a breaking change for other formats, with a recommendation to use a fallback mechanism instead.
| with safe_open(save_path, framework="pt", device="meta") as f: | ||
| for key in f.keys(): | ||
| tensor_meta = f.get_tensor(key) | ||
| tensors[key] = tensor_meta |
There was a problem hiding this comment.
Using safe_open with device="meta" and calling f.get_tensor(key) will raise a RuntimeError at runtime (e.g., RuntimeError: Cannot copy to/from a Meta tensor) because meta device tensors do not have storage and cannot have data copied into them. To load the tensors as meta tensors without loading their data into memory, we can parse the safetensors file header directly using standard file I/O and construct the meta tensors using str_to_torch_dtype.
import json
import struct
tensors = {}
with open(save_path, "rb") as f:
header_size_bytes = f.read(8)
if len(header_size_bytes) == 8:
header_size = struct.unpack("<Q", header_size_bytes)[0]
header_bytes = f.read(header_size)
header = json.loads(header_bytes.decode("utf-8"))
for key, metadata in header.items():
if key != "__metadata__":
shape = metadata["shape"]
dtype = str_to_torch_dtype(metadata["data_type"])
tensors[key] = torch.empty(shape, dtype=dtype, device="meta")| "ignore": ignore, | ||
| "quantization_status": QuantizationStatus.COMPRESSED, | ||
| "format": scheme.format, | ||
| "format": "float-quantized", |
There was a problem hiding this comment.
Hardcoding the format to "float-quantized" will break configuration generation for other quantization formats, such as integer quantization (e.g., "int-quantized"). If the goal is to provide a fallback or handle cases where scheme.format is not set, we should use a fallback instead of unconditionally overriding it.
| "format": "float-quantized", | |
| "format": scheme.format or "float-quantized", |
Merge Protections🔴 1 of 1 protections blocking · waiting on 👀 reviews
🔴 Require one maintainer reviewWaiting for any of
This rule is failing.All PRs must have at least one approving review from a maintainer before merging.
|
2 participants
Summary
process_file(): when the output shard already exists, loads tensor metadata viasafe_openon meta device instead of re-computing quantizationos.makedirsinprocess_file_microscale_scheme()to ensure output directory exists before savingfloat-quantizedinsave_utilsfor correct config generationCompanion PR: compressed-tensors#810 adds
_skip_meta_devicedecorator toNVFP4PackedCompressor.compress()for meta-device tensor handlingTest plan
config.jsonhas correct format field🤖 Generated with Claude Code