feat(training): let the camera BEV grid be set, so a consumer GPU can train

Model/tests/test_workflow_training_lifecycle.py

@@ -1688,3 +1688,46 @@ def test_resume_load_keeps_rng_tensors_on_cpu():
     keywords = {item.arg: item.value for item in resume_load.keywords}
     assert ast.literal_eval(keywords["map_location"]) == "cpu"
     assert ast.literal_eval(keywords["weights_only"]) is False
+
+
+def test_camera_bev_grid_override_changes_only_the_grid():
+    """The 6 GB workaround must not silently change what the BEV covers.
+
+    Shrinking the grid is safe only if pc_range is left alone: the BEV then
+    covers the same ground area with coarser cells, and the map BEV stays
+    aligned. If pc_range moved too, the model would be looking at a smaller
+    patch of the world and the metric would not be comparable.
+    """
+    from navigation.geometry import DEFAULT_NAVIGATION_GEOMETRY
+
+    base = DEFAULT_NAVIGATION_GEOMETRY.camera_bev_kwargs()
+    resized = workflows._camera_bev_kwargs_with_grid(base, 64)
+
+    assert resized["bev_h"] == 64
+    assert resized["bev_w"] == 64
+    changed = {k for k in base if base[k] != resized.get(k)}
+    assert changed == {"bev_h", "bev_w"}, (
+        f"resizing the grid must not touch anything else, but changed {changed}"
+    )
+    assert base["bev_h"] != 64, "fixture would be vacuous if the default were 64"
+    # The caller must not have its own dict mutated underneath it.
+    assert base["bev_h"] == DEFAULT_NAVIGATION_GEOMETRY.camera_bev_kwargs()["bev_h"]
+
+
+@pytest.mark.parametrize("bad", [0, -1, -256])
+def test_camera_bev_grid_override_rejects_non_positive(bad):
+    from navigation.geometry import DEFAULT_NAVIGATION_GEOMETRY
+
+    base = DEFAULT_NAVIGATION_GEOMETRY.camera_bev_kwargs()
+    with pytest.raises(ValueError, match="camera_bev_size must be positive"):
+        workflows._camera_bev_kwargs_with_grid(base, bad)
+
+
+def test_camera_bev_grid_defaults_to_the_geometry():
+    """With the parameter unset the run must be byte-identical to before."""
+    source = inspect.getsource(workflows.train_il.task_function)
+    assert "camera_bev_size: Optional[int] = None" in source
+    assert "if camera_bev_size is not None:" in source, (
+        "the override must be opt-in; an unconditional call would change the "
+        "default grid for every existing run"
+    )

Platform/pipelines/workflows.py

@@ -403,6 +403,27 @@ def _training_num_views_from_manifests(
     return max(dataset_views.values())
 
 
+def _camera_bev_kwargs_with_grid(
+    base: dict,
+    size: int,
+) -> dict:
+    """Return ``base`` with the camera BEV grid resized to ``size`` per side.
+
+    Only ``bev_h``/``bev_w`` change. ``pc_range`` is deliberately untouched, so
+    the BEV still covers the ground area the navigation geometry defines and the
+    map BEV stays aligned with it -- the grid is coarser, not smaller.
+
+    The KITScenes geometry pins the grid at 256, which allocates 65,536 BEV
+    queries and does not fit in 6 GB of VRAM.
+    """
+    if size < 1:
+        raise ValueError(f"camera_bev_size must be positive, got {size}")
+    resized = dict(base)
+    resized["bev_h"] = size
+    resized["bev_w"] = size
+    return resized
+
+
 def _training_source_revision(
     manifests: dict[str, dict],
     *,
@@ -3111,6 +3132,11 @@ def train_il(
     # GPU step. Effective parallelism is capped by shard count, so scale needs more
     # (smaller) shards too.
     num_workers: int = 0,
+    # Camera BEV grid, in cells per side. None keeps the value the KITScenes
+    # geometry defines (256), which allocates 65,536 queries and does not fit in
+    # 6 GB of VRAM. Only the grid changes: pc_range stays as the geometry
+    # defines it, so the BEV covers the same ground area with coarser cells.
+    camera_bev_size: Optional[int] = None,
     resume_from: Optional[FlyteFile] = None,
     early_stopping_patience: int = 5,
     allow_resume_policy_transition: bool = False,
@@ -3387,6 +3413,16 @@ def train_il(
         view_fusion_kwargs = (
             DEFAULT_NAVIGATION_GEOMETRY.camera_bev_kwargs()
         )
+        if camera_bev_size is not None:
+            view_fusion_kwargs = _camera_bev_kwargs_with_grid(
+                view_fusion_kwargs, camera_bev_size
+            )
+            print(
+                f"Camera BEV grid overridden: {camera_bev_size}x{camera_bev_size} "
+                f"({camera_bev_size ** 2} cells; geometry default is "
+                f"{DEFAULT_NAVIGATION_GEOMETRY.camera_bev_kwargs()['bev_h']}x"
+                f"{DEFAULT_NAVIGATION_GEOMETRY.camera_bev_kwargs()['bev_w']})"
+            )
 
     from Platform.pipelines.training_checkpoint import stable_digest
 
@@ -8698,6 +8734,7 @@ def wf_train_il(
     val_fraction: float = 0.1,
     validation_scope: str = "full",
     num_workers: int = 0,
+    camera_bev_size: Optional[int] = None,
     resume_from: Optional[FlyteFile] = None,
     early_stopping_patience: int = 5,
     allow_resume_policy_transition: bool = False,
@@ -8733,7 +8770,8 @@ def wf_train_il(
                    enable_reasoning=enable_reasoning, reasoning_mode=reasoning_mode,
                    enable_world_model=enable_world_model, val_fraction=val_fraction,
                    validation_scope=validation_scope,
-                   num_workers=num_workers, resume_from=resume_from,
+                   num_workers=num_workers, camera_bev_size=camera_bev_size,
+                   resume_from=resume_from,
                    early_stopping_patience=early_stopping_patience,
                    allow_resume_policy_transition=allow_resume_policy_transition)
     return evaluate_il_policy(checkpoint=out.checkpoint, shards=shards, dataset=dataset,