diff --git a/README.md b/README.md
index f2827e8..a01983a 100644
--- a/README.md
+++ b/README.md
@@ -59,6 +59,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]` and are still stored as v3 datasets with `_shard000000` files. ScaFFold only reuses matching v3 caches; older full-volume caches are ignored and a new v3 dataset is generated instead.
+
 ScaFFold benchmark training always uses PyTorch distributed execution with DistConv spatial parallelism. For a singleton run, launch one distributed rank rather than disabling distributed execution.
 
 `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.
@@ -168,6 +174,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.
+
 ### Performance Profiling
 
 #### 1. Profiling with the PyTorch Profiler
diff --git a/ScaFFold/cli.py b/ScaFFold/cli.py
index 130b3ef..599983f 100644
--- a/ScaFFold/cli.py
+++ b/ScaFFold/cli.py
@@ -115,6 +115,11 @@ def main():
     benchmark_parser.add_argument(
         "--base-run-dir", type=str, help="Subfolder of $(pwd) in which to run jobs."
     )
+    benchmark_parser.add_argument(
+        "--dataset-dir",
+        type=str,
+        help="Directory in which to store and query benchmark datasets.",
+    )
     benchmark_parser.add_argument(
         "--fract-base-dir",
         type=str,
@@ -161,6 +166,23 @@ def main():
         type=int,
         help="Number of DataLoader worker processes per rank.",
     )
