gaotianyu1350
commented
1 year ago Hi,
Can you provide more details about it? All Huggingface models should be supported (maybe with slight modification when using prefix/lora due to different naming.
Closed thistleknot closed 1 year ago
gaotianyu1350
commented
1 year ago Hi,
Can you provide more details about it? All Huggingface models should be supported (maybe with slight modification when using prefix/lora due to different naming.
thistleknot
commented
1 year ago I'm surprised I didn't put in more context. I'll give it another shot and output the results of the log. I thought I had pasted what I had done.
thistleknot
commented
1 year ago My apologies for not providing more context.
I did attempt to modify the configuration files and swap in a different model but was unclear and no luck
thistleknot
commented
1 year ago https://github.com/princeton-nlp/MeZO/blob/main/large_models/trainer.py#L272
how would I override this line for gpt-neo? What about llama and llama-2?
gaotianyu1350
commented
1 year ago This is for linear probing so you can just ignore it if you are not doing linear probing.
thistleknot
commented
1 year ago Hopefully this looks right https://gist.github.com/thistleknot/6c0349c29e81c71cb547574666918c01
thistleknot
commented
1 year ago 0%| | 0/128 [00:00<?, ?it/s]You're using a GPTNeoXTokenizerFast tokenizer. Please note that with a fast tokenizer, using the __call__ method is faster than using a method to encode the text followed by a call to the pad method to get a padded encoding.
0%| | 0/128 [00:00<?, ?it/s]
Traceback (most recent call last):
File "/home/user/TrainLLMv2/mezo.py", line 72, in
when trying to use dolly-v2-3b
in functions.py
from common_imports import *
import inspect
from dataclasses import dataclass, is_dataclass, asdict
logger = logging.get_logger(__name__)
@dataclass
class OurArguments(TrainingArguments):
# dataset and sampling strategy
task_name: str = "SST2" # task name should match the string before Dataset in the Dataset class name. We support the following task_name: SST2, RTE, CB, BoolQ, WSC, WIC, MultiRC, Copa, ReCoRD, SQuAD, DROP
# Number of examples
num_train: int = 0 # ICL mode: number of demonstrations; training mode: number of training samples
num_dev: int = None # (only enabled with training) number of development samples
num_eval: int = None # number of evaluation samples
num_train_sets: int = None # how many sets of training samples/demos to sample; if None and train_set_seed is None, then we will sample one set for each evaluation sample
train_set_seed: int = None # designated seed to sample training samples/demos
result_file: str = None # file name for saving performance; if None, then use the task name, model name, and config
# Model loading
model_name: str = "facebook/opt-125m" # HuggingFace model name
load_float16: bool = False # load model parameters as float16
load_bfloat16: bool = False # load model parameters as bfloat16
load_int8: bool = False # load model parameters as int8
max_length: int = 2048 # max length the model can take
no_auto_device: bool = False # do not load model by auto device; should turn this on when using FSDP
# Calibration
sfc: bool = False # whether to use SFC calibration
icl_sfc: bool = False # whether to use SFC calibration for ICL samples
# Training
trainer: str = "none"
## options
## - none: no training -- for zero-shot or in-context learning (ICL)
## - regular: regular huggingface trainer -- for fine-tuning
## - zo: zeroth-order (MeZO) training
only_train_option: bool = True # whether to only train the option part of the input
train_as_classification: bool = False # take the log likelihood of all options and train as classification
# MeZO
zo_eps: float = 1e-3 # eps in MeZO
# Prefix tuning
prefix_tuning: bool = False # whether to use prefix tuning
num_prefix: int = 5 # number of prefixes to use
no_reparam: bool = True # do not use reparameterization trick
prefix_init_by_real_act: bool = True # initialize prefix by real activations of random words
# LoRA
lora: bool = False # whether to use LoRA
lora_alpha: int = 16 # alpha in LoRA
lora_r: int = 8 # r in LoRA
# Generation
sampling: bool = False # whether to use sampling
temperature: float = 1.0 # temperature for generation
num_beams: int = 1 # number of beams for generation
top_k: int = None # top-k for generation
top_p: float = 0.95 # top-p for generation
max_new_tokens: int = 50 # max number of new tokens to generate
eos_token: str = "\n" # end of sentence token
# Saving
save_model: bool = False # whether to save the model
no_eval: bool = False # whether to skip evaluation
tag: str = "" # saving tag
# Linear probing
linear_probing: bool = False # whether to do linear probing
lp_early_stopping: bool = False # whether to do early stopping in linear probing
head_tuning: bool = False # head tuning: only tune the LM head
# Untie emb/lm_head weights
untie_emb: bool = False # untie the embeddings and LM head
# Display
verbose: bool = False # verbose output
# Non-diff objective
non_diff: bool = False # use non-differentiable objective (only support F1 for SQuAD for now)
# Auto saving when interrupted
save_on_interrupt: bool = False # save model when interrupted (useful for long training)
class MeZOTrainer(Trainer):
def _inner_training_loop(
self, batch_size=None, args=None, resume_from_checkpoint=None, trial=None, ignore_keys_for_eval=None
):
self.accelerator.free_memory()
self._train_batch_size = batch_size
logger.debug(f"Currently training with a batch size of: {self._train_batch_size}")
# Data loader and number of training steps
train_dataloader = self.get_train_dataloader()
# MeZO added: Linear probing
if self.args.linear_probing:
def _get_token_prediction_layer(model):
if model.config.model_type == "opt":
return model.lm_head
else:
raise NotImplementedError(model.config.model_type)
def _extract_features(model, *args, **kwargs):
"""some magic for getting features pre last layer"""
features = {}
def __hook(model_, input_, output_):
features["features"] = input_[0].detach()
_get_token_prediction_layer(model).register_forward_hook(__hook)
model.forward(*args, **kwargs)
return features["features"]
logger.info("Linear probing")
logger.info("Starting to get features for training dataset")
targets = []
features = []
with torch.inference_mode():
for step, inputs in enumerate(tqdm(train_dataloader)):
for k, v in inputs.items():
if isinstance(v, torch.Tensor):
inputs[k] = v.to(self.model.device)
feature = _extract_features(self.model, **inputs)
target = inputs["labels"]
# Shift the target (bc it's autoregressive LM) and add the corresponding part
assert not self.args.train_as_classification and self.args.only_train_option
feature, target = feature[:, :-1], target[:, 1:]
for _i, _len in enumerate(inputs["option_len"]):
features.append(feature[_i, -_len:])
targets.append(target[_i, -_len:])
logger.info("Finished getting features for training dataset")
features = torch.cat(features, dim=0).cpu().numpy()
targets = torch.cat(targets, dim=0).cpu().numpy()
# Whether to use bias
if self.model.config.model_type in ["opt", "gpt2"]:
use_bias = False
else:
raise NotImplementedError
# Set early stopping
tol = 0.01 if self.args.lp_early_stopping else 1e-4 # 1e-4 is scipy default
max_iter = 1000 if self.args.lp_early_stopping else 5000
logger.info("Fitting logistic regression...")
reg = LogisticRegressionCV(max_iter=max_iter, fit_intercept=use_bias, multi_class="multinomial", random_state=0, tol=tol, n_jobs=-1).fit(features, targets)
logger.info("Done")
logger.info("Assigning weights to model")
decoder = _get_token_prediction_layer(self.model)
coef_torch = torch.tensor(reg.coef_, device=decoder.weight.device, dtype=decoder.weight.dtype)
if use_bias:
bias_torch = torch.tensor(reg.intercept_, device=decoder.weight.device, dtype=decoder.weight.dtype)
if coef_torch.shape[0] == 1: # The regressor only detects two classes
assert len(reg.classes_) == 2
coef_torch = torch.cat([-coef_torch / 2, coef_torch / 2], dim=0)
if use_bias:
bias_torch = torch.cat([-bias_torch / 2, bias_torch / 2], dim=0)
for _i, token_id in enumerate(reg.classes_):
decoder.weight.data[token_id] = coef_torch[_i]
if use_bias:
decoder.bias.data[token_id] = bias_torch[_i]
return None
# Setting up training control variables:
# number of training epochs: num_train_epochs
# number of training steps per epoch: num_update_steps_per_epoch
# total number of training steps to execute: max_steps
total_train_batch_size = self._train_batch_size * args.gradient_accumulation_steps * args.world_size
len_dataloader = None
if has_length(train_dataloader):
len_dataloader = len(train_dataloader)
num_update_steps_per_epoch = len_dataloader // args.gradient_accumulation_steps
num_update_steps_per_epoch = max(num_update_steps_per_epoch, 1)
num_examples = self.num_examples(train_dataloader)
if args.max_steps > 0:
max_steps = args.max_steps
num_train_epochs = args.max_steps // num_update_steps_per_epoch + int(
args.max_steps % num_update_steps_per_epoch > 0
)
# May be slightly incorrect if the last batch in the training dataloader has a smaller size but it's
# the best we can do.
num_train_samples = args.max_steps * total_train_batch_size
else:
max_steps = math.ceil(args.num_train_epochs * num_update_steps_per_epoch)
num_train_epochs = math.ceil(args.num_train_epochs)
num_train_samples = self.num_examples(train_dataloader) * args.num_train_epochs
elif args.max_steps > 0: # Rely on max_steps when dataloader does not have a working size
max_steps = args.max_steps
# Setting a very large number of epochs so we go as many times as necessary over the iterator.
num_train_epochs = sys.maxsize
num_update_steps_per_epoch = max_steps
num_examples = total_train_batch_size * args.max_steps
num_train_samples = args.max_steps * total_train_batch_size
else:
raise ValueError(
"args.max_steps must be set to a positive value if dataloader does not have a length, was"
f" {args.max_steps}"
)
# Compute absolute values for logging, eval, and save if given as ratio
if args.logging_steps and args.logging_steps < 1:
args.logging_steps = math.ceil(max_steps * args.logging_steps)
if args.eval_steps and args.eval_steps < 1:
args.eval_steps = math.ceil(max_steps * args.eval_steps)
if args.save_steps and args.save_steps < 1:
args.save_steps = math.ceil(max_steps * args.save_steps)
if DebugOption.UNDERFLOW_OVERFLOW in self.args.debug:
if self.args.n_gpu > 1:
# nn.DataParallel(model) replicates the model, creating new variables and module
# references registered here no longer work on other gpus, breaking the module
raise ValueError(
"Currently --debug underflow_overflow is not supported under DP. Please use DDP"
" (torch.distributed.launch)."
)
else:
debug_overflow = DebugUnderflowOverflow(self.model) # noqa
delay_optimizer_creation = (
self.sharded_ddp is not None
and self.sharded_ddp != ShardedDDPOption.SIMPLE
or is_sagemaker_mp_enabled()
or self.fsdp is not None
)
# We need to reset the scheduler, as its parameters may be different on subsequent calls
if self._created_lr_scheduler:
self.lr_scheduler = None
self._created_lr_scheduler = False
if self.is_deepspeed_enabled:
self.optimizer, self.lr_scheduler = deepspeed_init(self, num_training_steps=max_steps)
if not delay_optimizer_creation:
self.create_optimizer_and_scheduler(num_training_steps=max_steps)
self.state = TrainerState()
self.state.is_hyper_param_search = trial is not None
# Activate gradient checkpointing if needed
if args.gradient_checkpointing:
self.model.gradient_checkpointing_enable()
model = self._wrap_model(self.model_wrapped)
if is_sagemaker_mp_enabled() and resume_from_checkpoint is not None:
self._load_from_checkpoint(resume_from_checkpoint, model)
# as the model is wrapped, don't use `accelerator.prepare`
# this is for unhandled cases such as
# Fairscale Sharded DDP, FSDP-XLA, SageMaker MP/DP, DataParallel, IPEX
use_accelerator_prepare = True if model is self.model else False
if delay_optimizer_creation:
self.create_optimizer_and_scheduler(num_training_steps=max_steps)
# prepare using `accelerator` prepare
if use_accelerator_prepare:
self.model.train()
if hasattr(self.lr_scheduler, "step"):
if self.use_apex:
model = self.accelerator.prepare(self.model)
else:
model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
else:
# to handle cases wherein we pass "DummyScheduler" such as when it is specified in DeepSpeed config.
model, self.optimizer, self.lr_scheduler = self.accelerator.prepare(
self.model, self.optimizer, self.lr_scheduler
)
if self.is_fsdp_enabled:
self.model = model
# for the rest of this function `model` is the outside model, whether it was wrapped or not
if model is not self.model:
self.model_wrapped = model
# backward compatibility
if self.is_deepspeed_enabled:
self.deepspeed = self.model_wrapped
# deepspeed ckpt loading
if resume_from_checkpoint is not None and self.is_deepspeed_enabled:
deepspeed_load_checkpoint(self.model_wrapped, resume_from_checkpoint)
# Check if saved optimizer or scheduler states exist
self._load_optimizer_and_scheduler(resume_from_checkpoint)
# important: at this point:
# self.model is the Transformers Model
# self.model_wrapped is DDP(Transformers Model), Deepspeed(Transformers Model), etc.
# Train!
logger.info("***** Running training *****")
logger.info(f" Num examples = {num_examples:,}")
logger.info(f" Num Epochs = {num_train_epochs:,}")
logger.info(f" Instantaneous batch size per device = {self.args.per_device_train_batch_size:,}")
if self.args.per_device_train_batch_size != self._train_batch_size:
logger.info(f" Training with DataParallel so batch size has been adjusted to: {self._train_batch_size:,}")
logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_train_batch_size:,}")
logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}")
logger.info(f" Total optimization steps = {max_steps:,}")
logger.info(f" Number of trainable parameters = {get_model_param_count(model, trainable_only=True):,}")
self.state.epoch = 0
start_time = time.time()
epochs_trained = 0
steps_trained_in_current_epoch = 0
steps_trained_progress_bar = None
# Check if continuing training from a checkpoint
if resume_from_checkpoint is not None and os.path.isfile(
os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME)
):
self.state = TrainerState.load_from_json(os.path.join(resume_from_checkpoint, TRAINER_STATE_NAME))
epochs_trained = self.state.global_step // num_update_steps_per_epoch
if not args.ignore_data_skip:
steps_trained_in_current_epoch = self.state.global_step % (num_update_steps_per_epoch)
steps_trained_in_current_epoch *= args.gradient_accumulation_steps
else:
steps_trained_in_current_epoch = 0
logger.info(" Continuing training from checkpoint, will skip to saved global_step")
logger.info(f" Continuing training from epoch {epochs_trained}")
logger.info(f" Continuing training from global step {self.state.global_step}")
if not args.ignore_data_skip:
logger.info(
f" Will skip the first {epochs_trained} epochs then the first"
f" {steps_trained_in_current_epoch} batches in the first epoch."
)
# Update the references
self.callback_handler.model = self.model
self.callback_handler.optimizer = self.optimizer
self.callback_handler.lr_scheduler = self.lr_scheduler
self.callback_handler.train_dataloader = train_dataloader
if self.hp_name is not None and self._trial is not None:
# use self._trial because the SigOpt/Optuna hpo only call `_hp_search_setup(trial)` instead of passing trial
# parameter to Train when using DDP.
self.state.trial_name = self.hp_name(self._trial)
if trial is not None:
assignments = trial.assignments if self.hp_search_backend == HPSearchBackend.SIGOPT else trial
self.state.trial_params = hp_params(assignments)
else:
self.state.trial_params = None
# This should be the same if the state has been saved but in case the training arguments changed, it's safer
# to set this after the load.
self.state.max_steps = max_steps
self.state.num_train_epochs = num_train_epochs
self.state.is_local_process_zero = self.is_local_process_zero()
self.state.is_world_process_zero = self.is_world_process_zero()
# tr_loss is a tensor to avoid synchronization of TPUs through .item()
tr_loss = torch.tensor(0.0).to(args.device)
# _total_loss_scalar is updated everytime .item() has to be called on tr_loss and stores the sum of all losses
self._total_loss_scalar = 0.0
self._globalstep_last_logged = self.state.global_step
model.zero_grad()
self.control = self.callback_handler.on_train_begin(args, self.state, self.control)
# Skip the first epochs_trained epochs to get the random state of the dataloader at the right point.
if not args.ignore_data_skip:
for epoch in range(epochs_trained):
for _ in train_dataloader:
break
total_batched_samples = 0
for epoch in range(epochs_trained, num_train_epochs):
epoch_iterator = train_dataloader
# Reset the past mems state at the beginning of each epoch if necessary.
if args.past_index >= 0:
self._past = None
steps_in_epoch = (
len(epoch_iterator)
if len_dataloader is not None
else args.max_steps * args.gradient_accumulation_steps
)
self.control = self.callback_handler.on_epoch_begin(args, self.state, self.control)
if epoch == epochs_trained and resume_from_checkpoint is not None and steps_trained_in_current_epoch == 0:
self._load_rng_state(resume_from_checkpoint)
rng_to_sync = False
steps_skipped = 0
if steps_trained_in_current_epoch > 0:
epoch_iterator = skip_first_batches(epoch_iterator, steps_trained_in_current_epoch)
steps_skipped = steps_trained_in_current_epoch
steps_trained_in_current_epoch = 0
rng_to_sync = True
step = -1
for step, inputs in enumerate(epoch_iterator):
total_batched_samples += 1
if rng_to_sync:
self._load_rng_state(resume_from_checkpoint)
rng_to_sync = False
# Skip past any already trained steps if resuming training
if steps_trained_in_current_epoch > 0:
steps_trained_in_current_epoch -= 1
if steps_trained_progress_bar is not None:
steps_trained_progress_bar.update(1)
if steps_trained_in_current_epoch == 0:
self._load_rng_state(resume_from_checkpoint)
continue
elif steps_trained_progress_bar is not None:
steps_trained_progress_bar.close()
steps_trained_progress_bar = None
if step % args.gradient_accumulation_steps == 0:
self.control = self.callback_handler.on_step_begin(args, self.state, self.control)
with self.accelerator.accumulate(model):
# MeZO added: estimate gradient
if args.trainer == "zo":
tr_loss_step = self.zo_step(model, inputs)
else:
if (
((step + 1) % args.gradient_accumulation_steps != 0)
and args.local_rank != -1
and args._no_sync_in_gradient_accumulation
):
# Avoid unnecessary DDP synchronization since there will be no backward pass on this example.
with model.no_sync():
tr_loss_step = self.training_step(model, inputs)
else:
tr_loss_step = self.training_step(model, inputs)
if (
args.logging_nan_inf_filter
and not is_torch_tpu_available()
and (torch.isnan(tr_loss_step) or torch.isinf(tr_loss_step))
):
# if loss is nan or inf simply add the average of previous logged losses
tr_loss += tr_loss / (1 + self.state.global_step - self._globalstep_last_logged)
else:
tr_loss += tr_loss_step
self.current_flos += float(self.floating_point_ops(inputs))
is_last_step_and_steps_less_than_grad_acc = (
steps_in_epoch <= args.gradient_accumulation_steps and (step + 1) == steps_in_epoch
)
if (
total_batched_samples % args.gradient_accumulation_steps == 0
or
# last step in epoch but step is always smaller than gradient_accumulation_steps
is_last_step_and_steps_less_than_grad_acc
):
# the `or` condition of `is_last_step_and_steps_less_than_grad_acc` is not covered
# in accelerate. So, explicitly enable sync gradients to True in that case.
if is_last_step_and_steps_less_than_grad_acc or (
version.parse(accelerate_version) <= version.parse("0.20.3")
):
self.accelerator.gradient_state._set_sync_gradients(True)
# MeZO added: update model with the estimated gradient
if args.trainer == "zo":
self.zo_update(model)
else:
# Gradient clipping
if args.max_grad_norm is not None and args.max_grad_norm > 0:
# deepspeed does its own clipping
if self.do_grad_scaling:
# Reduce gradients first for XLA
if is_torch_tpu_available():
gradients = xm._fetch_gradients(self.optimizer)
xm.all_reduce("sum", gradients, scale=1.0 / xm.xrt_world_size())
# AMP: gradients need unscaling
self.scaler.unscale_(self.optimizer)
if is_sagemaker_mp_enabled() and args.fp16:
self.optimizer.clip_master_grads(args.max_grad_norm)
elif hasattr(self.optimizer, "clip_grad_norm"):
# Some optimizers (like the sharded optimizer) have a specific way to do gradient clipping
self.optimizer.clip_grad_norm(args.max_grad_norm)
elif hasattr(model, "clip_grad_norm_"):
# Some models (like FullyShardedDDP) have a specific way to do gradient clipping
model.clip_grad_norm_(args.max_grad_norm)
elif self.use_apex:
# Revert to normal clipping otherwise, handling Apex or full precision
nn.utils.clip_grad_norm_(
amp.master_params(self.optimizer),
args.max_grad_norm,
)
else:
self.accelerator.clip_grad_norm_(
model.parameters(),
args.max_grad_norm,
)
# Optimizer step
optimizer_was_run = True
if is_torch_tpu_available():
if self.do_grad_scaling:
self.scaler.step(self.optimizer)
self.scaler.update()
else:
# tpu-comment: accelerate wrapped optimizers call xm.optimizer_step
self.optimizer.step()
elif self.do_grad_scaling:
scale_before = self.scaler.get_scale()
self.scaler.step(self.optimizer)
self.scaler.update()
scale_after = self.scaler.get_scale()
optimizer_was_run = scale_before <= scale_after
else:
self.optimizer.step()
optimizer_was_run = not self.accelerator.optimizer_step_was_skipped
if optimizer_was_run:
# Delay optimizer scheduling until metrics are generated
if not isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
self.lr_scheduler.step()
model.zero_grad()
self.state.global_step += 1
self.state.epoch = epoch + (step + 1 + steps_skipped) / steps_in_epoch
self.control = self.callback_handler.on_step_end(args, self.state, self.control)
self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
else:
self.control = self.callback_handler.on_substep_end(args, self.state, self.control)
if self.control.should_epoch_stop or self.control.should_training_stop:
break
if step < 0:
logger.warning(
"There seems to be not a single sample in your epoch_iterator, stopping training at step"
f" {self.state.global_step}! This is expected if you're using an IterableDataset and set"
f" num_steps ({max_steps}) higher than the number of available samples."
)
self.control.should_training_stop = True
self.control = self.callback_handler.on_epoch_end(args, self.state, self.control)
self._maybe_log_save_evaluate(tr_loss, model, trial, epoch, ignore_keys_for_eval)
if DebugOption.TPU_METRICS_DEBUG in self.args.debug:
if is_torch_tpu_available():
# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
xm.master_print(met.metrics_report())
else:
logger.warning(
"You enabled PyTorch/XLA debug metrics but you don't have a TPU "
"configured. Check your training configuration if this is unexpected."
)
if self.control.should_training_stop:
break
if args.past_index and hasattr(self, "_past"):
# Clean the state at the end of training
delattr(self, "_past")
logger.info("\n\nTraining completed. Do not forget to share your model on huggingface.co/models =)\n\n")
if args.load_best_model_at_end and self.state.best_model_checkpoint is not None:
# Wait for everyone to get here so we are sur the model has been saved by process 0.
if is_torch_tpu_available():
xm.rendezvous("load_best_model_at_end")
elif args.parallel_mode == ParallelMode.DISTRIBUTED:
dist.barrier()
elif is_sagemaker_mp_enabled():
smp.barrier()
self._load_best_model()
# add remaining tr_loss
self._total_loss_scalar += tr_loss.item()
train_loss = self._total_loss_scalar / self.state.global_step
metrics = speed_metrics("train", start_time, num_samples=num_train_samples, num_steps=self.state.max_steps)
self.store_flos()
metrics["total_flos"] = self.state.total_flos
metrics["train_loss"] = train_loss
self.is_in_train = False
self._memory_tracker.stop_and_update_metrics(metrics)
self.log(metrics)
run_dir = self._get_output_dir(trial)
checkpoints_sorted = self._sorted_checkpoints(use_mtime=False, output_dir=run_dir)
# Delete the last checkpoint when save_total_limit=1 if it's different from the best checkpoint and process allowed to save.
if self.args.should_save and self.state.best_model_checkpoint is not None and self.args.save_total_limit == 1:
for checkpoint in checkpoints_sorted:
if checkpoint != self.state.best_model_checkpoint:
logger.info(f"Deleting older checkpoint [{checkpoint}] due to args.save_total_limit")
shutil.rmtree(checkpoint)
self.control = self.callback_handler.on_train_end(args, self.state, self.control)
return TrainOutput(self.state.global_step, train_loss, metrics)
############## MeZO ##############
def zo_perturb_parameters(self, random_seed=None, scaling_factor=1):
"""
Perturb the parameters with random vector z.
Input:
- random_seed: random seed for MeZO in-place perturbation (if it's None, we will use self.zo_random_seed)
- scaling_factor: theta = theta + scaling_factor * z * eps
"""
# Set the random seed to ensure that we sample the same z for perturbation/update
torch.manual_seed(random_seed if random_seed is not None else self.zo_random_seed)
for name, param in self.named_parameters_to_optim:
z = torch.normal(mean=0, std=1, size=param.data.size(), device=param.data.device, dtype=param.data.dtype)
param.data = param.data + scaling_factor * z * self.args.zo_eps
def zo_forward(self, model, inputs):
"""
Get (no gradient) loss from the model. Dropout is turned off too.
"""
model.eval()
if self.args.non_diff:
# Non-differentiable objective (may require autoregressive generation)
return self.zo_forward_nondiff(model, inputs)
with torch.inference_mode():
inputs = self._prepare_inputs(inputs)
with self.compute_loss_context_manager():
loss = self.compute_loss(model, inputs)
if self.args.n_gpu > 1:
# Warning: this is copied from the original Huggingface Trainer. Untested.
loss = loss.mean() # mean() to average on multi-gpu parallel training
return loss.detach()
def zo_forward_nondiff(self, model, inputs):
"""
Get (no gradient) non-diffiable loss from the model.
"""
model.eval()
assert self.args.task_name == "SQuAD", "Non differentiable objective only supports SQuAD for now."
with torch.inference_mode():
inputs = self._prepare_inputs(inputs)
args = self.args
outputs = self.model.generate(
inputs["input_ids"], do_sample=args.sampling, temperature=args.temperature,
num_beams=args.num_beams, top_p=args.top_p, top_k=args.top_k, max_new_tokens=min(args.max_new_tokens, args.max_length - inputs["input_ids"].size(1)),
num_return_sequences=1, eos_token_id=[self.tokenizer.encode(args.eos_token, add_special_tokens=False)[0], self.tokenizer.eos_token_id],
)
output_text = []
for i in range(len(outputs)):
output_text.append(self.tokenizer.decode(outputs[i][inputs["input_ids"].size(1):], skip_special_tokens=True).strip())
f1s = [f1(output_text[i], inputs['gold'][i]) for i in range(len(output_text))]
return -torch.tensor(np.mean(f1s), dtype=torch.float32)
def zo_step(self, model, inputs):
"""
Estimate gradient by MeZO. Return the loss from f(theta + z)
"""
args = self.args
# What parameters to optimize
self.named_parameters_to_optim = []
for name, param in model.named_parameters():
if param.requires_grad:
self.named_parameters_to_optim.append((name, param))
# Sample the random seed for sampling z
self.zo_random_seed = np.random.randint(1000000000)
# First function evaluation
self.zo_perturb_parameters(scaling_factor=1)
loss1 = self.zo_forward(model, inputs)
# Second function evaluation
self.zo_perturb_parameters(scaling_factor=-2)
loss2 = self.zo_forward(model, inputs)
self.projected_grad = ((loss1 - loss2) / (2 * self.args.zo_eps)).item()
# No gradient accumulation support
assert self.args.gradient_accumulation_steps == 1
# Reset model back to its parameters at start of step
self.zo_perturb_parameters(scaling_factor=1)
return loss1
def zo_update(self, model):
"""
Update the parameters with the estimated gradients.
"""
args = self.args
# Reset the random seed for sampling zs
torch.manual_seed(self.zo_random_seed)
for name, param in self.named_parameters_to_optim:
# Resample z
z = torch.normal(mean=0, std=1, size=param.data.size(), device=param.data.device, dtype=param.data.dtype)
if "bias" not in name and "layer_norm" not in name and "layernorm" not in name:
param.data = param.data - self._get_learning_rate() * (self.projected_grad * z + args.weight_decay * param.data)
else:
param.data = param.data - self._get_learning_rate() * (self.projected_grad * z)
self.lr_scheduler.step()
############## Misc overload functions ##############
def _set_signature_columns_if_needed(self):
"""
We overload this function for non-differentiable objective training to pass "gold" -- the gold text for the task
"""
if self._signature_columns is None:
# Inspect model forward signature to keep only the arguments it accepts.
signature = inspect.signature(self.model.forward)
self._signature_columns = list(signature.parameters.keys())
# Labels may be named label or label_ids, the default data collator handles that.
self._signature_columns += list(set(["label", "label_ids"] + self.label_names))
self._signature_columns += ["gold"]
def save_model(self, output_dir: Optional[str] = None, _internal_call: bool = False):
"""
We overload this function to fix an FSDP saving bug (before fix, it will likely cause OOM)
"""
if output_dir is None:
output_dir = self.args.output_dir
if is_torch_tpu_available():
self._save_tpu(output_dir)
elif is_sagemaker_mp_enabled():
# Calling the state_dict needs to be done on the wrapped model and on all processes.
os.makedirs(output_dir, exist_ok=True)
state_dict = self.model_wrapped.state_dict()
if self.args.should_save:
self._save(output_dir, state_dict=state_dict)
if IS_SAGEMAKER_MP_POST_1_10:
# 'user_content.pt' indicates model state_dict saved with smp >= 1.10
Path(os.path.join(output_dir, "user_content.pt")).touch()
elif (
ShardedDDPOption.ZERO_DP_2 in self.args.sharded_ddp
or ShardedDDPOption.ZERO_DP_3 in self.args.sharded_ddp
or self.fsdp is not None
):
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, StateDictType, FullStateDictConfig
full_state_dict_config = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
# Fix the FSDP loading bug
with FSDP.state_dict_type(self.model, StateDictType.FULL_STATE_DICT, full_state_dict_config):
state_dict = self.model.state_dict()
# state_dict = self.model.state_dict()
if self.args.should_save:
self._save(output_dir, state_dict=state_dict)
elif self.deepspeed:
# this takes care of everything as long as we aren't under zero3
if self.args.should_save:
self._save(output_dir)
if is_deepspeed_zero3_enabled():
# It's too complicated to try to override different places where the weights dump gets
# saved, so since under zero3 the file is bogus, simply delete it. The user should
# either user deepspeed checkpoint to resume or to recover full weights use
# zero_to_fp32.py stored in the checkpoint.
if self.args.should_save:
file = os.path.join(output_dir, WEIGHTS_NAME)
if os.path.isfile(file):
# logger.info(f"deepspeed zero3: removing {file}, see zero_to_fp32.py to recover weights")
os.remove(file)
# now save the real model if stage3_gather_16bit_weights_on_model_save=True
# if false it will not be saved.
# This must be called on all ranks
if not self.deepspeed.save_16bit_model(output_dir, WEIGHTS_NAME):
logger.warning(
"deepspeed.save_16bit_model didn't save the model, since"
" stage3_gather_16bit_weights_on_model_save=false. Saving the full checkpoint instead, use"
" zero_to_fp32.py to recover weights"
)
self.deepspeed.save_checkpoint(output_dir)
elif self.args.should_save:
self._save(output_dir)
# Push to the Hub when `save_model` is called by the user.
if self.args.push_to_hub and not _internal_call:
self.push_to_hub(commit_message="Model save")
from transformers import TrainingArguments, Trainer, AutoTokenizer, AutoModelForCausalLM
from datasets import load_dataset
from functions import *
# Model and dataset settings
MODEL = "/home/user/oobabooga_linux/text-generation-webui/models/dolly-v2-3b"
TASK = "text" # Assuming you're treating this as a text generation task
INPUT_FILE = '0082 Encyclopedia of Classical Philosophy-OCR.txt'
# Hyperparameters and settings
MODEL_NAME = "dolly-v2-3b-pretrained"
TAG = "finetuning" # Modify as per your needs
SEED = 42
TRAIN = 2048 # Modify as per your needs
DEV = 256 # Modify as per your needs
EVAL = 512 # Modify as per your needs
STEPS = 10000 # Modify as per your needs
LR = 5e-5
EPS = 1e-8
BS = 8 # Batch size
EVAL_STEPS = 500 # Modify as per your needs
EXTRA_ARGS = "" # Any extra arguments
TASK_ARGS = "" # Task-specific arguments
# Tokenize the data
#tokenizer = GPT2Tokenizer.from_pretrained(MODEL)
tokenizer = AutoTokenizer.from_pretrained(MODEL)
tokenizer.pad_token = tokenizer.eos_token
block_size = 2048
stride = block_size // 2
def tokenize_function(examples):
result = tokenizer(examples["text"], truncation=True, padding='max_length', max_length=block_size, stride=stride, return_overflowing_tokens=True)
return result
dataset = load_dataset('text', data_files={'train': INPUT_FILE})
tokenized_datasets = dataset.map(tokenize_function, batched=True)
# Model and Training settings
#model = GPT2LMHeadModel.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL)
training_args = OurArguments(
per_device_train_batch_size=BS,
learning_rate=LR,
num_train_epochs=1,
logging_dir=f'./logs/{TASK}-{MODEL_NAME}-{TAG}',
logging_steps=10,
save_steps=EVAL_STEPS,
eval_steps=EVAL_STEPS,
max_steps=STEPS,
evaluation_strategy="steps",
save_strategy="steps",
save_total_limit=1,
load_best_model_at_end=True,
push_to_hub=False,
output_dir=f"./result/{TASK}-{MODEL_NAME}-{TAG}",
overwrite_output_dir=True,
linear_probing=True,
lr_scheduler_type="constant",
)
trainer = MeZOTrainer(
model=model,
args=training_args,
train_dataset=tokenized_datasets["train"],
eval_dataset=tokenized_datasets["train"], # Assuming you're evaluating on the same data, modify as needed
tokenizer=tokenizer,
)
# Train the model
trainer.train()
# Save the model and tokenizer
model.save_pretrained(f"./result/{TASK}-{MODEL_NAME}-{TAG}")
tokenizer.save_pretrained(f"./result/{TASK}-{MODEL_NAME}-{TAG}")
:/