[v1] upgrade batching (#9751)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

src/llamafactory/v1/utils/objects.py

@@ -0,0 +1,67 @@
+# Copyright 2025 Optuna, HuggingFace Inc. and the LlamaFactory team.
+#
+# This code is inspired by the HuggingFace's transformers library.
+# https://github.com/huggingface/transformers/blob/v5.0.0rc0/src/transformers/utils/logging.py
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .types import ModelInput
+
+
+class StatefulBuffer:
+    """A buffer that stores model inputs."""
+
+    def __init__(self, max_buffer_size: int = 1_000_000_000) -> None:
+        self._buffer: list[ModelInput] = []
+        self._buffer_size: int = 0
+        self._max_buffer_size: int = max_buffer_size
+
+    def __len__(self) -> int:
+        return len(self._buffer)
+
+    @property
+    def size(self) -> int:
+        return self._buffer_size
+
+    def put(self, samples: list[ModelInput]) -> None:
+        """Add samples to the buffer."""
+        num_tokens = sum(len(sample["input_ids"]) for sample in samples)
+        if self._buffer_size + num_tokens > self._max_buffer_size:
+            raise ValueError(f"Buffer size exceeds max buffer size {self._max_buffer_size}.")
+
+        self._buffer.extend(samples)
+        self._buffer_size += num_tokens
+
+    def get(self, value: int) -> list[ModelInput]:
+        """Get samples from the buffer and remove them."""
+        samples = self._buffer[:value]
+        self._buffer_size -= sum(len(sample["input_ids"]) for sample in samples)
+        del self._buffer[:value]
+        return samples
+
+    def clear(self) -> None:
+        """Clear the buffer."""
+        self._buffer = []
+        self._buffer_size = 0
+
+    def state_dict(self) -> dict:
+        """Returns the state of the buffer."""
+        return {
+            "buffer": self._buffer,
+            "buffer_size": self._buffer_size,
+        }
+
+    def load_state_dict(self, state_dict: dict) -> None:
+        """Loads the state into the buffer."""
+        self._buffer = state_dict["buffer"]
+        self._buffer_size = state_dict["buffer_size"]

src/llamafactory/v1/utils/plugin.py

@@ -15,6 +15,7 @@
 
 from collections import defaultdict
 from collections.abc import Callable
+from typing import Any
 
 from . import logging
 
@@ -26,33 +27,37 @@ class BasePlugin:
     """Base class for plugins.
 
     A plugin is a callable object that can be registered and called by name.
+
+    Example usage:
+    ```python
+    class PrintPlugin(BasePlugin):
+        def again(self):  # optional
+            self["again"]()
+
+
+    @PrintPlugin("hello").register()
+    def print_hello():
+        print("Hello world!")
+
+
+    @PrintPlugin("hello").register("again")
+    def print_hello_again():
+        print("Hello world! Again.")
+
+
+    PrintPlugin("hello")()
+    PrintPlugin("hello").again()
+    ```
     """
 
     _registry: dict[str, dict[str, Callable]] = defaultdict(dict)
 
-    def __init__(self, name: str | None = None):
-        """Initialize the plugin with a name.
-
-        Args:
-            name (str): The name of the plugin.
-        """
+    def __init__(self, name: str | None = None) -> None:
+        """Initialize the plugin with a name."""
         self.name = name
 
-    def register(self, method_name: str = "__call__"):
-        """Decorator to register a function as a plugin.
-
-        Example usage:
-        ```python
-        @PrintPlugin("hello").register()
-        def print_hello():
-            print("Hello world!")
-
-
-        @PrintPlugin("hello").register("again")
-        def print_hello_again():
-            print("Hello world! Again.")
-        ```
-        """
+    def register(self, method_name: str = "__call__") -> Callable:
+        """Decorator to register a function as a plugin."""
         if self.name is None:
             raise ValueError("Plugin name should be specified.")
 
@@ -65,27 +70,16 @@ class BasePlugin:
 
         return decorator
 
-    def __call__(self, *args, **kwargs):
-        """Call the registered function with the given arguments.
+    def __call__(self, *args, **kwargs) -> Any:
+        """Call the registered function with the given arguments."""
+        return self["__call__"](*args, **kwargs)
 
-        Example usage:
-        ```python
-        PrintPlugin("hello")()
-        ```
-        """
-        if "__call__" not in self._registry[self.name]:
-            raise ValueError(f"Method __call__ of plugin {self.name} is not registered.")
+    def __getattr__(self, method_name: str) -> Callable:
+        """Get the registered function with the given name."""
+        return self[method_name]
 
-        return self._registry[self.name]["__call__"](*args, **kwargs)
-
-    def __getattr__(self, method_name: str):
-        """Get the registered function with the given name.
-
-        Example usage:
-        ```python
-        PrintPlugin("hello").again()
-        ```
-        """
+    def __getitem__(self, method_name: str) -> Callable:
+        """Get the registered function with the given name."""
         if method_name not in self._registry[self.name]:
             raise ValueError(f"Method {method_name} of plugin {self.name} is not registered.")
 
@@ -98,7 +92,8 @@ if __name__ == "__main__":
     """
 
     class PrintPlugin(BasePlugin):
-        pass
+        def again(self):  # optional
+            self["again"]()
 
     @PrintPlugin("hello").register()
     def print_hello():

src/llamafactory/v1/utils/types.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from collections.abc import Iterator
 from typing import TYPE_CHECKING, Any, Literal, NotRequired, TypedDict, Union
 
 
@@ -161,3 +162,14 @@ class BatchInput(TypedDict, total=False):
     """Position ids for the model (optional)."""
     token_type_ids: NotRequired[Tensor]
     """Token type ids used in DPO, 0 represents the chosen messages, 1 represents the rejected messages."""
+
+
+class BatchInfo(TypedDict):
+    micro_batch_size: int
+    """Micro batch size."""
+    num_micro_batch: int
+    """Number of micro batches."""
+    cutoff_len: int
+    """Cutoff length."""
+    data_iter: Iterator[list[ModelInput]]
+    """Data iterator."""