olmo_tap.experiments.uncertainty.weights_handler

Helper class to handle cycling through frozen LLM heads during uncertainty finetuning.

Classes

FrozenHeadHandler(model, prod_config, ...)

During uncertainty head finetuning we cycle through randomly sampled frozen LLM heads (frozen meaning no grad or trainable LoRA weights).
class olmo_tap.experiments.uncertainty.weights_handler.FrozenHeadHandler(model: HydraTransformer, prod_config: HydraLoRAConfig, robust_config: HydraLoRAConfig, prod_dir: Path, robust_dir: Path, n_frozen: int)[source]

Bases: object

During uncertainty head finetuning we cycle through randomly sampled frozen LLM heads (frozen meaning no grad or trainable LoRA weights). This class manages the loading and unloading of different heads.

Parameters:

  • model – Hydra transformer model to be trained.

  • prod_config – config for Hydra production (security) LoRA weights.

  • robust_config – config for Hydra robustness LoRA weights.

  • prod_dir – directory storing production (security) LoRA weights.

  • robust_dir – directory storing robustness LoRA weights.

  • n_frozen – number of frozen LLM heads in total.

NOTE: only one frozen LLM head is ever loaded at a given time, n_frozen refers to the total number of heads available to cycle through.

swap_to_expert(frozen_idx: int)[source]

Restores the base head and merges the new frozen head LoRA weights.