def auto_configure_device_map(num_gpus: int, use_v2: bool) -> Dict[str, int]:
    r"""
    Configures device map for ChatGLM.

    Borrowed from: https://github.com/THUDM/ChatGLM-6B/blob/dev_multi_gpu/utils.py#L8
    """
    num_layers = 28
    layers_per_gpu = 30 / num_gpus
    if use_v2:
        device_map = {
            "transformer.embedding.word_embeddings": 0,
            "transformer.encoder.final_layernorm": 0,
            "transformer.output_layer": 0,
            "transformer.rotary_pos_emb": 0,
            "transformer.prefix_encoder": 0,
            "lm_head": 0
        }
    else:
        device_map = {
            "transformer.word_embeddings": 0,
            "transformer.final_layernorm": 0,
            "transformer.prefix_encoder": 0,
            "lm_head": 0
        }

    added_layers = 2
    target_gpu = 0

    for i in range(num_layers):
        if added_layers >= layers_per_gpu:
            target_gpu += 1
            added_layers = 0
        assert target_gpu < num_gpus
        if use_v2:
            device_map[f"transformer.encoder.layers.{i}"] = target_gpu
        else:
            device_map[f"transformer.layers.{i}"] = target_gpu
        added_layers += 1

    return device_map

    Configures device map for ChatGLM.

    Borrowed from: https://github.com/THUDM/ChatGLM-6B/blob/dev_multi_gpu/utils.py#L8
    
src/glmtuner/extras/misc.py