2023-12-21 09:01:01 +00:00
|
|
|
import enum
|
2023-07-11 10:03:13 +00:00
|
|
|
import logging
|
|
|
|
import os
|
2023-12-21 09:01:01 +00:00
|
|
|
import warnings
|
2023-06-05 07:58:31 +00:00
|
|
|
from functools import partial
|
2023-07-11 10:03:13 +00:00
|
|
|
from pathlib import Path
|
2023-09-05 07:04:02 +00:00
|
|
|
from types import MethodType
|
2023-10-31 07:19:37 +00:00
|
|
|
from typing import Callable, Dict, Iterator, List, Optional, Tuple
|
2023-04-26 06:37:25 +00:00
|
|
|
|
|
|
|
import torch
|
|
|
|
import torch.nn as nn
|
2023-12-21 09:01:01 +00:00
|
|
|
from torch.nn import Parameter
|
2023-04-26 06:37:25 +00:00
|
|
|
from torch.optim import Optimizer
|
|
|
|
from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
|
|
|
|
from torch.utils._pytree import tree_map
|
|
|
|
from torch.utils.data import DataLoader
|
|
|
|
|
2024-01-09 02:20:05 +00:00
|
|
|
from colossalai.accelerator import get_accelerator
|
2023-07-11 10:03:13 +00:00
|
|
|
from colossalai.checkpoint_io import CheckpointIndexFile, CheckpointIO
|
|
|
|
from colossalai.checkpoint_io.utils import (
|
|
|
|
get_optimizer_base_filenames,
|
|
|
|
get_shard_filename,
|
2023-09-04 03:26:45 +00:00
|
|
|
load_param_groups_into_optimizer,
|
|
|
|
load_shard_state_dict,
|
|
|
|
load_states_into_optimizer,
|
2023-07-11 10:03:13 +00:00
|
|
|
save_param_groups,
|
|
|
|
save_state_dict,
|
2023-09-04 03:26:45 +00:00
|
|
|
sharded_optimizer_loading_epilogue,
|
2023-07-11 10:03:13 +00:00
|
|
|
)
|
2023-09-05 07:04:02 +00:00
|
|
|
from colossalai.interface import AMPModelMixin, ModelWrapper, OptimizerWrapper
|
2024-04-17 07:03:31 +00:00
|
|
|
from colossalai.quantization import BnbQuantizationConfig, quantize_model
|
2023-09-05 07:04:02 +00:00
|
|
|
from colossalai.zero import LowLevelZeroOptimizer
|
2023-04-26 06:37:25 +00:00
|
|
|
|
2023-05-05 11:36:10 +00:00
|
|
|
from .dp_plugin_base import DPPluginBase
|
2023-04-26 06:37:25 +00:00
|
|
|
from .torch_ddp_plugin import TorchDDPCheckpointIO
|
|
|
|
|
2023-09-19 06:20:26 +00:00
|
|
|
__all__ = ["LowLevelZeroPlugin"]
|
2023-04-26 06:37:25 +00:00
|
|
|
|
|
|
|
|
2023-06-05 07:58:31 +00:00
|
|
|
def _convert_floating_point(x, dtype: torch.dtype = torch.float16):
|
2023-04-26 06:37:25 +00:00
|
|
|
if isinstance(x, torch.Tensor) and torch.is_floating_point(x):
|
2023-06-05 07:58:31 +00:00
|
|
|
return x.to(dtype)
|
2023-04-26 06:37:25 +00:00
|
|
|
return x
|
|
|
|
|
|
|
|
|
2023-09-19 06:20:26 +00:00
|
|
|
SUPPORTED_PRECISION = ["fp16", "bf16", "fp32"]
|
2023-06-05 07:58:31 +00:00
|
|
|
|
|
|
|
|
2023-12-21 09:01:01 +00:00
|
|
|
class OptimizerParamCheckState(enum.Enum):
|
|
|
|
ORIGIN_PARAM_FINDED = 0
|
|
|
|
ORIGIN_PARAM_NOT_FIND = -1
|
|
|
|
LORA_PARM_EXISTED = -2
|
|
|
|
|
|
|
|
|
2023-09-05 07:04:02 +00:00
|
|
|
class LowLevelZeroModel(ModelWrapper, AMPModelMixin):
|
|
|
|
def __init__(self, module: nn.Module, precision: str) -> None:
|
|
|
|
super().__init__(module)
|
|
|
|
self.dtype = None
|
2023-09-19 06:20:26 +00:00
|
|
|
if precision == "fp16":
|
2023-09-05 07:04:02 +00:00
|
|
|
self.dtype = torch.float16
|
2023-09-19 06:20:26 +00:00
|
|
|
elif precision == "bf16":
|
2023-09-05 07:04:02 +00:00
|
|
|
self.dtype = torch.bfloat16
|
|
|
|
if self.dtype is not None:
|
|
|
|
module = module.to(self.dtype)
|
2024-01-09 02:20:05 +00:00
|
|
|
module = module.to(get_accelerator().get_current_device())
|
2023-09-05 07:04:02 +00:00
|
|
|
self.module = module
|
|
|
|
self.convert_fn = None
|
|
|
|
if self.dtype is not None:
|
|
|
|
self.convert_fn = partial(_convert_floating_point, dtype=self.dtype)
|
|
|
|
|
|
|
|
def forward(self, *args, **kwargs):
|
|
|
|
if self.convert_fn is not None:
|
|
|
|
args = tree_map(self.convert_fn, args)
|
|
|
|
kwargs = tree_map(self.convert_fn, kwargs)
|
|
|
|
return super().forward(*args, **kwargs)
|
|
|
|
|
|
|
|
|
2023-04-26 06:37:25 +00:00
|
|
|
class LowLevelZeroCheckpointIO(TorchDDPCheckpointIO):
|
2023-07-11 10:03:13 +00:00
|
|
|
def save_unsharded_optimizer(self, optimizer: OptimizerWrapper, checkpoint: str, gather_dtensor: bool = False):
|
|
|
|
"""Save optimizer to checkpoint but only on master process.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
optimizer (OptimizerWrapper): Optimizer to save state_dict
|
|
|
|
checkpoint (str): Path to save checkpoint
|
|
|
|
gather_dtensor (bool): Whether to gather_dtensor, not used
|
2023-04-26 06:37:25 +00:00
|
|
|
"""
|
2023-09-20 10:29:37 +00:00
|
|
|
assert isinstance(optimizer, LowLevelZeroOptimizer), "Please boost the optimizer before saving!"
|
2023-07-11 10:03:13 +00:00
|
|
|
# the `state_dict` in LowLevelZeroOptimizer has communication
|
|
|
|
# if only the master rank collect state_dict and save,
|
|
|
|
# the communication on each rank would not match
|
|
|
|
state_dict = optimizer.state_dict()
|
|
|
|
if self.coordinator.is_master():
|
|
|
|
save_state_dict(state_dict, checkpoint, use_safetensors=False)
|
|
|
|
|
2023-09-19 06:20:26 +00:00
|
|
|
def save_sharded_optimizer(
|
|
|
|
self,
|
|
|
|
optimizer: OptimizerWrapper,
|
|
|
|
checkpoint: str,
|
|
|
|
gather_dtensor: bool = False,
|
|
|
|
prefix: str = None,
|
|
|
|
size_per_shard: int = 1024,
|
|
|
|
):
|
2023-07-11 10:03:13 +00:00
|
|
|
"""
|
|
|
|
Save sharded Zero-optimizer checkpoint under the given checkpointing path.
|
|
|
|
The following files will be created under the path:
|
|
|
|
- An index file (pytorch_optim.bin.index.json) containing a map between optimizer states and file names
|
|
|
|
- A group file (pytorch_optim_group.bin) recording information of param_groups
|
|
|
|
- Multiple files (pytorch_optim-000XX.bin) that store state tensors of optimizer in a sharding way
|
|
|
|
|
|
|
|
Args:
|
|
|
|
optimizer (OptimizerWrapper): Optimizer to save sharded state_dict
|
|
|
|
checkpoint (str): Path to save optimizer state_dict
|
|
|
|
gather_dtensor (bool): Whether to gather_dtensor, not used
|
|
|
|
prefix (str): Perfix of file to save
|
|
|
|
size_per_shard (int): Max file size of each file that store state tensors
|
|
|
|
"""
|
2023-09-20 10:29:37 +00:00
|
|
|
assert isinstance(optimizer, LowLevelZeroOptimizer), "Please boost the optimizer before saving!"
|
2023-07-11 10:03:13 +00:00
|
|
|
if os.path.isfile(checkpoint):
|
|
|
|
logging.error(f"Provided path ({checkpoint}) should be a directory, not a file")
|
|
|
|
return
|
|
|
|
|
|
|
|
Path(checkpoint).mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
|
|
# state_dict only provide only 'param_groups'
|
|
|
|
state_dict = optimizer.optim.state_dict()
|
|
|
|
# state shard would be handled by the low-level zero optimizer
|
|
|
|
sharded_state = optimizer.state_dict_shard(max_shard_size=size_per_shard)
|
|
|
|
|
|
|
|
# Preparing file paths and index file.
|
|
|
|
states_name, save_index_file, param_group_file = get_optimizer_base_filenames(prefix)
|
|
|
|
index_file = CheckpointIndexFile(checkpoint)
|
2023-10-31 06:48:01 +00:00
|
|
|
index_file.append_meta_data("param_groups", param_group_file)
|
2023-07-11 10:03:13 +00:00
|
|
|
|
|
|
|
# Store the information of param groups to param_group_file.
|
2023-10-31 06:48:01 +00:00
|
|
|
if self.coordinator.is_master():
|
|
|
|
group_file_path = os.path.join(checkpoint, param_group_file)
|
|
|
|
save_param_groups(state_dict, group_file_path)
|
2023-07-11 10:03:13 +00:00
|
|
|
|
|
|
|
# Save shards of optimizer states.
|
|
|
|
total_size = 0
|
|
|
|
for idx, shard_pair in enumerate(sharded_state):
|
|
|
|
shard, current_size = shard_pair
|
|
|
|
shard_file = get_shard_filename(states_name, idx)
|
|
|
|
total_size = total_size + current_size
|
|
|
|
for param_id in shard.keys():
|
|
|
|
index_file.append_weight_map(str(param_id), shard_file)
|
|
|
|
|
|
|
|
checkpoint_file_path = os.path.join(checkpoint, shard_file)
|
|
|
|
if self.coordinator.is_master():
|
|
|
|
save_state_dict(shard, checkpoint_file_path, use_safetensors=False)
|
|
|
|
|
|
|
|
# Wrap up index file.
|
|
|
|
index_file.append_meta_data("total_size", total_size)
|
|
|
|
if self.coordinator.is_master():
|
|
|
|
index_file.write_index_file(save_index_file)
|
2023-09-19 06:20:26 +00:00
|
|
|
logging.info(
|
|
|
|
f"The optimizer is going to be split to checkpoint shards. "
|
|
|
|
f"You can find where each parameters has been saved in the "
|
|
|
|
f"index located at {save_index_file}."
|
|
|
|
)
|
2023-07-11 10:03:13 +00:00
|
|
|
|
|
|
|
def load_sharded_optimizer(self, optimizer: OptimizerWrapper, index_file_path: str, prefix: str):
|
|
|
|
"""Load sharded optimizer with the given path to index file.
|
|
|
|
|
|
|
|
Args:
|
|
|
|
optimizer (OptimizerWrapper): Optimizer to load state_dict
|
|
|
|
index_file_path (str): Path to the index file
|
|
|
|
prefix (str): Not used.
|
|
|
|
"""
|
2023-09-20 10:29:37 +00:00
|
|
|
assert isinstance(optimizer, LowLevelZeroOptimizer), "Please boost the optimizer before Loading!"
|
|
|
|
optimizer = optimizer.unwrap()
|
2023-09-04 03:26:45 +00:00
|
|
|
|
|
|
|
# Read checkpoint index file.
|
|
|
|
ckpt_index_file = CheckpointIndexFile.from_file(index_file_path)
|
|
|
|
|
|
|
|
# Load param_groups
|
|
|
|
param_group_path = ckpt_index_file.get_param_group_filename()
|
|
|
|
if param_group_path is None:
|
2023-09-19 06:20:26 +00:00
|
|
|
raise RuntimeError(
|
|
|
|
f"Invalid index file path {index_file_path} for an optimizer. \
|
|
|
|
Lacking param group file under current directory."
|
|
|
|
)
|
2023-09-04 03:26:45 +00:00
|
|
|
id_map = load_param_groups_into_optimizer(optimizer, param_group_path)
|
|
|
|
|
|
|
|
checkpoint_files, _ = ckpt_index_file.get_checkpoint_filenames()
|
|
|
|
|
|
|
|
for shard_file in checkpoint_files:
|
|
|
|
state_dict = load_shard_state_dict(Path(shard_file), use_safetensors=False)
|
|
|
|
# shard state dict
|
|
|
|
for param_idx, state in state_dict.items():
|
|
|
|
for k, v in state.items():
|
2023-09-19 06:20:26 +00:00
|
|
|
if isinstance(v, torch.Tensor) and k != "step":
|
|
|
|
padding_size = (
|
|
|
|
self.coordinator.world_size - v.numel() % self.coordinator.world_size
|
|
|
|
) % self.coordinator.world_size
|
2023-09-04 03:26:45 +00:00
|
|
|
with torch.no_grad():
|
|
|
|
v = v.flatten()
|
|
|
|
if padding_size > 0:
|
|
|
|
v = torch.nn.functional.pad(v, [0, padding_size])
|
|
|
|
v_list = v.split(v.numel() // self.coordinator.world_size)
|
|
|
|
state_dict[param_idx][k] = v_list[self.coordinator.rank].detach().clone()
|
|
|
|
load_states_into_optimizer(optimizer, state_dict, id_map)
|
|
|
|
sharded_optimizer_loading_epilogue(optimizer)
|
2023-04-26 06:37:25 +00:00
|
|
|
|
2023-09-20 10:29:37 +00:00
|
|
|
def load_unsharded_model(self, model: ModelWrapper, checkpoint: str, strict: bool = True):
|
|
|
|
assert isinstance(model, LowLevelZeroModel), "Please boost the model before loading!"
|
|
|
|
super().load_unsharded_model(model, checkpoint, strict)
|
2023-09-05 07:04:02 +00:00
|
|
|
model.update_master_params()
|
|
|
|
|
2023-09-19 06:20:26 +00:00
|
|
|
def load_sharded_model(
|
|
|
|
self,
|
2023-09-20 10:29:37 +00:00
|
|
|
model: ModelWrapper,
|
2023-09-19 06:20:26 +00:00
|
|
|
checkpoint_index_file: Path,
|
|
|
|
strict: bool = False,
|
|
|
|
use_safetensors: bool = False,
|
|
|
|
load_sub_module: bool = True,
|
|
|
|
):
|
2023-09-20 10:29:37 +00:00
|
|
|
assert isinstance(model, LowLevelZeroModel), "Please boost the model before loading!"
|
|
|
|
super().load_sharded_model(model, checkpoint_index_file, strict, use_safetensors, load_sub_module)
|
2023-09-05 07:04:02 +00:00
|
|
|
model.update_master_params()
|
2023-04-26 06:37:25 +00:00
|
|
|
|
2023-12-21 09:01:01 +00:00
|
|
|
def save_lora_as_pretrained(self, model, checkpoint, use_safetensors):
|
|
|
|
if os.path.isfile(checkpoint):
|
|
|
|
logging.error(f"Provided path ({checkpoint}) should be a directory, not a file")
|
|
|
|
return
|
|
|
|
from peft import PeftModel
|
|
|
|
|
|
|
|
assert isinstance(model, ModelWrapper), "Please boost the model before saving!"
|
|
|
|
peft_model = model.unwrap()
|
|
|
|
assert isinstance(
|
|
|
|
peft_model, PeftModel
|
|
|
|
), "The model doesn't have lora adapters, please enable lora before saving."
|
|
|
|
return peft_model.save_pretrained(checkpoint, safe_serialization=use_safetensors)
|
|
|
|
|
2023-04-26 06:37:25 +00:00
|
|
|
|
2023-05-05 11:36:10 +00:00
|
|
|
class LowLevelZeroPlugin(DPPluginBase):
|
2023-04-26 06:37:25 +00:00
|
|
|
"""
|
|
|
|
Plugin for low level zero.
|
|
|
|
|
2023-09-26 02:57:47 +00:00
|
|
|
```python
|
|
|
|
from colossalai.booster import Booster
|
|
|
|
from colossalai.booster.plugin import LowLevelZeroPlugin
|
|
|
|
|
|
|
|
model, train_dataset, optimizer, criterion = ...
|
|
|
|
plugin = LowLevelZeroPlugin()
|
|
|
|
|
|
|
|
train_dataloader = plugin.prepare_dataloader(train_dataset, batch_size=8)
|
|
|
|
booster = Booster(plugin=plugin)
|
|
|
|
model, optimizer, train_dataloader, criterion = booster.boost(model, optimizer, train_dataloader, criterion)
|
|
|
|
```
|
2023-04-26 06:37:25 +00:00
|
|
|
|
|
|
|
Args:
|
2023-10-05 16:30:34 +00:00
|
|
|
stage (int, optional): ZeRO stage. Defaults to 1.
|
2023-06-05 07:58:31 +00:00
|
|
|
precision (str, optional): precision. Support 'fp16', 'bf16' and 'fp32'. Defaults to 'fp16'.
|
2023-04-26 06:37:25 +00:00
|
|
|
initial_scale (float, optional): Initial scale used by DynamicGradScaler. Defaults to 2**32.
|
|
|
|
min_scale (float, optional): Min scale used by DynamicGradScaler. Defaults to 1.
|
|
|
|
growth_factor (float, optional): growth_factor used by DynamicGradScaler. Defaults to 2.
|
|
|
|
backoff_factor (float, optional): backoff_factor used by DynamicGradScaler. Defaults to 0.5.
|
|
|
|
growth_interval (float, optional): growth_interval used by DynamicGradScaler. Defaults to 1000.
|
|
|
|
hysteresis (float, optional): hysteresis used by DynamicGradScaler. Defaults to 2.
|
|
|
|
max_scale (int, optional): max_scale used by DynamicGradScaler. Defaults to 2**32.
|
|
|
|
max_norm (float, optional): max_norm used for `clip_grad_norm`. You should notice that you shall not do
|
|
|
|
clip_grad_norm by yourself when using ZeRO DDP. The ZeRO optimizer will take care of clip_grad_norm.
|
|
|
|
norm_type (float, optional): norm_type used for `clip_grad_norm`.
|
|
|
|
reduce_bucket_size_in_m (int, optional): grad reduce bucket size in M. Defaults to 12.
|
|
|
|
communication_dtype (torch.dtype, optional): communication dtype. If not specified, the dtype of param will be used. Defaults to None.
|
|
|
|
overlap_communication (bool, optional): whether to overlap communication and computation. Defaults to True.
|
|
|
|
cpu_offload (bool, optional): whether to offload grad, master weight and optimizer state to cpu. Defaults to False.
|
|
|
|
verbose (bool, optional): verbose mode. Debug info including grad overflow will be printed. Defaults to False.
|
|
|
|
"""
|
|
|
|
|
|
|
|
def __init__(
|
|
|
|
self,
|
|
|
|
stage: int = 1,
|
2023-09-19 06:20:26 +00:00
|
|
|
precision: str = "fp16",
|
2023-04-26 06:37:25 +00:00
|
|
|
initial_scale: float = 2**32,
|
|
|
|
min_scale: float = 1,
|
|
|
|
growth_factor: float = 2,
|
|
|
|
backoff_factor: float = 0.5,
|
|
|
|
growth_interval: int = 1000,
|
|
|
|
hysteresis: int = 2,
|
|
|
|
max_scale: float = 2**32,
|
|
|
|
max_norm: float = 0.0,
|
|
|
|
norm_type: float = 2.0,
|
|
|
|
reduce_bucket_size_in_m: int = 12,
|
|
|
|
communication_dtype: Optional[torch.dtype] = None,
|
|
|
|
overlap_communication: bool = True,
|
|
|
|
cpu_offload: bool = False,
|
2023-10-13 07:57:45 +00:00
|
|
|
master_weights: bool = True,
|
2023-04-26 06:37:25 +00:00
|
|
|
verbose: bool = False,
|
|
|
|
) -> None:
|
2023-05-05 11:36:10 +00:00
|
|
|
super().__init__()
|
2023-09-19 06:20:26 +00:00
|
|
|
assert stage in (1, 2), f"LowLevelZeroPlugin only supports stage 1/2 training"
|
|
|
|
assert precision in SUPPORTED_PRECISION, f"LowLevelZeroPlugin only supports {SUPPORTED_PRECISION} training"
|
|
|
|
assert norm_type == 2.0, f"LowLevelZeroPlugin only supports norm_type=2.0 now"
|
2023-04-26 06:37:25 +00:00
|
|
|
self.stage = stage
|
|
|
|
self.precision = precision
|
2023-09-05 07:04:02 +00:00
|
|
|
self.zero_optim_kwargs = dict(
|
|
|
|
initial_scale=initial_scale,
|
2023-10-13 07:57:45 +00:00
|
|
|
min_scale=min_scale,
|
2023-09-05 07:04:02 +00:00
|
|
|
growth_factor=growth_factor,
|
|
|
|
backoff_factor=backoff_factor,
|
|
|
|
growth_interval=growth_interval,
|
|
|
|
hysteresis=hysteresis,
|
|
|
|
max_scale=max_scale,
|
|
|
|
clip_grad_norm=max_norm,
|
|
|
|
reduce_bucket_size=reduce_bucket_size_in_m * 1024 * 1024,
|
|
|
|
communication_dtype=communication_dtype,
|
|
|
|
overlap_communication=overlap_communication,
|
|
|
|
partition_grad=(stage == 2),
|
2023-10-13 07:57:45 +00:00
|
|
|
cpu_offload=cpu_offload,
|
|
|
|
master_weights=master_weights,
|
2023-09-05 07:04:02 +00:00
|
|
|
)
|
2023-12-21 09:01:01 +00:00
|
|
|
self.lora_enabled = False
|
2023-04-26 06:37:25 +00:00
|
|
|
self.verbose = verbose
|
|
|
|
|
2023-07-04 04:00:33 +00:00
|
|
|
# set class name with stage, for better error message
|
|
|
|
setattr(self.__class__, "__name__", f"LowLevelZeroPlugin_ZeRO-{stage}")
|
|
|
|
|
2023-04-26 06:37:25 +00:00
|
|
|
def support_no_sync(self) -> bool:
|
2023-07-04 04:00:33 +00:00
|
|
|
return self.stage == 1
|
2023-04-26 06:37:25 +00:00
|
|
|
|
2023-10-31 07:19:37 +00:00
|
|
|
def support_lora(self) -> bool:
|
|
|
|
return False
|
|
|
|
|
2023-04-26 06:37:25 +00:00
|
|
|
def control_precision(self) -> bool:
|
|
|
|
return True
|
|
|
|
|
|
|
|
def supported_precisions(self) -> List[str]:
|
2023-06-05 07:58:31 +00:00
|
|
|
return SUPPORTED_PRECISION
|
2023-04-26 06:37:25 +00:00
|
|
|
|
|
|
|
def control_device(self) -> bool:
|
|
|
|
return True
|
|
|
|
|
|
|
|
def supported_devices(self) -> List[str]:
|
2023-11-20 08:12:41 +00:00
|
|
|
return ["cuda", "npu"]
|
2023-04-26 06:37:25 +00:00
|
|
|
|
2023-12-21 09:01:01 +00:00
|
|
|
def support_lora(self) -> bool:
|
|
|
|
return True
|
|
|
|
|
|
|
|
def enable_lora(
|
2024-04-17 07:03:31 +00:00
|
|
|
self,
|
|
|
|
model: nn.Module,
|
|
|
|
pretrained_dir: Optional[str] = None,
|
|
|
|
lora_config: Optional[Dict] = None,
|
|
|
|
bnb_quantization_config: Optional[BnbQuantizationConfig] = None,
|
2023-12-21 09:01:01 +00:00
|
|
|
) -> nn.Module:
|
|
|
|
from peft import PeftModel, get_peft_model
|
|
|
|
|
|
|
|
assert not isinstance(model, LowLevelZeroModel), "Lora should be enabled before boosting the model."
|
|
|
|
self.lora_enabled = True
|
|
|
|
warnings.warn("You have enabled LoRa training. Please check the hyperparameters such as lr")
|
|
|
|
|
2024-04-17 07:03:31 +00:00
|
|
|
if bnb_quantization_config is not None:
|
|
|
|
model = quantize_model(model, bnb_quantization_config)
|
|
|
|
|
2023-12-21 09:01:01 +00:00
|
|
|
if pretrained_dir is None:
|
|
|
|
peft_model = get_peft_model(model, lora_config)
|
|
|
|
else:
|
|
|
|
peft_model = PeftModel.from_pretrained(model, pretrained_dir, is_trainable=True)
|
|
|
|
return peft_model
|
|
|
|
|
|
|
|
def get_param_group_id(self, optimizer: Optimizer, origin_param: Parameter):
|
|
|
|
origin_param_id = id(origin_param)
|
|
|
|
for group_id, param_group in enumerate(optimizer.param_groups):
|
|
|
|
for p in param_group["params"]:
|
|
|
|
if id(p) == origin_param_id:
|
|
|
|
return group_id
|
|
|
|
return -1
|
|
|
|
|
|
|
|
def get_param_group_id(self, optimizer: Optimizer, origin_param: Parameter, lora_param: Parameter):
|
|
|
|
origin_param_id = id(origin_param)
|
|
|
|
lora_param_id = id(lora_param)
|
|
|
|
target_group_id = None
|
|
|
|
for group_id, param_group in enumerate(optimizer.param_groups):
|
|
|
|
for p in param_group["params"]:
|
|
|
|
if id(p) == lora_param_id:
|
|
|
|
# check if the lora parameter exists.
|
|
|
|
return target_group_id, OptimizerParamCheckState.LORA_PARM_EXISTED
|
|
|
|
if id(p) == origin_param_id:
|
|
|
|
target_group_id = group_id
|
|
|
|
if target_group_id is not None:
|
|
|
|
return target_group_id, OptimizerParamCheckState.ORIGIN_PARAM_FINDED
|
|
|
|
else:
|
|
|
|
return target_group_id, OptimizerParamCheckState.ORIGIN_PARAM_NOT_FIND
|
|
|
|
|
|
|
|
def add_lora_params_to_optimizer(self, model, optimizer):
|
|
|
|
"""add lora parameters to optimizer"""
|
|
|
|
name2param = {}
|
|
|
|
for name, param in model.named_parameters():
|
|
|
|
name2param[name] = param
|
|
|
|
|
|
|
|
for name, param in name2param.items():
|
|
|
|
if "lora_A" in name or "lora_B" in name:
|
|
|
|
origin_key = name.replace("lora_A.", "")
|
|
|
|
origin_key = origin_key.replace("lora_B.", "")
|
|
|
|
origin_key = origin_key.replace(f"{model.active_adapter}", "base_layer")
|
|
|
|
origin_param = name2param[origin_key]
|
|
|
|
group_id, check_state = self.get_param_group_id(optimizer, origin_param, param)
|
|
|
|
if check_state == OptimizerParamCheckState.ORIGIN_PARAM_NOT_FIND:
|
|
|
|
warnings.warn(
|
|
|
|
"Origin parameter {origin_key} related to {name} doesn't exist in optimizer param_groups."
|
|
|
|
)
|
|
|
|
elif (
|
|
|
|
check_state == OptimizerParamCheckState.ORIGIN_PARAM_FINDED
|
|
|
|
and group_id is not None
|
|
|
|
and group_id >= 0
|
|
|
|
):
|
|
|
|
optimizer.param_groups[group_id]["params"].append(param)
|
|
|
|
|
2023-04-26 06:37:25 +00:00
|
|
|
def configure(
|
|
|
|
self,
|
|
|
|
model: nn.Module,
|
2023-06-15 09:38:42 +00:00
|
|
|
optimizer: Optional[Optimizer] = None,
|
|
|
|
criterion: Optional[Callable] = None,
|
|
|
|
dataloader: Optional[DataLoader] = None,
|
|
|
|
lr_scheduler: Optional[LRScheduler] = None,
|
|
|
|
) -> Tuple[nn.Module, OptimizerWrapper, Callable, DataLoader, LRScheduler]:
|
2023-12-21 09:01:01 +00:00
|
|
|
if self.lora_enabled:
|
|
|
|
from peft import PeftModel
|
|
|
|
|
|
|
|
assert isinstance(
|
|
|
|
model, PeftModel
|
|
|
|
), "The model should have been wrapped as a PeftModel when self.lora_enabled is True"
|
|
|
|
if optimizer is not None:
|
|
|
|
self.add_lora_params_to_optimizer(model, optimizer)
|
|
|
|
|
2023-04-26 06:37:25 +00:00
|
|
|
if not isinstance(model, ModelWrapper):
|
2023-09-05 07:04:02 +00:00
|
|
|
model = LowLevelZeroModel(model, self.precision)
|
2023-04-26 06:37:25 +00:00
|
|
|
|
2023-09-19 06:20:26 +00:00
|
|
|
if optimizer is not None and not isinstance(optimizer, OptimizerWrapper):
|
|
|
|
optimizer: LowLevelZeroOptimizer = LowLevelZeroOptimizer(
|
|
|
|
optimizer, **self.zero_optim_kwargs, verbose=self.verbose
|
|
|
|
)
|
2023-09-05 07:04:02 +00:00
|
|
|
# inject update_master_params
|
|
|
|
model.update_master_params = MethodType(optimizer.update_master_params, model)
|
2023-04-26 06:37:25 +00:00
|
|
|
|
|
|
|
return model, optimizer, criterion, dataloader, lr_scheduler
|
|
|
|
|
|
|
|
def control_checkpoint_io(self) -> bool:
|
|
|
|
return True
|
|
|
|
|
|
|
|
def get_checkpoint_io(self) -> CheckpointIO:
|
|
|
|
return LowLevelZeroCheckpointIO()
|
2023-05-09 03:10:02 +00:00
|
|
|
|
2023-07-04 04:00:33 +00:00
|
|
|
def no_sync(self, model: nn.Module, optimizer: OptimizerWrapper) -> Iterator[None]:
|
|
|
|
assert isinstance(optimizer, LowLevelZeroOptimizer)
|
2023-10-17 06:07:21 +00:00
|
|
|
return optimizer.no_sync()
|