| """ |
| Base wrapper for HF Trainer compatibility. |
| |
| Inherits PreTrainedModel so Trainer gets save_pretrained, from_pretrained, |
| tied weight handling, gradient checkpointing — all for free. |
| |
| Subclasses store the backbone as `self.backbone` and must implement |
| forward() and generate(). |
| """ |
|
|
| from abc import abstractmethod |
| from transformers import PreTrainedModel |
|
|
|
|
| class BaseModelWrapper(PreTrainedModel): |
| """Abstract base that wraps a CausalLM backbone for HF Trainer. |
| |
| Uses the backbone's config. The default state_dict strips the backbone |
| prefix for AR checkpoints; FlowModelWrapper overrides it to retain flow |
| layers in its checkpoints. |
| |
| Subclasses must implement forward() and generate(). |
| """ |
|
|
| base_model_prefix = "backbone" |
| supports_gradient_checkpointing = True |
|
|
| def __init__(self, backbone: PreTrainedModel): |
| import copy |
|
|
| config = copy.deepcopy(backbone.config) |
| config._attn_implementation_internal = "eager" |
| super().__init__(config) |
| self.backbone = backbone |
| if backbone._tied_weights_keys is not None: |
| self._tied_weights_keys = backbone._tied_weights_keys |
|
|
| def _init_weights(self, module): |
| pass |
|
|
| @abstractmethod |
| def forward(self, **kwargs): ... |
|
|
| @abstractmethod |
| def generate(self, **kwargs): ... |
|
|
| def state_dict(self, *args, **kwargs): |
| """Strip 'backbone.' prefix so checkpoints are loadable by AutoModel.""" |
| sd = super().state_dict(*args, **kwargs) |
| return {k.removeprefix("backbone."): v for k, v in sd.items()} |
|
|
| def load_state_dict(self, state_dict, *args, **kwargs): |
| """Load into backbone directly (handles keys with or without prefix).""" |
| if any(k.startswith("backbone.") for k in state_dict): |
| state_dict = {k.removeprefix("backbone."): v for k, v in state_dict.items()} |
| return self.backbone.load_state_dict(state_dict, *args, **kwargs) |
|
|
