[v1] add sft (#9752)

src/llamafactory/v1/utils/helper.py

@@ -16,6 +16,7 @@
 import torch
 from transformers import PreTrainedTokenizer
 
+from ..accelerator.interface import DistributedInterface
 from .constants import IGNORE_INDEX
 from .types import BatchInput, ModelInput, Processor, Tensor
 
@@ -73,3 +74,20 @@ def pad_and_truncate(samples: list[ModelInput], max_seqlen: int) -> list[BatchIn
         padded_samples.append(padded_sample)
 
     return padded_samples
+
+
+def compute_valid_tokens(batches: list[BatchInput]) -> int:
+    """Compute valid tokens in batches.
+
+    Args:
+        batches: Batches.
+
+    Returns:
+        Number of valid tokens.
+    """
+    device = DistributedInterface().current_device
+    return sum(
+        (batch["labels"].to(device, non_blocking=True) != IGNORE_INDEX).sum().item()
+        for batch in batches
+        if "labels" in batch
+    )

src/llamafactory/v1/utils/types.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from collections.abc import Iterator
-from typing import TYPE_CHECKING, Any, Literal, NotRequired, TypedDict, Union
+from typing import TYPE_CHECKING, Any, Literal, NamedTuple, NotRequired, TypedDict, Union
 
 
 if TYPE_CHECKING:
@@ -146,7 +146,7 @@ class ModelInput(TypedDict, total=False):
     position_ids: NotRequired[list[int] | list[list[int]]]
     """Position ids for the model (optional)."""
     token_type_ids: NotRequired[list[int]]
-    """Token type ids used in DPO, 0 represents the chosen messages, 1 represents the rejected messages."""
+    """Token type ids used in DPO, 1 represents the chosen messages, 2 represents the rejected messages."""
 
 
 class BatchInput(TypedDict, total=False):
@@ -161,7 +161,7 @@ class BatchInput(TypedDict, total=False):
     position_ids: NotRequired[Tensor]
     """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."""
+    """Token type ids used in DPO, 1 represents the chosen messages, 2 represents the rejected messages."""
 
 
 class BatchInfo(TypedDict):
@@ -173,3 +173,8 @@ class BatchInfo(TypedDict):
     """Cutoff length."""
     data_iter: Iterator[list[ModelInput]]
     """Data iterator."""
+
+
+class ModelOutput(NamedTuple):
+    logits: Tensor
+    """Logits for the model."""