mirror of
https://github.com/SakanaAI/doc-to-lora.git
synced 2026-07-23 17:01:04 +02:00
fix logging + l1norm div by num ctx (work correctly w/ grad accum)
This commit is contained in:
51616
parent
41a1cfb90a
commit
a17d2532c2
1 changed files with 151 additions and 133 deletions
|
|
@ -13,14 +13,10 @@ from ctx_to_lora.modeling.hypernet import ModulatedPretrainedModel
|
|||
logger = logging.getLogger()
|
||||
|
||||
|
||||
class DistillationTrainer(Trainer):
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0)
|
||||
self.use_per_ctx_average_loss = kwargs.pop("use_per_ctx_average_loss", False)
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
class ModulatedModelTrainer(Trainer):
|
||||
# modified from the base Trainer to support per-context average loss
|
||||
def get_batch_samples_ctx(self, epoch_iterator, num_batches, device):
|
||||
def get_batch_samples(self, epoch_iterator, num_batches, device):
|
||||
# only used with `use_per_ctx_average_loss=True`
|
||||
batch_samples = []
|
||||
num_items_in_batch = None
|
||||
|
||||
|
|
@ -33,27 +29,17 @@ class DistillationTrainer(Trainer):
|
|||
count_num_items_in_batch = (
|
||||
len(batch_samples) > 0
|
||||
and "labels" in batch_samples[0]
|
||||
and "ctx_ids" in batch_samples[0]
|
||||
and (
|
||||
# num_items_in_batch is passed to model forward
|
||||
# https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3757
|
||||
self.model_accepts_loss_kwargs
|
||||
# num_items_in_batch is passed to compute_loss_func
|
||||
# https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3773
|
||||
or self.compute_loss_func is not None
|
||||
# num_items_in_batch is also verified if (self.model_accepts_loss_kwargs or self.compute_loss_func)
|
||||
# https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/trainer.py#L3790
|
||||
)
|
||||
and "n_ctx_chunks" in batch_samples[0]
|
||||
)
|
||||
|
||||
if count_num_items_in_batch:
|
||||
# For now we don't support object detection
|
||||
try:
|
||||
num_items_in_batch = sum(
|
||||
[(batch["ctx_position_ids"] == 0).sum() for batch in batch_samples]
|
||||
)
|
||||
except (TypeError, AttributeError):
|
||||
pass
|
||||
num_items_in_batch = dict()
|
||||
num_items_in_batch["ctx"] = torch.tensor(
|
||||
sum([batch["n_ctx_chunks"].numel() for batch in batch_samples])
|
||||
)
|
||||
num_items_in_batch["labels"] = sum(
|
||||
[(batch["labels"].ne(-100)).sum() for batch in batch_samples]
|
||||
)
|
||||
|
||||
if num_items_in_batch is not None:
|
||||
if self.args.average_tokens_across_devices:
|
||||
|
|
@ -68,117 +54,138 @@ class DistillationTrainer(Trainer):
|
|||
|
||||
return batch_samples, num_items_in_batch
|
||||
|
||||
|
||||
class DistillationTrainer(ModulatedModelTrainer):
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0)
|
||||
self.use_per_ctx_average_loss = kwargs.pop("use_per_ctx_average_loss", False)
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def compute_loss(
|
||||
self, model, inputs, return_outputs=False, num_items_in_batch=None
|
||||
):
|
||||
if self.model_accepts_loss_kwargs:
|
||||
loss_kwargs = {}
|
||||
if num_items_in_batch is not None:
|
||||
loss_kwargs["num_items_in_batch"] = num_items_in_batch
|
||||
inputs = {**inputs, **loss_kwargs}
|
||||
use_kl = "logprobs_vals" in inputs
|
||||
if use_kl:
|
||||
labels = inputs.pop("labels", None)
|
||||
label_pos = torch.where(labels != -100)
|
||||
# NOTE: the loss output from this fn will be ***added***
|
||||
# meaning that we should always scale the loss wrt `num_items_in_batch`
|
||||
# (average over the number of items in the accumulated batch)
|
||||
|
||||
labels = inputs.pop("labels", None)
|
||||
label_pos = torch.where(labels != -100)
|
||||
outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True)
|
||||
|
||||
if not use_kl:
|
||||
loss = outputs["loss"]
|
||||
target_logp = inputs.pop("logprobs_vals").squeeze(0)
|
||||
indices = inputs.pop("logprobs_indices").squeeze(0)
|
||||
|
||||
else:
|
||||
target_logp = inputs.pop("logprobs_vals").squeeze(0)
|
||||
indices = inputs.pop("logprobs_indices").squeeze(0)
|
||||
assert label_pos[0].shape[0] == target_logp.shape[0], (
|
||||
"Label positions and target log probabilities should have the same # tokens."
|
||||
f"Got : {label_pos[0].shape[0]=} and {target_logp.shape[0]=}"
|
||||
)
|
||||
|
||||
assert label_pos[0].shape[0] == target_logp.shape[0], (
|
||||
"Label positions and target log probabilities should have the same # tokens."
|
||||
f"Got : {label_pos[0].shape[0]=} and {target_logp.shape[0]=}"
|
||||
##### KL loss
|
||||
outputs_logits = outputs.logits[label_pos[0], label_pos[1] - 1] # shift back 1
|
||||
|
||||
teacher_logp = torch.full_like(outputs_logits, -torch.inf)
|
||||
teacher_logp.scatter_(1, indices, target_logp)
|
||||
# reduction = "batchmean" if num_items_in_batch is None else "sum"
|
||||
|
||||
p = teacher_logp.exp()
|
||||
logq = nn.functional.log_softmax(outputs_logits, dim=-1)
|
||||
loss = -torch.sum(p * logq, dim=-1)
|
||||
|
||||
if self.use_per_ctx_average_loss:
|
||||
n_queries_per_ctx = inputs["n_queries"].tolist()
|
||||
|
||||
position_ids = inputs["position_ids"].squeeze(0)
|
||||
# account only label positions
|
||||
label_mask = labels.squeeze(0) != -100
|
||||
label_pos_ids = label_mask * position_ids
|
||||
label_pos_ids_diff = label_pos_ids.diff(
|
||||
append=torch.tensor([0], device=position_ids.device)
|
||||
)
|
||||
start_label_pos = torch.where((label_pos_ids_diff > 0) * ~label_mask)[0]
|
||||
end_label_pos = torch.where((label_pos_ids_diff < 0) * label_mask)[0]
|
||||
|
||||
label_seq_lens = end_label_pos - start_label_pos
|
||||
cu_label_seq_lens = torch.cumsum(label_seq_lens, dim=0)
|
||||
start_indices = torch.cat(
|
||||
(
|
||||
torch.tensor([0], device=cu_label_seq_lens.device),
|
||||
cu_label_seq_lens[:-1],
|
||||
)
|
||||
)
|
||||
|
||||
##### KL loss
|
||||
outputs_logits = outputs.logits[
|
||||
label_pos[0], label_pos[1] - 1
|
||||
] # shift back 1
|
||||
|
||||
teacher_logp = torch.full_like(outputs_logits, -torch.inf)
|
||||
teacher_logp.scatter_(1, indices, target_logp)
|
||||
reduction = "batchmean" if num_items_in_batch is None else "sum"
|
||||
|
||||
p = teacher_logp.exp()
|
||||
logq = nn.functional.log_softmax(outputs_logits, dim=-1)
|
||||
loss = -torch.sum(p * logq, dim=-1)
|
||||
|
||||
if self.use_per_ctx_average_loss:
|
||||
n_queries_per_ctx = inputs["n_queries"].tolist()
|
||||
|
||||
position_ids = inputs["position_ids"].squeeze(0)
|
||||
# account only label positions
|
||||
label_mask = labels.squeeze(0) != -100
|
||||
label_pos_ids = label_mask * position_ids
|
||||
label_pos_ids_diff = label_pos_ids.diff(
|
||||
append=torch.tensor([0], device=position_ids.device)
|
||||
)
|
||||
start_label_pos = torch.where((label_pos_ids_diff > 0) * ~label_mask)[0]
|
||||
end_label_pos = torch.where((label_pos_ids_diff < 0) * label_mask)[0]
|
||||
|
||||
label_seq_lens = end_label_pos - start_label_pos
|
||||
cu_label_seq_lens = torch.cumsum(label_seq_lens, dim=0)
|
||||
start_indices = torch.cat(
|
||||
(
|
||||
torch.tensor([0], device=cu_label_seq_lens.device),
|
||||
cu_label_seq_lens[:-1],
|
||||
)
|
||||
)
|
||||
|
||||
# these stack and split can be optimized but let's keep it simple
|
||||
# mean across tokens of each q
|
||||
qa_losses = torch.stack(
|
||||
[
|
||||
loss[start:end].mean()
|
||||
for start, end in zip(start_indices, cu_label_seq_lens)
|
||||
]
|
||||
)
|
||||
|
||||
# mean across queries of each ctx
|
||||
per_ctx_losses = [
|
||||
ctx.mean() for ctx in torch.split(qa_losses, n_queries_per_ctx)
|
||||
# these stack and split can be optimized but let's keep it simple
|
||||
# mean across tokens of each q
|
||||
qa_losses = torch.stack(
|
||||
[
|
||||
loss[start:end].mean()
|
||||
for start, end in zip(start_indices, cu_label_seq_lens)
|
||||
]
|
||||
)
|
||||
|
||||
# mean across contexts
|
||||
loss = torch.stack(per_ctx_losses)
|
||||
# mean across queries of each ctx
|
||||
per_ctx_losses = [
|
||||
ctx.mean() for ctx in torch.split(qa_losses, n_queries_per_ctx)
|
||||
]
|
||||
|
||||
if reduction == "batchmean":
|
||||
loss = loss.mean()
|
||||
elif reduction == "sum" and num_items_in_batch is not None:
|
||||
loss = loss.sum() / num_items_in_batch
|
||||
#####
|
||||
# mean across contexts
|
||||
loss = torch.stack(per_ctx_losses)
|
||||
|
||||
if num_items_in_batch is not None:
|
||||
if self.use_per_ctx_average_loss:
|
||||
loss = loss.sum() / num_items_in_batch["ctx"]
|
||||
else:
|
||||
loss = loss.sum() / num_items_in_batch["labels"]
|
||||
else:
|
||||
# eval
|
||||
loss = loss.mean()
|
||||
|
||||
# if reduction == "batchmean":
|
||||
# loss = loss.mean()
|
||||
# elif reduction == "sum":
|
||||
# # loss does not scale with grad acc
|
||||
# # num_items_in_batch does
|
||||
# # this works for both token-avg and ctx-avg
|
||||
# # loss = loss.sum() / num_items_in_batch
|
||||
|
||||
# `num_items_in_batch` is # tokens if `args.use_ctx_average_loss=False``
|
||||
# loss = loss.sum() / num_items_in_batch
|
||||
#####
|
||||
|
||||
##### unpack gen lora dict and compute regularization loss
|
||||
l1_norm = 0
|
||||
n_modules = len(gen_loras)
|
||||
for module, lora in gen_loras.items():
|
||||
l1_norm += lora["A"].abs().mean() + lora["B"].abs().mean()
|
||||
l1_norm += lora["A"].abs().sum(0).mean() + lora["B"].abs().sum(0).mean()
|
||||
l1_norm /= n_modules
|
||||
if num_items_in_batch is not None:
|
||||
# during eval `num_items_in_batch` will be None
|
||||
l1_norm /= num_items_in_batch["ctx"]
|
||||
|
||||
total_loss = loss + self.gen_lora_l1_reg_coef * l1_norm
|
||||
#####
|
||||
|
||||
scaler = self.args.gradient_accumulation_steps
|
||||
if self.args.average_tokens_across_devices and num_items_in_batch is not None:
|
||||
total_loss *= self.accelerator.num_processes
|
||||
scaler *= self.accelerator.num_processes
|
||||
|
||||
# rough estimate of the losses (we only log the values from one step)
|
||||
if (self.state.global_step == 1 and self.args.logging_first_step) or (
|
||||
self.args.logging_strategy == IntervalStrategy.STEPS
|
||||
and self.state.global_step % self.state.logging_steps == 0
|
||||
):
|
||||
loss_field = "kl_loss" if use_kl else "ce_loss"
|
||||
self.log({loss_field: loss.item(), "gen_lora_l1_norm": l1_norm.item()})
|
||||
loss += self.gen_lora_l1_reg_coef * l1_norm
|
||||
#####
|
||||
# compensate `num_items_in_batch` division
|
||||
self.log(
|
||||
{
|
||||
"kl_loss": loss.item() * scaler,
|
||||
"gen_lora_l1_norm": l1_norm.item() * scaler,
|
||||
}
|
||||
)
|
||||
|
||||
if (
|
||||
self.args.average_tokens_across_devices
|
||||
and self.model_accepts_loss_kwargs
|
||||
and num_items_in_batch is not None
|
||||
):
|
||||
loss *= self.accelerator.num_processes
|
||||
|
||||
return (loss, outputs) if return_outputs else loss
|
||||
return (total_loss, outputs) if return_outputs else total_loss
|
||||
|
||||
|
||||
class ModulatedModelTrainer(Trainer):
|
||||
class CrossEntropyTrainer(ModulatedModelTrainer):
|
||||
def __init__(self, *args, **kwargs):
|
||||
self.gen_lora_l1_reg_coef = kwargs.pop("gen_lora_l1_reg_coef", 0.0)
|
||||
super().__init__(*args, **kwargs)
|
||||
|
|
@ -197,11 +204,11 @@ class ModulatedModelTrainer(Trainer):
|
|||
labels = inputs.pop("labels")
|
||||
else:
|
||||
labels = None
|
||||
if self.model_accepts_loss_kwargs:
|
||||
loss_kwargs = {}
|
||||
if num_items_in_batch is not None:
|
||||
loss_kwargs["num_items_in_batch"] = num_items_in_batch
|
||||
inputs = {**inputs, **loss_kwargs}
|
||||
|
||||
loss_kwargs = {}
|
||||
if num_items_in_batch is not None:
|
||||
loss_kwargs["num_items_in_batch"] = num_items_in_batch["labels"]
|
||||
inputs = {**inputs, **loss_kwargs}
|
||||
outputs, (gen_loras, _) = model(**inputs, return_generated_lora=True)
|
||||
# Save past state if it exists
|
||||
if self.args.past_index >= 0:
|
||||
|
|
@ -216,7 +223,7 @@ class ModulatedModelTrainer(Trainer):
|
|||
# User-defined compute_loss function
|
||||
if self.compute_loss_func is not None:
|
||||
loss = self.compute_loss_func(
|
||||
outputs, labels, num_items_in_batch=num_items_in_batch
|
||||
outputs, labels, num_items_in_batch=num_items_in_batch["labels"]
|
||||
)
|
||||
elif model_name in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES.values():
|
||||
loss = self.label_smoother(outputs, labels, shift_labels=True)
|
||||
|
|
@ -236,25 +243,36 @@ class ModulatedModelTrainer(Trainer):
|
|||
|
||||
##### unpack gen lora dict and compute regularization loss
|
||||
l1_norm = 0
|
||||
n_modules = len(gen_loras)
|
||||
for module, lora in gen_loras.items():
|
||||
l1_norm += lora["A"].abs().mean() + lora["B"].abs().mean()
|
||||
l1_norm += lora["A"].abs().sum(0).mean() + lora["B"].abs().sum(0).mean()
|
||||
l1_norm /= n_modules
|
||||
if num_items_in_batch is not None:
|
||||
# during eval `num_items_in_batch` will be None
|
||||
l1_norm /= num_items_in_batch["ctx"]
|
||||
|
||||
total_loss = loss + self.gen_lora_l1_reg_coef * l1_norm
|
||||
#####
|
||||
|
||||
scaler = self.args.gradient_accumulation_steps
|
||||
if self.args.average_tokens_across_devices and num_items_in_batch is not None:
|
||||
total_loss *= self.accelerator.num_processes
|
||||
scaler *= self.accelerator.num_processes
|
||||
|
||||
# rough estimate of the losses (we only log the values from one step)
|
||||
if (self.state.global_step == 1 and self.args.logging_first_step) or (
|
||||
self.args.logging_strategy == IntervalStrategy.STEPS
|
||||
and self.state.global_step % self.state.logging_steps == 0
|
||||
):
|
||||
self.log({"ce_loss": loss.item(), "gen_lora_l1_norm": l1_norm.item()})
|
||||
loss += self.gen_lora_l1_reg_coef * l1_norm
|
||||
#####
|
||||
# compensate `num_items_in_batch` division
|
||||
self.log(
|
||||
{
|
||||
"ce_loss": loss.item() * scaler,
|
||||
"gen_lora_l1_norm": l1_norm.item() * scaler,
|
||||
}
|
||||
)
|
||||
|
||||
if (
|
||||
self.args.average_tokens_across_devices
|
||||
and self.model_accepts_loss_kwargs
|
||||
and num_items_in_batch is not None
|
||||
):
|
||||
loss *= self.accelerator.num_processes
|
||||
|
||||
return (loss, outputs) if return_outputs else loss
|
||||
return (total_loss, outputs) if return_outputs else total_loss
|
||||
|
||||
|
||||
def get_decay_parameter_names(model) -> list[str]:
|
||||
|
|
@ -269,7 +287,7 @@ def get_decay_parameter_names(model) -> list[str]:
|
|||
decay_parameters = get_parameter_names(
|
||||
model,
|
||||
[nn.Embedding, nn.LayerNorm],
|
||||
["bias", "layernorm", "rmsnorm", "latents_q"],
|
||||
["scaler", "bias", "layernorm", "rmsnorm", "latents_q"],
|
||||
)
|
||||
return decay_parameters
|
||||
|
||||
|
|
@ -299,8 +317,8 @@ def train_model(
|
|||
is_modulated_model = isinstance(model, ModulatedPretrainedModel)
|
||||
trainer_cls = Trainer
|
||||
if is_modulated_model:
|
||||
logger.info("Training with modulated model. Using ModulatedModelTrainer.")
|
||||
trainer_cls = ModulatedModelTrainer
|
||||
logger.info("Training with modulated model. Using CrossEntropyTrainer.")
|
||||
trainer_cls = CrossEntropyTrainer
|
||||
trainer_kwargs["gen_lora_l1_reg_coef"] = training_args.gen_lora_l1_reg_coef
|
||||
del training_args.gen_lora_l1_reg_coef
|
||||
|
||||
|
|
@ -319,8 +337,8 @@ def train_model(
|
|||
training_args.per_device_train_batch_size = 128
|
||||
|
||||
trainer = trainer_cls(**trainer_kwargs)
|
||||
if getattr(trainer, "use_per_ctx_average_loss", False):
|
||||
trainer.get_batch_samples = trainer.get_batch_samples_ctx
|
||||
# if getattr(trainer, "use_per_ctx_average_loss", False):
|
||||
# trainer.get_batch_samples = trainer.get_batch_samples_ctx
|
||||
|
||||
# MONKEY PATCH: remove embedding layers from weight decay
|
||||
trainer.get_decay_parameter_names = get_decay_parameter_names
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue