From 9d94c475ad8648cbb97c79ecebbfa09ab61a06e8 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Tue, 11 Feb 2025 14:22:19 -0800
Subject: [PATCH] max seq length

---
 unsloth/models/llama.py |  6 +++---
 unsloth/models/rl.py    | 19 +++++++++++++++++++
 2 files changed, 22 insertions(+), 3 deletions(-)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index c50f65e4bd..5583702e7e 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -1952,13 +1952,13 @@ class FastLlamaModel:
         Trainer._inner_training_loop = _fast_inner_training_loop
 
         # Save max_seq_length
-        model.max_seq_length = max_position_embeddings
+        model.max_seq_length = max_seq_length
         internal_model = model
         while hasattr(internal_model, "model"):
-            internal_model.max_seq_length = max_position_embeddings
+            internal_model.max_seq_length = max_seq_length
             internal_model = internal_model.model
         pass
-        internal_model.max_seq_length = max_position_embeddings
+        internal_model.max_seq_length = max_seq_length
 
         # We check the tokenizer first for errors
         if fix_tokenizer:
diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py
index 5ec418dda8..dad658170e 100644
--- a/unsloth/models/rl.py
+++ b/unsloth/models/rl.py
@@ -287,6 +287,25 @@ def _patch_trl_rl_trainers(trainer_file = "grpo_trainer"):
         extra_args += saving_check
     pass
 
+    # Edit dataset_num_proc
+    if "dataset_num_proc" in call_args:
+        num_proc_check = \
+        "if dataset_num_proc is None:\n"\
+        "    from multiprocessing import cpu_count\n"\
+        "    dataset_num_proc = cpu_count()\n"
+        extra_args += num_proc_check
+    pass
+
+    # Check max_seq_length
+    if "max_seq_length" in call_args:
+        length_check = \
+        "if hasattr(model, 'max_seq_length') and model.max_seq_length > max_seq_length:\n"\
+        "    print('Unsloth: You set `max_seq_length` as ' + str(max_seq_length) + ' but the\\n'"\
+        "          'model maximum sequence length is ' + str(model.max_seq_length) + '. We will reduce it.')\n"
+        "    max_seq_length = model.max_seq_length\n"
+        extra_args += length_check
+    pass
+
     # Create RLConfig args
     extra_args = extra_args.split("\n")
     extra_args = "\n".join(" "*8 + x for x in extra_args)