Repository navigation
feat: share the Block-AP batch builder and first-step guard with any HF Trainer - #2
Merged
Merged
Conversation
AGENTS.md is the file every coding agent reads; CLAUDE.md points at it, so Claude Code still loads the same text. No content change.
Every future change should apply it: small interfaces that hide a lot, complexity pulled down into the module, no pass-through layers, and errors defined out of existence only where invariant 2 (refuse, do not approximate) still holds.
… Trainer The axolotl plugin (later steps) drives a Hugging Face Trainer it did not build and needs two things QATTrainer kept private: the Block-AP calibration batches and the "first step must move parameters" guard. Both now live in qpeft/hf_trainer.py, and QATTrainer calls them. block_ap_batches(trainer, *, train_size, seqlen, batch_size) trainer.add_callback(FirstStepMustMoveParams()) Shape (deep modules): each piece takes only what it cannot find itself. - block_ap_batches takes the trainer, not a list of batches. The hidden part is the Trainer-specific work: get_train_dataloader, the private _prepare_inputs (device placement), label removal, the sequence cut, the re-split and the exact row count. Taking pre-made batches would push _prepare_inputs into every caller. - FirstStepMustMoveParams no longer holds the trainer. It reads model and optimizer from the callback events, which is what the Trainer hands every callback, and it reads the fp16 skip from the optimizer (accelerate's step_was_skipped, the flag that accelerator.optimizer_step_was_skipped itself reads). So it works with any Trainer through the plain callback contract. Not merged into one "Block-AP phase" call (batches + run_block_ap + set_phase): that is plan step 4's call to make, when the plugin is the second user. hf_trainer.py needs transformers, so `import qpeft` stays torch-only. No behavior change: QATTrainer with the old and the new code gives a bit-identical merged model (Block-AP + E2E-QP on a tiny Llama, CPU). Two existing tests change with the move: the fp16-skip test now patches AcceleratedOptimizer.step_was_skipped, and the batch-size test checks that QATTrainer's block_ap_* args reach the batches. The batch builder and the guard get their own tests on a plain Trainer.
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
Plan step 1 of the qpeft axolotl integration (plan in firstmate report qpeft-axolotl-s7, section 8): extract the reusable training pieces out of QATTrainer so the future axolotl plugin (shipped inside qpeft, later steps) can use them without QATTrainer. Pieces: the Block-AP calibration batch builder (was private QATTrainer._block_ap_batches) and the 'first optimizer step with lr > 0 must move a trainable parameter' guard callback (was private _FirstStepMustMoveParams). Make them a public, deep module with a tiny interface that both QATTrainer and the future plugin call. Result: new qpeft/hf_trainer.py with block_ap_batches(trainer, , train_size, seqlen, batch_size) and FirstStepMustMoveParams() (no constructor args). Captain priority: design after Ousterhout, A Philosophy of Software Design - deep modules, small interfaces hiding complexity, no shallow wrappers or pass-through methods, no leaked internals; keep invariant 2 (refuse instead of approximate). Deliberate decisions: (1) block_ap_batches takes the HF Trainer, not pre-made batches, so the private Trainer._prepare_inputs and get_train_dataloader stay hidden in one place instead of in every caller. (2) The callback no longer holds the trainer; it reads model and optimizer from the HF callback event kwargs (verified in transformers: the callback handler's model/optimizer are set to trainer.model/trainer.optimizer before on_train_begin). Because of that, the existing fp16-skip test no longer patches the accelerator property, and the old test that called the removed private block_ap_batches now checks that QATTrainer's block_ap args reach the batches (spying run_block_ap); the builder itself is tested on a plain HF Trainer in tests/test_hf_trainer.py. (3) Deliberately NOT merged into one 'Block-AP phase' call (batches + run_block_ap + set_phase e2e_qp): that is plan step 4's decision once the plugin is the second user. (4) hf_trainer.py needs transformers and is not exported from qpeft/init so 'import qpeft' / star-import stays torch-only; import path is qpeft.hf_trainer. (5) Block-AP args have no defaults in block_ap_batches on purpose (no second source of truth beside QATTrainingArguments). Must be a pure refactor: QATTrainer behavior unchanged - verified by the existing suite plus a run of old vs new code that gave a bit-identical merged model. Known pre-existing gap, fix explicitly WANTED by the captain: a previous no-mistakes test step found that the fp16 GradScaler skip check (accelerate's step_was_skipped, and the old accelerator.optimizer_step_was_skipped) is not set with HF's default fused AdamW, so a skipped first fp16 step can falsely abort training; its fix (detect the skip by a drop in the GradScaler scale instead of the accelerate flag) was lost only because the commit failed on a missing git identity, now fixed. If the test step finds it again, fixing it in this PR is in scope and is the one allowed behavior change (it makes the guard correct, never weakens fake_quant == merge). Also requested: add a short 'Design principle: deep modules (Ousterhout)' section (2-4 lines) to qpeft's CLAUDE.md. The firstmate helper fm-ensure-agents-md.sh moved CLAUDE.md to AGENTS.md with CLAUDE.md as a symlink (own commit, no content change); the section lives in AGENTS.md. README structure tree and docs/specs/qat_trainer.md got one-line pointers. Keep the PR small and reviewable; the captain reads all code. Never weaken exactness (fake_quant == merge). Tests run locally only (pytest); no GitHub Actions dependency.
What Changed
qpeft/hf_trainer.py(needstransformers, not exported fromqpeft/__init__) with two public pieces moved out ofQATTrainer.block_ap_batches(trainer, *, train_size, seqlen, batch_size)was the privateQATTrainer._block_ap_batches.FirstStepMustMoveParams()was the private_FirstStepMustMoveParams; it no longer holds the trainer and readsmodelandoptimizerfrom the HF callback arguments.QATTrainernow calls both, with the same Block-AP args.step_was_skippedflag. Before, with HF's default fused AdamW, a skipped first fp16 step could stop training by mistake.tests/test_hf_trainer.pytests the batch builder on a plain HF Trainer, and tests the guard with stub and real Trainer runs (including real fp16 skips).tests/test_qat_trainer.pynow checks thatQATTrainer'sblock_ap_*args reach the batches.CLAUDE.mdmoved toAGENTS.md(CLAUDE.mdis now a symlink to it), andAGENTS.mdgets a short "Design principle: deep modules (Ousterhout)" section. The README structure tree anddocs/specs/qat_trainer.mdnow point to the new module.Risk Assessment
✅ Low: The extraction matches the old QATTrainer code. The combined skip check (accelerate step_was_skipped OR a GradScaler scale drop) now covers non-fused AMP, fused AdamW and DeepSpeed. Checked against transformers 5.19 and accelerate 1.15: the callback handler gets the accelerate-wrapped optimizer before on_train_begin, and that wrapper exposes
.scalerandstep_was_skipped.Testing
I ran the targeted test files
test_hf_trainer.pyandtest_qat_trainer.py; all passed, and the fp16 tests really ran on MPS. Then I drove the product live on the locally cached tiny Qwen3 model. QATTrainer end to end (Block-AP, E2E-QP, merge check, save, load) on base and HEAD code gave the same artifact sha256 (980a066a...) and the same logits after load, with merge equivalence max|delta| = 0. A plain HF Trainer with the newqpeft.hf_trainerpieces worked as an integration would use it. With fp16 on MPS and an overflowing first step, the real GradScaler skipped steps 1-3. Fused AdamW never set accelerate'sstep_was_skipped. The base guard falsely stopped training; the new guard did not, for both fused and non-fused AdamW. The new guard still stopped lr=0 runs and runs where the step moves nothing.import qpeftstays torch-only. Only the DeepSpeed fp16 skip path was not driven live, because there is no DeepSpeed and no CUDA here. All evidence logs and scripts are in the evidence directory. The worktree is clean. One temp dir under $TMPDIR (base snapshot plus run outputs) stays: a permission prompt blocked its removal.Evidence: QATTrainer run, HEAD code (sha256 + merge check + load)
Evidence: QATTrainer run, base commit code (same sha256)
Evidence: Old vs new merged artifact
Evidence: Plain Trainer fp16 skip, fused AdamW, NEW guard
Evidence: Plain Trainer fp16 skip, fused AdamW, BASE guard (false abort reproduced)
Evidence: fp16 skip probe (fused AdamW, MPS)
Evidence: Plain Trainer fp16 skip, non-fused AdamW, NEW guard
Evidence: Dead run lr=0 stopped
Evidence: Dead run no-move stopped (fused)
Evidence: Dead run no-move stopped (non-fused)
Evidence: torch-only import + public interface check
Evidence: Targeted pytest output
~/.no-mistakes/evidence/01M4GRYTPPD7M16ZV1QYMQ72Q5/scripts)Pipeline
Updates from git push no-mistakes
✅ **intent** - passed
✅ No issues found.
✅ **Rebase** - passed
✅ No issues found.
🔧 **Review** - 1 issue found → auto-fixed (2) ✅
qpeft/hf_trainer.py:66- This is the known gap the captain asked to have fixed. FirstStepMustMoveParams now says it works for 'any HF Trainer', but it detects a skipped fp16 step only through accelerate'soptimizer.step_was_skipped. That flag is never set when the optimizer is fused AdamW. How it fails: accelerate swaps in its ownoptimizer.step, which sets_accelerate_step_called = True. Fused AdamW has_step_supports_amp_scaling, so the GradScaler always calls thatstepand lets the fused kernel skip the update when there is an overflow. So_is_overflowstays False. Concrete sequence: a plain HF Trainer (the planned axolotl path, and tests/test_hf_trainer.py) with the default optim (adamw_torch_fusedin transformers 5.x), fp16=True, on CUDA. The first step overflows at the initial scale of 65536, which is common, so the GradScaler skips it. The params do not move, and line 70 raises 'no trainable parameter changed', so a healthy run stops by mistake. QATTrainer is not hit today because it builds a non-fusedtorch.optim.AdamW(groups). Fix, as the intent says: in on_step_begin, recordoptimizer.scaler.get_scale()whengetattr(optimizer, 'scaler', None)is set. In on_step_end, treat a lower scale as a skipped step (the GradScaler halves the scale only on overflow). Remove thestep_was_skippedread. Places that depend on the same rule: qpeft/hf_trainer.py:57-58 (on_step_begin must take the scale snapshot); tests/test_hf_trainer.py:91-92 and :124-129 (the fake optimizer and the skip test usestep_was_skippedand must use a scaler whose scale drops); tests/test_qat_trainer.py:287-288 (patchesAcceleratedOptimizer.step_was_skippedand must instead lower the scaler scale on the skipped step). Add a regression test that fails before the fix: a real or fake GradScaler that skips the step with no accelerate flag set. This is the one behavior change the intent allows. It does not weaken fake_quant == merge.🔧 Fix applied.
2 warnings still open:
qpeft/hf_trainer.py:66- This is the known gap the captain asked to have fixed. FirstStepMustMoveParams now says it works for 'any HF Trainer', but it detects a skipped fp16 step only through accelerate'soptimizer.step_was_skipped. That flag is never set when the optimizer is fused AdamW. How it fails: accelerate swaps in its ownoptimizer.step, which sets_accelerate_step_called = True. Fused AdamW has_step_supports_amp_scaling, so the GradScaler always calls thatstepand lets the fused kernel skip the update when there is an overflow. So_is_overflowstays False. Concrete sequence: a plain HF Trainer (the planned axolotl path, and tests/test_hf_trainer.py) with the default optim (adamw_torch_fusedin transformers 5.x), fp16=True, on CUDA. The first step overflows at the initial scale of 65536, which is common, so the GradScaler skips it. The params do not move, and line 70 raises 'no trainable parameter changed', so a healthy run stops by mistake. QATTrainer is not hit today because it builds a non-fusedtorch.optim.AdamW(groups). Fix, as the intent says: in on_step_begin, recordoptimizer.scaler.get_scale()whengetattr(optimizer, 'scaler', None)is set. In on_step_end, treat a lower scale as a skipped step (the GradScaler halves the scale only on overflow). Remove thestep_was_skippedread. Places that depend on the same rule: qpeft/hf_trainer.py:57-58 (on_step_begin must take the scale snapshot); tests/test_hf_trainer.py:91-92 and :124-129 (the fake optimizer and the skip test usestep_was_skippedand must use a scaler whose scale drops); tests/test_qat_trainer.py:287-288 (patchesAcceleratedOptimizer.step_was_skippedand must instead lower the scaler scale on the skipped step). Add a regression test that fails before the fix: a real or fake GradScaler that skips the step with no accelerate flag set. This is the one behavior change the intent allows. It does not weaken fake_quant == merge.qpeft/hf_trainer.py:66- Round 1 fixed fused AdamW, but it removed the accelerate skip flag completely, as that round told it to. This opened a gap for DeepSpeed. Now FirstStepMustMoveParams sees a skipped step only when the GradScaler scale goes down (_grad_scalereadsoptimizer.scaler). Under DeepSpeed, accelerate creates no GradScaler (accelerator.py: native_amp/scaler stay None for DistributedType.DEEPSPEED), and DeepSpeedOptimizerWrapper is built withscaler=None. So_grad_scalealways returns 1.0, and a DeepSpeed fp16 overflow skip is never seen. The base code handled this case:accelerator.optimizer_step_was_skippedreadsDeepSpeedOptimizerWrapper.step_was_skipped, which isoptimizer.overflow. How it fails: a plain HF Trainer or QATTrainer with a DeepSpeed ZeRO-1/2 config and fp16 (the multi-GPU path HF provides, and a common setup for the planned axolotl plugin). The first step overflows at the default initial loss scale 2^16. DeepSpeed skips it, so no param moves, and line 73 raises 'no trainable parameter changed'. A healthy run stops by mistake. Fix: count a step as skipped when either signal is set:getattr(optimizer, 'step_was_skipped', False) or _grad_scale(optimizer) < self.scale_before. The flag covers DeepSpeed and non-fused AMP, and the scale drop covers fused AdamW. Keep the comment on why the scale check is needed. Add a unit test with a fake optimizer that hasstep_was_skipped=Trueand no scaler. Same rule elsewhere: tests/test_hf_trainer.py:85-86 (_optimizerfake has nostep_was_skipped).🔧 Fix applied.
✅ Re-checked - no issues remain.
✅ **Test** - passed
✅ No issues found.
.venv/bin/python -m pytest tests/test_hf_trainer.py tests/test_qat_trainer.py -q -rs(49 passed, fp16 tests ran on MPS, none skipped)PYTHONPATH=<worktree> python scripts/qat_run.pyvsPYTHONPATH=<base b01059f qpeft> python scripts/qat_run.py: QATTrainer Block-AP -> E2E-QP -> merge check -> save_model -> QuantModel.from_pretrained on cached axolotl-ai-co/tiny-qwen3-129m, CPU, offline; sha256 of merged artifact + loaded logits comparedscripts/plain_trainer_run.py fp16_skip_fused/fp16_skip_torch: plain HF Trainer + QuantModel, block_ap_batches -> run_block_ap -> set_phase(e2e_qp) -> train fp16 on MPS with an inf first gradient; per-step probe of GradScaler scale, accelerate step_was_skipped, params movedscripts/plain_trainer_run.py old_callback_fp16_skip_fused/old_callback_fp16_skip_torch: same run with the base commit's _FirstStepMustMoveParams(trainer) to reproduce the pre-fix false abortscripts/plain_trainer_run.py lr0,nomove_fused,nomove_torch: adversarial dead runs in the real Trainer looppython -c 'import sys, qpeft; from qpeft import *; ...': transformers not imported, hf_trainer not exported, public signatures checked✅ **Document** - passed
✅ No issues found.
✅ **Lint** - passed
✅ No issues found.
✅ **Push** - passed
✅ No issues found.