pandora_llm.features.GradNorm

Module Contents

class pandora_llm.features.GradNorm.GradNorm(*args, **kwargs)[source]

Bases: pandora_llm.features.base.FeatureComputer, pandora_llm.features.base.LLMHandler

GradNorm thresholding attack

compute_features(dataloader, accelerator, norms=[1, 2, float('inf')], gradient_checkpointing=False)[source]

Compute the layerwise gradient norms for a given dataloader, using the specified norms.

Parameters:

  • dataloader (torch.utils.data.DataLoader) – input data to compute statistic over

  • accelerator (accelerate.Accelerator) – accelerator object

  • norms (Iterable[Union[int, float]]) – list of norm orders

  • gradient_checkpointing (bool) – whether to use gradient checkpointing to save memory

Returns:

grad norm of input IDs

Return type:

torch.Tensor or list

pandora_llm.features.GradNorm.compute_gradnorms_dl(model, dataloader, accelerator, norms=[1, 2, float('inf')], gradient_checkpointing=False)[source]

Computes gradient norms of dataloader.

Parameters:
Returns:

Gradient norms of each layer, in a tensor of shape [n x layer x norm]

Return type:

jaxtyping.Float[torch.Tensor, n layer norm]

pandora_llm.features.GradNorm.compute_gradnorms(model, input_ids, accelerator, norms=[1, 2, float('inf')], attention_mask=None)[source]

Computes gradient norm of each layer Note: takes advantage of the fact that norm([a,b])=norm([norm(a),norm(b)])

Parameters:

  • model (transformers.AutoModelForCausalLM) – HuggingFace model

  • input_ids (jaxtyping.Integer[torch.Tensor, batch seq]) – tensor of input IDs

  • accelerator (accelerate.Accelerator) – Accelerator object

  • norms (Iterable[Union[int, float]]) – Gradient norm types

  • attention_mask (jaxtyping.Bool[torch.Tensor, batch seq]) – Attention mask

Returns:

Gradient norms of each layer, in a tensor of shape [batch x layer x norm]

Return type:

jaxtyping.Float[torch.Tensor, batch layer norm]