feat/fix(ara): Save-time model tensor integrity check

[hpnyaggerman]

Note: turns out it was not as straightforward as I thought it would be.

The PR made in response to #211 (comment). It is probably overengineered but my aims here were:

1. To make sure that what works already mostly keeps working.
2. To make sure that things that break, break loudly, and with exceptions only when the reason for what appears to be a breakage is well-understood and is explicitly considered benign.

This is obviously just the initial version of the feature/fix, and I am willing to explain any unclear parts of the design and to correct any parts which have problems.

What the check does is compare the tensor names and shapes the save actually wrote against the ones in the source checkpoint, reading safetensors headers only, so it costs no measurable time and never loads the artifact back. The invariant it rests on is that Heretic rewrites values and never renames anything. Where a correct save legitimately does differ from its source, the difference is not forgiven by loosening the comparison but by an explicit entry that names the mechanism responsible and claims the specific tensor names that mechanism accounts for, so a name is either explained by exactly one entry or it is unexplained, and unexplained means the save is rejected.

It runs at three points: once before optimization begins, as a trial save on the freshly loaded model so that an environment which cannot serialize the model is caught before hours are spent rather than after, after a local save, and before an upload, on a staged copy, so that a rejected model cannot be published by choosing the upload action after the save action has already reported the problem. Only the first is skippable, with --no-preflight-save-check, since the other two are the ones that actually protect anything.

I have tested my design with the following models and got the following results.

Columns are transformers versions. Ten releases were measured, 5.3.0, 5.4.0, 5.5.0, 5.5.1, 5.5.2, 5.5.3, 5.5.4, 5.6.0, 5.8.0 and 5.14.1, and every model behaved identically within the bands shown, so the table is banded rather than repeating ten identical columns. Pass means every tensor in the source checkpoint was written under the same name at the same shape. Fail means at least one name or shape could not be accounted for, and the run stops. Warn means everything expected was written correctly and the artifact additionally carries names the source did not, which (in my testing) loads and computes correctly, so the run continues. The counts are of unexplained differences only, after every applicable exception has been applied, so a non-zero count is one that nothing could account for.

Model	Tensors in source	5.3.0	5.4.0 - 5.5.4	5.6.0 - 5.14.1
Qwen/Qwen2.5-0.5B-Instruct	290	pass	pass	pass
tiny-random/gpt-oss	36	pass	pass	pass
tiny-random/qwen3-moe	46	pass	pass	pass
tiny-random/glm-4-moe	148	pass	pass	pass
tiny-random/glm-4-moe-lite	433	pass	pass	pass
tiny-random/deepseek-v3.1	241	pass	pass	pass
hf-internal-testing/tiny-random-Qwen3VLMoeForConditionalGeneration	376	pass	pass	pass
tiny-random/gemma-3	67	pass	fail: 67 absent, 67 extra	pass
hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration	58	pass	fail: 31 absent, 31 extra	pass
OpenGVLab/InternVL3-1B-hf	733	pass	fail: 732 absent, 732 extra	pass
hf-internal-testing/tiny-random-MixtralForCausalLM	21	fail: 6 absent, 26 extra	fail: 6 absent, 26 extra	pass
PrimeIntellect/GLM-0.5B	429	warn: 0 absent, 552 extra	warn: 0 absent, 552 extra	pass

Two of those rows need a word before they are read as the check misbehaving, and the Mixtral one especially, since it is the only row that fails on the version uv.lock currently pins. hf-internal-testing/tiny-random-MixtralForCausalLM is a test fixture stored in the post-conversion naming rather than the on-disk naming everything else uses, so a correct save legitimately un-fuses its 21 tensors into 41, which is where the 6 absent and 26 extra come from. That rejection is, as far as I can tell, a false positive, and covering this class of mismatch under the whitelist is something I plan to do later down the line.

And gpt-oss specifically, checked against the base snapshot's 459 tensors:

Artifact	Verdict	Absent	Extra	Shape
saved by 5.3.0, quantized	pass	0	0	0
saved by 5.14.1, quantized	pass	0	0	0
saved by 5.4.0, quantized	fail	96	96	0
saved by 5.5.1, dequantize	fail	96	48	48
saved by 5.5.2, dequantize	fail	96	2	0
saved by 5.14.1, dequantize	fail	96	2	0
saved by 5.3.0, dequantize	fail	none written	none written	none written
saved by 5.4.0, dequantize	fail	none written	none written	none written
p-e-w/gpt-oss-20b-heretic-ara-v4	fail	96	96	0

The absent and extra columns are names: absent counts names the source has that the artifact does not, and extra counts names the artifact has that the source does not. The shape counts tensors that appear under the same name in both the source and the artifact, with different dimensions, which means a loader finds exactly the name it expects and gets the wrong geometry behind it. The forty-eight on the 5.5.1 dequantize row are the expert blocks going from a four-dimensional (32, 2880, 90, 16) in the source to a three-dimensional (32, 2880, 2880) in the artifact, and the gate-up half from (32, 5760, 90, 16) to (32, 2880, 5760). Shape is evaluated before names and always fails, since there is no reading under which a correctly named tensor with the wrong geometry is benign.

The two none written rows are the case where save_pretrained raises part way through, leaving config.json and generation_config.json and no weights at all, so there is nothing to count and nothing to compare. Those are the rows that only the trial save before optimization can catch, since a save that does not return never reaches a post-save check.

Interestingly enough, tiny-random/gemma-3, hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration and OpenGVLab/InternVL3-1B-hf fail to pass for a reason which is neither of the defects identified by me in #211 (comment), which means that there are likely more bugs in transformers than previously believed, and the two I have previously identified might be only the tip of the iceberg.

[hpnyaggerman]

@claude review?

[hpnyaggerman]

Guess not. Was worth a try. In truth, I have no idea how this thing works. @p-e-w, I am awaiting your criticism.

Once more, worth remembering that this is essentially a draft which has the rough shape of what I believe the right way of handling this to be and it can undergo as many cuts and revisions as it needs.

[p-e-w]

Thanks for the PR, and the super detailed investigation!

I must admit I'm a little shocked by the size of this change. I had expected 20-50 lines of code, not 1000+. It seems like you've implemented this with manual binary reads. Isn't this precisely what the safetensors library is for? Surely getting the names and shapes of all tensors doesn't require that much code?

I don't think the reason why something broke is of any real interest. The "reason" is always that Transformers/safetensors has a bug, as such discrepancies from a simple load/save cycle are never acceptable.

[hpnyaggerman]

@p-e-w, appreciate your feedback. I understand quite well your issue with the size and I myself was and remain quite concerned with it. But first of all, the procedure you are referring to is one of the smallest parts of the broader check procedure and using safetensors for it would, therefore, not reduce the amount of lines by much. There is something to be noted here, though. I was trying to minimize the amount of dependencies introduced, and looking at safe_open(filename, framework, device), I only now realize that framework did not have to be torch but could in fact be numpy, so the argument of minimizing dependencies does not really work here, so my bad. Though this method does not expose byte offsets. But that aside, I also looked and found that huggingface_hub also hand-rolls its own version of the same conceptual operation and even exposes it publicly through parse_local_safetensors_file_metadata and get_local_safetensors_metadata, both on the version heretic pins, and it does expose byte offsets. So either of those options (safetensors or hub) should work, though it would not reduce the size by that much.

This, of course, raises a reasonable question, where does the size come from? And this is where I have to bring some nuance into your second point, because most of the tensor check code is actually about being nuanced and precise, and not rejecting cases where tensors differ between the loaded and the saved model in a manner which is legitimate and is not attributable to a bug. Some cases, all of them saves that are bit-for-bit lossless (line numbers below hold for transformers 5.3.0):

• A checkpoint carrying legacy per-layer rotary_emb.inv_freq keys. Transformers does not write them.
• Keys the model never registers, and parameters the checkpoint lacks that transformers initialises itself. These are reported, in unexpected_keys and missing_keys, but they are model-namespace names while both sides of my comparison are checkpoint-namespace, so they have to be matched across the two before they can subtract anything. A reported name that appears in neither the source nor the artifact cannot be an expectation about what gets written at all.
• MoE checkpoints carrying per-expert tensors for a module the model does not build. The fusing converter collapses all of them into one reported name, which is itself in neither the source nor the artifact, so the individual source names have to be attributed to it by module path. On a byte-correct save that is 48 legitimately unwritten source names for tiny-random/glm-4-moe, 96 for deepseek-v3.1 and 192 for glm-4-moe-lite.
• qwen3_vl_moe declares converters whose source and target patterns are identical and which carry a Transpose (conversion_mapping.py:151-162), so on save the reverse op compares the tensor against the parameter it already is, finds it matching, and declines to transpose back. All 376 names round-trip, two expert tensors written (4, 128, 64) against a source (4, 64, 128), reload bit-identical. A lossless save with a permuted shape in it.

I handle each of these by naming the mechanism and having it claim the specific names that mechanism accounts for, rather than by loosening the comparison anywhere. A name is claimed by exactly one mechanism or it is unexplained, and unexplained is a rejection. The length comes from seven mechanisms of this whitelist, each needing its evidence gathered from the loaded model, several of them only correct under a precondition that has to be checked per name.

And that's roughly 50% of the file. And the rest of the code in the file is the check itself, and the auxiliary methods such as failure and warning message composition, etc. And that (the check file) is about 50% of the size of the whole PR. Much of the rest of the size comes from my decision to introduce an optional pre-flight check which evaluates whether a model can even be saved properly after loading before the optimization runs, which I thought makes sense because there are bugs in transformers I have found (and reported in my previous report) which break the model on save in ways which are not recoverable. I judged that given how broken transformers can be, letting people waste their time on abliterating a model they will not be able to save would be quite cruel, especially given how some (including myself) pay for their compute time when abliterating models. And then there was also the decision to restructure the upload path because I judged that it would be a mistake to let potentially broken models be uploaded to the hub, and preventing that outcome while making the error traceable and diagnosis easy required making a number of changes, which can be seen in the diff.

The PR is still bloated though, and that is a consequence of me rushing to get a working prototype out because I was testing my change on a cloud GPU VPS which burned through my pockets at a rate which did not leave me with more than 48 hours of time to get everything working. Now that I know that the change works, I can start optimizing the code. But before I do that, I need your judgment on the decisions I've made because on your judgment about which aspects of this PR, if any, you want trimmed, depends how I will go about optimizing it.

[p-e-w]

I actually consider all of the cases you describe to be Transformers bugs. Like with any other software, I expect loading and then saving again to produce the same file, or at least a file with the same content. A program like Heretic should be able to reproduce the model exactly, with the exception of the content of the tensors it modifies. It's very disappointing that Transformers does not guarantee this in so many cases, and I don't think any of them are acceptable.

I also don't think Heretic should contain logic for trying to guess which of those unacceptable cases has occurred. It should compare the structure of the model files, and warn the user about any differences. It's up to the user then to figure out whether those differences are relevant to them or not. Logic for identifying which specific version of a bug has happened in another library is unmaintainable.

[hpnyaggerman]

@p-e-w Fair enough. My initial thought was also that doing it the simple way would be the way to go until I realized that there exist cases in which the differences do not technically break the model. If you are fine with not discriminating based on whether the discrepancy is benign or not, then I will simplify the code and leave only the essential logic.

[p-e-w]

until I realized that there exist cases in which the differences do not technically break the model.

The problem is that we can't decide that, because not everyone uses the model only for inference. Every observable feature of an object will be depended on by someone. Any structural difference is a potential problem, even if inference (with Transformers) still appears to work the same way.

IMO, a simple structural check is what we want, and any structural difference should be flagged so the user can make an informed decision.

[teleprint-me]

Just my 2 cents since I have limited free time because Im in term right now and finals are coming up.

@p-e-w is correct in that this is a huggingface library issue. This should be reported to them, though I wouldnt expect it to be fixed anytime soon unless you work on it directly yourself and have them review the changes. The libraries are complex due to obtuse abstractions and delegated assumptions in the pipeline.

As you stated, a simple approach is desirable here (always is). Having experienced issues with transformers and other hf libraries, I think a simple diff as a sanity check with sane defaults is the right way to go. You already isolated the issue, so you can delegate what the differences are based on the outliers you observed.

If the checkpoint does not align with the expected output weights and heretic and ara seem fine, then report the difference and potential area of interest and exit the optimization run. You could do this in a few hundred lines or less using pure torch (which is what I did with my inital local draft for investigating ara v4).

For example, my original proof of concept was less than 100 LoC. Anything over 300 LoC is a sign that it has become too complex and is doing too much as it has become out of scope for heretic (considering the projects automation goal).

Hopefully this is helpful. Great detective work btw. Issues like this are not easy to find, isolate, and are challenging to patch.