fix(model): preserve seq-cls metadata across save round trips - #10228
Open
Excelius-Wang wants to merge 5 commits into
Open
Excelius-Wang wants to merge 5 commits into
Excelius-Wang wants to merge 5 commits into
Conversation
…on (modelscope#9704) When fine-tuning a VLM or LM with task_type=seq_cls (or reranker), the seq_cls patcher monkey-patches a 'score' head onto the generation model class without swapping the class itself. transformers' PreTrainedModel.save_pretrained writes 'model.__class__.__name__' into config.json['architectures'], so the on-disk checkpoint advertises the generation architecture (e.g. Qwen3VLForConditionalGeneration) while shipping a score head, num_labels, id2label, and problem_type. Downstream vLLM deployment reads architectures to pick the model class, finds the generation class, and rejects the checkpoint — even though inference through PtEngine works correctly. The fix rewrites model.config.architectures in-place to the matching *ForSequenceClassification class at the same place the score head is attached, so every save path (trainer, save_checkpoint, export, peft merge) writes the right value. A new helper _seq_cls_architectures() handles the suffix rewrite idempotently and leaves unknown/custom architectures untouched. Unit tests cover the rewrite, idempotence, multi-arch lists, and empty/None inputs. Refs: modelscope#9704
PreTrainedModel.save_pretrained() resets config.architectures to the model class name right before serializing the config, so the seq_cls name assigned in _patch_sequence_classification never reached config.json. Restore it in PretrainedConfig.save_pretrained, the last hook before the file is written, and cover the save path with a regression test that reads config.json back. Also fix the yapf formatting that failed the lint job.
The seq_cls architectures hook was bound onto the model instance as a bare closure capturing the original model and config, so copy.deepcopy(model) silently saved the original weights and pickle.dumps(model) failed outright. Install the wrapper on the model class instead, gated by an instance flag, so deep copies and pickled models keep saving their own weights. Also replace the == [class_name] check with a pre-save snapshot comparison: the restore now only undoes a change made during the save, so a value assigned after patching is kept, and a config without architectures no longer gets null written into config.json. Covered by regression tests for deepcopy, pickle, the architectures=None case, and post-patch assignment.
Follow up on modelscope#10063: bind a pickle-safe instance save wrapper and restore architecture metadata even when config saving is skipped or fails. Cover non-main saves, fresh-process pickle, failure cleanup and existing save wrappers.
Excelius-Wang
marked this pull request as ready for review
September 22, 2026 14:18
Transformers 5.16 checks is_remote_code through a classmethod. Patch the class _auto_class attribute temporarily so the pre-config failure regression exercises the export path on both old and new Transformers versions.
This branch has not been deployed
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
2 participants
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.
PR type
PR information
A patched sequence classifier can still save generation metadata after a non-main save, or after a pickle round trip to a fresh interpreter. This follow-up preserves classifier metadata in those cases and restores it when saving fails before config serialization.
Dependency: #10063. This branch retains RerankerGuo's three original commits through
c121140unchanged, followed by the incremental save fix and its test compatibility follow-up. Please review the follow-up commits separately while #10063 is pending; this follow-up is not intended to replace that PR or duplicate its credit. I can rebase onto main once the dependency lands.Reproduced cases at c121140
model.save_pretrained(first_dir, is_main_process=False)followed by an ordinary save writesQwen2ForCausalLMinto the second config. Transformers resets architectures but the non-main call skips the restoring config callback.Bind a module-level wrapper to the model instance with
functools.partial, retaining the original save implementation. Restore metadata before config serialization and infinally. Normalize bound methods before retaining them so pickle cannot resolve the original method name to the replacement wrapper. This uses ordinary Python serialization and does not modify the model class or reimplement checkpoint writing.Five added tests cover non-main then main saving, fresh-process pickle and saved score weights, failures during/before config saving, and an existing save wrapper with repeated installation. Existing deepcopy, custom architecture, missing architecture and unrelated-instance tests remain covered.
Experiment results
python -m pytest tests/general/test_model.py -q -k 'not test_qwen2 and not test_modelscope_hub'.pre-commit run --all-files: all 10 hooks passed.git diff --checkpassed.These are save-API edge cases. The training smoke deliberately injects the non-main call and pickle transport; it does not establish that ordinary Swift/DDP saving naturally triggers them. GPU/BF16 training, FSDP/DeepSpeed, full serialized-model forward semantics, and vLLM loading/inference are not validated. This follow-up does not claim to resolve the complete 20-label vLLM deployment problem in #9704.
The first NPU CI run used Transformers 5.16.1 and exposed a test setup issue:
is_remote_code()reads the class_auto_class, while the test set the instance attribute. The follow-up patches the class attribute in a restoring context. The original failure is reproduced on CPU with 5.16.1, the corrected test still detects the pre-fix defect, and all three local version selections pass. Updated NPU CI results are pending.