Conversation
![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 functionality to restore missing .weight_scale_inv tensors for FP8 modules during model loading, along with helper utilities and comprehensive unit tests. The feedback suggests improving the robustness of optional imports in unsloth/models/loader_utils.py by using absolute imports and catching ModuleNotFoundError specifically instead of a broad exception.
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.
![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: 0ca7ca40e1
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
5c896a1 to
a2640a0
Compare
2b83177 to
a8e4474
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0e4c298613
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
0e4c298 to
cf0e821
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 78433573d7
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c958f1f854
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a8f586362c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e8a19f50d3
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bfae1be1c9
ℹ️ 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".
bfae1be to
41c9918
Compare
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7022b7063b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
unsloth/unsloth/models/loader.py
Line 627 in 7eb0d47
Clear FP8 request after PEFT base remap
When a PEFT adapter's base is remapped to a pre-quantized FP8 repo and fast_inference=True, this keeps load_in_fp8 true after selecting the FP8 checkpoint. The later vLLM path turns that true flag into fp8_mode and passes TorchAO quantization overrides, so it attempts on-the-fly FP8 quantization on an already-FP8 base; mirror the non-PEFT branch by clearing the flag when the base name changes here and in the duplicated FastModel path.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
unsloth/unsloth/models/loader.py
Lines 625 to 631 in c619ad4
Disable FP8 requests after PEFT base remaps
When load_in_fp8=True is used with a PEFT adapter whose base maps to a pre-quantized FP8 repo, this remaps the base name but keeps load_in_fp8 set. The later dispatch still receives load_in_fp8=True, so fast-inference/vLLM PEFT loads can try to quantize an already-FP8 checkpoint instead of just loading it; mirror the non-PEFT branch that clears the flag when the mapper changes the model name, and apply the same fix in the duplicated FastModel PEFT block.
ℹ️ About Codex in GitHub
Codex has been enabled to automatically 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 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bc71f4124c
ℹ️ 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".
| f"*.{variant}.safetensors", | ||
| f"*.{variant}-*.safetensors", |
There was a problem hiding this comment.
When restoring scales from a remote repo with an explicit variant, these wildcard patterns still match every *.{variant}.safetensors / *.{variant}-*.safetensors artifact in the repo, including files that from_pretrained did not select (for example extra adapter or alternate component safetensors). In those repos the post-load restore step can download large unused artifacts despite the selected-artifact restriction; use the exact default model/pytorch variant names here and rely on the index-extra pass for custom shard filenames.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 7650ed7: the variant branch now lists the default model/pytorch_model variant single-file, sharded, and index names instead of bare *.{variant} globs, so alternate-component variant safetensors (adapters, vision towers) are no longer pulled; custom shard filenames still come from the index-extra pass. Added a test asserting the default variant names match but a stray adapter.{variant}.safetensors does not.
| target_device = ( | ||
| projection_device if projection_device is not None else scale_tensor.device | ||
| ) | ||
| if target_device.type == "meta": |
There was a problem hiding this comment.
When a missing scale placeholder remains as a meta tensor but the owning FP8 weight has already been materialized (for example low-memory/device-map loading with a checkpoint alias such as .scale), target_device is taken from that meta placeholder and this branch skips the restore. That leaves the FP8 layer with a meta/missing scale even though the real weight device is available; fall back to the materialized weight/projection device before deciding to skip.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 7650ed7: the device selection now prefers a materialized device, so a scale placeholder left on meta next to an already-loaded weight restores onto the live weight/projection device instead of skipping on the placeholder's meta device. Added a test with a meta weight_scale_inv placeholder and a live weight asserting the scale is restored on the weight device.
| ("gate_up_proj_scale", ".gate_up_proj_scale"), | ||
| ("gate_up_proj_scale_inv", ".gate_up_proj_scale_inv"), | ||
| ("down_proj_scale", ".down_proj_scale"), | ||
| ("down_proj_scale_inv", ".down_proj_scale_inv"), |
There was a problem hiding this comment.
Static-activation FP8 checkpoints also carry activation scale parameters such as activation_scale, gate_up_proj_activation_scale, and down_proj_activation_scale, and FP8 expert forwards read them when activation_scheme == "static". Because the restore scanner only admits the weight-scale suffixes in this allow-list, a static checkpoint whose activation-scale placeholders are dropped will never have those tensors fetched or restored, leaving the layer with missing/default activation scales even though the weight scales were repaired.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 7650ed7: added the static-activation scale suffixes (activation_scale plus the per-projection gate/up/gate_up/down variants) to the scanner allow-list, so a static FP8 checkpoint whose activation-scale placeholders were dropped has them restored alongside the weight scales. This stays checkpoint-driven (only restored when the checkpoint actually carries those tensors, i.e. a static-activation model), so dynamic-activation loads are unaffected. Added a test restoring a dropped activation_scale.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b3165b3208
ℹ️ 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".
| new_model_name = os.path.join(temp_dir, cache_name) | ||
| print(f"Unsloth: Quantizing '{model_name}' to fp8, using model_name='{new_model_name}' instead") | ||
|
|
There was a problem hiding this comment.
When a user already has a bin-only FP8 cache from a previous Unsloth version and the original source model cannot be re-quantized (for example an offline job or a gated/removed remote model), this removes the only usable cached directory before the new save_pretrained(..., safe_serialization=True) path has succeeded. Previously that cache path was returned and could still be loaded; now a failed regeneration leaves the user with no cache at all. Please regenerate into a separate temp directory and atomically replace the old cache only after the safetensors save completes, or keep the bin cache as a fallback.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 0d6c49c: re-quantization now writes into a staging dir and only replaces an existing bin-only cache after the safetensors save succeeds (via os.replace), so a failed regeneration from an offline, gated, or removed source leaves the previous loadable cache intact instead of deleting it. Added a test that a from_pretrained failure preserves the bin cache and leaves no staging dir.
| # Update in place so runtime metadata (e.g. `block_size` fp8 dequant reads via | ||
| # getattr) survives; a fresh Parameter would drop it and force default geometry. #6749 | ||
| with torch.no_grad(): | ||
| existing.data = restored_scale |
There was a problem hiding this comment.
When low_cpu_mem_usage/device_map leaves an existing FP8 scale placeholder as a Parameter on the meta device while the weight has already been materialized, the restore path chooses the live weight device but then assigns that CPU/CUDA tensor through existing.data. PyTorch rejects set_data from a meta parameter to a materialized tensor with an incompatible tensor type, so these loads crash instead of restoring the scale; meta Parameter placeholders need to be replaced or otherwise materialized rather than updated via .data.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 0d6c49c: when the existing scale parameter is on meta, the restore now replaces it with a materialized Parameter instead of assigning through existing.data (which torch rejects for a meta to concrete set_data), so device-map/low-memory loads no longer crash. In-place data update is kept for a live parameter so its metadata survives. Added a test with a meta Parameter placeholder and a live weight.
| "model-*-of-*.safetensors", | ||
| "pytorch_model-*-of-*.safetensors", |
There was a problem hiding this comment.
For a remote sharded FP8 checkpoint that is not already present in the selected cache, or when the user passed force_download=True, these shard globs cause the restore helper to fetch every model shard before it has parsed the index, even though only the shard(s) containing scale tensors are needed. On large models this can repeat a full model download after from_pretrained has already loaded the model and can exhaust disk/network; download the index first and then request only the filenames mapped from scale keys.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 0d6c49c: the first snapshot pass now fetches only the index and the single-file default names, never the shard globs. The index-extra pass then reads the index and downloads only the shards the scale keys map to, so a large sharded checkpoint is no longer fully re-downloaded during restore (and force_download no longer refetches every shard). Repos without an index return no extra shards, matching the scanner which only reads scales from a single-file or index-listed layout. Updated the pattern tests accordingly.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0d6c49c98c
ℹ️ 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".
| elif isinstance(reference, torch.Tensor): | ||
| target_device = reference.device |
There was a problem hiding this comment.
When a projection-only FP8 expert has a materialized gate_up_proj/down_proj but its scale placeholder is still a meta tensor, this branch selects the meta reference and immediately skips the restore before _fp8_projection_device can use the live projection device. Fresh evidence in the current code is that the projection fallback only runs after both reference and weight are absent, so low-memory/device-map loads can still leave weightless experts without their checkpoint scale.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in c1eed80: the device selection now resolves the projection device up front and prefers it over a meta reference, so a projection-only expert with a materialized gate_up_proj but a meta scale placeholder restores onto the live projection device instead of skipping. Added a test for that case.
| module._parameters[attr_name] = torch.nn.Parameter( | ||
| restored_scale, | ||
| requires_grad = bool(getattr(existing, "requires_grad", False)), | ||
| ) |
There was a problem hiding this comment.
When an existing FP8 scale parameter is on the meta device, replacing it with a fresh Parameter only preserves requires_grad and drops runtime attributes such as block_size, which the FP8 kernels read to validate/dequantize non-default block layouts. Fresh evidence after the metadata fix is that this meta-materialization branch still creates a new parameter without copying the existing parameter's custom attributes, so checkpoints using non-default FP8 block geometry can be restored with the wrong layout metadata.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in c1eed80: when replacing a meta placeholder Parameter the restore now copies the existing parameter's custom attributes (vars(existing), e.g. block_size) onto the new Parameter, so non-default FP8 block geometry survives materialization. Added a test asserting block_size is preserved.
| ) | ||
| needs_refresh = os.path.isdir(new_model_name) and not has_safetensors_cache | ||
|
|
||
| if not os.path.isdir(new_model_name) or needs_refresh: |
There was a problem hiding this comment.
When needs_refresh is true but the source model cannot be re-quantized (for example an offline, gated, or removed remote model), the refresh branch still propagates the from_pretrained error and never returns the bin-only cache that the staging fix deliberately preserves. Fresh evidence in the current code is that the stale cache survives on disk but there is no exception path that loads it without forcing safetensors, so users with a previously usable /tmp/*-fp8-* cache still fail after upgrading.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in c1eed80: the regeneration is now wrapped so that if from_pretrained fails and a previous bin-only cache is still on disk (the needs_refresh case), the function returns that preserved cache instead of propagating the error, letting an already-usable cache keep loading after an upgrade. If there is no cache to fall back to it still raises. Updated the cache test and added a no-fallback-raises test.
| if load_in_fp8 != False and new_model_name != old_model_name: | ||
| restore_fp8_scales = True | ||
| load_in_fp8 = False |
There was a problem hiding this comment.
When load_in_fp8 remaps a base model to a pre-quantized FP8 sibling and the caller supplied use_safetensors=False, this branch clears load_in_fp8 but leaves that kwarg unchanged. The later from_pretrained can therefore look for .bin weights in safetensors-only FP8 repos and fail before the scale-restore path runs; mirror the offline FP8 path by forcing safetensors whenever the mapper selects an FP8 sibling.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in c1eed80: both mapped pre-quantized FP8 sibling branches now set kwargs['use_safetensors'] = True right after clearing load_in_fp8, mirroring the offline path, so a caller-supplied use_safetensors=False no longer makes from_pretrained look for absent .bin weights in a safetensors-only FP8 repo. Added an AST test asserting both branches force safetensors.
…lback, force safetensors for FP8 siblings in PR unslothai#6749
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c1eed80f3a
ℹ️ 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".
| load_in_16bit, | ||
| ) | ||
| model_name = _offline_quantize_to_fp8(model_name, fp8_mode, text_only = text_only) | ||
| kwargs["use_safetensors"] = True |
There was a problem hiding this comment.
When _offline_quantize_to_fp8 falls back to a preserved bin-only FP8 cache because regeneration failed (for example offline, gated, or removed source), this unconditional assignment makes the subsequent from_pretrained look only for safetensors in that bin-only directory, so the fallback returned from loader_utils.py is still unusable. Fresh evidence in the current code is this post-helper kwargs["use_safetensors"] = True assignment; only force safetensors when the helper actually produced a safetensors cache, or have the helper signal that it returned the legacy bin cache.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in d5b499f: the offline path now forces safetensors only when the returned cache dir actually contains a .safetensors file. When _offline_quantize_to_fp8 falls back to a preserved bin-only cache, the guard is false so from_pretrained loads the bin cache normally instead of failing on absent safetensors. restore_fp8_scales stays set (the scanner just finds nothing in a bin dir). Applied to both loader paths, with an AST test.
| if load_in_fp8 != False and model_name != base_before_remap: | ||
| load_in_fp8 = False |
There was a problem hiding this comment.
When a PEFT adapter's base is remapped to a pre-quantized FP8 sibling and the caller passed use_safetensors=False, this branch only clears load_in_fp8; unlike the top-level FP8 remap, it leaves the kwarg unchanged. FastBaseModel.from_pretrained can then search for .bin weights in a safetensors-only FP8 base repo and fail before scale restoration runs, so set kwargs["use_safetensors"] = True in this PEFT remap path as well.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in d5b499f: the PEFT base-remap branch now also sets kwargs['use_safetensors'] = True when it remaps to a pre-quantized FP8 sibling and clears load_in_fp8, mirroring the top-level FP8 remap, so FastBaseModel.from_pretrained does not search for .bin weights in a safetensors-only FP8 base repo. Added a test asserting the PEFT remap forces safetensors.
… for PEFT FP8 base remaps in PR unslothai#6749
|
Codex Review: Didn't find any major issues. Chef's kiss. 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". |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9a57062b57
ℹ️ 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".
| if isinstance(reference, torch.Tensor) and reference.device.type != "meta": | ||
| target_device = reference.device |
There was a problem hiding this comment.
When a scale placeholder survives as a regular CPU tensor while the FP8 weight/projection has already been materialized on CUDA (plain tensor attributes are not moved by device_map), this branch chooses the placeholder's CPU device and re-registers the restored scale there. The subsequent FP8 forward then sees CUDA weights with CPU scales and can hit a device mismatch; prefer the materialized weight/projection device whenever the reference tensor is on a different non-meta device.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in da86965: the device selection now prefers the live weight/projection compute device over a reference placeholder, so a scale placeholder left on CPU (device_map does not move plain tensor attributes) while the weight is on CUDA restores onto the weight device instead of registering a CPU scale next to CUDA weights. Meta-placeholder and projection-only cases still resolve to the live device. Added a CUDA cross-device test asserting the restored scale lands on the weight device.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: da86965f00
ℹ️ 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".
| and scale_tensor.numel() != 1 | ||
| and scale_tensor.ndim > 0 | ||
| and weight.ndim > 0 | ||
| and scale_tensor.shape[0] > weight.shape[0] |
There was a problem hiding this comment.
For blockwise FP8 checkpoints the scale tensor's first dimension is the number of row blocks, not the number of weight rows. In tensor-parallel/local-shard loads where the scale placeholder is missing, a full checkpoint scale such as (256, 64) can pass this raw weight.shape[0] check for a smaller local weight shard like (16384, 8192), get registered, and then the FP8 block kernel rejects it because it compares scale.shape to ceil(weight.shape / block_size) on first forward. Please validate block scales against the expected block grid before accepting missing-placeholder scales.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
This restore path runs after from_pretrained with device_map, which places whole Linear modules on a device; it does not tensor-parallel row-shard a single Linear, so weight here is always the full weight and its block grid matches the full checkpoint scale (scale.shape[0] = ceil(rows/block) < rows, which this guard accepts). The (256,64)-scale-on-a-16384-row-local-shard case needs intra-Linear TP sharding this loader does not perform, so there is no reachable mismatch to validate against here.
| else: | ||
| if hasattr(module, attr_name): | ||
| delattr(module, attr_name) | ||
| module.register_buffer(attr_name, restored_scale) |
There was a problem hiding this comment.
When this missing-attribute path restores a blockwise FP8 scale for a module whose placeholder was completely dropped, the new tensor is registered without the module's block_size metadata. Unsloth's patched FP8 linear path reads weight_scale.block_size and otherwise falls back to [128, 128], so checkpoints using any other weight_block_size will run the first forward with the wrong block geometry and either fail shape validation or dequantize with incorrect scales.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Fixed in 88985b5: before registering a restored scale whose placeholder was dropped, the code now carries the module's block_size (getattr(module, 'block_size') then the weight's) onto it, so the blockwise FP8 forward reads the correct geometry instead of defaulting to [128, 128]. The _parameters path already preserves the old attributes. Added a test with a non-default [64, 64] block_size.
| scale_tensor = scale_tensor.reshape(target_shape) | ||
| if ( | ||
| not isinstance(reference, torch.Tensor) | ||
| and isinstance(weight, torch.Tensor) |
There was a problem hiding this comment.
For weightless expert modules there is no top-level weight, so this compatibility guard never runs even though the restore path still attaches scales for gate_up_proj_scale(_inv) or down_proj_scale(_inv) based on the projection tensors. In tensor-parallel/local-shard loads where the scale placeholder was dropped, the full checkpoint expert scale can be registered next to a sharded projection and the expert forward then receives incompatible projection/scale shapes. Please validate missing expert scales against the corresponding projection tensor before registering them.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Same as the weight case: the restore runs post from_pretrained/device_map, which keeps a whole expert module (and its projection tensors) on one device rather than row-sharding a projection across ranks, so the projection seen here is the full projection and the full checkpoint expert scale matches it. A sharded-projection-with-full-scale mismatch would require intra-module TP sharding this loader does not do, so there is no reachable case to validate here.
|
@codex review |
|
Codex Review: Didn't find any major issues. More of your lovely PRs please. 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". |
|
Closing in favor of #6978, which merged the FP8 weight_scale_inv restore on July 9, 2026. |
3 participants
Problem
FastLanguageModel.from_pretrained(..., load_in_fp8=True, ...)can drop checkpoint scale tensors while loading FP8 checkpoints. ForQwen/Qwen3.6-27B-FP8, the reported failure drops tensors such asmlp.gate_proj.weight_scale_inv, so FP8 weights remain loaded without the inverse scale tensors Unsloth's FP8 runtime expects.Fix
Restore missing FP8 scale tensors only from the checkpoint artifacts selected for the actual load. The scanner now reads safetensors only, honors
subfolder,variant,use_safetensors,revision,token,local_files_only,cache_dir, andforce_download, uses Transformers-style sharded variant index names, and skips legacy.binshards rather than materializing a full shard after model load.Offline FP8 cache writes now use safe serialization, stale bin-only offline FP8 caches are regenerated before the loader forces
use_safetensors=True, and offline-generated FP8 caches forceuse_safetensors=Truebefore the laterfrom_pretraineddispatch. The restore loop stays limited to real local FP8 modules, skips meta-device weights, reshapes same-numel checkpoint scales to existing local scale placeholders, skips true shape-incompatible local shards, skips oversized checkpoint scales for smaller local weight shards, and can restore missing multi-element scales. It supports FineGrainedFP8 dense.scalealiases asweight_scale_inv, directweight_scale_inv, FBGEMMweight_scale, inverse expert scales, unfused expert aliases such asup_proj_scale(_inv), and FBGEMM expert scales such asgate_up_proj_scaleanddown_proj_scale. It also recognizes explicitly FP8 weightless expert owners so missing expert-scale placeholders can be restored. The optional FP8 owner probe is non-fatal if kernel imports are broken.FastModel PEFT base remapping now forwards
load_in_fp8, matching the text loader path, and clearsload_in_fp8when the remap selects a different FP8 sibling so vLLM does not apply on-the-fly FP8 to an already-FP8 base.Testing
Observed on this Windows host:
git diff --check origin/main...HEADpassed.The tests use synthetic modules and tiny local safetensors checkpoints, so they prove restore selection, variant and subfolder isolation, selected cache/download kwargs, Transformers-style sharded variant indexes, skipped
.binreloads, safe offline cache serialization, stale bin-only offline cache regeneration, forced safetensors loading for generated FP8 caches, non-fatal optional kernel import, meta and shape skips, same-numel local placeholder reshape, oversized local-shard scale skips, multi-element missing-scale restore, dense.scalemapping toweight_scale_inv, FBGEMMweight_scale, PEFT FP8 mapping, PEFT remap flag clearing, inverse expert scales, missing weightless expert placeholders, unfused expert aliases, and FBGEMM expert-scale restore without downloading the 27B model.Scope
This is limited to FP8 model loading and runtime scale restoration. It does not change FP8 export, Studio routes, frontend code, or generic quantization policy.
Closes #6200