Fix Objects365 DDP stream cardinality

Author: Boris Peyriguere

complexity/generative/detection/data.py

@@ -23,6 +23,20 @@
 IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
 
 
+def _bound_and_split_stream(
+    stream: Any,
+    *,
+    local_examples: int,
+    rank: int,
+    world_size: int,
+    splitter: Any,
+) -> Any:
+    """Bound a lazy HF stream globally before attaching its rank split."""
+
+    bounded_stream = stream.take(local_examples * world_size)
+    return splitter(bounded_stream, rank=rank, world_size=world_size)
+
+
 def _normalize_image(image: Image.Image) -> torch.Tensor:
     pixels = torch.from_numpy(np.array(image)).float().permute(2, 0, 1) / 255.0
     return (pixels - 0.5) / 0.5
@@ -284,12 +298,18 @@ def _stream(self, *, metadata_only: bool = False):
             seed=self.seed + self.epoch,
             buffer_size=self.shuffle_buffer,
         )
-        stream = split_dataset_by_node(
+        # Hugging Face applies the distributed split lazily at iteration time.
+        # Calling ``take(local_examples)`` after attaching that split still
+        # limits the *global* stream first, leaving only local/world examples
+        # on each rank. Bound the global stream to an exactly divisible size,
+        # then attach the rank split so every rank yields ``local_examples``.
+        return _bound_and_split_stream(
             stream,
+            local_examples=self.local_examples,
             rank=self.rank,
             world_size=self.world_size,
+            splitter=split_dataset_by_node,
         )
-        return stream.take(self.local_examples)
 
     def _targets(self, row: Dict[str, Any]) -> torch.Tensor:
         annotations = row[self.annotations_column]

complexity/generative/detection/training.py

@@ -865,6 +865,39 @@ def _average_precision_from_matches(
     return float(((recall[changing + 1] - recall[changing]) * precision[changing + 1]).sum())
 
 
+def verify_loader_cardinality(
+    loader: DataLoader,
+    processed_batches: int,
+    distributed: DistributedContext | None,
+    *,
+    phase: str,
+) -> None:
+    """Refuse silently truncated iterable streams on any DDP rank."""
+
+    expected_batches = len(loader)
+    local = {
+        "rank": 0 if distributed is None else distributed.rank,
+        "processed": processed_batches,
+        "expected": expected_batches,
+    }
+    reports = (
+        [local]
+        if distributed is None
+        else distributed.all_gather_objects(local)
+    )
+    invalid = [
+        report
+        for report in reports
+        if report["processed"] != report["expected"]
+    ]
+    if invalid:
+        details = ", ".join(
+            f"rank {report['rank']}: {report['processed']}/{report['expected']}"
+            for report in invalid
+        )
+        raise RuntimeError(f"{phase} loader ended before its declared cardinality ({details})")
+
+
 @torch.inference_mode()
 def evaluate_detector(
     model: TRHashObjectDetector,
@@ -892,7 +925,9 @@ def evaluate_detector(
         leave=False,
         disable=False if show_progress else True,
     )
+    processed_batches = 0
     for pixel_values, targets in progress:
+        processed_batches += 1
         autocast = torch.autocast("cuda", dtype=torch.bfloat16) if use_amp else nullcontext()
         with autocast:
             model_inputs = pixel_values.to(
@@ -921,6 +956,12 @@ def evaluate_detector(
                 detection["labels"],
                 image_targets,
             )
+    verify_loader_cardinality(
+        loader,
+        processed_batches,
+        distributed,
+        phase="validation",
+    )
     if distributed is not None and distributed.enabled:
         states = distributed.all_gather_objects(metrics.state_dict())
         metrics = DetectionMetricsAccumulator(
@@ -1655,7 +1696,9 @@ def write_checkpoint(
             leave=False,
             disable=not distributed.is_main,
         )
+        processed_batches = batches_to_skip
         for batch_index, (pixel_values, targets) in enumerate(progress, start=batches_to_skip):
+            processed_batches = batch_index + 1
             pixel_values = pixel_values.to(device, non_blocking=device.type == "cuda")
             if args.multi_scale_min:
                 choices = range(
@@ -1742,6 +1785,13 @@ def write_checkpoint(
             if args.save_steps and step % args.save_steps == 0:
                 write_checkpoint(epoch=epoch, batch_in_epoch=batch_index + 1)
 
+        verify_loader_cardinality(
+            loader,
+            processed_batches,
+            distributed,
+            phase=f"training epoch {epoch + 1}",
+        )
+
         should_validate = validation_loader is not None and should_validate_epoch(
             epoch,
             args.epochs,

configs/supervisor/tr_hash_detector_objects365_v06_specialized.conf

@@ -12,4 +12,4 @@ killasgroup=true
 redirect_stderr=true
 stdout_logfile=/workspace/complexity-framework/artifacts/tr_hash_detector_objects365_v06_specialized.log
 stdout_logfile_maxbytes=0
-environment=HF_HOME="/workspace/.hf_home",HF_TOKEN_PATH="/root/.cache/huggingface/token",HF_XET_HIGH_PERFORMANCE="1",HF_XET_CHUNK_CACHE_SIZE_BYTES="0",HF_DOWNLOAD_WORKERS="16",OBJECTS365_LOCAL_DIR="/workspace/datasets/object365",OUTPUT="artifacts/detector_objects365_v06_specialized",EPOCHS="12",NPROC_PER_NODE="8",BATCH_SIZE_PER_GPU="24"
+environment=HF_HOME="/workspace/.hf_home",HF_TOKEN_PATH="/root/.cache/huggingface/token",HF_XET_HIGH_PERFORMANCE="1",HF_XET_CHUNK_CACHE_SIZE_BYTES="0",HF_DOWNLOAD_WORKERS="16",PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True",OBJECTS365_LOCAL_DIR="/workspace/datasets/object365",OUTPUT="artifacts/detector_objects365_v06_specialized",EPOCHS="12",NPROC_PER_NODE="8",BATCH_SIZE_PER_GPU="16"

scripts/detector_checkpoint_status.py

@@ -30,7 +30,11 @@ def latest_resumable_checkpoint(root: Path) -> Path | None:
     return max(candidates, key=_step_number, default=None)
 
 
-def checkpoint_status(root: Path, expected_epochs: int) -> tuple[int, Path | None]:
+def checkpoint_status(
+    root: Path,
+    expected_epochs: int,
+    expected_steps: int | None = None,
+) -> tuple[int, Path | None]:
     checkpoint = latest_resumable_checkpoint(root)
     if checkpoint is None:
         return NOT_FOUND, None
@@ -56,6 +60,13 @@ def checkpoint_status(root: Path, expected_epochs: int) -> tuple[int, Path | Non
         print(f"invalid training cursor in {checkpoint}", file=sys.stderr)
         return INCOMPATIBLE, checkpoint
     if epoch == expected_epochs and batch_in_epoch == 0:
+        if expected_steps is not None and state.get("step") != expected_steps:
+            print(
+                f"step budget mismatch in {checkpoint}: "
+                f"saved={state.get('step')!r}, requested={expected_steps}",
+                file=sys.stderr,
+            )
+            return INCOMPATIBLE, checkpoint
         return COMPLETE, checkpoint
     if 0 <= epoch < expected_epochs and batch_in_epoch >= 0:
         return INCOMPLETE, checkpoint
@@ -71,9 +82,14 @@ def main() -> int:
     parser = argparse.ArgumentParser()
     parser.add_argument("root", type=Path)
     parser.add_argument("--expected-epochs", type=int, required=True)
+    parser.add_argument("--expected-steps", type=int, default=None)
     args = parser.parse_args()
 
-    status, checkpoint = checkpoint_status(args.root, args.expected_epochs)
+    status, checkpoint = checkpoint_status(
+        args.root,
+        args.expected_epochs,
+        args.expected_steps,
+    )
     if checkpoint is not None:
         print(checkpoint)
     return status

scripts/vast_pretrain_detector_specialized_objects365.sh

@@ -23,7 +23,9 @@ export COPY_PASTE=0
 export RANDOM_ERASING=0
 export WORKERS=0
 export EPOCHS="${EPOCHS:-12}"
+export BATCH_SIZE_PER_GPU="${BATCH_SIZE_PER_GPU:-16}"
 export LR="${LR:-5.4e-3}"
 export WARMUP_STEPS="${WARMUP_STEPS:-1000}"
+export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True}"
 
 exec scripts/vast_train_detector_specialized_coco.sh

scripts/vast_smoke_detector_specialized_objects365.sh

@@ -4,8 +4,12 @@ set -euo pipefail
 cd "${REPO_ROOT:-/workspace/complexity-framework}"
 
 OBJECTS365_LOCAL_DIR="${OBJECTS365_LOCAL_DIR:-/workspace/datasets/object365}"
-MARKER="${OBJECTS365_SMOKE_MARKER:-$OBJECTS365_LOCAL_DIR/.tr-hash-specialized-smoke-v1-complete}"
-OUTPUT="${SMOKE_OUTPUT:-artifacts/detector_objects365_v06_specialized_local_smoke}"
+MARKER="${OBJECTS365_SMOKE_MARKER:-$OBJECTS365_LOCAL_DIR/.tr-hash-specialized-smoke-v2-complete}"
+OUTPUT="${SMOKE_OUTPUT:-artifacts/detector_objects365_v06_specialized_local_smoke_v2}"
+NPROC_PER_NODE="${NPROC_PER_NODE:-8}"
+BATCH_SIZE_PER_GPU="${BATCH_SIZE_PER_GPU:-16}"
+TRAIN_EXAMPLES="${SMOKE_TRAIN_EXAMPLES:-512}"
+EXPECTED_STEPS=$(( (TRAIN_EXAMPLES / NPROC_PER_NODE + BATCH_SIZE_PER_GPU - 1) / BATCH_SIZE_PER_GPU ))
 
 if [[ ! -f "$OBJECTS365_LOCAL_DIR/.download-complete" ]]; then
   echo "Objects365 snapshot must be complete before the local smoke test" >&2
@@ -20,12 +24,18 @@ echo "[smoke] validating local Objects365 decode, 8-GPU DDP and specialized loss
 OBJECTS365_LOCAL_DIR="$OBJECTS365_LOCAL_DIR" \
 OUTPUT="$OUTPUT" \
 EPOCHS=1 \
-HF_DETECTION_TRAIN_EXAMPLES="${SMOKE_TRAIN_EXAMPLES:-512}" \
+NPROC_PER_NODE="$NPROC_PER_NODE" \
+BATCH_SIZE_PER_GPU="$BATCH_SIZE_PER_GPU" \
+HF_DETECTION_TRAIN_EXAMPLES="$TRAIN_EXAMPLES" \
 HF_DETECTION_VALIDATION_EXAMPLES="${SMOKE_VALIDATION_EXAMPLES:-128}" \
 HF_DETECTION_SHUFFLE_BUFFER="${SMOKE_SHUFFLE_BUFFER:-512}" \
 EVAL_EVERY=1 \
 SAVE_STEPS=1000 \
 scripts/vast_run_detector_specialized_objects365.sh
 
+PYTHON_BIN="${PYTHON_BIN:-/venv/main/bin/python}"
+"$PYTHON_BIN" scripts/detector_checkpoint_status.py \
+  "$OUTPUT" --expected-epochs 1 --expected-steps "$EXPECTED_STEPS" >/dev/null
+
 touch "$MARKER"
 echo "[smoke] passed: $MARKER"

tests/test_detector_pretraining_pipeline.py

@@ -16,6 +16,11 @@
     collect,
     write_reports,
 )
+from scripts.detector_checkpoint_status import (
+    COMPLETE,
+    INCOMPATIBLE,
+    checkpoint_status,
+)
 
 PROJECT_ROOT = Path(__file__).parents[1]
 OBJECTS365_CHECKPOINT = "artifacts/detector_objects365_v06_specialized/best"
@@ -48,6 +53,7 @@ def test_objects365_stage_starts_the_complete_detector_from_scratch() -> None:
     assert "--hf-detection-num-classes 365" in command
     assert "--output artifacts/detector_objects365_v06_specialized" in command
     assert "--epochs 12" in command
+    assert "--batch-size 16" in command
     assert "--mosaic 0" in command
     assert "--workers 0" in command
     assert "--detector-checkpoint" not in command
@@ -69,9 +75,11 @@ def test_objects365_pipeline_requires_local_ddp_smoke_before_long_run() -> None:
     assert pipeline.index("vast_smoke_detector_specialized_objects365.sh") < pipeline.index(
         "vast_run_detector_specialized_objects365.sh"
     )
-    assert "HF_DETECTION_TRAIN_EXAMPLES=\"${SMOKE_TRAIN_EXAMPLES:-512}\"" in smoke
+    assert 'TRAIN_EXAMPLES="${SMOKE_TRAIN_EXAMPLES:-512}"' in smoke
     assert "HF_DETECTION_VALIDATION_EXAMPLES=\"${SMOKE_VALIDATION_EXAMPLES:-128}\"" in smoke
     assert "EVAL_EVERY=1" in smoke
+    assert "--expected-steps \"$EXPECTED_STEPS\"" in smoke
+    assert "smoke-v2-complete" in smoke
     assert "touch \"$MARKER\"" in smoke
 
     resumable = (
@@ -81,6 +89,26 @@ def test_objects365_pipeline_requires_local_ddp_smoke_before_long_run() -> None:
     assert "PYTHON_BIN=/venv/main/bin/python" in resumable
 
 
+def test_detector_checkpoint_status_can_enforce_exact_step_budget(tmp_path: Path) -> None:
+    checkpoint = tmp_path / "step_000004"
+    checkpoint.mkdir()
+    torch.save(
+        {
+            "epoch": 1,
+            "batch_in_epoch": 0,
+            "total_epochs": 1,
+            "step": 4,
+        },
+        checkpoint / "training_state.pt",
+    )
+
+    assert checkpoint_status(tmp_path, 1, expected_steps=4) == (COMPLETE, checkpoint)
+    assert checkpoint_status(tmp_path, 1, expected_steps=3) == (
+        INCOMPATIBLE,
+        checkpoint,
+    )
+
+
 def test_coco_stage_refines_the_objects365_detector_by_default() -> None:
     command = _dry_run("scripts/vast_train_detector_specialized_coco.sh")
 

tests/test_tr_hash_detector_training.py

@@ -4,6 +4,7 @@
 import torch
 from PIL import Image
 from safetensors.torch import save_file
+from torch.utils.data import DataLoader, TensorDataset
 
 from complexity.generative.detection import (
     CocoDetectionDataset,
@@ -22,6 +23,7 @@
     load_training_state,
     save_training_state,
 )
+from complexity.generative.detection.data import _bound_and_split_stream
 from complexity.generative.detection.distributed import DistributedEvalSampler
 from complexity.generative.detection.training import (
     _average_precision_from_matches,
@@ -33,6 +35,7 @@
     resize_detector_inputs,
     save_object_bucket_count_cache,
     should_validate_epoch,
+    verify_loader_cardinality,
     vision_backend_summary,
 )
 
@@ -322,6 +325,39 @@ def test_hf_detection_dataset_resolves_local_parquet_glob(tmp_path):
     )
 
 
+def test_hf_detection_stream_bounds_globally_before_rank_sharding():
+    operations = []
+
+    class FakeStream:
+        def take(self, count):
+            operations.append(("take", count))
+            return self
+
+    def split(stream, *, rank, world_size):
+        operations.append(("split", rank, world_size))
+        return stream
+
+    stream = FakeStream()
+    result = _bound_and_split_stream(
+        stream,
+        local_examples=4,
+        rank=2,
+        world_size=4,
+        splitter=split,
+    )
+
+    assert result is stream
+    assert operations == [("take", 16), ("split", 2, 4)]
+
+
+def test_loader_cardinality_guard_rejects_silent_truncation():
+    loader = DataLoader(TensorDataset(torch.arange(4)), batch_size=2)
+
+    verify_loader_cardinality(loader, 2, None, phase="test")
+    with pytest.raises(RuntimeError, match="test loader ended"):
+        verify_loader_cardinality(loader, 1, None, phase="test")
+
+
 def test_synthetic_dataset_deterministic_across_instances():
     first = SyntheticShapesDataset(length=4, image_size=64, seed=7)
     second = SyntheticShapesDataset(length=4, image_size=64, seed=7)