Function invert_module_mapping

tensorrt_llm/lora_manager.py:211–230

Invert module mapping from TensorRT LLM -> HF to HF -> TensorRT-LLM. Args: trtllm_modules_to_hf_modules: Mapping from TensorRT LLM module names to HF module names (values can be strings or lists of strings) Returns: Dictionary mapping HF

(
    trtllm_modules_to_hf_modules: Dict[str, Union[str, List[str]]],
)

def invert_module_mapping(
    trtllm_modules_to_hf_modules: Dict[str, Union[str, List[str]]],
) -> Dict[str, str]:
    """Invert module mapping from TensorRT LLM -> HF to HF -> TensorRT-LLM.

    Args:
        trtllm_modules_to_hf_modules: Mapping from TensorRT LLM module names to HF module names
                                     (values can be strings or lists of strings)

    Returns:
        Dictionary mapping HF module names to TensorRT LLM module names
    """
    hf_modules_to_trtllm_modules: Dict[str, str] = {}
    for k, hf_modules in trtllm_modules_to_hf_modules.items():
        if isinstance(hf_modules, list):
            for hf_module in hf_modules:
                hf_modules_to_trtllm_modules[hf_module] = k
        else:
            hf_modules_to_trtllm_modules[hf_modules] = k
    return hf_modules_to_trtllm_modules


def norm_dora_magnitude(