aquin.compute.hf_llm_shim Namespace Reference

Classes
class	_ShimCfg
class	_IdentityNorm
class	_UnembedShim
class	HfLlmShim

Functions
bool	_names_filter_match (Any names_filter, str name)
str|dict[str, Any]	_hf_load_device_map (dict[str, Any] cfg, str device)
bool	_model_has_device_map (Any hf_model)
tuple[Any, Any, Any, Any]	_resolve_modules (Any hf_model)
torch.Tensor	project_residual_to_logits (Any model, torch.Tensor resid)

Variables
	_RESID_POST = re.compile(r"^blocks\.(\d+)\.hook_resid_post$")
	_RESID_PRE = re.compile(r"^blocks\.(\d+)\.hook_resid_pre$")

Function Documentation

_hf_load_device_map()

str | dict[str, Any] _hf_load_device_map ( dict[str, Any] cfg, str device )

protected

Use device_map='auto' for large HF-native models (e.g. Sarvam 30B MoE).

Definition at line 46 of file hf_llm_shim.py.

Referenced by aquin.compute.hf_llm_shim.HfLlmShim.from_pretrained().

_model_has_device_map()

bool _model_has_device_map ( Any hf_model )

protected

Definition at line 56 of file hf_llm_shim.py.

Referenced by aquin.compute.hf_llm_shim.HfLlmShim.to().

_names_filter_match()

bool _names_filter_match ( Any names_filter, str name )

protected

Definition at line 34 of file hf_llm_shim.py.

Referenced by aquin.compute.hf_llm_shim.HfLlmShim._hooks_from_filter().

_resolve_modules()

tuple[Any, Any, Any, Any] _resolve_modules ( Any hf_model )

protected

Return (backbone, layer_list, embed_module, lm_head).

Definition at line 60 of file hf_llm_shim.py.

project_residual_to_logits()

torch.Tensor project_residual_to_logits	(	Any	model,
		torch.Tensor	resid )

Map a residual-space direction to vocab logits (TL or HF lm_head layout).

Definition at line 111 of file hf_llm_shim.py.

Variable Documentation

_RESID_POST

aquin.compute.hf_llm_shim._RESID_POST = re.compile(r"^blocks\.(\d+)\.hook_resid_post$")

protected

Definition at line 22 of file hf_llm_shim.py.

_RESID_PRE

aquin.compute.hf_llm_shim._RESID_PRE = re.compile(r"^blocks\.(\d+)\.hook_resid_pre$")

protected

Definition at line 23 of file hf_llm_shim.py.