fix sampler save when dataloader num_workers > 0

Author: jayhenry

tests/datasets/test_dataloader.py

@@ -1,12 +1,21 @@
 from pathlib import Path
 import os
 import pickle
+import socket
 
 import torch
 
-from xtuner.v1.datasets import build_dataloader, build_datasets, get_dataloader_state, load_dataloader_state, FTDPTokenizeFnConfig, DatasetConfig, DataloaderConfig
+from xtuner.v1.datasets import (
+    DataloaderConfig,
+    DatasetConfig,
+    FTDPTokenizeFnConfig,
+    build_dataloader,
+    build_datasets,
+    get_dataloader_state,
+    load_dataloader_state,
+)
 from xtuner.v1.train.toy_tokenizer import UTF8ByteTokenizer
-from torch.multiprocessing import spawn, get_context
+from torch.multiprocessing import spawn
 from torch.distributed.device_mesh import init_device_mesh
 import pytest
 
@@ -15,6 +24,12 @@
 from itertools import repeat, chain
 
 
+def _alloc_master_port() -> None:
+    """Bind an ephemeral TCP port so concurrent test runs avoid EADDRINUSE on a fixed port."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("127.0.0.1", 0))
+        os.environ["MASTER_PORT"] = str(s.getsockname()[1])
+
 
 
 class RandomDataset:
@@ -282,65 +297,53 @@ def _test_resume_spmd(
     rank: int,
     world_size: int,
     dataloader_config: DataloaderConfig,
-    dataset_configs: list[dict],
     global_batch_size: int,
     micro_batch_size: int,
-    step:int,
+    step: int,
     seed: int,
     save_path: Path,
     dataloader_state: dict | None = None,
-    consumed_samples: int = 0,
 ):
     os.environ["RANK"] = str(rank)
     os.environ["LOCAL_RANK"] = str(rank)
     os.environ["WORLD_SIZE"] = str(world_size)
-    os.environ["MASTER_ADDR"] = "localhost"
-    os.environ["MASTER_PORT"] = "29505"
-
+    os.environ.setdefault("MASTER_ADDR", "localhost")
+    if "MASTER_PORT" not in os.environ:
+        raise RuntimeError("tests must call _alloc_master_port() before torch.multiprocessing.spawn")
 
     torch.distributed.init_process_group(backend="nccl", rank=rank, world_size=world_size)
     torch.cuda.set_device(rank)
     data_mesh = init_device_mesh(
         device_type="cuda",
-        mesh_shape=(world_size,)
+        mesh_shape=(world_size,),
     )
     tokenizer = UTF8ByteTokenizer()
 
-    datasets = build_datasets(
-        dataset_config=dataset_configs,
+    dataloader = dataloader_config.build(
         tokenizer=tokenizer,
-    )
-    dataloader = build_dataloader(
-        dataloader_config=dataloader_config,
-        datasets=datasets,
+        dp_mesh=data_mesh,
         global_batch_size=global_batch_size,
         micro_batch_size=micro_batch_size,
         seed=seed,
-        dp_mesh=data_mesh,
     )
 
     if dataloader_state is not None:
-        load_dataloader_state(dataloader, dataloader_state)
+        dataloader.load_state_dict(dataloader_state)
 
     data_iter = iter(dataloader)
     data_list = []
     for _ in range(step):
         batch = next(data_iter)
         data_list.append(batch)
-        consumed_samples += len(batch)
 
-    consumed_samples_list = [None for _ in range(world_size)]
-    torch.distributed.all_gather_object(consumed_samples_list, consumed_samples)
-    global_consumed_samples = sum(consumed_samples_list)
+    # Snapshot after the first `step` batches so total_consumed_steps matches resume intent.
+    dataloader_state = dataloader.get_state_dict()
 
     expected_data = []
-
     for _ in range(step):
         batch = next(data_iter)
         expected_data.append(batch)
 
-    dataloader_state = get_dataloader_state(dataloader, global_consumed_samples)
-
     all_data_list = [None for _ in range(world_size)]
     torch.distributed.all_gather_object(all_data_list, list(chain(*data_list)))
 
@@ -372,7 +375,6 @@ def _test_resume_spmd(
                         "dataloader_state": dataloader_state,
                         "data_list": all_data_list,
                         "expected_data": all_expected_data,
-                        "consumed_samples": consumed_samples
                     }
                 )
             )
@@ -389,7 +391,6 @@ def _test_resume_spmd(
         ("none", 0, False),
         ("soft", 0, True),
         ("soft", 4, True),
-        ("soft", 4, True),
     ]
 )
 def test_dataloader_resume_multi_process(tmp_path, pack_level, num_workers, group_by_length):
@@ -402,36 +403,36 @@ def test_dataloader_resume_multi_process(tmp_path, pack_level, num_workers, grou
     _create_fake_dataset(data_dir1 / f"depth3", dataset_num=3, max_depth=3, dup_times=9)
 
     # 1. Test resuming with the same world size
+    dataset_configs = [
+        {
+            "dataset": DatasetConfig(anno_path=str(data_dir1)),
+            "tokenize_fn": FTDPTokenizeFnConfig(max_length=1024),
+        },
+    ]
+
     dataloader_config = DataloaderConfig(
+        dataset_config_list=dataset_configs,
         pack_max_length=1024,
         pack_level=pack_level,
         num_workers=num_workers,
         group_by_length=group_by_length,
-        collator="fake_collator"
+        collator="fake_collator",
     )
-    dataset_configs = [
-        {
-            "dataset": DatasetConfig(anno_path=str(data_dir1)),
-            "tokenize_fn": FTDPTokenizeFnConfig(max_length=1024)
-        },
-    ]
 
-    ctx = get_context("spawn")
     world_size = 2
     save_path1 = tmp_path / "dataloader_state.pkl"
+    _alloc_master_port()
     spawn(
         _test_resume_spmd,
         args=(
             world_size,
             dataloader_config,
-            dataset_configs,
             16,
             BATCH_SIZE,
             TOTAL_STEP,
             10,
             save_path1,
             None,
-            0,
         ),
         nprocs=2,
         join=True,
@@ -443,19 +444,18 @@ def test_dataloader_resume_multi_process(tmp_path, pack_level, num_workers, grou
 
     # 2. tet Rsume with same world size
     save_path2 = tmp_path / "dataloader_state2.pkl"
+    _alloc_master_port()
     spawn(
         _test_resume_spmd,
         args=(
             world_size,
             dataloader_config,
-            dataset_configs,
             16,
             BATCH_SIZE,
             TOTAL_STEP,
             10,
             save_path2,
             result1["dataloader_state"],
-            result1["consumed_samples"],
         ),
         nprocs=world_size,
         join=True,
@@ -470,19 +470,18 @@ def test_dataloader_resume_multi_process(tmp_path, pack_level, num_workers, grou
 
     world_size = 4
     save_path3 = tmp_path / "dataloader_state3.pkl"
+    _alloc_master_port()
     spawn(
         _test_resume_spmd,
         args=(
             world_size,
             dataloader_config,
-            dataset_configs,
             16,
             BATCH_SIZE,
             TOTAL_STEP,
             10,
             save_path3,
             result1["dataloader_state"],
-            result1["consumed_samples"],
         ),
         nprocs=world_size,
         join=True,

xtuner/v1/datasets/dataloader.py

@@ -4,6 +4,7 @@
 import torch
 
 from xtuner.v1.datasets.collator import ColateItem
+from xtuner.v1.datasets.consumed_steps import ConsumedStepsTracker
 from xtuner.v1.datasets.resume import get_dataloader_state, load_dataloader_state
 from xtuner.v1.utils import get_logger
 
@@ -42,6 +43,11 @@ def load_state_dict(
         state_dict: dict,
         train_state_total_consumed_samples: int | None = None,
     ) -> None:
+        if train_state_total_consumed_samples is not None:
+            logger.warning(
+                "Dataloader.load_state_dict(train_state_total_consumed_samples=...) is deprecated; "
+                "use the default (None). Consumed samples are tracked on the sampler."
+            )
         load_dataloader_state(
             self,
             state_dict,
@@ -58,7 +64,17 @@ def get_state_dict(self, consumed_samples: int = -1) -> dict:
         dataloader_state = get_dataloader_state(self, consumed_samples)
         return cast(dict, dataloader_state)
 
-    # __iter__ is inherited from torch.utils.data.DataLoader
+    def __iter__(self) -> Iterator[list[ColateItem]]:  # type: ignore[override]
+        # Override to count delivered batches, not prefetched indices.
+        # With num_workers > 0 the sampler is iterated ahead by DataLoader's prefetch queue,
+        # so recording inside sampler.__iter__ would count too many samples.  Instead we
+        # increment _consumed exactly once per batch that reaches the caller.
+        sampler = self.sampler
+        consumed: ConsumedStepsTracker | None = getattr(sampler, "_consumed", None)
+        for batch in super().__iter__():
+            if consumed is not None:
+                consumed.record(len(batch))
+            yield batch
 
     # Streaming dataloader may not have `set_epoch` and `__len__` method, so we add here.
     def set_epoch(self, epoch: int) -> None:

xtuner/v1/datasets/preset_sampler.py

@@ -163,9 +163,7 @@ def __init__(
     def __iter__(self) -> Iterator[int]:
         # load order from npy → global_order → rank_view 类型均为 memmap, 子视图 的路径仍然保持
         # memmap 语义（视图、按需分页、文件后端）；单机多进程可共享同一份文件页缓存
-        for idx in self.global_order[self.step + self.rank : self.total_size : self.world_size]:
-            self._consumed.record(1)
-            yield int(idx)
+        yield from (int(idx) for idx in self.global_order[self.step + self.rank : self.total_size : self.world_size])
         self.step = 0
 
     def __len__(self) -> int:

xtuner/v1/datasets/sampler.py

@@ -113,9 +113,7 @@ def __iter__(self) -> Iterator[int]:
         # subsample
         indices = indices[self.step + self.rank : self.total_size : self.world_size]
 
-        for idx in indices:
-            self._consumed.record(1)
-            yield idx
+        yield from indices
         self.step = 0
 
     def __len__(self) -> int:
@@ -268,9 +266,7 @@ def __iter__(self) -> Iterator[int]:
         assert len(indices) == self.total_size
         indices = indices[self.step + self.rank : self.total_size : self.world_size]
         assert len(indices) == self.num_samples - self.step // self.world_size
-        for idx in indices:
-            self._consumed.record(1)
-            yield idx
+        yield from indices
         self.step = 0
 
     def __len__(self) -> int: