wehub-resource-sync
370 lines
13 KiB
Python
370 lines
13 KiB
Python
import glob
|
|
import os
|
|
import random
|
|
import re
|
|
from typing import Dict, List
|
|
|
|
import hydra
|
|
import numpy as np
|
|
import omegaconf
|
|
import torch
|
|
import torch.distributed as dist
|
|
from omegaconf import DictConfig
|
|
from torch.utils.data import ConcatDataset
|
|
from tqdm import tqdm
|
|
from transformers import PreTrainedTokenizer
|
|
|
|
from general_util.logger import get_child_logger
|
|
|
|
logger = get_child_logger("TrainingUtils")
|
|
|
|
|
|
def set_seed(args):
|
|
random.seed(args.seed)
|
|
np.random.seed(args.seed)
|
|
torch.manual_seed(args.seed)
|
|
if args.n_gpu > 0:
|
|
torch.cuda.manual_seed_all(args.seed)
|
|
|
|
|
|
def set_seed_int(seed):
|
|
random.seed(seed)
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
|
|
|
|
def get_rank():
|
|
if dist.is_initialized():
|
|
return dist.get_rank()
|
|
else:
|
|
return -1
|
|
|
|
|
|
def to_list(tensor):
|
|
return tensor.detach().cpu().tolist()
|
|
|
|
|
|
def unwrap_model(model: torch.nn.Module) -> torch.nn.Module:
|
|
"""
|
|
Recursively unwraps a model from potential containers (as used in distributed training).
|
|
Args:
|
|
model (:obj:`torch.nn.Module`): The model to unwrap.
|
|
"""
|
|
# since there could be multiple levels of wrapping, unwrap recursively
|
|
if hasattr(model, "module"):
|
|
return unwrap_model(model.module)
|
|
else:
|
|
return model
|
|
|
|
|
|
def get_zero_stage(cfg: DictConfig):
|
|
if hasattr(cfg, "zero_optimization"):
|
|
return int(getattr(cfg.zero_optimization, "stage", 0))
|
|
return 0
|
|
|
|
|
|
def return_torch_dtype(dtype: str):
|
|
if dtype == "float16":
|
|
return torch.float16
|
|
elif dtype == "bfloat16":
|
|
return torch.bfloat16
|
|
elif dtype == "float32":
|
|
return torch.float32
|
|
elif dtype == "int8":
|
|
return torch.int8
|
|
else:
|
|
return dtype
|
|
|
|
|
|
def batch_to_device(batch: Dict[str, torch.Tensor], device):
|
|
if "meta_data" in batch:
|
|
batch.pop("meta_data")
|
|
if "index" in batch:
|
|
batch.pop("index")
|
|
|
|
batch_on_device = {}
|
|
for k, v in batch.items():
|
|
if isinstance(v, torch.Tensor):
|
|
batch_on_device[k] = v.to(device)
|
|
else:
|
|
batch_on_device[k] = v
|
|
return batch_on_device
|
|
|
|
|
|
def initialize_dataset(cfg: DictConfig, file_path: str, tokenizer: PreTrainedTokenizer):
|
|
if "_target_" in cfg:
|
|
return hydra.utils.call(cfg, file_path=file_path, tokenizer=tokenizer)
|
|
else:
|
|
datasets = [initialize_dataset(cfg[key], file_path, tokenizer) for key in cfg.keys()]
|
|
assert len(datasets)
|
|
datasets = ConcatDataset(datasets)
|
|
return datasets
|
|
|
|
|
|
def load_and_cache_examples(cfg, tokenizer: PreTrainedTokenizer, _split="train", _file: str = None):
|
|
if_barrier = False
|
|
|
|
if _file is not None:
|
|
input_file = _file
|
|
if_barrier = True
|
|
else:
|
|
if _split == "train":
|
|
input_file = cfg.train_file
|
|
if_barrier = True
|
|
elif _split == "dev":
|
|
input_file = cfg.dev_file
|
|
if cfg.ddp_eval and cfg.local_rank != -1:
|
|
if_barrier = True
|
|
elif _split == "test":
|
|
input_file = cfg.test_file
|
|
if cfg.ddp_eval and cfg.local_rank != -1:
|
|
if_barrier = True
|
|
else:
|
|
raise RuntimeError(_split)
|
|
|
|
if getattr(cfg, "dist_load_data_barrier", True) and if_barrier and cfg.local_rank not in [-1, 0]:
|
|
dist.barrier() # Make sure only the first process in distributed training process the dataset, and the others will use the cache
|
|
|
|
logger.info("Creating features from dataset file at %s", input_file)
|
|
sub_config = f"read_tensor_{_split}"
|
|
if sub_config in cfg:
|
|
dataset = initialize_dataset(cfg[sub_config], file_path=input_file, tokenizer=tokenizer)
|
|
else:
|
|
dataset = initialize_dataset(cfg.read_tensor, file_path=input_file, tokenizer=tokenizer)
|
|
|
|
if getattr(cfg, "dist_load_data_barrier", True) and if_barrier and cfg.local_rank == 0:
|
|
dist.barrier() # Make sure only the first process in distributed training process the dataset, and the others will use the cache
|
|
|
|
return dataset
|
|
|
|
|
|
def organize_multiple_dataset(cfg, tokenizer: PreTrainedTokenizer, _split="train"):
|
|
if "_target_" in cfg.train_file:
|
|
files = hydra.utils.instantiate(cfg.train_file)
|
|
elif isinstance(cfg.train_file, omegaconf.ListConfig):
|
|
files = list(cfg.train_file)
|
|
elif cfg.train_file.startswith("hf:"):
|
|
files = [cfg.train_file[3:]]
|
|
elif cfg.train_file.startswith("list:"):
|
|
files = [cfg.train_file[5:]]
|
|
elif os.path.exists(cfg.train_file):
|
|
files = [cfg.train_file]
|
|
else:
|
|
files = list(glob.glob(cfg.train_file))
|
|
logger.info(files)
|
|
|
|
if getattr(cfg, "total_dataset_len", -1) > 0:
|
|
total_dataset_len = cfg.total_dataset_len
|
|
else:
|
|
total_dataset_len = 0
|
|
if dist.is_initialized() and dist.get_rank() != 0:
|
|
dist.barrier()
|
|
|
|
if not dist.is_initialized() or dist.get_rank() == 0:
|
|
for _file in tqdm(files, total=len(files)):
|
|
sub_train_dataset = load_and_cache_examples(cfg, tokenizer, _split="train", _file=_file)
|
|
total_dataset_len += len(sub_train_dataset)
|
|
del sub_train_dataset
|
|
|
|
if dist.is_initialized():
|
|
dist.barrier()
|
|
|
|
if dist.is_initialized():
|
|
if dist.get_rank() == 0:
|
|
objects = [total_dataset_len for _ in range(dist.get_world_size())]
|
|
else:
|
|
objects = [None for _ in range(dist.get_world_size())]
|
|
output_list = [None]
|
|
dist.scatter_object_list(output_list, objects, src=0)
|
|
if dist.get_rank() != 0:
|
|
total_dataset_len = output_list[0]
|
|
|
|
assert total_dataset_len > 0
|
|
|
|
logger.warning(f"Rank No. {cfg.local_rank} has {total_dataset_len} samples.")
|
|
cfg.total_dataset_len = total_dataset_len
|
|
|
|
return files, total_dataset_len
|
|
|
|
|
|
def if_cancel_sync(cfg: DictConfig, step: int):
|
|
if getattr(cfg, "forward_sync", False) is False and (
|
|
step + 1) % cfg.gradient_accumulation_steps != 0 and cfg.local_rank != -1:
|
|
return True
|
|
return False
|
|
|
|
|
|
def initialize_optimizer(cfg: DictConfig, grouped_parameters: List[Dict] = None, model: torch.nn.Module = None):
|
|
if grouped_parameters is None:
|
|
assert model is not None, "Either ``grouped_parameters`` or ``model`` must be specified."
|
|
no_decay = ['bias', 'LayerNorm.weight', 'layer_norm.weight']
|
|
grouped_parameters = [
|
|
{
|
|
'params': [p for n, p in model.named_parameters() if
|
|
(not any(nd in n for nd in no_decay)) and p.requires_grad],
|
|
'weight_decay': cfg.weight_decay
|
|
},
|
|
{
|
|
'params': [p for n, p in model.named_parameters() if
|
|
(any(nd in n for nd in no_decay)) and p.requires_grad],
|
|
'weight_decay': 0.0
|
|
}
|
|
]
|
|
|
|
if "optimizer" in cfg and cfg.optimizer and 'lamb' in cfg.optimizer:
|
|
if "bit_training" in cfg and cfg.bit_training:
|
|
from bitsandbytes.optim import LAMB8bit
|
|
|
|
optimizer = LAMB8bit(grouped_parameters,
|
|
lr=cfg.learning_rate,
|
|
betas=eval(cfg.adam_betas),
|
|
eps=cfg.adam_epsilon,
|
|
max_unorm=cfg.max_grad_norm)
|
|
else:
|
|
if cfg.optimizer == 'fused_lamb':
|
|
try:
|
|
from apex.optimizers.fused_mixed_precision_lamb import FusedMixedPrecisionLamb as FusedLAMB
|
|
except ImportError:
|
|
from apex.optimizers.fused_lamb import FusedLAMB
|
|
else:
|
|
from apex.optimizers.fused_lamb import FusedLAMB
|
|
|
|
optimizer = FusedLAMB(grouped_parameters,
|
|
lr=cfg.learning_rate,
|
|
betas=eval(cfg.adam_betas),
|
|
eps=cfg.adam_epsilon,
|
|
use_nvlamb=(cfg.use_nvlamb if "use_nvlamb" in cfg else False),
|
|
max_grad_norm=cfg.max_grad_norm)
|
|
elif "optimizer" in cfg and cfg.optimizer and "adafactor" in cfg.optimizer:
|
|
from transformers.optimization import Adafactor
|
|
|
|
optimizer = Adafactor(
|
|
grouped_parameters,
|
|
lr=cfg.learning_rate,
|
|
eps=(1e-30, 1e-3),
|
|
clip_threshold=1.0,
|
|
beta1=None,
|
|
weight_decay=0.0,
|
|
relative_step=False,
|
|
scale_parameter=False,
|
|
warmup_init=False
|
|
)
|
|
else:
|
|
if "bit_training" in cfg and cfg.bit_training:
|
|
from bitsandbytes.optim import AdamW8bit
|
|
|
|
optimizer = AdamW8bit(grouped_parameters, lr=cfg.learning_rate, eps=cfg.adam_epsilon,
|
|
betas=(eval(cfg.adam_betas)))
|
|
else:
|
|
if hasattr(cfg, "multi_tensor") and cfg.multi_tensor:
|
|
from torch.optim._multi_tensor import AdamW
|
|
else:
|
|
from torch.optim.adamw import AdamW
|
|
|
|
optimizer = AdamW(grouped_parameters, lr=cfg.learning_rate, eps=cfg.adam_epsilon,
|
|
betas=(eval(cfg.adam_betas)))
|
|
|
|
return optimizer
|
|
|
|
|
|
def get_optimizer_grouped_parameters(
|
|
model,
|
|
weight_decay,
|
|
lora_lr=5e-4,
|
|
no_decay_name_list=("bias", "LayerNorm.weight", "layernorm.weight"),
|
|
# lora_name_list=("lora_right_weight", "lora_left_weight"),
|
|
):
|
|
optimizer_grouped_parameters = [
|
|
{
|
|
"params": [
|
|
p
|
|
for n, p in model.named_parameters()
|
|
if (
|
|
not any(nd in n for nd in no_decay_name_list)
|
|
and p.requires_grad
|
|
# and not any(nd in n for nd in lora_name_list)
|
|
)
|
|
],
|
|
"weight_decay": weight_decay,
|
|
},
|
|
# {
|
|
# "params": [
|
|
# p
|
|
# for n, p in model.named_parameters()
|
|
# if (
|
|
# not any(nd in n for nd in no_decay_name_list)
|
|
# and p.requires_grad
|
|
# and any(nd in n for nd in lora_name_list)
|
|
# )
|
|
# ],
|
|
# "weight_decay": weight_decay,
|
|
# "lr": lora_lr,
|
|
# },
|
|
{
|
|
"params": [
|
|
p
|
|
for n, p in model.named_parameters()
|
|
if (any(nd in n for nd in no_decay_name_list) and p.requires_grad)
|
|
],
|
|
"weight_decay": 0.0,
|
|
},
|
|
]
|
|
if not optimizer_grouped_parameters[1]["params"]:
|
|
optimizer_grouped_parameters.pop(1)
|
|
return optimizer_grouped_parameters
|
|
|
|
|
|
def initialize_lr_scheduler(cfg: DictConfig, optimizer, num_warmup_steps: int, num_training_steps: int):
|
|
if hasattr(cfg, "lr_scheduler"):
|
|
if cfg.lr_scheduler == "linear":
|
|
from transformers import get_linear_schedule_with_warmup
|
|
|
|
lr_scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps)
|
|
elif cfg.lr_scheduler == "cosine":
|
|
from transformers import get_cosine_schedule_with_warmup
|
|
|
|
lr_scheduler = get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps)
|
|
elif cfg.lr_scheduler == "constant":
|
|
from transformers import get_constant_schedule_with_warmup
|
|
|
|
lr_scheduler = get_constant_schedule_with_warmup(optimizer, num_warmup_steps)
|
|
elif cfg.lr_scheduler == "poly":
|
|
from transformers import get_polynomial_decay_schedule_with_warmup
|
|
|
|
lr_scheduler = get_polynomial_decay_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps)
|
|
else:
|
|
raise NotImplementedError()
|
|
else:
|
|
from transformers import get_linear_schedule_with_warmup
|
|
|
|
lr_scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps)
|
|
|
|
return lr_scheduler
|
|
|
|
|
|
def note_best_checkpoint(cfg: DictConfig, results: Dict[str, float], sub_path: str):
|
|
metric = results[cfg.prediction_cfg.metric]
|
|
if (not cfg.prediction_cfg.best_result) or (
|
|
cfg.prediction_cfg.measure > 0 and metric > cfg.prediction_cfg.best_result) or (
|
|
cfg.prediction_cfg.measure < 0 and metric < cfg.prediction_cfg.best_result):
|
|
cfg.prediction_cfg.best_result = metric
|
|
cfg.prediction_cfg.best_checkpoint = sub_path
|
|
return True
|
|
return False
|
|
|
|
|
|
PREFIX_CHECKPOINT_DIR = "checkpoint"
|
|
_re_checkpoint = re.compile(r"^" + PREFIX_CHECKPOINT_DIR + r"-(\d+)$")
|
|
|
|
|
|
def get_last_checkpoint(folder):
|
|
content = os.listdir(folder)
|
|
checkpoints = [
|
|
path
|
|
for path in content
|
|
if _re_checkpoint.search(path) is not None and os.path.isdir(os.path.join(folder, path))
|
|
]
|
|
if len(checkpoints) == 0:
|
|
return None
|
|
return os.path.join(folder, max(checkpoints, key=lambda x: int(_re_checkpoint.search(x).groups()[0])))
|