aquin.compute.weight_rank Namespace Reference

Functions
bool	_is_heavy_weight_model (model)
bool	_should_skip_weight_param (str name, *, bool heavy)
int	_layer_from_key (str key)
str	_hf_matrix_label (str key)
object	_weight_root (model)
bool	_is_hf_weight_model (model)
float	_scalar (torch.Tensor t)
torch.Tensor	_to_2d (torch.Tensor w)
tuple[float, float]	_matrix_ranks (torch.Tensor w)
dict	run_weight_rank_hf (model, float collapse_threshold=0.1)
dict	run_weight_rank_tl (model, float collapse_threshold=0.1)
dict	run_weight_rank_analysis (model, float collapse_threshold=0.1)

Variables
	_LAYER_RE = re.compile(r"(?:^|\.)(?:layers?|layer|block)\.?(\d+)(?:\.|$)", re.I)
	_EXPERT_WEIGHT_RE = re.compile(r"\.experts?\.\d+", re.I)

Function Documentation

_hf_matrix_label()

str _hf_matrix_label ( str key )

protected

Definition at line 54 of file weight_rank.py.

Referenced by run_weight_rank_hf().

_is_heavy_weight_model()

bool _is_heavy_weight_model ( model )

protected

Large / MoE models fill VRAM :  weight analysis must run on CPU and skip experts.

Definition at line 22 of file weight_rank.py.

Referenced by run_weight_rank_hf().

_is_hf_weight_model()

bool _is_hf_weight_model ( model )

protected

Definition at line 107 of file weight_rank.py.

Referenced by run_weight_rank_analysis().

_layer_from_key()

int _layer_from_key ( str key )

protected

Definition at line 49 of file weight_rank.py.

Referenced by run_weight_rank_hf().

_matrix_ranks()

tuple[float, float] _matrix_ranks ( torch.Tensor w )

protected

Definition at line 133 of file weight_rank.py.

References _scalar(), and _to_2d().

Referenced by run_weight_rank_hf(), and run_weight_rank_tl().

_scalar()

float _scalar ( torch.Tensor t )

protected

Definition at line 119 of file weight_rank.py.

Referenced by _matrix_ranks().

_should_skip_weight_param()

bool _should_skip_weight_param ( str name, * , bool heavy )

protected

Skip tensors that are redundant or too large for in-GPU analysis.

Definition at line 35 of file weight_rank.py.

Referenced by run_weight_rank_hf().

_to_2d()

torch.Tensor _to_2d ( torch.Tensor w )

protected

HookedTransformer attn weights are often 3D (n_heads, d_in, d_out).

Definition at line 123 of file weight_rank.py.

Referenced by _matrix_ranks().

_weight_root()

object _weight_root ( model )

protected

Definition at line 101 of file weight_rank.py.

Referenced by run_weight_rank_hf().

run_weight_rank_analysis()

dict run_weight_rank_analysis	(		model,
		float	collapse_threshold = 0.1 )

Compute stable rank for attention/MLP (TL) or HF-native weights.

Definition at line 237 of file weight_rank.py.

References _is_hf_weight_model(), run_weight_rank_hf(), and run_weight_rank_tl().

run_weight_rank_hf()

dict run_weight_rank_hf	(		model,
		float	collapse_threshold = 0.1 )

Stable rank over HuggingFace parameter names (LFM, Sarvam, etc.).

Definition at line 154 of file weight_rank.py.

References _hf_matrix_label(), _is_heavy_weight_model(), _layer_from_key(), _matrix_ranks(), _should_skip_weight_param(), and _weight_root().

Referenced by run_weight_rank_analysis().

run_weight_rank_tl()

dict run_weight_rank_tl	(		model,
		float	collapse_threshold = 0.1 )

Compute stable rank for Q/K/V/O and MLP matrices across all layers.

Definition at line 197 of file weight_rank.py.

References _matrix_ranks().

Referenced by run_weight_rank_analysis().

Variable Documentation

_EXPERT_WEIGHT_RE

aquin.compute.weight_rank._EXPERT_WEIGHT_RE = re.compile(r"\.experts?\.\d+", re.I)

protected

Definition at line 19 of file weight_rank.py.

_LAYER_RE

aquin.compute.weight_rank._LAYER_RE = re.compile(r"(?:^|\.)(?:layers?|layer|block)\.?(\d+)(?:\.|$)", re.I)

protected

Definition at line 18 of file weight_rank.py.