The cost of not validating

The per-batch validation overhead tutorial measures what validating every batch costs. This one measures the other side: what happens to training when corrupt batches reach the model. The overhead benchmark’s “without pandera” mode can’t answer that, because it skips corrupt batches with a precomputed mask, for free. A real pipeline has no mask, so here the unvalidated loop trains on whatever it gets.

What it measures

The benchmark trains a REINFORCE policy (a 4-64-2 MLP, Adam at lr 1e-2) on CartPole-v1 from gymnasium. Each update collects a batch of 8 episodes into a TensorDict (observation, action, reward, episode). Before the update, a share of batches is corrupted in one of two ways:

Each cell runs in two modes with the same seeds and the same corruption schedule:

The cost is environment steps to reach a 475 return, CartPole-v1’s “solved” threshold, averaged over the last 20 episodes. Every collected step counts, including steps thrown away by a rollback and steps in a batch the validated loop skipped. Steps are deterministic for a given seed, so the result doesn’t depend on how busy the node is. Wall-clock time and time spent in validate() are recorded too.

Setting up the environment

Each run is single-threaded and small, so cells run in parallel as separate tasks rather than as threads in one pod. The constants set the target return, the checkpoint interval, and the 1M-step budget after which a run that hasn’t solved CartPole stops and counts at the budget.

image = (
    flyte.Image.from_debian_base(python_version=(3, 12), name="pandera-cost-of-not-validating")
    .with_pip_packages("torch==2.*", index_url="https://download.pytorch.org/whl/cpu")
    .with_pip_packages("pandera[torch]==0.34.0", "gymnasium==1.*", "matplotlib")
    .with_apt_packages("fonts-inter")
)

# Each run is single-threaded and small, so cells run in parallel as separate
# tasks rather than as threads in one pod.
env = flyte.TaskEnvironment(
    name="pandera_cost_of_not_validating",
    image=image,
    resources=flyte.Resources(cpu=(1, 2), memory="2Gi"),
)

OBS_BOUND = 10.0  # CartPole terminates long before any observation gets near this
TARGET_RETURN = 475.0
RETURN_WINDOW = 20
EPISODES_PER_UPDATE = 8
CKPT_EVERY = 10  # updates between checkpoints
MAX_ENV_STEPS = 1_000_000  # a run that hasn't reached the target by here counts as failed
CORRUPTIONS = ("nan_reward", "obs_scale")

The schema

The schema checks dtypes and shapes on every key, observations within ±10 (CartPole terminates long before any observation gets near that), actions in {0, 1}, and rewards in [0, 1], which also rejects NaN:

def build_schema():
    import torch
    import pandera.tensordict as pa

    class Transition(pa.TensorDictModel):
        observation: torch.float32 = pa.Field(shape=(None, 4), ge=-OBS_BOUND, le=OBS_BOUND)
        action: torch.int64 = pa.Field(shape=(None,), isin=[0, 1])
        reward: torch.float32 = pa.Field(shape=(None,), ge=0.0, le=1.0)
        episode: torch.int64 = pa.Field(shape=(None,))

        class Config:
            batch_size = (None,)

    return Transition

Corrupting and validating a batch

Corruption draws from its own random stream, so both modes see the same schedule. In the validated mode, a batch that fails is skipped, and its steps still count toward the cost:

            if corruption is not None and corrupt_rng.random() < rate:
                corrupted += 1
                if corruption == "nan_reward":
                    td["reward"][corrupt_rng.randrange(n)] = float("nan")
                elif corruption == "obs_scale":
                    td["observation"] = td["observation"] * 100.0

            if schema is not None:
                t0 = time.perf_counter()
                try:
                    schema.validate(td, inplace=True)
                    is_valid = True
                except pa.SchemaError:
                    is_valid = False
                validate_s += time.perf_counter() - t0
                if not is_valid:
                    skipped += 1
                    skipped_steps += n
                    continue

Rolling back after a crash

A NaN update makes the next rollout raise. The loop checkpoints every 10 updates and, on a crash, rolls back to the newest checkpoint whose weights and optimizer state are finite. A checkpoint saved right after a NaN update is poisoned, and restoring it would crash again forever, so the loop steps back past it, the way an operator would.

        except (ValueError, RuntimeError):
            # NaN weights: the rollout raised. Roll back to the newest checkpoint with
            # finite weights. A checkpoint saved right after a NaN update is poisoned,
            # and restoring it would crash again forever, so step back past it the way
            # an operator would.
            crashes += 1
            while not _finite(checkpoints[-1]):
                checkpoints.pop()
                poisoned_checkpoints += 1
            model_state, opt_state, mark = checkpoints[-1]
            wasted_steps += env_steps - max(mark, restored_at)
            restored_at = env_steps
            # Load copies. Optimizer.load_state_dict keeps references to the tensors it
            # is given, so loading the stored state directly lets the next NaN update
            # write into the checkpoint itself and poison it for every later rollback.
            policy.load_state_dict(copy.deepcopy(model_state))
            optimizer.load_state_dict(copy.deepcopy(opt_state))

Optimizer.load_state_dict keeps references to the tensors it’s given. If you restore an optimizer from a stored state dict without copying it, the next NaN update writes into the stored checkpoint: its weights stay finite while its optimizer state goes NaN, and every later rollback to it crashes again. The loop loads copies and checks optimizer state, not only weights, for finiteness.

Fanning out the cells

run_cell runs every seed of one (corruption, rate) cell in both modes. main launches one run_cell per cell with asyncio.gather, so the clean cell and each corruption at each rate run in parallel, then summarizes the medians and quartiles and renders the report.

@env.task
async def run_cell(corruption: str | None, rate: float, seeds: int) -> str:
    """Every seed of one (corruption, rate) cell, in both modes."""
    import torch

    torch.set_num_threads(1)
    runs = [train(corruption, rate, validate, seed) for seed in range(seeds) for validate in (False, True)]
    for row in runs:
        print(
            f"{corruption}@{rate:.0%} seed={row['seed']} validate={row['validate']} reached={row['reached']} "
            f"steps={row['env_steps']:,} crashes={row['crashes']} skipped={row['skipped_batches']}",
            flush=True,
        )
    return json.dumps(runs)

@env.task
async def main(rates: list[float] = [0.01, 0.05, 0.2], seeds: int = 30) -> str:
    import os
    import platform

    cells = [(None, 0.0)] + [(c, r) for c in CORRUPTIONS for r in rates]
    outs = await asyncio.gather(*(run_cell(c, r, seeds) for c, r in cells))
    runs = [row for out in outs for row in json.loads(out)]

    import gymnasium
    import pandera
    import tensordict
    import torch

    result = {
        "config": {
            "rates": rates, "seeds": seeds, "target_return": TARGET_RETURN, "return_window": RETURN_WINDOW,
            "episodes_per_update": EPISODES_PER_UPDATE, "ckpt_every": CKPT_EVERY, "max_env_steps": MAX_ENV_STEPS,
            "obs_bound": OBS_BOUND,
        },
        "environment": {
            "torch": torch.__version__, "pandera": pandera.__version__, "tensordict": tensordict.__version__,
            "gymnasium": gymnasium.__version__, "python": platform.python_version(),
            "machine": platform.machine(), "cpu_count": os.cpu_count(),
        },
        "summary": summarize(runs),
        "runs": runs,
    }
    out = json.dumps(result, indent=2)
    await plot(out)
    return out

Corruption rates are per batch: 1%, 5%, and 20% of updates. A per-row rate would be much harsher, since a 0.1% bad-row rate puts at least one bad row in almost every 8-episode batch. Each cell runs 30 seeds, reported as the median and interquartile range, because REINFORCE on CartPole is noisy: clean runs alone range from 39K to 114K steps.

Results

The run used 30 seeds per cell, 420 runs in all, on x86_64 with torch 2.14.1+cpu, pandera 0.34.0, tensordict 0.14.3, and gymnasium 1.4.0.

Median environment steps to reach a 475 return, with the interquartile range in parentheses. A clean run took 69.9K.

How to read it:

• Validation cost 0.5% of wall-clock time in every cell (0.47–0.51%, validate() time over total run time). Here most of the time goes to collecting episodes, so this isn’t directly comparable to the overhead benchmark, which times the update step alone.
• The silent corruption is the expensive one. Scaled observations never crash anything, so nothing tells the loop to stop. At 5% of batches the unvalidated runs needed twice the compute; at 20%, half of them hadn’t solved the task by 1M steps, and the median of ≥744K is a lower bound. The validated runs solved it every time, at 1.0–1.2x a clean run’s compute.
• The crash is cheaper, because it’s loud. A NaN reward kills the run within one update, so the loop loses at most a checkpoint interval. At 5% that was 2 crashes and 15K wasted steps per run (median), and the total cost was close to the validated loop’s; at 20% it was 25 crashes, 89K wasted steps, and 2.1x the compute. 103 checkpoints across all runs were saved right after a NaN update and had to be skipped on rollback.
• At 1% there’s no measurable difference, which is expected: at about 45 updates per run, a 1% rate corrupts fewer than one batch per run on average.
• The validated loop pays for skipped batches, since their steps count too. At 20% that’s why its median rises to 1.2x: a fifth of the collected data is thrown away rather than trained on.

Caveats

• The cost is conditional on how often your data goes bad. The benchmark can’t know that rate for any real pipeline, so it reports the cost at a few rates rather than an expected value.
• Crash cost depends on checkpoint spacing. Rolling back 10 updates is cheap; rolling back an hour of GPU time isn’t. The nan_reward numbers scale with the checkpoint interval.
• CPU only. On a GPU, an out-of-range action becomes a delayed device-side assert that points at the wrong line, and float64 inputs double memory. Neither shows up here, and debugging time isn’t measured at all.
• One task and one algorithm. CartPole with REINFORCE is small and recovers quickly; a larger model can take longer to recover from a bad update, or not recover.

Run it

# From v2/tutorials/pandera_cost_of_not_validating in your unionai-examples clone.
# On the cluster in your Flyte config; builds the image remotely, then writes results/ locally
uv run main.py

# Locally, with a short configuration
uv run main.py --local --seeds 3 --rates 0.05