From ed5178cd3d3ebe48dcabb45097565e02ac49dcbd Mon Sep 17 00:00:00 2001 From: Dan Saunders Date: Wed, 26 Feb 2025 21:03:44 +0000 Subject: [PATCH] update --- src/axolotl/train.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/src/axolotl/train.py b/src/axolotl/train.py index 008b27add..ae10c2800 100644 --- a/src/axolotl/train.py +++ b/src/axolotl/train.py @@ -28,17 +28,11 @@ from axolotl.utils.freeze import freeze_layers_except from axolotl.utils.models import load_model, load_processor, load_tokenizer from axolotl.utils.trainer import setup_trainer -# Optional imports with graceful fallbacks try: from optimum.bettertransformer import BetterTransformer except ImportError: BetterTransformer = None -# Project setup -project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) -src_dir = os.path.join(project_root, "src") -sys.path.insert(0, src_dir) - configure_logging() LOG = get_logger(__name__) @@ -233,9 +227,10 @@ def save_trained_model( # final model weights have already been saved by `ReLoRACallback.on_train_end` return - # TODO do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading - # only save on rank 0, otherwise it corrupts output on multi-GPU when multiple processes attempt to write the same file if cfg.fsdp: + # TODO: do we need this fix? https://huggingface.co/docs/accelerate/usage_guides/fsdp#saving-and-loading + # only save on rank 0, otherwise it corrupts output on multi-GPU when multiple + # processes attempt to write the same file if ( state_dict_type == "SHARDED_STATE_DICT" and cfg.fsdp_config.fsdp_state_dict_type == "SHARDED_STATE_DICT" @@ -267,7 +262,6 @@ def save_trained_model( os.remove(os.path.join(cfg.output_dir, "model.safetensors")) except FileNotFoundError: pass - elif cfg.local_rank == 0: if cfg.flash_optimum and BetterTransformer: model = BetterTransformer.reverse(model)