fix: pass attention_mask directly to model.generate(), not via GenerationConfig

fix_attention_mask_gen.py

@@ -0,0 +1,100 @@
+"""
+Fix: attention_mask is incorrectly added to gen_kwargs before GenerationConfig,
+but GenerationConfig does not accept attention_mask. It must be extracted and
+passed directly to model.generate(), just like the synced_gpus fix.
+"""
+import os, re
+
+REPOS = ["root_gainlora/src", "improve_gainlora/src"]
+
+# Pattern to find and fix
+OLD_BLOCK = (
+    '        if "attention_mask" in inputs:\n'
+    '            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)\n'
+    '\n'
+    '        generation_config = GenerationConfig(**gen_kwargs)\n'
+)
+NEW_BLOCK = (
+    '        attention_mask = inputs.get("attention_mask", None)\n'
+    '\n'
+    '        generation_config = GenerationConfig(**gen_kwargs)\n'
+)
+
+ALREADY = 'attention_mask = inputs.get("attention_mask", None)'
+
+# All three model.generate() patterns that need attention_mask added
+# Pattern A: encoder-decoder branch (first if)
+OLD_GEN_A = (
+    '            generated_tokens = self.model.generate(\n'
+    '                input_ids=generation_inputs, \n'
+    '                generation_config=generation_config,\n'
+    '                synced_gpus=synced_gpus,\n'
+    '            )\n'
+)
+NEW_GEN_A = (
+    '            generated_tokens = self.model.generate(\n'
+    '                input_ids=generation_inputs, \n'
+    '                generation_config=generation_config,\n'
+    '                attention_mask=attention_mask,\n'
+    '                synced_gpus=synced_gpus,\n'
+    '            )\n'
+)
+
+# Pattern B: LLaMA branch with input_ids_wo_label
+OLD_GEN_B = (
+    '                generated_tokens = self.model.generate(\n'
+    '                    input_ids=generation_inputs,\n'
+    '                    input_ids_wo_label=inputs["input_ids_wo_label"],\n'
+    '                    generation_config=generation_config,\n'
+    '                    synced_gpus=synced_gpus,\n'
+    '                )\n'
+)
+NEW_GEN_B = (
+    '                generated_tokens = self.model.generate(\n'
+    '                    input_ids=generation_inputs,\n'
+    '                    input_ids_wo_label=inputs["input_ids_wo_label"],\n'
+    '                    generation_config=generation_config,\n'
+    '                    attention_mask=attention_mask,\n'
+    '                    synced_gpus=synced_gpus,\n'
+    '                )\n'
+)
+
+# Pattern C: T5 (else branch, no input_ids_wo_label)
+OLD_GEN_C = (
+    '                generated_tokens = self.model.generate(\n'
+    '                    input_ids=generation_inputs,\n'
+    '                    generation_config=generation_config,\n'
+    '                    synced_gpus=synced_gpus,\n'
+    '                )\n'
+)
+NEW_GEN_C = (
+    '                generated_tokens = self.model.generate(\n'
+    '                    input_ids=generation_inputs,\n'
+    '                    generation_config=generation_config,\n'
+    '                    attention_mask=attention_mask,\n'
+    '                    synced_gpus=synced_gpus,\n'
+    '                )\n'
+)
+
+for repo in REPOS:
+    for fname in sorted(os.listdir(repo)):
+        if not fname.startswith("cl_trainer_") or not fname.endswith(".py"):
+            continue
+        fpath = os.path.join(repo, fname)
+        with open(fpath) as f:
+            src = f.read()
+        if ALREADY in src:
+            print(f"SKIP (already fixed): {fpath}")
+            continue
+        if OLD_BLOCK not in src:
+            print(f"SKIP (no old block): {fpath}")
+            continue
+        new_src = src.replace(OLD_BLOCK, NEW_BLOCK, 1)
+        new_src = new_src.replace(OLD_GEN_A, NEW_GEN_A)
+        new_src = new_src.replace(OLD_GEN_B, NEW_GEN_B)
+        new_src = new_src.replace(OLD_GEN_C, NEW_GEN_C)
+        with open(fpath, "w") as f:
+            f.write(new_src)
+        print(f"FIXED: {fpath}")
+
+print("Done.")

