Sharded dataset generation

README.md

@@ -60,6 +60,12 @@ The model is trained from a random initialization until convergence, which is de
 1. Once fractal generation completes, run the benchmark:  
     `torchrun-hpc -N 1 -n 4 --gpus-per-proc 1 $(which scaffold) benchmark -c ScaFFold/configs/benchmark_default.yml`
 
+### Dataset cache and sharded datagen
+
+`benchmark` creates or reuses datasets under `dataset_dir`. New datasets are written in the v3 format, which stores one volume and mask file per logical sample per physical shard. The physical layout is controlled by `dc_num_shards` and `dc_shard_dims`; for example, `dc_num_shards: [1, 1, 2]` writes two physical shards per logical volume, with filenames such as `120_shard000000.npy` and `120_shard000001.npy`. Datasets are generated with the same sharding configuration used for model training.
+
+Unsharded runs use `dc_num_shards: [1, 1, 1]`. For those runs, ScaFFold can still reuse an existing v2 full-volume dataset cache. Sharded runs require a matching v3 cache or generate a new v3 dataset.
+
 `benchmark` creates a folder for the benchmark run(s) at `base_run_dir` set in the config file. For reproducibility, we store a copy of the benchmark run config yml. Within each run subfolder, `benchmark` creates a yml config for that specific run.
 
 After each run completes, statistics from the run are stored in `train_stats.csv`. Additionally, users can inspect plots of the training and validation losses over time in `<base_run_dir/figures`.
@@ -69,14 +75,18 @@ Parameters are set in a `.yml` config file and can be modified by the user:
 ```yml
 # External/user-facing
 base_run_dir: "benchmark_runs"     # Subfolder of $(pwd) in which to run jobs.
+dataset_dir: "datasets"            # Directory in which to store and query for datasets.
 fract_base_dir: "fractals"         # Base directory for fractal IFS and instances.
 n_categories: 5                    # Number of fractal categories present in the dataset.
 n_instances_used_per_fractal: 145  # Number of unique instances to pull from each fractal class. There are 145 unique; exceeding this number will reuse some instances.
 problem_scale: 6                   # Determines dataset resolution and number of unet layers. Default is 6.
 unet_bottleneck_dim: 3             # Power of 2 of the unet bottleneck layer dimension. Default of 3 -> bottleneck layer of size 8.
 seed: 42                           # Random seed.
-batch_size: 1                      # Batch sizes for each vol size.
+batch_size: 1                      # Batch size per rank.
+dataloader_num_workers: 1          # Number of DataLoader worker processes per rank.
 optimizer: "ADAM"                  # "ADAM" is preferred option, otherwise training defautls to RMSProp.
+dc_num_shards: [1, 1, 1]           # Physical data shards per sample for DistConv.
+dc_shard_dims: [2, 3, 4]           # Tensor dimensions used for physical sharding.
 
 # Internal/dev use only
 variance_threshold: 0.15           # Variance threshold for valid fractals. Default is 0.15.
@@ -97,6 +107,7 @@ framework: "torch"                 # The DL framework to train with. Only valid
 checkpoint_dir: "checkpoints"      # Subfolder in which to save training checkpoints.
 checkpoint_interval: 1             # Number of epochs between saving training checkpoints.
 loss_freq: 1                       # Number of epochs between logging the overall loss.
+warmup_batches: 64                 # Training and validation warmup batches per DDP rank.
 ```
 
 ## How the benchmark works
