Support v1 raw multimodal image offload

docs/advanced.md

@@ -88,7 +88,7 @@ To add a new model family permanently, append an entry to `VLM_REGISTRY` in `src
 ### Limitations
 
 - **Vision encoder frozen by default.** Set `freeze_vision_encoder = false` to fine-tune it; in that case it's FSDP-sharded per block. The combination `freeze_vision_encoder = false` + LoRA is rejected by a config validator — LoRA freezes everything non-adapter, so unfreezing the encoder under LoRA would be a silent no-op.
-- **No multimodal-safe truncation.** Token sequences are truncated to `seq_len`, but `pixel_values` and `image_grid_thw` pass through unchanged. If a sample's tokens overflow, image tokens may get dropped while image tensors still describe the full image set. Set `seq_len` to cover your longest sample.
+- **No multimodal truncation.** Raw multimodal samples that exceed `seq_len` are rejected before packing, because truncating image placeholder spans would desynchronize tokens from the image descriptors. Set `seq_len` to cover your longest sample.
 - **bfloat16 mandatory.** The trainer config validator refuses any other `optimization_dtype` / `reduce_dtype` for VLMs — vLLM serves VLMs in bfloat16 and a mismatch breaks the importance ratio.
 - **Higher KL mismatch with multi-image inputs.** Expect noisier `mismatch_kl` than text-only; this is from minor numerical differences between the trainer's and vLLM's image processing.
 - **Images aren't logged to monitors.** Sample logging captures the prompt text but not the actual images.
@@ -146,4 +146,4 @@ salloc -N 1 --gres=gpu:1 bash -c 'bash scripts/install_nixl_from_source.sh'
 uv pip install --reinstall --no-deps deps/nixl_cu12-*.whl
 ```
 
-The script writes UCX 1.19 to `third_party/ucx/`; the bundled sbatch templates prepend it to `LD_LIBRARY_PATH` so it overrides the system version. Re-run both commands after every `uv sync`, since the lock pins the wheel.
+The script writes UCX 1.19 to `third_party/ucx/`; the bundled sbatch templates prepend it to `LD_LIBRARY_PATH` so it overrides the system version. Re-run both commands after every `uv sync`, since the lock pins the wheel.

docs/configuration.md

@@ -195,7 +195,7 @@ The `rl` launcher applies these the same way in both single-node and multi-node
 1. The launcher's own defaults — **your `env_vars` override these**.
 2. Your top-level `[env_vars]`.
 3. Your `[component.env_vars]`.
-4. Orchestration-critical vars the launcher always sets last — `CUDA_VISIBLE_DEVICES` (GPU partitioning) and `WANDB_SHARED_*` (the single shared W&B run) — **these cannot be overridden** from `env_vars`.
+4. Orchestration-critical vars the launcher always sets last — `CUDA_VISIBLE_DEVICES` (GPU partitioning), `WANDB_SHARED_*` (the single shared W&B run), and `VF_RENDERER_IMAGE_OFFLOAD_DIR` (raw multimodal image asset path) — **these cannot be overridden** from `env_vars`.
 
 For standalone `sft` and `inference` configs, `[env_vars]` applies to that entrypoint's process(es). For disaggregated P/D inference, the role-specific [`deployment.{prefill,decode}_env_vars`](inference.md) layer on top of any shared inference env vars.
 

packages/prime-rl-configs/src/prime_rl/configs/inference.py

@@ -5,7 +5,7 @@
 from pydantic import Field, model_validator
 from pydantic_config import BaseConfig
 
-from prime_rl.configs.shared import BaseModelConfig, EnvVars, LogConfig, SlurmConfig
+from prime_rl.configs.shared import BaseModelConfig, EnvVars, LogConfig, MultimodalConfig, SlurmConfig
 from prime_rl.utils.config import find_package_resource, rgetattr, rsetattr
 from prime_rl.utils.parsers import resolve_reasoning_parser, resolve_tool_call_parser
 
@@ -438,6 +438,9 @@ class InferenceConfig(BaseConfig):
 
     experimental: InferenceExperimentalConfig = InferenceExperimentalConfig()
 
+    multimodal: MultimodalConfig = MultimodalConfig()
+    """Raw multimodal image offload settings shared with trainer and orchestrator."""
+
     @model_validator(mode="after")
     def validate_multi_node_requires_slurm(self):
         if self.deployment.type in ("multi_node", "disaggregated") and self.slurm is None:

packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py

@@ -17,6 +17,7 @@
     FileSystemTransportConfig,
     HeartbeatConfig,
     LogConfig,
+    MultimodalConfig,
     PrimeMonitorConfig,
     TransportConfig,
     WandbWithExtrasConfig,
