From 5da887c366ecf7af27ec9c928823b52b35c2eb99 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 10 Mar 2025 00:31:49 -0700 Subject: [PATCH] Update _utils.py --- unsloth/models/_utils.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 50dbe7cae6..a63aaccc28 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -957,9 +957,13 @@ def _unsloth_get_batch_samples(self, epoch_iterator, num_batches): num_items_in_batch = None # Check if model allows **kwargs - model = self.model - f = model.base_model.model.forward if hasattr(model, "base_model") else model.forward - has_kwargs = tuple(inspect.signature(f).parameters.values())[-1].kind == inspect._VAR_KEYWORD + m = self.model + while hasattr(m, "model"): + # Stop at last model entry + if not hasattr(m, "model") or not hasattr(m, "forward"): break + m = m.model + signature = inspect.signature(m.forward).parameters.values() + has_kwargs = tuple(signature)[-1].kind == inspect._VAR_KEYWORD # Iterate to find all batches for _ in range(num_batches):