def group_qkv_layers(state_dict: dict) -> dict:
    """Find corresponding layer names for query, key and value layers and stack them in a single layer"""

    state_dict = state_dict.copy()  # shallow copy

    result = defaultdict(dict)
    layer_names = list(state_dict.keys())
    qkv_names = ["query", "key", "value"]
    for layer_name in layer_names:
        for pattern in qkv_names:
            if pattern in layer_name:
                new_key = layer_name.replace(pattern, "qkv")
                result[new_key][pattern] = state_dict.pop(layer_name)
                break

    # merge them all
    for new_key, patterns in result.items():
        state_dict[new_key] = torch.cat(
            [patterns[qkv_name] for qkv_name in qkv_names], dim=0
        )

    return state_dict
Find corresponding layer names for query, key and value layers and stack them in a single layer
scripts/models-conversions/upernet-hf-to-smp.py