Fix setting max_seq_length when using longlora

Lightning-AI · Apr 25, 2024 · 2dfa7a5 · 2dfa7a5
1 parent 9facaf3
commit 2dfa7a5
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/litgpt/finetune/full.py b/litgpt/finetune/full.py
@@ -232,7 +232,9 @@ def fit(
     scheduler = state["scheduler"]
     tokenizer = Tokenizer(checkpoint_dir)
     longest_seq_length, longest_seq_ix = get_longest_seq_length(train_dataloader.dataset)
-    model.max_seq_length = min(longest_seq_length, train.max_seq_length or float("inf"))
+    if longlora.use_longlora:
+        longest_seq_length = find_multiple(longest_seq_length, longlora.n_groups)
+    model.max_seq_length = longest_seq_length
     fabric.print(
         f"The longest sequence length in the train data is {longest_seq_length}, the model's maximum sequence length is"
         f" {model.max_seq_length} and context length is {model.config.block_size}"

diff --git a/litgpt/finetune/lora.py b/litgpt/finetune/lora.py
@@ -302,7 +302,10 @@ def fit(
 ) -> None:
     tokenizer = Tokenizer(checkpoint_dir)
     longest_seq_length, longest_seq_ix = get_longest_seq_length(train_dataloader.dataset)
-    model.max_seq_length = min(longest_seq_length, train.max_seq_length or float("inf"))
+    longest_seq_length = min(longest_seq_length, train.max_seq_length or float("inf"))
+    if longlora.use_longlora:
+        longest_seq_length = find_multiple(longest_seq_length, longlora.n_groups)
+    model.max_seq_length = longest_seq_length
     fabric.print(
         f"The longest sequence length in the train data is {longest_seq_length}, the model's maximum sequence length is"
         f" {model.max_seq_length} and context length is {model.config.block_size}"