flax-duplicate-shard-shadow-poc

Transformers Flax Sharded Checkpoint Lets Shadow Shards Override Indexed Weights

Severity

Medium, 6.8/10.

Rationale: this is artifact-carried model weight and inference-output manipulation through a normal FlaxBertModel.from_pretrained() load path. It is not arbitrary code execution, and it is scoped to the deprecated-but-still-used Transformers 4.x Flax/JAX loader line, but the model artifact is self-contained and the loader returns weights that contradict the checkpoint index.

Summary

This repository contains a benign proof of concept for a Transformers 4.x Flax sharded checkpoint mismatch.

The checkpoint index says the target tensor embeddings/LayerNorm/bias is stored in aa_declared.msgpack:

"embeddings/LayerNorm/bias": "aa_declared.msgpack"

That declared shard contains the marker value -1111.0 for the target tensor.

The index maps zz_shadow.bin only for a different tensor, embeddings/LayerNorm/scale:

"embeddings/LayerNorm/scale": "zz_shadow.bin"

However, zz_shadow.bin is also a valid Flax msgpack shard and includes a duplicate copy of embeddings/LayerNorm/bias with marker value 8888.0.

When Transformers 4.x loads the model, the Flax sharded loader deserializes every tensor in every referenced shard and merges the dictionaries. The later shadow shard overrides the target tensor, so the runtime model loads 8888.0 even though the index assigns that tensor to aa_declared.msgpack.

No code execution payload is included. The PoC only demonstrates deterministic model weight and output manipulation.

Impact

• Artifact-carried Flax checkpoint weight manipulation.
• Reproduces through normal FlaxBertModel.from_pretrained(<repo>).
• The runtime-loaded target value is 8888.0, while the index-declared shard contains -1111.0.
• Inference changes: the verifier records a first-eight-output max absolute delta of about 3.46409 between the declared-shard expectation and the runtime-loaded model.
• ModelScan 0.8.8 reports no issues. It routes zz_shadow.bin through the PyTorch scanner and skips it as invalid magic, while Transformers consumes it as Flax msgpack because the shard is referenced by the Flax index.

Affected Versions Tested

• Python 3.12.3
• transformers==4.57.6
• huggingface_hub==0.36.2
• flax==0.12.7
• jax==0.10.0
• jaxlib==0.10.0
• modelscan==0.8.8
• picklescan==1.0.4

Transformers v5 has removed the JAX/Flax backend. This report is therefore scoped to the latest tested 4.x Flax-capable loader, not the v5 API line.

Files

config.json
  Tiny Flax BERT config.

flax_model.msgpack.index.json
  Sharded Flax checkpoint index. It assigns the target tensor to
  aa_declared.msgpack and assigns only a different tensor to zz_shadow.bin.

aa_declared.msgpack
  Declared Flax msgpack shard. Contains embeddings/LayerNorm/bias == -1111.0.

zz_shadow.bin
  Shadow Flax msgpack shard. It is mapped for embeddings/LayerNorm/scale but
  also contains a duplicate embeddings/LayerNorm/bias == 8888.0.

verify_flax_duplicate_shard_shadow_poc.py
  End-to-end verifier. Loads this repository with Transformers, checks the
  shard contents, compares inference output, and records scanner behavior.

scripts/build_and_verify_duplicate_shadow.py
  Original builder/verifier used to regenerate the local PoC artifact.

evidence/lab_verify.json
  Original lab verification output.

evidence/fresh_verify.json
  Fresh verification output from this staged repository.

evidence/hub_download_verify.json
  Verification output after uploading to Hugging Face, downloading the public
  repository with `hf download`, and running the verifier against that
  downloaded copy.

evidence/sha256.txt
  Hashes for uploaded files.

Reproduction

Install matching dependencies in an isolated environment:

pip install "transformers==4.57.6" "huggingface_hub==0.36.2" "flax==0.12.7" "jax==0.10.0" "jaxlib==0.10.0" "modelscan==0.8.8" "picklescan==1.0.4"

Then run:

python verify_flax_duplicate_shard_shadow_poc.py

Expected output highlights:

{
  "target_key": "embeddings/LayerNorm/bias",
  "weight_map_target_value": "aa_declared.msgpack",
  "declared_shard_target_value": -1111.0,
  "weight_map_dummy_value": "zz_shadow.bin",
  "shadow_shard_duplicate_target_value": 8888.0,
  "runtime_loaded_target_value": 8888.0,
  "first_8_max_abs_delta_declared_vs_runtime": 3.464093565940857,
  "verified": true
}

The critical observation is that runtime_loaded_target_value comes from zz_shadow.bin, despite the index assigning the target tensor to aa_declared.msgpack.

The uploaded Hugging Face repository was also downloaded into a fresh local directory and verified with the same script. That output is captured in evidence/hub_download_verify.json.

Scanner Behavior

ModelScan 0.8.8 reports no issues in this environment. The captured output shows that zz_shadow.bin is treated as a PyTorch candidate and skipped:

Scanning .../zz_shadow.bin using modelscan.scanners.PyTorchUnsafeOpScan model scan
No issues found
...
zz_shadow.bin was skipped during a pytorch scan:
Invalid magic number

PickleScan 1.0.4 also reports no infected files or dangerous globals:

Scanned files: 1
Infected files: 0
Dangerous globals: 0

This scanner output is supporting context. The core issue is the runtime loader trusting the full contents of every referenced shard instead of enforcing the authoritative key-to-shard mapping from flax_model.msgpack.index.json.

Root Cause

The Flax sharded checkpoint index is used to decide which shard files should be opened, but the loader does not restrict each shard to the tensor names assigned to that shard by weight_map. Instead, every tensor present in every referenced shard is deserialized and merged into a single state dictionary. If a later referenced shard contains a duplicate tensor key, it can override the value from the shard named by the index for that key.

In this PoC:

weight_map target assignment:
  embeddings/LayerNorm/bias -> aa_declared.msgpack

declared shard value:
  embeddings/LayerNorm/bias = -1111.0

shadow shard is mapped for:
  embeddings/LayerNorm/scale -> zz_shadow.bin

shadow shard also includes:
  embeddings/LayerNorm/bias = 8888.0

runtime loaded value:
  embeddings/LayerNorm/bias = 8888.0

Safety Notes

• No arbitrary code execution is used.
• No network callbacks, credential access, persistence, or destructive payloads are present.
• Marker values are numeric tensor values chosen only to make the loader mismatch obvious.

Limitations

• This is output/weight manipulation, not ACE.
• The validated affected loader is Transformers 4.x Flax/JAX support. The Transformers v5 API line removed the Flax backend.
• The finding is strongest for workflows that still load Flax msgpack sharded checkpoints with Transformers 4.x.

Duplicate Checks

Local research did not find an obvious published advisory for this specific Flax sharded-index key enforcement issue. This is distinct from generic unsafe pickle loading because both shards are Flax msgpack artifacts and the behavior is triggered by the model's flax_model.msgpack.index.json metadata.