From c3aa9b49bd95b53a984532375ae40fbe1d39769a Mon Sep 17 00:00:00 2001 From: ver217 Date: Thu, 7 Sep 2023 20:16:59 +0800 Subject: [PATCH] [checkpointio] fix save hf config --- colossalai/checkpoint_io/utils.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/colossalai/checkpoint_io/utils.py b/colossalai/checkpoint_io/utils.py index 6dadaba3e64f..3441eca38ce7 100644 --- a/colossalai/checkpoint_io/utils.py +++ b/colossalai/checkpoint_io/utils.py @@ -11,8 +11,6 @@ import torch import torch.nn as nn from torch.optim import Optimizer -from transformers.modeling_utils import PreTrainedModel, get_parameter_dtype -from transformers.modeling_utils import unwrap_model as unwrap_huggingface_model from colossalai.interface import ModelWrapper, OptimizerWrapper from colossalai.nn.optimizer import ColossalaiOptimizer @@ -383,6 +381,11 @@ def save_config_file(model: nn.Module, checkpoint_path: str, is_master: bool = T checkpoint_path (str): Path to the checkpoint directory. is_master (bool): Whether current rank is main process. """ + try: + from transformers.modeling_utils import PreTrainedModel, get_parameter_dtype + from transformers.modeling_utils import unwrap_model as unwrap_huggingface_model + except ImportError: + return if not isinstance(model, PreTrainedModel): return