+    benchmark_parser.add_argument(
+        "--dataloader-pin-memory",
+        type=int,
+        choices=[0, 1],
+        help="Whether DataLoader should pin host memory before device transfer.",
+    )
+    benchmark_parser.add_argument(
+        "--dataloader-persistent-workers",
+        type=int,
+        choices=[0, 1],
+        help="Whether DataLoader workers should persist across epochs.",
+    )
+    benchmark_parser.add_argument(
+        "--dataloader-prefetch-factor",
+        type=int,
+        help="Number of batches prefetched by each DataLoader worker.",
+    )
     benchmark_parser.add_argument(
         "--optimizer",
         type=str,
diff --git a/ScaFFold/datagen/category_search.py b/ScaFFold/datagen/category_search.py
index 3e40607..58dbd63 100644
--- a/ScaFFold/datagen/category_search.py
+++ b/ScaFFold/datagen/category_search.py
@@ -198,8 +198,10 @@ def main(config: Config) -> None:
     log = setup_mpi_logger(__file__, getattr(config, "verbose", 0))
 
     datagen_batch_size = int(getattr(config, "datagen_batch_size", 10000))
-    if datagen_batch_size <= 0:
-        raise ValueError("datagen_batch_size must be positive")
+    if datagen_batch_size < 1:
+        raise ValueError(
+            f"datagen_batch_size must be positive, got {datagen_batch_size}"
+        )
 
     # FIXME anything else to ensure determinism?
     np.random.seed(config.seed + rank)
diff --git a/ScaFFold/datagen/get_dataset.py b/ScaFFold/datagen/get_dataset.py
index f7d8861..b5f2f41 100644
--- a/ScaFFold/datagen/get_dataset.py
+++ b/ScaFFold/datagen/get_dataset.py
@@ -21,7 +21,7 @@
 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
@@ -30,7 +30,7 @@
 from ScaFFold.utils.utils import setup_mpi_logger
 
 META_FILENAME = "meta.yaml"
-DATASET_FORMAT_VERSION = 2
+DATASET_FORMAT_VERSION = 3
 INCLUDE_KEYS = [
     "dataset_format_version",
     "n_categories",
@@ -40,6 +40,8 @@
     "variance_threshold",
     "n_fracts_per_vol",
     "val_split",
+    "dc_num_shards",
+    "dc_shard_dims",
 ]
 
 
@@ -72,11 +74,45 @@ 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(config_dict, canonicalize_shard_layout=True):
+    """Build the hashable v3 dataset config subset."""
+
+    versioned_config = config_dict.copy()
+    versioned_config["dataset_format_version"] = DATASET_FORMAT_VERSION
+    volume_config = _get_required_keys_dict(
+        config=versioned_config,
+        include_keys=INCLUDE_KEYS,
+    )
+    if canonicalize_shard_layout:
+        volume_config = _canonicalize_v3_shard_layout(volume_config)
+    return volume_config
+
+
 def _git_commit_short(log) -> str:
     try:
         return (
@@ -101,6 +137,37 @@ def _git_commit_short(log) -> str:
         return "no-commit-id"
 
 
+def _find_reusable_dataset(
+    root: Path,
+    config_id: str,
+    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") != 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
@@ -122,40 +189,33 @@ 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(config_dict)
+    metadata_volume_config = _volume_config(
+        config_dict,
+        canonicalize_shard_layout=False,
     )
     config_id = _hash_volume_config(volume_config)
     commit = _git_commit_short(log)
 
-    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
+    # Reuse only matching v3 physical-shard datasets.
+    dataset_path = _find_reusable_dataset(
+        root,
+        config_id,
+        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
-        log.info("Reusing existing dataset at %s", dataset_path)
+    if dataset_path is not None:
+        log.info("Reusing existing v3 dataset at %s", dataset_path)
         return dataset_path
 
     # Otherwise, generate a new dataset
+    base = root / config_id
     log.info("No valid existing dataset found at %s. Generating new dataset.", base)
     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}"
@@ -178,34 +238,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)
 
-    errs = comm.gather(err, root=0) if not all_ok else None
-    if not all_ok:
-        if rank == 0:
+    failure_msg = None
+    if rank == 0:
+        if not all_ok:
             try:
                 shutil.rmtree(tmp, ignore_errors=True)
             except Exception:
                 pass
             msgs = "; ".join(e for e in errs if e)
-            raise RuntimeError(f"dataset generation failed: {msgs or 'unknown error'}")
-        raise RuntimeError("dataset generation failed on another rank")
+            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:
-        # 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)
+        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()
diff --git a/ScaFFold/datagen/volumegen.py b/ScaFFold/datagen/volumegen.py
index 1184e28..98ae9a4 100644
--- a/ScaFFold/datagen/volumegen.py
+++ b/ScaFFold/datagen/volumegen.py
@@ -18,13 +18,20 @@
 import random
 import time
 from math import ceil
-from typing import Dict
+from typing import Callable, Dict
 
 import numpy as np
 from mpi4py import MPI
 
 from ScaFFold.utils.config_utils import Config
 from ScaFFold.utils.data_types import DEFAULT_NP_DTYPE, MASK_DTYPE, VOLUME_DTYPE
+from ScaFFold.utils.spatial_sharding import (
+    normalize_sharding,
+    shard_file_suffix,
+    shard_id_to_indices,
+    spatial_slices,
+    total_shards,
+)
 from ScaFFold.utils.utils import setup_mpi_logger
 
 
@@ -43,7 +50,24 @@ def points_to_voxelgrid(
     Convert an (N,3) float64 point cloud directly into a boolean voxel grid
     of shape (grid_size, grid_size, grid_size).
     """
-    # 1) Axis‐aligned bounding box in float64
+    idx = points_to_voxel_indices(points, grid_size, eps=eps)
+
+    # Scatter into a boolean grid
+    grid = np.zeros((grid_size, grid_size, grid_size), dtype=bool)
+    grid[idx[:, 0], idx[:, 1], idx[:, 2]] = True
+
+    return grid
+
+
+def points_to_voxel_indices(
+    points: np.ndarray, grid_size: int, eps: float = 1e-6
+) -> np.ndarray:
+    """
+    Convert an (N,3) point cloud into global voxel indices using the same
+    math as points_to_voxelgrid(), without allocating a full boolean grid.
+    """
+
+    # 1) Axis-aligned bounding box in float64
     mins = points.min(axis=0)
     maxs = points.max(axis=0)
 
@@ -52,16 +76,111 @@ def points_to_voxelgrid(
 
     # 3) Map points into [0,grid_size) indices
     scaled = (points - mins) / voxel_size
-    idx = np.floor(scaled).astype(int)
+    np.floor(scaled, out=scaled)
+    idx = scaled.astype(int)
 
     # 4) Clip to valid range
-    idx = np.clip(idx, 0, grid_size - 1)
+    np.clip(idx, 0, grid_size - 1, out=idx)
 
-    # 5) Scatter into a boolean grid
-    grid = np.zeros((grid_size, grid_size, grid_size), dtype=bool)
-    grid[idx[:, 0], idx[:, 1], idx[:, 2]] = True
+    return idx
 
-    return grid
+
+def _physical_sharding(config):
+    """Return normalized physical sharding from the generation config."""
+
+    return normalize_sharding(config.dc_num_shards, config.dc_shard_dims)
+
+
+def _validate_generation_config(config):
+    """Validate sharded generation settings and return normalized layout data."""
+
+    num_shards, shard_dims = _physical_sharding(config)
+    n_total_shards = total_shards(num_shards)
+
+    grid_size = math.floor(config.vol_size * config.scale)
+    if grid_size != config.vol_size:
+        raise ValueError(
+            "Sharded volume generation currently requires config.scale == 1 so shard files tile the full volume"
+        )
+
+    return num_shards, shard_dims, n_total_shards, grid_size
+
+
+def _voxelized_fractals_for_volume(
+    config,
+    curr_vol: np.ndarray,
+    fractal_colors: np.ndarray,
+    point_cloud_loader: Callable[[str], np.ndarray] = load_np_ptcloud,
+):
+    """Load and voxelize all fractals needed for one logical volume."""
+
+    n_fracts_per_vol = config.n_fracts_per_vol
+    grid_size = math.floor(config.vol_size * config.scale)
+    voxelized_fractals = []
+
+    for curr_fract in range(n_fracts_per_vol):
+        curr_category = int(curr_vol[1 + 2 * curr_fract])
+        curr_instance = int(curr_vol[1 + 2 * curr_fract + 1])
+        fractal_color = fractal_colors[curr_category]
+
+        instances_dir = f"var{config.variance_threshold}/instances/np{config.point_num}"
+        point_cloud_path = os.path.join(
+            str(config.fract_base_dir),
+            instances_dir,
+            f"{curr_category:06d}",
+            f"{curr_category:06d}_{curr_instance:04d}.npy",
+        )
+        if point_cloud_loader is load_np_ptcloud and not os.path.exists(
+            point_cloud_path
+        ):
+            raise FileNotFoundError(
+                f"File {point_cloud_path} does not exist. Ensure you have run 'scaffold generate_fractals ...'"
+            )
+
+        points = point_cloud_loader(point_cloud_path)
+        idx = points_to_voxel_indices(points, grid_size)
+        voxelized_fractals.append((curr_category, fractal_color, idx))
+
+    return voxelized_fractals
+
+
+def _render_volume_shard(config, voxelized_fractals, shard_id: int):
+    """Render one physical shard from precomputed global voxel indices."""
+
+    num_shards, shard_dims = _physical_sharding(config)
+    shard_indices = shard_id_to_indices(shard_id, num_shards)
+    slices = spatial_slices(
+        (config.vol_size, config.vol_size, config.vol_size),
+        shard_dims,
+        num_shards,
+        shard_indices,
+    )
+    local_shape = tuple(s.stop - s.start for s in slices)
+
+    volume = np.full((3, *local_shape), 0, dtype=VOLUME_DTYPE)
+    mask = np.full(local_shape, 0, dtype=MASK_DTYPE)
+
+    for curr_category, fractal_color, idx in voxelized_fractals:
+        keep = np.ones(idx.shape[0], dtype=bool)
+        for axis, axis_slice in enumerate(slices):
+            keep &= idx[:, axis] >= axis_slice.start
+            keep &= idx[:, axis] < axis_slice.stop
+
+        if not np.any(keep):
+            continue
+
+        local_idx = idx[keep]
+        local_idx[:, 0] -= slices[0].start
+        local_idx[:, 1] -= slices[1].start
+        local_idx[:, 2] -= slices[2].start
+        d = local_idx[:, 0]
+        h = local_idx[:, 1]
+        w = local_idx[:, 2]
+
+        volume[:, d, h, w] = fractal_color[:, None]
+        mask[d, h, w] = curr_category + 1
+
+    return volume, mask
 
 
 def main(config: Dict):
@@ -78,176 +197,202 @@ def main(config: Dict):
     volumes_contents_path = os.path.join(dataset_dir, "volumes_contents.csv")
 
     n_fracts_per_vol = config.n_fracts_per_vol
+    _, _, n_total_shards, _ = _validate_generation_config(config)
 
     random.seed(config.seed)  # Python
     np.random.seed(config.seed)  # NumPy
 
     # Set up directories and select instances from each category
     volumes_contents = None
+    setup_err = ""
 
     if rank == 0:
-        if not os.path.exists(dataset_dir):
-            os.makedirs(dataset_dir)
-        for subdir in ["training", "validation"]:
-            os.makedirs(os.path.join(vol_path, subdir), exist_ok=True)
-            os.makedirs(os.path.join(mask_path, subdir), exist_ok=True)
-
-        # Force n_instances_used_per_fractal to be multiple of n_fracts_per_vol
-        if config.n_instances_used_per_fractal % n_fracts_per_vol != 0:
-            log.warning(
-                "n_instances_used_per_fractal (%s) is not a multiple of "
-                "n_fracts_per_vol=%s. Rounding down.",
-                config.n_instances_used_per_fractal,
-                n_fracts_per_vol,
-            )
-            config.n_instances_used_per_fractal = (
-                config.n_instances_used_per_fractal
-                // n_fracts_per_vol
-                * n_fracts_per_vol
-            )
-
-        # Randomly select n_instances_used_per_fractal instances from each fractal class.
-        instances_list = []
-        for category in range(config.n_categories):
-            instances_remaining = config.n_instances_used_per_fractal
-            random_instances = []
-            while instances_remaining > 0:
-                random_instances.extend(
-                    random.sample(range(145), min(145, instances_remaining))
+        try:
+            if not os.path.exists(dataset_dir):
+                os.makedirs(dataset_dir)
+            for subdir in ["training", "validation"]:
+                os.makedirs(os.path.join(vol_path, subdir), exist_ok=True)
+                os.makedirs(os.path.join(mask_path, subdir), exist_ok=True)
+
+            # Force n_instances_used_per_fractal to be multiple of n_fracts_per_vol
+            if config.n_instances_used_per_fractal % n_fracts_per_vol != 0:
+                log.warning(
+                    "n_instances_used_per_fractal (%s) is not a multiple of "
+                    "n_fracts_per_vol=%s. Rounding down.",
+                    config.n_instances_used_per_fractal,
+                    n_fracts_per_vol,
+                )
+                config.n_instances_used_per_fractal = (
+                    config.n_instances_used_per_fractal
+                    // n_fracts_per_vol
+                    * n_fracts_per_vol
                 )
-                instances_remaining -= min(145, instances_remaining)
 
-            category_instance_pairs = [
-                [category, instance] for instance in random_instances
-            ]
-            instances_list.extend(category_instance_pairs)
+            # Randomly select n_instances_used_per_fractal instances from each fractal class.
+            instances_list = []
+            for category in range(config.n_categories):
+                instances_remaining = config.n_instances_used_per_fractal
+                random_instances = []
+                while instances_remaining > 0:
+                    random_instances.extend(
+                        random.sample(range(145), min(145, instances_remaining))
+                    )
+                    instances_remaining -= min(145, instances_remaining)
 
-        instances_list = np.array(instances_list, dtype=int)
-        np.random.shuffle(instances_list)
+                category_instance_pairs = [
+                    [category, instance] for instance in random_instances
+                ]
+                instances_list.extend(category_instance_pairs)
 
-        volumes_contents = instances_list.reshape(-1, 2 * n_fracts_per_vol)
+            instances_list = np.array(instances_list, dtype=int)
+            np.random.shuffle(instances_list)
 
-        indices = np.arange(volumes_contents.shape[0]).reshape(-1, 1)
-        volumes_contents = np.hstack([indices, volumes_contents])
+            volumes_contents = instances_list.reshape(-1, 2 * n_fracts_per_vol)
 
-        with open(volumes_contents_path, "wb") as f:
-            np.savetxt(f, volumes_contents.astype(int), fmt="%i", delimiter=",")
-        log.info(
-            "Finished writing volumes_contents with shape %s", volumes_contents.shape
-        )
+            indices = np.arange(volumes_contents.shape[0]).reshape(-1, 1)
+            volumes_contents = np.hstack([indices, volumes_contents])
+
+            with open(volumes_contents_path, "wb") as f:
+                np.savetxt(f, volumes_contents.astype(int), fmt="%i", delimiter=",")
+            log.info(
+                "Finished writing volumes_contents with shape %s",
+                volumes_contents.shape,
+            )
+        except Exception as e:
+            setup_err = f"setup failed: rank {rank}: {type(e).__name__}: {e}"
 
     # Broadcast to all ranks
-    volumes_contents = comm.bcast(volumes_contents, root=0)
+    volumes_contents, setup_err = comm.bcast((volumes_contents, setup_err), root=0)
+    if setup_err:
+        raise RuntimeError(setup_err)
+
+    # Rank 0 creates shared metadata above; wait before local writer setup.
+    comm.Barrier()
+
+    for subdir in ["training", "validation"]:
+        os.makedirs(os.path.join(vol_path, subdir), exist_ok=True)
+        os.makedirs(os.path.join(mask_path, subdir), exist_ok=True)
+
+    # Wait until every rank has ensured the writer directories exist.
+    comm.Barrier()
 
     # Determine train/val split globally so all ranks know where to save
     num_volumes = len(volumes_contents)
-    random.seed(config.seed)  # Reset seed to ensure all ranks get same split
+    random.seed(config.seed)
     val_indices = set(
         random.sample(range(num_volumes), int(num_volumes * config.val_split / 100))
     )
 
-    # Work distribution
-    num_volumes = len(volumes_contents)
-    stride = ceil(num_volumes / size)
+    # Work distribution: each task renders one physical shard of one logical volume.
+    total_tasks = num_volumes * n_total_shards
+    stride = ceil(total_tasks / size)
     start_idx = rank * stride
-    end_idx = min(((rank + 1) * stride), num_volumes)
-
-    if start_idx >= end_idx:
-        log.debug("Rank %s given no volumes to generate", rank)
-
-    else:
-        volumes_contents_subset = volumes_contents[start_idx:end_idx]
-        log.debug(
-            "Rank %s responsible for volumes %s through %s",
-            rank,
-            volumes_contents_subset[0][0],
-            volumes_contents_subset[-1][0],
-        )
-
-        np.random.seed(config.seed)
-        fractal_colors = np.random.rand(config.n_categories, 3)
-
-        grid_size = math.floor(config.vol_size * config.scale)
-        fract_base_dir = str(config.fract_base_dir)
+    end_idx = min(((rank + 1) * stride), total_tasks)
 
-        # Generation loop
-        start_time = time.time()
-        for i, curr_vol in enumerate(volumes_contents_subset):
-            if i % 10 == 0:
-                log.debug("Rank %s processing local volume %s", rank, i)
+    generation_err = ""
+    n_generated_shards = 0
+    try:
+        if start_idx >= end_idx:
+            log.info("Rank %s given no physical shard tasks to generate", rank)
 
-            volume = np.full(
-                (config.vol_size, config.vol_size, config.vol_size, 3),
-                0,
-                dtype=VOLUME_DTYPE,
-            )
-            mask = np.full(
-                (config.vol_size, config.vol_size, config.vol_size), 0, dtype=MASK_DTYPE
+        else:
+            log.info(
+                "Rank %s responsible for physical shard tasks %s through %s",
+                rank,
+                start_idx,
+                end_idx - 1,
             )
 
-            global_vol_idx = curr_vol[0]
-            vol_seed = config.seed + int(global_vol_idx)
-            random.seed(vol_seed)
-            np.random.seed(vol_seed)
+            np.random.seed(config.seed)
+            fractal_colors = np.random.rand(
+                max(config.n_categories, n_fracts_per_vol), 3
+            )
 
-            for curr_fract in range(n_fracts_per_vol):
-                curr_category = curr_vol[1 + 2 * curr_fract]
-                curr_instance = curr_vol[1 + 2 * curr_fract + 1]
-                fractal_color = fractal_colors[curr_category]
+            # Generation loop
+            start_time = time.time()
+            cached_volume_idx = None
+            cached_global_vol_idx = None
+            cached_voxelized_fractals = None
+
+            for i, task_idx in enumerate(range(start_idx, end_idx)):
+                if i % 10 == 0:
+                    log.debug(
+                        "Rank %s processing local physical shard task %s",
+                        rank,
+                        i,
+                    )
 
-                instances_dir = (
-                    f"var{config.variance_threshold}/instances/np{config.point_num}"
-                )
+                volume_idx = task_idx // n_total_shards
+                shard_id = task_idx % n_total_shards
 
-                point_cloud_path = os.path.join(
-                    fract_base_dir,
-                    instances_dir,
-                    f"{curr_category:06d}",
-                    f"{curr_category:06d}_{curr_instance:04d}.npy",
-                )
+                if cached_volume_idx != volume_idx:
+                    curr_vol = volumes_contents[volume_idx]
+                    global_vol_idx = int(curr_vol[0])
+                    vol_seed = config.seed + global_vol_idx
+                    random.seed(vol_seed)
+                    np.random.seed(vol_seed)
 
-                if not os.path.exists(point_cloud_path):
-                    raise FileNotFoundError(
-                        f"File {point_cloud_path} does not exist. "
-                        "Ensure you have run 'scaffold generate_fractals ...'"
+                    cached_voxelized_fractals = _voxelized_fractals_for_volume(
+                        config,
+                        curr_vol,
+                        fractal_colors,
                     )
+                    cached_volume_idx = volume_idx
+                    cached_global_vol_idx = global_vol_idx
 
-                points = load_np_ptcloud(point_cloud_path)
-                mask3d = points_to_voxelgrid(points, grid_size)
-
-                assert mask3d.shape == volume.shape[:3], (
-                    f"mask3d {mask3d.shape} != volume spatial dims {volume.shape[:3]}"
+                volume_to_save, mask_to_save = _render_volume_shard(
+                    config,
+                    cached_voxelized_fractals,
+                    shard_id,
                 )
 
-                volume[mask3d] = fractal_color
-                mask[mask3d] = curr_category + 1
-
-            # Determine destination folder
-            subdir = "validation" if global_vol_idx in val_indices else "training"
-            # Tensors must logically be channels-first, later we will change striding/storage to channels-last on GPU (metadata will always stay channels-first).
-            volume_channels_first = volume.transpose((3, 0, 1, 2))
-            volume_to_save = np.ascontiguousarray(
-                volume_channels_first, dtype=VOLUME_DTYPE
-            )
-            mask_to_save = np.ascontiguousarray(mask, dtype=MASK_DTYPE)
+                # Determine destination folder
+                subdir = (
+                    "validation" if cached_global_vol_idx in val_indices else "training"
+                )
+                shard_suffix = shard_file_suffix(shard_id)
 
-            vol_file = os.path.join(vol_path, subdir, f"{global_vol_idx}.npy")
-            with open(vol_file, "wb") as f:
-                np.save(f, volume_to_save)
+                vol_file = os.path.join(
+                    vol_path, subdir, f"{cached_global_vol_idx}{shard_suffix}.npy"
+                )
+                with open(vol_file, "wb") as f:
+                    np.save(f, volume_to_save)
 
-            mask_file = os.path.join(mask_path, subdir, f"{global_vol_idx}_mask.npy")
-            with open(mask_file, "wb") as f:
-                np.save(f, mask_to_save)
+                mask_file = os.path.join(
+                    mask_path,
+                    subdir,
+                    f"{cached_global_vol_idx}{shard_suffix}_mask.npy",
+                )
+                with open(mask_file, "wb") as f:
+                    np.save(f, mask_to_save)
+                n_generated_shards += 1
+
+            end_time = time.time()
+            total_time = end_time - start_time
+            if rank == 0:
+                shard_rate = n_generated_shards / total_time
+                log.info(
+                    "Rank 0 generated %s volume shards from %s physical shard "
+                    "tasks in %.2f seconds | %.2f shards per second",
+                    n_generated_shards,
+                    end_idx - start_idx,
+                    total_time,
+                    shard_rate,
+                )
+    except Exception as e:
+        generation_err = (
+            f"volume shard generation failed: rank {rank}: {type(e).__name__}: {e}"
+        )
 
-        end_time = time.time()
-        total_time = end_time - start_time
-        if rank == 0:
-            log.info(
-                "Rank 0 generated %s volumes in %.2f seconds | %.2f volumes per second",
-                len(volumes_contents_subset),
-                total_time,
-                len(volumes_contents_subset) / total_time,
-            )
+    all_generated = comm.allreduce(1 if not generation_err else 0, op=MPI.MIN) == 1
+    generation_errs = comm.gather(generation_err, root=0)
+    generation_failure = ""
+    if rank == 0 and not all_generated:
+        msgs = "; ".join(e for e in generation_errs if e)
+        generation_failure = msgs or "unknown volume shard generation error"
+    generation_failure = comm.bcast(generation_failure, root=0)
+    if generation_failure:
+        raise RuntimeError(generation_failure)
 
     # Barrier to ensure all ranks are finished writing
     comm.Barrier()
diff --git a/ScaFFold/utils/config_utils.py b/ScaFFold/utils/config_utils.py
index c13d611..8136d4b 100644
--- a/ScaFFold/utils/config_utils.py
+++ b/ScaFFold/utils/config_utils.py
@@ -62,6 +62,15 @@ def __init__(self, config_dict):
         self.scale = 1
         self.local_batch_size = config_dict["local_batch_size"]
         self.dataloader_num_workers = config_dict["dataloader_num_workers"]
+        self.dataloader_pin_memory = bool(
+            config_dict.get("dataloader_pin_memory", True)
+        )
+        self.dataloader_persistent_workers = bool(
+            config_dict.get("dataloader_persistent_workers", True)
+        )
+        self.dataloader_prefetch_factor = config_dict.get(
+            "dataloader_prefetch_factor", 2
+        )
         self.epochs = config_dict["epochs"]
         self.optimizer = config_dict["optimizer"]
         self.disable_scheduler = bool(config_dict["disable_scheduler"])
diff --git a/ScaFFold/utils/data_loading.py b/ScaFFold/utils/data_loading.py
index 688f329..493759e 100644
--- a/ScaFFold/utils/data_loading.py
+++ b/ScaFFold/utils/data_loading.py
@@ -13,11 +13,12 @@
 # SPDX-License-Identifier: (Apache-2.0)
 
 import pickle
+import re
 from dataclasses import dataclass
 from os import listdir
-from os.path import isfile, join, splitext
+from os.path import isfile, join
 from pathlib import Path
-from typing import Dict, Optional, Tuple
+from typing import Optional, Tuple
 
 import numpy as np
 import torch
@@ -25,10 +26,15 @@
 from torch.utils.data import Dataset
 
 from ScaFFold.utils.data_types import MASK_DTYPE, VOLUME_DTYPE
+from ScaFFold.utils.spatial_sharding import (
+    normalize_sharding,
+    shard_file_suffix,
+    shard_indices_to_id,
+    total_shards,
+)
 from ScaFFold.utils.utils import customlog
 
-DATASET_FORMAT_VERSION = 2
-LEGACY_DATASET_FORMAT_VERSION = 1
+DATASET_FORMAT_VERSION = 3
 META_FILENAME = "meta.yaml"
 
 
@@ -65,42 +71,6 @@ def __post_init__(self):
                     f"Invalid shard_index {shard_index} for shard_dim {shard_dim} with {num_shards} shards"
                 )
 
-    @staticmethod
-    def _chunk_slice(size: int, num_shards: int, shard_index: int) -> slice:
-        """Match torch.chunk-style uneven shard boundaries."""
-
-        chunk_size = (size + num_shards - 1) // num_shards
-        start = shard_index * chunk_size
-        if start >= size:
-            raise ValueError(
-                f"Empty local shard: dim size {size}, num_shards {num_shards}, shard_index {shard_index}"
-            )
-        stop = min(size, start + chunk_size)
-        return slice(start, stop)
-
-    def slice_array(
-        self, array: np.ndarray, axis_map: Dict[int, int], array_label: str
-    ) -> np.ndarray:
-        if not self.shard_dims:
-            return array
-
-        slices = [slice(None)] * array.ndim
-        for shard_dim, num_shards, shard_index in zip(
-            self.shard_dims, self.num_shards, self.shard_indices
-        ):
-            if shard_dim not in axis_map:
-                raise ValueError(
-                    f"No axis mapping defined for {array_label} shard_dim {shard_dim}"
-                )
-            axis = axis_map[shard_dim]
-            if axis >= array.ndim:
-                raise ValueError(
-                    f"Axis {axis} out of range for {array_label} with shape {array.shape}"
-                )
-            slices[axis] = self._chunk_slice(array.shape[axis], num_shards, shard_index)
-
-        return array[tuple(slices)]
-
 
 class BasicDataset(Dataset):
     def __init__(
@@ -116,13 +86,26 @@ def __init__(
         self.mask_suffix = mask_suffix
         self.spatial_shard_spec = spatial_shard_spec
         self.dataset_root = self.images_dir.parents[1]
-        self.dataset_format_version = self._load_dataset_format_version()
+        self.dataset_meta = self._load_dataset_metadata()
+        if "dataset_format_version" not in self.dataset_meta:
+            raise RuntimeError(
+                f"Dataset metadata at {self.dataset_root / META_FILENAME} is missing "
+                "dataset_format_version. Only v3 datasets are supported."
+            )
+        self.dataset_format_version = int(self.dataset_meta["dataset_format_version"])
+        if self.dataset_format_version != DATASET_FORMAT_VERSION:
+            raise RuntimeError(
+                f"Unsupported dataset format version {self.dataset_format_version}; "
+                f"expected {DATASET_FORMAT_VERSION}"
+            )
+        self.physical_num_shards, self.physical_shard_dims = (
+            self._load_physical_sharding()
+        )
+        self.physical_total_shards = total_shards(self.physical_num_shards)
+        self.shard_id = self._select_physical_shard_id()
+        self.shard_suffix = shard_file_suffix(self.shard_id)
 
-        self.ids = [
-            splitext(file)[0]
-            for file in listdir(images_dir)
-            if isfile(join(images_dir, file)) and not file.startswith(".")
-        ]
+        self.ids = self._list_ids(images_dir)
         if not self.ids:
             raise RuntimeError(
                 f"No input file found in {images_dir}, make sure you put your images there"
@@ -136,97 +119,155 @@ def __init__(
         self.mask_values = data["mask_values"]
         customlog(f"Unique mask values: {self.mask_values}")
         customlog(f"Dataset format version: {self.dataset_format_version}")
+        customlog(
+            f"Loading physical shard files with suffix {self.shard_suffix}; "
+            f"dc_num_shards={self.physical_num_shards}, "
+            f"dc_shard_dims={self.physical_shard_dims}"
+        )
 
     def __len__(self):
         return len(self.ids)
 
     @staticmethod
-    def _load_numpy_array(path, mmap_mode=None):
-        return np.load(path, allow_pickle=False, mmap_mode=mmap_mode)
+    def _load_numpy_array(path):
+        return np.load(path, allow_pickle=False)
+
+    def _list_ids(self, images_dir):
+        """List logical sample IDs visible to this dataset instance."""
+
+        pattern = re.compile(rf"^(?P<id>.+){re.escape(self.shard_suffix)}\.npy$")
+        ids = []
+        for file in listdir(images_dir):
+            if file.startswith(".") or not isfile(join(images_dir, file)):
+                continue
+            match = pattern.match(file)
+            if match is not None:
+                ids.append(match.group("id"))
+        return sorted(ids)
+
+    def _load_dataset_metadata(self):
+        """Load required v3 dataset metadata."""
 
-    def _load_dataset_format_version(self):
         meta_path = self.dataset_root / META_FILENAME
         if not meta_path.exists():
-            return LEGACY_DATASET_FORMAT_VERSION
+            raise RuntimeError(
+                f"Dataset metadata not found at {meta_path}. Only v3 datasets are supported."
+            )
 
         try:
             with open(meta_path, "r") as meta_file:
                 meta = yaml.safe_load(meta_file) or {}
         except Exception as exc:
-            customlog(
-                f"Failed to read dataset metadata from {meta_path}: {exc}. Falling back to legacy loader."
+            raise RuntimeError(
+                f"Failed to read dataset metadata from {meta_path}: {exc}"
+            ) from exc
+
+        if not isinstance(meta, dict):
+            raise RuntimeError(f"Dataset metadata at {meta_path} must be a mapping")
+
+        return meta
+
+    def _load_physical_sharding(self):
+        """Load and normalize the physical shard layout from metadata."""
+
+        config_subset = self.dataset_meta.get("config_subset") or {}
+        num_shards = config_subset.get("dc_num_shards")
+        shard_dims = config_subset.get("dc_shard_dims")
+        if num_shards is None or shard_dims is None:
+            raise RuntimeError(
+                "Physical dataset is missing shard metadata. Expected "
+                "config_subset.dc_num_shards/config_subset.dc_shard_dims in meta.yaml."
             )
-            return LEGACY_DATASET_FORMAT_VERSION
 
-        return int(meta.get("dataset_format_version", LEGACY_DATASET_FORMAT_VERSION))
+        return normalize_sharding(num_shards, shard_dims)
 
     @staticmethod
-    def _prepare_legacy_image(img):
-        return np.ascontiguousarray(img.transpose((3, 0, 1, 2)), dtype=VOLUME_DTYPE)
+    def _layout_by_dim(num_shards, shard_dims):
+        """Map each sharded dimension to its shard count."""
 
-    @staticmethod
-    def _prepare_legacy_mask(mask_values, mask):
-        remapped = np.zeros(
-            (mask.shape[0], mask.shape[1], mask.shape[2]), dtype=MASK_DTYPE
+        return {int(dim): int(num) for num, dim in zip(num_shards, shard_dims)}
+
+    def _physical_layout_matches_spatial_spec(self):
+        """Return whether dataset shards match the requested spatial layout."""
+
+        if self.spatial_shard_spec is None:
+            return False
+        return self._layout_by_dim(
+            self.physical_num_shards, self.physical_shard_dims
+        ) == self._layout_by_dim(
+            self.spatial_shard_spec.num_shards,
+            self.spatial_shard_spec.shard_dims,
         )
-        for i, value in enumerate(mask_values):
-            if mask.ndim == 3:
-                remapped[mask == value] = i
-            else:
-                remapped[(mask == value).all(-1)] = i
 
-        return remapped
+    def _physical_shard_id_for_spatial_spec(self):
+        """Return the physical shard id selected by the spatial shard spec."""
+
+        spec_indices_by_dim = {
+            int(dim): int(index)
+            for dim, index in zip(
+                self.spatial_shard_spec.shard_dims,
+                self.spatial_shard_spec.shard_indices,
+            )
+        }
+        shard_indices = tuple(
+            spec_indices_by_dim[int(dim)] for dim in self.physical_shard_dims
+        )
+        return shard_indices_to_id(shard_indices, self.physical_num_shards)
+
+    def _select_physical_shard_id(self):
+        """Select the physical shard file this dataset instance should read."""
+
+        if self.spatial_shard_spec is None:
+            if self.physical_total_shards == 1:
+                return 0
+            raise RuntimeError(
+                "Physical dataset has multiple shard files, but no SpatialShardSpec "
+                "was provided. Use a DistConv layout matching the v3 dataset."
+            )
+        if not self._physical_layout_matches_spatial_spec():
+            raise RuntimeError(
+                "V3 physical dataset shard layout does not match the requested "
+                "DistConv layout. V3 requires physical dataset layout and "
+                "DistConv layout to match. "
+                f"dataset dc_num_shards={self.physical_num_shards}, "
+                f"dataset dc_shard_dims={self.physical_shard_dims}, "
+                f"dc_num_shards={self.spatial_shard_spec.num_shards}, "
+                f"dc_shard_dims={self.spatial_shard_spec.shard_dims}"
+            )
+
+        return self._physical_shard_id_for_spatial_spec()
 
     @staticmethod
-    def _prepare_optimized_image(img):
+    def _prepare_image(img):
         return np.array(img, dtype=VOLUME_DTYPE, copy=True, order="C")
 
     @staticmethod
-    def _prepare_optimized_mask(mask):
+    def _prepare_mask(mask):
         return np.array(mask, dtype=MASK_DTYPE, copy=True, order="C")
 
-    def _slice_image_array(self, img):
-        if self.spatial_shard_spec is None:
-            return img
-
-        if self.dataset_format_version >= DATASET_FORMAT_VERSION:
-            axis_map = {2: 1, 3: 2, 4: 3}
-        else:
-            axis_map = {2: 0, 3: 1, 4: 2}
-        return self.spatial_shard_spec.slice_array(img, axis_map, "image")
+    def _resolve_sample_files(self, name):
+        """Resolve image and mask file paths for a logical sample ID."""
 
-    def _slice_mask_array(self, mask):
-        if self.spatial_shard_spec is None:
-            return mask
-
-        axis_map = {2: 0, 3: 1, 4: 2}
-        return self.spatial_shard_spec.slice_array(mask, axis_map, "mask")
+        img_file = self.images_dir / f"{name}{self.shard_suffix}.npy"
+        mask_file = (
+            self.mask_dir / f"{name}{self.shard_suffix}{self.mask_suffix}.npy"
+        )
+        assert img_file.is_file(), (
+            f"No image found for ID {name}, shard {self.shard_id}: {img_file}"
+        )
+        assert mask_file.is_file(), (
+            f"No mask found for ID {name}, shard {self.shard_id}: {mask_file}"
+        )
+        return img_file, mask_file
 
     def __getitem__(self, idx):
         name = self.ids[idx]
-        mask_file = list(self.mask_dir.glob(name + self.mask_suffix + ".*"))
-        img_file = list(self.images_dir.glob(name + ".*"))
+        img_file, mask_file = self._resolve_sample_files(name)
 
-        assert len(img_file) == 1, (
-            f"Either no image or multiple images found for the ID {name}: {img_file}"
-        )
-        assert len(mask_file) == 1, (
-            f"Either no mask or multiple masks found for the ID {name}: {mask_file}"
-        )
-        mmap_mode = "r" if self.spatial_shard_spec is not None else None
-        # Memmap lets each rank slice out just its local shard without eagerly
-        # reading the full sample into process memory first.
-        mask = self._load_numpy_array(mask_file[0], mmap_mode=mmap_mode)
-        img = self._load_numpy_array(img_file[0], mmap_mode=mmap_mode)
-        mask = self._slice_mask_array(mask)
-        img = self._slice_image_array(img)
-
-        if self.dataset_format_version >= DATASET_FORMAT_VERSION:
-            img = self._prepare_optimized_image(img)
-            mask = self._prepare_optimized_mask(mask)
-        else:
-            img = self._prepare_legacy_image(img)
-            mask = self._prepare_legacy_mask(self.mask_values, mask)
+        mask = self._load_numpy_array(mask_file)
+        img = self._load_numpy_array(img_file)
+        img = self._prepare_image(img)
+        mask = self._prepare_mask(mask)
 
         return {
             "image": torch.from_numpy(img).contiguous().float(),
diff --git a/ScaFFold/utils/spatial_sharding.py b/ScaFFold/utils/spatial_sharding.py
new file mode 100644
index 0000000..91664dc
--- /dev/null
+++ b/ScaFFold/utils/spatial_sharding.py
@@ -0,0 +1,133 @@
+# Copyright (c) 2014-2026, Lawrence Livermore National Security, LLC.
+# Produced at the Lawrence Livermore National Laboratory.
+# Written by the LBANN Research Team (B. Van Essen, et al.) listed in
+# the CONTRIBUTORS file. See the top-level LICENSE file for details.
+#
+# LLNL-CODE-697807.
+# All rights reserved.
+#
+# This file is part of LBANN: Livermore Big Artificial Neural Network
+# Toolkit. For details, see http://software.llnl.gov/LBANN or
+# https://github.com/LBANN and https://github.com/LBANN/ScaFFold.
+#
+# SPDX-License-Identifier: (Apache-2.0)
+
+from math import prod
+from typing import Iterable, Tuple
+
+
+def normalize_sharding(num_shards: Iterable[int], shard_dims: Iterable[int]):
+    """Validate and normalize spatial sharding config."""
+
+    num_shards = tuple(int(x) for x in num_shards)
+    shard_dims = tuple(int(x) for x in shard_dims)
+
+    if len(num_shards) != len(shard_dims):
+        raise ValueError(
+            f"num_shards {num_shards} must have same length as shard_dims {shard_dims}"
+        )
+    if len(set(shard_dims)) != len(shard_dims):
+        raise ValueError(f"Shard dimensions must be unique: {shard_dims}")
+
+    for num_shards_i, shard_dim_i in zip(num_shards, shard_dims):
+        if num_shards_i < 1:
+            raise ValueError(f"Invalid num_shards value {num_shards_i}")
+        if shard_dim_i not in (2, 3, 4):
+            raise ValueError(
+                f"Invalid shard_dim {shard_dim_i}: only 3D spatial dimensions 2, 3, and 4 are supported"
+            )
+
+    return num_shards, shard_dims
+
+
+def total_shards(num_shards: Iterable[int]) -> int:
+    """Return the total number of shards in a multi-dimensional layout."""
+
+    return prod(tuple(int(x) for x in num_shards))
+
+
+def shard_id_to_indices(shard_id: int, num_shards: Iterable[int]) -> Tuple[int, ...]:
+    """Convert row-major linear shard id to multi-dimensional shard indices."""
+
+    num_shards = tuple(int(x) for x in num_shards)
+    total = total_shards(num_shards)
+    if shard_id < 0 or shard_id >= total:
+        raise ValueError(
+            f"shard_id {shard_id} out of range for num_shards={num_shards}"
+        )
+
+    indices = []
+    linear_idx = int(shard_id)
+    stride = total
+    for num_shards_i in num_shards:
+        stride //= num_shards_i
+        indices.append(linear_idx // stride)
+        linear_idx %= stride
+    return tuple(indices)
+
+
+def shard_indices_to_id(shard_indices: Iterable[int], num_shards: Iterable[int]) -> int:
+    """Convert multi-dimensional shard indices to row-major linear shard id."""
+
+    shard_indices = tuple(int(x) for x in shard_indices)
+    num_shards = tuple(int(x) for x in num_shards)
+    if len(shard_indices) != len(num_shards):
+        raise ValueError(
+            f"shard_indices {shard_indices} must match num_shards {num_shards}"
+        )
+
+    shard_id = 0
+    stride = 1
+    for shard_index_i, num_shards_i in zip(
+        reversed(shard_indices), reversed(num_shards)
+    ):
+        if shard_index_i < 0 or shard_index_i >= num_shards_i:
+            raise ValueError(
+                f"Invalid shard index {shard_index_i} for num_shards={num_shards}"
+            )
+        shard_id += shard_index_i * stride
+        stride *= num_shards_i
+    return shard_id
+
+
+def chunk_slice(size: int, num_shards: int, shard_index: int) -> slice:
+    """Match torch.chunk-style uneven shard boundaries."""
+
+    chunk_size = (size + num_shards - 1) // num_shards
+    start = shard_index * chunk_size
+    if start >= size:
+        raise ValueError(
+            f"Empty local shard: dim size {size}, num_shards {num_shards}, shard_index {shard_index}"
+        )
+    stop = min(size, start + chunk_size)
+    return slice(start, stop)
+
+
+def spatial_slices(
+    spatial_shape: Iterable[int],
+    shard_dims: Iterable[int],
+    num_shards: Iterable[int],
+    shard_indices: Iterable[int],
+) -> Tuple[slice, slice, slice]:
+    """Return local D/H/W slices for DistConv spatial dims 2/3/4."""
+
+    spatial_shape = tuple(int(x) for x in spatial_shape)
+    if len(spatial_shape) != 3:
+        raise ValueError(f"Expected 3D spatial shape, got {spatial_shape}")
+
+    slices = [slice(0, size) for size in spatial_shape]
+    for shard_dim, num_shards_i, shard_index_i in zip(
+        shard_dims, num_shards, shard_indices
+    ):
+        spatial_axis = int(shard_dim) - 2
+        slices[spatial_axis] = chunk_slice(
+            spatial_shape[spatial_axis], int(num_shards_i), int(shard_index_i)
+        )
+
+    return tuple(slices)
+
+
+def shard_file_suffix(shard_id: int) -> str:
+    """Return the filename suffix for a physical shard id."""
+
+    return f"_shard{int(shard_id):06d}"
diff --git a/ScaFFold/utils/trainer.py b/ScaFFold/utils/trainer.py
index d8d6547..b497954 100644
--- a/ScaFFold/utils/trainer.py
+++ b/ScaFFold/utils/trainer.py
@@ -42,7 +42,7 @@
     compute_sharded_cross_entropy_loss,
 )
 from ScaFFold.utils.perf_measure import adiak_value, begin_code_region, end_code_region
-from ScaFFold.utils.utils import gather_and_print_mem
+from ScaFFold.utils.utils import gather_and_print_cpu_mem, gather_and_print_mem
 
 
 class BaseTrainer:
@@ -159,16 +159,31 @@ def create_dataloaders(self):
         self.create_sampler()
 
         num_workers = self.config.dataloader_num_workers
+        pin_memory = bool(getattr(self.config, "dataloader_pin_memory", True))
+        persistent_workers = bool(
+            getattr(self.config, "dataloader_persistent_workers", True)
+        )
+        prefetch_factor = getattr(self.config, "dataloader_prefetch_factor", 2)
         loader_args = dict(
             batch_size=self.config.local_batch_size,
             num_workers=num_workers,
-            pin_memory=True,
+            pin_memory=pin_memory,
         )
         if num_workers > 0:
-            loader_args["persistent_workers"] = True
-            loader_args["prefetch_factor"] = 2
-        self.log.debug(
-            f"dataloader num_workers={loader_args['num_workers']}, prefetch_factor={loader_args.get('prefetch_factor')}, persistent_workers={loader_args.get('persistent_workers', False)}, os.cpu_count()={os.cpu_count()}, self.world_size={self.world_size} "
+            loader_args["persistent_workers"] = persistent_workers
+            if prefetch_factor is not None:
+                prefetch_factor = int(prefetch_factor)
+                if prefetch_factor < 1:
+                    raise ValueError(
+                        f"dataloader_prefetch_factor must be >= 1, got {prefetch_factor}"
+                    )
+                loader_args["prefetch_factor"] = prefetch_factor
+        self.log.info(
+            f"DataLoader config: num_workers={loader_args['num_workers']}, "
+            f"pin_memory={loader_args['pin_memory']}, "
+            f"prefetch_factor={loader_args.get('prefetch_factor')}, "
+            f"persistent_workers={loader_args.get('persistent_workers', False)}, "
+            f"os.cpu_count()={os.cpu_count()}, world_size={self.world_size}"
         )
         self.train_loader = DataLoader(
             self.train_set, sampler=self.train_sampler, **loader_args
@@ -183,6 +198,7 @@ def create_dataloaders(self):
                 f"data_num_replicas={self.data_num_replicas}. "
                 "Reduce local_batch_size or adjust validation sharding."
             )
+        gather_and_print_cpu_mem(self.log, "after_dataloader_creation")
 
     def setup_training_components(self):
         """Set up the optimizer, scheduler, gradient scaler, and loss function."""
@@ -520,11 +536,12 @@ def _run_training_batch(
 
         batch_size = images_dc.shape[0]
         detached_loss = loss.detach()
+        detached_dice = batch_dice_score.detach()
 
         # Free memory aggressively
         del images_dc, true_masks_dc, masks_pred_dc
         del local_preds, local_labels, local_preds_softmax, local_labels_one_hot
-        del loss_ce, loss
+        del loss_ce, loss, batch_dice_score
 
         if log_peak_mem and self.world_rank == 0:
             peak_alloc = torch.cuda.max_memory_allocated() / (1024**3)
@@ -533,7 +550,7 @@ def _run_training_batch(
                 f"[MEM-PEAK] Peak alloc: {peak_alloc:.2f} GiB | Peak reserved: {peak_reserved:.2f} GiB",
             )
 
-        return batch_size, detached_loss, batch_dice_score
+        return batch_size, detached_loss, detached_dice
 
     def _sync_gather_minibatch_timer(self, minibatch_events):
         minibatch_events[-1][1].synchronize()
@@ -610,6 +627,7 @@ def warmup(self):
 
         torch.distributed.barrier()
         self.log.info(f"Done warmup. Took {int(time.time() - start_warmup)}s")
+        gather_and_print_cpu_mem(self.log, "after_warmup", epoch=0)
 
     def train(self):
         """
@@ -643,6 +661,7 @@ def train(self):
                 self.val_loader.sampler.set_epoch(epoch)
                 self.model.train()
                 self.optimizer.zero_grad(set_to_none=False)
+                gather_and_print_cpu_mem(self.log, "epoch_start", epoch=epoch)
 
                 estr = (
                     f"{epoch}"
@@ -698,6 +717,9 @@ def train(self):
                 # Calculate overall loss as average of per-batch loss
                 overall_loss = epoch_loss.item() / len(self.train_loader)
                 self.total_optimizer_steps += epoch_optimizer_steps
+                gather_and_print_cpu_mem(
+                    self.log, "after_training_epoch", epoch=epoch
+                )
 
                 #
                 # Evaluate model on validation set, update LR if necessary
@@ -727,6 +749,7 @@ def train(self):
                     dice_info, op=torch.distributed.ReduceOp.SUM
                 )
                 val_score = dice_info[0].item() / max(dice_info[1].item(), 1)
+                gather_and_print_cpu_mem(self.log, "after_validation", epoch=epoch)
                 if not self.config.disable_scheduler:
                     self.scheduler.step()
                 else:
diff --git a/ScaFFold/utils/utils.py b/ScaFFold/utils/utils.py
index 46288d0..30a8ba8 100644
--- a/ScaFFold/utils/utils.py
+++ b/ScaFFold/utils/utils.py
@@ -13,7 +13,9 @@
 # SPDX-License-Identifier: (Apache-2.0)
 
 import logging
+import os
 import random
+import socket
 import sys
 
 import torch
@@ -153,3 +155,361 @@ def gather_and_print_mem(log, tag=""):
                 )
     else:
         log.debug(stats)
+
+
+def _rss_bytes_from_proc(pid):
+    """Return VmRSS for pid from procfs, or 0 if it is unavailable."""
+    try:
+        with open(f"/proc/{pid}/status", "r") as status_file:
+            for line in status_file:
+                if line.startswith("VmRSS:"):
+                    parts = line.split()
+                    if len(parts) >= 2:
+                        return int(parts[1]) * 1024
+    except OSError:
+        pass
+    return 0
+
+
+def _child_pids_from_proc(pid):
+    try:
+        with open(f"/proc/{pid}/task/{pid}/children", "r") as children_file:
+            return [int(child) for child in children_file.read().split()]
+    except OSError:
+        return []
+
+
+def _descendant_pids_from_proc(pid):
+    descendants = []
+    stack = list(_child_pids_from_proc(pid))
+    while stack:
+        child = stack.pop()
+        descendants.append(child)
+        stack.extend(_child_pids_from_proc(child))
+    return descendants
+
+
+def _cpu_mem_stats_procfs():
+    pid = os.getpid()
+    child_pids = _descendant_pids_from_proc(pid)
+    child_rss = sum(_rss_bytes_from_proc(child_pid) for child_pid in child_pids)
+    return _rss_bytes_from_proc(pid), child_rss, len(child_pids)
+
+
+def _cpu_mem_stats_psutil():
+    try:
+        import psutil
+    except ImportError:
+        return None
+
+    try:
+        process = psutil.Process(os.getpid())
+        rss = process.memory_info().rss
+        child_rss = 0
+        child_count = 0
+        for child in process.children(recursive=True):
+            try:
+                child_rss += child.memory_info().rss
+                child_count += 1
+            except (psutil.NoSuchProcess, psutil.AccessDenied):
+                continue
+        return rss, child_rss, child_count
+    except (psutil.NoSuchProcess, psutil.AccessDenied):
+        return None
+
+
+def _read_text_file(path):
+    try:
+        with open(path, "r") as file:
+            return file.read().strip()
+    except OSError:
+        return None
+
+
+def _read_int_file(path):
+    text = _read_text_file(path)
+    if text is None or text == "" or text == "max":
+        return None
+    try:
+        return int(text)
+    except ValueError:
+        return None
+
+
+def _parse_key_value_file(path):
+    text = _read_text_file(path)
+    values = {}
+    if not text:
+        return values
+    for line in text.splitlines():
+        parts = line.split()
+        if len(parts) != 2:
+            continue
+        try:
+            values[parts[0]] = int(parts[1])
+        except ValueError:
+            continue
+    return values
+
+
+def _unescape_mountinfo_path(path):
+    return (
+        path.replace("\\040", " ")
+        .replace("\\011", "\t")
+        .replace("\\012", "\n")
+        .replace("\\134", "\\")
+    )
+
+
+def _find_cgroup_mount(fs_type, controller=None):
+    try:
+        with open("/proc/self/mountinfo", "r") as file:
+            for line in file:
+                fields = line.split()
+                if "-" not in fields:
+                    continue
+                sep = fields.index("-")
+                if sep + 3 > len(fields):
+                    continue
+                mount_point = _unescape_mountinfo_path(fields[4])
+                current_fs_type = fields[sep + 1]
+                super_options = fields[sep + 3] if sep + 3 < len(fields) else ""
+                if current_fs_type != fs_type:
+                    continue
+                if controller and controller not in super_options.split(","):
+                    continue
+                return mount_point
+    except OSError:
+        return None
+    return None
+
+
+def _proc_self_cgroup_entries():
+    entries = []
+    try:
+        with open("/proc/self/cgroup", "r") as file:
+            for line in file:
+                parts = line.strip().split(":", 2)
+                if len(parts) == 3:
+                    entries.append(tuple(parts))
+    except OSError:
+        pass
+    return entries
+
+
+def _cgroup_v2_path():
+    mount = _find_cgroup_mount("cgroup2")
+    if not mount:
+        return None, None
+    for hierarchy, controllers, rel_path in _proc_self_cgroup_entries():
+        if hierarchy == "0" and controllers == "":
+            return mount, rel_path or "/"
+    return None, None
+
+
+def _cgroup_v1_memory_path():
+    mount = _find_cgroup_mount("cgroup", "memory")
+    if not mount:
+        return None, None
+    for _hierarchy, controllers, rel_path in _proc_self_cgroup_entries():
+        if "memory" in controllers.split(","):
+            return mount, rel_path or "/"
+    return None, None
+
+
+def _join_cgroup_path(mount, rel_path):
+    return os.path.normpath(os.path.join(mount, rel_path.lstrip("/")))
+
+
+def _cgroup_mem_stats():
+    mount, rel_path = _cgroup_v2_path()
+    if mount:
+        cg_path = _join_cgroup_path(mount, rel_path)
+        if os.path.exists(os.path.join(cg_path, "memory.current")):
+            return {
+                "version": "v2",
+                "path": rel_path,
+                "current": _read_int_file(os.path.join(cg_path, "memory.current")),
+                "peak": _read_int_file(os.path.join(cg_path, "memory.peak")),
+                "max": _read_int_file(os.path.join(cg_path, "memory.max")),
+                "events": _parse_key_value_file(
+                    os.path.join(cg_path, "memory.events")
+                ),
+                "stat": _parse_key_value_file(os.path.join(cg_path, "memory.stat")),
+            }
+
+    mount, rel_path = _cgroup_v1_memory_path()
+    if mount:
+        cg_path = _join_cgroup_path(mount, rel_path)
+        if os.path.exists(os.path.join(cg_path, "memory.usage_in_bytes")):
+            return {
+                "version": "v1",
+                "path": rel_path,
+                "current": _read_int_file(
+                    os.path.join(cg_path, "memory.usage_in_bytes")
+                ),
+                "peak": _read_int_file(
+                    os.path.join(cg_path, "memory.max_usage_in_bytes")
+                ),
+                "max": _read_int_file(
+                    os.path.join(cg_path, "memory.limit_in_bytes")
+                ),
+                "events": {
+                    "failcnt": _read_int_file(os.path.join(cg_path, "memory.failcnt"))
+                    or 0
+                },
+                "stat": _parse_key_value_file(os.path.join(cg_path, "memory.stat")),
+            }
+
+    return None
+
+
+def cpu_mem_stats(phase, epoch=None):
+    stats = _cpu_mem_stats_psutil()
+    if stats is None:
+        stats = _cpu_mem_stats_procfs()
+
+    rss, child_rss, child_count = stats
+    rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0
+    return {
+        "rank": rank,
+        "hostname": socket.gethostname(),
+        "pid": os.getpid(),
+        "phase": phase,
+        "epoch": epoch,
+        "rss": rss,
+        "child_rss": child_rss,
+        "total_rss": rss + child_rss,
+        "child_count": child_count,
+        "cgroup": _cgroup_mem_stats(),
+    }
+
+
+def _gib(num_bytes):
+    return num_bytes / (1024**3)
+
+
+def _gib_or_na(num_bytes):
+    if num_bytes is None:
+        return "NA"
+    return f"{_gib(num_bytes):.2f}"
+
+
+def _shorten_cgroup_path(path, max_len=90):
+    if path is None:
+        return "NA"
+    if len(path) <= max_len:
+        return path
+    return "..." + path[-(max_len - 3) :]
+
+
+def _sum_cgroup_stat(unique_cgroups, key):
+    return sum((item.get("stat") or {}).get(key, 0) for item in unique_cgroups)
+
+
+def _sum_cgroup_event(unique_cgroups, key):
+    return sum((item.get("events") or {}).get(key, 0) for item in unique_cgroups)
+
+
+def _unique_cgroups(gathered):
+    cgroups = {}
+    for item in gathered:
+        cgroup = item.get("cgroup")
+        if not cgroup or cgroup.get("current") is None:
+            continue
+        key = (item["hostname"], cgroup.get("version"), cgroup.get("path"))
+        existing = cgroups.setdefault(
+            key,
+            {
+                "hostname": item["hostname"],
+                "version": cgroup.get("version"),
+                "path": cgroup.get("path"),
+                "current": cgroup.get("current"),
+                "peak": cgroup.get("peak"),
+                "max": cgroup.get("max"),
+                "events": cgroup.get("events") or {},
+                "stat": cgroup.get("stat") or {},
+                "ranks": [],
+            },
+        )
+        existing["ranks"].append(item["rank"])
+        # Refresh these values in case later ranks read a slightly newer sample.
+        for field in ("current", "peak", "max", "events", "stat"):
+            existing[field] = cgroup.get(field) or existing[field]
+    return list(cgroups.values())
+
+
+def _log_cgroup_mem(log, phase, epoch, gathered):
+    unique_cgroups = _unique_cgroups(gathered)
+    if not unique_cgroups:
+        log.info(
+            f"[CGROUP-MEM] phase={phase} epoch={epoch if epoch is not None else '-'} "
+            "available=0"
+        )
+        return
+
+    currents = [item["current"] for item in unique_cgroups]
+    peaks = [item["peak"] for item in unique_cgroups if item.get("peak") is not None]
+    maxes = [item["max"] for item in unique_cgroups if item.get("max") is not None]
+    peak_max = f"{_gib(max(peaks)):.2f}" if peaks else "NA"
+    max_min = f"{_gib(min(maxes)):.2f}" if maxes else "NA"
+    details = " ".join(
+        f"cg@{item['hostname']}:{_shorten_cgroup_path(item['path'])}:"
+        f"cur={_gib_or_na(item.get('current'))}GiB,"
+        f"peak={_gib_or_na(item.get('peak'))}GiB,"
+        f"max={_gib_or_na(item.get('max'))}GiB,"
+        f"anon={_gib_or_na((item.get('stat') or {}).get('anon'))}GiB,"
+        f"file={_gib_or_na((item.get('stat') or {}).get('file'))}GiB,"
+        f"slab={_gib_or_na((item.get('stat') or {}).get('slab'))}GiB,"
+        f"oom={(item.get('events') or {}).get('oom', 'NA')},"
+        f"oom_kill={(item.get('events') or {}).get('oom_kill', 'NA')},"
+        f"ranks={','.join(str(rank) for rank in item['ranks'])}"
+        for item in unique_cgroups
+    )
+    log.info(
+        f"[CGROUP-MEM] phase={phase} epoch={epoch if epoch is not None else '-'} "
+        f"cgroups={len(unique_cgroups)} "
+        f"current_sum={_gib(sum(currents)):.2f} "
+        f"current_max={_gib(max(currents)):.2f} "
+        f"peak_max={peak_max} "
+        f"max_min={max_min} "
+        f"anon_sum={_gib(_sum_cgroup_stat(unique_cgroups, 'anon')):.2f} "
+        f"file_sum={_gib(_sum_cgroup_stat(unique_cgroups, 'file')):.2f} "
+        f"slab_sum={_gib(_sum_cgroup_stat(unique_cgroups, 'slab')):.2f} "
+        f"oom_sum={_sum_cgroup_event(unique_cgroups, 'oom')} "
+        f"oom_kill_sum={_sum_cgroup_event(unique_cgroups, 'oom_kill')} | "
+        f"{details}"
+    )
+
+
+def gather_and_print_cpu_mem(log, phase, epoch=None):
+    """Log concise CPU RSS for each rank and its child worker processes."""
+    stats = cpu_mem_stats(phase, epoch)
+    if dist.is_available() and dist.is_initialized():
+        world = dist.get_world_size()
+        gathered = [None for _ in range(world)]
+        dist.all_gather_object(gathered, stats)
+        if dist.get_rank() != 0:
+            return
+    else:
+        gathered = [stats]
+
+    gathered = sorted(gathered, key=lambda item: item["rank"])
+    rss_values = [item["rss"] for item in gathered]
+    child_values = [item["child_rss"] for item in gathered]
+    total_values = [item["total_rss"] for item in gathered]
+    details = " ".join(
+        f"r{item['rank']}@{item['hostname']}:rss={_gib(item['rss']):.2f}GiB,"
+        f"child={_gib(item['child_rss']):.2f}GiB,"
+        f"total={_gib(item['total_rss']):.2f}GiB,"
+        f"children={item['child_count']}"
+        for item in gathered
+    )
+    log.info(
+        f"[CPU-MEM] phase={phase} epoch={epoch if epoch is not None else '-'} "
+        f"rss_gib min={_gib(min(rss_values)):.2f} "
+        f"max={_gib(max(rss_values)):.2f} sum={_gib(sum(rss_values)):.2f} "
+        f"child_sum={_gib(sum(child_values)):.2f} "
+        f"total_sum={_gib(sum(total_values)):.2f} | {details}"
+    )
+    _log_cgroup_mem(log, phase, epoch, gathered)
diff --git a/ScaFFold/worker.py b/ScaFFold/worker.py
index 938206c..b48f5f5 100644
--- a/ScaFFold/worker.py
+++ b/ScaFFold/worker.py
@@ -19,12 +19,16 @@
 from argparse import Namespace
 
 import numpy as np
-import psutil
 import torch
 import torch.distributed as dist
 from distconv import DistConvDDP, ParallelStrategy
 from torch.distributed.tensor import Replicate, Shard
 
+try:
+    import psutil
+except ImportError:
+    psutil = None
+
 from ScaFFold.datagen.get_dataset import get_dataset
 from ScaFFold.unet import UNet
 from ScaFFold.utils.distributed import (
@@ -52,8 +56,12 @@ def check_resource_utilization(log, rank, world_size):
 
     # CPU
     log.debug(f"rank {rank}, world_size {world_size}")
-    log.debug(f"Number of Physical Cores: {psutil.cpu_count(logical=False)}")
-    log.debug(f"Number of Logical Cores: {psutil.cpu_count(logical=True)}")
+    if psutil is not None:
+        log.debug(f"Number of Physical Cores: {psutil.cpu_count(logical=False)}")
+        log.debug(f"Number of Logical Cores: {psutil.cpu_count(logical=True)}")
+    else:
+        log.debug("psutil is unavailable; falling back to os.cpu_count() only")
+        log.debug(f"Number of Logical Cores: {os.cpu_count()}")
     # GPU
     if torch.cuda.is_available() and rank == 0:
         # Get the number of GPUs available
@@ -125,6 +133,13 @@ def main(kwargs_dict: dict = {}):
     )
     log.info(f"rank={rank}, world_size={world_size}")
 
+    total_distconv_shards = math.prod(config.dc_num_shards)
+    if world_size % total_distconv_shards != 0:
+        raise ValueError(
+            f"world_size={world_size} must be divisible by total number of "
+            f"distconv shards = {total_distconv_shards}"
+        )
+
     # Generate or retrieve dataset
     begin_code_region("get_dataset")
     dataset_dir = get_dataset(
@@ -148,8 +163,6 @@ def main(kwargs_dict: dict = {}):
         group_norm_groups=config.group_norm_groups,
     )
     # DDP + DistConv setup
-    # Ensure world_size is divisible by total distconv shards
-    total_distconv_shards = math.prod(config.dc_num_shards)
     if world_size % total_distconv_shards != 0:
         raise ValueError(
             f"world_size={world_size} must be divisible by total number of "