wehub-resource-sync
45 lines
1.6 KiB
Python
45 lines
1.6 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
|
|
def parse_config_yaml(yaml_data):
|
|
# Initialize to default options.
|
|
quantization_options = {
|
|
"n_centroids": {
|
|
"Linear": ["in_features", {"*": 256}],
|
|
"Embedding": ["embedding_dim", {"*": 256}],
|
|
},
|
|
"block_sizes": {
|
|
"Linear": ["fuzzy_name", {"fc": 8, "attn": 4, "emb": 4}],
|
|
"Embedding": ["fuzzy_name", {"emb": 8}],
|
|
},
|
|
"layers_to_quantize": [
|
|
"decoder\\.layers\\.\\d+\\.fc[12]",
|
|
"decoder\\.embed_tokens\\.embeddings\\.[012]\\.[01]",
|
|
"decoder\\.layers\\.\\d+\\.self_attn\\.(k_proj|v_proj|q_proj|out_proj)",
|
|
],
|
|
}
|
|
|
|
if "n_centroids" in yaml_data:
|
|
quantization_options["n_centroids"] = {
|
|
layer: convert_yaml_to_tuple(layer_data)
|
|
for layer, layer_data in yaml_data["n_centroids"].items()
|
|
}
|
|
if "block_sizes" in yaml_data:
|
|
quantization_options["block_sizes"] = {
|
|
layer: convert_yaml_to_tuple(layer_data)
|
|
for layer, layer_data in yaml_data["block_sizes"].items()
|
|
}
|
|
if "layers_to_quantize" in yaml_data:
|
|
quantization_options["layers_to_quantize"] = yaml_data["layers_to_quantize"]
|
|
|
|
return quantization_options
|
|
|
|
|
|
def convert_yaml_to_tuple(yaml_dictionary):
|
|
"""Converts a yaml dictionary with two keys: `key` and `value` into a two
|
|
argument tuple of those values."""
|
|
return (yaml_dictionary["key"], yaml_dictionary["value"])
|