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

move the plugin post trainer create to the setup trainer (#2907)

Browse Source 
* move the plugin post trainer create to the setup trainer

* move post-train plugins to execute-training fn
This commit is contained in:
Wing Lian
2025-07-14 20:11:33 -04:00
committed by GitHub
parent 38359a8997
commit 5cc16040a8
1 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
src/axolotl/train.py
View File

@@ -224,6 +224,9 @@ def execute_training(
# torch.set_default_dtype(torch.bfloat16) # torch.set_default_dtype(torch.bfloat16)
trainer.train(resume_from_checkpoint=resume_from_checkpoint) trainer.train(resume_from_checkpoint=resume_from_checkpoint)
plugin_manager = PluginManager.get_instance()
plugin_manager.post_train(cfg, trainer.model)
def save_trained_model( def save_trained_model(
cfg: DictDefault, cfg: DictDefault,
@@ -510,6 +513,9 @@ def setup_model_and_trainer(cfg: DictDefault, dataset_meta: TrainDatasetMeta) ->
peft_config=peft_config, peft_config=peft_config,
) )
plugin_manager = PluginManager.get_instance()
plugin_manager.post_trainer_create(cfg, trainer)
return ( return (
trainer, trainer,
model, model,
@@ -541,9 +547,6 @@ def train(
processor, processor,
) = setup_model_and_trainer(cfg, dataset_meta) ) = setup_model_and_trainer(cfg, dataset_meta)
plugin_manager = PluginManager.get_instance()
plugin_manager.post_trainer_create(cfg, trainer)
# Handle untrained tokens if configured # Handle untrained tokens if configured
safe_serialization = cfg.save_safetensors is True safe_serialization = cfg.save_safetensors is True
train_dataset = dataset_meta.train_dataset train_dataset = dataset_meta.train_dataset
@@ -566,6 +569,4 @@ def train(
if not cfg.use_ray: if not cfg.use_ray:
cleanup_distributed() cleanup_distributed()
plugin_manager.post_train(cfg, model)
return model, tokenizer, trainer return model, tokenizer, trainer