From 95a2627bf6d96efebd64681a555bf8b9ab90e022 Mon Sep 17 00:00:00 2001
From: Irakli <39024518+IrakliXYZ@users.noreply.github.com>
Date: Fri, 12 Jun 2026 13:30:53 +0400
Subject: [PATCH] Fix step count mismatch when sequence packing is enabled
 (#5967)

* Fix step count mismatch when sequence packing is enabled

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Emit a single step-0 progress event and guard applyStatus totalSteps

Merge the two consecutive _update_progress calls before train() so the
step-0 gate in _on_progress fires once instead of twice, avoiding a
duplicate startup event and a null-metric step-0 row in training_metrics.

Apply the same positive-number guard to applyStatus that applyProgress
uses, so a stale or startup status poll can no longer overwrite the
packed step count with 0 or replace it with a stale total.

* Log debug message when train_dataset length is unavailable

The TypeError fallback for length-less datasets (e.g. streaming
IterableDataset) was silent, leaving no trace that the step estimate
came from the raw dataset rather than the packed one.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com>
---
 studio/backend/core/training/trainer.py        | 18 ++++++++++++++----
 studio/backend/core/training/worker.py         |  2 +-
 .../training/stores/training-runtime-store.ts  |  9 ++++++---
 3 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py
index 57342b2453..085f999dd6 100644
--- a/studio/backend/core/training/trainer.py
+++ b/studio/backend/core/training/trainer.py
@@ -3367,7 +3367,19 @@ class UnslothTrainer:
             # ========== PROGRESS TRACKING ==========
             self.trainer.add_callback(self._create_progress_callback())
 
-            num_samples = len(dataset["dataset"] if isinstance(dataset, dict) else dataset)
+            num_samples = None
+            if hasattr(self.trainer, "train_dataset") and self.trainer.train_dataset is not None:
+                try:
+                    num_samples = len(self.trainer.train_dataset)
+                except TypeError:
+                    logger.debug(
+                        "train_dataset does not support len(); falling back to "
+                        "raw dataset size for step estimation."
+                    )
+
+            if num_samples is None:
+                num_samples = len(dataset["dataset"] if isinstance(dataset, dict) else dataset)
+
             batch_size = training_args.get("batch_size", 2)
             total_steps = self._calculate_total_steps(
                 num_samples,
@@ -3376,10 +3388,8 @@ class UnslothTrainer:
                 training_args.get("num_epochs", 3),
                 training_args.get("max_steps", 0),
             )
-            self._update_progress(total_steps = total_steps)
-
             # ========== START TRAINING ==========
-            self._update_progress(status_message = "Starting training...")
+            self._update_progress(total_steps = total_steps, status_message = "Starting training...")
             logger.info("Starting training...\n")
             self.trainer.train(resume_from_checkpoint = training_args.get("resume_from_checkpoint"))
 
diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py
index 18b25cb4fe..c7c9003a04 100644
--- a/studio/backend/core/training/worker.py
+++ b/studio/backend/core/training/worker.py
@@ -2462,7 +2462,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
     def _on_progress(progress: TrainingProgress):
         has_train_loss = progress.step > 0 and progress.loss is not None
         has_eval_loss = progress.eval_loss is not None
-        if has_train_loss or has_eval_loss:
+        if (progress.step == 0 and progress.total_steps > 0) or has_train_loss or has_eval_loss:
             event_queue.put(
                 {
                     "type": "progress",
diff --git a/studio/frontend/src/features/training/stores/training-runtime-store.ts b/studio/frontend/src/features/training/stores/training-runtime-store.ts
index 97fbd32d57..9eaaa98c0e 100644
--- a/studio/frontend/src/features/training/stores/training-runtime-store.ts
+++ b/studio/frontend/src/features/training/stores/training-runtime-store.ts
@@ -209,8 +209,8 @@ export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((set) => (
         currentStep:
           typeof detailStep === "number" ? Math.max(detailStep, 0) : state.currentStep,
         totalSteps:
-          typeof detailTotal === "number"
-            ? Math.max(detailTotal, 0)
+          typeof detailTotal === "number" && detailTotal > 0
+            ? detailTotal
             : state.totalSteps,
         currentLoss:
           typeof detailLoss === "number" ? detailLoss : state.currentLoss,
@@ -273,7 +273,10 @@ export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((set) => (
         ...state,
         jobId: payload.job_id || state.jobId,
         currentStep: step,
-        totalSteps: Math.max(payload.total_steps, state.totalSteps),
+        totalSteps:
+          typeof payload.total_steps === "number" && payload.total_steps > 0
+            ? payload.total_steps
+            : state.totalSteps,
         // A null loss at a new step means the backend reported a non-finite
         // loss; clear the display instead of keeping the stale value.
         currentLoss: