Restore dropped FP8 weight_scale_inv tensors on load - #6978
danielhanchen merged 9 commits into
Conversation
Some block-scale FP8 checkpoints (for example Qwen3.6-27B-FP8, issue unslothai#6200) load with transformers leaving an mlp.gate_proj as a plain bf16 Linear instead of an fp8 module. Its raw quantized values are read into the bf16 weight and the weight_scale_inv is dropped as an unexpected key, so the weight is used un-scaled and the base model is garbage (perplexity around 2 million). After load, for every checkpoint weight_scale_inv whose live weight is not fp8, dequantize the orphaned weight in place using the block scale from the checkpoint index. Modules that were converted correctly keep an fp8 weight and are skipped, so healthy checkpoints and single-file checkpoints are a no-op. Verified on Qwen3.6-27B-FP8: 64 gate_proj scales restored, perplexity 2028902 to 8.9. No-op on Qwen3-8B-FP8 (all scales already live).
for more information, see https://pre-commit.ci
![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 a mechanism to restore dropped block-FP8 weight_scale_inv tensors on load, addressing an issue where some block-scale FP8 checkpoints leave Linear layers unconverted, resulting in un-scaled weights. The changes include adding utility functions in loader_utils.py to parse the FP8 block size, load the weight map, resolve shards, and dequantize orphaned weights in place, as well as integrating this recovery step into the vision model loader. Additionally, a comprehensive test suite is added to verify the restoration logic. Feedback on the changes suggests making the parsing of weight_block_size more robust to handle cases where it is specified as a scalar integer rather than a list or tuple (preventing a potential TypeError), and using rank-aware reshaping when broadcasting scale tensors over weight tensors during dequantization to defensively support weights of any rank.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if not weight_map: | ||
| return (0, 0) | ||
|
|
||
| scale_keys = {k: v for k, v in weight_map.items() if k.endswith(".weight_scale_inv")} | ||
| if not scale_keys: | ||
| return (0, 0) | ||
|
|
||
| bs0, bs1 = block | ||
| restored = 0 | ||
| skipped = 0 | ||
| shard_cache = {} | ||
| for scale_key, shard in scale_keys.items(): | ||
| module_name = scale_key[: -len(".weight_scale_inv")] | ||
| try: | ||
| module = model.get_submodule(module_name) | ||
| except AttributeError: | ||
| continue | ||
| weight = getattr(module, "weight", None) | ||
| if not isinstance(weight, torch.Tensor) or weight.device.type == "meta": | ||
| continue | ||
| if weight.dtype in _FP8_DTYPES: | ||
| # Correctly converted fp8 module: the scale is handled by the fp8 path already. | ||
| skipped += 1 | ||
| continue | ||
|
|
||
| if shard not in shard_cache: |
There was a problem hiding this comment.
When broadcasting scale tensors over weight tensors during dequantization, we should use rank-aware reshaping (e.g., scale.view(-1, *([1] * (weight.ndim - 1))) in PyTorch) to defensively support tensors of any rank (such as 4D convolutional weights) even if current checkpoints are expected to be 2D, rather than assuming or enforcing that the weight is 2D.
for scale_key, shard in scale_keys.items():
try:
module_name = scale_key[: -len(".weight_scale_inv")]
try:
module = model.get_submodule(module_name)
except AttributeError:
continue
weight = getattr(module, "weight", None)
if not isinstance(weight, torch.Tensor) or weight.device.type == "meta":
continue
if weight.dtype in _FP8_DTYPES:
# Correctly converted fp8 module: the scale is handled by the fp8 path already.
skipped += 1
continue
if shard not in shard_cache:
from safetensors import safe_open
shard_path = _resolve_fp8_shard(model_name, shard, local_files_only, token)
shard_cache[shard] = safe_open(shard_path, framework = "pt")
scale = shard_cache[shard].get_tensor(scale_key).to(torch.float32)
scale_reshaped = scale.view(-1, *([1] * (weight.ndim - 1)))
with torch.no_grad():
module.weight.data = (weight.to(torch.float32) * scale_reshaped).to(weight.dtype)
restored += 1
except Exception:
continueReferences
- When broadcasting scale tensors over weight tensors during dequantization, use rank-aware reshaping (e.g., scale.view(-1, *([1] * (weight.ndim - 1))) in PyTorch) to defensively support tensors of any rank (such as 4D convolutional weights) even if current checkpoints are expected to be 2D.
There was a problem hiding this comment.
Block FP8 weight_scale_inv only exists for 2D Linear weights, and the restore loop now explicitly skips anything that is not 2D (weight.ndim != 2). There are no 4D convolution weights with block scales in these checkpoints, so rank aware reshaping is not needed here and the 2D block grid expansion is the correct dequantization. Keeping the shape handling 2D also lets it validate the block grid and detect a transposed layout, which a generic reshape could not.
| return None | ||
| block = quant.get("weight_block_size") | ||
| if not block: | ||
| return None | ||
| if len(block) == 1: | ||
| block = [block[0], block[0]] |
There was a problem hiding this comment.
If weight_block_size is specified as a single integer (e.g., 128) instead of a list or tuple, calling len(block) will raise a TypeError. To make this more robust and prevent unexpected failures, we should check the type of block and handle both scalar integers and lists/tuples.
| return None | |
| block = quant.get("weight_block_size") | |
| if not block: | |
| return None | |
| if len(block) == 1: | |
| block = [block[0], block[0]] | |
| block = quant.get("weight_block_size") | |
| if not block: | |
| return None | |
| if isinstance(block, (int, float)): | |
| block = [block, block] | |
| elif isinstance(block, (list, tuple)): | |
| if len(block) == 1: | |
| block = [block[0], block[0]] | |
| else: | |
| return None | |
| return [int(block[0]), int(block[1])] |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0fd71835ac
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| _restore_dropped_fp8_scales( | ||
| model, model_name, | ||
| local_files_only = local_files_only, token = token, | ||
| ) |
There was a problem hiding this comment.
When text_only=True for VLMs, this loader already adds a key_mapping that strips prefixes like model.language_model. before from_pretrained builds the text decoder, but the restore pass is called with only model_name. _restore_dropped_fp8_scales then resolves the original index keys with get_submodule, so keys such as model.language_model.layers.*.mlp.gate_proj.weight_scale_inv no longer match the text-only module names and the dropped FP8 scales are silently skipped, leaving the text-only Qwen3.6-FP8 load unscaled. Pass/apply the same key mapping when deriving module_name.
Useful? React with 👍 / 👎.
| index_path = hf_hub_download( | ||
| model_name, index_file, | ||
| local_files_only = local_files_only, token = token, | ||
| ) |
There was a problem hiding this comment.
For remote models this download ignores the hub arguments used by the actual from_pretrained call, notably revision, subfolder, and cache_dir. If a user loads a pinned revision or a checkpoint stored under a subfolder, the model weights come from that location but the restore pass reads the index/scales from the default repo root, which can either no-op or multiply weights by scales from a different checkpoint. Thread the same hub kwargs through both the index and shard downloads.
Useful? React with 👍 / 👎.
| scale_expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(bs1, dim = 1) | ||
| scale_expanded = scale_expanded[:out_features, :in_features].to(weight.device) | ||
| with torch.no_grad(): | ||
| module.weight.data = (weight.to(torch.float32) * scale_expanded).to(weight.dtype) |
There was a problem hiding this comment.
When load_in_16bit=True is used with an FP8 checkpoint, FastBaseModel.from_pretrained asks the HF FP8 quantizer to dequantize=True, so live Linear weights are already real bf16/fp16 values rather than raw FP8 codes. This restore path treats every non-FP8 dtype with a weight_scale_inv key as dropped and multiplies it by the scale again, corrupting 16-bit loads of FP8 models. Skip the restore when the load path intentionally dequantized the checkpoint.
Useful? React with 👍 / 👎.
| except Exception: | ||
| return (0, 0) |
There was a problem hiding this comment.
If any later tensor raises after earlier layers have already been scaled, for example from a malformed scale shape or a CUDA OOM during the temporary fp32 multiply, this outer except reports (0, 0) and leaves those earlier in-place mutations in the model. That is not a safe no-op: the checkpoint can end up partially dequantized while the caller continues. Validate before mutating, handle errors per tensor, or roll back on failure.
Useful? React with 👍 / 👎.
| out_features, in_features = weight.shape | ||
| scale_expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(bs1, dim = 1) | ||
| scale_expanded = scale_expanded[:out_features, :in_features].to(weight.device) |
There was a problem hiding this comment.
For block-FP8 checkpoints whose weight_scale_inv is stored in the transposed block layout, the normal FP8 forward path detects this by comparing the block counts and transposes the scale before use, but this restore path always treats scale.shape[0] as output blocks. Those checkpoints either no-op via the outer exception when the expanded shape does not broadcast, or silently apply the wrong per-block scales when the block grid is square, leaving restored weights corrupted. Add the same shape check/transpose before expanding the scale.
Useful? React with 👍 / 👎.
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 23c998de8b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| fast_inference = fast_inference, | ||
| ) | ||
| # Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200). | ||
| _restore_dropped_fp8_scales( |
There was a problem hiding this comment.
This call only runs in FastBaseModel.from_pretrained's non-vLLM path, but the supported text model path dispatches through FastLanguageModel to FastLlamaModel.from_pretrained (unsloth/models/loader.py:838, then the CausalLM load in unsloth/models/llama.py) and never reaches this restore. For text FP8 checkpoints with dropped weight_scale_inv tensors, the newly added helper therefore is not invoked and those weights remain unscaled, so the reported FP8 load corruption still occurs outside the vision/FastModel fallback path.
Useful? React with 👍 / 👎.
| index_path = hf_hub_download( | ||
| model_name, | ||
| index_file, | ||
| local_files_only = local_files_only, | ||
| token = token, | ||
| ) |
There was a problem hiding this comment.
The restore path downloads model.safetensors.index.json with only model_name, local_files_only, and token, while the actual from_pretrained call just above can be loading a different snapshot or location via kwargs such as revision, subfolder, or cache_dir. In those cases the model weights come from the requested snapshot but the scale map/shards are fetched from the default root/main repo (or missed entirely with local_files_only), which can silently skip the fix or apply mismatched FP8 scales.
Useful? React with 👍 / 👎.
| if weight.dtype in _FP8_DTYPES: | ||
| # Correctly converted fp8 module: the scale is handled by the fp8 path already. | ||
| skipped += 1 | ||
| continue |
There was a problem hiding this comment.
When load_in_16bit=True is used for a pre-quantized FP8 checkpoint, the vision loader asks the HF FP8 quantizer to dequantize before from_pretrained, so the loaded bf16/fp16 weights have already had weight_scale_inv folded in. This check treats every non-float8 weight as an orphaned raw FP8 tensor and multiplies by the scale again, corrupting intentionally dequantized FP8 loads; the restore should skip when the quantization config is in dequantize mode or when the caller requested 16-bit dequantization.
Useful? React with 👍 / 👎.
| if os.path.isdir(model_name): | ||
| index_path = os.path.join(model_name, index_file) | ||
| if not os.path.exists(index_path): | ||
| return None |
There was a problem hiding this comment.
For an FP8 checkpoint saved as a single model.safetensors file, which from_pretrained can load normally, there is no model.safetensors.index.json, so this branch returns None and the restore is skipped even if the file contains dropped *.weight_scale_inv tensors. This leaves unsharded local or small Hub FP8 checkpoints with the same unscaled weights this helper is meant to repair.
Useful? React with 👍 / 👎.
| module_name = scale_key[: -len(".weight_scale_inv")] | ||
| try: | ||
| module = model.get_submodule(module_name) |
There was a problem hiding this comment.
When the load uses key_mapping (for example the text_only=True VLM path remaps language_model.model.* checkpoint keys onto model.* modules), deriving module_name directly from the raw checkpoint scale key means get_submodule looks for the pre-remap name and silently skips the scale. In that context the weights were loaded under the mapped module names, but the dropped FP8 scales are never restored.
Useful? React with 👍 / 👎.
- Skip restore when the model has no fp8 weights, so an intentionally dequantized load (load_in_16bit) is never re-scaled and corrupted. - Thread revision, subfolder and cache_dir through the index and shard downloads so scales come from the same snapshot as the weights. - Cover unsharded single-file model.safetensors checkpoints (no index). - Handle transposed block-scale layouts and skip on a true grid mismatch instead of applying a wrong scale. - Match text-only VLM loads where the language_model prefix was stripped. - Restore on the FastLanguageModel text path too, not only vision. - Handle a scalar weight_block_size; per-tensor error handling so one bad tensor cannot abort the rest or hide a partial mutation.
for more information, see https://pre-commit.ci
|
Thanks for the thorough review. Pushed a hardening commit that addresses the genuine items:
The one suggestion I did not take is the rank aware reshape for arbitrary tensor ranks: block FP8 scales are 2D only and the loop already skips non-2D weights, so the 2D handling is intentional (it also enables the grid and transpose checks). Re-verified end to end on |
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e4449e65cc
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| scale_expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(bs1, dim = 1) | ||
| scale_expanded = scale_expanded[:out_features, :in_features].to(weight.device) | ||
| with torch.no_grad(): | ||
| module.weight.data = (weight.to(torch.float32) * scale_expanded).to( | ||
| weight.dtype |
There was a problem hiding this comment.
When this runs on a large FP8 model already near the device_map VRAM limit, the restore path temporarily allocates a full expanded scale tensor plus a full fp32 copy/product of the projection weight on the target GPU. For a dropped MLP projection this can be hundreds of MB to multiple GB per layer, so loads that fit before this fix can now OOM during the post-load repair; restoring in row/block chunks or applying scales without expanding the whole matrix would keep the peak memory bounded.
Useful? React with 👍 / 👎.
| # Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200). | ||
| _restore_dropped_fp8_scales( |
There was a problem hiding this comment.
When FastLlamaModel.from_pretrained(..., num_labels=...) loads an FP8 checkpoint, execution stays in the AutoModelForSequenceClassification branch above and never reaches this restore call. The same dropped weight_scale_inv tensors can therefore remain un-applied for sequence-classification loads, leaving the base encoder/decoder weights raw and producing the same garbage activations this fix is meant to prevent.
Useful? React with 👍 / 👎.
| index_file = "model.safetensors.index.json" | ||
| single_file = "model.safetensors" |
There was a problem hiding this comment.
If the caller loads a safetensors variant (for example variant="fp8"), the main HF load reads model.<variant>.safetensors or its variant index, but this repair path only looks for the default model.safetensors.index.json / model.safetensors files. Variant-only FP8 repos will silently skip the restore, and repos that also have default weights can apply scales from the wrong checkpoint to the variant weights.
Useful? React with 👍 / 👎.
| if "language_model." in base: | ||
| candidate = base.replace("language_model.", "", 1) |
There was a problem hiding this comment.
For full VLM loads where checkpoint keys use language_model.model.* but the live module tree exposes model.language_model.*, this mapping strips the wrapper to model.* and then gives up. Those dropped scales are therefore never restored for that documented prefix mismatch, so affected VLM loads remain with raw FP8 values in their bf16 Linear weights; add the model.language_model.* alias before returning None.
Useful? React with 👍 / 👎.
- Bound peak memory: dequantize block views in place with the fp32 scale broadcast instead of materializing a full expanded scale and fp32 copy, so a near-VRAM-limit load is not pushed into OOM by the repair. - Restore on the sequence-classification load path too. - Cover more VLM key remappings (language_model.model.* to model.language_model.*) when matching modules. - Skip the restore for variant loads (variant=...) rather than risk applying default-checkpoint scales to variant weights.
|
Second round addressed:
Re-verified on |
for more information, see https://pre-commit.ci
|
@codex review |
|
Codex Review: Didn't find any major issues. Swish! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
Note on CI: the Backend CI Python 3.10 / 3.12 / 3.13 reds are a pre-existing studio issue unrelated to this change. They fail collecting |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a79a34b7ed
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| model_name, | ||
| local_files_only = kwargs.get("local_files_only", False), | ||
| token = token, | ||
| revision = revision, |
There was a problem hiding this comment.
When FastLlamaModel.from_pretrained(..., revision=...) is used, this repair now resolves the FP8 scale shards from that revision, but the CausalLM/SequenceClassification AutoModel...from_pretrained calls above do not forward revision at all (in this file it is otherwise only the wrapper argument/comment and these restore calls). In that scenario the weights are loaded from the default branch while the newly restored weight_scale_inv tensors come from the requested revision, so differing checkpoint revisions can silently rescale the wrong weights.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Correct. In llama.py the CausalLM/SequenceClassification weight loads resolve model_name on the default branch (revision is not forwarded there, unlike vision.py where it flows through kwargs). Changed the restore calls to read scales with revision=None so they come from the same default branch as the weights.
| weight = getattr(module, "weight", None) | ||
| if ( | ||
| not isinstance(weight, torch.Tensor) | ||
| or weight.device.type == "meta" |
There was a problem hiding this comment.
When a pre-quantized FP8 model is loaded with a device_map that offloads some layers to disk, those parameters remain on the meta device until the offload hook materializes them. This guard skips exactly those layers, so if one of the dropped-scale Linear weights is offloaded, the in-memory layers are repaired but the offloaded layer remains raw FP8 values and is loaded unscaled during inference; the restore pass needs to update the offloaded tensor or defer repair until the real weight is available.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Good point. A disk-offloaded layer keeps its weight on the meta device until the offload hook materializes it, so the scale cannot be applied in place at load time. The restore now skips such layers explicitly and prints a warning instead of silently leaving them unscaled; full in-place repair of disk-offloaded layers would need offload-hook integration and is a follow-up.
…disk-offloaded layers In llama.py the CausalLM/SequenceClassification weight loads resolve model_name on its default branch (revision is not forwarded there), so read the dropped weight_scale_inv tensors from the same default branch instead of the requested revision, avoiding rescaling default-branch weights with scales from another revision. In loader_utils.py a disk-offloaded layer keeps its weight on the meta device until the offload hook materializes it, so the scale cannot be applied in place. Skip such layers explicitly and print a warning rather than silently leaving them unscaled.
|
@codex review |
|
@codex review |
|
@codex review |
|
@codex review |
1 similar comment
|
@codex review |
|
Codex Review: Didn't find any major issues. 👍 Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
1 participant
Summary
Fixes #6200. Some block-scale FP8 checkpoints load as a garbage base model because a subset of Linear layers lose their
weight_scale_invon load.On
Qwen/Qwen3.6-27B-FP8(and models with the same nestedlanguage_modelblock-FP8 layout), transformers leaves everymlp.gate_projas a plain bf16Linearinstead of converting it to an FP8 module. Its raw quantized values get read into the bf16 weight and itsweight_scale_invis discarded as an unexpected key, so the weight is used without dequantization. The result is a weight roughly 200x too large and an unusable model (perplexity around 2 million).Fix
After the model is created in
FastBaseModel.from_pretrained,_restore_dropped_fp8_scaleswalks the checkpoint'smodel.safetensors.index.json. For every*.weight_scale_invwhose live weight is not fp8 (the dropped case), it loads the block scale, expands it to the weight shape, and multiplies it into the weight in place. Modules that were converted correctly keep an fp8 weight and are skipped, so there is never any double scaling. The bf16 weight already exists after load, so this adds no extra memory and is lossless (fp8 e4m3 values are exactly representable in bf16).The helper is self triggering (it only acts when the config is block FP8 and the index carries scale keys) and fail safe (any error returns a no-op), so single-file checkpoints, non-FP8 loads, and already-correct FP8 models are untouched.
Verification
Real fails-before / passes-after on
Qwen/Qwen3.6-27B-FP8(single B200,load_in_fp8=True):weight_scale_inv认真开展(gibberish)The(coherent)No regression on a healthy checkpoint:
unsloth/Qwen3-8B-FP8loads with all 252 scales already live, coherent generation, and the helper is a no-op (0 restored). Flat FP8 models (0.6B through 14B) are unaffected.New CPU unit test
tests/test_fp8_restore_dropped_scale.py(7 cases): dequant correctness, skip-when-already-fp8, non block divisible shape, nested submodule names, and the no-op guards for missing scale keys / missing index / non-FP8 config.Files
unsloth/models/loader_utils.py-_restore_dropped_fp8_scalesand helpers.unsloth/models/vision.py- call it once after the model is loaded.tests/test_fp8_restore_dropped_scale.py- new test.