fix_attention_mask_gen2.py

@@ -0,0 +1,57 @@
+"""
+Fix v2: For all files where attention_mask is still in gen_kwargs before GenerationConfig.
+Uses regex to handle all ordering variants.
+"""
+import os, re
+
+REPOS = ["root_gainlora/src", "improve_gainlora/src"]
+OLD_FLAG = 'gen_kwargs["attention_mask"]'
+ALREADY = 'attention_mask = inputs.get("attention_mask", None)'
+
+def fix_file(fpath):
+    with open(fpath) as f:
+        src = f.read()
+
+    if ALREADY in src:
+        print(f"SKIP (already): {fpath}")
+        return
+    if OLD_FLAG not in src:
+        print(f"SKIP (no flag): {fpath}")
+        return
+
+    # Step 1: Replace the "gen_kwargs["attention_mask"] = ..." block with extraction
+    step1 = re.sub(
+        r'        if "attention_mask" in inputs:\n            gen_kwargs\["attention_mask"\] = inputs\.get\("attention_mask", None\)\n\n',
+        '        attention_mask = inputs.get("attention_mask", None)\n\n',
+        src,
+    )
+    if step1 == src:
+        # Variant without blank line after
+        step1 = re.sub(
+            r'        if "attention_mask" in inputs:\n            gen_kwargs\["attention_mask"\] = inputs\.get\("attention_mask", None\)\n',
+            '        attention_mask = inputs.get("attention_mask", None)\n',
+            src,
+        )
+
+    # Step 2: Add attention_mask= to every model.generate() call that doesn't have it
+    result = re.sub(
+        r'(self\.model\.generate\(\n(?:(?!attention_mask)(?!synced_gpus)[^\n]*\n)*?)(\s*synced_gpus=synced_gpus,\n\s*\))',
+        r'\1                attention_mask=attention_mask,\n\2',
+        step1,
+    )
+
+    if result == src:
+        print(f"WARNING: no change for {fpath}")
+        return
+
+    with open(fpath, "w") as f:
+        f.write(result)
+    print(f"FIXED: {fpath}")
+
+for repo in REPOS:
+    for fname in sorted(os.listdir(repo)):
+        if not fname.startswith("cl_trainer_") or not fname.endswith(".py"):
+            continue
+        fix_file(os.path.join(repo, fname))
+
+print("Done.")

improve_gainlora/src/cl_trainer_gainlora_inflora.py

@@ -943,8 +943,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
                 
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         synced_gpus = gen_kwargs.pop("synced_gpus", False)
         generation_config = GenerationConfig(**gen_kwargs)
@@ -958,6 +957,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
             generated_tokens = self.model.generate(
                 input_ids=generation_inputs, 
                 generation_config=generation_config,
+                attention_mask=attention_mask,
                 synced_gpus=synced_gpus,
             )
         else:
@@ -968,6 +968,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
                     input_ids=generation_inputs,
                     input_ids_wo_label=inputs["input_ids_wo_label"],
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
             
