From d2b102cafb8f6165a727c87a818fc802432669c1 Mon Sep 17 00:00:00 2001
From: Zachery Aaron Shores-Chmielewski <zacheryasc@gmail.com>
Date: Wed, 29 Jul 2026 13:42:52 +0400
Subject: [PATCH] fix: use pre-processed dataset

---
 .skyignore                                    |   1 +
 configs/remote_minimal_scaling_node0.toml     | 119 +++++
 src/airfrans_frontier/cli.py                  |  37 +-
 src/airfrans_frontier/models/__init__.py      |   3 +-
 src/airfrans_frontier/models/mlp.py           |  46 ++
 src/airfrans_frontier/sweep.py                | 336 +++++++++++-
 src/airfrans_frontier/training/artifacts.py   |   2 +-
 src/airfrans_frontier/training/config.py      |  26 +-
 src/airfrans_frontier/training/data.py        |  15 +-
 src/airfrans_frontier/training/loop.py        | 165 +++++-
 .../training/streaming_data.py                | 505 ++++++++++++++++--
 tests/test_data_sources.py                    |   1 +
 tests/test_mlp.py                             |  19 +
 tests/test_streaming_data.py                  | 163 ++++++
 tests/test_sweep.py                           |  36 ++
 tests/test_training_config.py                 | 123 +++++
 tests/test_training_data.py                   |  21 +-
 tests/test_training_loop.py                   |  67 +++
 18 files changed, 1571 insertions(+), 114 deletions(-)
 create mode 100644 configs/remote_minimal_scaling_node0.toml

diff --git a/.skyignore b/.skyignore
index b5b3709..514486c 100644
--- a/.skyignore
+++ b/.skyignore
@@ -1,5 +1,6 @@
 /artifacts/remote_runs
 /artifacts/runs
+/artifacts/runs_minimal_scaling_20260728
 /artifacts/data_cache
 /artifacts/model_sanity
 /artifacts/current_run
diff --git a/configs/remote_minimal_scaling_node0.toml b/configs/remote_minimal_scaling_node0.toml
new file mode 100644
index 0000000..3977c89
--- /dev/null
+++ b/configs/remote_minimal_scaling_node0.toml
@@ -0,0 +1,119 @@
+[run]
+name = "minimal_scaling_sweep_node0"
+timeout_minutes = 1440
+local_artifact_dir = "artifacts/remote_runs"
+max_attempts = 1
+artifact_sync_interval_seconds = 60
+
+[provider]
+kind = "vastai"
+disk_gb = 192
+max_price_per_hour = 0.80
+image = "vastai/base:0.0.2"
+
+[provider.gpu]
+name = "RTX 4090"
+count = 1
+min_vram_gb = 20
+
+[selection]
+min_reliability = 0.95
+min_down_mbps = 100
+min_up_mbps = 25
+require_verified = true
+blocked_geos = ["CN"]
+blacklist_hosts = [59017, 1647, 92578, 1276, 75481, 1256, 85323, 34031, 528237]
+drop_cheap_frac = 0.30
+image_size_gb = 5.0
+base_url = "https://cloud.vast.ai"
+
+[workspace]
+workdir = "."
+exclude = [
+  "/artifacts/remote_runs",
+  "/artifacts/runs",
+  "/artifacts/runs_minimal_scaling_20260728",
+  "/artifacts/data_cache",
+  "/artifacts/model_sanity",
+  "/artifacts/current_run",
+  "/artifacts/public_airfrans",
+  "/artifacts/notes",
+  "/artifacts/preflight_20260725",
+  "/artifacts/remote_state_20260725",
+  "/artifacts/sweep_20260725",
+  "/artifacts/model_sanity_decisions",
+  "/artifacts/model_sanity_verify",
+  "/artifacts/*.json",
+  "/artifacts/*.yaml",
+  "/artifacts/*.md",
+  "/data/raw",
+  "/data/processed",
+  "/.venv",
+  "/.git",
+  "/.airfrans_hf_resume",
+  "/sky_logs",
+  "/notebooks",
+  "__pycache__",
+  "*.pyc",
+]
+
+[bootstrap]
+command = """
+uv sync --no-dev
+uv run --no-dev python -c "import torch; ok=torch.cuda.is_available(); print('torch_cuda_available=' + str(ok)); print(torch.cuda.get_device_name(0) if ok else 'no_cuda_device'); assert ok, 'torch CUDA unavailable'"
+uv run --no-dev python -c "import huggingface_hub, wandb; print('hf_wandb_import_ok')"
+"""
+
+[data]
+validation_command = """
+uv run --no-dev python - <<'PY'
+from airfrans_frontier.sweep import read_jobs
+
+jobs = read_jobs('artifacts/minimal_scaling_sweep_20260728/jobs.jsonl')
+pending = sum(1 for job in jobs if job['status'] == 'pending')
+phases = [job.get('phase') for job in jobs]
+transitions = []
+for phase in phases:
+    if not transitions or transitions[-1] != phase:
+        transitions.append(phase)
+assert len(jobs) == 23, len(jobs)
+assert pending > 0
+assert transitions == ['A_measurement_canary', 'B_family_viability', 'C_scaling_ladder'], transitions
+print('minimal_scaling_pending_jobs=' + str(pending))
+print('minimal_scaling_phase_order=' + ','.join(transitions))
+PY
+"""
+
+[job]
+command = """
+uv run --no-dev python scripts/aggressive_oom_node_wrapper.py --jobs artifacts/minimal_scaling_sweep_20260728/jobs.jsonl --node node-0 --artifact-dir artifacts/current_run --utilization artifacts/minimal_scaling_sweep_20260728/utilization_node0.jsonl --sample-interval-seconds 30 --stale-after-seconds 21600 --max-attempts 2 --require-success
+"""
+artifact_dir = "artifacts/current_run"
+heartbeat_file = "artifacts/current_run/heartbeat.json"
+metrics_file = "artifacts/current_run/metrics.jsonl"
+
+[artifacts]
+mode = "object_store_upload"
+required = [
+  "config.toml",
+  "job_manifest.json",
+  "metrics.jsonl",
+  "latest_metrics.json",
+  "heartbeat.json",
+  "checkpoint_latest.pt",
+  "checkpoint_best.pt",
+  "checkpoint_final.pt",
+  "final_metrics.json",
+  "run_manifest.json",
+  "environment_manifest.json",
+  "utilization.jsonl",
+  "node_summary.json",
+  "sweep_collect.json",
+  "sweep_jobs.jsonl",
+  "sweep_attempts.jsonl",
+  "sweep_utilization.jsonl",
+]
+
+[cleanup]
+on_success = "sky_down"
+on_failure = "collect_then_keep"
diff --git a/src/airfrans_frontier/cli.py b/src/airfrans_frontier/cli.py
index 20bad88..2030828 100644
--- a/src/airfrans_frontier/cli.py
+++ b/src/airfrans_frontier/cli.py
@@ -73,7 +73,9 @@ def build_parser() -> argparse.ArgumentParser:
     sweep_generate.add_argument("--data-root", default="artifacts/data_cache/airfrans_processed/processed/full")
     sweep_generate.add_argument("--artifact-dir", default="artifacts/runs")
     sweep_generate.add_argument("--hf-repo-id", default="zacheryasc/airfrans-frontier-checkpoints")
-    sweep_generate.add_argument("--group", default="aggressive_oom_sweep_01")
+    sweep_generate.add_argument("--group", default="minimal_scaling_sweep_20260728")
+    sweep_generate.add_argument("--plan", choices=("minimal-scaling", "legacy-aggressive"), default="minimal-scaling")
+    sweep_generate.add_argument("--include-proxies", action="store_true")
     sweep_generate.add_argument("--bands", nargs="*", default=["100m", "700m"])
     sweep_generate.add_argument("--families", nargs="*")
     sweep_generate.add_argument("--encodings", nargs="*")
@@ -250,21 +252,30 @@ def main(argv: list[str] | None = None) -> int:
         print(f"report: {resolve_path(args.artifact_dir) / 'model_sanity_results.json'}")
         print(f"families: {len(result['families'])}")
         return 0
-
     if args.command == "sweep-generate":
-        from airfrans_frontier.sweep import DEFAULT_ENCODINGS, DEFAULT_FAMILIES, generate_jobs
+        from airfrans_frontier.sweep import DEFAULT_ENCODINGS, DEFAULT_FAMILIES, generate_jobs, generate_minimal_scaling_jobs
 
         try:
-            jobs = generate_jobs(
-                output_dir=resolve_path(args.output_dir),
-                data_root=args.data_root,
-                artifact_dir=args.artifact_dir,
-                hf_repo_id=args.hf_repo_id,
-                group=args.group,
-                bands=tuple(args.bands),
-                families=tuple(args.families) if args.families else DEFAULT_FAMILIES,
-                encodings=tuple(args.encodings) if args.encodings else DEFAULT_ENCODINGS,
-            )
+            if args.plan == "minimal-scaling":
+                jobs = generate_minimal_scaling_jobs(
+                    output_dir=resolve_path(args.output_dir),
+                    data_root=args.data_root,
+                    artifact_dir=args.artifact_dir,
+                    hf_repo_id=args.hf_repo_id,
+                    group=args.group,
+                    include_proxies=args.include_proxies,
+                )
+            else:
+                jobs = generate_jobs(
+                    output_dir=resolve_path(args.output_dir),
+                    data_root=args.data_root,
+                    artifact_dir=args.artifact_dir,
+                    hf_repo_id=args.hf_repo_id,
+                    group=args.group,
+                    bands=tuple(args.bands),
+                    families=tuple(args.families) if args.families else DEFAULT_FAMILIES,
+                    encodings=tuple(args.encodings) if args.encodings else DEFAULT_ENCODINGS,
+                )
         except (FileNotFoundError, NotADirectoryError, ValueError, RuntimeError) as exc:
             print(f"error: {exc}", file=sys.stderr)
             return 1
diff --git a/src/airfrans_frontier/models/__init__.py b/src/airfrans_frontier/models/__init__.py
index b9eba70..e6b0319 100644
--- a/src/airfrans_frontier/models/__init__.py
+++ b/src/airfrans_frontier/models/__init__.py
@@ -10,12 +10,13 @@ from airfrans_frontier.models.frontier import (
     RasterFNOUNet,
     SirenConditionedINR,
 )