@@ -194,6 +205,8 @@ For n  in n_volumes:
     3. Save volume and mask  to files
 ```
 
+In the current v3 dataset format, this save step writes each logical sample as one or more physical shard files, matching the requested `dc_num_shards` layout. The dataloader then reads only the shard file needed by the current DistConv rank instead of loading a full volume and slicing it locally.
+
 ### 1. Profiling with the PyTorch Profiler
 
 Set `PROFILE_TORCH=ON` to generate a PyTorch profiling trace that can be read into [Perfetto](https://ui.perfetto.dev/).

ScaFFold/cli.py

@@ -182,6 +182,12 @@ def main():
         nargs=3,
         help="DistConv param: number of shards to divide the tensor into. It's best to choose the fewest ranks needed to fit one sample in GPU memory, since that keeps communication at a minimum",
     )
+    benchmark_parser.add_argument(
+        "--dc-shard-dims",
+        type=int,
+        nargs=3,
+        help="DistConv param: dimensions on which to shard.",
+    )
     benchmark_parser.add_argument(
         "--epochs",
         type=int,

ScaFFold/configs/benchmark_testing.yml

@@ -10,8 +10,8 @@ seed: 42                           # Random seed.
 batch_size: 1                      # Batch sizes for each vol size.
 dataloader_num_workers: 4          # Number of DataLoader worker processes per rank.
 optimizer: "ADAM"                  # "ADAM" is preferred option, otherwise training defautls to RMSProp.
-num_shards: [1, 1, 1]              # DistConv param: number of shards to divide the tensor into. It's best to choose the fewest ranks needed to fit one sample in GPU memory, since that keeps communication at a minimum
-shard_dim: [2, 3, 4]               # DistConv param: dimension on which to shard
+dc_num_shards: [1, 1, 1]           # DistConv param: number of shards to divide the tensor into. It's best to choose the fewest ranks needed to fit one sample in GPU memory, since that keeps communication at a minimum
+dc_shard_dims: [2, 3, 4]           # DistConv param: dimension on which to shard
 checkpoint_interval: -1            # Checkpoint every C epochs; set to -1 to disable checkpointing entirely.
 
 # Internal/dev use only

ScaFFold/datagen/get_dataset.py

@@ -21,16 +21,17 @@
 import time
 from argparse import Namespace
 from pathlib import Path
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 import yaml
 from mpi4py import MPI
 
 from ScaFFold.datagen import volumegen
 
 META_FILENAME = "meta.yaml"
-DATASET_FORMAT_VERSION = 2
-INCLUDE_KEYS = [
+DATASET_FORMAT_VERSION = 3
+V2_DATASET_FORMAT_VERSION = 2
+V2_INCLUDE_KEYS = [
     "dataset_format_version",
     "n_categories",
     "n_instances_used_per_fractal",
@@ -40,6 +41,10 @@
     "n_fracts_per_vol",
     "val_split",
 ]
+INCLUDE_KEYS = V2_INCLUDE_KEYS + [
+    "dc_num_shards",
+    "dc_shard_dims",
+]
 
 
 def canonicalize(input):
@@ -71,11 +76,63 @@ def _get_required_keys_dict(
     return canonicalize(required)
 
 
+def _canonicalize_v3_shard_layout(volume_config: Dict[str, Any]) -> Dict[str, Any]:
+    """Normalize shard layout ordering so equivalent v3 layouts share cache IDs."""
+
+    canonical_config = volume_config.copy()
+    num_shards = canonical_config["dc_num_shards"]
+    shard_dims = canonical_config["dc_shard_dims"]
+    if len(num_shards) != len(shard_dims):
+        raise ValueError(
+            f"dc_num_shards {num_shards} must have same length as dc_shard_dims {shard_dims}"
+        )
+
+    shard_layout = sorted(
+        (int(shard_dim), int(num_shard))
+        for num_shard, shard_dim in zip(num_shards, shard_dims)
+    )
+    canonical_config["dc_shard_dims"] = [shard_dim for shard_dim, _ in shard_layout]
+    canonical_config["dc_num_shards"] = [num_shard for _, num_shard in shard_layout]
+    return canonical_config
+
+
 def _hash_volume_config(volume_config: Dict[str, Any]) -> str:
     s = json.dumps(volume_config, separators=(",", ":"), sort_keys=True).encode()
     return hashlib.sha256(s).hexdigest()[:12]
 
 
+def _volume_config_for_version(
+    config_dict, dataset_format_version, canonicalize_v3_shard_layout=True
+):
+    """Build the hashable dataset config subset for a format version."""
+
+    versioned_config = config_dict.copy()
+    versioned_config["dataset_format_version"] = dataset_format_version
+    if dataset_format_version == DATASET_FORMAT_VERSION:
+        include_keys = INCLUDE_KEYS
+    else:
+        include_keys = V2_INCLUDE_KEYS
+    volume_config = _get_required_keys_dict(
+        config=versioned_config,
+        include_keys=include_keys,
+    )
+    if (
+        dataset_format_version == DATASET_FORMAT_VERSION
+        and canonicalize_v3_shard_layout
+    ):
+        volume_config = _canonicalize_v3_shard_layout(volume_config)
+    return volume_config
+
+
+def _requested_unsharded_layout(config_dict: Dict[str, Any]) -> bool:
+    """Return whether the requested layout has exactly one physical shard."""
+
+    total_shards = 1
+    for value in config_dict["dc_num_shards"]:
+        total_shards *= int(value)
+    return total_shards == 1
+
+
 def _git_commit_short() -> str:
     try:
         return (
@@ -98,6 +155,38 @@ def _git_commit_short() -> str:
         return "no-commit-id"
 
 
+def _find_reusable_dataset(
+    root: Path,
+    config_id: str,
+    dataset_format_version: int,
+    commit: str,
+    require_commit: bool,
+) -> Optional[Path]:
+    """Find the newest reusable dataset matching format, config, and commit."""
+
+    base = root / config_id
+    if not base.exists():
+        return None
+
+    candidates = sorted(
+        (p for p in base.iterdir() if p.is_dir()), key=lambda p: p.name, reverse=True
+    )
+    for dataset_path in candidates:
+        meta_path = dataset_path / META_FILENAME
+        if not meta_path.exists():
+            continue
+        meta = yaml.safe_load(meta_path.read_text()) or {}
+        if meta.get("config_id") != config_id:
+            continue
+        if meta.get("dataset_format_version", 1) != dataset_format_version:
+            continue
+        if require_commit and meta.get("code_commit") != commit:
+            continue
+        return dataset_path
+
+    return None
+
+
 def get_dataset(
     config: Namespace,
     require_commit: bool = False,  # default: ignore commit mismatches for reuse
@@ -118,42 +207,58 @@ def get_dataset(
     root = Path(config.dataset_dir)
     root.mkdir(exist_ok=True)
 
-    # Get dict of required keys and compute config_id
+    # V3 is the current physical-shard format. The physical dataset layout is
+    # defined by dc_num_shards/dc_shard_dims, matching the DistConv layout.
     config_dict = vars(config).copy()
-    config_dict["dataset_format_version"] = DATASET_FORMAT_VERSION
-    volume_config = _get_required_keys_dict(
-        config=config_dict, include_keys=INCLUDE_KEYS
+    volume_config = _volume_config_for_version(config_dict, DATASET_FORMAT_VERSION)
+    metadata_volume_config = _volume_config_for_version(
+        config_dict,
+        DATASET_FORMAT_VERSION,
+        canonicalize_v3_shard_layout=False,
     )
     config_id = _hash_volume_config(volume_config)
+    v2_volume_config = _volume_config_for_version(
+        config_dict, V2_DATASET_FORMAT_VERSION
+    )
+    v2_config_id = _hash_volume_config(v2_volume_config)
     commit = _git_commit_short()
 
-    base = root / config_id
-    base.mkdir(parents=True, exist_ok=True)
-
-    # Try to reuse latest candidate dataset
-    candidates = sorted(
-        (p for p in base.iterdir() if p.is_dir()), key=lambda p: p.name, reverse=True
+    # Prefer a matching V3 physical-shard dataset.
+    dataset_path = _find_reusable_dataset(
+        root,
+        config_id,
+        DATASET_FORMAT_VERSION,
+        commit,
+        require_commit,
     )
-    for dataset_path in candidates:
-        meta_path = dataset_path / META_FILENAME
-        if not meta_path.exists():
-            continue
-        meta = yaml.safe_load(meta_path.read_text())
-        if meta.get("config_id") != config_id:
-            continue
-        if meta.get("dataset_format_version", 1) != DATASET_FORMAT_VERSION:
-            continue
-        if require_commit and meta.get("code_commit") != commit:
-            continue
-        # If we pass the above checks, this dataset can be reused
+    if dataset_path is not None:
         print(
-            "Valid existing dataset found. Reusing this dataset..."
+            "Valid existing v3 sharded dataset found. Reusing this dataset..."
         )  # FIXME replace with updated logging
         return dataset_path
 
+    # V2 datasets are full-volume files without shard suffixes. Reuse them only
+    # for unsharded requests so sharded generation never silently returns a
+    # cache that lacks the requested shard files.
+    if _requested_unsharded_layout(config_dict):
+        dataset_path = _find_reusable_dataset(
+            root,
+            v2_config_id,
+            V2_DATASET_FORMAT_VERSION,
+            commit,
+            require_commit,
+        )
+        if dataset_path is not None:
+            print(
+                "Valid existing v2 full-volume dataset found. Reusing this dataset..."
+            )  # FIXME replace with updated logging
+            return dataset_path
+
     # Otherwise, generate a new dataset
+    base = root / config_id
     print(f"No valid existing dataset found at {base}. Generating new dataset...")
     if rank == 0:
+        base.mkdir(parents=True, exist_ok=True)
         ts = time.strftime("%Y%m%d-%H%M%S")
         dest = base / f"{ts}__{commit}"
         tmp = base / f".tmp_{ts}"
@@ -176,33 +281,52 @@ def get_dataset(
 
     # Check that all ranks succeeded in volumegen, then sync
     all_ok = comm.allreduce(1 if ok else 0, op=MPI.MIN) == 1
-    comm.Barrier()
+    errs = comm.gather(err, root=0)
 
-    # rank 0 has file write + move
+    failure_msg = None
     if rank == 0:
         if not all_ok:
             try:
                 shutil.rmtree(tmp, ignore_errors=True)
             except Exception:
                 pass
-            # collect & raise a representative error
-            errs = comm.gather(err, root=0)
             msgs = "; ".join(e for e in errs if e)
-            raise RuntimeError(f"dataset generation failed: {msgs or 'unknown error'}")
-
-        # Write to tmp, then move, so readers never see half-written dataset
-        meta = {
-            "config_id": config_id,
-            "dataset_format_version": DATASET_FORMAT_VERSION,
-            "config_subset": volume_config,
-            "include_keys": INCLUDE_KEYS,
-            "code_commit": commit,
-            "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
-        }
-        (tmp / META_FILENAME).write_text(
-            yaml.safe_dump(meta, sort_keys=True, default_flow_style=False)
-        )
-        tmp.rename(dest)
+            failure_msg = f"dataset generation failed: {msgs or 'unknown error'}"
+
+    failure_msg = comm.bcast(failure_msg, root=0)
+    if failure_msg:
+        raise RuntimeError(failure_msg)
+
+    # rank 0 has file write + move
+    finalize_err = ""
+    if rank == 0:
+        try:
+            # Write to tmp, then move, so readers never see half-written dataset
+            meta = {
+                "config_id": config_id,
+                "dataset_format_version": DATASET_FORMAT_VERSION,
+                "config_subset": metadata_volume_config,
+                "include_keys": INCLUDE_KEYS,
+                "code_commit": commit,
+                "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+            }
+            (tmp / META_FILENAME).write_text(
+                yaml.safe_dump(meta, sort_keys=True, default_flow_style=False)
+            )
+            tmp.rename(dest)
+        except Exception as e:
+            finalize_err = (
+                f"dataset finalization failed: rank 0: {type(e).__name__}: {e}"
+            )
+
+    finalize_err = comm.bcast(finalize_err, root=0)
+    if finalize_err:
+        if rank == 0:
+            try:
+                shutil.rmtree(tmp, ignore_errors=True)
+            except Exception:
+                pass
+        raise RuntimeError(finalize_err)
 
     # ensure the rename is visible everywhere before returning
     comm.Barrier()