fix typo.

alibaba · Jan 30, 2025 · bff3fdc · bff3fdc
1 parent e12d8fd
commit bff3fdc
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/chatlearn/synchronizer/megatron_vllm.py b/chatlearn/synchronizer/megatron_vllm.py
@@ -327,9 +327,9 @@ def regroup_qkv_tp_slices(self, name, param_data, tp_division):
         param_data_shape = param_data.shape
         # Regroup qkv tensors into different tp slices only for inference model which enables vLLM backend.
         to_fix_qkv_ordering_dict = self.sync_map.to_fix_qkv_ordering_dict
+        # pylint: disable=too-many-nested-blocks
         if "attention.query_key_value" in name or \
                 "self_attention.query_key_value" in name or \
-                "self_attention.linear_qkv" in name: # pylint: disable=too-many-nested-blocks
             src_tp_size = self.src_module_args.args_dict["tensor_model_parallel_size"]
             dst_tp_size = self.dst_module_args.args_dict["tensor_model_parallel_size"]
             heads = self.src_module_args.args_dict["num_attention_heads"] // src_tp_size