Fix/longlora (#294)

Author: tastelikefeet

swift/llm/infer.py

@@ -138,7 +138,8 @@ def prepare_model_template(
     logger.info(f'generation_config: {generation_config}')
     set_generation_config(model, generation_config)
     # Preparing LoRA
-    if args.sft_type == 'lora' and args.ckpt_dir is not None:
+    if args.sft_type in ('lora', 'qalora',
+                         'longlora') and args.ckpt_dir is not None:
         model = Swift.from_pretrained(
             model, args.ckpt_dir, inference_mode=True)
 

swift/trainers/trainers.py

@@ -104,7 +104,12 @@ def prediction_step(
         generate_inputs = inputs.copy()
         if has_labels:
             _labels = inputs['labels'][0]
-            n_mask = lower_bound(0, len(_labels), lambda i: _labels[i] != -100)
+            n_mask = 0
+            for i in range(len(_labels)):
+                if _labels[i] != -100:
+                    n_mask = i
+                    break
+
             for k in ['input_ids', 'attention_mask']:
                 generate_inputs[k] = generate_inputs[k][:, :n_mask]
             generate_inputs['labels'] = generate_inputs['labels'][:, n_mask:]

swift/tuners/longlora/llama.py

@@ -14,6 +14,10 @@
 from transformers.models.llama.modeling_llama import (apply_rotary_pos_emb,
                                                       repeat_kv, rotate_half)
 
+from swift.utils import get_logger
+
+logger = get_logger()
+
 
 def forward_flashattn(
     self,
@@ -306,8 +310,8 @@ def forward_flashattn_inference(
                ))  # noqa
 
     kv_seq_len = k.shape[1]
-    if past_key_value is not None:
-        past_kv_len = past_key_value[0].shape[2]
+    if past_key_value is not None and len(past_key_value):
+        past_kv_len = past_key_value.seen_tokens
         kv_seq_len += past_kv_len
 
     cos_sin = self.rotary_emb(v, seq_len=kv_seq_len)
@@ -316,15 +320,13 @@ def forward_flashattn_inference(
     q = q.transpose(1, 2)
     k = k.transpose(1, 2)
 
-    if past_key_value is not None:
-        assert (flash_attn_version >=
-                '2.1.0'), 'past_key_value support requires flash-attn >= 2.1.0'
-        # reuse k, v
-        k = torch.cat([past_key_value[0].transpose(1, 2), k], dim=1)
-        v = torch.cat([past_key_value[1].transpose(1, 2), v], dim=1)
-
-    past_key_value = (k.transpose(1, 2),
-                      v.transpose(1, 2)) if use_cache else None
+    if use_cache:
+        k, v = past_key_value.update(
+            k.transpose(1, 2), v.transpose(1, 2), layer_idx=self.idx)
+        k = k.transpose(1, 2)
+        v = v.transpose(1, 2)
+    else:
+        past_key_value = None
 
     if attention_mask is None:
         output = flash_attn_func(
@@ -405,12 +407,13 @@ def forward_flashattn_inference_s2_attn(
 
 def patch_llama_forward(model: nn.Module, forward_function) -> None:
     # Compatible with transformers device_map
-    for m in model.model.layers:
+    for idx, m in enumerate(model.model.layers):
         new_forward = MethodType(forward_function, m.self_attn)
         if hasattr(model, '_old_forward'):
             m.self_attn._old_forward = new_forward
         else:
             m.self_attn.forward = new_forward
+        m.self_attn.idx = idx
 
 
 def replace_llama_attn(model: nn.Module, use_flash_attn=True):
@@ -425,4 +428,7 @@ def replace_llama_attn(model: nn.Module, use_flash_attn=True):
             _prepare_decoder_attention_mask)
         patch_llama_forward(model, forward_flashattn_inference_s2_attn)
     else:
+        logger.warn(
+            'The source code of LongLoRA without flash '
+            'attention may has some problems, please use with careful.')
         patch_llama_forward(model, forward_noflashattn)

swift/tuners/longlora/longlora.py

@@ -8,6 +8,7 @@
 
 from swift import LoRA, LoRAConfig, SwiftOutput
 from swift.tuners.lora import lora_state_dict, mark_lora_as_trainable
+from swift.tuners.lora_layers import LoraModel
 
 
 class LongLoRAModelType:
@@ -59,17 +60,7 @@ class LongLoRA(LoRA):
     def prepare_model(model: nn.Module, config: LongLoRAConfig,
                       adapter_name: str):
         """Prepare a model with `LongLoRAConfig`"""
-        LoRA._dynamic_patch_lora(
-            model,
-            target_modules=config.target_modules,
-            r=config.r,
-            adapter_name=adapter_name,
-            lora_alpha=config.lora_alpha,
-            lora_dropout=config.lora_dropout,
-            merge_weights=config.merge_weights,
-            use_merged_linear=config.use_merged_linear,
-            enable_lora=config.enable_lora,
-            fan_in_fan_out=config.fan_in_fan_out)
+        LoraModel(model, config, adapter_name)
 
         def state_dict_callback(state_dict, adapter_name):
             _state_dict = lora_state_dict(state_dict, adapter_name,