[Help Wanted] Refactor module / parameter matching logic

[kylesayrs]

• There currently exist 7 different functions for matching modules/parameters. Let's consolidate this logic by using only the functions defined in utils/match.py

Function	Used for/by	Maps to
apply_quantization_config	Quantization	match_named_modules
expand_target_names	Compressor	match_named_parameters
is_target	Quantization	is_match
find_name_or_class_matches	Quantization	match_name or match_class
_find_matches	Quantization	Delete
get_nested_weight_mappings	Compressor	match_named_parameters
match_param_name	get_nested_weight_mappings	Delete,match_name

[fynnsu]

I can work on this

[fynnsu]

@kylesayrs Regarding get_nested_weight_mappings and match_param_name, I'm not sure there's much redundancy to remove using the match.py utilities.

get_nested_weight_mappings
Is mostly just building nested dicts and doesn't recurse through the model parameters so we can't use match_named_parameters. We also don't have access to the Module in either this function or the function it's dependent on get_weight_mappings:

compressed-tensors/src/compressed_tensors/utils/safetensors_load.py

Lines 183 to 249 in 3d49764

	def get_nested_weight_mappings(
	model_path: str,
	params_to_nest: Iterable[str],
	return_unmatched_params: bool = False,
	) -> Union[NestedWeightMappingType, Tuple[NestedWeightMappingType, WeightMappingType]]:
	"""
	Takes a path to a state dict saved in safetensors format and returns a nested
	mapping from uncompressed parameterized layer names to the file locations of
	each layer's compression parameters.
	Example of the nested mapping:
	layer: {
	bitmask: file_location,
	row_offsets: file_location,
	shape: file_location,
	compressed: file_location
	}
	If other parameters are found that do not match the nested parameters, they will
	be returned in a separate dictionary only if return_unmatched_params is True.
	This dictionary may be needed for cases where compressors are stacked (e.g.,
	quantization compression followed by sparse compression).
	Example of the unmatched params mapping:
	{
	layer.weight_scale: file_location,
	layer.input_scale: file_location
	}
	This generalizes to cases where the model is split into multiple safetensors
	files.
	:param model_path: Path to the safetensors state dict, must contain either a
	single safetensors file or multiple files with an index.
	:param params_to_nest: Iterable of parameter names to nest.
	:param return_unmatched_params: If True, return a second dictionary containing
	the remaining parameters that were not matched to the params_to_nest.
	:return:
	- If return_unmatched_params is False:
	NestedWeightMappingType: A nested mapping of parameterized layer names to
	file locations of each layer's compression parameters.
	- If return_unmatched_params is True:
	Tuple[NestedWeightMappingType, WeightMappingType]: A tuple containing:
	- NestedWeightMappingType: A nested mapping of parameterized layer
	names to file locations of each layer's compression parameters.
	- WeightMappingType: A mapping of the remaining parameter names to
	their file locations that were not matched to the params_to_nest.
	"""
	weight_mappings = get_weight_mappings(model_path)
	nested_weight_mappings = {}
	unmatched_params = {}
	for key, file_location in weight_mappings.items():
	matched = False
	for param_name in params_to_nest:
	module_path = match_param_name(key, param_name)
	if module_path:
	if module_path not in nested_weight_mappings:
	nested_weight_mappings[module_path] = {}
	nested_weight_mappings[module_path][param_name] = file_location
	matched = True
	if return_unmatched_params and not matched:
	unmatched_params[key] = file_location
	if return_unmatched_params:
	return nested_weight_mappings, unmatched_params
	return nested_weight_mappings

match_param_names is mostly a utility for extracting the submodule portion of a parameter name:

compressed-tensors/src/compressed_tensors/utils/safetensors_load.py

Lines 100 to 113 in 3d49764

	def match_param_name(full_name: str, param_name: str) -> Optional[str]:
	"""
	Helper function extracting the uncompressed parameterized layer name from a
	compressed name. Assumes the compressed name was merged using merge_names.
	:param full_name: full name of parameter in compressed model
	:param param_name: compression paramater name
	:return: uncompressed name of the uncompressed parameterized layer
	"""
	pattern = r"^(.*)\." + param_name + r"$"
	regex = re.findall(pattern, full_name)
	if len(regex) == 0:
	return None
	return regex[0]