micm_nlp.training.callbacks

Trainer callbacks.

• CustomEarlyStoppingCallback — early stopping decoupled from checkpoint selection and gated by an early_stopping_after floor, so a run cannot stop during its first fraction of steps.

• NormalizePromptEncoderEmbeddings — normalises prompt-encoder embeddings each step and logs the mean norm.

• ParamNormLogger — logs parameter and parameter-update norms to W&B.

• EmptyCudaCacheCallback — periodic torch.cuda.empty_cache().

• DownstreamFineTuningCallback — fine-tunes on downstream tasks at evaluation and save points.

NormalizePromptEncoderEmbeddings only registers when peft.encoder_embedding_normalize is set to 'unit' or 'clip'; without it the callback would be a no-op that still logged a zero norm every step.

Note

Until 0.2.0 this callback was never registered at all — the trainer looked for its settings under task.peft while peft is a top-level block. Results produced before that fix did not normalise, whatever their config said.

Classes

CustomEarlyStoppingCallback	Early stopping decoupled from model selection, gated by an
DownstreamFineTuningCallback	A custom callback that performs downstream fine-tuning on evaluation.
EmptyCudaCacheCallback	A custom callback that empties the CUDA cache at specified intervals.
NormalizePromptEncoderEmbeddings	A custom callback that normalizes the prompt encoder embeddings.
ParamNormLogger	Log the mean parameter norm and mean per-step update norm to wandb.

Module Contents

class micm_nlp.training.callbacks.CustomEarlyStoppingCallback(early_stopping_patience=5, early_stopping_threshold=0.0, early_stopping_after=0.5, early_stopping_metric='eval_loss')

Bases: transformers.EarlyStoppingCallback

Early stopping decoupled from model selection, gated by an early_stopping_after floor (fraction of max_steps before stopping is allowed).

The monitored signal is chosen by early_stopping_metric:

• 'metric_for_best_model' (sentinel): delegate to the Trainer’s args.metric_for_best_model + args.greater_is_better — i.e. stop on the same metric used to pick the best checkpoint.

• any other string (e.g. 'eval_loss'): treat it as a literal metric key; direction inferred ('loss' in the name → lower-is-better, else greater-is-better). This preserves the original eval_loss behavior and keeps early stopping SEPARABLE from selection.

Default 'eval_loss' reproduces the pre-existing behavior.

Parameters:

• early_stopping_patience – evaluations without improvement before stopping.

• early_stopping_threshold – how much counts as an improvement.

• early_stopping_after – fraction of max_steps that must elapse before stopping is allowed at all, so a slow start is not cut short.

• early_stopping_metric – metric to monitor – a literal key, or the sentinel 'metric_for_best_model' to follow the Trainer’s own selection metric.

on_evaluate(args, state, control, metrics, **kwargs)

Consider stopping, once past the early_stopping_after floor.

Evaluations before that step return unchanged, so patience is not spent while the model is still warming up.

SENTINEL_BEST = 'metric_for_best_model'

best_metric = None

early_stopping_after = 0.5

early_stopping_metric = 'eval_loss'

patience_counter = 0

class micm_nlp.training.callbacks.DownstreamFineTuningCallback(config, model_path)

Bases: transformers.TrainerCallback

A custom callback that performs downstream fine-tuning on evaluation.

Parameters:

• config – the run config; eval.downstream_tasks steers this callback.

• model_path – directory of the checkpoint to fine-tune from.

finetune_on_downstream_tasks(state_global_step)

Finetune the model on all downstream tasks.

on_evaluate(args: transformers.TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, **kwargs)

Run the downstream fine-tuning probe, if this evaluation warrants one.

Two cases qualify: the step-0 evaluation of an already-pretrained model (a baseline before any training), and any evaluation outside training.

on_save(args: transformers.TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, **kwargs)

Probe the just-saved checkpoint on the downstream tasks.

Skipped at step 0, and during training when eval.downstream_tasks.not_while_training is set – the probe is expensive, so a run can defer every one of them to the end.

on_train_begin(args: transformers.TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, **kwargs)

Mark training as in progress, so evaluations can tell which phase they are in.

on_train_end(args: transformers.TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, **kwargs)

Mark training as finished.

class micm_nlp.training.callbacks.EmptyCudaCacheCallback(empty_cache_steps=None)

Bases: transformers.TrainerCallback

A custom callback that empties the CUDA cache at specified intervals.

Parameters:

empty_cache_steps – free the cache every N steps; None disables.

on_step_end(args: transformers.TrainingArguments, state: transformers.TrainerState, control: transformers.TrainerControl, **kwargs)

Empty the CUDA allocator cache on the configured step interval.

device

empty_cache_steps = None

gb_coeff = 1073741824

class micm_nlp.training.callbacks.NormalizePromptEncoderEmbeddings

Bases: transformers.TrainerCallback

A custom callback that normalizes the prompt encoder embeddings.

on_step_end(args, state, control, **kwargs)

Renormalise the prompt encoder’s embeddings and log their mean norm.

Does nothing unless the active adapter’s prompt encoder is a CrossPromptEncoder; the encoder itself decides what normalisation means, from its own config.

class micm_nlp.training.callbacks.ParamNormLogger

Bases: transformers.TrainerCallback

Log the mean parameter norm and mean per-step update norm to wandb.

A training-stability probe: a norm that grows without bound, or an update norm that collapses, shows up here before it shows up in the loss. Only trainable parameters are measured – with PEFT that is the adapter, not the backbone.

It keeps a CPU copy of every trainable parameter between steps in order to difference them, so memory scales with the trainable parameter count.

Start with no previous step recorded; the first step logs no update norm.

on_step_end(args, state, control, **kwargs)

Measure this step’s parameter norms and log them.

prev_params