From 81036ddb6e5eee9e1edaafe7f16bd0abad81cc86 Mon Sep 17 00:00:00 2001
From: "mintlify[bot]" <109931778+mintlify[bot]@users.noreply.github.com>
Date: Tue, 28 Jul 2026 22:39:53 +0000
Subject: [PATCH 1/2] docs: document transform_parallelism option for
StreamingDataset
---
docs/training/index.mdx | 23 ++++++++++++++++++++++-
1 file changed, 22 insertions(+), 1 deletion(-)
diff --git a/docs/training/index.mdx b/docs/training/index.mdx
index 85078a0..1948976 100644
--- a/docs/training/index.mdx
+++ b/docs/training/index.mdx
@@ -115,7 +115,10 @@ ds = StreamingDataset(
Many model training workloads require a transformation step between loading the data and training the model. For
example, we may need to decode images, tokenize text, or normalize data. A transformation function can be provided
using the `transform` parameter. Transformations can be expensive, so `StreamingDataset` applies them with a
-`ThreadPoolExecutor` whose worker count equals the number of available CPUs.
+`ThreadPoolExecutor`. By default the worker count equals the number of available CPUs (falling back to one worker
+when the CPU count cannot be determined). Use the `transform_parallelism` parameter to override this limit — for
+example, to cap concurrency on shared machines or to raise it when transforms release the GIL and CPUs are plentiful.
+`transform_parallelism` must be a positive integer.
Transformations are applied to batches, not individual samples, to amortize per-batch overhead. A transformation
function receives a PyArrow `RecordBatch` and must return an iterable with exactly one output sample for every input
@@ -140,6 +143,24 @@ ds = StreamingDataset(table, shuffle_seed=42, transform=normalize)
```
+To pin the transform pool to a specific size, pass `transform_parallelism`:
+
+
+```py Python icon=Python
+ds = StreamingDataset(
+ table,
+ shuffle_seed=42,
+ transform=normalize,
+ transform_parallelism=4,
+)
+```
+
+
+
+`transform_parallelism` also bounds how many raw batches are held in flight for transformation, so a smaller value
+reduces peak memory use at the cost of throughput.
+
+
#### DataLoader workers
The thread-based transformation model that `StreamingDataset` uses by default is only effective when the transform
From 4f1da0aaaf1a688552d6b6130538099b628af35f Mon Sep 17 00:00:00 2001
From: "mintlify[bot]" <109931778+mintlify[bot]@users.noreply.github.com>
Date: Mon, 10 Aug 2026 16:24:27 +0000
Subject: [PATCH 2/2] docs: document on_transform_error and merge_state_dicts
for StreamingDataset
---
docs/training/index.mdx | 88 +++++++++++++++++++++++++++++++++++++++++
1 file changed, 88 insertions(+)
diff --git a/docs/training/index.mdx b/docs/training/index.mdx
index 1948976..85dc735 100644
--- a/docs/training/index.mdx
+++ b/docs/training/index.mdx
@@ -161,6 +161,57 @@ ds = StreamingDataset(
reduces peak memory use at the cost of throughput.
+#### Handling bad rows
+
+By default, any exception raised by the transform aborts iteration. When your data may contain the occasional bad
+row — for example, nulls or NaNs from incomplete surveys, or corrupt image bytes — use the `on_transform_error`
+parameter to keep training going. It accepts:
+
+- `"raise"` (the default): the exception propagates and iteration stops.
+- `"skip"`: the failing rows are dropped and iteration continues.
+- `"warn"`: like `"skip"`, plus a warning is logged for each failing batch.
+- A callable `handler(exc) -> bool`: return `True` to skip the failing rows, `False` to re-raise. Use this to skip
+ only specific error types (for example, `PIL.UnidentifiedImageError`) while letting unexpected errors surface.
+
+When a batch fails, the transform is re-run on each single-row slice so that only the rows that actually raised are
+dropped. Your transform must therefore be deterministic and must accept batches of any size, including one row.
+Read the running skip count from the `rows_skipped` property.
+
+
+```py Python icon=Python
+from PIL import UnidentifiedImageError
+
+def skip_bad_images(exc: Exception) -> bool:
+ return isinstance(exc, UnidentifiedImageError)
+
+ds = StreamingDataset(
+ table,
+ shuffle_seed=42,
+ transform=decode_images,
+ on_transform_error=skip_bad_images, # or "skip" / "warn"
+)
+
+for sample in ds:
+ train_step(sample)
+
+print(f"dropped {ds.rows_skipped} bad rows")
+```
+
+
+
+When bad rows are unevenly distributed across splits, each rank's iterator ends at the last cycle where every split
+it owns still has a row. That means one rank's iterator can yield noticeably fewer steps than another rank's in the
+same run. This is safe for asynchronous or single-rank training, but synchronous distributed training (for example,
+ranks that call `all_reduce` every step) can deadlock unless you broadcast a stop signal on `StopIteration`. The
+per-split sample sequence stays deterministic; only the epoch's tail can differ across topologies.
+
+
+
+Prefer the `filter` parameter when bad rows can be expressed as a SQL predicate (for example,
+`"image_bytes IS NOT NULL"`). Filtering happens before splits are built, so every determinism and resumability
+guarantee is fully preserved. Use `on_transform_error` for failures that can't be predicted from a predicate.
+
+
#### DataLoader workers
The thread-based transformation model that `StreamingDataset` uses by default is only effective when the transform
@@ -430,6 +481,43 @@ exact resume, also use the same table snapshot, `epoch`, `shuffle`, `filter`, an
`world_size` and number of DataLoader workers may change as long as the split and batch-size divisibility constraints
still hold.
+#### Elastic resume when rows are skipped
+
+When `on_transform_error` drops rows, each rank knows the exact permutation position only for the splits it owns.
+Save one `state_dict` per rank at a checkpoint boundary, then combine them with the `merge_state_dicts` static method
+before resuming. The merge is topology-agnostic: hand the same merged dict to every rank of the next run, regardless
+of whether the world size grew, shrank, or stayed the same.
+
+
+```py Python icon=Python
+# --- checkpointing on the old topology ---
+states = all_gather_object(dataset.state_dict()) # collect from every rank
+if rank == 0:
+ torch.save({"dataset_states": states, "model": model.state_dict()}, "ckpt.pt")
+
+# --- resuming on a new topology ---
+checkpoint = torch.load("ckpt.pt")
+merged = StreamingDataset.merge_state_dicts(checkpoint["dataset_states"])
+
+dataset = StreamingDataset(
+ table,
+ num_splits=merged["num_splits"],
+ shuffle_seed=merged["shuffle_seed"],
+ epoch=merged["epoch"],
+ rank=rank,
+ world_size=world_size, # may differ from the saved run
+ transform=decode_images,
+ on_transform_error="skip",
+)
+dataset.load_state_dict(merged) # every rank loads the same merged dict
+```
+
+
+`merge_state_dicts` raises `ValueError` if the states came from different runs (mismatched `shuffle_seed`,
+`num_splits`, `epoch`, or `samples_consumed_per_split`), so every rank must call `state_dict()` at the same
+global-step boundary. When no rows are ever skipped, merging is a no-op and you can keep using a single per-run
+`state_dict()` as before.
+
## Permutations
In more complicated scenarios, you may want the flexibility to shuffle, split, and select data without using the full