fix broken MX tests from transformers 5.8.1 upgrade#3679
Conversation
|
Important Review skippedAuto incremental 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 Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAdds a Transformers weight-initialization guard patch to prevent re-initialization failures on MXTensor and TorchAO quantized weights. The patch is applied during quantization, model loader setup, and validated via cross-process loading test. ChangesMX Quantization Initialization Guard
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
src/axolotl/utils/quantization.py (1)
170-179: ⚡ Quick winTrim this to a single short WHY comment per repo policy.
Please reduce this multi-line narrative to one short WHY-focused line to match the
src/axolotl/**/*.pycomment standard.As per coding guidelines, "Only add comments when explaining the WHY behind non-obvious logic, hidden constraints, or workarounds for specific bugs... Comments should be a maximum of one short line".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/axolotl/utils/quantization.py` around lines 170 - 179, Replace the multi-line explanatory docstring in src/axolotl/utils/quantization.py that describes why we skip re-initializing torchao-quantized weights with a single short WHY comment line; concisely state that we skip modules in from_pretrained/_init_weights to avoid re-initializing quantized tensors (e.g., MXTensor) which would drop the _is_hf_initialized flag and raise NotImplementedError.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/axolotl/loaders/model.py`:
- Around line 226-230: The unguarded import and call to
axolotl.utils.quantization.patch_transformers_skip_quantized_init() inside
_apply_pre_model_load_setup() can raise on platforms where torchao is
unavailable; change the code to avoid importing torchao-backed quantization
unconditionally by either (a) moving the import and call behind the existing
MX/quantization feature check or (b) wrapping the import/call in a try/except
ImportError that silently no-ops when axolotl.utils.quantization (or torchao) is
missing; reference patch_transformers_skip_quantized_init() and
_apply_pre_model_load_setup() when making the change.
---
Nitpick comments:
In `@src/axolotl/utils/quantization.py`:
- Around line 170-179: Replace the multi-line explanatory docstring in
src/axolotl/utils/quantization.py that describes why we skip re-initializing
torchao-quantized weights with a single short WHY comment line; concisely state
that we skip modules in from_pretrained/_init_weights to avoid re-initializing
quantized tensors (e.g., MXTensor) which would drop the _is_hf_initialized flag
and raise NotImplementedError.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 9d5a015b-c257-4359-b370-4b1a44c84bde
📒 Files selected for processing (3)
src/axolotl/loaders/model.pysrc/axolotl/utils/quantization.pytests/e2e/test_quantization.py
| from axolotl.utils.quantization import ( | ||
| patch_transformers_skip_quantized_init, | ||
| ) | ||
|
|
||
| patch_transformers_skip_quantized_init() |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== Check whether torchao is mandatory in project metadata =="
fd -i 'pyproject.toml|setup.py|requirements*.txt|uv.lock|poetry.lock' | while read -r f; do
echo "--- $f ---"
rg -n "torchao" "$f" || true
done
echo "== Confirm unconditional call site =="
rg -n -C3 'patch_transformers_skip_quantized_init|from axolotl.utils.quantization import' src/axolotl/loaders/model.py
echo "== Confirm torchao top-level imports in quantization module =="
rg -n -C2 '^from torchao|^import torchao' src/axolotl/utils/quantization.pyRepository: axolotl-ai-cloud/axolotl
Length of output: 1569
Avoid importing torchao-backed quantization unconditionally in the generic pre-load path
src/axolotl/loaders/model.py imports axolotl.utils.quantization and calls patch_transformers_skip_quantized_init() from _apply_pre_model_load_setup(). axolotl/utils/quantization.py imports torchao at module import time, and pyproject.toml declares torchao conditionally (missing on darwin/aarch64). This can cause import-time failures for non-quantized model loads in those environments—guard the import/call behind the MX/quantization condition (or wrap it in try/except ImportError with a no-op fallback).
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/axolotl/loaders/model.py` around lines 226 - 230, The unguarded import
and call to axolotl.utils.quantization.patch_transformers_skip_quantized_init()
inside _apply_pre_model_load_setup() can raise on platforms where torchao is
unavailable; change the code to avoid importing torchao-backed quantization
unconditionally by either (a) moving the import and call behind the existing
MX/quantization feature check or (b) wrapping the import/call in a try/except
ImportError that silently no-ops when axolotl.utils.quantization (or torchao) is
missing; reference patch_transformers_skip_quantized_init() and
_apply_pre_model_load_setup() when making the change.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
1 participant
Summary by CodeRabbit
Release Notes
Bug Fixes
Tests