fusion-bench 0.2.6__py3-none-any.whl → 0.2.8__py3-none-any.whl

fusion_bench/method/lm_finetune/fullfinetune_sft.py

@@ -1,3 +1,4 @@
+import functools
 import itertools
 import logging
 import os
@@ -18,7 +19,7 @@ from typing_extensions import TYPE_CHECKING, override
 
 from fusion_bench import BaseAlgorithm, BaseModelPool
 from fusion_bench.dataset.llama.collate import padded_collate_sft
-from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.mixins import FabricTrainingMixin
 from fusion_bench.modelpool import CausalLMPool
 from fusion_bench.utils import instantiate
 from fusion_bench.utils.dtype import get_dtype
@@ -34,7 +35,7 @@ if TYPE_CHECKING:
 log = logging.getLogger(__name__)
 
 
-class FullFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
+class FullFinetuneSFT(BaseAlgorithm, FabricTrainingMixin):
 
     model: Union[nn.Module, "_FabricModule", "LlamaForCausalLM"]
     optimizer: Union[torch.optim.Optimizer, "_FabricOptimizer"]
@@ -59,7 +60,10 @@ class FullFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
         gradient_clip_algorithm: Literal["value", "norm"] = "norm",
         save_optimizer_state: bool = False,
         save_full_model: bool = False,
+        save_ckpt_type: Literal["lightning", "hf"] = "lightning",
         ckpt_path: Optional[str] = None,
+        max_length: int = 6144,
+        fix_token_embedding: bool = True,
         **kwargs,
     ):
         """
@@ -81,7 +85,10 @@ class FullFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
             gradient_clip_algorithm(str): Algorithm to use for gradient clipping. Available options: 'value', 'norm'. If set to 'value', the gradients will be clipped to the specified value. If set to 'norm', the gradients will be clipped to the specified norm.
             save_optimizer_state(bool): Whether to save the optimizer and lr_scheduler state along with the model checkpoint.
             save_full_model(bool): Whether to save the full model or only the trainable parameters in the model checkpoint.
+            save_ckpt_type (str): Type of checkpoint to save. Available options: 'lightning', 'hf'. If set to 'lightning', the checkpoint will be saved in the lightning format. If set to 'hf', the checkpoint will be saved in the huggingface format.
             ckpt_path(str): Path to the checkpoint to load before training. If set to None, no checkpoint will be loaded.
+            max_length(int): Maximum input length to consider. If the input length exceeds this value, it will be truncated.
+            fix_token_embedding(bool): Whether to fix the token embeddings during training. If set to True, the token embeddings will not be updated during training.
         """
         self._optimizer = optimizer
         self._lr_scheduler = lr_scheduler
@@ -98,18 +105,28 @@ class FullFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
         self.gradient_clip_algorithm = gradient_clip_algorithm
         self.save_optimizer_state = save_optimizer_state
         self.save_full_model = save_full_model
+        self.save_ckpt_type = save_ckpt_type
         self.ckpt_path = ckpt_path
+        self.max_length = max_length
+        self.fix_token_embedding = fix_token_embedding
         super().__init__(**kwargs)
 
     def run(self, modelpool: CausalLMPool):
         self.modelpool = modelpool
         self.setup()
-        self.train()
+        self.train(self.model, self.optimizer, self.lr_scheduler)
         return self.model
 
     def setup_model(self):
+        self.tokenizer = self.modelpool.load_tokenizer()
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+
         model = self.modelpool.load_pretrained_model()
