Skip to content

llmcompressor.entrypoints.model_free.microscale

Functions:

build_microscale_inverse_weight_maps 

build_microscale_inverse_weight_maps(
    weight_map: dict[str, str],
    model_files: dict[str, str],
    converters: list[Converter],
) -> dict[str, InverseWeightMap]

This function replicates the logic of compressed_tensors.entrypoints.convert.build_inverse_weight_maps including the case of microscale partner shards, as defined in DEFAULT_FUSED_MAPPINGS

For a given output shard, precompute exactly which tensors to load from which source files — including required partner tensors from other shards.

This is necessary because some converters require that a set of tensors are accessible in order for them to be processed correctly.

Parameters:

  • shard_name –

    the shard filename this job will process and save

  • weight_map (dict[str, str]) –

    tensor name -> shard filename (from safetensors.index.json)

  • model_files (dict[str, str]) –

    shard filename -> resolved absolute path

Returns:

  • dict[str, InverseWeightMap] –

    {resolved_file_path: [tensor_names_to_load]}

Source code in src/llmcompressor/entrypoints/model_free/microscale.py 
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
def build_microscale_inverse_weight_maps(
    weight_map: dict[str, str],
    model_files: dict[str, str],
    converters: list[Converter],
) -> dict[str, InverseWeightMap]:
    """
    This function replicates the logic of
    `compressed_tensors.entrypoints.convert.build_inverse_weight_maps` including the
    case of microscale partner shards, as defined in DEFAULT_FUSED_MAPPINGS

    For a given output shard, precompute exactly which tensors to load from
    which source files — including required partner tensors from other shards.

    This is necessary because some converters require that a set of tensors are
    accessible in order for them to be processed correctly.

    :param shard_name: the shard filename this job will process and save
    :param weight_map: tensor name -> shard filename (from safetensors.index.json)
    :param model_files: shard filename -> resolved absolute path
    :return: {resolved_file_path: [tensor_names_to_load]}
    """

    # TODO move to top level, fulfill rest of Protocol contract
    #  - ideally remove the need for separate process_file and process_microscale_file
    #    functions
    #  - remove build_microscale_inverse_weight_maps entirely, replace with
    #    `compressed_tensors.entrypoints.convert.build_inverse_weight_maps`
    class MicroscaleConverter:
        def get_dependencies(self, weight_name: str) -> set[str]:
            deps = set()
            for primary_pattern, partner_templates in DEFAULT_FUSED_MAPPINGS.items():
                match = re.match(primary_pattern, weight_name)
                if match is None:
                    continue

                # Build partner names using named groups from the match
                for partner_template in partner_templates:
                    partner_name = partner_template.format(**match.groupdict())

                    deps.add(partner_name)
            return deps

    converters.append(MicroscaleConverter())
    return build_inverse_weight_maps(weight_map, model_files, converters)