-from airfrans_frontier.models.mlp import PointwiseMLP
+from airfrans_frontier.models.mlp import EncodedPointwiseMLP, PointwiseMLP
 
 __all__ = [
     "CoordinateEncoder",
     "CoordinateEncodingSpec",
     "DeepONetBranchTrunk",
+    "EncodedPointwiseMLP",
     "FourierFiLMMLP",
     "LocalPointTransformer",
     "NeRFCFDMultiRes",
diff --git a/src/airfrans_frontier/models/mlp.py b/src/airfrans_frontier/models/mlp.py
index d7b7f49..9163d32 100644
--- a/src/airfrans_frontier/models/mlp.py
+++ b/src/airfrans_frontier/models/mlp.py
@@ -3,6 +3,52 @@ from __future__ import annotations
 import torch
 from torch import nn
 
+from airfrans_frontier.models.coordinate_encoding import CoordinateEncoder, CoordinateEncodingSpec
+
+
+class EncodedPointwiseMLP(nn.Module):
+    def __init__(
+        self,
+        *,
+        feature_names: tuple[str, ...],
+        output_dim: int,
+        coordinate_features: tuple[str, ...],
+        coordinate_encoding_spec: CoordinateEncodingSpec,
+        hidden_width: int = 128,
+        depth: int = 4,
+        activation: str = "gelu",
+    ) -> None:
+        super().__init__()
+        if output_dim <= 0:
+            raise ValueError("output_dim must be positive")
+        if hidden_width <= 0:
+            raise ValueError("hidden_width must be positive")
+        if depth <= 0:
+            raise ValueError("depth must be positive")
+        name_to_index = {name: index for index, name in enumerate(feature_names)}
+        missing = [name for name in coordinate_features if name not in name_to_index]
+        if missing:
+            raise ValueError(f"coordinate features are missing from input schema: {missing}")
+        coord_indices = [name_to_index[name] for name in coordinate_features]
+        condition_indices = [index for index, name in enumerate(feature_names) if name not in set(coordinate_features)]
+        self.register_buffer("coord_indices", torch.tensor(coord_indices, dtype=torch.long), persistent=False)
+        self.register_buffer("condition_indices", torch.tensor(condition_indices, dtype=torch.long), persistent=False)
+        self.encoder = CoordinateEncoder(coordinate_encoding_spec, coordinate_dim=len(coord_indices))
+        input_dim = self.encoder.output_dim + len(condition_indices)
+        layers: list[nn.Module] = [nn.Linear(input_dim, hidden_width), _activation(activation)]
+        for _ in range(depth - 1):
+            layers.extend((nn.Linear(hidden_width, hidden_width), _activation(activation)))
+        layers.append(nn.Linear(hidden_width, output_dim))
+        self.network = nn.Sequential(*layers)
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        coordinates = features.index_select(dim=1, index=self.coord_indices)
+        encoded = self.encoder(coordinates)
+        if self.condition_indices.numel():
+            conditions = features.index_select(dim=1, index=self.condition_indices)
+            encoded = torch.cat((encoded, conditions), dim=-1)
+        return self.network(encoded)
+
 
 class PointwiseMLP(nn.Module):
     def __init__(
diff --git a/src/airfrans_frontier/sweep.py b/src/airfrans_frontier/sweep.py
index 9657c13..806fbf9 100644
--- a/src/airfrans_frontier/sweep.py
+++ b/src/airfrans_frontier/sweep.py
@@ -30,28 +30,37 @@ SUCCESS_ARTIFACTS = REQUIRED_JOB_ARTIFACTS + ("final_metrics.json",)
 FAILURE_ARTIFACTS = REQUIRED_JOB_ARTIFACTS + ("failure_report.json",)
 
 OOM_BANDS: dict[str, dict[str, int]] = {
-    "10m": {"hidden_width": 1024, "depth": 8, "steps": 300, "batch_size": 4096},
-    "100m": {"hidden_width": 2048, "depth": 16, "steps": 1500, "batch_size": 4096},
-    "700m": {"hidden_width": 4096, "depth": 40, "steps": 500, "batch_size": 2048},
+    "10m": {"hidden_width": 1024, "depth": 8, "steps": 5000, "batch_size": 4096},
+    "100m": {"hidden_width": 2048, "depth": 16, "steps": 2500, "batch_size": 4096},
+    "700m": {"hidden_width": 4096, "depth": 40, "steps": 1000, "batch_size": 2048},
 }
 FILM_DEPTH_OVERRIDES = {"700m": 32}
-DEFAULT_FAMILIES = (
-    "mlp",
+MAIN_SCALING_FAMILIES = (
+    "mlp_encoded_baseline",
     "nerf_cfd_multires",
     "deeponet_branch_trunk",
     "film_fourier_inr",
     "siren_conditioned_inr",
+)
+PROXY_SCALING_FAMILIES = (
     "raster_fno_unet",
     "point_context_perceiver",
     "meshgraphnet_or_point_transformer_local",
 )
-DEFAULT_ENCODINGS = (
+DEFAULT_FAMILIES = MAIN_SCALING_FAMILIES
+SUPPORTED_ENCODINGS = (
     "raw",
     "fixed_fourier",
     "nerf_multires",
     "random_fourier",
 )
+DEFAULT_ENCODINGS = (
+    "fixed_fourier",
+    "nerf_multires",
+    "random_fourier",
+)
 ENCODING_COMPATIBLE_FAMILIES = {
+    "mlp_encoded_baseline",
     "nerf_cfd_multires",
     "deeponet_branch_trunk",
     "film_fourier_mlp",
@@ -59,8 +68,8 @@ ENCODING_COMPATIBLE_FAMILIES = {
 }
 
 COORDINATE_ENCODING_EXCLUSION_REASONS = {
-    "mlp": "architecture_uses_full_raw_point_features",
-    "siren_conditioned_inr": "siren_uses_raw_coordinates_with_omega0",
+    "mlp": "legacy pointwise MLP uses full raw point features; use mlp_encoded_baseline for serious sweeps",
+    "siren_conditioned_inr": "siren_uses_sine_coordinate_basis_with_omega0",
     "raster_fno_unet": "raster_proxy_uses_xy_sampling_grid",
     "point_context_perceiver": "perceiver_proxy_uses_full_feature_projection",
     "meshgraphnet_or_point_transformer_local": "local_point_proxy_uses_knn_coordinates",
@@ -143,7 +152,7 @@ def write_default_pool(path: str | Path) -> Path:
     return target
 
 def coordinate_encoding_decision(family: str, encoding: str) -> dict[str, Any]:
-    if encoding not in DEFAULT_ENCODINGS:
+    if encoding not in SUPPORTED_ENCODINGS:
         raise ValueError(f"unsupported coordinate encoding: {encoding}")
     if family in ENCODING_COMPATIBLE_FAMILIES:
         return {
@@ -240,6 +249,235 @@ def generate_jobs(
     return jobs
 
 
+MINIMAL_SWEEP_RECIPES: dict[str, dict[str, Any]] = {
+    "mlp_encoded_baseline": {"encoding": "nerf_multires", "lr": 0.0001},
+    "nerf_cfd_multires": {"encoding": "nerf_multires", "lr": 0.0001},
+    "deeponet_branch_trunk": {"encoding": "nerf_multires", "lr": 0.0001},
+    "film_fourier_inr": {"encoding": "nerf_multires", "lr": 0.0001},
+    "siren_conditioned_inr": {"encoding": "raw", "lr": 0.00003},
+}
+
+
+def generate_minimal_scaling_jobs(
+    *,
+    output_dir: str | Path,
+    data_root: str = "artifacts/data_cache/airfrans_processed/processed/full",
+    artifact_dir: str = "artifacts/runs_minimal_scaling",
+    hf_repo_id: str = "zacheryasc/airfrans-frontier-checkpoints",
+    group: str = "minimal_scaling_sweep_20260728",
+    include_proxies: bool = False,
+    seed: int = 20260728,
+) -> list[dict[str, Any]]:
+    root = Path(output_dir)
+    config_dir = root / "configs"
+    config_dir.mkdir(parents=True, exist_ok=True)
+    jobs: list[dict[str, Any]] = []
+
+    def add_job(
+        *,
+        job_id: str,
+        phase: str,
+        purpose: str,
+        family: str,
+        band: str,
+        encoding: str,
+        lr: float,
+        steps: int,
+        batch_size: int,
+        train_cases: int,
+        val_cases: int,
+        test_cases: int,
+        log_interval: int,
+        checkpoint_policy: str,
+        include_optimizer_state: bool,
+        include_rng_state: bool,
+        dead_curve_patience_evals: int | None,
+        dead_curve_warmup_steps: int,
+        expected_terminal_status: str = "succeeded",
+        run_seed: int | None = None,
+    ) -> None:
+        decision = coordinate_encoding_decision(family, encoding)
+        metadata = model_metadata_for_family(family, coordinate_encoding_compatibility=str(decision["compatibility"]))
+        config_path = config_dir / f"{job_id}.toml"
+        config_path.write_text(
+            training_config_text(
+                run_name=job_id,
+                seed=seed + len(jobs) if run_seed is None else run_seed,
+                data_root=data_root,
+                artifact_dir=artifact_dir,
+                model_family=family,
+                band=band,
+                coordinate_encoding=encoding,
+                hf_repo_id=hf_repo_id,
+                group=group,
+                train_cases=train_cases,
+                val_cases=val_cases,
+                test_cases=test_cases,
+                streaming_normalization_cases=min(4, train_cases),
+                data_source="huggingface_streaming",
+                all_points_per_case=True,
+                points_per_case=None,
+                steps=steps,
+                batch_size=batch_size,
+                lr=lr,
+                log_interval=log_interval,
+                checkpoint_policy=checkpoint_policy,
+                checkpoint_include_optimizer_state=include_optimizer_state,
+                checkpoint_include_rng_state=include_rng_state,
+                dead_curve_patience_evals=dead_curve_patience_evals,
+                dead_curve_warmup_steps=dead_curve_warmup_steps,
+            )
+        )
+        jobs.append(
+            {
+                "job_id": job_id,
+                "status": "pending",
+                "phase": phase,
+                "purpose": purpose,
+                "band": band,
+                "model_family": family,
+                "requested_family": metadata["requested_family"],
+                "implementation_family": metadata["implementation_family"],
+                "reported_family": metadata["reported_family"],
+                "is_proxy": metadata["is_proxy"],
+                "proxy_for": metadata["proxy_for"],
+                "proxy_notes": metadata["proxy_notes"],
+                "coordinate_encoding": encoding,
+                "coordinate_encoding_decision": decision,
+                "coordinate_encoding_compatibility": decision["compatibility"],
+                "all_points_per_case": True,
+                "expected_terminal_status": expected_terminal_status,
+                "checkpoint_policy": checkpoint_policy,
+                "include_optimizer_state": include_optimizer_state,
+                "dead_curve_patience_evals": dead_curve_patience_evals,
+                "min_post_warmup_eval_points": 8,
+                "config_path": str(config_path),
+                "attempts": 0,
+            }
+        )
+
+    add_job(
+        job_id="A_data_all_points_smoke",
+        phase="A_measurement_canary",
+        purpose="all_points_data_path_smoke",
+        family="mlp_encoded_baseline",
+        band="10m",
+        encoding="fixed_fourier",
+        lr=0.0001,
+        steps=50,
+        batch_size=512,
+        train_cases=2,
+        val_cases=1,
+        test_cases=1,
+        log_interval=5,
+        checkpoint_policy="full",
+        include_optimizer_state=True,
+        include_rng_state=True,
+        dead_curve_patience_evals=None,
+        dead_curve_warmup_steps=0,
+    )
+    add_job(
+        job_id="A_checkpoint_700m_footprint_smoke",
+        phase="A_measurement_canary",
+        purpose="full_700m_checkpoint_smoke",
+        family="film_fourier_inr",
+        band="700m",
+        encoding="nerf_multires",
+        lr=0.0001,
+        steps=2,
+        batch_size=64,
+        train_cases=2,
+        val_cases=1,
+        test_cases=1,
+        log_interval=1,
+        checkpoint_policy="full",
+        include_optimizer_state=True,
+        include_rng_state=True,
+        dead_curve_patience_evals=None,
+        dead_curve_warmup_steps=0,
+    )
+    add_job(
+        job_id="A_plateau_interrupt_smoke",
+        phase="A_measurement_canary",
+        purpose="dead_curve_interrupt_smoke",
+        family="mlp_encoded_baseline",
+        band="10m",
+        encoding="fixed_fourier",
+        lr=1e-12,
+        steps=40,
+        batch_size=512,
+        train_cases=2,
+        val_cases=1,
+        test_cases=1,
+        log_interval=5,
+        checkpoint_policy="full",
+        include_optimizer_state=True,
+        include_rng_state=True,
+        dead_curve_patience_evals=3,
+        dead_curve_warmup_steps=5,
+        expected_terminal_status="failed",
+    )
+
+    families = tuple(MINIMAL_SWEEP_RECIPES)
+    if include_proxies:
+        families = families + PROXY_SCALING_FAMILIES
+    family_recipes = [
+        (family, MINIMAL_SWEEP_RECIPES.get(family, {"encoding": "fixed_fourier", "lr": 0.0001}))
+        for family in families
+    ]
+    for family, recipe in family_recipes:
+        add_job(
+            job_id=f"B_10m_{family}_{recipe['encoding']}_viability",
+            phase="B_family_viability",
+            purpose="viability",
+            family=family,
+            band="10m",
+            encoding=str(recipe["encoding"]),
+            lr=float(recipe["lr"]),
+            steps=300,
+            batch_size=2048,
+            train_cases=900,
+            val_cases=50,
+            test_cases=50,
+            log_interval=25,
+            checkpoint_policy="full",
+            include_optimizer_state=True,
+            include_rng_state=True,
+            dead_curve_patience_evals=8,
+            dead_curve_warmup_steps=50,
+        )
+    for family_index, (family, recipe) in enumerate(family_recipes):
+        for band in ("10m", "100m", "700m"):
+            shape = OOM_BANDS[band]
+            add_job(
+                job_id=f"C_{band}_{family}_{recipe['encoding']}_scale",
+                phase="C_scaling_ladder",
+                purpose="scaling_curve",
+                family=family,
+                band=band,
+                encoding=str(recipe["encoding"]),
+                lr=float(recipe["lr"]),
+                steps=shape["steps"],
+                batch_size=shape["batch_size"],
+                train_cases=900,
+                val_cases=50,
+                test_cases=50,
+                log_interval=50 if band == "10m" else 25,
+                checkpoint_policy="full",
+                include_optimizer_state=True,
+                include_rng_state=True,
+                dead_curve_patience_evals=8,
+                dead_curve_warmup_steps=100 if band == "10m" else 50,
+                run_seed=seed + 1000 + family_index,
+            )
+
+    write_jobs(root / "jobs.jsonl", jobs)
+    write_default_pool(root / "pool.toml")
+    (root / "budget_ledger.json").write_text(json.dumps({"budget_usd": 25.0, "spent_usd": 0.0, "reserved_usd": 0.0}, indent=2) + "\n")
+    write_rescue_templates(root / RESCUE_TEMPLATES_JSON)
+    return jobs
+
+
 def training_config_text(
     *,
     run_name: str,
@@ -251,11 +489,63 @@ def training_config_text(
     coordinate_encoding: str,
     hf_repo_id: str,
     group: str,
+    train_cases: int = 900,
+    val_cases: int = 50,
+    test_cases: int = 50,
+    streaming_normalization_cases: int | None = None,
+    data_source: str = "huggingface_streaming",
+    all_points_per_case: bool = False,
+    points_per_case: int | None = 8192,
+    steps: int | None = None,
+    batch_size: int | None = None,
+    lr: float = 0.0001,
+    log_interval: int = 25,
+    checkpoint_policy: str = "full",
+    checkpoint_include_optimizer_state: bool = True,
+    checkpoint_include_rng_state: bool = True,
+    dead_curve_patience_evals: int | None = None,
+    dead_curve_min_relative_improvement: float = 0.005,
+    dead_curve_warmup_steps: int = 0,
 ) -> str:
     shape = dict(OOM_BANDS[band])
     if model_family == "film_fourier_inr" and band in FILM_DEPTH_OVERRIDES:
         shape["depth"] = FILM_DEPTH_OVERRIDES[band]
     condition_width = max(512, shape["hidden_width"] // 2)
+    if steps is not None:
+        shape["steps"] = int(steps)
+    if batch_size is not None:
+        shape["batch_size"] = int(batch_size)
+    data_source_lines = {
+        "huggingface": (
+            "source = \"huggingface\"",
+            "hf_repo_id = \"zacheryasc/airfrans-processed\"",
+            "hf_repo_type = \"dataset\"",
+            "hf_path_prefix = \"processed/full\"",
+            "cache_dir = \"artifacts/data_cache/airfrans_processed\"",
+        ),
+        "huggingface_streaming": (
+            "source = \"huggingface_streaming\"",
+            "hf_repo_id = \"zacheryasc/airfrans-processed\"",
+            "hf_repo_type = \"dataset\"",
+            "hf_path_prefix = \"processed/full\"",
+            "cache_dir = \"artifacts/data_cache/airfrans_processed\"",
+            "streaming_queue_max_cases = 8",
+        ),
+        "public_zip_streaming": (
+            "source = \"public_zip_streaming\"",
+            "hf_repo_type = \"dataset\"",
+            "cache_dir = \"artifacts/data_cache/airfrans_public_streaming\"",
+            "streaming_scratch_dir = \"artifacts/data_cache/airfrans_public_streaming/_raw\"",
+            "streaming_cache_max_bytes = 34359738368",
+            "streaming_cache_high_water_bytes = 30064771072",
+            "streaming_cache_low_water_bytes = 21474836480",
+            "streaming_queue_max_cases = 8",
+        ),
+    }
+    if data_source not in data_source_lines:
+        raise ValueError(f"unsupported data source for sweep config: {data_source}")
+    if all_points_per_case and points_per_case is not None:
+        raise ValueError("all_points_per_case cannot be combined with points_per_case")
     condition_depth = 4 if band in {"100m", "700m"} else 2
     condition_dim = min(shape["hidden_width"], 1024)
     decision = coordinate_encoding_decision(model_family, coordinate_encoding)
@@ -272,17 +562,13 @@ def training_config_text(
             "",
             "[data]",
             f"root = {json.dumps(data_root)}",
-            "source = \"huggingface\"",
-            "hf_repo_id = \"zacheryasc/airfrans-processed\"",
-            "hf_repo_type = \"dataset\"",
-            "hf_path_prefix = \"processed/full\"",
-            "cache_dir = \"artifacts/data_cache/airfrans_processed\"",
-            "train_cases = 900",
-            "val_cases = 50",
-            "test_cases = 50",
-            "points_per_case = 8192",
+            *data_source_lines[data_source],
+            f"train_cases = {train_cases}",
+            f"val_cases = {val_cases}",
+            f"test_cases = {test_cases}",
+            "all_points_per_case = true" if all_points_per_case else f"points_per_case = {points_per_case}",
             f"batch_size = {shape['batch_size']}",
-            "streaming_normalization_cases = 16" if band == "700m" else None,
+            f"streaming_normalization_cases = {streaming_normalization_cases}" if streaming_normalization_cases is not None else None,
             "",
             "[coordinate_encoding]",
             _coordinate_encoding_toml(coordinate_encoding),
@@ -313,10 +599,10 @@ def training_config_text(
             f"coordinate_encoding_compatibility = {json.dumps(metadata['coordinate_encoding_compatibility'])}",
             "",
             "[optim]",
-            "lr = 0.0001",
+            f"lr = {lr}",
             "weight_decay = 0.0001",
             f"steps = {shape['steps']}",
-            "log_interval = 25",
+            f"log_interval = {log_interval}",
             "",
             "[device]",
             "type = \"cuda\"",
@@ -331,16 +617,22 @@ def training_config_text(
             "",
             "[checkpoint]",
             "interval_seconds = 900",
+            f"policy = {json.dumps(checkpoint_policy)}",
+            f"include_optimizer_state = {str(checkpoint_include_optimizer_state).lower()}",
+            f"include_rng_state = {str(checkpoint_include_rng_state).lower()}",
             "",
             "[stability]",
             "max_grad_norm = 1.0",
+            f"dead_curve_patience_evals = {dead_curve_patience_evals}" if dead_curve_patience_evals is not None else None,
+            f"dead_curve_min_relative_improvement = {dead_curve_min_relative_improvement}",
+            f"dead_curve_warmup_steps = {dead_curve_warmup_steps}",
             "",
             "[observability]",
             "backend = \"wandb\"",
             "entity = \"zacheryasc-personal\"",
             "project = \"airfRANS-model-sweep\"",
             f"group = {json.dumps(group)}",
-            f"tags = [\"airfrans\", \"aggressive-oom-sweep\", {json.dumps(band)}, {json.dumps(model_family)}, {json.dumps(coordinate_encoding)}]",
+            f"tags = [\"airfrans\", {json.dumps(group)}, {json.dumps(band)}, {json.dumps(model_family)}, {json.dumps(coordinate_encoding)}]",
             "",
             "[huggingface]",
             "enabled = true",
diff --git a/src/airfrans_frontier/training/artifacts.py b/src/airfrans_frontier/training/artifacts.py
index 68ed763..733bc6d 100644
--- a/src/airfrans_frontier/training/artifacts.py
+++ b/src/airfrans_frontier/training/artifacts.py
@@ -147,7 +147,7 @@ def _artifact_manifest_text(root: Path) -> tuple[str, str]:
     files = sorted(
         path
         for path in root.rglob("*")
-        if path.is_file() and path.name not in {"artifact_manifest.json", "checksums.txt", "verification_report.json"}
+        if path.is_file() and path.name not in {"artifact_manifest.json", "checksums.txt", "verification_report.json"} and not path.name.endswith(".tmp")
     )
     manifest = {
         "artifact_dir": str(root),
diff --git a/src/airfrans_frontier/training/config.py b/src/airfrans_frontier/training/config.py
index 064a7fb..b7ff55f 100644
--- a/src/airfrans_frontier/training/config.py
+++ b/src/airfrans_frontier/training/config.py
@@ -9,6 +9,7 @@ from airfrans_frontier.model_metadata import model_metadata_for_family
 
 _MODEL_TYPES = {
     "mlp",
+    "mlp_encoded_baseline",
     "film_fourier_mlp",
     "film_fourier_inr",
     "nerf_cfd_multires",
@@ -32,7 +33,8 @@ class DataConfig:
     train_cases: int
     val_cases: int
     test_cases: int
-    points_per_case: int
+    points_per_case: int | None
+    all_points_per_case: bool
     batch_size: int
     source: str
     hf_repo_id: str | None
@@ -49,7 +51,6 @@ class DataConfig:
     streaming_upload_batch_size: int = 8
     streaming_normalization_cases: int | None = None
 
-
 @dataclass(frozen=True)
 class ModelConfig:
     type: str
@@ -104,12 +105,17 @@ class OptimConfig:
 @dataclass(frozen=True)
 class CheckpointConfig:
     interval_seconds: int
+    include_optimizer_state: bool
+    include_rng_state: bool
+    policy: str
 
 
 @dataclass(frozen=True)
 class StabilityConfig:
     max_grad_norm: float | None
-
+    dead_curve_patience_evals: int | None
+    dead_curve_min_relative_improvement: float
+    dead_curve_warmup_steps: int
 
 @dataclass(frozen=True)
 class PrecisionConfig:
@@ -279,15 +285,19 @@ def load_training_config(path: str | Path) -> TrainingConfig:
         raise ValueError("data.streaming_cache_high_water_bytes must be <= data.streaming_cache_max_bytes")
     if streaming_low_water_bytes >= streaming_high_water_bytes:
         raise ValueError("data.streaming_cache_low_water_bytes must be < data.streaming_cache_high_water_bytes")
+    all_points_per_case = _boolean(data_raw, "all_points_per_case") if "all_points_per_case" in data_raw else False
+    points_per_case = None if all_points_per_case else _integer(data_raw, "points_per_case", minimum=1)
+
 
     data = DataConfig(
         root=_path(data_raw, "root"),
         train_cases=_integer(data_raw, "train_cases", minimum=1),
         val_cases=_integer(data_raw, "val_cases", minimum=0),
         test_cases=_integer(data_raw, "test_cases", minimum=0),
-        points_per_case=_integer(data_raw, "points_per_case", minimum=1),
+        points_per_case=points_per_case,
+        all_points_per_case=all_points_per_case,
         batch_size=_integer(data_raw, "batch_size", minimum=1),
-        source=_choice(_string(data_raw, "source", default="local").lower(), {"local", "huggingface", "public_zip_streaming"}, "data.source"),
+        source=_choice(_string(data_raw, "source", default="local").lower(), {"local", "huggingface", "huggingface_streaming", "public_zip_streaming"}, "data.source"),
         hf_repo_id=_optional_string(data_raw, "hf_repo_id"),
         hf_repo_type=_choice(_string(data_raw, "hf_repo_type", default="dataset"), {"dataset"}, "data.hf_repo_type"),
         hf_path_prefix=_string(data_raw, "hf_path_prefix", default=""),
@@ -354,9 +364,15 @@ def load_training_config(path: str | Path) -> TrainingConfig:
     loss = LossConfig(type=_choice(_string(loss_raw, "type"), {"normalized_mse"}, "loss.type"))
     checkpoint = CheckpointConfig(
         interval_seconds=_integer(checkpoint_raw, "interval_seconds", minimum=0, default=1800),
+        include_optimizer_state=_boolean(checkpoint_raw, "include_optimizer_state") if "include_optimizer_state" in checkpoint_raw else True,
+        include_rng_state=_boolean(checkpoint_raw, "include_rng_state") if "include_rng_state" in checkpoint_raw else True,
+        policy=_choice(_string(checkpoint_raw, "policy", default="full").lower(), {"full", "lightweight_scaling_probe"}, "checkpoint.policy"),
     )
     stability = StabilityConfig(
         max_grad_norm=_optional_number(stability_raw, "max_grad_norm", minimum=0.0, exclusive_minimum=True),
+        dead_curve_patience_evals=_optional_integer(stability_raw, "dead_curve_patience_evals", minimum=1),
+        dead_curve_min_relative_improvement=_number(stability_raw, "dead_curve_min_relative_improvement", minimum=0.0, default=0.005),
+        dead_curve_warmup_steps=_integer(stability_raw, "dead_curve_warmup_steps", minimum=0, default=0),
     )
     precision = PrecisionConfig(
         dtype=_choice(_string(precision_raw, "dtype", default="float32").lower(), {"float32", "bf16"}, "precision.dtype"),
diff --git a/src/airfrans_frontier/training/data.py b/src/airfrans_frontier/training/data.py
index 0301fb1..e720335 100644
--- a/src/airfrans_frontier/training/data.py
+++ b/src/airfrans_frontier/training/data.py
@@ -52,7 +52,7 @@ class SplitArrays:
     features: FloatArray
     targets: FloatArray
     case_ids: tuple[str, ...]
-
+    point_counts: tuple[int, ...] = ()
 
 @dataclass(frozen=True)
 class DatasetBundle:
@@ -177,7 +177,7 @@ def build_dataset_bundle(
     train_cases: int,
     val_cases: int,
     test_cases: int,
-    points_per_case: int,
+    points_per_case: int | None,
     seed: int,
 ) -> DatasetBundle:
     split = create_case_split(
@@ -194,7 +194,7 @@ def build_dataset_bundle_for_split(
     samples: list[SimulationSample],
     *,
     split: CaseSplit,
-    points_per_case: int,
+    points_per_case: int | None,
     seed: int,
 ) -> DatasetBundle:
     validate_common_schema(samples)
@@ -226,7 +226,7 @@ def build_dataset_bundle_for_split(
 def build_split_arrays(
     samples_by_id: dict[str, SimulationSample],
     case_ids: tuple[str, ...],
-    points_per_case: int,
+    points_per_case: int | None,
     *,
     seed: int,
 ) -> SplitArrays:
@@ -235,9 +235,11 @@ def build_split_arrays(
     rng = np.random.default_rng(seed)
     feature_parts: list[FloatArray] = []
     target_parts: list[FloatArray] = []
+    point_counts: list[int] = []
     for case_id in case_ids:
         sample = samples_by_id[case_id]
         indices = _sample_indices(sample.num_points, points_per_case, rng)
+        point_counts.append(int(indices.shape[0]))
         feature_parts.append(np.ascontiguousarray(sample.features[indices], dtype=np.float32))
         target_parts.append(np.ascontiguousarray(sample.targets[indices], dtype=np.float32))
 
@@ -245,13 +247,14 @@ def build_split_arrays(
         features=np.concatenate(feature_parts, axis=0),
         targets=np.concatenate(target_parts, axis=0),
         case_ids=case_ids,
+        point_counts=tuple(point_counts),
     )
 
 
-def _sample_indices(num_points: int, points_per_case: int, rng: np.random.Generator) -> NDArray[np.int64]:
+def _sample_indices(num_points: int, points_per_case: int | None, rng: np.random.Generator) -> NDArray[np.int64]:
     if num_points <= 0:
         raise ValueError("Cannot sample from an empty simulation")
-    if points_per_case >= num_points:
+    if points_per_case is None or points_per_case >= num_points:
         return np.arange(num_points, dtype=np.int64)
     return np.sort(rng.choice(num_points, size=points_per_case, replace=False)).astype(np.int64)
 
diff --git a/src/airfrans_frontier/training/loop.py b/src/airfrans_frontier/training/loop.py
index 9ca5cfb..7288812 100644
--- a/src/airfrans_frontier/training/loop.py
+++ b/src/airfrans_frontier/training/loop.py
@@ -16,6 +16,7 @@ from torch.nn import functional as F
 
 from airfrans_frontier.models import (
     DeepONetBranchTrunk,
+    EncodedPointwiseMLP,
     FourierFiLMMLP,
     LocalPointTransformer,
     NeRFCFDMultiRes,
@@ -69,6 +70,58 @@ class TrainingResult:
     final_metrics: dict[str, Any]
 
 
+class DeadCurveError(RuntimeError):
+    pass
+
+
+@dataclass
+class LossMovementMonitor:
+    patience_evals: int
+    min_relative_improvement: float
+    warmup_steps: int
+    baseline_loss: float | None = None
+    baseline_step: int | None = None
+    best_loss: float | None = None
+    evals_after_baseline: int = 0
+
+    @classmethod
+    def from_config(cls, config: TrainingConfig) -> "LossMovementMonitor | None":
+        if config.stability.dead_curve_patience_evals is None:
+            return None
+        return cls(
+            patience_evals=config.stability.dead_curve_patience_evals,
+            min_relative_improvement=config.stability.dead_curve_min_relative_improvement,
+            warmup_steps=config.stability.dead_curve_warmup_steps,
+        )
+
+    def observe(self, *, step: int, loss: float) -> dict[str, Any] | None:
+        if step < self.warmup_steps:
+            return None
+        if self.baseline_loss is None:
+            self.baseline_loss = float(loss)
+            self.baseline_step = int(step)
+            self.best_loss = float(loss)
+            self.evals_after_baseline = 0
+            return None
+        self.evals_after_baseline += 1
+        self.best_loss = min(float(self.best_loss), float(loss)) if self.best_loss is not None else float(loss)
+        denominator = max(abs(float(self.baseline_loss)), 1e-12)
+        relative_improvement = (float(self.baseline_loss) - float(self.best_loss)) / denominator
+        if self.evals_after_baseline >= self.patience_evals and relative_improvement < self.min_relative_improvement:
+            return {
+                "baseline_loss": float(self.baseline_loss),
+                "baseline_step": self.baseline_step,
+                "best_loss": float(self.best_loss),
+                "latest_loss": float(loss),
+                "relative_improvement": float(relative_improvement),
+                "min_relative_improvement": float(self.min_relative_improvement),
+                "patience_evals": int(self.patience_evals),
+                "evals_after_baseline": int(self.evals_after_baseline),
+                "warmup_steps": int(self.warmup_steps),
+            }
+        return None
+
+
 def train_from_config_path(path: str | Path, resume_path: str | Path | None = None) -> TrainingResult:
     config = load_training_config(path)
     return train(config, resume_path=resume_path or os.environ.get("AIRFRANS_RESUME_CHECKPOINT"))
@@ -204,8 +257,8 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
             }
         )
 
-    if config.data.source == "public_zip_streaming":
-        return _train_public_zip_streaming(
+    if config.data.source in {"public_zip_streaming", "huggingface_streaming"}:
+        return _train_streaming_data(
             config=config,
             resume=resume,
             resume_info=resume_info,
@@ -258,6 +311,9 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
                 "cache_dir": str(config.data.cache_dir) if config.data.cache_dir is not None else None,
                 "case_count": len(samples),
                 "total_points": sum(sample.num_points for sample in samples),
+                "selected_points": int(bundle.train.features.shape[0] + (bundle.val.features.shape[0] if bundle.val is not None else 0) + (bundle.test.features.shape[0] if bundle.test is not None else 0)),
+                "all_points_per_case": config.data.all_points_per_case,
+                "points_per_case": config.data.points_per_case,
                 "feature_names": list(bundle.feature_names),
                 "target_names": list(bundle.target_names),
                 "cases": [
@@ -356,6 +412,7 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
     start_step = 0
     best_val_loss: float | None = None
     initial_train_loss: float | None = None
+    loss_movement_monitor = LossMovementMonitor.from_config(config)
     full_train_eval_cases = len(bundle.train.case_ids)
     train_eval_cases = full_train_eval_cases
     train_eval_features = train_features
@@ -363,7 +420,12 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
     train_eval_case_limit = config.data.streaming_normalization_cases
     if train_eval_case_limit is not None and train_eval_case_limit < full_train_eval_cases:
         train_eval_cases = max(1, train_eval_case_limit)
-        train_eval_rows = min(train_features.shape[0], train_eval_cases * config.data.points_per_case)
+        if bundle.train.point_counts:
+            train_eval_rows = min(train_features.shape[0], sum(bundle.train.point_counts[:train_eval_cases]))
+        elif config.data.points_per_case is not None:
+            train_eval_rows = min(train_features.shape[0], train_eval_cases * config.data.points_per_case)
+        else:
+            train_eval_rows = train_features.shape[0]
         train_eval_features = train_features[:train_eval_rows]
         train_eval_targets = train_targets[:train_eval_rows]
 
@@ -678,6 +740,22 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
                         )
                     )
                 )
+                dead_curve = loss_movement_monitor.observe(step=step, loss=float(current_metric)) if loss_movement_monitor is not None else None
+                if dead_curve is not None:
+                    _write_failure(
+                        writer,
+                        phase="training",
+                        step=step,
+                        error_type="DeadCurveError",
+                        error_message="loss did not improve enough during the early measurement window",
+                        latest_loss=float(current_metric),
+                        latest_grad_norm=last_grad_norm,
+                        latest_checkpoint=LATEST_CHECKPOINT if (writer.run_dir / LATEST_CHECKPOINT).is_file() else None,
+                        config=config,
+                        failure_category="plateau_or_dead_recipe",
+                        dead_curve=dead_curve,
+                    )
+                    raise DeadCurveError("loss did not improve enough during the early measurement window")
                 last_log_at = time.perf_counter()
                 last_log_step = step
                 model.train()
@@ -778,6 +856,10 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
         "val_cases": len(bundle.split.val_ids),
         "test_cases": len(bundle.split.test_ids),
         "points_per_case": config.data.points_per_case,
+        "all_points_per_case": config.data.all_points_per_case,
+        "effective_point_updates": config.optim.steps * config.data.batch_size,
+        "checkpoint_policy": config.checkpoint.policy,
+        "checkpoint_include_optimizer_state": config.checkpoint.include_optimizer_state,
         "steps": config.optim.steps,
         "elapsed_seconds": elapsed,
         "points_per_sec": last_points_per_sec,
@@ -882,7 +964,7 @@ def train(config: TrainingConfig, *, resume_path: str | Path | None = None) -> T
     observer.finish(exit_code=0)
     return TrainingResult(run_dir=writer.run_dir, final_metrics=final_metrics)
 
-def _train_public_zip_streaming(
+def _train_streaming_data(
     *,
     config: TrainingConfig,
     resume: Path | None,
@@ -907,6 +989,7 @@ def _train_public_zip_streaming(
     last_points_per_sec: float | None = None
     first_streaming_metric_written = False
     first_streaming_checkpoint_written = False
+    loss_movement_monitor = LossMovementMonitor.from_config(config)
 
     def record_streaming_metrics(metrics: dict[str, Any]) -> None:
         nonlocal first_streaming_metric_written
@@ -950,7 +1033,7 @@ def _train_public_zip_streaming(
         run_manifest.update(
             {
                 "phase": "initialized",
-                "data_mode": "public_zip_streaming",
+                "data_mode": config.data.source,
                 "parameter_count": count_parameters(model),
                 **calibration_fields,
                 **protocol_fields,
@@ -1284,6 +1367,23 @@ def _train_public_zip_streaming(
                         latest_checkpoint=LATEST_CHECKPOINT if should_checkpoint else None,
                     )
                 )
+                dead_curve = loss_movement_monitor.observe(step=step, loss=float(current_metric)) if loss_movement_monitor is not None else None
+                if dead_curve is not None:
+                    _write_failure(
+                        writer,
+                        phase="streaming_training",
+                        step=step,
+                        error_type="DeadCurveError",
+                        error_message="loss did not improve enough during the early measurement window",
+                        latest_loss=float(current_metric),
+                        latest_grad_norm=last_grad_norm,
+                        latest_checkpoint=LATEST_CHECKPOINT if (writer.run_dir / LATEST_CHECKPOINT).is_file() else None,
+                        streaming_summary=streaming.telemetry_summary(),
+                        config=config,
+                        failure_category="plateau_or_dead_recipe",
+                        dead_curve=dead_curve,
+                    )
+                    raise DeadCurveError("loss did not improve enough during the early measurement window")
                 last_log_at = time.perf_counter()
                 last_log_step = step
                 model.train()
@@ -1341,6 +1441,10 @@ def _train_public_zip_streaming(
             "val_cases": len(bundle.split.val_ids),
             "test_cases": len(bundle.split.test_ids),
             "points_per_case": config.data.points_per_case,
+            "all_points_per_case": config.data.all_points_per_case,
+            "effective_point_updates": config.optim.steps * config.data.batch_size,
+            "checkpoint_policy": config.checkpoint.policy,
+            "checkpoint_include_optimizer_state": config.checkpoint.include_optimizer_state,
             "steps": config.optim.steps,
             "elapsed_seconds": elapsed,
             "points_per_sec": last_points_per_sec,
@@ -1348,7 +1452,7 @@ def _train_public_zip_streaming(
             "validation_runtime_seconds": validation_runtime_seconds,
             "checkpoint_interval_seconds": config.checkpoint.interval_seconds,
             "data_source": config.data.source,
-            "data_mode": "public_zip_streaming",
+            "data_mode": config.data.source,
             "data_public_source_url": config.data.public_source_url,
             "data_hf_repo_id": config.data.hf_repo_id,
             "data_hf_path_prefix": config.data.hf_path_prefix,
@@ -1410,7 +1514,7 @@ def _train_public_zip_streaming(
                 "phase": "completed",
                 "finished_at": time.time(),
                 "exit_code": 0,
-                "data_mode": "public_zip_streaming",
+                "data_mode": config.data.source,
                 "final_metrics_path": str(writer.run_dir / "final_metrics.json"),
                 "checkpoint_latest_path": str(writer.run_dir / LATEST_CHECKPOINT),
                 "checkpoint_best_path": str(writer.run_dir / BEST_CHECKPOINT),
@@ -1479,7 +1583,7 @@ def _train_public_zip_streaming(
                 config=config,
                 failure_category="streaming_training",
             )
-        run_manifest.update({"phase": "failed", "finished_at": time.time(), "exit_code": 1, "data_mode": "public_zip_streaming", **uploader.finalize(training_success=False)})
+        run_manifest.update({"phase": "failed", "finished_at": time.time(), "exit_code": 1, "data_mode": config.data.source, **uploader.finalize(training_success=False)})
         writer.write_json("run_manifest.json", run_manifest)
         writer.write_artifact_manifest()
         try:
@@ -1545,6 +1649,16 @@ def _build_model(config: TrainingConfig, bundle: DatasetBundle, *, output_dim: i
             depth=config.model.depth,
             activation=config.model.activation,
         )
+    if model_type == "mlp_encoded_baseline":
+        return EncodedPointwiseMLP(
+            feature_names=bundle.feature_names,
+            output_dim=output_dim,
+            coordinate_features=config.model.coordinate_features,
+            coordinate_encoding_spec=coord_spec,
+            hidden_width=config.model.hidden_width,
+            depth=config.model.depth,
+            activation=config.model.activation,
+        )
     if model_type in {"film_fourier_mlp", "film_fourier_inr"}:
         return FourierFiLMMLP(
             feature_names=bundle.feature_names,
@@ -1809,6 +1923,15 @@ def _job_manifest(*, config: TrainingConfig, run_id: str, run_dir: Path) -> dict
         "proxy_for": config.model_metadata.proxy_for,
         "proxy_notes": config.model_metadata.proxy_notes,
         "model_metadata": asdict(config.model_metadata),
+        "data": {
+            "source": config.data.source,
+            "train_cases": config.data.train_cases,
+            "val_cases": config.data.val_cases,
+            "test_cases": config.data.test_cases,
+            "all_points_per_case": config.data.all_points_per_case,
+            "points_per_case": config.data.points_per_case,
+            "batch_size": config.data.batch_size,
+        },
         "coordinate_encoding": {
             "type": config.coordinate_encoding.type,
             "features": list(config.coordinate_encoding.features),
@@ -1820,10 +1943,19 @@ def _job_manifest(*, config: TrainingConfig, run_id: str, run_dir: Path) -> dict
         },
         "checkpoint_policy": {
             "interval_seconds": config.checkpoint.interval_seconds,
+            "policy": config.checkpoint.policy,
+            "include_optimizer_state": config.checkpoint.include_optimizer_state,
+            "include_rng_state": config.checkpoint.include_rng_state,
             "latest": LATEST_CHECKPOINT,
             "best": BEST_CHECKPOINT,
             "final": FINAL_CHECKPOINT,
         },
+        "stability": {
+            "max_grad_norm": config.stability.max_grad_norm,
+            "dead_curve_patience_evals": config.stability.dead_curve_patience_evals,
+            "dead_curve_min_relative_improvement": config.stability.dead_curve_min_relative_improvement,
+            "dead_curve_warmup_steps": config.stability.dead_curve_warmup_steps,
+        },
         "observability": {
             "backend": config.observability.backend,
             "project": config.observability.project,
@@ -2111,6 +2243,8 @@ def _checkpoint_payload(
     initial_train_loss: float | None,
     final_metrics: dict[str, Any] | None,
 ) -> dict[str, Any]:
+    include_optimizer = config.checkpoint.include_optimizer_state
+    include_rng = config.checkpoint.include_rng_state
     return {
         "schema_version": CHECKPOINT_SCHEMA_VERSION,
         "run_id": os.environ.get("AIRFRANS_REMOTE_RUN_ID", config.run.name),
@@ -2123,18 +2257,21 @@ def _checkpoint_payload(
         "output_dim": bundle.train.targets.shape[1],
         "model_config": asdict(config.model),
         "model_state_dict": model.state_dict(),
-        "optimizer_state_dict": optimizer.state_dict(),
+        "optimizer_state_dict": optimizer.state_dict() if include_optimizer else None,
         "scheduler_state_dict": None,
         "config": config.config_text,
         "config_hash": _config_hash(config),
         "normalization": stats.to_dict(),
         "target_names": bundle.target_names,
         "feature_names": bundle.feature_names,
-        "rng_state": random.getstate(),
-        "numpy_rng_state": np.random.get_state(),
-        "torch_rng_state": torch.get_rng_state(),
-        "cuda_rng_state": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None,
-        "batch_rng_state": rng.bit_generator.state,
+        "rng_state": random.getstate() if include_rng else None,
+        "numpy_rng_state": np.random.get_state() if include_rng else None,
+        "torch_rng_state": torch.get_rng_state() if include_rng else None,
+        "cuda_rng_state": torch.cuda.get_rng_state_all() if include_rng and torch.cuda.is_available() else None,
+        "batch_rng_state": rng.bit_generator.state if include_rng else None,
+        "checkpoint_policy": config.checkpoint.policy,
+        "include_optimizer_state": include_optimizer,
+        "include_rng_state": include_rng,
         "final_metrics": final_metrics,
     }
 
diff --git a/src/airfrans_frontier/training/streaming_data.py b/src/airfrans_frontier/training/streaming_data.py
index a0d3d96..b8e1c59 100644
--- a/src/airfrans_frontier/training/streaming_data.py
+++ b/src/airfrans_frontier/training/streaming_data.py
@@ -1,9 +1,11 @@
 from __future__ import annotations
 
 import json
+import queue
 import os
 import shutil
 import time
+import threading
 from dataclasses import dataclass, field
 from pathlib import Path, PurePosixPath
 from typing import Any, Iterable, Iterator, Mapping
@@ -335,6 +337,8 @@ class PublicZipStreamingCache:
             "cache_high_water_bytes": self.data_config.streaming_cache_high_water_bytes,
             "cache_low_water_bytes": self.data_config.streaming_cache_low_water_bytes,
             "queue_max_cases": self.data_config.streaming_queue_max_cases,
+            "all_points_per_case": self.data_config.all_points_per_case,
+            "points_per_case": self.data_config.points_per_case,
             "ranged_bytes_read": self.reader.bytes_read if self.reader is not None else None,
             "cases": [
                 {
@@ -500,11 +504,246 @@ class PublicZipStreamingCache:
         return {"schema_version": 1, "cases": {}, "sampling": {"train": {}, "val": {}, "test": {}}}
 
 
+
+class HuggingFaceStreamingCache:
+    def __init__(self, data_config: DataConfig, *, run_dir: Path, recorder: StreamingEventRecorder) -> None:
+        if not data_config.hf_repo_id:
+            raise ValueError("data.hf_repo_id is required when data.source = 'huggingface_streaming'")
+        self.data_config = data_config
+        self.run_dir = run_dir
+        self.recorder = recorder
+        self.repo_id = data_config.hf_repo_id
+        self.repo_type = data_config.hf_repo_type
+        self.path_prefix = data_config.hf_path_prefix.strip("/")
+        self.cache_dir = data_config.cache_dir or data_config.root
+        self.data_root = self.cache_dir / self.path_prefix if self.path_prefix else self.cache_dir
+        self.state_path = run_dir / STREAMING_STATE
+        self.indices_dir = run_dir / "streaming_sample_indices"
+        self.cache_dir.mkdir(parents=True, exist_ok=True)
+        self.data_root.mkdir(parents=True, exist_ok=True)
+        self.indices_dir.mkdir(parents=True, exist_ok=True)
+        self.case_files: dict[str, str] = {}
+        self.state: dict[str, Any] = _read_json(self.state_path) if self.state_path.is_file() else self._empty_state()
+        self._active_cases: set[str] = set()
+        self._last_access: dict[str, float] = {}
+        self._loaded: dict[str, SimulationSample] = {}
+
+    def enumerate_cases(self) -> tuple[str, ...]:
+        self.recorder.emit(
+            "dataset_enumeration_start",
+            phase="data",
+            source="huggingface_streaming",
+            hf_repo_id=self.repo_id,
+            hf_repo_type=self.repo_type,
+            hf_path_prefix=self.path_prefix,
+        )
+        token = _resolve_optional_secret("HF_TOKEN")
+        try:
+            from huggingface_hub import HfApi
+        except ModuleNotFoundError as exc:
+            raise RuntimeError("huggingface_hub is required when data.source = 'huggingface_streaming'") from exc
+        api = HfApi(token=token)
+        files = api.list_repo_files(repo_id=self.repo_id, repo_type=self.repo_type)
+        prefix = f"{self.path_prefix}/" if self.path_prefix else ""
+        case_files = {
+            Path(name).stem: name
+            for name in files
+            if name.startswith(prefix) and name.endswith(".npz") and Path(name).name
+        }
+        if not case_files:
+            raise FileNotFoundError(f"No .npz cases found in HF dataset {self.repo_id!r} under {self.path_prefix!r}")
+        self.case_files = dict(sorted(case_files.items()))
+        case_ids = tuple(self.case_files)
+        self.state["source"] = "huggingface_streaming"
+        self.state["hf_repo_id"] = self.repo_id
+        self.state["hf_repo_type"] = self.repo_type
+        self.state["hf_path_prefix"] = self.path_prefix
+        self.state["enumerated_case_count"] = len(case_ids)
+        self._write_state()
+        self.recorder.emit(
+            "dataset_enumeration_end",
+            phase="data",
+            source="huggingface_streaming",
+            hf_repo_id=self.repo_id,
+            hf_repo_type=self.repo_type,
+            hf_path_prefix=self.path_prefix,
+            case_count=len(case_ids),
+        )
+        return case_ids
+
+    def ensure_case(self, case_id: str) -> SimulationSample:
+        cached = self._load_valid_cached_case(case_id)
+        if cached is not None:
+            self._active_cases.add(case_id)
+            self._last_access[case_id] = time.time()
+            return cached
+        filename = self.case_files.get(case_id)
+        if filename is None:
+            raise KeyError(f"Unknown Hugging Face AirfRANS case: {case_id}")
+        token = _resolve_optional_secret("HF_TOKEN")
+        self._mark_case(case_id, status="downloading", path=str(self.data_root / f"{case_id}.npz"), download_started_at=time.time())
+        started = time.perf_counter()
+        self.recorder.emit("hf_case_download_start", phase="data", case_id=case_id, hf_filename=filename)
+        try:
+            from huggingface_hub import hf_hub_download
+        except ModuleNotFoundError as exc:
+            raise RuntimeError("huggingface_hub is required when data.source = 'huggingface_streaming'") from exc
+        try:
+            downloaded_path = Path(
+                hf_hub_download(
+                    repo_id=self.repo_id,
+                    filename=filename,
+                    repo_type=self.repo_type,
+                    local_dir=str(self.cache_dir),
+                    token=token,
+                )
+            )
+            sample_path = self.data_root / f"{case_id}.npz"
+            if downloaded_path != sample_path and downloaded_path.is_file() and not sample_path.is_file():
+                sample_path.parent.mkdir(parents=True, exist_ok=True)
+                shutil.copy2(downloaded_path, sample_path)
+            sample = load_simulation_npz(sample_path)
+        except Exception as exc:
+            self._mark_case(case_id, status="failed", error_type=type(exc).__name__, error_message=str(exc), failed_at=time.time())
+            self.recorder.emit("hf_case_download_failure", phase="data", case_id=case_id, error_type=type(exc).__name__, error_message=str(exc))
+            raise
+        processed_bytes = sample.source_path.stat().st_size
+        self.recorder.add_downloaded_bytes(processed_bytes)
+        self._loaded[case_id] = sample
+        self._active_cases.add(case_id)
+        self._last_access[case_id] = time.time()
+        elapsed = time.perf_counter() - started
+        self._mark_case(
+            case_id,
+            status="processed",
+            path=str(sample.source_path),
+            processed_bytes=processed_bytes,
+            points=sample.num_points,
+            processed_at=time.time(),
+            processing_seconds=elapsed,
+            hf_filename=filename,
+        )
+        self.recorder.emit(
+            "hf_case_download_end",
+            phase="data",
+            case_id=case_id,
+            hf_filename=filename,
+            processed_bytes=processed_bytes,
+            points=sample.num_points,
+            download_seconds=elapsed,
+        )
+        self.recorder.emit(
+            "processing_end",
+            phase="data",
+            case_id=case_id,
+            processed_bytes=processed_bytes,
+            points=sample.num_points,
+            processing_seconds=elapsed,
+        )
+        self._observe_cache()
+        return sample
+
+    def release_case(self, case_id: str, *, consumed: bool = True) -> None:
+        self._active_cases.discard(case_id)
+        self._loaded.pop(case_id, None)
+        self._last_access[case_id] = time.time()
+        entry = self.state.setdefault("cases", {}).setdefault(case_id, {})
+        if consumed:
+            entry["consumed_count"] = int(entry.get("consumed_count", 0) or 0) + 1
+            entry["last_consumed_at"] = time.time()
+        self._write_state()
+        self._observe_cache()
+
+    def cache_size_bytes(self) -> int:
+        return _tree_size_bytes(self.data_root, suffix=".npz")
+
+    def data_manifest(self, *, split: CaseSplit, feature_names: tuple[str, ...], target_names: tuple[str, ...]) -> dict[str, Any]:
+        selected = set(split.train_ids) | set(split.val_ids) | set(split.test_ids)
+        cases_state = self.state.get("cases", {}) if isinstance(self.state.get("cases"), dict) else {}
+        return {
+            "source": "huggingface_streaming",
+            "hf_repo_id": self.repo_id,
+            "hf_repo_type": self.repo_type,
+            "hf_path_prefix": self.path_prefix,
+            "configured_root": str(self.data_config.root),
+            "cache_dir": str(self.cache_dir),
+            "data_root": str(self.data_root),
+            "case_count": self.state.get("enumerated_case_count"),
+            "selected_case_count": len(selected),
+            "train_cases": len(split.train_ids),
+            "val_cases": len(split.val_ids),
+            "test_cases": len(split.test_ids),
+            "feature_names": list(feature_names),
+            "target_names": list(target_names),
+            "queue_max_cases": self.data_config.streaming_queue_max_cases,
+            "all_points_per_case": self.data_config.all_points_per_case,
+            "points_per_case": self.data_config.points_per_case,
+            "cases": [
+                {
+                    "case_id": case_id,
+                    "split": _split_name_for_case(split, case_id),
+                    "status": cases_state.get(case_id, {}).get("status", "pending") if isinstance(cases_state.get(case_id), dict) else "pending",
+                    "points": cases_state.get(case_id, {}).get("points") if isinstance(cases_state.get(case_id), dict) else None,
+                    "processed_bytes": cases_state.get(case_id, {}).get("processed_bytes") if isinstance(cases_state.get(case_id), dict) else None,
+                    "hf_filename": self.case_files.get(case_id),
+                }
+                for case_id in sorted(selected)
+            ],
+        }
+
+    def _load_valid_cached_case(self, case_id: str) -> SimulationSample | None:
+        if case_id in self._loaded:
+            return self._loaded[case_id]
+        path = self.data_root / f"{case_id}.npz"
+        if not path.is_file():
+            return None
+        try:
+            sample = load_simulation_npz(path)
+        except Exception:
+            path.unlink(missing_ok=True)
+            self.recorder.emit("partial_unit_discarded", phase="data", case_id=case_id, path=str(path))
+            self._mark_case(case_id, status="pending", discarded_invalid_at=time.time())
+            return None
+        self._loaded[case_id] = sample
+        self._last_access[case_id] = time.time()
+        entry = self.state.setdefault("cases", {}).setdefault(case_id, {})
+        if entry.get("status") == "processed":
+            self.recorder.emit("resume_validated_unit_reused", phase="data", case_id=case_id, path=str(path), processed_bytes=path.stat().st_size)
+        self._mark_case(case_id, status="processed", path=str(path), processed_bytes=path.stat().st_size, points=sample.num_points)
+        self._observe_cache()
+        return sample
+
+    def _observe_cache(self) -> None:
+        self.recorder.observe_cache(self.data_root, cache_bytes=self.cache_size_bytes())
+
+    def _mark_case(self, case_id: str, **fields: Any) -> None:
+        cases = self.state.setdefault("cases", {})
+        entry = cases.setdefault(case_id, {})
+        entry.update(fields)
+        entry["updated_at"] = time.time()
+        self._write_state()
+
+    def _write_state(self) -> None:
+        self.state["schema_version"] = 1
+        self.state["cache_dir"] = str(self.cache_dir)
+        self.state["data_root"] = str(self.data_root)
+        self.state["updated_at"] = time.time()
+        _atomic_write_json(self.state_path, self.state)
+
+    @staticmethod
+    def _empty_state() -> dict[str, Any]:
+        return {"schema_version": 1, "cases": {}, "sampling": {"train": {}, "val": {}, "test": {}}}
+
 class StreamingTrainingData:
     def __init__(self, config: TrainingConfig, *, run_dir: Path, recorder: StreamingEventRecorder) -> None:
         self.config = config
         self.recorder = recorder
-        self.cache = PublicZipStreamingCache(config.data, run_dir=run_dir, recorder=recorder)
+        if config.data.source == "public_zip_streaming":
+            self.cache = PublicZipStreamingCache(config.data, run_dir=run_dir, recorder=recorder)
+        elif config.data.source == "huggingface_streaming":
+            self.cache = HuggingFaceStreamingCache(config.data, run_dir=run_dir, recorder=recorder)
+        else:
+            raise ValueError(f"StreamingTrainingData does not support data.source = {config.data.source!r}")
         self.split: CaseSplit | None = None
         self.feature_names: tuple[str, ...] | None = None
         self.target_names: tuple[str, ...] | None = None
@@ -514,11 +753,18 @@ class StreamingTrainingData:
         self._train_counts: list[int] = []
         self._train_offsets: NDArray[np.int64] | None = None
         self._upload_queue = ProcessedDataUploadQueue.from_config(config, run_dir=run_dir, recorder=recorder)
+        self._io_lock = threading.RLock()
+        self._batch_queue: queue.Queue[tuple[FloatArray, FloatArray] | BaseException] | None = None
+        self._batch_thread: threading.Thread | None = None
+        self._batch_stop = threading.Event()
+        self._acquisition_thread: threading.Thread | None = None
+        self._acquisition_stop = threading.Event()
+        self._acquisition_error: BaseException | None = None
 
     @classmethod
     def from_config(cls, config: TrainingConfig, *, run_dir: Path, recorder: StreamingEventRecorder) -> StreamingTrainingData:
-        if config.data.source != "public_zip_streaming":
-            raise ValueError(f"StreamingTrainingData requires data.source = 'public_zip_streaming', got {config.data.source!r}")
+        if config.data.source not in {"public_zip_streaming", "huggingface_streaming"}:
+            raise ValueError(f"StreamingTrainingData requires a streaming data source, got {config.data.source!r}")
         return cls(config, run_dir=run_dir, recorder=recorder)
 
     def prepare(self) -> None:
@@ -543,11 +789,12 @@ class StreamingTrainingData:
             seed=self.config.run.seed,
         )
         first_case = split.train_ids[0]
-        sample = self.cache.ensure_case(first_case)
-        self.feature_names = sample.feature_names
-        self.target_names = sample.target_names
-        self._upload_queue.enqueue(sample.source_path)
-        self.cache.release_case(first_case, consumed=False)
+        with self._io_lock:
+            sample = self.cache.ensure_case(first_case)
+            self.feature_names = sample.feature_names
+            self.target_names = sample.target_names
+            self._upload_queue.enqueue(sample.source_path)
+            self.cache.release_case(first_case, consumed=False)
 
     def load_or_compute_normalization(self) -> NormalizationStats:
         if self.split is None:
@@ -558,6 +805,8 @@ class StreamingTrainingData:
             self.stats = stats
             self._restore_sampling_specs("train")
             self._build_train_offsets()
+            if self._train_offsets is not None:
+                self._start_background_acquisition()
             self.recorder.emit("normalization_end", phase="normalization", reused=True, normalization_runtime_seconds=0.0)
             return stats
         assert self.feature_names is not None
@@ -576,15 +825,16 @@ class StreamingTrainingData:
         accumulator = _StatsAccumulator(feature_names=self.feature_names, target_names=self.target_names)
         rng = np.random.default_rng(self.config.run.seed + _SAMPLE_SEEDS["train"])
         for case_id in normalization_case_ids:
-            sample = self.cache.ensure_case(case_id)
-            self._validate_schema(sample)
-            indices = self._sampling_spec_for_case("train", sample, rng)
-            features = _selected_rows(sample.features, indices)
-            targets = _selected_rows(sample.targets, indices)
-            accumulator.update(features, targets)
-            self._upload_queue.enqueue(sample.source_path)
-            self.cache.release_case(case_id)
-            self._upload_queue.drain()
+            with self._io_lock:
+                sample = self.cache.ensure_case(case_id)
+                self._validate_schema(sample)
+                indices = self._sampling_spec_for_case("train", sample, rng)
+                features = _selected_rows(sample.features, indices)
+                targets = _selected_rows(sample.targets, indices)
+                accumulator.update(features, targets)
+                self._upload_queue.enqueue(sample.source_path)
+                self.cache.release_case(case_id)
+                self._upload_queue.drain()
         stats = accumulator.finish()
         self.stats = stats
         self._build_train_offsets()
@@ -597,6 +847,8 @@ class StreamingTrainingData:
             sample_count=accumulator.count,
             normalization_runtime_seconds=runtime,
         )
+        if self._train_offsets is not None:
+            self._start_background_acquisition()
         return stats
 
     def schema_bundle(self) -> DatasetBundle:
@@ -622,31 +874,167 @@ class StreamingTrainingData:
         self.recorder.set_upload_state(**self._upload_queue.telemetry_state())
         return self.recorder.to_dict()
 
+    def _start_background_acquisition(self) -> None:
+        if self.config.data.source != "huggingface_streaming" or self._acquisition_thread is not None:
+            return
+        if self.split is None:
+            raise RuntimeError("split is missing")
+        targets = tuple(dict.fromkeys(self.split.train_ids + self.split.val_ids + self.split.test_ids))
+        self._acquisition_stop.clear()
+        self._acquisition_error = None
+        self._acquisition_thread = threading.Thread(
+            target=self._background_acquisition_loop,
+            args=(targets,),
+            name="airfrans-hf-dataset-acquisition",
+            daemon=True,
+        )
+        self._acquisition_thread.start()
+        self.recorder.emit("background_acquisition_start", phase="data", source="huggingface_streaming", case_count=len(targets))
+
+    def _background_acquisition_loop(self, case_ids: tuple[str, ...]) -> None:
+        acquired = 0
+        try:
+            for case_id in case_ids:
+                if self._acquisition_stop.is_set():
+                    self.recorder.emit("background_acquisition_stopped", phase="data", source="huggingface_streaming", acquired_cases=acquired, target_cases=len(case_ids))
+                    return
+                with self._io_lock:
+                    sample = self.cache.ensure_case(case_id)
+                    self._validate_schema(sample)
+                    self.cache.release_case(case_id, consumed=False)
+                acquired += 1
+                self.recorder.emit("background_case_ready", phase="data", source="huggingface_streaming", case_id=case_id, acquired_cases=acquired, target_cases=len(case_ids))
+            self.recorder.emit("background_acquisition_complete", phase="data", source="huggingface_streaming", acquired_cases=acquired, target_cases=len(case_ids))
+        except BaseException as exc:
+            self._acquisition_error = exc
+            self.recorder.emit("background_acquisition_failure", phase="data", source="huggingface_streaming", error_type=type(exc).__name__, error_message=str(exc), acquired_cases=acquired, target_cases=len(case_ids))
+
+    def _check_background_acquisition(self) -> None:
+        if self._acquisition_error is not None:
+            raise self._acquisition_error
+
+    def stop_background_acquisition(self, *, wait: bool) -> None:
+        if self._acquisition_thread is None:
+            return
+        if not wait:
+            self._acquisition_stop.set()
+        self._acquisition_thread.join(timeout=None if wait else 2.0)
+        if self._acquisition_thread.is_alive():
+            return
+        self._acquisition_thread = None
+        self._check_background_acquisition()
+
+    def _ensure_train_batch_prefetch(self, *, batch_size: int) -> queue.Queue[tuple[FloatArray, FloatArray] | BaseException]:
+        if self._batch_queue is not None:
+            return self._batch_queue
+        max_batches = max(2, self.config.data.streaming_queue_max_cases * 4)
+        self._batch_queue = queue.Queue(maxsize=max_batches)
+        self._batch_stop.clear()
+        self._batch_thread = threading.Thread(
+            target=self._train_batch_prefetch_loop,
+            args=(batch_size,),
+            name="airfrans-streaming-batch-prefetch",
+            daemon=True,
+        )
+        self._batch_thread.start()
+        self.recorder.emit("train_batch_prefetch_start", phase="training", batch_size=batch_size, max_batches=max_batches)
+        return self._batch_queue
+
+    def _train_batch_prefetch_loop(self, batch_size: int) -> None:
+        assert self.stats is not None
+        assert self.split is not None
+        assert self._batch_queue is not None
+        rng = np.random.default_rng(self.config.run.seed + 505)
+        train_ids = list(self.split.train_ids)
+        batches_per_case = max(1, self.config.data.streaming_queue_max_cases)
+        try:
+            while not self._batch_stop.is_set():
+                for case_index in rng.permutation(len(train_ids)):
+                    if self._batch_stop.is_set():
+                        return
+                    case_id = train_ids[int(case_index)]
+                    batches: list[tuple[FloatArray, FloatArray]] = []
+                    with self._io_lock:
+                        sample = self.cache.ensure_case(case_id)
+                        self._validate_schema(sample)
+                        spec = self._sampling_spec_for_case("train", sample, rng)
+                        max_batches_for_case = max(1, min(batches_per_case, (spec.count + batch_size - 1) // batch_size))
+                        for _ in range(max_batches_for_case):
+                            batches.append(self._sample_normalized_batch(sample, spec, rng, batch_size=batch_size))
+                        self._upload_queue.enqueue(sample.source_path)
+                        self.cache.release_case(case_id)
+                        self._upload_queue.drain()
+                    for batch in batches:
+                        if not self._put_prefetched_batch(batch):
+                            return
+        except BaseException as exc:
+            self.recorder.emit(
+                "train_batch_prefetch_failure",
+                phase="training",
+                error_type=type(exc).__name__,
+                error_message=str(exc),
+            )
+            self._put_prefetched_batch(exc)
+
+    def _sample_normalized_batch(
+        self,
+        sample: SimulationSample,
+        spec: SamplingSpec,
+        rng: np.random.Generator,
+        *,
+        batch_size: int,
+    ) -> tuple[FloatArray, FloatArray]:
+        assert self.stats is not None
+        local_indices = rng.integers(0, spec.count, size=batch_size, dtype=np.int64)
+        source_indices = _source_indices_for_local(spec, local_indices)
+        selected_features = sample.features[source_indices]
+        selected_targets = sample.targets[source_indices]
+        features = ((selected_features - self.stats.feature_mean) / self.stats.feature_std).astype(np.float32, copy=False)
+        targets = ((selected_targets - self.stats.target_mean) / self.stats.target_std).astype(np.float32, copy=False)
+        return np.ascontiguousarray(features, dtype=np.float32), np.ascontiguousarray(targets, dtype=np.float32)
+
+    def _put_prefetched_batch(self, item: tuple[FloatArray, FloatArray] | BaseException) -> bool:
+        assert self._batch_queue is not None
+        while not self._batch_stop.is_set():
+            try:
+                self._batch_queue.put(item, timeout=0.1)
+                return True
+            except queue.Full:
+                continue
+        return False
+
+    def stop_train_batch_prefetch(self) -> None:
+        self._batch_stop.set()
+        if self._batch_thread is not None:
+            self._batch_thread.join(timeout=2.0)
+        self._batch_thread = None
+        self._batch_queue = None
+
+
     def sample_train_batch(self, rng: np.random.Generator, *, batch_size: int, step: int) -> tuple[FloatArray, FloatArray]:
         if self.stats is None or self.split is None:
             raise RuntimeError("Streaming normalization must be computed before sampling")
+        self._check_background_acquisition()
         started = time.perf_counter()
+        if self._train_offsets is None:
+            batch_queue = self._ensure_train_batch_prefetch(batch_size=batch_size)
+            item = batch_queue.get()
+            wait_seconds = time.perf_counter() - started
+            self.recorder.add_trainer_wait(wait_seconds, step=step)
+            self.recorder.mark_first_batch_ready()
+            if isinstance(item, BaseException):
+                raise item
+            self._start_background_acquisition()
+            return item
+
         features = np.empty((batch_size, len(self.stats.feature_names)), dtype=np.float32)
         targets = np.empty((batch_size, len(self.stats.target_names)), dtype=np.float32)
-        if self._train_offsets is None:
-            case_id = self.split.train_ids[int(rng.integers(0, len(self.split.train_ids)))]
-            sample = self.cache.ensure_case(case_id)
-            self._validate_schema(sample)
-            spec = self._sampling_spec_for_case("train", sample, rng)
-            local_indices = rng.integers(0, spec.count, size=batch_size, dtype=np.int64)
-            source_indices = _source_indices_for_local(spec, local_indices)
-            selected_features = sample.features[source_indices]
-            selected_targets = sample.targets[source_indices]
-            features[:] = ((selected_features - self.stats.feature_mean) / self.stats.feature_std).astype(np.float32, copy=False)
-            targets[:] = ((selected_targets - self.stats.target_mean) / self.stats.target_std).astype(np.float32, copy=False)
-            self._upload_queue.enqueue(sample.source_path)
-            self.cache.release_case(case_id)
-        else:
-            total = int(self._train_offsets[-1]) if self._train_offsets.size else 0
-            if total <= 0:
-                raise ValueError("Streaming train split has no sampled points")
-            global_indices = rng.integers(0, total, size=batch_size)
-            case_positions = np.searchsorted(self._train_offsets[1:], global_indices, side="right")
+        total = int(self._train_offsets[-1]) if self._train_offsets.size else 0
+        if total <= 0:
+            raise ValueError("Streaming train split has no sampled points")
+        global_indices = rng.integers(0, total, size=batch_size)
+        case_positions = np.searchsorted(self._train_offsets[1:], global_indices, side="right")
+        with self._io_lock:
             for case_position in np.unique(case_positions):
                 mask = case_positions == case_position
                 case_id = self.split.train_ids[int(case_position)]
@@ -661,10 +1049,10 @@ class StreamingTrainingData:
                 targets[mask] = ((selected_targets - self.stats.target_mean) / self.stats.target_std).astype(np.float32, copy=False)
                 self._upload_queue.enqueue(sample.source_path)
                 self.cache.release_case(case_id)
+            self._upload_queue.drain()
         wait_seconds = time.perf_counter() - started
         self.recorder.add_trainer_wait(wait_seconds, step=step)
         self.recorder.mark_first_batch_ready()
-        self._upload_queue.drain()
         return np.ascontiguousarray(features, dtype=np.float32), np.ascontiguousarray(targets, dtype=np.float32)
 
     def iter_split_batches(self, split_name: str, *, batch_size: int) -> Iterator[tuple[FloatArray, FloatArray]]:
@@ -673,22 +1061,26 @@ class StreamingTrainingData:
         case_ids = self._split_case_ids.get(split_name)
         if case_ids is None:
             raise ValueError(f"Unknown split: {split_name}")
+        self.stop_train_batch_prefetch()
         rng = np.random.default_rng(self.config.run.seed + _SAMPLE_SEEDS[split_name])
         for case_id in case_ids:
-            sample = self.cache.ensure_case(case_id)
-            self._validate_schema(sample)
-            spec = self._sampling_spec_for_case(split_name, sample, rng)
-            for start in range(0, spec.count, batch_size):
-                stop = min(start + batch_size, spec.count)
-                source_indices = _source_indices_for_local(spec, np.arange(start, stop, dtype=np.int64))
-                features = ((sample.features[source_indices] - self.stats.feature_mean) / self.stats.feature_std).astype(np.float32, copy=False)
-                targets = ((sample.targets[source_indices] - self.stats.target_mean) / self.stats.target_std).astype(np.float32, copy=False)
-                yield np.ascontiguousarray(features, dtype=np.float32), np.ascontiguousarray(targets, dtype=np.float32)
-            self._upload_queue.enqueue(sample.source_path)
-            self.cache.release_case(case_id)
-            self._upload_queue.drain()
+            with self._io_lock:
+                sample = self.cache.ensure_case(case_id)
+                self._validate_schema(sample)
+                spec = self._sampling_spec_for_case(split_name, sample, rng)
+                for start in range(0, spec.count, batch_size):
+                    stop = min(start + batch_size, spec.count)
+                    source_indices = _source_indices_for_local(spec, np.arange(start, stop, dtype=np.int64))
+                    features = ((sample.features[source_indices] - self.stats.feature_mean) / self.stats.feature_std).astype(np.float32, copy=False)
+                    targets = ((sample.targets[source_indices] - self.stats.target_mean) / self.stats.target_std).astype(np.float32, copy=False)
+                    yield np.ascontiguousarray(features, dtype=np.float32), np.ascontiguousarray(targets, dtype=np.float32)
+                self._upload_queue.enqueue(sample.source_path)
+                self.cache.release_case(case_id)
+                self._upload_queue.drain()
 
     def finish(self, *, success: bool) -> None:
+        self.stop_train_batch_prefetch()
+        self.stop_background_acquisition(wait=success)
         self._upload_queue.drain(force=True)
         self._upload_queue.finalize(training_success=success)
         self.cache.state["finished_at"] = time.time()
@@ -712,7 +1104,7 @@ class StreamingTrainingData:
         existing = self._sampling_specs[split_name].get(sample.case_id)
         if existing is not None:
             return existing
-        if self.config.data.points_per_case >= sample.num_points:
+        if self.config.data.all_points_per_case or self.config.data.points_per_case is None or self.config.data.points_per_case >= sample.num_points:
             spec = SamplingSpec(case_id=sample.case_id, split_name=split_name, count=sample.num_points, mode="all")
         else:
             indices = np.sort(rng.choice(sample.num_points, size=self.config.data.points_per_case, replace=False)).astype(np.int64)
@@ -1129,10 +1521,21 @@ def _read_json(path: Path) -> dict[str, Any]:
         raise ValueError(f"Expected JSON object in {path}")
     return data
 
+def _resolve_optional_secret(name: str) -> str | None:
+    value = os.environ.get(name)
+    if value:
+        return value
+    path = Path(".env") / name
+    if path.is_file():
+        value = path.read_text().strip()
+        if value:
+            return value
+    return None
+
 
 def _atomic_write_json(path: Path, payload: Mapping[str, Any]) -> None:
     path.parent.mkdir(parents=True, exist_ok=True)
-    tmp = path.with_name(f"{path.name}.tmp")
+    tmp = path.with_name(f"{path.name}.{os.getpid()}.{threading.get_ident()}.tmp")
     tmp.write_text(json.dumps(dict(payload), indent=2, sort_keys=True) + "\n")
     tmp.replace(path)
 
diff --git a/tests/test_data_sources.py b/tests/test_data_sources.py
index 07875d9..8be53e9 100644
--- a/tests/test_data_sources.py
+++ b/tests/test_data_sources.py
@@ -37,6 +37,7 @@ class DataSourceTests(unittest.TestCase):
                 val_cases=0,
                 test_cases=0,
                 points_per_case=1,
+                all_points_per_case=False,
                 batch_size=1,
                 source="huggingface",
                 hf_repo_id="owner/airfrans-processed",
diff --git a/tests/test_mlp.py b/tests/test_mlp.py
index a837475..450a64e 100644
--- a/tests/test_mlp.py
+++ b/tests/test_mlp.py
@@ -9,7 +9,9 @@ remove_pythonpath_entries()
 import torch
 
 from airfrans_frontier.models import (
+    CoordinateEncodingSpec,
     DeepONetBranchTrunk,
+    EncodedPointwiseMLP,
     LocalPointTransformer,
     NeRFCFDMultiRes,
     PointContextPerceiver,
@@ -28,6 +30,23 @@ class PointwiseMLPTests(unittest.TestCase):
 
         self.assertEqual(tuple(output.shape), (7, 4))
 
+    def test_encoded_pointwise_mlp_uses_selected_coordinate_features(self) -> None:
+        feature_names = ("x", "y", "sdf", "u_inf", "aoa_deg")
+        model = EncodedPointwiseMLP(
+            feature_names=feature_names,
+            output_dim=4,
+            coordinate_features=("x", "y", "sdf"),
+            coordinate_encoding_spec=CoordinateEncodingSpec(type="fixed_fourier", scales=(1.0, 2.0)),
+            hidden_width=16,
+            depth=2,
+            activation="gelu",
+        )
+        batch = torch.randn(7, len(feature_names))
+
+        output = model(batch)
+
+        self.assertEqual(tuple(output.shape), (7, 4))
+
     def test_frontier_models_return_batch_by_target_dim(self) -> None:
         feature_names = ("x", "y", "sdf", "u_inf", "log_re", "aoa_deg")
         batch = torch.randn(8, len(feature_names))
diff --git a/tests/test_streaming_data.py b/tests/test_streaming_data.py
index 89485b0..c4b3f54 100644
--- a/tests/test_streaming_data.py
+++ b/tests/test_streaming_data.py
@@ -2,9 +2,11 @@ from __future__ import annotations
 
 import gzip
 import json
+import shutil
 import os
 import sys
 import tempfile
+import time
 import types
 import unittest
 import zipfile
@@ -124,6 +126,122 @@ interval_seconds = 0
         + "\n"
     )
 
+def write_huggingface_streaming_config(
+    path: Path,
+    *,
+    cache_dir: Path,
+    artifact_dir: Path,
+    train_cases: int = 4,
+    val_cases: int = 1,
+    test_cases: int = 1,
+    steps: int = 1,
+    normalization_cases: int = 1,
+) -> None:
+    path.write_text(
+        f"""
+[run]
+name = "hf_streaming_test"
+seed = 7
+artifact_dir = "{artifact_dir}"
+
+[data]
+root = "{cache_dir / 'processed' / 'full'}"
+source = "huggingface_streaming"
+hf_repo_id = "owner/airfrans-processed"
+hf_repo_type = "dataset"
+hf_path_prefix = "processed/full"
+cache_dir = "{cache_dir}"
+train_cases = {train_cases}
+val_cases = {val_cases}
+test_cases = {test_cases}
+all_points_per_case = true
+batch_size = 2
+streaming_queue_max_cases = 2
+streaming_normalization_cases = {normalization_cases}
+
+[model]
+type = "mlp"
+hidden_width = 16
+depth = 2
+activation = "gelu"
+
+[optim]
+lr = 0.01
+weight_decay = 0.0
+steps = {steps}
+log_interval = 1
+
+[device]
+type = "cpu"
+allow_cpu_fallback = false
+benchmark_kernels = false
+
+[loss]
+type = "normalized_mse"
+
+[checkpoint]
+interval_seconds = 0
+policy = "full"
+include_optimizer_state = true
+include_rng_state = true
+""".strip()
+        + "\n"
+    )
+
+
+def write_processed_case(root: Path, relative_prefix: str, case_id: str, offset: float) -> None:
+    target = root / relative_prefix / f"{case_id}.npz"
+    target.parent.mkdir(parents=True, exist_ok=True)
+    features = np.asarray(
+        [
+            [offset + 0.0, 0.0, 1.0, 0.1],
+            [offset + 1.0, 1.0, 0.5, 0.2],
+            [offset + 2.0, 0.5, 0.25, 0.3],
+            [offset + 3.0, 0.25, 0.125, 0.4],
+        ],
+        dtype=np.float32,
+    )
+    targets = np.asarray(
+        [
+            [offset + 0.0, 0.1, 0.2],
+            [offset + 0.2, 0.3, 0.4],
+            [offset + 0.4, 0.5, 0.6],
+            [offset + 0.6, 0.7, 0.8],
+        ],
+        dtype=np.float32,
+    )
+    np.savez(
+        target,
+        features=features,
+        targets=targets,
+        feature_names=np.asarray(["x", "y", "sdf", "alpha"], dtype="U"),
+        target_names=np.asarray(["u", "v", "p"], dtype="U"),
+    )
+
+
+def fake_huggingface_module(source_root: Path, calls: list[str]) -> types.ModuleType:
+    module = types.ModuleType("huggingface_hub")
+
+    class FakeHfApi:
+        def __init__(self, token: str | None = None) -> None:
+            self.token = token
+
+        def list_repo_files(self, *, repo_id: str, repo_type: str) -> list[str]:
+            return sorted(str(path.relative_to(source_root)) for path in source_root.rglob("*") if path.is_file())
+
+    def hf_hub_download(*, repo_id: str, filename: str, repo_type: str, local_dir: str, token: str | None = None) -> str:
+        calls.append(filename)
+        source = source_root / filename
+        destination = Path(local_dir) / filename
+        time.sleep(0.01)
+        destination.parent.mkdir(parents=True, exist_ok=True)
+        shutil.copy2(source, destination)
+        return str(destination)
+
+    module.HfApi = FakeHfApi
+    module.hf_hub_download = hf_hub_download
+    return module
+
 
 def read_events(run_dir: Path) -> list[dict[str, object]]:
     return [json.loads(line) for line in (run_dir / "streaming_events.jsonl").read_text().splitlines() if line.strip()]
@@ -215,6 +333,51 @@ class FullDataBackpressureStreamingTests(unittest.TestCase):
             normalization_end = next(event for event in events if event["event"] == "normalization_end")
             self.assertEqual(normalization_end["normalization_cases"], 1)
 
+    def test_huggingface_streaming_fast_start_downloads_remaining_selected_cases(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            tmp_path = Path(tmp)
+            hf_source = tmp_path / "hf_source"
+            case_names = [f"case_{index:04d}" for index in range(6)]
+            for index, case_name in enumerate(case_names):
+                write_processed_case(hf_source, "processed/full", case_name, float(index))
+            config_path = tmp_path / "hf_streaming.toml"
+            cache_dir = tmp_path / "cache"
+            artifact_dir = tmp_path / "artifacts"
+            write_huggingface_streaming_config(
+                config_path,
+                cache_dir=cache_dir,
+                artifact_dir=artifact_dir,
+                train_cases=4,
+                val_cases=1,
+                test_cases=1,
+                steps=1,
+                normalization_cases=1,
+            )
+            calls: list[str] = []
+            fake_module = fake_huggingface_module(hf_source, calls)
+
+            with patch.dict(sys.modules, {"huggingface_hub": fake_module}):
+                result = train(load_training_config(config_path))
+
+            self.assertEqual(result.final_metrics["data_mode"], "huggingface_streaming")
+            events = read_events(result.run_dir)
+            event_names = {event["event"] for event in events}
+            self.assertIn("background_acquisition_start", event_names)
+            self.assertIn("hf_case_download_end", event_names)
+            self.assertIn("background_acquisition_complete", event_names)
+            first_gpu = next(index for index, event in enumerate(events) if event["event"] == "first_gpu_batch_consumed")
+            processed_before_gpu = {
+                str(event["case_id"])
+                for event in events[:first_gpu]
+                if event["event"] == "processing_end"
+            }
+            self.assertLess(len(processed_before_gpu), 4)
+            manifest = json.loads((result.run_dir / "data_manifest.json").read_text())
+            self.assertEqual(manifest["source"], "huggingface_streaming")
+            cached_cases = sorted(path.stem for path in (cache_dir / "processed" / "full").glob("*.npz"))
+            self.assertEqual(cached_cases, case_names)
+            self.assertEqual(sorted(set(calls)), [f"processed/full/{case_name}.npz" for case_name in case_names])
+
     def test_backpressure_pauses_resumes_and_bounds_cache_with_inflight_slack(self) -> None:
         with tempfile.TemporaryDirectory() as tmp:
             tmp_path = Path(tmp)
diff --git a/tests/test_sweep.py b/tests/test_sweep.py
index 4a8aeaa..8aa0dd0 100644
--- a/tests/test_sweep.py
+++ b/tests/test_sweep.py
@@ -13,6 +13,7 @@ from airfrans_frontier.sweep import (
     collect_job_status,
     complete_job_attempt,
     generate_jobs,
+    generate_minimal_scaling_jobs,
     read_jobs,
     rescue_templates,
     run_node,
@@ -51,6 +52,41 @@ class SweepFoundationTests(unittest.TestCase):
             self.assertEqual(config.model.hidden_width, 2048)
             self.assertEqual(config.model.depth, 16)
 
+
+    def test_generate_minimal_scaling_jobs_are_all_points_and_non_raw_for_main_coordinate_models(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            jobs = generate_minimal_scaling_jobs(output_dir=root, include_proxies=False)
+
+            self.assertEqual(len(jobs), 23)
+            self.assertEqual([job["job_id"] for job in jobs[:3]], ["A_data_all_points_smoke", "A_checkpoint_700m_footprint_smoke", "A_plateau_interrupt_smoke"])
+            serious_raw = [
+                job["job_id"]
+                for job in jobs
+                if job["phase"] in {"B_family_viability", "C_scaling_ladder"}
+                and job["coordinate_encoding"] == "raw"
+                and job["model_family"] != "siren_conditioned_inr"
+            ]
+            self.assertEqual(serious_raw, [])
+            self.assertTrue(all(job["all_points_per_case"] for job in jobs))
+            c10 = load_training_config(root / "configs" / "C_10m_film_fourier_inr_nerf_multires_scale.toml")
+            c100 = load_training_config(root / "configs" / "C_100m_film_fourier_inr_nerf_multires_scale.toml")
+            config = load_training_config(root / "configs" / "C_700m_film_fourier_inr_nerf_multires_scale.toml")
+            self.assertTrue(config.data.all_points_per_case)
+            self.assertEqual(config.data.source, "huggingface_streaming")
+            self.assertEqual(config.data.hf_repo_id, "zacheryasc/airfrans-processed")
+            self.assertEqual(config.data.hf_path_prefix, "processed/full")
+            self.assertEqual(config.data.train_cases, 900)
+            self.assertEqual(config.data.val_cases, 50)
+            self.assertEqual(config.data.test_cases, 50)
+            self.assertEqual(config.data.streaming_normalization_cases, 4)
+            self.assertEqual(config.data.streaming_queue_max_cases, 8)
+            self.assertEqual(c10.run.seed, c100.run.seed)
+            self.assertEqual(c100.run.seed, config.run.seed)
+            self.assertEqual(config.checkpoint.policy, "full")
+            self.assertTrue(config.checkpoint.include_optimizer_state)
+            self.assertEqual(config.stability.dead_curve_patience_evals, 8)
+
     def test_generate_jobs_only_crosses_encoding_axis_for_compatible_families(self) -> None:
         with tempfile.TemporaryDirectory() as tmp:
             root = Path(tmp)
diff --git a/tests/test_training_config.py b/tests/test_training_config.py
index c610449..a92ee65 100644
--- a/tests/test_training_config.py
+++ b/tests/test_training_config.py
@@ -87,6 +87,129 @@ type = "normalized_mse"
             self.assertTrue(parsed.model_metadata.is_proxy)
             self.assertEqual(parsed.model_metadata.coordinate_encoding_compatibility, "raw_only")
 
+    def test_config_loader_accepts_all_points_dead_curve_and_lightweight_checkpoint(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            config_path = Path(tmp) / "all_points.toml"
+            config_path.write_text(
+                f"""
+[run]
+name = "all_points"
+seed = 0
+artifact_dir = "{Path(tmp) / "runs"}"
+
+[data]
+root = "{Path(tmp) / "data"}"
+train_cases = 1
+val_cases = 0
+test_cases = 0
+all_points_per_case = true
+batch_size = 1
+
+[coordinate_encoding]
+type = "nerf_multires"
+features = ["x", "y", "sdf"]
+levels = 4
+
+[model]
+type = "mlp_encoded_baseline"
+hidden_width = 8
+depth = 1
+activation = "gelu"
+coordinate_features = ["x", "y", "sdf"]
+
+[optim]
+lr = 0.001
+weight_decay = 0.0
+steps = 1
+
+[device]
+type = "cpu"
+allow_cpu_fallback = false
+benchmark_kernels = false
+
+[loss]
+type = "normalized_mse"
+
+[checkpoint]
+policy = "lightweight_scaling_probe"
+include_optimizer_state = false
+include_rng_state = false
+interval_seconds = 0
+
+[stability]
+dead_curve_patience_evals = 3
+dead_curve_min_relative_improvement = 0.01
+dead_curve_warmup_steps = 5
+""".strip()
+                + "\n"
+            )
+
+            parsed = load_training_config(config_path)
+
+        self.assertTrue(parsed.data.all_points_per_case)
+        self.assertIsNone(parsed.data.points_per_case)
+        self.assertEqual(parsed.model.type, "mlp_encoded_baseline")
+        self.assertEqual(parsed.checkpoint.policy, "lightweight_scaling_probe")
+        self.assertFalse(parsed.checkpoint.include_optimizer_state)
+        self.assertEqual(parsed.stability.dead_curve_patience_evals, 3)
+
+
+    def test_config_loader_accepts_huggingface_streaming_data_source(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            config_path = Path(tmp) / "hf_streaming.toml"
+            config_path.write_text(
+                f"""
+[run]
+name = "hf_streaming"
+seed = 0
+artifact_dir = "{Path(tmp) / "runs"}"
+
+[data]
+root = "{Path(tmp) / "cache" / "processed" / "full"}"
+source = "huggingface_streaming"
+hf_repo_id = "zacheryasc/airfrans-processed"
+hf_repo_type = "dataset"
+hf_path_prefix = "processed/full"
+cache_dir = "{Path(tmp) / "cache"}"
+train_cases = 1
+val_cases = 0
+test_cases = 0
+all_points_per_case = true
+batch_size = 1
+streaming_queue_max_cases = 8
+streaming_normalization_cases = 1
+
+[model]
+type = "mlp"
+hidden_width = 8
+depth = 1
+activation = "gelu"
+
+[optim]
+lr = 0.001
+weight_decay = 0.0
+steps = 1
+
+[device]
+type = "cpu"
+allow_cpu_fallback = false
+benchmark_kernels = false
+
+[loss]
+type = "normalized_mse"
+
+[checkpoint]
+interval_seconds = 0
+""".strip()
+                + "\n"
+            )
+
+            parsed = load_training_config(config_path)
+
+        self.assertEqual(parsed.data.source, "huggingface_streaming")
+        self.assertEqual(parsed.data.hf_repo_id, "zacheryasc/airfrans-processed")
+        self.assertTrue(parsed.data.all_points_per_case)
+
     def test_config_loader_rejects_missing_section(self) -> None:
         with tempfile.TemporaryDirectory() as tmp:
             config_path = Path(tmp) / "bad.toml"
diff --git a/tests/test_training_data.py b/tests/test_training_data.py
index d42df56..47145e7 100644
--- a/tests/test_training_data.py
+++ b/tests/test_training_data.py
@@ -10,7 +10,7 @@ remove_pythonpath_entries()
 
 import numpy as np
 
-from airfrans_frontier.training.data import create_case_split, load_processed_dataset, load_simulation_npz
+from airfrans_frontier.training.data import build_dataset_bundle, create_case_split, load_processed_dataset, load_simulation_npz
 from airfrans_frontier.training.normalize import compute_normalization_stats, normalize_targets
 
 
@@ -74,6 +74,25 @@ class TrainingDataTests(unittest.TestCase):
         self.assertEqual(len(first.val_ids), 2)
         self.assertEqual(len(first.test_ids), 2)
 
+    def test_all_points_mode_keeps_every_row_per_selected_case(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            write_case(root / "case_a.npz", offset=0.0)
+            write_case(root / "case_b.npz", offset=10.0)
+
+            bundle = build_dataset_bundle(
+                load_processed_dataset(root),
+                train_cases=1,
+                val_cases=1,
+                test_cases=0,
+                points_per_case=None,
+                seed=0,
+            )
+
+        self.assertEqual(bundle.train.features.shape[0], 3)
+        self.assertEqual(bundle.val.features.shape[0], 3)
+        self.assertEqual(bundle.train.point_counts, (3,))
+
     def test_normalization_uses_train_split_only(self) -> None:
         train_features = np.array([[0.0], [2.0]], dtype=np.float32)
         train_targets = np.array([[10.0], [14.0]], dtype=np.float32)
diff --git a/tests/test_training_loop.py b/tests/test_training_loop.py
index 350edc9..4dc10f5 100644
--- a/tests/test_training_loop.py
+++ b/tests/test_training_loop.py
@@ -67,6 +67,9 @@ def write_training_config(
     log_interval: int = 250,
     checkpoint_interval_seconds: int = 1800,
     hidden_width: int = 128,
+    checkpoint_include_optimizer_state: bool = True,
+    checkpoint_include_rng_state: bool = True,
+    dead_curve_patience_evals: int | None = None,
 ) -> None:
     path.write_text(
         f"""
@@ -105,6 +108,14 @@ type = "normalized_mse"
 
 [checkpoint]
 interval_seconds = {checkpoint_interval_seconds}
+include_optimizer_state = {str(checkpoint_include_optimizer_state).lower()}
+include_rng_state = {str(checkpoint_include_rng_state).lower()}
+{f'''
+[stability]
+dead_curve_patience_evals = {dead_curve_patience_evals}
+dead_curve_min_relative_improvement = 0.01
+dead_curve_warmup_steps = 0
+''' if dead_curve_patience_evals is not None else ''}
 """.strip()
         + "\n"
     )
@@ -421,6 +432,62 @@ class TrainingLoopTests(unittest.TestCase):
             self.assertTrue((run_dir / "artifact_manifest.json").is_file())
             self.assertTrue((run_dir / "checksums.txt").is_file())
 
+    def test_dead_curve_interruption_writes_plateau_failure(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            tmp_path = Path(tmp)
+            data_root = tmp_path / "data"
+            artifact_dir = tmp_path / "artifacts"
+            config_path = tmp_path / "config.toml"
+            write_toy_simulator_dataset(data_root, cases=4, points=64)
+            write_training_config(
+                config_path,
+                data_root=data_root,
+                artifact_dir=artifact_dir,
+                device_type="cpu",
+                steps=6,
+                log_interval=1,
+                checkpoint_interval_seconds=0,
+                hidden_width=32,
+                dead_curve_patience_evals=2,
+            )
+            config = load_training_config(config_path)
+            with patch("airfrans_frontier.training.loop.torch.optim.AdamW.step", lambda self: None):
+                with self.assertRaisesRegex(RuntimeError, "loss did not improve"):
+                    train(config)
+
+            run_dir = next(path for path in artifact_dir.iterdir() if path.is_dir())
+            report = json.loads((run_dir / "failure_report.json").read_text())
+            self.assertEqual(report["failure_category"], "plateau_or_dead_recipe")
+            self.assertEqual(report["error_type"], "DeadCurveError")
+            self.assertIn("dead_curve", report)
+
+    def test_lightweight_checkpoint_omits_optimizer_payload(self) -> None:
+        with tempfile.TemporaryDirectory() as tmp:
+            tmp_path = Path(tmp)
+            data_root = tmp_path / "data"
+            artifact_dir = tmp_path / "artifacts"
+            config_path = tmp_path / "config.toml"
+            write_toy_simulator_dataset(data_root, cases=4, points=64)
+            write_training_config(
+                config_path,
+                data_root=data_root,
+                artifact_dir=artifact_dir,
+                device_type="cpu",
+                steps=2,
+                log_interval=1,
+                checkpoint_interval_seconds=0,
+                hidden_width=32,
+                checkpoint_include_optimizer_state=False,
+                checkpoint_include_rng_state=False,
+            )
+
+            result = train(load_training_config(config_path))
+
+            checkpoint = torch.load(result.run_dir / "checkpoint_final.pt", map_location="cpu", weights_only=False)
+            self.assertIsNone(checkpoint["optimizer_state_dict"])
+            self.assertIsNone(checkpoint["rng_state"])
+            self.assertFalse(result.final_metrics["checkpoint_include_optimizer_state"])
+
     def test_nonfinite_gradient_writes_failure_report(self) -> None:
         with tempfile.TemporaryDirectory() as tmp:
             tmp_path = Path(tmp)