searcher

Searcher that calibrates sparse attention thresholds.

Classes

SparseAttentionSearcher

Calibrate sparse attention thresholds with the forward loop of the search config.

class SparseAttentionSearcher

Bases: BaseSearcher

Calibrate sparse attention thresholds with the forward loop of the search config.

Calibration is configured by sparse_cfg["calibration"] of the model’s sparse_attention mode config and is skipped without it. Without a forward_loop or data_loader, a RULER calibration dataset is generated from the model’s tokenizer.

property default_state_dict: dict[str, Any]

Return default state dict.

run_search()

Calibrate the sparse attention modules of the model.

Return type:

None