-        self.model = model
+        self.model: "LlamaForCausalLM" = model
+
+        if self.fix_token_embedding:
+            self.model.model.embed_tokens.requires_grad_(False)
 
         if self.fabric.strategy == "fsdp" or isinstance(
             self.fabric.strategy, FSDPStrategy
@@ -125,17 +142,7 @@ class FullFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
 
     def configure_optimizer(self):
         # compute expected total steps
-        self.expected_total_steps = []
-        if self.max_steps > 0:
-            self.expected_total_steps.append(self.max_steps)
-        if self.max_steps_per_epoch > 0 and self.max_epochs > 0:
-            self.expected_total_steps.append(self.max_steps_per_epoch * self.max_epochs)
-        if self.max_epochs > 0:
-            self.expected_total_steps.append(
-                len(self.train_dataloader) * self.max_epochs
-            )
-        self.expected_total_steps = min(self.expected_total_steps)
-        log.info(f"Expected total steps: {self.expected_total_steps}")
+        self.compute_expected_total_steps(self.train_dataloader)
 
         optimizer = instantiate(self._optimizer, self.model.parameters())
         if self._lr_scheduler is not None:
@@ -174,7 +181,9 @@ class FullFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
             train_dataset,
             **self.dataloader_kwargs,
             shuffle=True,
-            collate_fn=padded_collate_sft,
+            collate_fn=functools.partial(
+                padded_collate_sft, pad_token_id=self.tokenizer.pad_token_id
+            ),
         )
         self.train_dataloader = fabric.setup_dataloaders(self.train_dataloader)
 
@@ -190,25 +199,15 @@ class FullFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
         self.model, self.optimizer = fabric.setup(self.model, optimizer)
         self.lr_scheduler = lr_scheduler
 
-    def _clip_gradients_if_needed(self):
+    @override
+    def train_epoch(self, *args, **kwargs):
         fabric = self.fabric
 
-        if self.gradient_clip_val is not None:
-            if self.gradient_clip_algorithm == "value":
-                fabric.clip_gradients(self.model, clip_val=self.gradient_clip_val)
-            elif self.gradient_clip_algorithm == "norm":
-                fabric.clip_gradients(self.model, max_norm=self.gradient_clip_val)
-            else:
-                raise ValueError(
-                    f"Unknown gradient clip algorithm: {self.gradient_clip_algorithm}. Available options: 'value', 'norm'"
-                )
-
-    def train_epoch(self):
-        fabric = self.fabric
+        accumulated_loss = 0
         for step_idx, batch in enumerate(
             pbar := tqdm(
                 self.train_dataloader,
-                desc="Training Steps",
+                desc="Training Batches",
                 dynamic_ncols=True,
                 leave=False,
                 disable=not fabric.is_global_zero,
@@ -216,6 +215,14 @@ class FullFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
         ):
             is_accumulating = (step_idx + 1) % self.accumulate_grad_batches != 0
 
+            if self.max_length > 0 and batch["input_ids"].shape[1] > self.max_length:
+                log.warning(
+                    f"Input length exceeds max_length: {batch['input_ids'].shape[1]} > {self.max_length}. Truncating input."
+                )
+                batch["input_ids"] = batch["input_ids"][:, : self.max_length]
+                batch["attention_mask"] = batch["attention_mask"][:, : self.max_length]
+                batch["labels"] = batch["labels"][:, : self.max_length]
+
             # disable gradient synchronization if accumulating gradients across steps for improved performance
             with fabric.no_backward_sync(self.model, enabled=is_accumulating):
                 # use_cache=True is not compatible with gradient checkpointing, so we disable it here
@@ -225,20 +232,13 @@ class FullFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
                     labels=batch["labels"],
                     use_cache=self.use_cache,
                 )
-                loss = output["loss"]
+                loss = output["loss"] / self.accumulate_grad_batches
 
                 fabric.backward(loss)
-
-            metrics = {
-                "train/loss": loss.item(),
-                "train/epoch_idx": self.epoch_idx,
-                "train/lr": self.optimizer.param_groups[0]["lr"],
-            }
-            fabric.log_dict(metrics, step=self.global_step_idx)
-            pbar.set_postfix(metrics)
+                accumulated_loss += loss.item()
 
             if not is_accumulating:
-                self._clip_gradients_if_needed()
+                self.clip_gradients_if_needed(self.model, self.optimizer)
 
                 # run lr_scheduler at the end of the step if interval is set to "step"
                 if (
@@ -251,105 +251,30 @@ class FullFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
                 self.optimizer.step()
                 self.optimizer.zero_grad()
 
-            # save the model at the end of the step if interval is set to "step" and frequency is met
-            self._try_save_checkpoint(stage="end_of_step")
+                metrics = {
+                    "train/loss": accumulated_loss,
+                    "train/epoch_idx": self.epoch_idx,
+                    "train/lr": self.optimizer.param_groups[0]["lr"],
+                }
+                fabric.log_dict(metrics, step=self.global_step_idx)
+                pbar.set_postfix(metrics)
 
-            # break if max_steps_per_epoch is set, and exit epoch
-            if (
-                self.max_steps_per_epoch > 0
-                and step_idx + 1 >= self.max_steps_per_epoch
-            ):
-                break
-            # break if max_steps is set, and exit training
-            if self.max_steps > 0 and self.global_step_idx >= self.max_steps - 1:
-                self.is_training = False
-                break
+                # save the model at the end of the step if interval is set to "step" and frequency is met
+                self.conditional_checkpoint_save(stage="end_of_step")
 
-            self.global_step_idx += 1
+                # break if max_steps_per_epoch is set, and exit epoch
+                if (
+                    self.max_steps_per_epoch > 0
+                    and step_idx + 1 >= self.max_steps_per_epoch
+                ):
+                    break
+                # break if max_steps is set, and exit training
+                if self.max_steps > 0 and self.global_step_idx >= self.max_steps - 1:
+                    self.is_training = False
+                    break
 
-    def train(self):
-        fabric = self.fabric
-        self.is_training = True
-        self.global_step_idx = 0
-        self.model.train()
-        for epoch_idx in tqdm(
-            range(self.max_epochs) if self.max_epochs > 0 else itertools.count(0),
-            "Training Epoch",
-            dynamic_ncols=True,
-            leave=False,
-            disable=not fabric.is_global_zero,
-        ):
-            self.epoch_idx = epoch_idx
-            self.train_epoch()
-            # run lr_scheduler at the end of the epoch if interval is set to "epoch"
-            if (
-                self.lr_scheduler_interval == "epoch"
-                and (epoch_idx + 1) % self.lr_scheduler_frequency == 0
-            ):
-                self.lr_scheduler.step()
-
-            # save the model at the end of the epoch if interval is set to "epoch" and frequency is met
-            self._try_save_checkpoint(stage="end_of_epoch")
-
-            if not self.is_training:
-                break
-
-        # save the model at the end of training
-        self._try_save_checkpoint(stage="end_of_training")
-
-    def _try_save_checkpoint(
-        self, stage: Literal["end_of_step", "end_of_epoch", "end_of_training"]
-    ):
-        if stage == "end_of_step":
-            if (
-                self.checkpoint_save_interval == "step"
-                and (self.global_step_idx + 1) % self.checkpoint_save_frequency == 0
-            ):
-                self.save_checkpoint(
-                    os.path.join(
-                        self.log_dir, "checkpoints", f"step={self.global_step_idx}.ckpt"
-                    )
-                )
-        elif stage == "end_of_epoch":
-            if (
-                self.checkpoint_save_interval == "epoch"
-                and (self.epoch_idx + 1) % self.checkpoint_save_frequency == 0
-            ):
-                self.save_checkpoint(
-                    os.path.join(
-                        self.log_dir, "checkpoints", f"epoch={self.epoch_idx}.ckpt"
-                    )
-                )
-        elif stage == "end_of_training":
-            # if the checkpoint has not been saved yet, save it
-            if self.global_step_idx > self._latest_saved_checkpoint_global_step:
-                self.save_checkpoint(
-                    os.path.join(
-                        self.log_dir,
-                        "checkpoints",
-                        f"epoch={self.epoch_idx}_step={self.global_step_idx}.ckpt",
-                    )
-                )
-                try:
-                    os.symlink(
-                        os.path.join(
-                            self.log_dir,
-                            "checkpoints",
-                            "latest_model.ckpt",
-                        ),
-                        dst := os.path.join(
-                            self.log_dir,
-                            "checkpoints",
-                            f"epoch={self.epoch_idx}_step={self.global_step_idx}.ckpt",
-                        ),
-                        target_is_directory=os.path.isdir(dst),
-                    )
-                except Exception as e:
-                    log.error(f"Failed to create symlink: {e}")
-        else:
-            raise ValueError(
-                f"Unknown stage: {stage}. Available options: 'end_of_step', 'end_of_epoch', 'end_of_training'"
-            )
+                self.global_step_idx += 1
+                accumulated_loss = 0
 
     def save_checkpoint(
         self,
@@ -361,31 +286,36 @@ class FullFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
             return log.warning(f"Checkpoint already exists at {path}. Skipping save.")
 
         fabric = self.fabric
-        state = {"model": self.model}
 
-        # save the optimizer and lr_scheduler state if needed
-        if self.save_optimizer_state and save_optimizer_state is not False:
-            state.update(
-                {
-                    "optimizer": self.optimizer,
-                    "lr_scheduler": self.lr_scheduler,
-                    "global_step_idx": self.global_step_idx,
-                    "epoch_idx": self.epoch_idx,
-                }
+        if self.save_ckpt_type == "lightning":
+            state = {"model": self.model}
+
+            # save the optimizer and lr_scheduler state if needed
+            if self.save_optimizer_state and save_optimizer_state is not False:
+                state.update(
+                    {
+                        "optimizer": self.optimizer,
+                        "lr_scheduler": self.lr_scheduler,
+                        "global_step_idx": self.global_step_idx,
+                        "epoch_idx": self.epoch_idx,
+                    }
+                )
+
+            trainable_param_names = set(
+                name
+                for name, param in self.model.state_dict(keep_vars=True).items()
+                if param.requires_grad
+            )
+            filter = (
+                None
+                if self.save_full_model
+                else {"model": lambda k, p: k in trainable_param_names}
             )
 
-        trainable_param_names = set(
-            name
-            for name, param in self.model.state_dict(keep_vars=True).items()
-            if param.requires_grad
-        )
-        filter = (
-            None
-            if self.save_full_model
-            else {"model": lambda k, p: k in trainable_param_names}
-        )
+            fabric.save(path, state=state, filter=filter)
+        else:
+            self.model.save_pretrained(path, is_main_process=fabric.is_global_zero)
 
-        fabric.save(path, state=state, filter=filter)
         self._latest_saved_checkpoint_global_step = self.global_step_idx
 
     def load_checkpoint(self, path: Union[str, Path]):
@@ -425,9 +355,9 @@ if __name__ == "__main__":
     import argparse
 
     parser = argparse.ArgumentParser()
-    parser.add_argument("--base_model_path", type=str)
-    parser.add_argument("--ckpt_path", type=str)
-    parser.add_argument("--output_path", type=str)
+    parser.add_argument("--base-model-path", type=str)
+    parser.add_argument("--ckpt-path", type=str)
+    parser.add_argument("--output-path", type=str)
 
     args = parser.parse_args()
 

fusion_bench/method/lm_finetune/peftfinetune_sft.py

@@ -1,3 +1,4 @@
+import functools
 import itertools
 import logging
 import os
@@ -11,7 +12,7 @@ import torch
 from lightning.fabric.strategies.fsdp import FSDPStrategy
 from lightning.fabric.utilities import rank_zero_only
 from omegaconf import DictConfig, OmegaConf
-from peft import PeftModel, get_peft_config, get_peft_model
+from peft import LoraConfig, PeftModel, get_peft_config, get_peft_model
 from torch import nn
 from torch.utils.data import ConcatDataset, DataLoader
 from tqdm.auto import tqdm
@@ -19,7 +20,7 @@ from typing_extensions import TYPE_CHECKING, override
 
 from fusion_bench import BaseAlgorithm, BaseModelPool
 from fusion_bench.dataset.llama.collate import padded_collate_sft
-from fusion_bench.mixins import LightningFabricMixin
+from fusion_bench.mixins import FabricTrainingMixin
 from fusion_bench.modelpool import CausalLMPool
 from fusion_bench.utils import instantiate
 from fusion_bench.utils.dtype import get_dtype
@@ -35,7 +36,7 @@ if TYPE_CHECKING:
 log = logging.getLogger(__name__)
 
 
-class PeftFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
+class PeftFinetuneSFT(BaseAlgorithm, FabricTrainingMixin):
 
     model: Union[
         nn.Module, "_FabricModule", "LlamaForCausalLM", PeftModel, peft.LoraModel
@@ -67,7 +68,7 @@ class PeftFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
         save_full_model: bool = False,
         save_ckpt_type: Literal["lightning", "peft"] = "peft",
         ckpt_path: Optional[str] = None,
-        max_length: int = 6150,
+        max_length: int = 6144,
         **kwargs,
     ):
         """
@@ -121,17 +122,23 @@ class PeftFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
     def run(self, modelpool: CausalLMPool):
         self.modelpool = modelpool
         self.setup()
-        self.train()
+        self.train(self.model, self.optimizer, self.lr_scheduler)
 
         if self.merge_and_unload:
             self.model = self.model.merge_and_unload()
         return self.model
 
     def setup_model(self):
+        # https://github.com/Lightning-AI/litgpt/blob/main/litgpt/finetune/lora.py
+        self.tokenizer = self.modelpool.load_tokenizer()
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+
         model = self.modelpool.load_pretrained_model()
 
         # get the PEFT model
         peft_config = instantiate(self._peft_config, _convert_="all")
+        peft_config.save_pretrained(os.path.join(self.log_dir, "peft_config"))
         peft_model = get_peft_model(model, peft_config, self.adapter_name)
         peft_model.print_trainable_parameters()
 
@@ -149,20 +156,11 @@ class PeftFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
             self.use_cache = True
 
         self.model_dtype = get_dtype(self.model)
+        self.model = self.model.to(dtype=self.model_dtype)
 
     def configure_optimizer(self):
         # compute expected total steps
-        self.expected_total_steps = []
-        if self.max_steps > 0:
-            self.expected_total_steps.append(self.max_steps)
-        if self.max_steps_per_epoch > 0 and self.max_epochs > 0:
-            self.expected_total_steps.append(self.max_steps_per_epoch * self.max_epochs)
-        if self.max_epochs > 0:
-            self.expected_total_steps.append(
-                len(self.train_dataloader) * self.max_epochs
-            )
-        self.expected_total_steps = min(self.expected_total_steps)
-        log.info(f"Expected total steps: {self.expected_total_steps}")
+        self.compute_expected_total_steps(self.train_dataloader)
 
         optimizer = instantiate(self._optimizer, self.model.parameters())
         if self._lr_scheduler is not None:
@@ -201,7 +199,9 @@ class PeftFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
             train_dataset,
             **self.dataloader_kwargs,
             shuffle=True,
-            collate_fn=padded_collate_sft,
+            collate_fn=functools.partial(
+                padded_collate_sft, pad_token_id=self.tokenizer.pad_token_id
+            ),
         )
         self.train_dataloader = fabric.setup_dataloaders(self.train_dataloader)
 
@@ -214,28 +214,19 @@ class PeftFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
         optimizer = self.configure_optimizer()
         optimizer, lr_scheduler = optimizer["optimizer"], optimizer["lr_scheduler"]
 
-        self.model, self.optimizer = fabric.setup(self.model, optimizer)
+        self.model = self.fabric.setup_module(self.model)
+        self.optimizer = self.fabric.setup_optimizers(optimizer)
         self.lr_scheduler = lr_scheduler
 
-    def _clip_gradients_if_needed(self):
+    @override
+    def train_epoch(self, *args, **kwargs):
         fabric = self.fabric
 
-        if self.gradient_clip_val is not None:
-            if self.gradient_clip_algorithm == "value":
-                fabric.clip_gradients(self.model, clip_val=self.gradient_clip_val)
-            elif self.gradient_clip_algorithm == "norm":
-                fabric.clip_gradients(self.model, max_norm=self.gradient_clip_val)
-            else:
-                raise ValueError(
-                    f"Unknown gradient clip algorithm: {self.gradient_clip_algorithm}. Available options: 'value', 'norm'"
-                )
-
-    def train_epoch(self):
-        fabric = self.fabric
+        accumulated_loss = 0
         for step_idx, batch in enumerate(
             pbar := tqdm(
                 self.train_dataloader,
-                desc="Training Steps",
+                desc="Training Batches",
                 dynamic_ncols=True,
                 leave=False,
                 disable=not fabric.is_global_zero,
@@ -250,6 +241,7 @@ class PeftFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
                 batch["input_ids"] = batch["input_ids"][:, : self.max_length]
                 batch["attention_mask"] = batch["attention_mask"][:, : self.max_length]
                 batch["labels"] = batch["labels"][:, : self.max_length]
+
             # disable gradient synchronization if accumulating gradients across steps for improved performance
             with fabric.no_backward_sync(self.model, enabled=is_accumulating):
                 # use_cache=True is not compatible with gradient checkpointing, so we disable it here
@@ -259,20 +251,13 @@ class PeftFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
                     labels=batch["labels"],
                     use_cache=self.use_cache,
                 )
-                loss = output["loss"]
+                loss = output["loss"] / self.accumulate_grad_batches
 
                 fabric.backward(loss)
-
-            metrics = {
-                "train/loss": loss.item(),
-                "train/epoch_idx": self.epoch_idx,
-                "train/lr": self.optimizer.param_groups[0]["lr"],
-            }
-            fabric.log_dict(metrics, step=self.global_step_idx)
-            pbar.set_postfix(metrics)
+                accumulated_loss += loss.item()
 
             if not is_accumulating:
-                self._clip_gradients_if_needed()
+                self.clip_gradients_if_needed(self.model, self.optimizer)
 
                 # run lr_scheduler at the end of the step if interval is set to "step"
                 if (
@@ -285,105 +270,30 @@ class PeftFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
                 self.optimizer.step()
                 self.optimizer.zero_grad()
 
-            # save the model at the end of the step if interval is set to "step" and frequency is met
-            self._try_save_checkpoint(stage="end_of_step")
+                metrics = {
+                    "train/loss": accumulated_loss,
+                    "train/epoch_idx": self.epoch_idx,
+                    "train/lr": self.optimizer.param_groups[0]["lr"],
+                }
+                fabric.log_dict(metrics, step=self.global_step_idx)
+                pbar.set_postfix(metrics)
 
-            # break if max_steps_per_epoch is set, and exit epoch
-            if (
-                self.max_steps_per_epoch > 0
-                and step_idx + 1 >= self.max_steps_per_epoch
-            ):
-                break
-            # break if max_steps is set, and exit training
-            if self.max_steps > 0 and self.global_step_idx >= self.max_steps - 1:
-                self.is_training = False
-                break
+                # save the model at the end of the step if interval is set to "step" and frequency is met
+                self.conditional_checkpoint_save(stage="end_of_step")
 
-            self.global_step_idx += 1
+                # break if max_steps_per_epoch is set, and exit epoch
+                if (
+                    self.max_steps_per_epoch > 0
+                    and step_idx + 1 >= self.max_steps_per_epoch
+                ):
+                    break
+                # break if max_steps is set, and exit training
+                if self.max_steps > 0 and self.global_step_idx >= self.max_steps - 1:
+                    self.is_training = False
+                    break
 
-    def train(self):
-        fabric = self.fabric
-        self.is_training = True
-        self.global_step_idx = 0
-        self.model.train()
-        for epoch_idx in tqdm(
-            range(self.max_epochs) if self.max_epochs > 0 else itertools.count(0),
-            "Training Epoch",
-            dynamic_ncols=True,
-            leave=False,
-            disable=not fabric.is_global_zero,
-        ):
-            self.epoch_idx = epoch_idx
-            self.train_epoch()
-            # run lr_scheduler at the end of the epoch if interval is set to "epoch"
-            if (
-                self.lr_scheduler_interval == "epoch"
-                and (epoch_idx + 1) % self.lr_scheduler_frequency == 0
-            ):
-                self.lr_scheduler.step()
-
-            # save the model at the end of the epoch if interval is set to "epoch" and frequency is met
-            self._try_save_checkpoint(stage="end_of_epoch")
-
-            if not self.is_training:
-                break
-
-        # save the model at the end of training
-        self._try_save_checkpoint(stage="end_of_training")
-
-    def _try_save_checkpoint(
-        self, stage: Literal["end_of_step", "end_of_epoch", "end_of_training"]
-    ):
-        if stage == "end_of_step":
-            if (
-                self.checkpoint_save_interval == "step"
-                and (self.global_step_idx + 1) % self.checkpoint_save_frequency == 0
-            ):
-                self.save_checkpoint(
-                    os.path.join(
-                        self.log_dir, "checkpoints", f"step={self.global_step_idx}.ckpt"
-                    )
-                )
-        elif stage == "end_of_epoch":
-            if (
-                self.checkpoint_save_interval == "epoch"
-                and (self.epoch_idx + 1) % self.checkpoint_save_frequency == 0
-            ):
-                self.save_checkpoint(
-                    os.path.join(
-                        self.log_dir, "checkpoints", f"epoch={self.epoch_idx}.ckpt"
-                    )
-                )
-        elif stage == "end_of_training":
-            # if the checkpoint has not been saved yet, save it
-            if self.global_step_idx > self._latest_saved_checkpoint_global_step:
-                self.save_checkpoint(
-                    os.path.join(
-                        self.log_dir,
-                        "checkpoints",
-                        f"epoch={self.epoch_idx}_step={self.global_step_idx}.ckpt",
-                    )
-                )
-                try:
-                    os.symlink(
-                        os.path.join(
-                            self.log_dir,
-                            "checkpoints",
-                            "latest_model.ckpt",
-                        ),
-                        dst := os.path.join(
-                            self.log_dir,
-                            "checkpoints",
-                            f"epoch={self.epoch_idx}_step={self.global_step_idx}.ckpt",
-                        ),
-                        target_is_directory=os.path.isdir(dst),
-                    )
-                except Exception as e:
-                    log.error(f"Failed to create symlink: {e}")
-        else:
-            raise ValueError(
-                f"Unknown stage: {stage}. Available options: 'end_of_step', 'end_of_epoch', 'end_of_training'"
-            )
+                self.global_step_idx += 1
+                accumulated_loss = 0
 
     def save_checkpoint(
         self,
@@ -418,7 +328,7 @@ class PeftFinetuneSFT(BaseAlgorithm, LightningFabricMixin):
                 if self.save_full_model
                 else {"model": lambda k, p: k in trainable_param_names}
             )
-
+            os.makedirs(os.path.dirname(path), exist_ok=True)
             fabric.save(path, state=state, filter=filter)
         elif self.save_ckpt_type == "peft":
             self.model.save_pretrained(path, is_main_process=fabric.is_global_zero)

fusion_bench/method/pruning/llama_magnitude_prune.py

@@ -1,7 +1,7 @@
-from typing import Literal, Optional, Union
+from typing import Dict, Literal, Optional, Union
 
 import torch
-from torch import Dict, nn
+from torch import nn
 from tqdm.auto import tqdm
 from transformers import LlamaForCausalLM, LlamaModel