@@ -975,6 +976,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
                 generated_tokens = self.model.generate(
                     input_ids=generation_inputs,
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
 

improve_gainlora/src/cl_trainer_gainlora_inflora_llama.py

@@ -974,8 +974,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
                 
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         synced_gpus = gen_kwargs.pop("synced_gpus", False)
         generation_config = GenerationConfig(**gen_kwargs)
@@ -989,6 +988,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
             generated_tokens = self.model.generate(
                 input_ids=generation_inputs, 
                 generation_config=generation_config,
+                attention_mask=attention_mask,
                 synced_gpus=synced_gpus,
             )
         else:
@@ -999,6 +999,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
                     input_ids=generation_inputs,
                     input_ids_wo_label=inputs["input_ids_wo_label"],
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
             
@@ -1006,6 +1007,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
                 generated_tokens = self.model.generate(
                     input_ids=generation_inputs,
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
 

improve_gainlora/src/cl_trainer_gainlora_olora.py

@@ -982,8 +982,7 @@ class GainLoRA_OLoRA_Trainer(Seq2SeqTrainer):
                 
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         synced_gpus = gen_kwargs.pop("synced_gpus", False)
         generation_config = GenerationConfig(**gen_kwargs)
@@ -997,6 +996,7 @@ class GainLoRA_OLoRA_Trainer(Seq2SeqTrainer):
             generated_tokens = self.model.generate(
                 input_ids=generation_inputs, 
                 generation_config=generation_config,
+                attention_mask=attention_mask,
                 synced_gpus=synced_gpus,
             )
         else:
@@ -1007,6 +1007,7 @@ class GainLoRA_OLoRA_Trainer(Seq2SeqTrainer):
                     input_ids=generation_inputs,
                     input_ids_wo_label=inputs["input_ids_wo_label"],
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
             
@@ -1014,6 +1015,7 @@ class GainLoRA_OLoRA_Trainer(Seq2SeqTrainer):
                 generated_tokens = self.model.generate(
                     input_ids=generation_inputs,
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
 

improve_gainlora/src/cl_trainer_gainlora_olora_llama.py

@@ -965,8 +965,7 @@ class GainLoRA_OLoRA_Trainer(Seq2SeqTrainer):
                 
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         synced_gpus = gen_kwargs.pop("synced_gpus", False)
         generation_config = GenerationConfig(**gen_kwargs)
@@ -980,6 +979,7 @@ class GainLoRA_OLoRA_Trainer(Seq2SeqTrainer):
             generated_tokens = self.model.generate(
                 input_ids=generation_inputs, 
                 generation_config=generation_config,
+                attention_mask=attention_mask,
                 synced_gpus=synced_gpus,
             )
         else:
@@ -990,6 +990,7 @@ class GainLoRA_OLoRA_Trainer(Seq2SeqTrainer):
                     input_ids=generation_inputs,
                     input_ids_wo_label=inputs["input_ids_wo_label"],
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
             
@@ -997,6 +998,7 @@ class GainLoRA_OLoRA_Trainer(Seq2SeqTrainer):
                 generated_tokens = self.model.generate(
                     input_ids=generation_inputs,
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
 

improve_gainlora/src/cl_trainer_inflora.py

@@ -738,8 +738,7 @@ class InfLoRATrainer(Seq2SeqTrainer):
                 
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         synced_gpus = gen_kwargs.pop("synced_gpus", False)
         generation_config = GenerationConfig(**gen_kwargs)
@@ -753,6 +752,7 @@ class InfLoRATrainer(Seq2SeqTrainer):
             generated_tokens = self.model.generate(
                 input_ids=generation_inputs, 
                 generation_config=generation_config,
+                attention_mask=attention_mask,
                 synced_gpus=synced_gpus,
             )
         else:
@@ -763,6 +763,7 @@ class InfLoRATrainer(Seq2SeqTrainer):
                     input_ids=generation_inputs,
                     input_ids_wo_label=inputs["input_ids_wo_label"],
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
             
@@ -770,6 +771,7 @@ class InfLoRATrainer(Seq2SeqTrainer):
                 generated_tokens = self.model.generate(
                     input_ids=generation_inputs,
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
 

improve_gainlora/src/cl_trainer_inflora_llama.py

@@ -765,8 +765,7 @@ class InfLoRATrainer(Seq2SeqTrainer):
                 
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         synced_gpus = gen_kwargs.pop("synced_gpus", False)
         generation_config = GenerationConfig(**gen_kwargs)
@@ -780,6 +779,7 @@ class InfLoRATrainer(Seq2SeqTrainer):
             generated_tokens = self.model.generate(
                 input_ids=generation_inputs, 
                 generation_config=generation_config,
+                attention_mask=attention_mask,
                 synced_gpus=synced_gpus,
             )
         else:
@@ -790,6 +790,7 @@ class InfLoRATrainer(Seq2SeqTrainer):
                     input_ids=generation_inputs,
                     input_ids_wo_label=inputs["input_ids_wo_label"],
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
             
@@ -797,6 +798,7 @@ class InfLoRATrainer(Seq2SeqTrainer):
                 generated_tokens = self.model.generate(
                     input_ids=generation_inputs,
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
 

improve_gainlora/src/cl_trainer_olora.py

@@ -580,8 +580,7 @@ class OLoRATrainer(Seq2SeqTrainer):
                 
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         generation_config = GenerationConfig(**gen_kwargs)
 

improve_gainlora/src/cl_trainer_olora_llama.py

@@ -582,8 +582,7 @@ class OLoRATrainer(Seq2SeqTrainer):
                 
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         generation_config = GenerationConfig(**gen_kwargs)
 

improve_gainlora/src/cl_trainer_specroute.py

@@ -585,12 +585,11 @@ class SpecRoute_Trainer(Seq2SeqTrainer):
                 }
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         # synced_gpus and attention_mask must be passed to generate(), not GenerationConfig
         _synced_gpus = gen_kwargs.pop("synced_gpus", False)
-        _attention_mask = gen_kwargs.pop("attention_mask", None)
+        _attention_mask = inputs.get("attention_mask", None)  # from inputs, not gen_kwargs
 
         generation_config = GenerationConfig(**gen_kwargs)
 

root_gainlora/src/cl_trainer_gainlora_inflora.py

@@ -952,8 +952,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
                 
             synced_gpus = gen_kwargs.pop("synced_gpus", False)
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         generation_config = GenerationConfig(**gen_kwargs)
 
@@ -966,6 +965,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
             generated_tokens = self.model.generate(
                 input_ids=generation_inputs, 
                 generation_config=generation_config,
+                attention_mask=attention_mask,
                 synced_gpus=synced_gpus,
             )
         else:
@@ -976,6 +976,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
                     input_ids=generation_inputs,
                     input_ids_wo_label=inputs["input_ids_wo_label"],
                     generation_config=generation_config,
+                    attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
             
@@ -983,6 +984,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
                 generated_tokens = self.model.generate(
                     input_ids=generation_inputs,
                     generation_config=generation_config,
+                    attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
 

root_gainlora/src/cl_trainer_gainlora_inflora_llama.py

@@ -975,8 +975,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
                 
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         synced_gpus = gen_kwargs.pop("synced_gpus", False)
         generation_config = GenerationConfig(**gen_kwargs)
@@ -990,6 +989,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
             generated_tokens = self.model.generate(
                 input_ids=generation_inputs, 
                 generation_config=generation_config,
+                attention_mask=attention_mask,
                 synced_gpus=synced_gpus,
             )
         else:
@@ -1000,6 +1000,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
                     input_ids=generation_inputs,
                     input_ids_wo_label=inputs["input_ids_wo_label"],
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
             
@@ -1007,6 +1008,7 @@ class GainLoRA_InfLoRA_Trainer(Seq2SeqTrainer):
                 generated_tokens = self.model.generate(
                     input_ids=generation_inputs,
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
 

root_gainlora/src/cl_trainer_gainlora_olora.py

@@ -983,8 +983,7 @@ class GainLoRA_OLoRA_Trainer(Seq2SeqTrainer):
                 
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         generation_config = GenerationConfig(**gen_kwargs)
 

root_gainlora/src/cl_trainer_gainlora_olora_llama.py

@@ -966,8 +966,7 @@ class GainLoRA_OLoRA_Trainer(Seq2SeqTrainer):
                 
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         synced_gpus = gen_kwargs.pop("synced_gpus", False)
         generation_config = GenerationConfig(**gen_kwargs)
@@ -981,6 +980,7 @@ class GainLoRA_OLoRA_Trainer(Seq2SeqTrainer):
             generated_tokens = self.model.generate(
                 input_ids=generation_inputs, 
                 generation_config=generation_config,
+                attention_mask=attention_mask,
                 synced_gpus=synced_gpus,
             )
         else:
@@ -991,6 +991,7 @@ class GainLoRA_OLoRA_Trainer(Seq2SeqTrainer):
                     input_ids=generation_inputs,
                     input_ids_wo_label=inputs["input_ids_wo_label"],
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
             
@@ -998,6 +999,7 @@ class GainLoRA_OLoRA_Trainer(Seq2SeqTrainer):
                 generated_tokens = self.model.generate(
                     input_ids=generation_inputs,
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
 

root_gainlora/src/cl_trainer_inflora.py

@@ -739,8 +739,7 @@ class InfLoRATrainer(Seq2SeqTrainer):
                 
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         synced_gpus = gen_kwargs.pop("synced_gpus", False)
         generation_config = GenerationConfig(**gen_kwargs)
@@ -754,6 +753,7 @@ class InfLoRATrainer(Seq2SeqTrainer):
             generated_tokens = self.model.generate(
                 input_ids=generation_inputs, 
                 generation_config=generation_config,
+                attention_mask=attention_mask,
                 synced_gpus=synced_gpus,
             )
         else:
@@ -764,6 +764,7 @@ class InfLoRATrainer(Seq2SeqTrainer):
                     input_ids=generation_inputs,
                     input_ids_wo_label=inputs["input_ids_wo_label"],
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
             
@@ -771,6 +772,7 @@ class InfLoRATrainer(Seq2SeqTrainer):
                 generated_tokens = self.model.generate(
                     input_ids=generation_inputs,
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
 

root_gainlora/src/cl_trainer_inflora_llama.py

@@ -766,8 +766,7 @@ class InfLoRATrainer(Seq2SeqTrainer):
                 
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         generation_config = GenerationConfig(**gen_kwargs)
 

root_gainlora/src/cl_trainer_olora.py

@@ -581,8 +581,7 @@ class OLoRATrainer(Seq2SeqTrainer):
                 
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         synced_gpus = gen_kwargs.pop("synced_gpus", False)
         generation_config = GenerationConfig(**gen_kwargs)
@@ -596,6 +595,7 @@ class OLoRATrainer(Seq2SeqTrainer):
             generated_tokens = self.model.generate(
                 input_ids=generation_inputs, 
                 generation_config=generation_config,
+                attention_mask=attention_mask,
                 synced_gpus=synced_gpus,
             )
         else:
@@ -606,6 +606,7 @@ class OLoRATrainer(Seq2SeqTrainer):
                     input_ids=generation_inputs,
                     input_ids_wo_label=inputs["input_ids_wo_label"],
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
             
@@ -613,6 +614,7 @@ class OLoRATrainer(Seq2SeqTrainer):
                 generated_tokens = self.model.generate(
                     input_ids=generation_inputs,
                     generation_config=generation_config,
+                attention_mask=attention_mask,
                     synced_gpus=synced_gpus,
                 )
 

root_gainlora/src/cl_trainer_olora_llama.py

@@ -583,8 +583,7 @@ class OLoRATrainer(Seq2SeqTrainer):
                 
             gen_kwargs["synced_gpus"] = False
 
-        if "attention_mask" in inputs:
-            gen_kwargs["attention_mask"] = inputs.get("attention_mask", None)
+        attention_mask = inputs.get("attention_mask", None)
 
         generation_config = GenerationConfig(**gen_kwargs)