@@ -565,6 +566,9 @@ class OrchestratorConfig(BaseConfig):
 
     experimental: OrchestratorExperimentalConfig = OrchestratorExperimentalConfig()
 
+    multimodal: MultimodalConfig = MultimodalConfig()
+    """Raw multimodal image offload settings shared with trainer and inference."""
+
     @model_validator(mode="before")
     @classmethod
     def _env_to_train(cls, data: Any) -> Any:

packages/prime-rl-configs/src/prime_rl/configs/rl.py

@@ -17,6 +17,7 @@
 )
 from prime_rl.configs.shared import (
     EnvVars,
+    MultimodalConfig,
     SlurmConfig,
     VLMConfig,
 )
@@ -223,6 +224,9 @@ class RLConfig(BaseConfig):
 
     weight_broadcast: SharedWeightBroadcastConfig | None = None
 
+    multimodal: MultimodalConfig = MultimodalConfig()
+    """Shared raw multimodal image offload settings. Propagated to trainer, orchestrator, and inference."""
+
     bench: bool = False
     """Benchmark mode. Sets trainer and orchestrator to benchmark mode and, when set, suffixes the W&B project with ``-bench``."""
 

packages/prime-rl-configs/src/prime_rl/configs/shared.py

@@ -10,7 +10,13 @@
 # and the single shared W&B run. The launcher always sets these last, so allowing them in
 # `env_vars` would be a silent no-op (or, on multi-node, a footgun) — reject them instead.
 PROTECTED_ENV_VARS = frozenset(
-    {"CUDA_VISIBLE_DEVICES", "WANDB_SHARED_MODE", "WANDB_SHARED_RUN_ID", "WANDB_SHARED_LABEL"}
+    {
+        "CUDA_VISIBLE_DEVICES",
+        "VF_RENDERER_IMAGE_OFFLOAD_DIR",
+        "WANDB_SHARED_MODE",
+        "WANDB_SHARED_RUN_ID",
+        "WANDB_SHARED_LABEL",
+    }
 )
 
 
@@ -82,6 +88,11 @@ def resolve_project_dir(self):
 ServerType = Literal["vllm", "openai"]
 
 
+class MultimodalConfig(BaseConfig):
+    offload_dir: Path | None = None
+    """Directory for offloaded image assets. Supports environment expansion such as ``/data/outputs/run_${RUN_ID}/assets/images``. When unset, prime-rl resolves a run-scoped default."""
+
+
 class VLMConfig(BaseConfig):
     vision_encoder_attr: str
     """Dotted attribute path to the vision encoder module (e.g. ``model.visual``)."""

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -10,6 +10,7 @@
     FileSystemTransportConfig,
     HeartbeatConfig,
     MetricsServerConfig,
+    MultimodalConfig,
     TrainerLogConfig,
     TransportConfig,
     WandbConfig,
@@ -20,6 +21,7 @@
 
 AttnImplementation: TypeAlias = Literal["eager", "sdpa", "flash_attention_2", "flash_attention_3", "fa4"]
 EPCommBackend: TypeAlias = Literal["torch", "deepep"]
+MissingMMImagePolicy: TypeAlias = Literal["error", "placeholder_zero_loss"]
 
 # User-facing name -> internal name. Users set `flash_attention_4` in configs,
 # which gets rewritten to `fa4` before pydantic validation.
@@ -566,6 +568,12 @@ class TrainerConfig(BaseConfig):
     max_concurrent_runs: int = Field(1, ge=1)
     """Maximum number of concurrent runs to allow. If 1, only one run may run at a time."""
 
+    missing_mm_image_policy: MissingMMImagePolicy = "placeholder_zero_loss"
+    """Policy when raw multimodal image files disappear before trainer materialization. ``placeholder_zero_loss`` warns, synthesizes zero-valued image tensors with the original descriptor geometry, and masks out the affected microbatch loss; ``error`` preserves fail-fast behavior."""
+
+    multimodal: MultimodalConfig = MultimodalConfig()
+    """Raw multimodal image offload settings shared with orchestrator and inference."""
+
     enable_token_export: bool = False
     """Opt-in per-token JSONL export for rollout debugging. When enabled, writes token ids and aligned trainer metrics after each forward pass."""
 

packages/prime-rl-configs/src/prime_rl/utils/validation.py

@@ -126,6 +126,7 @@ def propagate(shared_path: str, *targets: str) -> None:
     # Top-level scalars.
     propagate("max_steps", "trainer.max_steps", "orchestrator.max_steps")
     propagate("seq_len", "trainer.model.seq_len", "orchestrator.seq_len")
