Fix transformers 4.45

unsloth_zoo/__init__.py

@@ -14,7 +14,7 @@
 # You should have received a copy of the GNU Lesser General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
-__version__ = "2025.3.15"
+__version__ = "2025.3.16"
 
 from importlib.util import find_spec
 if find_spec("unsloth") is None:

unsloth_zoo/loss_utils.py

@@ -287,19 +287,11 @@ def _unsloth_get_batch_samples(self, epoch_iterator, num_batches, device = None,
                     [((x["labels"][..., 1:] != -100) & (x["attention_mask"][..., 1:] != 0))\
                     .sum() for x in batch_samples]
                 )
-            if device is None: # transformers < 4.50.0 path
-                if self.args.average_tokens_across_devices:
-                    num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum().item()
-                if torch.is_tensor(num_items_in_batch):
-                    num_items_in_batch = num_items_in_batch.item()
-                pass
-            else: # transformers >= 4.50.0 path
-                if self.args.average_tokens_across_devices:
-                    num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum()
-                if torch.is_tensor(num_items_in_batch):
-                    num_items_in_batch = num_items_in_batch.to(device)
-                pass
 
+            if self.args.average_tokens_across_devices:
+                num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum()
+            if device is not None and torch.is_tensor(num_items_in_batch):
+                num_items_in_batch = num_items_in_batch.to(device)
         except Exception as exception:
             raise RuntimeError(exception)
     pass

unsloth_zoo/vllm_utils.py

@@ -386,7 +386,7 @@ def get_vllm_state_dict(llm, return_state_dict = False, config = None):
     # All Unsloth Zoo code licensed under LGPLv3
     # Unmerges vLLM modules and returns HF equivalent state_dict
     try:
-        llm_engine = getattr(llm, "llm_engine", llm)
+        llm_engine = getattr(llm, "llm_engine", getattr(llm, "engine", llm))
         vllm_internals = llm_engine.model_executor.driver_worker.model_runner.model
     except:
         raise RuntimeError("Unsloth: Failed to access llm.llm_engine.model_executor.driver_worker.model_runner.model")
@@ -919,7 +919,7 @@ def load_vllm(
     #         os.environ["VLLM_USE_V1"] = "1"
     # pass
 
-    from vllm import LLM, LLMEngine, AsyncLLMEngine, EngineArgs
+    from vllm import LLM, LLMEngine, AsyncLLMEngine, EngineArgs, AsyncEngineArgs
 
     # Default vLLM max_num_seqs is 256
     approx_max_num_seqs = 256
@@ -1007,7 +1007,7 @@ def load_vllm(
         swap_space             = swap_space, # Low memory devices like Colab (13GB) default 4GB
         device                 = device,
     )
-    good_keys = inspect.signature(EngineArgs).parameters.keys()
+    good_keys = inspect.signature(AsyncEngineArgs if use_async else EngineArgs).parameters.keys()
     old_keys = engine_args.keys()
     for key in old_keys:
         if key not in good_keys:
@@ -1021,7 +1021,7 @@ def load_vllm(
     while True:
         try:
             if use_async:
-                llm = AsyncLLMEngine.from_engine_args(EngineArgs(**engine_args))
+                llm = AsyncLLMEngine.from_engine_args(AsyncEngineArgs(**engine_args))
             elif use_engine:
                 llm = LLMEngine.from_engine_args(EngineArgs(**engine_args))
             else: