def get_relevant_parameter_names()

in optimum/neuron/models/training/transformations_utils.py [0:0]

• 46 lines of code
• 14 McCabe index (conditional complexity)

    def get_relevant_parameter_names(self, module_fully_qualified_name: str) -> set[str]:
        param_names = ["weight", "bias"] if self.bias else ["weight"]
        original_proj_names = [
            self.query_projection_name,
            self.key_projection_name,
            self.value_projection_name,
            self.output_projection_name,
        ]
        original_names = {
            f"{module_fully_qualified_name}.{proj_name}.{param_name}"
            for proj_name in original_proj_names
            for param_name in param_names
        }

        if self.fuse_qkv:
            gqa_names = {
                f"{module_fully_qualified_name}.{self.gqa_qkv_projection_name}.{param_name}_qkv"
                for param_name in param_names
            }
        else:
            gqa_names = {
                f"{module_fully_qualified_name}.{self.gqa_qkv_projection_name}.{param_name}_{suffix}"
                for param_name in param_names
                for suffix in ["q", "k", "v"]
            }

        lora_param_names = set()
        # LoRA for original layers
        for param_name in param_names:
            for proj_name in original_proj_names:
                lora_param_names.add(f"{module_fully_qualified_name}.{proj_name}.lora_A.{param_name}")
                lora_param_names.add(f"{module_fully_qualified_name}.{proj_name}.lora_B.{param_name}")

        # LoRA for GQA layer
        for param_name in param_names:
            lora_param_names.add(f"{module_fully_qualified_name}.{self.gqa_qkv_projection_name}.lora_A.{param_name}")
            if self.fuse_qkv:
                lora_param_names.add(
                    f"{module_fully_qualified_name}.{self.gqa_qkv_projection_name}.lora_B.{param_name}"
                )
            else:
                lora_param_names.add(
                    f"{module_fully_qualified_name}.{self.gqa_qkv_projection_name}.lora_B.{param_name}_q"
                )
                lora_param_names.add(
                    f"{module_fully_qualified_name}.{self.gqa_qkv_projection_name}.lora_B.{param_name}_k"
                )
                lora_param_names.add(
                    f"{module_fully_qualified_name}.{self.gqa_qkv_projection_name}.lora_B.{param_name}_v"
                )
        return original_names | gqa_names | lora_param_names