+    propagate("multimodal", "trainer.multimodal", "orchestrator.multimodal", "inference.multimodal")
 
     # [slurm] → inference: a multi-node RL run drives its inference deployment under
     # the same SLURM allocation, so the nested inference inherits [slurm]. This is

src/prime_rl/entrypoints/orchestrator.py

@@ -14,11 +14,13 @@
 from prime_rl.configs.orchestrator import OrchestratorConfig
 from prime_rl.utils.config import cli
 from prime_rl.utils.process import set_proc_title
+from prime_rl.utils.run_assets import apply_run_asset_env
 
 
 def main():
     set_proc_title("Orchestrator")
     config = cli(OrchestratorConfig)
+    apply_run_asset_env(config.output_dir, config.multimodal)
     from prime_rl.orchestrator.orchestrator import run_orchestrator
 
     asyncio.run(run_orchestrator(config))

src/prime_rl/entrypoints/rl.py

@@ -35,6 +35,7 @@
     monitor_process,
     set_proc_title,
 )
+from prime_rl.utils.run_assets import build_run_asset_env
 
 RL_TOML = "rl.toml"
 RL_SBATCH = "rl.sbatch"
@@ -125,6 +126,8 @@ def rl_local(config: RLConfig):
         "WANDB_SHARED_MODE": "1",
         "WANDB_SHARED_RUN_ID": os.environ.get("WANDB_SHARED_RUN_ID", uuid.uuid4().hex),
     }
+    inherited_env = dict(os.environ)
+    writer_run_asset_env = build_run_asset_env(config.orchestrator.output_dir, multimodal=config.multimodal)
 
     # Validate client port matches inference server port
     if config.inference is not None and not config.orchestrator.model.client.is_elastic:
@@ -170,7 +173,7 @@ def sigterm_handler(signum, frame):
                 inference_process = Popen(
                     inference_cmd,
                     env={
-                        **os.environ,
+                        **inherited_env,
                         **DEFAULT_COMMON_ENV_VARS,
                         **DEFAULT_INFERENCE_ENV_VARS,
                         **config.env_vars,
@@ -225,7 +228,7 @@ def sigterm_handler(signum, frame):
                 stdout=log_file,
                 stderr=log_file,
                 env={
-                    **os.environ,
+                    **writer_run_asset_env,
                     **DEFAULT_COMMON_ENV_VARS,
                     "LOGURU_FORCE_COLORS": "1",
                     "WANDB_PROGRAM": "uv run rl",
@@ -274,7 +277,7 @@ def sigterm_handler(signum, frame):
             trainer_process = Popen(
                 trainer_cmd,
                 env={
-                    **os.environ,
+                    **inherited_env,
                     **DEFAULT_COMMON_ENV_VARS,
                     **DEFAULT_TRAINER_ENV_VARS,
                     "LOGURU_FORCE_COLORS": "1",
@@ -372,6 +375,9 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) ->
         kv_offload_disk_path=str(offload.disk.path) if (is_mooncake and offload.disk is not None) else "",
         kv_offload_device_name=offload.device_name if is_mooncake else "",
     )
+    image_offload_dir = (
+        os.path.expanduser(str(config.multimodal.offload_dir)) if config.multimodal.offload_dir is not None else ""
+    )
 
     # Per-component env vars: launcher defaults (shared + multi-node-specific) with the
     # user's config merged on top. Runtime wiring stays in the template.
@@ -393,6 +399,7 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) ->
             **config.slurm.template_vars,
             config_path=config_dir / RL_TOML,
             output_dir=config.output_dir,
+            image_offload_dir=image_offload_dir,
             gpus_per_node=config.deployment.gpus_per_node,
         )
     elif config.inference is not None and config.inference.deployment.type == "disaggregated":
@@ -404,6 +411,7 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) ->
             config_dir=config_dir,
             output_dir=config.output_dir,
             orchestrator_output_dir=config.orchestrator.output_dir,
+            image_offload_dir=image_offload_dir,
             num_train_nodes=config.deployment.num_train_nodes,
             num_infer_nodes=infer_deploy.num_nodes * config.deployment.num_infer_replicas,
             nodes_per_infer_replica=infer_deploy.num_nodes,
@@ -439,6 +447,7 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) ->
             config_dir=config_dir,  # TODO: should prob have each subconfig path separately
             output_dir=config.output_dir,
             orchestrator_output_dir=config.orchestrator.output_dir,
+            image_offload_dir=image_offload_dir,
             num_train_nodes=config.deployment.num_train_nodes,
             num_infer_nodes=config.deployment.total_infer_nodes,
             nodes_per_infer_replica=config.deployment.num_infer_nodes,