tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

add hub_revision support for specifying branch when pushing checkpoints (#3387) [skip ci]

Browse Source
This commit is contained in:
madScientist10
2026-02-10 12:53:09 +02:00
committed by GitHub
parent ed7105dba7
commit 37e9da7a53
2 changed files with 9 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/axolotl/core/builders/base.py
View File

@@ -409,6 +409,9 @@ class TrainerBuilderBase(abc.ABC):
if self.cfg.hub_strategy: if self.cfg.hub_strategy:
training_args_kwargs["hub_strategy"] = self.cfg.hub_strategy training_args_kwargs["hub_strategy"] = self.cfg.hub_strategy
if self.cfg.hub_revision:
training_args_kwargs["hub_revision"] = self.cfg.hub_revision
def _configure_save_and_eval_strategy(self, training_args_kwargs: dict): def _configure_save_and_eval_strategy(self, training_args_kwargs: dict):
# save_strategy and save_steps # save_strategy and save_steps
if self.cfg.save_steps: if self.cfg.save_steps:

6
src/axolotl/utils/schemas/model.py
View File

@@ -120,6 +120,12 @@ class ModelOutputConfig(BaseModel):
default=None, default=None,
json_schema_extra={"description": "how to push checkpoints to hub"}, json_schema_extra={"description": "how to push checkpoints to hub"},
) )
hub_revision: str | None = Field(
default=None,
json_schema_extra={
"description": "branch/revision to push to on hub (default: main)"
},
)
save_safetensors: bool | None = Field( save_safetensors: bool | None = Field(
default=True, default=True,
json_schema_extra={ json_schema_extra={