Repository navigation
feat(integrations): add the axolotl QPeftPlugin for QA-LoRA and PEQA - #4
Merged
Merged
Conversation
Step 3 of the axolotl integration: `adapter: qpeft` in an axolotl YAML trains QA-LoRA or PEQA and leaves the merged integer artifact (qpeft_config.json + qpeft_model.pt) in output_dir. qpeft/integrations/axolotl/plugin.py, QPeftPlugin: - Found through the `axolotl.plugins` entry point; every hook is a no-op unless `adapter: qpeft`, and the import adds only qpeft modules. - register refuses method efficient_qat before axolotl validates the config (Block-AP is step 4). - load_adapter builds the config through args.QPeftBlock and calls get_quant_model in place; it refuses inference / config-only loads. - post_trainer_create sets trainer.optimizer from param_groups with the trainer's own AdamW class and kwargs (axolotl never calls the plugin create_optimizer hook; HF keeps a preset optimizer), and adds FirstStepMustMoveParams. - post_train runs verify_quant_model, merges and saves; idempotent, because axolotl calls it twice. Logs under axolotl.integrations.qpeft. - Interim guard until axolotl has a plugin save hook (D1): the plugin owns save_pretrained on the model instance. A save into output_dir writes no transformers weights (the qpeft artifact is the model; after Ctrl-C it writes nothing and says so), so no model.safetensors that loads with random weights is left. Checkpoints still save. The Ctrl-C message uses os.write: axolotl's stdout tee holds a non-reentrant lock that a log call in its signal handler deadlocks on. Core: get_quant_model now refuses target_modules that match no layer, and names every target layer that group_size or packing cannot split, before the first layer is swapped (it used to stop at the first layer with the model half converted). Also: args.py refuses `inference` (axolotl inference would quantize the base again); `axolotl` extra; numpy in the dev extra. Tests: tests/test_axolotl_plugin.py (hooks on a tiny Llama with a real HF Trainer); tests/test_axolotl_e2e.py (QPEFT_RUN_SLOW=1: the real `axolotl train` CLI for both methods, fresh-process reload, resume, Ctrl-C).
…s; fix Ctrl-C logging
gapsong
force-pushed
the
fm/qpeft-axo-plugin
branch
from
October 12, 2026 12:58
7fd016f to
ce38373
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Intent
Update PR 4 (step 3 of the qpeft axolotl integration: QPeftPlugin for qa_lora and peqa) after PR 3 (GPTQ export, qpeft/export/gptq.py) was squash-merged into main and PR 4 now conflicts in README.md and AGENTS.md. Firstmate asked: rebase fm/qpeft-axo-plugin onto origin/main, resolve the README.md / AGENTS.md conflicts by keeping both sides' content, rerun the full test suite, push the update so PR 4 is clean and mergeable again; CI is skipped (no GitHub Actions; tests run locally only). Expected resolution: AGENTS.md Commands keeps both the tests/test_gptq_load.py line and the QPEFT_RUN_SLOW axolotl e2e line, then the 'GPTQ export facts' section, then the 'axolotl integration' section; README.md structure tests line reads '... tinygemm, GPTQ export, axolotl args and plugin'. No other content changes: the plugin commit and the two earlier no-mistakes fix commits stay as they are. Earlier captain decisions still hold and were already applied and reviewed: D1 (a) a later upstream axolotl save hook, not this step; D2 (a) QA-LoRA refuses GPTQ export in v1; the plugin refuses qpeft.export != qpeft at register() because wiring qpeft.export.export_gptq into the plugin is a later step (step 6), even though export_gptq now exists on main; README documents the manual plugins: line for axolotl <= 0.20.x and keeps axolotl>=0.20; efficient_qat is refused in the plugin (step 4). Known and out of scope: tests/test_examples.py::test_qat_trainer_example_runs (QPEFT_RUN_SLOW only) runs out of MPS memory on this 64 GB Mac in the E2E-QP backward; it fails identically on clean origin/main, and this branch does not touch examples/qat_trainer.py, qpeft/trainer.py or qpeft/block_ap.py. Do not add unrelated changes.
What Changed
qpeft/integrations/axolotl/plugin.pywithQPeftPlugin, registered under theaxolotl.pluginsentry point, plus a new.[axolotl]extra (axolotl>=0.20).With
adapter: qpeft, it quantizes the model inload_adapter, sets qpeft'sparam_groupsoptimizer and aFirstStepMustMoveParamscallback inpost_trainer_create, and in an idempotentpost_trainrunsverify_quant_model, merges, and saves the integer artifact tooutput_dir.It refuses
method: efficient_qat, anyqpeft.exportother thanqpeft, and inference/config-only loads.It also replaces
save_pretrainedon the model instance, so axolotl's final save and its Ctrl-C save intooutput_dirwrite no transformers weights, while Trainer checkpoints still save.args.py, theqpeft:block now also refusesinference,loraplus_lr_ratio,loraplus_lr_embedding,embedding_lr,embedding_lr_scaleandlr_groups.In
tuners_utils.py,get_quant_modelnow refuses before it changes the model whentarget_modulesmatches nonn.Linear, or when a matched layer'sin_featuresdoes not fitgroup_size/bits, and the error lists every such layer.QPEFT_RUN_SLOWend-to-end test that runs the realaxolotl trainCLI (tests/test_axolotl_e2e.py,tests/axolotl_probe.py).Updates the README (install line, the manual
plugins:line for axolotl <= 0.20.x, plugin usage, injection refusals, structure) and the AGENTS.md notes on axolotl.Risk Assessment
✅ Low: I compared the old and new commits with
git range-diff. The rebase changed only the AGENTS.md and README.md conflict regions, exactly as the intent asks. Commits f7531c0 and ce38373 are identical to their pre-rebase versions (11b463e, 7fd016f), and the plugin and tuner code was already reviewed. I traced the new "no target matched" refusal ininject_adapters: an already-quantized model returns before the check, so phase switches andQuantModel.from_pretrainedare not affected.Testing
In a throwaway venv with axolotl main d17227d, I drove the rebased plugin through the real
axolotl trainCLI, and all five end-to-end tests pass. QA-LoRA and PEQA artifacts reload in a fresh process and match the trained logits (PEQA bit-exact). Checkpoint resume reproduces the uninterrupted losses for both methods. Ctrl-C exits cleanly and writes no model file. All four refusal configs stopped before training with the expected messages. Targeted plugin and GPTQ export unit tests pass on the rebased code, but they were not a live run, so that scenario is marked untested. A Python 3.10 venv without axolotl confirms the plugin tests skip instead of breaking collection. The README.md and AGENTS.md resolution was checked only by reading the diff, so that scenario is also marked untested. The first e2e run failed only because of a venv setup gap (bitsandbytes is missing on macOS), which I fixed before re-running. There is no UI surface; the evidence is CLI logs. The only leftovers in the worktree are qpeft.egg-info and pycache directories from the editable install. git ignores them; my cleanup command was blocked by the permission prompt.git diff origin/main HEAD -- AGENTS.md README.mdwith the a…Evidence: axolotl e2e log: 5 passed
test_the_artifact_reloads_in_a_fresh_process_as_trained[qa_lora] {'max_diff': 0.0166, 'max_abs_logit': 12.98, 'argmax_agree': 1.0} PASSED test_the_artifact_reloads_in_a_fresh_process_as_trained[peqa] {'max_diff': 0.0, 'argmax_agree': 1.0, 'plain_transformers_load': 'refused: Error no file named model.safetensors ...'} PASSED test_resume_from_a_checkpoint_gives_the_same_losses[qa_lora] PASSED test_resume_from_a_checkpoint_gives_the_same_losses[peqa] PASSED test_ctrl_c_leaves_no_model_file_that_loads_with_random_weights PASSED 5 passed in 149.75sEvidence: axolotl train refusal: qpeft.export gptq
UnsupportedSchemeError: adapter: qpeft does not run qpeft.export: gptq in axolotl yet: the plugin writes only the qpeft artifact (qpeft_config.json + qpeft_model.pt), GPTQ export is not wired in. Use export: qpeft. (QA-LoRA refuses GPTQ export by design.)Evidence: axolotl train refusal: efficient_qat
UnsupportedSchemeError: adapter: qpeft does not run method efficient_qat in axolotl yet: its Block-AP phase is not wired into the plugin. Supported here: ['qa_lora', 'peqa'].Evidence: axolotl train refusal: loraplus_lr_ratio
adapter: qpeft refuses this config: -loraplus_lr_ratio: qpeft builds its own optimizer groups (param_groups) and would ignore it.Evidence: axolotl train refusal: lr_groups
adapter: qpeft refuses this config: -lr_groups: qpeft builds its own optimizer groups (param_groups) and would ignore it.Pipeline
Updates from git push no-mistakes
✅ **intent** - passed
✅ No issues found.
🔧 **Rebase** - 2 issues found → auto-fixed ✅
AGENTS.md- merge conflict rebasing onto origin/mainREADME.md- merge conflict rebasing onto origin/main🔧 Fix applied.
✅ Re-checked - no issues remain.
✅ **Review** - passed
✅ No issues found.
AGENTS.md- On macOS, axolotl main d17227d does not install bitsandbytes but imports it inaxolotl train. Sopip install -e ".[axolotl]"alone gives an e2e run that crashes atimport bitsandbytesuntil you install bitsandbytes by hand. AGENTS.md does not mention this. It is outside the scope of this rebase, so I only report it.git diff origin/main HEAD -- AGENTS.md README.mdwith the a…QPEFT_RUN_SLOW=1 /tmp/qpeft-test/venv/bin/python -m pytest tests/test_axolotl_e2e.py -v -s(1st run hit a venv gap: axolotl imports bitsandbytes but does not install it on macOS; after installing it, the 2nd run passed 5 of 5 in 2.5 min)python -m axolotl.cli.train export_gptq.yml(adapter: qpeft, qpeft.export: gptq) -> refused at register()python -m axolotl.cli.train efficient_qat.yml-> refused at register()python -m axolotl.cli.train loraplus.yml(loraplus_lr_ratio: 16) -> refused at config validationpython -m axolotl.cli.train lr_groups.yml-> refused at config validationpytest tests/test_axolotl_plugin.py tests/test_axolotl_args.py tests/test_tuners_utils.py tests/test_gptq_export.pyin the Python 3.12 venv with axolotlpytest tests/test_axolotl_plugin.py tests/test_axolotl_args.py tests/test_axolotl_e2e.pyin a Python 3.10 venv without axolotl (no collection error; the axolotl-dependent tests skip)git diff origin/main HEAD -- AGENTS.md README.md(checked the conflict resolution against the agreed text)✅ **Document** - passed
✅ No issues found.
✅ **Lint** - passed
✅ No issues found.
✅ **Push** - passed
✅ No issues found.