From 678f84229ebfe2e154abf34b4069095dd1e1efb9 Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Sat, 1 Aug 2026 04:47:22 +0000
Subject: [PATCH 01/13] feat: add TAO WTS fine-tuning support

---
 Dockerfile                                    |   8 +-
 cosmos_framework/callbacks/tao_status.py      | 381 ++++++++++++++++++
 cosmos_framework/callbacks/tao_status_test.py |  86 ++++
 cosmos_framework/checkpoint/base.py           |   5 +-
 cosmos_framework/checkpoint/dcp.py            |   9 +-
 cosmos_framework/checkpoint/dcp_distill.py    |  13 +-
 cosmos_framework/checkpoint/dummy.py          |   1 +
 .../base/reasoner/defaults/callbacks.py       |  17 +-
 .../base/reasoner/experiment/wts_vlm.py       | 294 ++++++++++++++
 .../base/reasoner/experiment/wts_vlm_test.py  |  56 +++
 .../configs/toml_config/sft_config.py         | 103 +++--
 .../configs/toml_config/sft_config_test.py    |  18 +
 cosmos_framework/scripts/train.py             |  26 +-
 cosmos_framework/trainer/__init__.py          |  83 +++-
 cosmos_framework/utils/callback.py            |   9 +-
 cosmos_framework/utils/checkpointer.py        |  10 +-
 cosmos_framework/utils/config.py              |  10 +
 .../utils/reasoner/dcp_checkpointer.py        |   7 +-
 examples/launch_sft_wts.sh                    |  32 ++
 examples/toml/sft_config/wts_vlm.toml         |  99 +++++
 examples/toml/sft_config/wts_vlm_edge.toml    |  95 +++++
 uv.lock                                       |  10 +-
 22 files changed, 1292 insertions(+), 80 deletions(-)
 create mode 100644 cosmos_framework/callbacks/tao_status.py
 create mode 100644 cosmos_framework/callbacks/tao_status_test.py
 create mode 100644 cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py
 create mode 100644 cosmos_framework/configs/base/reasoner/experiment/wts_vlm_test.py
 create mode 100755 examples/launch_sft_wts.sh
 create mode 100644 examples/toml/sft_config/wts_vlm.toml
 create mode 100644 examples/toml/sft_config/wts_vlm_edge.toml

diff --git a/Dockerfile b/Dockerfile
index 529ddb2c..ea474be5 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -49,9 +49,15 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,source=pyproject.toml,target=pyproject.toml \
     --mount=type=bind,source=.python-version,target=.python-version \
     --mount=type=bind,source=packages,target=packages \
-    uv sync --locked --no-install-project --no-editable --all-extras --group=$(cat /root/.cuda-name) --group=vllm
+    uv sync --locked --no-install-project --no-editable --all-extras --group=$(cat /root/.cuda-name)-train
 ENV PATH="/workspace/.venv/bin:$PATH"
 
+# Package the exact source state into the image so it can run on managed
+# platforms without a host bind mount.
+COPY . /workspace
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --no-deps .
+
 # Triton bundled ptxas doesn't support latest GPU architectures
 ENV TRITON_PTXAS_PATH="/usr/local/cuda/bin/ptxas"
 
diff --git a/cosmos_framework/callbacks/tao_status.py b/cosmos_framework/callbacks/tao_status.py
new file mode 100644
index 00000000..d318d076
--- /dev/null
+++ b/cosmos_framework/callbacks/tao_status.py
@@ -0,0 +1,381 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""TAO-compatible lifecycle and metric logging for Cosmos Framework training."""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+from datetime import datetime, timedelta
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.distributed as dist
+
+from cosmos_framework.utils import distributed, log, misc
+from cosmos_framework.utils.callback import Callback
+
+
+def _to_json_value(value: Any) -> Any:
+    """Recursively convert tensors and array-like values to JSON-safe values."""
+    if isinstance(value, dict):
+        return {str(key): _to_json_value(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_to_json_value(item) for item in value]
+    if isinstance(value, torch.Tensor):
+        value = value.detach()
+        if value.numel() == 1:
+            return value.item()
+        return value.cpu().tolist()
+    if hasattr(value, "item"):
+        try:
+            return value.item()
+        except (TypeError, ValueError):
+            pass
+    if hasattr(value, "tolist"):
+        try:
+            return value.tolist()
+        except (TypeError, ValueError):
+            pass
+    return value
+
+
+class _TAOStatusWriter:
+    """Use TAO Core when available, with a compatible JSON-lines fallback."""
+
+    def __init__(self, filename: str) -> None:
+        self.filename = filename
+        Path(filename).parent.mkdir(parents=True, exist_ok=True)
+        self._tao_logger = None
+        self._tao_status = None
+        self._tao_verbosity = None
+
+        try:
+            from nvidia_tao_core.loggers.logging import Status, StatusLogger, Verbosity
+        except ImportError:
+            log.warning(f"nvidia_tao_core is not installed; writing TAO-compatible JSON records directly to {filename}")
+        else:
+            self._tao_status = Status
+            self._tao_verbosity = Verbosity
+            self._tao_logger = StatusLogger(
+                filename=filename,
+                is_master=True,
+                verbosity=Verbosity.INFO,
+                append=True,
+            )
+
+    def write(
+        self,
+        *,
+        status: str,
+        message: str,
+        data: dict[str, Any] | None = None,
+        kpi: dict[str, Any] | None = None,
+        verbosity: str = "INFO",
+    ) -> None:
+        data = _to_json_value(data or {})
+        kpi = _to_json_value(kpi or {})
+
+        if self._tao_logger is not None:
+            status_level = getattr(self._tao_status, status)
+            verbosity_level = getattr(self._tao_verbosity, verbosity)
+            self._tao_logger.kpi = kpi
+            self._tao_logger.write(
+                data=data,
+                status_level=status_level,
+                verbosity_level=verbosity_level,
+                message=message,
+            )
+            return
+
+        now = datetime.now()
+        payload: dict[str, Any] = {
+            **data,
+            "date": f"{now.month}/{now.day}/{now.year}",
+            "time": f"{now.hour}:{now.minute}:{now.second}",
+            "status": status,
+            "verbosity": verbosity,
+            "message": message,
+        }
+        if kpi:
+            payload["kpi"] = kpi
+        with open(self.filename, "a", encoding="utf-8") as status_file:
+            status_file.write(json.dumps(payload, default=str) + "\n")
+
+
+class TAOStatusCallback(Callback):
+    """Write TAO lifecycle, training, and validation records from rank zero.
+
+    The output path is resolved in this order:
+
+    1. ``status_file_path`` when explicitly configured.
+    2. ``$TAO_RESULTS_ROOT/$TAO_JOB_ID/status.json`` (TAO SDK).
+    3. ``$TAO_API_RESULTS_DIR/$TAO_API_JOB_ID/status.json`` (TAO API).
+    4. ``<job.path_local>/status.json`` for direct launches.
+    """
+
+    def __init__(
+        self,
+        enabled: bool = False,
+        status_file_path: str | None = None,
+        experiment_name: str = "",
+        logging_interval: int = 1,
+        validation_heartbeat_interval: int = 1,
+    ) -> None:
+        if logging_interval < 1:
+            raise ValueError("logging_interval must be >= 1")
+        if validation_heartbeat_interval < 1:
+            raise ValueError("validation_heartbeat_interval must be >= 1")
+        self.enabled = enabled
+        self.status_file_path = status_file_path
+        self.experiment_name = experiment_name
+        self.logging_interval = logging_interval
+        self.validation_heartbeat_interval = validation_heartbeat_interval
+        self._writer: _TAOStatusWriter | None = None
+        self._train_start_time = 0.0
+        self._step_start_time = 0.0
+        self._validation_batches = 0
+        self._validation_loss_sum = 0.0
+        self._validation_sample_count = 0
+        self.last_validation_loss: float | None = None
+
+    def _is_rank_zero(self) -> bool:
+        if dist.is_available() and dist.is_initialized():
+            return distributed.is_rank0()
+        return int(os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0"))) == 0
+
+    def _component_name(self) -> str:
+        if self.experiment_name:
+            return self.experiment_name
+        return getattr(self.config.job, "name", "Cosmos Framework SFT")
+
+    def _resolve_status_file(self) -> str:
+        if self.status_file_path:
+            return self.status_file_path
+
+        job_id = os.environ.get("TAO_JOB_ID")
+        if job_id:
+            results_root = os.environ.get("TAO_RESULTS_ROOT", "/results")
+            return os.path.join(results_root, job_id, "status.json")
+
+        api_job_id = os.environ.get("TAO_API_JOB_ID")
+        if api_job_id:
+            results_root = os.environ.get("TAO_API_RESULTS_DIR", "/results")
+            return os.path.join(results_root, api_job_id, "status.json")
+
+        return os.path.join(self.config.job.path_local, "status.json")
+
+    def _get_writer(self) -> _TAOStatusWriter | None:
+        if not self.enabled or not self._is_rank_zero():
+            return None
+        if self._writer is None:
+            self._writer = _TAOStatusWriter(self._resolve_status_file())
+        return self._writer
+
+    def _progress_data(self, iteration: int, seconds_per_step: float | None = None) -> dict[str, Any]:
+        trainer = getattr(self, "trainer", None)
+        max_step = int(getattr(trainer, "max_iterations", self.config.trainer.max_iter))
+        if seconds_per_step is None:
+            elapsed = max(time.monotonic() - self._train_start_time, 0.0)
+            seconds_per_step = elapsed / max(iteration, 1)
+        eta_seconds = max(max_step - iteration, 0) * seconds_per_step
+        data = {
+            "component": self._component_name(),
+            "step": iteration,
+            "max_step": max_step,
+            "time_per_step": str(timedelta(seconds=seconds_per_step)),
+            "eta": str(timedelta(seconds=eta_seconds)),
+        }
+        steps_per_epoch = getattr(trainer, "steps_per_epoch", None)
+        num_epochs = getattr(trainer, "num_epochs", None)
+        if steps_per_epoch and num_epochs:
+            completed_epochs = iteration // steps_per_epoch
+            if iteration == 0:
+                epoch = 1
+                step_in_epoch = 0
+            elif iteration % steps_per_epoch == 0:
+                epoch = min(completed_epochs, num_epochs)
+                step_in_epoch = steps_per_epoch
+            else:
+                epoch = min(completed_epochs + 1, num_epochs)
+                step_in_epoch = iteration % steps_per_epoch
+            data.update(
+                {
+                    "epoch": epoch,
+                    "max_epoch": num_epochs,
+                    "completed_epochs": min(completed_epochs, num_epochs),
+                    "step_in_epoch": step_in_epoch,
+                    "steps_per_epoch": steps_per_epoch,
+                    "time_per_epoch": str(timedelta(seconds=seconds_per_step * steps_per_epoch)),
+                }
+            )
+        return data
+
+    def _epoch_label(self, iteration: int) -> str:
+        progress = self._progress_data(iteration, seconds_per_step=0.0)
+        if "epoch" not in progress:
+            return f"training step {iteration}/{progress['max_step']}"
+        return f"epoch {progress['epoch']}/{progress['max_epoch']}"
+
+    @staticmethod
+    def _global_average_loss(loss: torch.Tensor, data_batch: dict[str, Any]) -> tuple[float, int]:
+        sample_count = misc.get_data_batch_size(data_batch)
+        loss_sum = loss.detach() * sample_count
+        count = torch.tensor(sample_count, device=loss.device, dtype=torch.long)
+        if dist.is_available() and dist.is_initialized():
+            dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM)
+            dist.all_reduce(count, op=dist.ReduceOp.SUM)
+        global_count = int(count.item())
+        average = float(loss_sum.item()) / max(global_count, 1)
+        return average, global_count
+
+    def on_train_start(self, model: Any, iteration: int = 0) -> None:
+        self._train_start_time = time.monotonic()
+        writer = self._get_writer()
+        if writer is not None:
+            writer.write(
+                status="STARTED",
+                message=f"Starting {self._component_name()} training",
+                data=self._progress_data(iteration, seconds_per_step=0.0),
+            )
+            log.info(f"TAO status will be logged to {writer.filename}")
+
+    def on_training_step_start(self, model: Any, data: dict[str, Any], iteration: int = 0) -> None:
+        self._step_start_time = time.monotonic()
+
+    def on_training_step_end(
+        self,
+        model: Any,
+        data_batch: dict[str, Any],
+        output_batch: dict[str, Any],
+        loss: torch.Tensor,
+        iteration: int = 0,
+    ) -> None:
+        interval = int(self.config.trainer.logging_iter) * self.logging_interval
+        if iteration % interval != 0:
+            return
+
+        average_loss, _ = self._global_average_loss(loss, data_batch)
+        writer = self._get_writer()
+        if writer is None:
+            return
+        seconds_per_step = max(time.monotonic() - self._step_start_time, 0.0)
+        kpi = {"train/loss": average_loss, "train/loss_avg": average_loss}
+        progress = self._progress_data(iteration, seconds_per_step=seconds_per_step)
+        if "epoch" in progress:
+            message = (
+                f"Training epoch {progress['epoch']}/{progress['max_epoch']}, "
+                f"step {progress['step_in_epoch']}/{progress['steps_per_epoch']} "
+                f"(global step {iteration}/{progress['max_step']}) - Loss: {average_loss:.6f}"
+            )
+        else:
+            message = f"Training step {iteration}/{progress['max_step']} - Loss: {average_loss:.6f}"
+        writer.write(
+            status="RUNNING",
+            message=message,
+            data=progress,
+            kpi=kpi,
+        )
+
+    def on_validation_start(self, model: Any, dataloader_val: Any, iteration: int = 0) -> None:
+        self._validation_batches = 0
+        self._validation_loss_sum = 0.0
+        self._validation_sample_count = 0
+        writer = self._get_writer()
+        if writer is not None:
+            writer.write(
+                status="RUNNING",
+                message=f"Starting validation for {self._epoch_label(iteration)}",
+                data={**self._progress_data(iteration), "phase": "validation_starting"},
+            )
+            log.info(f"Starting validation for {self._epoch_label(iteration)}")
+
+    def on_validation_step_end(
+        self,
+        model: Any,
+        data_batch: dict[str, Any],
+        output_batch: dict[str, Any],
+        loss: torch.Tensor,
+        iteration: int = 0,
+    ) -> None:
+        average_loss, sample_count = self._global_average_loss(loss, data_batch)
+        self._validation_batches += 1
+        self._validation_loss_sum += average_loss * sample_count
+        self._validation_sample_count += sample_count
+
+        if self._validation_batches % self.validation_heartbeat_interval != 0:
+            return
+        writer = self._get_writer()
+        max_validation_batches = getattr(self.config.trainer, "max_val_iter", None)
+        batch_progress = (
+            f"{self._validation_batches}/{max_validation_batches}"
+            if max_validation_batches is not None
+            else str(self._validation_batches)
+        )
+        if writer is not None:
+            writer.write(
+                status="RUNNING",
+                message=(
+                    f"Validation {self._epoch_label(iteration)}, batch {batch_progress} - Loss: {average_loss:.6f}"
+                ),
+                data={
+                    **self._progress_data(iteration),
+                    "phase": "validation_batch_complete",
+                    "validation_batch": self._validation_batches,
+                    "max_validation_batches": max_validation_batches,
+                },
+                kpi={"val/batch_loss": average_loss},
+            )
+            log.info(f"Validation {self._epoch_label(iteration)}, batch {batch_progress} - Loss: {average_loss:.6f}")
+
+    def on_validation_end(self, model: Any, iteration: int = 0) -> None:
+        if self._validation_sample_count == 0:
+            log.warning("TAO validation logging saw zero samples; no val/loss record was written")
+            return
+
+        self.last_validation_loss = self._validation_loss_sum / self._validation_sample_count
+        writer = self._get_writer()
+        if writer is not None:
+            writer.write(
+                status="RUNNING",
+                message=(
+                    f"Validation complete for {self._epoch_label(iteration)} - Loss: {self.last_validation_loss:.6f}"
+                ),
+                data={
+                    **self._progress_data(iteration),
+                    "phase": "validation_complete",
+                    "validation_batches": self._validation_batches,
+                    "validation_samples": self._validation_sample_count,
+                },
+                kpi={"val/loss": self.last_validation_loss, "val/avg_loss": self.last_validation_loss},
+            )
+        log.info(f"Validation loss ({self._epoch_label(iteration)}): {self.last_validation_loss:.6f}")
+
+    def on_app_end(self) -> None:
+        writer = self._get_writer()
+        if writer is not None:
+            writer.write(
+                status="SUCCESS",
+                message=f"{self._component_name()} training completed successfully",
+                data=self._progress_data(
+                    int(getattr(getattr(self, "trainer", None), "max_iterations", self.config.trainer.max_iter))
+                ),
+                kpi=(
+                    {"val/loss": self.last_validation_loss, "val/avg_loss": self.last_validation_loss}
+                    if self.last_validation_loss is not None
+                    else None
+                ),
+            )
+
+    def on_exception(self, error: BaseException) -> None:
+        writer = self._get_writer()
+        if writer is not None:
+            writer.write(
+                status="FAILURE",
+                verbosity="ERROR",
+                message=f"{self._component_name()} training failed: {error}",
+                data={"component": self._component_name(), "error_type": type(error).__name__},
+            )
diff --git a/cosmos_framework/callbacks/tao_status_test.py b/cosmos_framework/callbacks/tao_status_test.py
new file mode 100644
index 00000000..f862487c
--- /dev/null
+++ b/cosmos_framework/callbacks/tao_status_test.py
@@ -0,0 +1,86 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+from __future__ import annotations
+
+import json
+from types import SimpleNamespace
+
+import torch
+
+from cosmos_framework.callbacks.tao_status import TAOStatusCallback
+
+
+def _callback(tmp_path) -> TAOStatusCallback:
+    callback = TAOStatusCallback(
+        enabled=True,
+        status_file_path=str(tmp_path / "status.json"),
+        experiment_name="test",
+    )
+    callback.config = SimpleNamespace(
+        job=SimpleNamespace(name="job", path_local=str(tmp_path)),
+        trainer=SimpleNamespace(max_iter=6, max_val_iter=2, logging_iter=1),
+    )
+    callback.trainer = SimpleNamespace(max_iterations=6, num_epochs=2, steps_per_epoch=3)
+    return callback
+
+
+def _records(tmp_path) -> list[dict]:
+    return [json.loads(line) for line in (tmp_path / "status.json").read_text(encoding="utf-8").splitlines()]
+
+
+def test_tao_status_callback_writes_training_validation_and_success(tmp_path) -> None:
+    callback = _callback(tmp_path)
+    callback.on_train_start(model=None, iteration=0)
+    callback.on_training_step_start(model=None, data={}, iteration=0)
+    callback.on_training_step_end(
+        model=None,
+        data_batch={"input_ids": torch.zeros(2, 3)},
+        output_batch={},
+        loss=torch.tensor(0.5),
+        iteration=3,
+    )
+    callback.on_validation_start(model=None, dataloader_val=None, iteration=3)
+    callback.on_validation_step_end(
+        model=None,
+        data_batch={"input_ids": torch.zeros(2, 3)},
+        output_batch={},
+        loss=torch.tensor(0.25),
+        iteration=3,
+    )
+    callback.on_validation_end(model=None, iteration=3)
+    callback.on_app_end()
+
+    records = _records(tmp_path)
+    assert [record["status"] for record in records] == [
+        "STARTED",
+        "RUNNING",
+        "RUNNING",
+        "RUNNING",
+        "RUNNING",
+        "SUCCESS",
+    ]
+    assert records[1]["kpi"]["train/loss_avg"] == 0.5
+    assert records[1]["epoch"] == 1
+    assert records[1]["max_epoch"] == 2
+    assert records[1]["step_in_epoch"] == 3
+    assert records[1]["steps_per_epoch"] == 3
+    assert records[1]["max_step"] == 6
+    assert records[2]["phase"] == "validation_starting"
+    assert records[3]["max_validation_batches"] == 2
+    assert "epoch 1/2" in records[3]["message"]
+    assert records[4]["kpi"]["val/avg_loss"] == 0.25
+    assert "epoch 1/2" in records[4]["message"]
+    assert records[-1]["kpi"]["val/loss"] == 0.25
+    assert records[-1]["epoch"] == 2
+    assert records[-1]["completed_epochs"] == 2
+
+
+def test_tao_status_callback_writes_failure(tmp_path) -> None:
+    callback = _callback(tmp_path)
+    callback.on_train_start(model=None, iteration=0)
+    callback.on_exception(RuntimeError("boom"))
+
+    records = _records(tmp_path)
+    assert records[-1]["status"] == "FAILURE"
+    assert records[-1]["error_type"] == "RuntimeError"
diff --git a/cosmos_framework/checkpoint/base.py b/cosmos_framework/checkpoint/base.py
index f17ebd2d..e22f2360 100644
--- a/cosmos_framework/checkpoint/base.py
+++ b/cosmos_framework/checkpoint/base.py
@@ -7,11 +7,11 @@
 
 import torch
 
-from cosmos_framework.utils.config import CheckpointConfig, JobConfig
-from cosmos_framework.utils.flags import INTERNAL
 from cosmos_framework.model._base import ImaginaireModel
 from cosmos_framework.utils import callback
+from cosmos_framework.utils.config import CheckpointConfig, JobConfig
 from cosmos_framework.utils.easy_io import easy_io
+from cosmos_framework.utils.flags import INTERNAL
 
 
 class AbstractCheckpointer(ABC):
@@ -92,6 +92,7 @@ def save(
         scheduler: torch.optim.lr_scheduler.LRScheduler,
         grad_scaler: torch.amp.GradScaler,
         iteration: int,
+        epoch: int | None = None,
     ) -> None:
         pass
 
diff --git a/cosmos_framework/checkpoint/dcp.py b/cosmos_framework/checkpoint/dcp.py
index ba5f996a..a4ab9398 100644
--- a/cosmos_framework/checkpoint/dcp.py
+++ b/cosmos_framework/checkpoint/dcp.py
@@ -71,9 +71,9 @@
 
 from cosmos_framework.checkpoint.base import AbstractCheckpointer
 from cosmos_framework.checkpoint.s3_filesystem import S3StorageReader, S3StorageWriter
-from cosmos_framework.utils.config import CheckpointConfig, JobConfig
 from cosmos_framework.model._base import ImaginaireModel
 from cosmos_framework.utils import callback, distributed, log, misc
+from cosmos_framework.utils.config import CheckpointConfig, JobConfig
 from cosmos_framework.utils.easy_io import easy_io
 from cosmos_framework.utils.generator.rand_state import get_rand_state_dict, set_rand_state_dict
 
@@ -748,7 +748,7 @@ def keys_to_resume_during_load(self) -> tuple[set[str], str | None, bool | None]
 
                     # If the path doesn't end with specific checkpoint, read the latest
                     # checkpoint file to determine the most recent checkpoint iteration.
-                    if not re.search(r"/checkpoints/iter_\d{9}/?$", checkpoint_path):
+                    if not re.search(r"/checkpoints/(?:iter_\d{9}|epoch_\d+)/?$", checkpoint_path):
                         old_ckpt_path = checkpoint_path
                         latest_ckpt_path = os.path.join(checkpoint_path, "checkpoints/latest_checkpoint.txt")
 
@@ -1090,6 +1090,7 @@ def save(
         scheduler: torch.optim.lr_scheduler.LRScheduler,
         grad_scaler: torch.amp.GradScaler,
         iteration: int,
+        epoch: int | None = None,
     ) -> None:
         """Save network weights, optimizer parameters, scheduler parameters to a checkpoint.
 
@@ -1106,7 +1107,7 @@ def save(
         if self.callbacks is not None:
             self.callbacks.on_save_checkpoint_start(model, iteration)
 
-        checkpoint_file = f"iter_{iteration:09}"
+        checkpoint_file = f"epoch_{epoch}" if epoch is not None else f"iter_{iteration:09}"
 
         # Use rank-specific key for RNG state to ensure each rank saves its own state
         rng_key = f"rng_state_{dist.get_rank()}"
@@ -1129,7 +1130,7 @@ def save(
             self.callbacks.on_save_checkpoint(model, state_dict=to_save_dict)
 
         for k in to_save_dict.keys():
-            output_dirname = os.path.join(self.save_dirname, f"iter_{iteration:09}/{k}")
+            output_dirname = os.path.join(self.save_dirname, checkpoint_file, k)
             to_save_dict[k] = (to_save_dict[k], output_dirname)
 
         if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
diff --git a/cosmos_framework/checkpoint/dcp_distill.py b/cosmos_framework/checkpoint/dcp_distill.py
index a1b5279c..a6921e3a 100644
--- a/cosmos_framework/checkpoint/dcp_distill.py
+++ b/cosmos_framework/checkpoint/dcp_distill.py
@@ -26,9 +26,6 @@
 from torch.distributed.checkpoint.stateful import Stateful
 from torch.nn.modules.module import _IncompatibleKeys
 
-from cosmos_framework.model._base import ImaginaireModel
-from cosmos_framework.utils import log, misc
-from cosmos_framework.utils.easy_io import easy_io
 from cosmos_framework.checkpoint.dcp import (
     AsyncMode,
     CustomLoadPlanner,
@@ -38,8 +35,11 @@
     DistributedCheckpointer as _DistributedCheckpointer,
 )
 from cosmos_framework.checkpoint.dcp import ModelWrapper as VFMModelWrapper
-from cosmos_framework.utils.generator.rand_state import get_rand_state_dict, set_rand_state_dict
+from cosmos_framework.model._base import ImaginaireModel
 from cosmos_framework.model.generator.distillation.optimizer import OptimizerContainerLike, is_optimizer_container
+from cosmos_framework.utils import log, misc
+from cosmos_framework.utils.easy_io import easy_io
+from cosmos_framework.utils.generator.rand_state import get_rand_state_dict, set_rand_state_dict
 
 __all__: tuple[str, ...] = (
     "DistributedCheckpointer",
@@ -324,6 +324,7 @@ def save(
         scheduler: Any = None,
         grad_scaler: torch.amp.GradScaler | None = None,
         iteration: int = 0,
+        epoch: int | None = None,
     ) -> None:
         if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
             self._wait_for_previous_async_checkpoint()
@@ -332,7 +333,7 @@ def save(
             self.callbacks.on_save_checkpoint_start(model, iteration)
 
         model_dict = model.model_dict()
-        checkpoint_file = f"iter_{iteration:09}"
+        checkpoint_file = f"epoch_{epoch}" if epoch is not None else f"iter_{iteration:09}"
 
         rng_key = f"rng_state_{dist.get_rank()}"
         to_save_dict: dict[str, Any] = {
@@ -360,7 +361,7 @@ def save(
             to_save_dict["dataloader"] = dataloader_wrapper.state_dict()
 
         for key in list(to_save_dict.keys()):
-            output_dirname = os.path.join(self.save_dirname, f"iter_{iteration:09}/{key}")
+            output_dirname = os.path.join(self.save_dirname, checkpoint_file, key)
             to_save_dict[key] = (to_save_dict[key], output_dirname)
 
         if self.callbacks is not None:
diff --git a/cosmos_framework/checkpoint/dummy.py b/cosmos_framework/checkpoint/dummy.py
index 3cfa1bab..443241d9 100644
--- a/cosmos_framework/checkpoint/dummy.py
+++ b/cosmos_framework/checkpoint/dummy.py
@@ -22,6 +22,7 @@ def save(
         scheduler: torch.optim.lr_scheduler.LRScheduler,
         grad_scaler: torch.amp.GradScaler,
         iteration: int,
+        epoch: int | None = None,
     ) -> None:
         pass
 
diff --git a/cosmos_framework/configs/base/reasoner/defaults/callbacks.py b/cosmos_framework/configs/base/reasoner/defaults/callbacks.py
index 7d1911f5..37c7a26a 100644
--- a/cosmos_framework/configs/base/reasoner/defaults/callbacks.py
+++ b/cosmos_framework/configs/base/reasoner/defaults/callbacks.py
@@ -7,22 +7,22 @@
 
 from hydra.core.config_store import ConfigStore
 
-from cosmos_framework.callbacks.manual_gc import ManualGarbageCollection
-from cosmos_framework.utils.lazy_config import PLACEHOLDER
-from cosmos_framework.utils.lazy_config import LazyCall as L
-from cosmos_framework.utils.callback import LowPrecisionCallback, WandBCallback
 from cosmos_framework.callbacks.dataloader_state import DataLoaderStateCallback
-
 from cosmos_framework.callbacks.grad_clip import GradClip
 from cosmos_framework.callbacks.hf_export import HFExportCallback
 from cosmos_framework.callbacks.iter_speed import IterSpeed
 from cosmos_framework.callbacks.learning_rate_logger import LearningRateLogger
 from cosmos_framework.callbacks.log_tensor_shape import LogTensorShapeCallback
+from cosmos_framework.callbacks.manual_gc import ManualGarbageCollection
 from cosmos_framework.callbacks.param_count import ParamCount
+from cosmos_framework.callbacks.tao_status import TAOStatusCallback
 from cosmos_framework.callbacks.tokens_per_sec import VLMTokensPerSec
 from cosmos_framework.callbacks.wandb_log import WandbCallback as WandBCallbackMultiplier
 from cosmos_framework.callbacks.wandb_vis import VisualizationLoggingCallback
 from cosmos_framework.configs.base.defaults.callbacks import JOB_MONITOR_CALLBACKS
+from cosmos_framework.utils.callback import LowPrecisionCallback, WandBCallback
+from cosmos_framework.utils.lazy_config import PLACEHOLDER
+from cosmos_framework.utils.lazy_config import LazyCall as L
 
 # from cosmos_framework.utils.callback import NVTXCallback
 
@@ -52,6 +52,13 @@ def register_callbacks():
             config=PLACEHOLDER,
             trainer=PLACEHOLDER,
         ),  # reads model.precision; no extra kwarg needed
+        tao=L(TAOStatusCallback)(
+            enabled=False,
+            status_file_path=None,
+            experiment_name="",
+            logging_interval=1,
+            validation_heartbeat_interval=1,
+        ),
         # nvtx=L(NVTXCallback)(synchronize=True),
     )
 
diff --git a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py
new file mode 100644
index 00000000..6238a3db
--- /dev/null
+++ b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py
@@ -0,0 +1,294 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Woven Traffic Safety (WTS) video-QA SFT on the Cosmos3-Nano VLM path."""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+from collections import OrderedDict
+from copy import deepcopy
+from typing import Any
+
+import torch
+from hydra.core.config_store import ConfigStore
+from PIL import Image
+from torch.utils.data import Dataset
+
+from cosmos_framework.callbacks.cosmos_dataloader_state import CosmosDataLoaderStateCallback
+from cosmos_framework.configs.base.reasoner.experiment.dataflow_roles import VLMCollator, VLMProcessor
+from cosmos_framework.data.generator.dataflow import CosmosDataLoader, MapDistributor, PoolPackingBatcher
+from cosmos_framework.data.generator.processors import build_processor
+from cosmos_framework.utils.generator.torchcodec_video import TorchCodecVideoReader
+from cosmos_framework.utils.lazy_config import LazyCall as L
+from cosmos_framework.utils.lazy_config import LazyDict
+from cosmos_framework.utils.reasoner.constant import IGNORE_INDEX
+
+
+class WTSLlavaDataset(Dataset):
+    """Map-style loader for WTS LLaVA JSON annotations."""
+
+    def __init__(
+        self,
+        annotation_path: str,
+        media_path: str,
+        limit: int | str | None = None,
+    ) -> None:
+        self.annotation_path = os.path.abspath(os.path.expanduser(annotation_path))
+        self.media_path = os.path.abspath(os.path.expanduser(media_path))
+        if limit in ("", None):
+            parsed_limit = None
+        else:
+            parsed_limit = int(limit)
+            if parsed_limit < 1:
+                parsed_limit = None
+
+        with open(self.annotation_path, encoding="utf-8") as annotation_file:
+            records = json.load(annotation_file)
+        if not isinstance(records, list):
+            raise TypeError(f"WTS annotations must be a JSON array, got {type(records).__name__}")
+        self.records = records[:parsed_limit] if parsed_limit is not None else records
+        if not self.records:
+            raise ValueError(f"WTS annotation file contains no usable records: {self.annotation_path}")
+
+        for index, record in enumerate(self.records):
+            if not isinstance(record, dict) or not isinstance(record.get("video"), str):
+                raise ValueError(f"WTS record {index} must contain a string 'video' field")
+            conversations = record.get("conversations")
+            if not isinstance(conversations, list) or len(conversations) < 2:
+                raise ValueError(f"WTS record {index} must contain at least two conversation turns")
+
+    def __len__(self) -> int:
+        return len(self.records)
+
+    def __getitem__(self, index: int) -> dict[str, Any]:
+        record = dict(self.records[index])
+        video_path = record["video"]
+        if not os.path.isabs(video_path):
+            video_path = os.path.join(self.media_path, video_path)
+        if not os.path.isfile(video_path):
+            raise FileNotFoundError(f"WTS video does not exist: {video_path}")
+        record["video"] = video_path
+        return record
+
+
+class WTSProcessor(VLMProcessor):
+    """Convert WTS ShareGPT records and uniformly sample each video to PIL frames."""
+
+    def __init__(
+        self,
+        processor: Any,
+        ignore_index: int = IGNORE_INDEX,
+        num_video_frames: int = 8,
+        video_cache_size: int = 8,
+        system_prompt: str = "",
+    ) -> None:
+        super().__init__(processor=processor, ignore_index=ignore_index)
+        if num_video_frames < 1:
+            raise ValueError("num_video_frames must be >= 1")
+        if video_cache_size < 0:
+            raise ValueError("video_cache_size must be >= 0")
+        self.num_video_frames = num_video_frames
+        self.video_cache_size = video_cache_size
+        self.system_prompt = system_prompt
+        self._video_cache: OrderedDict[str, tuple[list[Image.Image], float]] = OrderedDict()
+
+    def _decode_video(self, video_path: str) -> tuple[list[Image.Image], float]:
+        cached = self._video_cache.get(video_path)
+        if cached is not None:
+            self._video_cache.move_to_end(video_path)
+            return cached
+
+        reader = TorchCodecVideoReader(video_path, num_threads=2)
+        total_frames = len(reader)
+        if total_frames < 1:
+            raise ValueError(f"WTS video has zero frames: {video_path}")
+        sample_count = min(self.num_video_frames, total_frames)
+        if sample_count == 1:
+            indices = [0]
+        else:
+            indices = torch.linspace(0, total_frames - 1, steps=sample_count).round().to(dtype=torch.long).tolist()
+        frames_np = reader.get_frames_nhwc_uint8(indices)
+        frames = [Image.fromarray(frame) for frame in frames_np]
+
+        source_fps = reader.get_avg_fps()
+        average_stride = (indices[-1] - indices[0]) / max(len(indices) - 1, 1) if len(indices) > 1 else 1.0
+        effective_fps = source_fps / max(average_stride, 1.0)
+        decoded = (frames, float(effective_fps))
+        if self.video_cache_size > 0:
+            self._video_cache[video_path] = decoded
+            self._video_cache.move_to_end(video_path)
+            while len(self._video_cache) > self.video_cache_size:
+                self._video_cache.popitem(last=False)
+        return decoded
+
+    def _sharegpt_to_openai(self, item: dict) -> list[dict]:
+        conversations = item.get("conversations", [])
+        video_path = item.get("video")
+        frames, fps = self._decode_video(video_path)
+        messages: list[dict] = []
+        video_inserted = False
+        if self.system_prompt:
+            messages.append({"role": "system", "content": self.system_prompt})
+
+        for turn in conversations:
+            role = "user" if turn["from"] == "human" else "assistant"
+            text = re.sub(r"(\n)?</?(image|video)>(\n)?", "", turn["value"]).strip()
+            if role == "user" and not video_inserted:
+                content: Any = [
+                    {"type": "video", "video": frames, "fps": fps},
+                    {"type": "text", "text": text},
+                ]
+                video_inserted = True
+            else:
+                content = text
+            messages.append({"role": role, "content": content})
+        return messages
+
+
+def _wts_dataloader(
+    *,
+    annotation_env: str,
+    media_env: str,
+    limit_env: str,
+    shuffle: bool,
+) -> LazyDict:
+    return L(CosmosDataLoader)(
+        distributor=L(MapDistributor)(
+            dataset=L(WTSLlavaDataset)(
+                annotation_path=f"${{oc.env:{annotation_env}}}",
+                media_path=f"${{oc.env:{media_env}}}",
+                limit=f"${{oc.env:{limit_env},''}}",
+            ),
+            shuffle=shuffle,
+            seed=42,
+            name="train" if shuffle else "val",
+        ),
+        processor=L(WTSProcessor)(
+            processor=L(build_processor)(
+                tokenizer_type="${model.config.policy.backbone.model_name}",
+                config_variant="hf",
+            ),
+            ignore_index=IGNORE_INDEX,
+            num_video_frames=8,
+            video_cache_size=8,
+            system_prompt=("You are a helpful assistant that can answer questions about a street-view CCTV footage."),
+        ),
+        batcher=L(PoolPackingBatcher)(
+            max_tokens="${data_setting.max_tokens}",
+            pool_size=16 if shuffle else 1,
+            max_batch_size=1,
+            long_threshold=6400,
+        ),
+        collator=L(VLMCollator)(),
+        num_workers=0,
+        prefetch_factor=None,
+        persistent_workers=False,
+        pin_memory=False,
+    )
+
+
+wts_vlm = LazyDict(
+    dict(
+        defaults=[
+            {"override /checkpoint": "local"},
+            {"override /data_train": None},
+            {"override /data_val": None},
+            {"override /model": "vlm_fsdp"},
+            {"override /vlm_policy": "qwen3_vl_8b_instruct"},
+            {"override /callbacks": ["basic_vlm", "basic_log"]},
+            "_self_",
+        ],
+        job=dict(
+            project="cosmos3_reasoner",
+            group="wts_sft",
+            wandb_mode="disabled",
+        ),
+        trainer=dict(
+            callbacks=dict(
+                dataloader_state=L(CosmosDataLoaderStateCallback)(),
+                tao=dict(
+                    enabled=True,
+                    logging_interval=1,
+                    validation_heartbeat_interval=1,
+                ),
+            ),
+            max_iter=10,
+            logging_iter=1,
+            run_validation=True,
+            validation_iter=10,
+            max_val_iter=10,
+            run_validation_on_start=False,
+            grad_accum_iter=1,
+        ),
+        optimizer=dict(
+            lr=1.0e-4,
+            fused=True,
+            weight_decay=0.01,
+            betas=[0.9, 0.999],
+            lr_multipliers={"model.visual": 1.0},
+        ),
+        model=dict(
+            config=dict(
+                policy=dict(
+                    model_max_length=81920,
+                    qwen_max_video_token_length=8192,
+                ),
+                freeze=dict(trainable_params=[".*"]),
+                parallelism=dict(
+                    data_parallel_shard_degree=4,
+                    data_parallel_replicate_degree=1,
+                ),
+            ),
+        ),
+        data_setting=dict(
+            max_tokens=81920,
+            qwen_max_video_token_length=8192,
+        ),
+        checkpoint=dict(
+            save_iter=100,
+            load_from_object_store=dict(enabled=False, credentials="", bucket=""),
+            save_to_object_store=dict(enabled=False, credentials="", bucket=""),
+        ),
+        dataloader_train=_wts_dataloader(
+            annotation_env="WTS_TRAIN_ANNOTATION",
+            media_env="WTS_TRAIN_MEDIA",
+            limit_env="WTS_TRAIN_LIMIT",
+            shuffle=True,
+        ),
+        dataloader_val=_wts_dataloader(
+            annotation_env="WTS_VAL_ANNOTATION",
+            media_env="WTS_VAL_MEDIA",
+            limit_env="WTS_VAL_LIMIT",
+            shuffle=False,
+        ),
+        upload_reproducible_setup=False,
+    ),
+    flags={"allow_objects": True},
+)
+
+ConfigStore.instance().store(
+    group="experiment",
+    package="_global_",
+    name="wts_vlm",
+    node=wts_vlm,
+)
+
+
+# Edge uses the same WTS data contract but a different native HF backbone.
+# Keep it in this module so the dataloader worker callables retain a single
+# pickle identity when the reasoner config loader reloads experiment modules.
+wts_vlm_edge = deepcopy(wts_vlm)
+wts_vlm_edge["defaults"][4] = {"override /vlm_policy": "cosmos3_edge_reasoner"}
+wts_vlm_edge["job"]["group"] = "wts_edge_sft"
+wts_vlm_edge["optimizer"].pop("lr_multipliers", None)
+wts_vlm_edge["model"]["config"]["policy"]["model_max_length"] = 16000
+
+ConfigStore.instance().store(
+    group="experiment",
+    package="_global_",
+    name="wts_vlm_edge",
+    node=wts_vlm_edge,
+)
diff --git a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm_test.py b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm_test.py
new file mode 100644
index 00000000..3257a3b0
--- /dev/null
+++ b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm_test.py
@@ -0,0 +1,56 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+from __future__ import annotations
+
+import json
+
+import pytest
+
+from cosmos_framework.configs.base.reasoner.experiment.wts_vlm import WTSLlavaDataset, wts_vlm_edge
+
+
+def test_wts_dataset_resolves_video_paths_and_limit(tmp_path) -> None:
+    media = tmp_path / "videos"
+    media.mkdir()
+    (media / "clip.mp4").write_bytes(b"video")
+    annotations = tmp_path / "annotations.json"
+    annotations.write_text(
+        json.dumps(
+            [
+                {
+                    "video": "clip.mp4",
+                    "conversations": [
+                        {"from": "human", "value": "<video> what happens?"},
+                        {"from": "gpt", "value": "A car stops."},
+                    ],
+                },
+                {
+                    "video": "clip.mp4",
+                    "conversations": [
+                        {"from": "human", "value": "<video> what happens next?"},
+                        {"from": "gpt", "value": "Traffic moves."},
+                    ],
+                },
+            ]
+        ),
+        encoding="utf-8",
+    )
+
+    dataset = WTSLlavaDataset(str(annotations), str(media), limit="1")
+    assert len(dataset) == 1
+    assert dataset[0]["video"] == str(media / "clip.mp4")
+
+
+def test_wts_dataset_rejects_invalid_record(tmp_path) -> None:
+    annotations = tmp_path / "annotations.json"
+    annotations.write_text(json.dumps([{"video": "clip.mp4"}]), encoding="utf-8")
+
+    with pytest.raises(ValueError, match="conversation"):
+        WTSLlavaDataset(str(annotations), str(tmp_path))
+
+
+def test_wts_edge_uses_native_edge_policy() -> None:
+    assert wts_vlm_edge["defaults"][4] == {"override /vlm_policy": "cosmos3_edge_reasoner"}
+    assert "lr_multipliers" not in wts_vlm_edge["optimizer"]
+    assert wts_vlm_edge["model"]["config"]["policy"]["model_max_length"] == 16000
diff --git a/cosmos_framework/configs/toml_config/sft_config.py b/cosmos_framework/configs/toml_config/sft_config.py
index 04d0efbc..7fda51e3 100644
--- a/cosmos_framework/configs/toml_config/sft_config.py
+++ b/cosmos_framework/configs/toml_config/sft_config.py
@@ -357,8 +357,7 @@ class ModelConfig(BaseModel):
     lora_rank: int = Field(
         default=16,
         description=(
-            "LoRA rank `r`. Adapter shape is (rank × hidden_dim) per target "
-            "module. Standard values are 4, 8, 16, 32."
+            "LoRA rank `r`. Adapter shape is (rank × hidden_dim) per target module. Standard values are 4, 8, 16, 32."
         ),
     )
     lora_alpha: int = Field(
@@ -379,9 +378,7 @@ class ModelConfig(BaseModel):
     ema: EMAConfig = Field(default_factory=EMAConfig)
     parallelism: ParallelismConfig = Field(default_factory=ParallelismConfig)
     compile: CompileConfig = Field(default_factory=CompileConfig)
-    activation_checkpointing: ActivationCheckpointingConfig = Field(
-        default_factory=ActivationCheckpointingConfig
-    )
+    activation_checkpointing: ActivationCheckpointingConfig = Field(default_factory=ActivationCheckpointingConfig)
     tokenizer: ModelTokenizerConfig = Field(default_factory=ModelTokenizerConfig)
     backbone: BackboneConfig = Field(default_factory=BackboneConfig)
 
@@ -473,15 +470,12 @@ class SchedulerConfig(BaseModel):
     )
     f_start: list[float] = Field(
         default_factory=lambda: [1.0e-6],
-        description=(
-            "Initial LR multiplier at step 0, before warmup ramps up."
-        ),
+        description=("Initial LR multiplier at step 0, before warmup ramps up."),
     )
     verbosity_interval: int = Field(
         default=0,
         description=(
-            "How often the scheduler logs the current LR (in optimizer "
-            "steps). 0 = silent. VFM only — skipped on VLM."
+            "How often the scheduler logs the current LR (in optimizer steps). 0 = silent. VFM only — skipped on VLM."
         ),
     )
     warm_up_steps: list[int] = Field(
@@ -533,8 +527,7 @@ class GradClipCallback(BaseModel):
     clip_norm: float = Field(
         default=1.0,
         description=(
-            "Maximum global L2 norm of the gradient. Steps with a larger "
-            "norm are rescaled so ||grad|| ≤ clip_norm."
+            "Maximum global L2 norm of the gradient. Steps with a larger norm are rescaled so ||grad|| ≤ clip_norm."
         ),
     )
     force_finite: bool = Field(
@@ -547,8 +540,37 @@ class GradClipCallback(BaseModel):
     )
 
 
+class TAOStatusCallbackConfig(BaseModel):
+    """TAO-compatible lifecycle and training/validation metric logging."""
+
+    model_config = _PYDANTIC_MODEL_CONFIG
+
+    enabled: bool = Field(default=False, description="Enable TAO status.json logging.")
+    status_file_path: Optional[str] = Field(
+        default=None,
+        description=(
+            "Explicit status.json path. When unset, TAO_JOB_ID/TAO_RESULTS_ROOT, "
+            "legacy TAO API variables, then job.path_local are used."
+        ),
+    )
+    experiment_name: str = Field(
+        default="",
+        description="TAO component name. Empty uses job.name.",
+    )
+    logging_interval: int = Field(
+        default=1,
+        ge=1,
+        description="Multiplier applied to trainer.logging_iter for TAO training records.",
+    )
+    validation_heartbeat_interval: int = Field(
+        default=1,
+        ge=1,
+        description="Write validation progress every N validation batches.",
+    )
+
+
 class TrainerCallbacksConfig(BaseModel):
-    """Only the two callbacks the schema currently surfaces. The full
+    """Callbacks surfaced by the structured TOML schema. The full
     callbacks dict (norm_monitor, mfu, heart_beat, …) stays in the
     experiment Python.
     """
@@ -557,6 +579,7 @@ class TrainerCallbacksConfig(BaseModel):
 
     compile_tokenizer: CompileTokenizerCallback = Field(default_factory=CompileTokenizerCallback)
     grad_clip: GradClipCallback = Field(default_factory=GradClipCallback)
+    tao: TAOStatusCallbackConfig = Field(default_factory=TAOStatusCallbackConfig)
 
 
 class TrainerConfig(BaseModel):
@@ -567,8 +590,7 @@ class TrainerConfig(BaseModel):
     distributed_parallelism: str = Field(
         default="fsdp",
         description=(
-            "Distributed strategy. 'fsdp' (the only supported value today) "
-            "routes through cosmos's FSDP wrapper."
+            "Distributed strategy. 'fsdp' (the only supported value today) routes through cosmos's FSDP wrapper."
         ),
     )
     grad_accum_iter: int = Field(
@@ -587,6 +609,39 @@ class TrainerConfig(BaseModel):
         default=500,
         description="Total number of optimizer steps the run will execute.",
     )
+    num_epochs: Optional[int] = Field(
+        default=None,
+        ge=1,
+        description="Number of complete training epochs. Requires steps_per_epoch and takes priority over max_iter.",
+    )
+    steps_per_epoch: Optional[int] = Field(
+        default=None,
+        ge=1,
+        description="Optimizer updates in one training epoch.",
+    )
+    max_val_iter: Optional[int] = Field(
+        default=None,
+        ge=1,
+        description="Maximum validation batches per validation pass. None consumes a finite validation loader.",
+    )
+    run_validation: bool = Field(
+        default=False,
+        description="Enable validation during training.",
+    )
+    validation_iter: int = Field(
+        default=100,
+        ge=1,
+        description="Run validation every N optimizer steps.",
+    )
+    validation_freq_in_epoch: int = Field(
+        default=0,
+        ge=0,
+        description="Validate every N completed epochs; 0 uses validation_iter.",
+    )
+    run_validation_on_start: bool = Field(
+        default=False,
+        description="Run one validation pass before the first training step.",
+    )
     callbacks: TrainerCallbacksConfig = Field(default_factory=TrainerCallbacksConfig)
 
 
@@ -618,6 +673,11 @@ class CheckpointConfig(BaseModel):
         default=100,
         description="Save a new checkpoint every N optimizer steps.",
     )
+    save_freq_in_epoch: int = Field(
+        default=0,
+        ge=0,
+        description="Save every N completed epochs; 0 uses save_iter.",
+    )
 
 
 # ---------------------------------------------------------------- dataloader_train
@@ -658,10 +718,7 @@ class DataloaderTrainConfig(BaseModel):
     )
     seed: int = Field(
         default=42,
-        description=(
-            "Dataloader RNG seed. Skipped on VLM (CosmosDataLoader has "
-            "no seed ctor kwarg there)."
-        ),
+        description=("Dataloader RNG seed. Skipped on VLM (CosmosDataLoader has no seed ctor kwarg there)."),
     )
 
 
@@ -746,8 +803,7 @@ def load_experiment_from_toml(
         base_config_path = TASK_TO_BASE_CONFIG[task]
     except KeyError as e:
         raise ValueError(
-            f"{toml_path}: [job].task={task!r} is not supported. "
-            f"Valid values: {sorted(TASK_TO_BASE_CONFIG)}"
+            f"{toml_path}: [job].task={task!r} is not supported. Valid values: {sorted(TASK_TO_BASE_CONFIG)}"
         ) from e
 
     overrides = build_hydra_overrides(raw)
@@ -759,10 +815,7 @@ def load_experiment_from_toml(
             if not o or o == "--":
                 continue
             if "=" not in o:
-                raise ValueError(
-                    f"extra override {o!r} must be Hydra dotted-path syntax "
-                    f"(e.g. 'optimizer.lr=1e-5')."
-                )
+                raise ValueError(f"extra override {o!r} must be Hydra dotted-path syntax (e.g. 'optimizer.lr=1e-5').")
             overrides.append(o)
 
     # Import lazily so this module stays cheap to import in non-training contexts.
diff --git a/cosmos_framework/configs/toml_config/sft_config_test.py b/cosmos_framework/configs/toml_config/sft_config_test.py
index 8ff8f3f0..95f5b110 100644
--- a/cosmos_framework/configs/toml_config/sft_config_test.py
+++ b/cosmos_framework/configs/toml_config/sft_config_test.py
@@ -76,6 +76,24 @@ def test_custom_does_not_loosen_sibling_validation(self) -> None:
                 }
             )
 
+    def test_native_epoch_schedule_fields_validate(self) -> None:
+        cfg = SFTExperimentConfig.model_validate(
+            {
+                "job": {"task": "vlm", "experiment": "wts_vlm"},
+                "trainer": {
+                    "num_epochs": 2,
+                    "steps_per_epoch": 108,
+                    "validation_freq_in_epoch": 1,
+                },
+                "checkpoint": {"save_freq_in_epoch": 2},
+            }
+        )
+
+        assert cfg.trainer.num_epochs == 2
+        assert cfg.trainer.steps_per_epoch == 108
+        assert cfg.trainer.validation_freq_in_epoch == 1
+        assert cfg.checkpoint.save_freq_in_epoch == 2
+
 
 # --------------------------------------------------------------------------- #
 # 2. build_hydra_overrides must NOT emit [custom] as per-leaf overrides        #
diff --git a/cosmos_framework/scripts/train.py b/cosmos_framework/scripts/train.py
index 7e137503..8155ef6c 100644
--- a/cosmos_framework/scripts/train.py
+++ b/cosmos_framework/scripts/train.py
@@ -31,15 +31,14 @@
 import torch
 from loguru import logger as logging
 
-from cosmos_framework.utils.config import Config
-from cosmos_framework.utils.lazy_config import LazyConfig, instantiate
-from cosmos_framework.utils.serialization import to_yaml
+from cosmos_framework.configs.toml_config.sft_config import load_experiment_from_toml
 from cosmos_framework.utils import distributed
+from cosmos_framework.utils.config import Config
 from cosmos_framework.utils.context_managers import data_loader_init, distributed_init, model_init
 from cosmos_framework.utils.launch import log_reproducible_setup
+from cosmos_framework.utils.lazy_config import LazyConfig, instantiate
+from cosmos_framework.utils.serialization import to_yaml
 from cosmos_framework.utils.training_telemetry import telemetry
-from cosmos_framework.configs.toml_config.sft_config import load_experiment_from_toml
-
 
 # ---------------------------------------------------------------------------
 # --deterministic: mirrors launch_vfm.sh determinism settings.
@@ -217,12 +216,17 @@ def launch(config: Config, args: argparse.Namespace) -> None:
         dataloader_train = instantiate(config.dataloader_train)
         dataloader_val = instantiate(config.dataloader_val)
 
-    # Start training
-    trainer.train(
-        model,
-        dataloader_train,
-        dataloader_val,
-    )
+    # Start training. Give callbacks a final failure hook so external
+    # orchestrators (including TAO) receive a terminal status record.
+    try:
+        trainer.train(
+            model,
+            dataloader_train,
+            dataloader_val,
+        )
+    except BaseException as error:
+        trainer.callbacks.on_exception(error)
+        raise
 
 
 if __name__ == "__main__":
diff --git a/cosmos_framework/trainer/__init__.py b/cosmos_framework/trainer/__init__.py
index cafdf441..2cf55ecb 100644
--- a/cosmos_framework/trainer/__init__.py
+++ b/cosmos_framework/trainer/__init__.py
@@ -12,9 +12,13 @@
 import torch.distributed as dist
 import torch.utils.data
 
-from cosmos_framework.utils.flags import INTERNAL
 from cosmos_framework.utils.context_managers import distributed_init
-from cosmos_framework.utils.profiling import maybe_enable_memory_snapshot, maybe_enable_nsys_profiling, maybe_enable_profiling
+from cosmos_framework.utils.flags import INTERNAL
+from cosmos_framework.utils.profiling import (
+    maybe_enable_memory_snapshot,
+    maybe_enable_nsys_profiling,
+    maybe_enable_profiling,
+)
 
 try:
     from megatron.core import parallel_state
@@ -24,14 +28,13 @@
     USE_MEGATRON = False
 
 
-from cosmos_framework.utils.lazy_config import LazyConfig, instantiate
 from cosmos_framework.model._base import ImaginaireModel
 from cosmos_framework.utils import callback, distributed, ema, log, misc
 from cosmos_framework.utils.checkpointer import Checkpointer
+from cosmos_framework.utils.lazy_config import LazyConfig, instantiate
 from cosmos_framework.utils.misc import StragglerDetectorV2
 
 
-
 @dataclass
 class ContextParallelDataWindow:
     """Caches one dataloader batch across a ``cp_size``-step CP data window.
@@ -74,7 +77,6 @@ def assert_synced_with_model(self, model: Any) -> None:
                 f"CP data-window desync: trainer offset {self.offset} != model window slot {model_window_slot}."
             )
 
-
 class ImaginaireTrainer:
     """The base trainer class of Imaginaire.
 
@@ -95,6 +97,21 @@ def __init__(self, config):
         """
         super().__init__()
         self.config = config
+        steps_per_epoch = getattr(config.trainer, "steps_per_epoch", None)
+        num_epochs = getattr(config.trainer, "num_epochs", None)
+        self.steps_per_epoch = int(steps_per_epoch) if steps_per_epoch else None
+        self.num_epochs = int(num_epochs) if num_epochs else None
+        if (self.steps_per_epoch is None) != (self.num_epochs is None):
+            raise ValueError("trainer.steps_per_epoch and trainer.num_epochs must be configured together")
+        self.max_iterations = (
+            self.steps_per_epoch * self.num_epochs
+            if self.steps_per_epoch is not None and self.num_epochs is not None
+            else int(config.trainer.max_iter)
+        )
+        if int(getattr(config.checkpoint, "save_freq_in_epoch", 0)) > 0 and self.steps_per_epoch is None:
+            raise ValueError("checkpoint.save_freq_in_epoch requires trainer.num_epochs and trainer.steps_per_epoch")
+        if int(getattr(config.trainer, "validation_freq_in_epoch", 0)) > 0 and self.steps_per_epoch is None:
+            raise ValueError("trainer.validation_freq_in_epoch requires trainer.num_epochs and trainer.steps_per_epoch")
         # Set up the distributed computing environment.
         with distributed_init():
             distributed.init()
@@ -325,7 +342,7 @@ def train(
                     finally:
                         self.callbacks.on_after_dataloading(iteration)
                     # If max_iter is reached, exit the training loop.
-                    if iteration >= self.config.trainer.max_iter:
+                    if iteration >= self.max_iterations:
                         _end_training = True
                         break
                     # Move all tensors in the data batch to GPU device.
@@ -357,11 +374,39 @@ def train(
                     # Do the following when an actual optimizer (update) step has been made.
                     iteration += 1
                     # Save checkpoint.
-                    if iteration % self.config.checkpoint.save_iter == 0:
-                        self.checkpointer.save(model, optimizer, scheduler, grad_scaler, iteration=iteration)
+                    completed_epoch = (
+                        iteration // self.steps_per_epoch
+                        if self.steps_per_epoch and iteration % self.steps_per_epoch == 0
+                        else None
+                    )
+                    save_freq_in_epoch = int(getattr(self.config.checkpoint, "save_freq_in_epoch", 0))
+                    epoch_save_due = bool(
+                        completed_epoch is not None
+                        and save_freq_in_epoch > 0
+                        and completed_epoch % save_freq_in_epoch == 0
+                    )
+                    iter_save_due = save_freq_in_epoch == 0 and iteration % self.config.checkpoint.save_iter == 0
+                    if epoch_save_due or iter_save_due:
+                        self.checkpointer.save(
+                            model,
+                            optimizer,
+                            scheduler,
+                            grad_scaler,
+                            iteration=iteration,
+                            epoch=completed_epoch if epoch_save_due else None,
+                        )
                     self.callbacks.on_training_step_end(model, data_batch, output_batch, loss, iteration=iteration)
                     # Validation.
-                    if self.config.trainer.run_validation and iteration % self.config.trainer.validation_iter == 0:
+                    validation_freq_in_epoch = int(getattr(self.config.trainer, "validation_freq_in_epoch", 0))
+                    epoch_validation_due = bool(
+                        completed_epoch is not None
+                        and validation_freq_in_epoch > 0
+                        and completed_epoch % validation_freq_in_epoch == 0
+                    )
+                    iter_validation_due = (
+                        validation_freq_in_epoch == 0 and iteration % self.config.trainer.validation_iter == 0
+                    )
+                    if self.config.trainer.run_validation and (epoch_validation_due or iter_validation_due):
                         self.validate(model, dataloader_val, iteration=iteration)
                     # This iteration is successful; reset the timeout signal.
                     signal.alarm(self.config.trainer.timeout_period)
@@ -377,8 +422,24 @@ def train(
         log.success("Done with training.")
         if sm_carveout:
             torch._C._set_sm_carveout_experimental(None)
-        if iteration % self.config.checkpoint.save_iter != 0:
-            self.checkpointer.save(model, optimizer, scheduler, grad_scaler, iteration=iteration)
+        final_epoch = (
+            iteration // self.steps_per_epoch
+            if self.steps_per_epoch and iteration % self.steps_per_epoch == 0
+            else None
+        )
+        save_freq_in_epoch = int(getattr(self.config.checkpoint, "save_freq_in_epoch", 0))
+        final_already_saved = (
+            final_epoch is not None and save_freq_in_epoch > 0 and final_epoch % save_freq_in_epoch == 0
+        ) or (save_freq_in_epoch == 0 and iteration % self.config.checkpoint.save_iter == 0)
+        if not final_already_saved:
+            self.checkpointer.save(
+                model,
+                optimizer,
+                scheduler,
+                grad_scaler,
+                iteration=iteration,
+                epoch=final_epoch if save_freq_in_epoch > 0 else None,
+            )
         self.callbacks.on_train_end(model, iteration=iteration)
         self.checkpointer.finalize()
         distributed.barrier()
diff --git a/cosmos_framework/utils/callback.py b/cosmos_framework/utils/callback.py
index e2b96ccd..c0c7a3da 100644
--- a/cosmos_framework/utils/callback.py
+++ b/cosmos_framework/utils/callback.py
@@ -15,11 +15,10 @@
 import tqdm
 import wandb
 
-from cosmos_framework.utils.lazy_config import instantiate
 from cosmos_framework.utils import distributed, log, misc, wandb_util
+from cosmos_framework.utils.lazy_config import instantiate
 from cosmos_framework.utils.misc import get_local_tensor_if_DTensor
 
-
 try:
     from megatron.core import parallel_state
 except ImportError:
@@ -27,9 +26,9 @@
 
 
 if TYPE_CHECKING:
-    from cosmos_framework.utils.config import Config
     from cosmos_framework.model._base import ImaginaireModel
     from cosmos_framework.trainer import ImaginaireTrainer
+    from cosmos_framework.utils.config import Config
 
 
 class CallBackGroup:
@@ -305,6 +304,10 @@ def on_train_end(self, model: ImaginaireModel, iteration: int = 0) -> None:
     def on_app_end(self) -> None:
         pass
 
+    def on_exception(self, error: BaseException) -> None:
+        """Called when the application exits because of an uncaught exception."""
+        pass
+
 
 def _missing_callback_hooks(callback: object) -> list[str]:
     """Return required callback hooks that are absent or non-callable."""
diff --git a/cosmos_framework/utils/checkpointer.py b/cosmos_framework/utils/checkpointer.py
index 71cd2021..cb28534f 100644
--- a/cosmos_framework/utils/checkpointer.py
+++ b/cosmos_framework/utils/checkpointer.py
@@ -61,6 +61,7 @@ def save(
         scheduler: torch.optim.lr_scheduler.LRScheduler,
         grad_scaler: torch.amp.GradScaler,
         iteration: int,
+        epoch: int | None = None,
     ) -> None:
         """Save network weights, optimizer parameters, scheduler parameters to a checkpoint.
 
@@ -73,7 +74,7 @@ def save(
         """
         self.callbacks.on_save_checkpoint_start(model, iteration)
 
-        checkpoint_file = f"iter_{iteration:09}.pt"
+        checkpoint_file = f"epoch_{epoch}.pt" if epoch is not None else f"iter_{iteration:09}.pt"
 
         if distributed.get_rank() == 0:
             state_dict = dict(
@@ -116,7 +117,7 @@ def _save_worker_local(self, state_dict: dict[str, torch.Tensor], checkpoint_fil
             if rank == 0:
                 self._write_latest_checkpoint_file(checkpoint_file)
             log.success(f"Saved checkpoint (local): {checkpoint_path}")
-            iteration = int(checkpoint_file.replace("iter_", "").replace(".pt", ""))
+            iteration = int(state_dict["iteration"])
             self.callbacks.on_save_checkpoint_success(iteration=iteration)
         except Exception as e:  # noqa: BLE001
             log.exception(f"Checkpoint failed to save (local): {e}")
@@ -138,7 +139,7 @@ def _save_worker_object_store(
             if rank == 0:
                 self._write_latest_checkpoint_file(checkpoint_file)
             log.success(f"Saved checkpoint (object store): {checkpoint_path}")
-            iteration = int(checkpoint_file.replace("iter_", "").replace(".pt", ""))
+            iteration = int(state_dict["iteration"])
             self.callbacks.on_save_checkpoint_success(iteration=iteration)
         except Exception as e:  # noqa: BLE001
             log.exception(f"Checkpoint failed to upload (object store): {e}")
@@ -305,6 +306,7 @@ def save(
         scheduler: torch.optim.lr_scheduler.LRScheduler,
         grad_scaler: torch.amp.GradScaler,
         iteration: int,
+        epoch: int | None = None,
     ) -> None:
         """Save network weights, optimizer parameters, scheduler parameters to a checkpoint.
 
@@ -317,7 +319,7 @@ def save(
         """
         # checkpoint_file = f"iter_{iteration:09}.pt"
         postfix, _, total_ema_num = model.get_ckpt_postfix()
-        checkpoint_file = f"iter_{iteration:09}{postfix}.pt"
+        checkpoint_file = f"epoch_{epoch}{postfix}.pt" if epoch is not None else f"iter_{iteration:09}{postfix}.pt"
         save_ranks = list(range(total_ema_num))
         for _rank in save_ranks:
             if distributed.get_rank() == _rank:
diff --git a/cosmos_framework/utils/config.py b/cosmos_framework/utils/config.py
index d886b92c..b951628b 100644
--- a/cosmos_framework/utils/config.py
+++ b/cosmos_framework/utils/config.py
@@ -288,6 +288,9 @@ class CheckpointConfig:
 
     # Save the checkpoint every N iterations.
     save_iter: int = 999999999
+    # Save every N completed epochs when trainer.steps_per_epoch is configured.
+    # A positive value takes priority over save_iter.
+    save_freq_in_epoch: int = 0
 
     # Load state_dict to the models in strict mode. If True, `allow_partial_load` in dcp
     # planner will be set to False. DCP will raise an error if there are missing keys.
@@ -453,6 +456,10 @@ class TrainerConfig:
     grad_scaler_args: dict = attrs.field(factory=lambda: dict(enabled=False))
     # Maximum number of iterations to train the model.
     max_iter: int = 999999999
+    # Native epoch schedule. When both values are positive, the trainer runs
+    # num_epochs * steps_per_epoch optimizer updates and reports epoch progress.
+    num_epochs: int | None = None
+    steps_per_epoch: int | None = None
     # Maximum number of iterations to validate the model. If None, validate on the entire dataset.
     max_val_iter: int | None = None
     # How often we log the training stats.
@@ -461,6 +468,9 @@ class TrainerConfig:
     run_validation: bool = True
     # How often we evaluate on the validation set.
     validation_iter: int = 999999999
+    # Validate every N completed epochs. A positive value takes priority over
+    # validation_iter when steps_per_epoch is configured.
+    validation_freq_in_epoch: int = 0
     # Whether to run the validation on the start of the training.
     run_validation_on_start: bool = False
     # Kill the process after N seconds since the last iteration (usually means dead job).
diff --git a/cosmos_framework/utils/reasoner/dcp_checkpointer.py b/cosmos_framework/utils/reasoner/dcp_checkpointer.py
index 442b8a4c..806c3afa 100644
--- a/cosmos_framework/utils/reasoner/dcp_checkpointer.py
+++ b/cosmos_framework/utils/reasoner/dcp_checkpointer.py
@@ -55,8 +55,8 @@
 
 from cosmos_framework.checkpoint.base import AbstractCheckpointer
 from cosmos_framework.checkpoint.s3_filesystem import S3StorageReader, S3StorageWriter
-from cosmos_framework.utils.config import CheckpointConfig, JobConfig
 from cosmos_framework.utils import callback, distributed, log, misc
+from cosmos_framework.utils.config import CheckpointConfig, JobConfig
 from cosmos_framework.utils.easy_io import easy_io
 from cosmos_framework.utils.reasoner.model_wrapper import ModelWrapper
 from cosmos_framework.utils.reasoner.optimizer import OptimizersContainer
@@ -458,6 +458,7 @@ def save(
         scheduler: torch.optim.lr_scheduler.LRScheduler,
         grad_scaler: torch.amp.GradScaler,
         iteration: int,
+        epoch: int | None = None,
     ) -> None:
         """Save network weights, optimizer parameters, scheduler parameters to a checkpoint.
 
@@ -471,7 +472,7 @@ def save(
         del grad_scaler
         self.callbacks.on_save_checkpoint_start(model_parts, iteration)
 
-        checkpoint_file = f"iter_{iteration:09}"
+        checkpoint_file = f"epoch_{epoch}" if epoch is not None else f"iter_{iteration:09}"
         to_save_dict = {
             "model": ModelWrapper(model_parts).state_dict(),
             "optim": optimizer.state_dict(),
@@ -482,7 +483,7 @@ def save(
         self.callbacks.on_save_checkpoint(model_parts, state_dict=to_save_dict)
 
         for k in to_save_dict.keys():
-            output_dirname = os.path.join(self.save_dirname, f"iter_{iteration:09}/{k}")
+            output_dirname = os.path.join(self.save_dirname, checkpoint_file, k)
             to_save_dict[k] = (to_save_dict[k], output_dirname)
 
         if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
diff --git a/examples/launch_sft_wts.sh b/examples/launch_sft_wts.sh
new file mode 100755
index 00000000..bdc6ac7c
--- /dev/null
+++ b/examples/launch_sft_wts.sh
@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+set -u
+
+: "${WTS_TRAIN_ANNOTATION:?Set WTS_TRAIN_ANNOTATION to the training JSON file}"
+: "${WTS_TRAIN_MEDIA:?Set WTS_TRAIN_MEDIA to the training video root}"
+: "${WTS_VAL_ANNOTATION:?Set WTS_VAL_ANNOTATION to the validation JSON file}"
+: "${WTS_VAL_MEDIA:?Set WTS_VAL_MEDIA to the validation video root}"
+: "${VLM_SAFETENSORS_PATH:?Set VLM_SAFETENSORS_PATH to converted Cosmos3-Nano safetensors}"
+
+export WTS_TRAIN_ANNOTATION WTS_TRAIN_MEDIA
+export WTS_VAL_ANNOTATION WTS_VAL_MEDIA
+export VLM_SAFETENSORS_PATH
+export WTS_TRAIN_LIMIT="${WTS_TRAIN_LIMIT:-}"
+export WTS_VAL_LIMIT="${WTS_VAL_LIMIT:-}"
+
+TOML_FILE="examples/toml/sft_config/wts_vlm.toml"
+EXTRA_DATASET_CHECK='
+[[ -f "$WTS_TRAIN_ANNOTATION" ]] || { echo "ERROR: training annotations not found: $WTS_TRAIN_ANNOTATION" >&2; exit 1; }
+[[ -d "$WTS_TRAIN_MEDIA" ]] || { echo "ERROR: training media root not found: $WTS_TRAIN_MEDIA" >&2; exit 1; }
+[[ -f "$WTS_VAL_ANNOTATION" ]] || { echo "ERROR: validation annotations not found: $WTS_VAL_ANNOTATION" >&2; exit 1; }
+[[ -d "$WTS_VAL_MEDIA" ]] || { echo "ERROR: validation media root not found: $WTS_VAL_MEDIA" >&2; exit 1; }
+[[ -d "$VLM_SAFETENSORS_PATH" ]] || { echo "ERROR: converted checkpoint not found: $VLM_SAFETENSORS_PATH" >&2; exit 1; }
+'
+
+TAIL_OVERRIDES=(
+    ${EXTRA_TAIL_OVERRIDES:-}
+)
+
+source "$(dirname "${BASH_SOURCE[0]}")/_sft_launcher_common.sh"
diff --git a/examples/toml/sft_config/wts_vlm.toml b/examples/toml/sft_config/wts_vlm.toml
new file mode 100644
index 00000000..767142b0
--- /dev/null
+++ b/examples/toml/sft_config/wts_vlm.toml
@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+# WTS video-QA SFT through the migrated Cosmos Framework VLM trainer.
+# Required environment:
+#   WTS_TRAIN_ANNOTATION, WTS_TRAIN_MEDIA
+#   WTS_VAL_ANNOTATION, WTS_VAL_MEDIA
+#   VLM_SAFETENSORS_PATH (Cosmos3-Nano converted to Qwen3-VL safetensors)
+# This checked-in recipe is a 10-step-per-epoch smoke run. Set
+# trainer.steps_per_epoch to the distributed update count for a production dataset.
+
+[job]
+task       = "vlm"
+experiment = "wts_vlm"
+project    = "cosmos3_reasoner"
+group      = "wts_sft"
+name       = "wts_vlm"
+wandb_mode = "disabled"
+
+[model]
+attn_implementation = "cosmos"
+precision = "bfloat16"
+
+[model.backbone]
+model_name = "${oc.env:VLM_SAFETENSORS_PATH}"
+safetensors_path = "${oc.env:VLM_SAFETENSORS_PATH}"
+
+[model.ema]
+enabled = false
+rate = 0.1
+iteration_shift = 0
+
+[model.parallelism]
+data_parallel_shard_degree = 4
+data_parallel_replicate_degree = 1
+context_parallel_shard_degree = 1
+cfg_parallel_shard_degree = 1
+
+[model.compile]
+enabled = false
+compile_dynamic = true
+
+[model.activation_checkpointing]
+mode = "full"
+save_ops_regex = ["fmha"]
+preserve_rng_state = true
+determinism_check = "default"
+
+[optimizer]
+betas = [0.9, 0.999]
+eps = 1.0e-8
+fused = true
+lr = 1.0e-4
+weight_decay = 0.01
+
+[scheduler]
+cycle_lengths = [20]
+f_max = [1.0]
+f_min = [0.0]
+f_start = [0.0]
+verbosity_interval = 0
+warm_up_steps = [0]
+
+[trainer]
+distributed_parallelism = "fsdp"
+grad_accum_iter = 1
+logging_iter = 1
+max_iter = 20
+num_epochs = 2
+steps_per_epoch = 10
+max_val_iter = 10
+run_validation = true
+validation_iter = 10
+validation_freq_in_epoch = 1
+run_validation_on_start = false
+
+[trainer.callbacks.compile_tokenizer]
+compile_after_iterations = 3
+enabled = false
+
+[trainer.callbacks.grad_clip]
+clip_norm = 1.0
+force_finite = false
+
+[trainer.callbacks.tao]
+enabled = true
+experiment_name = "Cosmos Framework WTS SFT"
+logging_interval = 1
+validation_heartbeat_interval = 1
+
+[checkpoint]
+keys_to_skip_loading = []
+load_path = "???"
+save_iter = 100
+save_freq_in_epoch = 2
+
+[dataloader_train]
+max_samples_per_batch = 1
+max_sequence_length = 81920
diff --git a/examples/toml/sft_config/wts_vlm_edge.toml b/examples/toml/sft_config/wts_vlm_edge.toml
new file mode 100644
index 00000000..ffdffde6
--- /dev/null
+++ b/examples/toml/sft_config/wts_vlm_edge.toml
@@ -0,0 +1,95 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+# WTS video-QA SFT on the native Cosmos3-Edge reasoner checkpoint.
+# This checked-in recipe is a 10-step-per-epoch smoke run. Set
+# trainer.steps_per_epoch to the distributed update count for a production dataset.
+
+[job]
+task       = "vlm"
+experiment = "wts_vlm_edge"
+project    = "cosmos3_reasoner"
+group      = "wts_edge_sft"
+name       = "wts_vlm_edge"
+wandb_mode = "disabled"
+
+[model]
+attn_implementation = "flash_attention_2"
+precision = "bfloat16"
+
+[model.backbone]
+model_name = "${oc.env:VLM_SAFETENSORS_PATH}"
+safetensors_path = "${oc.env:VLM_SAFETENSORS_PATH}"
+
+[model.ema]
+enabled = false
+rate = 0.1
+iteration_shift = 0
+
+[model.parallelism]
+data_parallel_shard_degree = 4
+data_parallel_replicate_degree = 1
+context_parallel_shard_degree = 1
+cfg_parallel_shard_degree = 1
+
+[model.compile]
+enabled = false
+compile_dynamic = true
+
+[model.activation_checkpointing]
+mode = "full"
+save_ops_regex = ["fmha"]
+preserve_rng_state = true
+determinism_check = "default"
+
+[optimizer]
+betas = [0.9, 0.999]
+eps = 1.0e-8
+fused = true
+lr = 1.0e-4
+weight_decay = 0.01
+
+[scheduler]
+cycle_lengths = [20]
+f_max = [1.0]
+f_min = [0.0]
+f_start = [0.0]
+verbosity_interval = 0
+warm_up_steps = [0]
+
+[trainer]
+distributed_parallelism = "fsdp"
+grad_accum_iter = 1
+logging_iter = 1
+max_iter = 20
+num_epochs = 2
+steps_per_epoch = 10
+max_val_iter = 10
+run_validation = true
+validation_iter = 10
+validation_freq_in_epoch = 1
+run_validation_on_start = false
+
+[trainer.callbacks.compile_tokenizer]
+compile_after_iterations = 3
+enabled = false
+
+[trainer.callbacks.grad_clip]
+clip_norm = 1.0
+force_finite = false
+
+[trainer.callbacks.tao]
+enabled = true
+experiment_name = "Cosmos Framework WTS Edge SFT"
+logging_interval = 1
+validation_heartbeat_interval = 1
+
+[checkpoint]
+keys_to_skip_loading = []
+load_path = "???"
+save_iter = 100
+save_freq_in_epoch = 2
+
+[dataloader_train]
+max_samples_per_batch = 1
+max_sequence_length = 16000
diff --git a/uv.lock b/uv.lock
index 42ce3eef..f8f862b8 100644
--- a/uv.lock
+++ b/uv.lock
@@ -146,22 +146,22 @@ conflicts = [[
     { package = "cosmos-framework", group = "cu130-train" },
     { package = "cosmos-framework", group = "vllm" },
 ], [
-    { package = "cosmos-framework", group = "cu128" },
     { package = "cosmos-framework", group = "cu128-train" },
     { package = "cosmos-framework", group = "cu130" },
+    { package = "cosmos-framework", group = "cu130-torch213" },
     { package = "cosmos-framework", group = "cu130-torch213-train" },
     { package = "cosmos-framework", group = "cu130-train" },
     { package = "cosmos-framework", group = "vllm" },
 ], [
+    { package = "cosmos-framework", group = "cu128" },
     { package = "cosmos-framework", group = "cu128-train" },
     { package = "cosmos-framework", group = "cu130" },
-    { package = "cosmos-framework", group = "cu130-torch213" },
     { package = "cosmos-framework", group = "cu130-torch213-train" },
     { package = "cosmos-framework", group = "cu130-train" },
     { package = "cosmos-framework", group = "vllm" },
 ], [
+    { package = "cosmos-framework", group = "cu128" },
     { package = "cosmos-framework", group = "cu128-train" },
-    { package = "cosmos-framework", group = "cu130" },
     { package = "cosmos-framework", group = "cu130-torch213-train" },
     { package = "cosmos-framework", group = "cu130-train" },
     { package = "cosmos-framework", group = "vllm" },
@@ -177,15 +177,15 @@ conflicts = [[
     { package = "cosmos-framework", group = "cu130-train" },
     { package = "cosmos-framework", group = "vllm" },
 ], [
-    { package = "cosmos-framework", group = "cu128" },
     { package = "cosmos-framework", group = "cu128-train" },
-    { package = "cosmos-framework", group = "cu130-torch213" },
+    { package = "cosmos-framework", group = "cu130" },
     { package = "cosmos-framework", group = "cu130-torch213-train" },
     { package = "cosmos-framework", group = "cu130-train" },
     { package = "cosmos-framework", group = "vllm" },
 ], [
     { package = "cosmos-framework", group = "cu128" },
     { package = "cosmos-framework", group = "cu128-train" },
+    { package = "cosmos-framework", group = "cu130-torch213" },
     { package = "cosmos-framework", group = "cu130-torch213-train" },
     { package = "cosmos-framework", group = "cu130-train" },
     { package = "cosmos-framework", group = "vllm" },

From 949688a25b98bb20269733b56048ca89a5e532e5 Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Sat, 1 Aug 2026 05:06:10 +0000
Subject: [PATCH 02/13] fix: allow non-root training containers

---
 Dockerfile           | 3 ++-
 docker/entrypoint.sh | 4 +++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index ea474be5..1d730f5d 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -40,7 +40,8 @@ WORKDIR /workspace
 # Install python
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,source=.python-version,target=.python-version \
-    uv python install
+    uv python install && \
+    chmod 0711 /root
 
 # Install into virtual environment
 RUN echo "$CUDA_VERSION" | sed -E 's/^([0-9]+)\.([0-9]+).*/cu\1\2/' > /root/.cuda-name
diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh
index 55f01aee..8c0f7bbb 100755
--- a/docker/entrypoint.sh
+++ b/docker/entrypoint.sh
@@ -7,6 +7,8 @@
 
 set -e
 
-uv pip install --no-deps -e . || true
+if [[ -w /workspace/.venv ]]; then
+    uv pip install --no-deps -e . || true
+fi
 
 exec "$@"

From 37c593a9eb357f5a5d7f2d72aeba2ccf04dce328 Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Sat, 1 Aug 2026 05:09:50 +0000
Subject: [PATCH 03/13] build: exclude Python bytecode from image context

---
 .dockerignore | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/.dockerignore b/.dockerignore
index 0dfe444b..2c1ae741 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -1,5 +1,7 @@
 .venv
 .git
+**/__pycache__
+**/*.pyc
 /checkpoints
 /datasets
 /output

From bf98eef3df69436ba72f4db7034096540fdd275a Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Tue, 4 Aug 2026 02:18:03 +0000
Subject: [PATCH 04/13] feat: add DAFT-backed AETC SFT

---
 .../base/reasoner/experiment/wts_vlm.py       |  83 ++++++++++++
 .../base/reasoner/experiment/wts_vlm_test.py  |  71 +++++++++-
 .../generator/local_datasets/tao_vl_reason.py | 124 ++++++++++++++++++
 examples/toml/sft_config/aetc_daft_vlm.toml   |  99 ++++++++++++++
 4 files changed, 376 insertions(+), 1 deletion(-)
 create mode 100644 cosmos_framework/data/generator/local_datasets/tao_vl_reason.py
 create mode 100644 examples/toml/sft_config/aetc_daft_vlm.toml

diff --git a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py
index 6238a3db..7143aea5 100644
--- a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py
+++ b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py
@@ -20,6 +20,10 @@
 from cosmos_framework.callbacks.cosmos_dataloader_state import CosmosDataLoaderStateCallback
 from cosmos_framework.configs.base.reasoner.experiment.dataflow_roles import VLMCollator, VLMProcessor
 from cosmos_framework.data.generator.dataflow import CosmosDataLoader, MapDistributor, PoolPackingBatcher
+from cosmos_framework.data.generator.local_datasets.tao_vl_reason import (
+    TaoVlReasonDaftDataset,
+    apply_daft_chat_template,
+)
 from cosmos_framework.data.generator.processors import build_processor
 from cosmos_framework.utils.generator.torchcodec_video import TorchCodecVideoReader
 from cosmos_framework.utils.lazy_config import LazyCall as L
@@ -84,6 +88,7 @@ def __init__(
         num_video_frames: int = 8,
         video_cache_size: int = 8,
         system_prompt: str = "",
+        use_daft_chat_template: bool = False,
     ) -> None:
         super().__init__(processor=processor, ignore_index=ignore_index)
         if num_video_frames < 1:
@@ -93,6 +98,9 @@ def __init__(
         self.num_video_frames = num_video_frames
         self.video_cache_size = video_cache_size
         self.system_prompt = system_prompt
+        self.use_daft_chat_template = use_daft_chat_template
+        if self.use_daft_chat_template:
+            apply_daft_chat_template(processor)
         self._video_cache: OrderedDict[str, tuple[list[Image.Image], float]] = OrderedDict()
 
     def _decode_video(self, video_path: str) -> tuple[list[Image.Image], float]:
@@ -125,6 +133,23 @@ def _decode_video(self, video_path: str) -> tuple[list[Image.Image], float]:
         return decoded
 
     def _sharegpt_to_openai(self, item: dict) -> list[dict]:
+        if "messages" in item:
+            messages = deepcopy(item["messages"])
+            for message in messages:
+                content = message.get("content")
+                if not isinstance(content, list):
+                    continue
+                for part in content:
+                    if part.get("type") != "video":
+                        continue
+                    video_path = part.get("video")
+                    if not isinstance(video_path, str):
+                        raise TypeError("DAFT video content must contain a string path")
+                    frames, fps = self._decode_video(video_path)
+                    part["video"] = frames
+                    part["fps"] = fps
+            return messages
+
         conversations = item.get("conversations", [])
         video_path = item.get("video")
         frames, fps = self._decode_video(video_path)
@@ -190,6 +215,49 @@ def _wts_dataloader(
     )
 
 
+def _aetc_daft_dataloader(*, split: str, shuffle: bool) -> LazyDict:
+    annotation_env = f"AETC_{split.upper()}_ANNOTATIONS"
+    media_env = f"AETC_{split.upper()}_MEDIA"
+    limit_env = f"AETC_{split.upper()}_LIMIT"
+    return L(CosmosDataLoader)(
+        distributor=L(MapDistributor)(
+            dataset=L(TaoVlReasonDaftDataset)(
+                annotation_paths=f"${{oc.env:{annotation_env}}}",
+                media_root=f"${{oc.env:{media_env}}}",
+                response_mode="hybrid" if split == "train" else "answer",
+                system_prompt="${oc.env:AETC_SYSTEM_PROMPT,''}",
+                vision_kwargs={},
+                max_samples=f"${{oc.env:{limit_env},''}}",
+            ),
+            shuffle=shuffle,
+            seed=42,
+            name=split,
+        ),
+        processor=L(WTSProcessor)(
+            processor=L(build_processor)(
+                tokenizer_type="${model.config.policy.backbone.model_name}",
+                config_variant="hf",
+            ),
+            ignore_index=IGNORE_INDEX,
+            num_video_frames=8,
+            video_cache_size=8,
+            system_prompt="",
+            use_daft_chat_template=True,
+        ),
+        batcher=L(PoolPackingBatcher)(
+            max_tokens="${data_setting.max_tokens}",
+            pool_size=16 if shuffle else 1,
+            max_batch_size=1,
+            long_threshold=6400,
+        ),
+        collator=L(VLMCollator)(),
+        num_workers=0,
+        prefetch_factor=None,
+        persistent_workers=False,
+        pin_memory=False,
+    )
+
+
 wts_vlm = LazyDict(
     dict(
         defaults=[
@@ -277,6 +345,21 @@ def _wts_dataloader(
 )
 
 
+# Internal TAO AETC path: keep Framework's trainer/model implementation while
+# consuming the same DAFT dataset and Qwen chat-template contract as Cosmos-RL.
+aetc_daft_vlm = deepcopy(wts_vlm)
+aetc_daft_vlm["job"]["group"] = "aetc_daft_sft"
+aetc_daft_vlm["dataloader_train"] = _aetc_daft_dataloader(split="train", shuffle=True)
+aetc_daft_vlm["dataloader_val"] = _aetc_daft_dataloader(split="val", shuffle=False)
+
+ConfigStore.instance().store(
+    group="experiment",
+    package="_global_",
+    name="aetc_daft_vlm",
+    node=aetc_daft_vlm,
+)
+
+
 # Edge uses the same WTS data contract but a different native HF backbone.
 # Keep it in this module so the dataloader worker callables retain a single
 # pickle identity when the reasoner config loader reloads experiment modules.
diff --git a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm_test.py b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm_test.py
index 3257a3b0..13d44175 100644
--- a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm_test.py
+++ b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm_test.py
@@ -4,10 +4,46 @@
 from __future__ import annotations
 
 import json
+import sys
+import types
 
 import pytest
 
-from cosmos_framework.configs.base.reasoner.experiment.wts_vlm import WTSLlavaDataset, wts_vlm_edge
+from cosmos_framework.configs.base.reasoner.experiment.wts_vlm import (
+    WTSLlavaDataset,
+    aetc_daft_vlm,
+    wts_vlm_edge,
+)
+from cosmos_framework.data.generator.local_datasets.tao_vl_reason import (
+    TaoVlReasonDaftDataset,
+    apply_daft_chat_template,
+    parse_path_list,
+)
+
+
+def _install_fake_daft(monkeypatch) -> list[object]:
+    calls: list[object] = []
+
+    class FakeDataset:
+        def __init__(self, **kwargs) -> None:
+            calls.append(kwargs)
+            self._raw_length = 2
+
+        def __getitem__(self, index: int) -> list[dict]:
+            return [{"role": "assistant", "content": f"daft-{index}"}]
+
+    def fake_template(processor) -> None:
+        calls.append(processor)
+
+    package = types.ModuleType("nvidia_tao_daft")
+    datasets = types.ModuleType("nvidia_tao_daft.datasets")
+    module = types.ModuleType("nvidia_tao_daft.datasets.tao_vl_reason_v1_0")
+    module.TaoVlReasonV1_0CosmosRLConversationDataset = FakeDataset
+    module.apply_chat_template_override = fake_template
+    monkeypatch.setitem(sys.modules, "nvidia_tao_daft", package)
+    monkeypatch.setitem(sys.modules, "nvidia_tao_daft.datasets", datasets)
+    monkeypatch.setitem(sys.modules, "nvidia_tao_daft.datasets.tao_vl_reason_v1_0", module)
+    return calls
 
 
 def test_wts_dataset_resolves_video_paths_and_limit(tmp_path) -> None:
@@ -54,3 +90,36 @@ def test_wts_edge_uses_native_edge_policy() -> None:
     assert wts_vlm_edge["defaults"][4] == {"override /vlm_policy": "cosmos3_edge_reasoner"}
     assert "lr_multipliers" not in wts_vlm_edge["optimizer"]
     assert wts_vlm_edge["model"]["config"]["policy"]["model_max_length"] == 16000
+
+
+def test_daft_dataset_matches_internal_hybrid_index_order(monkeypatch) -> None:
+    calls = _install_fake_daft(monkeypatch)
+    dataset = TaoVlReasonDaftDataset(
+        annotation_paths='["bcq.json", "mcq.json"]',
+        media_root="/data/aetc",
+        response_mode="hybrid",
+        system_prompt="system",
+    )
+
+    assert len(dataset) == 4
+    assert [dataset[index]["messages"][0]["content"] for index in range(4)] == [
+        "daft-0",
+        "daft-2",
+        "daft-1",
+        "daft-3",
+    ]
+    assert calls[0]["annotation_paths"] == ["bcq.json", "mcq.json"]
+    assert calls[0]["media_roots"] == "/data/aetc"
+
+
+def test_daft_chat_template_targets_wrapped_hf_processor(monkeypatch) -> None:
+    calls = _install_fake_daft(monkeypatch)
+    wrapped = types.SimpleNamespace(processor=object())
+    apply_daft_chat_template(wrapped)
+    assert calls == [wrapped.processor]
+
+
+def test_daft_recipe_and_json_path_parser() -> None:
+    assert parse_path_list('["a.json", "b.json"]') == ["a.json", "b.json"]
+    assert aetc_daft_vlm["job"]["group"] == "aetc_daft_sft"
+    assert aetc_daft_vlm["dataloader_train"]["processor"]["use_daft_chat_template"]
diff --git a/cosmos_framework/data/generator/local_datasets/tao_vl_reason.py b/cosmos_framework/data/generator/local_datasets/tao_vl_reason.py
new file mode 100644
index 00000000..398d1d1d
--- /dev/null
+++ b/cosmos_framework/data/generator/local_datasets/tao_vl_reason.py
@@ -0,0 +1,124 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""DAFT-backed ``tao-vl-reason-v1.0`` dataset adapter for Framework SFT."""
+
+from __future__ import annotations
+
+import json
+from typing import Any, Literal
+
+from torch.utils.data import Dataset
+
+ResponseMode = Literal["think", "answer", "hybrid"]
+
+
+def parse_path_list(value: str | list[str]) -> list[str]:
+    """Resolve a path or JSON-encoded path list supplied through Hydra/env."""
+    if isinstance(value, list):
+        paths = value
+    elif isinstance(value, str):
+        stripped = value.strip()
+        if not stripped:
+            paths = []
+        elif stripped.startswith("["):
+            paths = json.loads(stripped)
+        else:
+            paths = [stripped]
+    else:
+        raise TypeError(f"annotation_paths must be a string or list, got {type(value).__name__}")
+    if not paths or not all(isinstance(path, str) and path.strip() for path in paths):
+        raise ValueError("annotation_paths must contain at least one non-empty path")
+    return [path.strip() for path in paths]
+
+
+def _optional_positive_int(value: int | str | None) -> int | None:
+    if value in (None, ""):
+        return None
+    parsed = int(value)
+    return parsed if parsed > 0 else None
+
+
+class TaoVlReasonDaftDataset(Dataset):
+    """Expose DAFT conversations through Cosmos Framework's map dataflow.
+
+    The index mapping intentionally matches the internal Cosmos-RL DAFT hook:
+    hybrid mode interleaves answer and reasoning targets for every raw item.
+    """
+
+    def __init__(
+        self,
+        annotation_paths: str | list[str],
+        media_root: str | list[str] | None = None,
+        response_mode: ResponseMode = "answer",
+        system_prompt: str = "",
+        vision_kwargs: dict[str, Any] | None = None,
+        max_samples: int | str | None = None,
+        sample_stride: int = 1,
+        sample_offset: int = 0,
+    ) -> None:
+        try:
+            from nvidia_tao_daft.datasets.tao_vl_reason_v1_0 import (
+                TaoVlReasonV1_0CosmosRLConversationDataset,
+            )
+        except ImportError as exc:
+            raise ImportError("nvidia-tao-daft>=2.9.1 is required for the AETC DAFT recipe") from exc
+
+        paths = parse_path_list(annotation_paths)
+        if isinstance(media_root, str) and media_root.strip().startswith("["):
+            media_root = parse_path_list(media_root)
+        if response_mode not in {"think", "answer", "hybrid"}:
+            raise ValueError(f"unsupported response_mode: {response_mode!r}")
+        if sample_stride < 1:
+            raise ValueError("sample_stride must be positive")
+        if sample_offset < 0:
+            raise ValueError("sample_offset must be non-negative")
+
+        self.dataset = TaoVlReasonV1_0CosmosRLConversationDataset(
+            annotation_paths=paths,
+            media_roots=media_root,
+            system_prompt=system_prompt,
+            vision_kwargs=vision_kwargs or {},
+            response_mode=response_mode,
+        )
+        self.response_mode = response_mode
+        self.sample_stride = sample_stride
+        self.sample_offset = sample_offset
+        self.max_samples = _optional_positive_int(max_samples)
+        self.raw_length = int(getattr(self.dataset, "_raw_length"))
+
+    def _subsampled_raw_length(self) -> int:
+        if self.sample_offset >= self.raw_length:
+            return 0
+        return ((self.raw_length - self.sample_offset - 1) // self.sample_stride) + 1
+
+    def __len__(self) -> int:
+        length = self._subsampled_raw_length()
+        if self.response_mode == "hybrid":
+            length *= 2
+        if self.max_samples is not None:
+            length = min(length, self.max_samples)
+        return length
+
+    def __getitem__(self, index: int) -> dict[str, list[dict[str, Any]]]:
+        if not 0 <= index < len(self):
+            raise IndexError(index)
+        if self.response_mode == "hybrid":
+            raw_index = self.sample_offset + (index // 2) * self.sample_stride
+            daft_index = raw_index if index % 2 == 0 else self.raw_length + raw_index
+        else:
+            daft_index = self.sample_offset + index * self.sample_stride
+        return {"messages": self.dataset[daft_index]}
+
+
+def apply_daft_chat_template(processor: Any) -> None:
+    """Apply DAFT's Qwen3-VL Instruct template to a Framework processor."""
+    try:
+        from nvidia_tao_daft.datasets.tao_vl_reason_v1_0 import (
+            apply_chat_template_override,
+        )
+    except ImportError as exc:
+        raise ImportError("nvidia-tao-daft>=2.9.1 is required for the AETC DAFT recipe") from exc
+
+    huggingface_processor = getattr(processor, "processor", processor)
+    apply_chat_template_override(huggingface_processor)
diff --git a/examples/toml/sft_config/aetc_daft_vlm.toml b/examples/toml/sft_config/aetc_daft_vlm.toml
new file mode 100644
index 00000000..f9d456f2
--- /dev/null
+++ b/examples/toml/sft_config/aetc_daft_vlm.toml
@@ -0,0 +1,99 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+# AETC dense SFT through the internal NVIDIA TAO DAFT data contract.
+# Required environment:
+#   AETC_TRAIN_ANNOTATIONS: JSON array of tao-vl-reason-v1.0 files
+#   AETC_TRAIN_MEDIA: shared media root for the training annotations
+#   AETC_VAL_ANNOTATIONS: JSON array of validation files
+#   AETC_VAL_MEDIA: shared media root for validation
+#   VLM_SAFETENSORS_PATH: Cosmos3-Nano checkpoint
+
+[job]
+task       = "vlm"
+experiment = "aetc_daft_vlm"
+project    = "cosmos3_reasoner"
+group      = "aetc_daft_sft"
+name       = "aetc_daft_vlm"
+wandb_mode = "disabled"
+
+[model]
+attn_implementation = "cosmos"
+precision = "bfloat16"
+
+[model.backbone]
+model_name = "${oc.env:VLM_SAFETENSORS_PATH}"
+safetensors_path = "${oc.env:VLM_SAFETENSORS_PATH}"
+
+[model.ema]
+enabled = false
+rate = 0.1
+iteration_shift = 0
+
+[model.parallelism]
+data_parallel_shard_degree = 4
+data_parallel_replicate_degree = 1
+context_parallel_shard_degree = 1
+cfg_parallel_shard_degree = 1
+
+[model.compile]
+enabled = false
+compile_dynamic = true
+
+[model.activation_checkpointing]
+mode = "full"
+save_ops_regex = ["fmha"]
+preserve_rng_state = true
+determinism_check = "default"
+
+[optimizer]
+betas = [0.9, 0.999]
+eps = 1.0e-6
+fused = true
+lr = 1.0e-5
+weight_decay = 0.01
+
+[scheduler]
+cycle_lengths = [10]
+f_max = [1.0]
+f_min = [1.0]
+f_start = [1.0]
+verbosity_interval = 0
+warm_up_steps = [0]
+
+[trainer]
+distributed_parallelism = "fsdp"
+grad_accum_iter = 1
+logging_iter = 1
+max_iter = 10
+num_epochs = 1
+steps_per_epoch = 10
+max_val_iter = 10
+run_validation = true
+validation_iter = 10
+validation_freq_in_epoch = 1
+run_validation_on_start = false
+
+[trainer.callbacks.compile_tokenizer]
+compile_after_iterations = 3
+enabled = false
+
+[trainer.callbacks.grad_clip]
+clip_norm = 1.0
+force_finite = false
+
+[trainer.callbacks.tao]
+enabled = true
+experiment_name = "Cosmos Framework AETC DAFT SFT"
+logging_interval = 1
+validation_heartbeat_interval = 1
+
+[checkpoint]
+keys_to_skip_loading = []
+load_path = "???"
+save_iter = 10
+save_freq_in_epoch = 1
+
+[dataloader_train]
+max_samples_per_batch = 1
+max_sequence_length = 40960

From 7f51d4900c80a7ef0379b09639929a0eb4fa238c Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Wed, 5 Aug 2026 02:11:53 +0000
Subject: [PATCH 05/13] feat: add reproducible Cosmos3 Nano TAO training
 support

---
 Dockerfile                                    |  25 +-
 cosmos_framework/callbacks/tao_status.py      | 222 +++++++++++++---
 cosmos_framework/callbacks/tao_status_test.py |  35 ++-
 .../base/reasoner/defaults/policy_config.py   |  37 ++-
 .../base/reasoner/experiment/wts_vlm.py       |  31 ++-
 .../configs/toml_config/sft_config.py         |  11 +-
 .../configs/toml_config/sft_config_test.py    |  41 +++
 .../configs/toml_config/toml_config_helper.py |  22 +-
 .../generator/algorithm/loss/cross_entropy.py |  28 +-
 .../loss/cross_entropy_stats_test.py          |  38 +++
 .../model/generator/qwen3_vl_compat.py        |  62 +++++
 .../model/generator/qwen3_vl_compat_test.py   |  50 ++++
 cosmos_framework/model/generator/vlm_model.py |  73 ++++-
 cosmos_framework/scripts/export_vlm_dcp.py    | 250 ++++++++++++++++++
 .../scripts/export_vlm_dcp_test.py            |  43 +++
 cosmos_framework/scripts/train.py             |  86 +++---
 .../utils/easy_io/backends/__init__.py        |  19 +-
 .../utils/easy_io/backends/registry_utils.py  |  16 +-
 cosmos_framework/utils/easy_io/file_client.py |  13 +-
 cosmos_framework/utils/generator/lora.py      | 158 +++++++++--
 cosmos_framework/utils/generator/lora_test.py |  60 +++++
 docker/write_image_provenance.py              |  88 ++++++
 22 files changed, 1265 insertions(+), 143 deletions(-)
 create mode 100644 cosmos_framework/model/generator/algorithm/loss/cross_entropy_stats_test.py
 create mode 100644 cosmos_framework/model/generator/qwen3_vl_compat.py
 create mode 100644 cosmos_framework/model/generator/qwen3_vl_compat_test.py
 create mode 100644 cosmos_framework/scripts/export_vlm_dcp.py
 create mode 100644 cosmos_framework/scripts/export_vlm_dcp_test.py
 create mode 100644 cosmos_framework/utils/generator/lora_test.py
 create mode 100644 docker/write_image_provenance.py

diff --git a/Dockerfile b/Dockerfile
index 1d730f5d..c2e70fba 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -7,6 +7,23 @@ ARG CUDA_VERSION=13.0.2
 ARG BASE_IMAGE=nvidia/cuda:${CUDA_VERSION}-cudnn-devel-ubuntu24.04
 FROM ${BASE_IMAGE}
 
+ARG SOURCE_COMMIT
+ARG SOURCE_TREE
+ARG SOURCE_DIRTY=1
+ARG BUILD_TIMESTAMP
+ARG BASE_IMAGE
+ARG CUDA_VERSION
+LABEL org.opencontainers.image.revision="${SOURCE_COMMIT}" \
+      org.opencontainers.image.created="${BUILD_TIMESTAMP}" \
+      com.nvidia.tao.source-tree="${SOURCE_TREE}" \
+      com.nvidia.tao.backend="cosmos-framework"
+ENV SOURCE_COMMIT="${SOURCE_COMMIT}" \
+    SOURCE_TREE="${SOURCE_TREE}" \
+    SOURCE_DIRTY="${SOURCE_DIRTY}" \
+    BUILD_TIMESTAMP="${BUILD_TIMESTAMP}" \
+    PROVENANCE_BASE_IMAGE="${BASE_IMAGE}" \
+    CUDA_VERSION="${CUDA_VERSION}"
+
 # Set the DEBIAN_FRONTEND environment variable to avoid interactive prompts during apt operations.
 ENV DEBIAN_FRONTEND=noninteractive
 
@@ -28,7 +45,8 @@ COPY --from=ghcr.io/astral-sh/uv:0.11.28 /uv /uvx /usr/local/bin/
 # Copy from the cache instead of linking since it's a mounted volume
 ENV UV_LINK_MODE=copy
 # Cache python downloads
-ENV UV_PYTHON_CACHE_DIR=/root/.cache/uv/python
+ENV UV_PYTHON_CACHE_DIR=/opt/uv-python-cache \
+    UV_PYTHON_INSTALL_DIR=/opt/uv-python
 
 # Install just: https://just.systems/man/en/pre-built-binaries.html
 RUN curl --proto '=https' --tlsv1.2 -sSf https://just.systems/install.sh | bash -s -- --to /usr/local/bin --tag 1.46.0
@@ -41,7 +59,7 @@ WORKDIR /workspace
 RUN --mount=type=cache,target=/root/.cache/uv \
     --mount=type=bind,source=.python-version,target=.python-version \
     uv python install && \
-    chmod 0711 /root
+    chmod -R a+rX /opt/uv-python /opt/uv-python-cache
 
 # Install into virtual environment
 RUN echo "$CUDA_VERSION" | sed -E 's/^([0-9]+)\.([0-9]+).*/cu\1\2/' > /root/.cuda-name
@@ -59,6 +77,9 @@ COPY . /workspace
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --no-deps .
 
+RUN /workspace/.venv/bin/python /workspace/docker/write_image_provenance.py && \
+    chmod -R a+rX /opt/tao /workspace/.venv /workspace/cosmos_framework
+
 # Triton bundled ptxas doesn't support latest GPU architectures
 ENV TRITON_PTXAS_PATH="/usr/local/cuda/bin/ptxas"
 
diff --git a/cosmos_framework/callbacks/tao_status.py b/cosmos_framework/callbacks/tao_status.py
index d318d076..2013a363 100644
--- a/cosmos_framework/callbacks/tao_status.py
+++ b/cosmos_framework/callbacks/tao_status.py
@@ -15,7 +15,7 @@
 import torch
 import torch.distributed as dist
 
-from cosmos_framework.utils import distributed, log, misc
+from cosmos_framework.utils import distributed, log
 from cosmos_framework.utils.callback import Callback
 
 
@@ -106,6 +106,34 @@ def write(
             status_file.write(json.dumps(payload, default=str) + "\n")
 
 
+def write_early_failure(error: BaseException) -> bool:
+    """Write a terminal TAO record before the callback/config exists.
+
+    The orchestration layer supplies ``TAO_STATUS_FILE`` for direct launches,
+    or the normal TAO job/result variables. No implicit host path is used.
+    """
+    status_path = os.environ.get("TAO_STATUS_FILE")
+    if not status_path:
+        job_id = os.environ.get("TAO_JOB_ID")
+        results_root = os.environ.get("TAO_RESULTS_ROOT")
+        if job_id and results_root:
+            status_path = os.path.join(results_root, job_id, "status.json")
+    if not status_path:
+        api_job_id = os.environ.get("TAO_API_JOB_ID")
+        results_root = os.environ.get("TAO_API_RESULTS_DIR")
+        if api_job_id and results_root:
+            status_path = os.path.join(results_root, api_job_id, "status.json")
+    if not status_path:
+        return False
+    _TAOStatusWriter(status_path).write(
+        status="FAILURE",
+        verbosity="ERROR",
+        message=f"Cosmos Framework training failed before callback initialization: {error}",
+        data={"phase": "preflight_or_initialization", "error_type": type(error).__name__},
+    )
+    return True
+
+
 class TAOStatusCallback(Callback):
     """Write TAO lifecycle, training, and validation records from rank zero.
 
@@ -138,8 +166,11 @@ def __init__(
         self._train_start_time = 0.0
         self._step_start_time = 0.0
         self._validation_batches = 0
-        self._validation_loss_sum = 0.0
-        self._validation_sample_count = 0
+        self._validation_loss_numerator = 0.0
+        self._validation_loss_denominator = 0
+        self._train_loss_numerator = 0.0
+        self._train_loss_denominator = 0
+        self.last_training_loss: float | None = None
         self.last_validation_loss: float | None = None
 
     def _is_rank_zero(self) -> bool:
@@ -158,12 +189,16 @@ def _resolve_status_file(self) -> str:
 
         job_id = os.environ.get("TAO_JOB_ID")
         if job_id:
-            results_root = os.environ.get("TAO_RESULTS_ROOT", "/results")
+            results_root = os.environ.get("TAO_RESULTS_ROOT")
+            if not results_root:
+                raise RuntimeError("TAO_RESULTS_ROOT is required when TAO_JOB_ID is set")
             return os.path.join(results_root, job_id, "status.json")
 
         api_job_id = os.environ.get("TAO_API_JOB_ID")
         if api_job_id:
-            results_root = os.environ.get("TAO_API_RESULTS_DIR", "/results")
+            results_root = os.environ.get("TAO_API_RESULTS_DIR")
+            if not results_root:
+                raise RuntimeError("TAO_API_RESULTS_DIR is required when TAO_API_JOB_ID is set")
             return os.path.join(results_root, api_job_id, "status.json")
 
         return os.path.join(self.config.job.path_local, "status.json")
@@ -221,31 +256,88 @@ def _epoch_label(self, iteration: int) -> str:
         return f"epoch {progress['epoch']}/{progress['max_epoch']}"
 
     @staticmethod
-    def _global_average_loss(loss: torch.Tensor, data_batch: dict[str, Any]) -> tuple[float, int]:
-        sample_count = misc.get_data_batch_size(data_batch)
-        loss_sum = loss.detach() * sample_count
-        count = torch.tensor(sample_count, device=loss.device, dtype=torch.long)
+    def _local_token_stats(
+        loss: torch.Tensor,
+        data_batch: dict[str, Any],
+        output_batch: dict[str, Any],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        numerator = output_batch.get("loss_numerator")
+        denominator = output_batch.get("loss_denominator")
+        if numerator is not None and denominator is not None:
+            return numerator.detach(), denominator.detach().to(dtype=torch.long)
+
+        # Backward-compatible fallback for non-VLM models. It is deliberately
+        # sample-weighted and is never used by the Cosmos3 VLM path, which
+        # always emits exact token statistics.
+        sample_count = next(
+            (
+                int(value.shape[0])
+                for value in data_batch.values()
+                if isinstance(value, torch.Tensor) and value.ndim > 0
+            ),
+            1,
+        )
+        return (
+            loss.detach() * sample_count,
+            torch.tensor(sample_count, device=loss.device, dtype=torch.long),
+        )
+
+    @classmethod
+    def _global_token_average(
+        cls,
+        loss: torch.Tensor,
+        data_batch: dict[str, Any],
+        output_batch: dict[str, Any],
+    ) -> tuple[float, float, int]:
+        numerator, denominator = cls._local_token_stats(loss, data_batch, output_batch)
+        numerator = numerator.clone()
+        denominator = denominator.clone()
+        if dist.is_available() and dist.is_initialized():
+            dist.all_reduce(numerator, op=dist.ReduceOp.SUM)
+            dist.all_reduce(denominator, op=dist.ReduceOp.SUM)
+        global_denominator = int(denominator.item())
+        global_numerator = float(numerator.item())
+        return global_numerator / max(global_denominator, 1), global_numerator, global_denominator
+
+    @staticmethod
+    def _reduce_accumulator(numerator: float, denominator: int) -> tuple[float, int]:
+        device = "cuda" if dist.is_available() and dist.is_initialized() else "cpu"
+        values = torch.tensor([numerator, float(denominator)], dtype=torch.float64, device=device)
         if dist.is_available() and dist.is_initialized():
-            dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM)
-            dist.all_reduce(count, op=dist.ReduceOp.SUM)
-        global_count = int(count.item())
-        average = float(loss_sum.item()) / max(global_count, 1)
-        return average, global_count
+            dist.all_reduce(values, op=dist.ReduceOp.SUM)
+        return float(values[0].item()), int(values[1].item())
 
     def on_train_start(self, model: Any, iteration: int = 0) -> None:
         self._train_start_time = time.monotonic()
+        self._train_loss_numerator = 0.0
+        self._train_loss_denominator = 0
         writer = self._get_writer()
         if writer is not None:
             writer.write(
                 status="STARTED",
                 message=f"Starting {self._component_name()} training",
-                data=self._progress_data(iteration, seconds_per_step=0.0),
+                data={
+                    **self._progress_data(iteration, seconds_per_step=0.0),
+                    "parameter_summary": getattr(model, "parameter_summary", None),
+                },
             )
             log.info(f"TAO status will be logged to {writer.filename}")
 
     def on_training_step_start(self, model: Any, data: dict[str, Any], iteration: int = 0) -> None:
         self._step_start_time = time.monotonic()
 
+    def on_training_step_batch_end(
+        self,
+        model: Any,
+        data_batch: dict[str, Any],
+        output_batch: dict[str, Any],
+        loss: torch.Tensor,
+        iteration: int = 0,
+    ) -> None:
+        numerator, denominator = self._local_token_stats(loss, data_batch, output_batch)
+        self._train_loss_numerator += float(numerator.item())
+        self._train_loss_denominator += int(denominator.item())
+
     def on_training_step_end(
         self,
         model: Any,
@@ -258,12 +350,16 @@ def on_training_step_end(
         if iteration % interval != 0:
             return
 
-        average_loss, _ = self._global_average_loss(loss, data_batch)
+        average_loss, numerator, denominator = self._global_token_average(loss, data_batch, output_batch)
         writer = self._get_writer()
         if writer is None:
             return
         seconds_per_step = max(time.monotonic() - self._step_start_time, 0.0)
-        kpi = {"train/loss": average_loss, "train/loss_avg": average_loss}
+        kpi = {
+            "train/step_loss": average_loss,
+            "train/step_loss_numerator": numerator,
+            "train/step_loss_denominator": denominator,
+        }
         progress = self._progress_data(iteration, seconds_per_step=seconds_per_step)
         if "epoch" in progress:
             message = (
@@ -282,8 +378,8 @@ def on_training_step_end(
 
     def on_validation_start(self, model: Any, dataloader_val: Any, iteration: int = 0) -> None:
         self._validation_batches = 0
-        self._validation_loss_sum = 0.0
-        self._validation_sample_count = 0
+        self._validation_loss_numerator = 0.0
+        self._validation_loss_denominator = 0
         writer = self._get_writer()
         if writer is not None:
             writer.write(
@@ -301,10 +397,11 @@ def on_validation_step_end(
         loss: torch.Tensor,
         iteration: int = 0,
     ) -> None:
-        average_loss, sample_count = self._global_average_loss(loss, data_batch)
+        local_numerator, local_denominator = self._local_token_stats(loss, data_batch, output_batch)
+        average_loss, _, _ = self._global_token_average(loss, data_batch, output_batch)
         self._validation_batches += 1
-        self._validation_loss_sum += average_loss * sample_count
-        self._validation_sample_count += sample_count
+        self._validation_loss_numerator += float(local_numerator.item())
+        self._validation_loss_denominator += int(local_denominator.item())
 
         if self._validation_batches % self.validation_heartbeat_interval != 0:
             return
@@ -332,11 +429,14 @@ def on_validation_step_end(
             log.info(f"Validation {self._epoch_label(iteration)}, batch {batch_progress} - Loss: {average_loss:.6f}")
 
     def on_validation_end(self, model: Any, iteration: int = 0) -> None:
-        if self._validation_sample_count == 0:
+        numerator, denominator = self._reduce_accumulator(
+            self._validation_loss_numerator, self._validation_loss_denominator
+        )
+        if denominator == 0:
             log.warning("TAO validation logging saw zero samples; no val/loss record was written")
             return
 
-        self.last_validation_loss = self._validation_loss_sum / self._validation_sample_count
+        self.last_validation_loss = numerator / denominator
         writer = self._get_writer()
         if writer is not None:
             writer.write(
@@ -348,12 +448,65 @@ def on_validation_end(self, model: Any, iteration: int = 0) -> None:
                     **self._progress_data(iteration),
                     "phase": "validation_complete",
                     "validation_batches": self._validation_batches,
-                    "validation_samples": self._validation_sample_count,
+                    "validation_loss_numerator": numerator,
+                    "validation_valid_label_count": denominator,
+                },
+                kpi={
+                    "val/loss": self.last_validation_loss,
+                    "val/avg_loss": self.last_validation_loss,
+                    "val/loss_numerator": numerator,
+                    "val/valid_label_count": denominator,
                 },
-                kpi={"val/loss": self.last_validation_loss, "val/avg_loss": self.last_validation_loss},
             )
         log.info(f"Validation loss ({self._epoch_label(iteration)}): {self.last_validation_loss:.6f}")
 
+    def on_save_checkpoint_success(self, iteration: int = 0, elapsed_time: float = 0) -> None:
+        checkpoint_path = None
+        checkpoint_root = Path(self.config.job.path_local) / "checkpoints"
+        latest = checkpoint_root / "latest_checkpoint.txt"
+        if latest.is_file():
+            checkpoint_name = latest.read_text(encoding="utf-8").strip()
+            if checkpoint_name:
+                checkpoint_path = str((checkpoint_root / checkpoint_name).resolve())
+        writer = self._get_writer()
+        if writer is not None:
+            writer.write(
+                status="RUNNING",
+                message=f"Checkpoint saved successfully at step {iteration}",
+                data={
+                    **self._progress_data(iteration),
+                    "phase": "checkpoint_complete",
+                    "checkpoint_iteration": iteration,
+                    "checkpoint_elapsed_seconds": elapsed_time,
+                    "checkpoint_path": checkpoint_path,
+                },
+            )
+
+    def on_train_end(self, model: Any, iteration: int = 0) -> None:
+        numerator, denominator = self._reduce_accumulator(
+            self._train_loss_numerator, self._train_loss_denominator
+        )
+        if denominator == 0:
+            raise RuntimeError("TAO metric collection observed zero valid training labels")
+        self.last_training_loss = numerator / denominator
+        writer = self._get_writer()
+        if writer is not None:
+            writer.write(
+                status="RUNNING",
+                message=f"Training complete - token-weighted loss: {self.last_training_loss:.6f}",
+                data={
+                    **self._progress_data(iteration),
+                    "phase": "training_complete",
+                    "train_loss_numerator": numerator,
+                    "train_valid_label_count": denominator,
+                },
+                kpi={
+                    "train/avg_loss": self.last_training_loss,
+                    "train/loss_numerator": numerator,
+                    "train/valid_label_count": denominator,
+                },
+            )
+
     def on_app_end(self) -> None:
         writer = self._get_writer()
         if writer is not None:
@@ -363,11 +516,18 @@ def on_app_end(self) -> None:
                 data=self._progress_data(
                     int(getattr(getattr(self, "trainer", None), "max_iterations", self.config.trainer.max_iter))
                 ),
-                kpi=(
-                    {"val/loss": self.last_validation_loss, "val/avg_loss": self.last_validation_loss}
-                    if self.last_validation_loss is not None
-                    else None
-                ),
+                kpi={
+                    **(
+                        {"train/avg_loss": self.last_training_loss}
+                        if self.last_training_loss is not None
+                        else {}
+                    ),
+                    **(
+                        {"val/loss": self.last_validation_loss, "val/avg_loss": self.last_validation_loss}
+                        if self.last_validation_loss is not None
+                        else {}
+                    ),
+                },
             )
 
     def on_exception(self, error: BaseException) -> None:
diff --git a/cosmos_framework/callbacks/tao_status_test.py b/cosmos_framework/callbacks/tao_status_test.py
index f862487c..9c6e87a7 100644
--- a/cosmos_framework/callbacks/tao_status_test.py
+++ b/cosmos_framework/callbacks/tao_status_test.py
@@ -33,10 +33,17 @@ def test_tao_status_callback_writes_training_validation_and_success(tmp_path) ->
     callback = _callback(tmp_path)
     callback.on_train_start(model=None, iteration=0)
     callback.on_training_step_start(model=None, data={}, iteration=0)
+    callback.on_training_step_batch_end(
+        model=None,
+        data_batch={"input_ids": torch.zeros(2, 3)},
+        output_batch={"loss_numerator": torch.tensor(6.0), "loss_denominator": torch.tensor(12)},
+        loss=torch.tensor(0.5),
+        iteration=3,
+    )
     callback.on_training_step_end(
         model=None,
         data_batch={"input_ids": torch.zeros(2, 3)},
-        output_batch={},
+        output_batch={"loss_numerator": torch.tensor(6.0), "loss_denominator": torch.tensor(12)},
         loss=torch.tensor(0.5),
         iteration=3,
     )
@@ -44,11 +51,12 @@ def test_tao_status_callback_writes_training_validation_and_success(tmp_path) ->
     callback.on_validation_step_end(
         model=None,
         data_batch={"input_ids": torch.zeros(2, 3)},
-        output_batch={},
+        output_batch={"loss_numerator": torch.tensor(2.0), "loss_denominator": torch.tensor(8)},
         loss=torch.tensor(0.25),
         iteration=3,
     )
     callback.on_validation_end(model=None, iteration=3)
+    callback.on_train_end(model=None, iteration=6)
     callback.on_app_end()
 
     records = _records(tmp_path)
@@ -58,9 +66,10 @@ def test_tao_status_callback_writes_training_validation_and_success(tmp_path) ->
         "RUNNING",
         "RUNNING",
         "RUNNING",
+        "RUNNING",
         "SUCCESS",
     ]
-    assert records[1]["kpi"]["train/loss_avg"] == 0.5
+    assert records[1]["kpi"]["train/step_loss"] == 0.5
     assert records[1]["epoch"] == 1
     assert records[1]["max_epoch"] == 2
     assert records[1]["step_in_epoch"] == 3
@@ -70,12 +79,32 @@ def test_tao_status_callback_writes_training_validation_and_success(tmp_path) ->
     assert records[3]["max_validation_batches"] == 2
     assert "epoch 1/2" in records[3]["message"]
     assert records[4]["kpi"]["val/avg_loss"] == 0.25
+    assert records[4]["kpi"]["val/loss_numerator"] == 2.0
+    assert records[4]["kpi"]["val/valid_label_count"] == 8
     assert "epoch 1/2" in records[4]["message"]
+    assert records[5]["phase"] == "training_complete"
+    assert records[5]["kpi"]["train/avg_loss"] == 0.5
+    assert records[5]["train_loss_numerator"] == 6.0
+    assert records[5]["train_valid_label_count"] == 12
+    assert records[5]["kpi"]["train/valid_label_count"] == 12
     assert records[-1]["kpi"]["val/loss"] == 0.25
     assert records[-1]["epoch"] == 2
     assert records[-1]["completed_epochs"] == 2
 
 
+def test_tao_status_callback_reports_checkpoint_event(tmp_path) -> None:
+    callback = _callback(tmp_path)
+    checkpoint = tmp_path / "checkpoints" / "epoch_1"
+    checkpoint.mkdir(parents=True)
+    (tmp_path / "checkpoints" / "latest_checkpoint.txt").write_text("epoch_1\n")
+    callback.on_save_checkpoint_success(iteration=3, elapsed_time=1.25)
+    record = _records(tmp_path)[-1]
+    assert record["phase"] == "checkpoint_complete"
+    assert record["checkpoint_iteration"] == 3
+    assert record["checkpoint_elapsed_seconds"] == 1.25
+    assert record["checkpoint_path"] == str(checkpoint.resolve())
+
+
 def test_tao_status_callback_writes_failure(tmp_path) -> None:
     callback = _callback(tmp_path)
     callback.on_train_start(model=None, iteration=0)
diff --git a/cosmos_framework/configs/base/reasoner/defaults/policy_config.py b/cosmos_framework/configs/base/reasoner/defaults/policy_config.py
index 8c5adccf..526d8241 100644
--- a/cosmos_framework/configs/base/reasoner/defaults/policy_config.py
+++ b/cosmos_framework/configs/base/reasoner/defaults/policy_config.py
@@ -30,7 +30,21 @@ class PolicyConfig:
     #   0 < exponent < 1 -> interpolation; e.g. exponent=0.5 gives square-root per-token loss (Qwen3-VL)
     weighted_ce_exponent: float = 1.0
 
-    # Extra model config
+    # Parameter-efficient fine-tuning.  These fields intentionally live on
+    # the VLM policy instead of reusing the VFM/MoT model config: the two
+    # backends have different module names and checkpoint layouts.
+    lora_enabled: bool = False
+    lora_rank: int = 16
+    lora_alpha: float = 32.0
+    lora_dropout: float = 0.0
+    lora_target_modules: str = "q_proj,k_proj,v_proj,o_proj"
+    lora_bias: str = "none"
+    lora_use_rslora: bool = False
+    lora_modules_to_save: str = ""
+    lora_precision: str | None = None
+
+    # Legacy free-form field retained for config compatibility. New recipes
+    # must use the explicit fields above so PEFT equivalence can be validated.
     lora: Union[str, None] = None
     enable_liger_kernel: bool = False
     trainable_map: Union[str, None] = None
@@ -41,6 +55,27 @@ class PolicyConfig:
     # "sdpa", or "eager" for fallback.
     attn_implementation: str = "cosmos"
 
+    # Qwen3-VL's patch projection is a non-overlapping Conv3d and is therefore
+    # algebraically equivalent to a linear projection. ``auto`` selects the
+    # linear implementation on A100 (SM80), where large FSDP runs have exposed
+    # a cuDNN Conv3d backward failure.  No environment-time monkey patch is
+    # required.
+    qwen3_vl_patch_embed: str = "auto"
+
+    def __attrs_post_init__(self) -> None:
+        if self.lora_rank <= 0:
+            raise ValueError("lora_rank must be positive")
+        if self.lora_alpha <= 0:
+            raise ValueError("lora_alpha must be positive")
+        if not 0.0 <= self.lora_dropout < 1.0:
+            raise ValueError("lora_dropout must be in [0, 1)")
+        if self.lora_bias not in {"none", "all", "lora_only"}:
+            raise ValueError("lora_bias must be one of: none, all, lora_only")
+        if self.lora_precision not in {None, "float32", "float16", "bfloat16"}:
+            raise ValueError("lora_precision must be float32, float16, bfloat16, or unset")
+        if self.qwen3_vl_patch_embed not in {"auto", "linear", "conv3d"}:
+            raise ValueError("qwen3_vl_patch_embed must be auto, linear, or conv3d")
+
 
 @attrs.define(slots=False)
 class VLMModelConfig:
diff --git a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py
index 7143aea5..fc2dec1e 100644
--- a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py
+++ b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py
@@ -87,16 +87,23 @@ def __init__(
         ignore_index: int = IGNORE_INDEX,
         num_video_frames: int = 8,
         video_cache_size: int = 8,
+        video_device: str = "cuda",
+        video_num_threads: int = 1,
         system_prompt: str = "",
         use_daft_chat_template: bool = False,
     ) -> None:
         super().__init__(processor=processor, ignore_index=ignore_index)
+        num_video_frames = int(num_video_frames)
+        video_cache_size = int(video_cache_size)
+        video_num_threads = int(video_num_threads)
         if num_video_frames < 1:
             raise ValueError("num_video_frames must be >= 1")
         if video_cache_size < 0:
             raise ValueError("video_cache_size must be >= 0")
         self.num_video_frames = num_video_frames
         self.video_cache_size = video_cache_size
+        self.video_device = video_device
+        self.video_num_threads = video_num_threads
         self.system_prompt = system_prompt
         self.use_daft_chat_template = use_daft_chat_template
         if self.use_daft_chat_template:
@@ -109,7 +116,11 @@ def _decode_video(self, video_path: str) -> tuple[list[Image.Image], float]:
             self._video_cache.move_to_end(video_path)
             return cached
 
-        reader = TorchCodecVideoReader(video_path, num_threads=2)
+        reader = TorchCodecVideoReader(
+            video_path,
+            num_threads=self.video_num_threads,
+            device=self.video_device,
+        )
         total_frames = len(reader)
         if total_frames < 1:
             raise ValueError(f"WTS video has zero frames: {video_path}")
@@ -188,7 +199,7 @@ def _wts_dataloader(
                 limit=f"${{oc.env:{limit_env},''}}",
             ),
             shuffle=shuffle,
-            seed=42,
+            seed="${oc.env:TAO_DATALOADER_SEED,42}",
             name="train" if shuffle else "val",
         ),
         processor=L(WTSProcessor)(
@@ -197,9 +208,11 @@ def _wts_dataloader(
                 config_variant="hf",
             ),
             ignore_index=IGNORE_INDEX,
-            num_video_frames=8,
-            video_cache_size=8,
-            system_prompt=("You are a helpful assistant that can answer questions about a street-view CCTV footage."),
+            num_video_frames="${oc.env:WTS_NUM_VIDEO_FRAMES,8}",
+            video_cache_size="${oc.env:WTS_VIDEO_CACHE_SIZE,8}",
+            video_device="${oc.env:TAO_VIDEO_DECODER_DEVICE,cuda}",
+            video_num_threads="${oc.env:TAO_VIDEO_DECODER_THREADS,1}",
+            system_prompt="${oc.env:WTS_SYSTEM_PROMPT,''}",
         ),
         batcher=L(PoolPackingBatcher)(
             max_tokens="${data_setting.max_tokens}",
@@ -230,7 +243,7 @@ def _aetc_daft_dataloader(*, split: str, shuffle: bool) -> LazyDict:
                 max_samples=f"${{oc.env:{limit_env},''}}",
             ),
             shuffle=shuffle,
-            seed=42,
+            seed="${oc.env:TAO_DATALOADER_SEED,42}",
             name=split,
         ),
         processor=L(WTSProcessor)(
@@ -239,8 +252,10 @@ def _aetc_daft_dataloader(*, split: str, shuffle: bool) -> LazyDict:
                 config_variant="hf",
             ),
             ignore_index=IGNORE_INDEX,
-            num_video_frames=8,
-            video_cache_size=8,
+            num_video_frames="${oc.env:AETC_NUM_VIDEO_FRAMES,8}",
+            video_cache_size="${oc.env:AETC_VIDEO_CACHE_SIZE,8}",
+            video_device="${oc.env:TAO_VIDEO_DECODER_DEVICE,cuda}",
+            video_num_threads="${oc.env:TAO_VIDEO_DECODER_THREADS,1}",
             system_prompt="",
             use_daft_chat_template=True,
         ),
diff --git a/cosmos_framework/configs/toml_config/sft_config.py b/cosmos_framework/configs/toml_config/sft_config.py
index 7fda51e3..e17036fc 100644
--- a/cosmos_framework/configs/toml_config/sft_config.py
+++ b/cosmos_framework/configs/toml_config/sft_config.py
@@ -11,7 +11,7 @@
 from __future__ import annotations
 
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Literal, Optional
 
 import tomllib
 from pydantic import BaseModel, ConfigDict, Field
@@ -374,6 +374,15 @@ class ModelConfig(BaseModel):
             "adapter. Defaults target the four MoE-gen projection matrices."
         ),
     )
+    lora_dropout: float = Field(default=0.0, ge=0.0, lt=1.0)
+    lora_bias: Literal["none", "all", "lora_only"] = "none"
+    lora_use_rslora: bool = False
+    lora_modules_to_save: str = Field(
+        default="",
+        description="Comma-separated module-name suffixes trained and saved with LoRA adapters.",
+    )
+    lora_precision: Literal["float32", "float16", "bfloat16"] | None = None
+    qwen3_vl_patch_embed: Literal["auto", "linear", "conv3d"] = "auto"
 
     ema: EMAConfig = Field(default_factory=EMAConfig)
     parallelism: ParallelismConfig = Field(default_factory=ParallelismConfig)
diff --git a/cosmos_framework/configs/toml_config/sft_config_test.py b/cosmos_framework/configs/toml_config/sft_config_test.py
index 95f5b110..fe4bb9d8 100644
--- a/cosmos_framework/configs/toml_config/sft_config_test.py
+++ b/cosmos_framework/configs/toml_config/sft_config_test.py
@@ -94,6 +94,47 @@ def test_native_epoch_schedule_fields_validate(self) -> None:
         assert cfg.trainer.validation_freq_in_epoch == 1
         assert cfg.checkpoint.save_freq_in_epoch == 2
 
+    def test_vlm_peft_fields_are_explicit_and_validated(self) -> None:
+        cfg = SFTExperimentConfig.model_validate(
+            {
+                "job": {"task": "vlm", "experiment": "wts_vlm"},
+                "model": {
+                    "lora_enabled": True,
+                    "lora_rank": 8,
+                    "lora_alpha": 16,
+                    "lora_dropout": 0.05,
+                    "lora_target_modules": "q_proj,v_proj",
+                    "lora_bias": "none",
+                    "lora_use_rslora": True,
+                    "lora_modules_to_save": "lm_head",
+                    "lora_precision": "bfloat16",
+                    "qwen3_vl_patch_embed": "linear",
+                },
+            }
+        )
+        assert cfg.model.lora_enabled is True
+        assert cfg.model.lora_rank == 8
+        assert cfg.model.lora_use_rslora is True
+
+    def test_vlm_peft_fields_reach_policy_overrides(self) -> None:
+        overrides = build_hydra_overrides(
+            {
+                "job": {"task": "vlm", "experiment": "wts_vlm"},
+                "model": {
+                    "lora_enabled": True,
+                    "lora_rank": 8,
+                    "lora_dropout": 0.05,
+                    "lora_bias": "lora_only",
+                    "qwen3_vl_patch_embed": "auto",
+                },
+            }
+        )
+        assert "model.config.policy.lora_enabled=true" in overrides
+        assert "model.config.policy.lora_rank=8" in overrides
+        assert "model.config.policy.lora_dropout=0.05" in overrides
+        assert "model.config.policy.lora_bias=lora_only" in overrides
+        assert "model.config.policy.qwen3_vl_patch_embed=auto" in overrides
+
 
 # --------------------------------------------------------------------------- #
 # 2. build_hydra_overrides must NOT emit [custom] as per-leaf overrides        #
diff --git a/cosmos_framework/configs/toml_config/toml_config_helper.py b/cosmos_framework/configs/toml_config/toml_config_helper.py
index 4d1535c5..27f02ac0 100644
--- a/cosmos_framework/configs/toml_config/toml_config_helper.py
+++ b/cosmos_framework/configs/toml_config/toml_config_helper.py
@@ -52,6 +52,12 @@
         ("job", "upload_reproducible_setup"): ("upload_reproducible_setup",),
         ("model", "attn_implementation"): None,
         ("model", "backbone"): None,                                           # VLM-only — VFM has no model.config.backbone
+        ("model", "lora_dropout"): None,
+        ("model", "lora_bias"): None,
+        ("model", "lora_use_rslora"): None,
+        ("model", "lora_modules_to_save"): None,
+        ("model", "lora_precision"): None,
+        ("model", "qwen3_vl_patch_embed"): None,
         # Per-caption token cap lives on the nested SFT dataset, not a top-level
         # dataloader scalar — route it to the get_sft_dataset node.
         ("dataloader_train", "max_caption_tokens"): (
@@ -72,10 +78,16 @@
         # No VLM analog — skip these leaves
         ("model", "max_num_tokens_after_packing"): None,
         ("model", "joint_attn_implementation"): None,
-        ("model", "lora_enabled"): None,
-        ("model", "lora_rank"): None,
-        ("model", "lora_alpha"): None,
-        ("model", "lora_target_modules"): None,
+        ("model", "lora_enabled"): ("model", "config", "policy", "lora_enabled"),
+        ("model", "lora_rank"): ("model", "config", "policy", "lora_rank"),
+        ("model", "lora_alpha"): ("model", "config", "policy", "lora_alpha"),
+        ("model", "lora_dropout"): ("model", "config", "policy", "lora_dropout"),
+        ("model", "lora_target_modules"): ("model", "config", "policy", "lora_target_modules"),
+        ("model", "lora_bias"): ("model", "config", "policy", "lora_bias"),
+        ("model", "lora_use_rslora"): ("model", "config", "policy", "lora_use_rslora"),
+        ("model", "lora_modules_to_save"): ("model", "config", "policy", "lora_modules_to_save"),
+        ("model", "lora_precision"): ("model", "config", "policy", "lora_precision"),
+        ("model", "qwen3_vl_patch_embed"): ("model", "config", "policy", "qwen3_vl_patch_embed"),
         ("model", "tokenizer"): None,                                          # blocks model.tokenizer.*
         ("dataloader_train", "seed"): None,
         ("optimizer", "eps"): None,                                            # VLM_OPTIMIZER_KWARGS has no eps field
@@ -208,5 +220,3 @@ def _hydra_format(v: Any, in_list: bool = False) -> str:
             return f"'{v}'"
         return v
     return str(v)
-
-
diff --git a/cosmos_framework/model/generator/algorithm/loss/cross_entropy.py b/cosmos_framework/model/generator/algorithm/loss/cross_entropy.py
index 283b1366..fd12e091 100644
--- a/cosmos_framework/model/generator/algorithm/loss/cross_entropy.py
+++ b/cosmos_framework/model/generator/algorithm/loss/cross_entropy.py
@@ -42,7 +42,8 @@ def cross_entropy_loss(
     dp_group: dist.ProcessGroup | None = None,
     cp_group: dist.ProcessGroup | None = None,
     ignore_index: int = IGNORE_INDEX,
-) -> torch.Tensor:
+    return_stats: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Next-token-prediction CE loss with DP/CP group reduction.
 
     Matches the behavior of cosmos_rl.policy.trainer.llm_trainer.sft_trainer.async_safe_ce
@@ -79,7 +80,14 @@ def cross_entropy_loss(
             reduction="mean",
         )
         loss = torch.nan_to_num(loss, nan=0.0)
-        return loss * loss_scaling_factor
+        result = loss * loss_scaling_factor
+        if return_stats:
+            per_token_loss = F.cross_entropy(
+                shifted_logits.float(), shifted_labels, ignore_index=ignore_index, reduction="none"
+            )
+            valid = shifted_labels != ignore_index
+            return result, per_token_loss[valid].sum().detach(), valid.sum().detach()
+        return result
 
     # No-CP path: per-token loss, then normalize over the global valid-token count.
     # Reference: async_safe_ce:89-109
@@ -89,13 +97,18 @@ def cross_entropy_loss(
         ignore_index=ignore_index,
         reduction="none",
     )
-    n_valid_tokens = (shifted_labels != ignore_index).sum()
+    valid_mask = shifted_labels != ignore_index
+    local_numerator = per_token_loss[valid_mask].sum()
+    local_denominator = valid_mask.sum()
+    n_valid_tokens = local_denominator.detach().clone()
     num_dp_workers = 1
     if dp_group is not None:
         dist.all_reduce(n_valid_tokens, op=dist.ReduceOp.SUM, group=dp_group)
         num_dp_workers = dist.get_world_size(group=dp_group)
 
-    loss = per_token_loss.sum() / (n_valid_tokens + 1e-8) * (num_dp_workers * loss_scaling_factor)
+    loss = local_numerator / (n_valid_tokens + 1e-8) * (num_dp_workers * loss_scaling_factor)
+    if return_stats:
+        return loss, local_numerator.detach(), local_denominator.detach()
     return loss
 
 
@@ -107,7 +120,8 @@ def weighted_cross_entropy_loss(
     dp_group: dist.ProcessGroup | None = None,
     cp_group: dist.ProcessGroup | None = None,
     ignore_index: int = IGNORE_INDEX,
-) -> torch.Tensor:
+    return_stats: bool = False,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Next-token-prediction CE loss interpolated between per-token and per-sample reductions.
 
     Matches ``cosmos_rl.policy.trainer.llm_trainer.sft_trainer.async_safe_weighted_ce``
@@ -156,4 +170,8 @@ def weighted_cross_entropy_loss(
         num_dp_workers = dist.get_world_size()
 
     loss = local_loss_sum / local_exp_weight_sum.clamp(min=1) * (num_dp_workers * loss_scaling_factor)  # []
+    if return_stats:
+        token_numerator = (per_token_loss * valid_mask).sum().detach()
+        token_denominator = valid_mask.sum().detach()
+        return loss, token_numerator, token_denominator
     return loss
diff --git a/cosmos_framework/model/generator/algorithm/loss/cross_entropy_stats_test.py b/cosmos_framework/model/generator/algorithm/loss/cross_entropy_stats_test.py
new file mode 100644
index 00000000..96cf2f20
--- /dev/null
+++ b/cosmos_framework/model/generator/algorithm/loss/cross_entropy_stats_test.py
@@ -0,0 +1,38 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+from cosmos_framework.model.generator.algorithm.loss.cross_entropy import (
+    cross_entropy_loss,
+    weighted_cross_entropy_loss,
+)
+
+
+def test_cross_entropy_returns_token_numerator_and_denominator() -> None:
+    torch.manual_seed(7)
+    logits = torch.randn(2, 4, 5, requires_grad=True)
+    labels = torch.tensor([[0, 1, 2, -100], [0, 2, 3, 4]])
+    loss, numerator, denominator = cross_entropy_loss(logits, labels, return_stats=True)
+    shifted_logits = logits[:, :-1].reshape(-1, 5)
+    shifted_labels = labels[:, 1:].reshape(-1)
+    expected = F.cross_entropy(shifted_logits.float(), shifted_labels, ignore_index=-100, reduction="none")
+    valid = shifted_labels != -100
+    assert denominator.item() == int(valid.sum())
+    torch.testing.assert_close(numerator, expected[valid].sum())
+    torch.testing.assert_close(loss, numerator / denominator)
+
+
+def test_weighted_objective_keeps_unweighted_token_stats() -> None:
+    torch.manual_seed(11)
+    logits = torch.randn(2, 4, 5)
+    labels = torch.tensor([[0, 1, -100, -100], [0, 2, 3, 4]])
+    objective, numerator, denominator = weighted_cross_entropy_loss(
+        logits, labels, exponent=1.0, return_stats=True
+    )
+    assert objective.ndim == 0
+    assert denominator.item() == 4
+    assert numerator.item() > 0
diff --git a/cosmos_framework/model/generator/qwen3_vl_compat.py b/cosmos_framework/model/generator/qwen3_vl_compat.py
new file mode 100644
index 00000000..0cf254c6
--- /dev/null
+++ b/cosmos_framework/model/generator/qwen3_vl_compat.py
@@ -0,0 +1,62 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Repository-owned Qwen3-VL training compatibility helpers."""
+
+from __future__ import annotations
+
+import types
+
+import torch
+import torch.nn.functional as F
+
+
+def _linear_patch_embed_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    patch_elements = self.in_channels * self.temporal_patch_size * self.patch_size * self.patch_size
+    flat_input = hidden_states.reshape(-1, patch_elements).to(dtype=self.proj.weight.dtype)
+    flat_weight = self.proj.weight.reshape(self.embed_dim, patch_elements)
+    return F.linear(flat_input, flat_weight, self.proj.bias)
+
+
+_linear_patch_embed_forward._tao_channels_last_3d = True  # type: ignore[attr-defined]
+
+
+def should_use_linear_patch_embed(mode: str, *, device_capability: tuple[int, int] | None = None) -> bool:
+    """Resolve ``auto|linear|conv3d`` without consulting environment shims."""
+    if mode not in {"auto", "linear", "conv3d"}:
+        raise ValueError(f"Unsupported Qwen3-VL patch-embed mode: {mode!r}")
+    if mode == "linear":
+        return True
+    if mode == "conv3d":
+        return False
+    if device_capability is None and torch.cuda.is_available():
+        device_capability = torch.cuda.get_device_capability()
+    return device_capability == (8, 0)
+
+
+def apply_qwen3_vl_patch_embed_compat(
+    model: torch.nn.Module,
+    *,
+    model_type: str,
+    mode: str,
+    device_capability: tuple[int, int] | None = None,
+) -> bool:
+    """Use the A100-safe algebraic projection for Qwen3-VL patch embeds.
+
+    Returns whether a replacement was applied. The original Conv3d parameters
+    and state-dict keys are unchanged, so checkpoint compatibility is exact.
+    """
+    if model_type not in {"qwen3_vl", "qwen3_vl_moe"}:
+        return False
+    if not should_use_linear_patch_embed(mode, device_capability=device_capability):
+        return False
+
+    patched = 0
+    for module in model.modules():
+        required = ("in_channels", "temporal_patch_size", "patch_size", "embed_dim", "proj")
+        if all(hasattr(module, name) for name in required) and isinstance(module.proj, torch.nn.Conv3d):
+            module.forward = types.MethodType(_linear_patch_embed_forward, module)
+            patched += 1
+    if patched != 1:
+        raise RuntimeError(f"Expected one Qwen3-VL PatchEmbed module, found {patched}")
+    return True
diff --git a/cosmos_framework/model/generator/qwen3_vl_compat_test.py b/cosmos_framework/model/generator/qwen3_vl_compat_test.py
new file mode 100644
index 00000000..73c5f8c1
--- /dev/null
+++ b/cosmos_framework/model/generator/qwen3_vl_compat_test.py
@@ -0,0 +1,50 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+from __future__ import annotations
+
+import torch
+
+from cosmos_framework.model.generator.qwen3_vl_compat import (
+    apply_qwen3_vl_patch_embed_compat,
+    should_use_linear_patch_embed,
+)
+
+
+class _PatchEmbed(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.in_channels = 3
+        self.temporal_patch_size = 2
+        self.patch_size = 2
+        self.embed_dim = 5
+        self.proj = torch.nn.Conv3d(3, 5, kernel_size=(2, 2, 2), stride=(2, 2, 2), bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.view(-1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size)
+        return self.proj(x).view(-1, self.embed_dim)
+
+
+def test_auto_selects_linear_only_on_sm80() -> None:
+    assert should_use_linear_patch_embed("auto", device_capability=(8, 0))
+    assert not should_use_linear_patch_embed("auto", device_capability=(9, 0))
+    assert should_use_linear_patch_embed("linear", device_capability=(9, 0))
+    assert not should_use_linear_patch_embed("conv3d", device_capability=(8, 0))
+
+
+def test_linear_patch_is_algebraically_equivalent_and_preserves_keys() -> None:
+    model = torch.nn.Sequential(_PatchEmbed())
+    values = torch.randn(4, 3 * 2 * 2 * 2)
+    expected = model(values)
+    keys_before = set(model.state_dict())
+
+    changed = apply_qwen3_vl_patch_embed_compat(
+        model,
+        model_type="qwen3_vl",
+        mode="linear",
+        device_capability=(8, 0),
+    )
+
+    assert changed is True
+    torch.testing.assert_close(model(values), expected)
+    assert set(model.state_dict()) == keys_before
diff --git a/cosmos_framework/model/generator/vlm_model.py b/cosmos_framework/model/generator/vlm_model.py
index 2d33fa4f..51bdf37d 100644
--- a/cosmos_framework/model/generator/vlm_model.py
+++ b/cosmos_framework/model/generator/vlm_model.py
@@ -301,8 +301,17 @@ def __init__(self, config: VLMModelConfig, checkpoint):
                 self.model.model.sound_und_model.projector.eval()
             n_trainable = sum(parameter.requires_grad for parameter in self.model.model.parameters())
             assert n_trainable > 0, "audio freeze policy left 0 trainable parameters — check freeze patterns"
+        trainable_parameters = sum(parameter.numel() for parameter in self.model.parameters() if parameter.requires_grad)
+        total_parameters = sum(parameter.numel() for parameter in self.model.parameters())
+        self.parameter_summary = {
+            "training_mode": "peft" if self.config.policy.lora_enabled else "dense_sft",
+            "trainable_parameters": trainable_parameters,
+            "total_parameters": total_parameters,
+            "frozen_parameters": total_parameters - trainable_parameters,
+            "trainable_parameter_tensors": n_trainable,
+        }
         log.info(
-            f"freeze config applied (model_type={self.hf_config.model_type}): {n_trainable} trainable parameter tensors"
+            f"freeze config applied (model_type={self.hf_config.model_type}): {self.parameter_summary}"
         )
 
         dp_group = None
@@ -381,6 +390,32 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None:
             # a cache directory and no longer identifies the Edge Reasoner.
             configured_model_name_or_path=policy.backbone.model_name,
         )
+
+        from cosmos_framework.model.generator.qwen3_vl_compat import apply_qwen3_vl_patch_embed_compat
+
+        patch_embed_changed = apply_qwen3_vl_patch_embed_compat(
+            hf_model.model,
+            model_type=hf_model.hf_config.model_type,
+            mode=policy.qwen3_vl_patch_embed,
+        )
+        if patch_embed_changed:
+            log.info("Using repository-owned linear Qwen3-VL PatchEmbed compatibility path")
+
+        if policy.lora_enabled:
+            from cosmos_framework.utils.generator.lora import inject_lora_pre_fsdp
+
+            inject_lora_pre_fsdp(
+                hf_model.model,
+                lora_rank=policy.lora_rank,
+                lora_alpha=policy.lora_alpha,
+                lora_dropout=policy.lora_dropout,
+                lora_target_modules=policy.lora_target_modules,
+                lora_bias=policy.lora_bias,
+                lora_use_rslora=policy.lora_use_rslora,
+                lora_modules_to_save=policy.lora_modules_to_save,
+                lora_precision=policy.lora_precision,
+            )
+
         # ── b.1. Early family-gate for backbone.pretrained_weights ──
         # Fail-fast on unsupported VLM families BEFORE any expensive work
         # (parallelize, materialize, base-weight load, overlay download).
@@ -479,10 +514,12 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None:
             else:
                 safetensors_local_path = local_path
 
+            base_skip_patterns = [r".*\.lora_[AB]\.weight"] if policy.lora_enabled else None
             hf_model.load_weights(
                 checkpoint_path=safetensors_local_path,
                 credential_path=None,  # local path after download
                 parallel_dims=parallel_dims if torch.distributed.is_initialized() else None,
+                extra_skip_patterns=base_skip_patterns,
             )
 
             # ── g.2. Optional LLM overlay (backbone.pretrained_weights) ──
@@ -508,7 +545,7 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None:
                     checkpoint_path=llm_local_path,
                     credential_path=None,
                     parallel_dims=parallel_dims if torch.distributed.is_initialized() else None,
-                    extra_skip_patterns=overlay_skip_patterns,
+                    extra_skip_patterns=overlay_skip_patterns + (base_skip_patterns or []),
                 )
                 lm_loaded = {k for k in keys_loaded if is_lm_key(k)}
                 if not lm_loaded:
@@ -537,6 +574,11 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None:
                 f"{audio_config.encoder_checkpoint_path}"
             )
 
+        if policy.lora_enabled:
+            from cosmos_framework.utils.generator.lora import init_lora_weights_post_materialization
+
+            init_lora_weights_post_materialization(hf_model.model)
+
         # ── i. Gradient checkpointing ──
         # HF backbone supports only binary on/off via gradient_checkpointing_enable,
         # so VLMActivationCheckpointingConfig.mode is restricted to {"full", "none"}.
@@ -548,6 +590,14 @@ def _init_vlm(self, config: VLMModelConfig, checkpoint) -> None:
         self.model_name_or_path = local_path
         self.hf_config = hf_model.hf_config
 
+        trainable = sum(p.numel() for p in hf_model.parameters() if p.requires_grad)
+        total = sum(p.numel() for p in hf_model.parameters())
+        log.info(
+            "parameter summary: "
+            f"mode={'peft' if policy.lora_enabled else 'dense_sft'}, "
+            f"trainable={trainable}, total={total}, frozen={total - trainable}"
+        )
+
     def on_train_start(self, memory_format) -> None:
         """Called by trainer after model.to("cuda"). No device move needed here."""
 
@@ -638,7 +688,7 @@ def training_step(self, data: dict, iteration: int) -> tuple[dict, torch.Tensor]
         labels = data.pop("labels")
         data.pop("attention_mask", None)
         logits = self.model(**data)
-        loss = self._loss_fn(logits, labels)
+        loss, loss_numerator, loss_denominator = self._loss_fn(logits, labels, return_stats=True)
 
         # loss_avg: DP-averaged loss for logging (matches cosmos-rl ReduceOp.AVG).
         # Does not affect the backward scalar. Pick the same 1-D sub-mesh the
@@ -656,7 +706,13 @@ def training_step(self, data: dict, iteration: int) -> tuple[dict, torch.Tensor]
         if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
             log.info(f"train/loss_avg: {loss_avg.item():.5f} (iteration {iteration})")
 
-        return {"loss": loss, "loss_avg": loss_avg, "labels": labels}, loss
+        return {
+            "loss": loss,
+            "loss_avg": loss_avg,
+            "loss_numerator": loss_numerator,
+            "loss_denominator": loss_denominator,
+            "labels": labels,
+        }, loss
 
     @torch.no_grad()
     def validation_step(self, data: dict, iteration: int) -> tuple[dict, torch.Tensor]:
@@ -675,5 +731,10 @@ def validation_step(self, data: dict, iteration: int) -> tuple[dict, torch.Tenso
         labels = data.pop("labels")
         data.pop("attention_mask", None)
         logits = self.model(**data)
-        loss = self._loss_fn(logits, labels)
-        return {"loss": loss, "labels": labels}, loss
+        loss, loss_numerator, loss_denominator = self._loss_fn(logits, labels, return_stats=True)
+        return {
+            "loss": loss,
+            "loss_numerator": loss_numerator,
+            "loss_denominator": loss_denominator,
+            "labels": labels,
+        }, loss
diff --git a/cosmos_framework/scripts/export_vlm_dcp.py b/cosmos_framework/scripts/export_vlm_dcp.py
new file mode 100644
index 00000000..5cdef6e9
--- /dev/null
+++ b/cosmos_framework/scripts/export_vlm_dcp.py
@@ -0,0 +1,250 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Export a Framework Cosmos3 VLM DCP checkpoint to native HF safetensors.
+
+This is intentionally repository-owned: action images and the TAO skill call
+this entry point, but do not carry a private copy or patch checkpoint keys.
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+from pathlib import Path
+from typing import Any
+
+_VLM_MODEL_TARGET = "cosmos_framework.model.generator.vlm_model.VLMModel"
+
+
+def _read_config(config_file: str | Path) -> dict[str, Any]:
+    path = Path(config_file)
+    if path.suffix.lower() == ".json":
+        return json.loads(path.read_text(encoding="utf-8"))
+    import yaml
+
+    value = yaml.safe_load(path.read_text(encoding="utf-8"))
+    if not isinstance(value, dict):
+        raise ValueError(f"Expected an object in Framework config {path}")
+    return value
+
+
+def is_vlm_training_config(config_file: str | Path) -> bool:
+    try:
+        return _VLM_MODEL_TARGET in Path(config_file).read_text(encoding="utf-8")
+    except OSError:
+        return False
+
+
+def infer_vlm_base_model(config_file: str | Path) -> str | None:
+    try:
+        backbone = _read_config(config_file)["model"]["config"]["policy"]["backbone"]
+    except (OSError, ValueError, TypeError, KeyError):
+        return None
+    for key in ("model_name", "safetensors_path"):
+        value = backbone.get(key)
+        if isinstance(value, str) and value.strip():
+            return value.strip()
+    return None
+
+
+def infer_vlm_lora_config(config_file: str | Path) -> dict[str, Any]:
+    try:
+        policy = _read_config(config_file)["model"]["config"]["policy"]
+    except (OSError, ValueError, TypeError, KeyError):
+        return {"enabled": False}
+    return {
+        "enabled": bool(policy.get("lora_enabled", False)),
+        "rank": int(policy.get("lora_rank", 16)),
+        "alpha": float(policy.get("lora_alpha", 32.0)),
+        "dropout": float(policy.get("lora_dropout", 0.0)),
+        "target_modules": str(policy.get("lora_target_modules", "q_proj,k_proj,v_proj,o_proj")),
+        "bias": str(policy.get("lora_bias", "none")),
+        "use_rslora": bool(policy.get("lora_use_rslora", False)),
+        "modules_to_save": str(policy.get("lora_modules_to_save", "")),
+        "precision": policy.get("lora_precision"),
+    }
+
+
+def _checkpoint_model_dir(checkpoint_path: Path) -> Path:
+    for candidate in (checkpoint_path / "model", checkpoint_path):
+        if (candidate / ".metadata").is_file():
+            return candidate
+    raise ValueError(f"No Framework model DCP metadata found below {checkpoint_path}")
+
+
+def _sha256_file(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as stream:
+        for block in iter(lambda: stream.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def fingerprint_model_files(path: str | Path) -> dict[str, Any]:
+    root = Path(path).expanduser()
+    if not root.exists():
+        return {"source": str(path), "kind": "uri"}
+    root = root.resolve()
+    names = (
+        "config.json",
+        "generation_config.json",
+        "tokenizer.json",
+        "tokenizer_config.json",
+        "processor_config.json",
+        "preprocessor_config.json",
+        "chat_template.json",
+    )
+    files = {
+        name: _sha256_file(root / name)
+        for name in names
+        if (root / name).is_file()
+    }
+    weights = sorted(root.glob("*.safetensors"))
+    files.update({item.name: _sha256_file(item) for item in weights})
+    manifest_bytes = json.dumps(files, sort_keys=True, separators=(",", ":")).encode()
+    return {
+        "source": str(path),
+        "resolved": str(root),
+        "kind": "local",
+        "files": files,
+        "sha256": hashlib.sha256(manifest_bytes).hexdigest(),
+    }
+
+
+def export_vlm_dcp(
+    checkpoint_path: str | Path,
+    *,
+    config_file: str | Path,
+    output_dir: str | Path,
+    base_model_path_or_uri: str | None = None,
+    base_model_revision: str | None = None,
+    dtype: str = "bfloat16",
+) -> dict[str, Any]:
+    import torch
+    import torch.distributed.checkpoint as dcp
+    from torch.distributed.checkpoint.filesystem import FileSystemReader
+    from torch.distributed.checkpoint.state_dict import get_model_state_dict
+    from transformers import AutoConfig, AutoModelForImageTextToText, AutoProcessor
+
+    checkpoint_root = Path(checkpoint_path).expanduser().resolve(strict=True)
+    dcp_dir = _checkpoint_model_dir(checkpoint_root)
+    config_path = Path(config_file).expanduser().resolve(strict=True)
+    export_path = Path(output_dir).expanduser().resolve()
+    model_source = base_model_path_or_uri or infer_vlm_base_model(config_path)
+    if not model_source:
+        raise ValueError("base_model_path_or_uri is required when it cannot be inferred from the saved config")
+    local_source = Path(model_source).expanduser()
+    if local_source.exists():
+        model_source = str(local_source.resolve())
+
+    import cosmos_framework.model.generator.reasoner.cosmos3_edge  # noqa: F401
+
+    kwargs: dict[str, Any] = {
+        "dtype": getattr(torch, dtype),
+        "device_map": "cpu",
+        "trust_remote_code": True,
+    }
+    if base_model_revision:
+        kwargs["revision"] = base_model_revision
+    model = AutoModelForImageTextToText.from_pretrained(model_source, **kwargs)
+
+    lora = infer_vlm_lora_config(config_path)
+    if lora["enabled"]:
+        from cosmos_framework.utils.generator.lora import inject_lora_pre_fsdp
+
+        inject_lora_pre_fsdp(
+            model,
+            lora_rank=lora["rank"],
+            lora_alpha=lora["alpha"],
+            lora_dropout=lora["dropout"],
+            lora_target_modules=lora["target_modules"],
+            lora_bias=lora["bias"],
+            lora_use_rslora=lora["use_rslora"],
+            lora_modules_to_save=lora["modules_to_save"],
+            lora_precision=lora["precision"],
+        )
+
+    outer = torch.nn.Module()
+    inner = torch.nn.Module()
+    inner.model = model
+    outer.model = inner
+    state_dict = get_model_state_dict(outer)
+    reader = FileSystemReader(str(dcp_dir))
+    checkpoint_keys = set(reader.read_metadata().state_dict_metadata)
+    model_keys = set(state_dict)
+    if checkpoint_keys != model_keys:
+        missing = sorted(model_keys - checkpoint_keys)
+        unexpected = sorted(checkpoint_keys - model_keys)
+        raise RuntimeError(
+            "Framework VLM DCP exact-key mismatch: "
+            f"missing={missing[:10]}, unexpected={unexpected[:10]}, "
+            f"model_keys={len(model_keys)}, checkpoint_keys={len(checkpoint_keys)}"
+        )
+    dcp.load(state_dict=state_dict, storage_reader=reader)
+
+    merged_adapters = 0
+    if lora["enabled"]:
+        from cosmos_framework.utils.generator.lora import merge_and_strip_lora_adapters_
+
+        merged_adapters = merge_and_strip_lora_adapters_(model)
+        if merged_adapters == 0:
+            raise RuntimeError("PEFT checkpoint was restored but no LoRA adapters were merged")
+
+    export_path.mkdir(parents=True, exist_ok=True)
+    model.save_pretrained(export_path, safe_serialization=True, max_shard_size="4GB")
+    config_kwargs = {"trust_remote_code": True}
+    if base_model_revision:
+        config_kwargs["revision"] = base_model_revision
+    model_type = AutoConfig.from_pretrained(model_source, **config_kwargs).model_type
+    if model_type == "cosmos3_edge":
+        from cosmos_framework.data.generator.processors.cosmos3_edge_processing import build_cosmos3_edge_processor
+
+        processor = build_cosmos3_edge_processor(model_source)
+    else:
+        processor = AutoProcessor.from_pretrained(model_source, **config_kwargs)
+    processor.save_pretrained(export_path)
+
+    manifest = {
+        "format": "cosmos-framework-vlm-dcp",
+        "checkpoint": str(checkpoint_root),
+        "checkpoint_metadata_sha256": _sha256_file(dcp_dir / ".metadata"),
+        "config": str(config_path),
+        "config_sha256": _sha256_file(config_path),
+        "base_model_path_or_uri": base_model_path_or_uri or model_source,
+        "base_model_revision": base_model_revision,
+        "base_model_fingerprint": fingerprint_model_files(model_source),
+        "tensor_count": len(state_dict),
+        "lora": lora,
+        "merged_adapters": merged_adapters,
+    }
+    (export_path / "export_manifest.json").write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
+    (export_path / "checkpoint.json").write_text(
+        json.dumps({"checkpoint_path": str(checkpoint_root), "checkpoint_type": "vlm_dcp"}, indent=2) + "\n"
+    )
+    return manifest
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint-path", required=True)
+    parser.add_argument("--config-file", required=True)
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--base-model-path-or-uri")
+    parser.add_argument("--base-model-revision")
+    parser.add_argument("--dtype", default="bfloat16", choices=("float32", "float16", "bfloat16"))
+    args = parser.parse_args()
+    export_vlm_dcp(
+        args.checkpoint_path,
+        config_file=args.config_file,
+        output_dir=args.output_dir,
+        base_model_path_or_uri=args.base_model_path_or_uri,
+        base_model_revision=args.base_model_revision,
+        dtype=args.dtype,
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/cosmos_framework/scripts/export_vlm_dcp_test.py b/cosmos_framework/scripts/export_vlm_dcp_test.py
new file mode 100644
index 00000000..98905b77
--- /dev/null
+++ b/cosmos_framework/scripts/export_vlm_dcp_test.py
@@ -0,0 +1,43 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+from __future__ import annotations
+
+import json
+
+from cosmos_framework.scripts.export_vlm_dcp import (
+    fingerprint_model_files,
+    infer_vlm_base_model,
+    infer_vlm_lora_config,
+)
+
+
+def test_export_inputs_are_inferred_without_rewriting_user_path(tmp_path) -> None:
+    supplied = tmp_path / "model-source"
+    supplied.mkdir()
+    (supplied / "config.json").write_text('{"model_type":"qwen3_vl"}\n')
+    config = tmp_path / "config.json"
+    config.write_text(
+        json.dumps(
+            {
+                "_target_": "cosmos_framework.model.generator.vlm_model.VLMModel",
+                "model": {
+                    "config": {
+                        "policy": {
+                            "backbone": {"model_name": str(supplied)},
+                            "lora_enabled": True,
+                            "lora_rank": 8,
+                            "lora_alpha": 16,
+                            "lora_target_modules": "q_proj,v_proj",
+                        }
+                    }
+                },
+            }
+        )
+    )
+    assert infer_vlm_base_model(config) == str(supplied)
+    assert infer_vlm_lora_config(config)["rank"] == 8
+    fingerprint = fingerprint_model_files(supplied)
+    assert fingerprint["source"] == str(supplied)
+    assert fingerprint["resolved"] == str(supplied.resolve())
+    assert fingerprint["files"]["config.json"]
diff --git a/cosmos_framework/scripts/train.py b/cosmos_framework/scripts/train.py
index 8155ef6c..fd4bef43 100644
--- a/cosmos_framework/scripts/train.py
+++ b/cosmos_framework/scripts/train.py
@@ -229,7 +229,7 @@ def launch(config: Config, args: argparse.Namespace) -> None:
         raise
 
 
-if __name__ == "__main__":
+def main() -> int:
     parser = argparse.ArgumentParser(description="SFT training (structured TOML)")
     parser.add_argument(
         "--sft-toml",
@@ -283,49 +283,41 @@ def launch(config: Config, args: argparse.Namespace) -> None:
     if args.deterministic:
         _setup_deterministic_env_and_backends()
 
-    config = load_experiment_from_toml(args.sft_toml, extra_overrides=args.opts)
-
-    # log_reproducible_setup reads args.config for telemetry; this entrypoint
-    # only takes --sft-toml, so alias it so the launch info records the TOML.
-    args.config = args.sft_toml
-
-    if args.dryrun:
-        logging.info("Config:\n" + config.pretty_print(use_color=True))
-        os.makedirs(config.job.path_local, exist_ok=True)
-        try:
-            to_yaml(config, f"{config.job.path_local}/config.yaml")
-        except Exception:
-            logging.error("to_yaml failed, falling back to LazyConfig.save_yaml:")
-            logging.error(f"Traceback: {traceback.format_exc()}")
-            LazyConfig.save_yaml(config, f"{config.job.path_local}/config.yaml")
-        print(f"{config.job.path_local}/config.yaml")
-    else:
-        launch(config, args)
-        # Training finished. Exit without running interpreter finalization.
-        #
-        # A native thread outlives training: at ``atexit`` time
-        # ``threading.enumerate()`` reports only MainThread (plus wandb's daemon
-        # thread), yet finalization then faults with
-        #   Fatal Python error: PyGILState_Release: thread state ... must be
-        #   current when releasing
-        # on some hosts, and blocks forever on others -- in both cases *after*
-        # "Done with training." has been logged and every result produced. The
-        # streaming HuggingFace dataloader is what brings the thread in (a
-        # map-style dataset shuts down cleanly), but nothing reachable from
-        # Python owns it, so there is no handle to close. Not finalizing is what
-        # makes shutdown deterministic.
-        #
-        # ``os._exit`` skips atexit hooks and buffered-file flushing, so flush
-        # explicitly first. Only reached on success: an exception from
-        # ``launch`` propagates normally and keeps the usual traceback and
-        # non-zero exit.
-        #
-        # Opt-in, because skipping atexit is not free: anything registered there
-        # stops running, most visibly wandb's final flush/upload. Off by default
-        # so ordinary training keeps its normal shutdown; the launch regression
-        # sets it because that job is otherwise killed by the CI timeout.
-        if os.environ.get("COSMOS_EXIT_WITHOUT_FINALIZE", "").lower() in ("1", "true"):
-            logging.info("COSMOS_EXIT_WITHOUT_FINALIZE set: exiting without interpreter finalization.")
-            sys.stdout.flush()
-            sys.stderr.flush()
-            os._exit(0)
+    try:
+        config = load_experiment_from_toml(args.sft_toml, extra_overrides=args.opts)
+
+        # log_reproducible_setup reads args.config for telemetry; this entrypoint
+        # only takes --sft-toml, so alias it so the launch info records the TOML.
+        args.config = args.sft_toml
+
+        if args.dryrun:
+            logging.info("Config:\n" + config.pretty_print(use_color=True))
+            os.makedirs(config.job.path_local, exist_ok=True)
+            try:
+                to_yaml(config, f"{config.job.path_local}/config.yaml")
+            except Exception:
+                logging.error("to_yaml failed, falling back to LazyConfig.save_yaml:")
+                logging.error(f"Traceback: {traceback.format_exc()}")
+                LazyConfig.save_yaml(config, f"{config.job.path_local}/config.yaml")
+            print(f"{config.job.path_local}/config.yaml")
+        else:
+            launch(config, args)
+            # Some streaming dataloaders leave an unreachable native thread
+            # alive after successful training. This opt-in path preserves the
+            # established workaround while keeping normal atexit behavior by
+            # default. Failures still propagate through the exception handler.
+            if os.environ.get("COSMOS_EXIT_WITHOUT_FINALIZE", "").lower() in ("1", "true"):
+                logging.info("COSMOS_EXIT_WITHOUT_FINALIZE set: exiting without interpreter finalization.")
+                sys.stdout.flush()
+                sys.stderr.flush()
+                os._exit(0)
+    except BaseException as error:
+        from cosmos_framework.callbacks.tao_status import write_early_failure
+
+        write_early_failure(error)
+        raise
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/cosmos_framework/utils/easy_io/backends/__init__.py b/cosmos_framework/utils/easy_io/backends/__init__.py
index 8a5a806d..5211dedb 100644
--- a/cosmos_framework/utils/easy_io/backends/__init__.py
+++ b/cosmos_framework/utils/easy_io/backends/__init__.py
@@ -17,10 +17,15 @@
 ]
 
 if TRAINING:
-    from cosmos_framework.utils.easy_io.backends.boto3_backend import Boto3Backend
-    from cosmos_framework.utils.easy_io.backends.msc_backend import MSCBackend
-
-    __all__ += [
-        "Boto3Backend",
-        "MSCBackend",
-    ]
+    try:
+        from cosmos_framework.utils.easy_io.backends.boto3_backend import Boto3Backend
+    except ImportError:
+        pass
+    else:
+        __all__.append("Boto3Backend")
+    try:
+        from cosmos_framework.utils.easy_io.backends.msc_backend import MSCBackend
+    except ImportError:
+        pass
+    else:
+        __all__.append("MSCBackend")
diff --git a/cosmos_framework/utils/easy_io/backends/registry_utils.py b/cosmos_framework/utils/easy_io/backends/registry_utils.py
index 94cd1853..11581cf5 100644
--- a/cosmos_framework/utils/easy_io/backends/registry_utils.py
+++ b/cosmos_framework/utils/easy_io/backends/registry_utils.py
@@ -115,8 +115,14 @@ def _register(backend_cls):
 register_backend("http", HTTPBackend, prefixes=["http", "https"])
 
 if TRAINING:
-    from cosmos_framework.utils.easy_io.backends.msc_backend import MSCBackend
-
-    # To avoid breaking backward Compatibility, 's3' is also used as a
-    # prefix for MSCBackend
-    register_backend("s3", MSCBackend, prefixes=["s3"])
+    try:
+        from cosmos_framework.utils.easy_io.backends.msc_backend import MSCBackend
+    except ImportError as error:
+        # multistorageclient is an optional internal S3 dependency. Local,
+        # HTTP, Docker, and Pyxis training must remain importable without it.
+        # A later s3:// access fails clearly because no S3 backend is present.
+        MSC_IMPORT_ERROR = error
+    else:
+        # To avoid breaking backward Compatibility, 's3' is also used as a
+        # prefix for MSCBackend.
+        register_backend("s3", MSCBackend, prefixes=["s3"])
diff --git a/cosmos_framework/utils/easy_io/file_client.py b/cosmos_framework/utils/easy_io/file_client.py
index 650489ff..cb3af4c5 100644
--- a/cosmos_framework/utils/easy_io/file_client.py
+++ b/cosmos_framework/utils/easy_io/file_client.py
@@ -73,11 +73,14 @@ class FileClient:
     }
 
     if TRAINING:
-        from cosmos_framework.utils.easy_io.backends.msc_backend import MSCBackend
-
-        _backends["s3"] = MSCBackend
-        _backends["msc"] = MSCBackend
-        _prefix_to_backends["s3"] = MSCBackend
+        try:
+            from cosmos_framework.utils.easy_io.backends.msc_backend import MSCBackend
+        except ImportError:
+            pass
+        else:
+            _backends["s3"] = MSCBackend
+            _backends["msc"] = MSCBackend
+            _prefix_to_backends["s3"] = MSCBackend
 
     _instances: dict = {}
 
diff --git a/cosmos_framework/utils/generator/lora.py b/cosmos_framework/utils/generator/lora.py
index bbd3215d..9bc95e1b 100644
--- a/cosmos_framework/utils/generator/lora.py
+++ b/cosmos_framework/utils/generator/lora.py
@@ -34,14 +34,24 @@ class LoraInjectedLinear(nn.Linear):
     Forward computes ``y = base(x) + (alpha / r) * lora_B(lora_A(x))``.
     """
 
-    def __init__(self, base: nn.Linear, rank: int, alpha: int) -> None:
+    def __init__(
+        self,
+        base: nn.Linear,
+        rank: int,
+        alpha: float,
+        *,
+        dropout: float = 0.0,
+        use_rslora: bool = False,
+        adapter_dtype: torch.dtype | None = None,
+    ) -> None:
         # Reuse base's geometry. Inherit nn.Linear so ``super().forward(x)``
         # dispatches to the standard F.linear path.
         super().__init__(
             base.in_features,
             base.out_features,
             bias=base.bias is not None,
-            device="meta",
+            device=base.weight.device,
+            dtype=base.weight.dtype,
         )
         # Replace nn.Linear's freshly-allocated parameters with the base's
         # existing ones. On meta device this is a no-op for memory; on a
@@ -51,20 +61,43 @@ def __init__(self, base: nn.Linear, rank: int, alpha: int) -> None:
             self.bias = base.bias
         # Sibling submodules. Using bias=False nn.Linear gives us a single
         # ``weight`` parameter and a clean ``.lora_A.weight`` state-dict key.
-        self.lora_A = nn.Linear(base.in_features, rank, bias=False, device="meta")
-        self.lora_B = nn.Linear(rank, base.out_features, bias=False, device="meta")
+        adapter_dtype = adapter_dtype or base.weight.dtype
+        self.lora_A = nn.Linear(
+            base.in_features,
+            rank,
+            bias=False,
+            device=base.weight.device,
+            dtype=adapter_dtype,
+        )
+        self.lora_B = nn.Linear(
+            rank,
+            base.out_features,
+            bias=False,
+            device=base.weight.device,
+            dtype=adapter_dtype,
+        )
+        self.lora_dropout = nn.Dropout(dropout) if dropout else nn.Identity()
         self._lora_rank = int(rank)
         self._lora_alpha = float(alpha)
+        self._lora_use_rslora = bool(use_rslora)
 
     @property
     def _lora_scale(self) -> float:
-        return self._lora_alpha / self._lora_rank
+        denominator = math.sqrt(self._lora_rank) if self._lora_use_rslora else self._lora_rank
+        return self._lora_alpha / denominator
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         base_out = F.linear(x, self.weight, self.bias)
-        lora_out = self.lora_B(self.lora_A(x))
+        adapter_input = self.lora_dropout(x).to(dtype=self.lora_A.weight.dtype)
+        lora_out = self.lora_B(self.lora_A(adapter_input)).to(dtype=base_out.dtype)
         return base_out + self._lora_scale * lora_out
 
+    @torch.no_grad()
+    def merge_adapter_(self) -> None:
+        """Merge the adapter update into the base weight in place."""
+        delta = self.lora_B.weight.float() @ self.lora_A.weight.float()
+        self.weight.add_(delta.to(dtype=self.weight.dtype), alpha=self._lora_scale)
+
 
 def _target_matches(full_child_path: str, child_name: str, target: str) -> bool:
     """Return True if ``target`` selects the child at ``full_child_path``.
@@ -90,7 +123,10 @@ def _inject_lora_inplace(
     network: nn.Module,
     target_modules: list[str],
     rank: int,
-    alpha: int,
+    alpha: float,
+    dropout: float = 0.0,
+    use_rslora: bool = False,
+    adapter_dtype: torch.dtype | None = None,
 ) -> int:
     """Replace each targeted ``nn.Linear`` child in-place with ``LoraInjectedLinear``.
 
@@ -109,7 +145,18 @@ def _inject_lora_inplace(
                 continue
             full_child_path = f"{parent_name}.{child_name}" if parent_name else child_name
             if any(_target_matches(full_child_path, child_name, t) for t in target_modules):
-                setattr(parent, child_name, LoraInjectedLinear(child, rank, alpha))
+                setattr(
+                    parent,
+                    child_name,
+                    LoraInjectedLinear(
+                        child,
+                        rank,
+                        alpha,
+                        dropout=dropout,
+                        use_rslora=use_rslora,
+                        adapter_dtype=adapter_dtype,
+                    ),
+                )
                 replaced += 1
     return replaced
 
@@ -118,8 +165,13 @@ def inject_lora_pre_fsdp(
     network: torch.nn.Module,
     *,
     lora_rank: int,
-    lora_alpha: int,
+    lora_alpha: float,
     lora_target_modules: str,
+    lora_dropout: float = 0.0,
+    lora_bias: str = "none",
+    lora_use_rslora: bool = False,
+    lora_modules_to_save: str = "",
+    lora_precision: str | None = None,
 ) -> torch.nn.Module:
     """Inject LoRA adapters into ``network`` BEFORE FSDP wrap on meta device.
 
@@ -142,6 +194,12 @@ def inject_lora_pre_fsdp(
         raise ValueError(f"LoRA rank must be positive, got {lora_rank}")
     if lora_alpha <= 0:
         raise ValueError(f"LoRA alpha must be positive, got {lora_alpha}")
+    if not 0.0 <= lora_dropout < 1.0:
+        raise ValueError(f"LoRA dropout must be in [0, 1), got {lora_dropout}")
+    if lora_bias not in {"none", "all", "lora_only"}:
+        raise ValueError(f"Unsupported LoRA bias mode: {lora_bias!r}")
+    if lora_precision not in {None, "float32", "float16", "bfloat16"}:
+        raise ValueError(f"Unsupported LoRA precision: {lora_precision!r}")
 
     target_modules_list = [m.strip() for m in lora_target_modules.split(",") if m.strip()]
     if not target_modules_list:
@@ -161,20 +219,48 @@ def _target_exists(t: str) -> bool:
     if invalid_modules:
         log.warning(f"LoRA target modules not found in model: {invalid_modules}")
 
-    log.info(f"Injecting LoRA on meta device: rank={lora_rank}, alpha={lora_alpha}, targets={target_modules_list}")
+    modules_to_save = {m.strip() for m in lora_modules_to_save.split(",") if m.strip()}
+    adapter_dtype = getattr(torch, lora_precision) if lora_precision else None
+    log.info(
+        "Injecting LoRA before FSDP: "
+        f"rank={lora_rank}, alpha={lora_alpha}, dropout={lora_dropout}, "
+        f"rslora={lora_use_rslora}, bias={lora_bias}, targets={target_modules_list}, "
+        f"modules_to_save={sorted(modules_to_save)}"
+    )
 
     try:
-        replaced = _inject_lora_inplace(network, target_modules_list, lora_rank, lora_alpha)
+        replaced = _inject_lora_inplace(
+            network,
+            target_modules_list,
+            lora_rank,
+            lora_alpha,
+            dropout=lora_dropout,
+            use_rslora=lora_use_rslora,
+            adapter_dtype=adapter_dtype,
+        )
     except Exception as e:
         raise RuntimeError(f"Failed to inject LoRA adapters into model: {e}") from e
 
     if replaced == 0:
-        log.warning(f"LoRA injection replaced 0 modules — check lora_target_modules={lora_target_modules!r}")
+        raise ValueError(f"LoRA injection replaced 0 modules; check lora_target_modules={lora_target_modules!r}")
 
     lora_params = 0
     frozen_params = 0
     for name, param in network.named_parameters():
-        if "lora_" in name:
+        module_name = name.rsplit(".", 1)[0]
+        parent_name = module_name.rsplit(".", 1)[0]
+        is_module_to_save = any(
+            module_name == suffix
+            or module_name.endswith(f".{suffix}")
+            or parent_name == suffix
+            or parent_name.endswith(f".{suffix}")
+            for suffix in modules_to_save
+        )
+        is_target_bias = name.endswith(".bias") and any(
+            f".{target}." in f".{name}" for target in target_modules_list
+        )
+        is_trainable_bias = lora_bias == "all" and name.endswith(".bias")
+        if "lora_" in name or is_module_to_save or is_trainable_bias or (lora_bias == "lora_only" and is_target_bias):
             param.requires_grad_(True)
             lora_params += param.numel()
         else:
@@ -203,8 +289,48 @@ def init_lora_weights_post_materialization(network: torch.nn.Module) -> None:
     for module in network.modules():
         if not isinstance(module, LoraInjectedLinear):
             continue
-        base_dtype = module.weight.dtype
+        adapter_dtype = module.lora_A.weight.dtype
         torch.nn.init.kaiming_uniform_(module.lora_A.weight, a=math.sqrt(5))
-        module.lora_A.weight.data = module.lora_A.weight.data.to(base_dtype)
+        module.lora_A.weight.data = module.lora_A.weight.data.to(adapter_dtype)
         torch.nn.init.zeros_(module.lora_B.weight)
-        module.lora_B.weight.data = module.lora_B.weight.data.to(base_dtype)
+        module.lora_B.weight.data = module.lora_B.weight.data.to(adapter_dtype)
+
+
+@torch.no_grad()
+def merge_lora_adapters_(network: torch.nn.Module) -> int:
+    """Merge every injected adapter and return the number merged.
+
+    The modules remain present so a DCP state dict keeps the same exact keys;
+    callers exporting to a dense Hugging Face model should copy only base
+    parameters after this operation.
+    """
+    merged = 0
+    for module in network.modules():
+        if isinstance(module, LoraInjectedLinear):
+            module.merge_adapter_()
+            merged += 1
+    return merged
+
+
+@torch.no_grad()
+def merge_and_strip_lora_adapters_(network: torch.nn.Module) -> int:
+    """Merge adapters and restore plain ``nn.Linear`` modules for HF export."""
+    merged = 0
+    for parent in list(network.modules()):
+        for child_name, child in list(parent.named_children()):
+            if not isinstance(child, LoraInjectedLinear):
+                continue
+            child.merge_adapter_()
+            dense = nn.Linear(
+                child.in_features,
+                child.out_features,
+                bias=child.bias is not None,
+                device=child.weight.device,
+                dtype=child.weight.dtype,
+            )
+            dense.weight = child.weight
+            if child.bias is not None:
+                dense.bias = child.bias
+            setattr(parent, child_name, dense)
+            merged += 1
+    return merged
diff --git a/cosmos_framework/utils/generator/lora_test.py b/cosmos_framework/utils/generator/lora_test.py
new file mode 100644
index 00000000..3632f794
--- /dev/null
+++ b/cosmos_framework/utils/generator/lora_test.py
@@ -0,0 +1,60 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+from __future__ import annotations
+
+import torch
+
+from cosmos_framework.utils.generator.lora import (
+    LoraInjectedLinear,
+    inject_lora_pre_fsdp,
+    merge_lora_adapters_,
+)
+
+
+class _Tiny(torch.nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.q_proj = torch.nn.Linear(4, 3, bias=True)
+        self.lm_head = torch.nn.Linear(3, 2, bias=False)
+
+
+def test_vlm_lora_preserves_base_keys_and_reports_trainable_scope() -> None:
+    model = _Tiny()
+    base_keys = set(model.state_dict())
+    inject_lora_pre_fsdp(
+        model,
+        lora_rank=2,
+        lora_alpha=4,
+        lora_target_modules="q_proj",
+        lora_dropout=0.0,
+        lora_bias="lora_only",
+        lora_modules_to_save="lm_head",
+    )
+
+    assert isinstance(model.q_proj, LoraInjectedLinear)
+    assert base_keys <= set(model.state_dict())
+    assert model.q_proj.weight.requires_grad is False
+    assert model.q_proj.bias.requires_grad is True
+    assert model.lm_head.weight.requires_grad is True
+    assert model.q_proj.lora_A.weight.requires_grad is True
+
+
+def test_rslora_merge_matches_unmerged_forward() -> None:
+    model = _Tiny()
+    inject_lora_pre_fsdp(
+        model,
+        lora_rank=2,
+        lora_alpha=4,
+        lora_target_modules="q_proj",
+        lora_use_rslora=True,
+    )
+    torch.nn.init.normal_(model.q_proj.lora_A.weight)
+    torch.nn.init.normal_(model.q_proj.lora_B.weight)
+    values = torch.randn(2, 4)
+    expected = model.q_proj(values)
+    assert merge_lora_adapters_(model) == 1
+    with torch.no_grad():
+        model.q_proj.lora_A.weight.zero_()
+        model.q_proj.lora_B.weight.zero_()
+    torch.testing.assert_close(model.q_proj(values), expected)
diff --git a/docker/write_image_provenance.py b/docker/write_image_provenance.py
new file mode 100644
index 00000000..fa6b6eb4
--- /dev/null
+++ b/docker/write_image_provenance.py
@@ -0,0 +1,88 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Write deterministic source and runtime provenance during an image build."""
+
+from __future__ import annotations
+
+import hashlib
+import importlib.metadata
+import importlib
+import json
+import os
+import platform
+import sys
+from pathlib import Path
+
+
+def _sha256(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as stream:
+        for block in iter(lambda: stream.read(1024 * 1024), b""):
+            digest.update(block)
+    return digest.hexdigest()
+
+
+def _version(distribution: str) -> str | None:
+    try:
+        return importlib.metadata.version(distribution)
+    except importlib.metadata.PackageNotFoundError:
+        return None
+
+
+def main() -> int:
+    workspace = Path(os.environ.get("PROVENANCE_WORKSPACE", "/workspace")).resolve(strict=True)
+    output_dir = Path(os.environ.get("PROVENANCE_OUTPUT_DIR", "/opt/tao")).resolve()
+    output_dir.mkdir(parents=True, exist_ok=True)
+    roots = [workspace / "cosmos_framework", workspace / "pyproject.toml", workspace / "uv.lock", workspace / "Dockerfile"]
+    files: list[Path] = []
+    for root in roots:
+        if root.is_file():
+            files.append(root)
+        elif root.is_dir():
+            files.extend(path for path in root.rglob("*") if path.is_file() and "__pycache__" not in path.parts)
+    entries = [(str(path.relative_to(workspace)), _sha256(path)) for path in sorted(files)]
+    manifest_text = "".join(f"{digest}  {name}\n" for name, digest in entries)
+    (output_dir / "source-manifest.sha256").write_text(manifest_text, encoding="utf-8")
+    manifest_sha256 = hashlib.sha256(manifest_text.encode()).hexdigest()
+    payload = {
+        "schema_version": 1,
+        "repository": "cosmos-framework",
+        "repository_commit": os.environ.get("SOURCE_COMMIT"),
+        "repository_tree": os.environ.get("SOURCE_TREE"),
+        "source_dirty": os.environ.get("SOURCE_DIRTY") == "1",
+        "repositories": {
+            "cosmos-framework": {
+                "commit": os.environ.get("SOURCE_COMMIT"),
+                "tree": os.environ.get("SOURCE_TREE"),
+                "dirty": os.environ.get("SOURCE_DIRTY") == "1",
+            }
+        },
+        "source_manifest_sha256": manifest_sha256,
+        "dependency_lock_sha256": _sha256(workspace / "uv.lock"),
+        "dockerfile_sha256": _sha256(workspace / "Dockerfile"),
+        "build_timestamp": os.environ.get("BUILD_TIMESTAMP"),
+        "base_image": os.environ.get("PROVENANCE_BASE_IMAGE"),
+        "cuda_build_version": os.environ.get("CUDA_VERSION"),
+        "python": {"executable": sys.executable, "version": platform.python_version()},
+        "packages": {
+            name: _version(name)
+            for name in ("cosmos-framework", "torch", "transformers", "torchcodec", "PyNvVideoCodec", "peft")
+        },
+        "package_locations": {
+            name: str(Path(importlib.import_module(name).__file__).resolve())
+            for name in ("cosmos_framework", "torch", "transformers")
+        },
+    }
+    (output_dir / "image-provenance.json").write_text(
+        json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8"
+    )
+    if payload["source_dirty"]:
+        raise RuntimeError("A reproducibility image cannot be built from a dirty source tree")
+    if not payload["repository_commit"] or not payload["repository_tree"] or not payload["build_timestamp"]:
+        raise RuntimeError("SOURCE_COMMIT, SOURCE_TREE, and BUILD_TIMESTAMP build arguments are required")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

From e92d622d0d3de52128c76caf8cb2651c5dd2cf6d Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Wed, 5 Aug 2026 02:18:02 +0000
Subject: [PATCH 06/13] fix: allow non-root framework container execution

---
 Dockerfile                |  5 ++++-
 docker/dockerfile_test.py | 14 ++++++++++++++
 2 files changed, 18 insertions(+), 1 deletion(-)
 create mode 100644 docker/dockerfile_test.py

diff --git a/Dockerfile b/Dockerfile
index c2e70fba..4b901681 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -78,7 +78,10 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv pip install --no-deps .
 
 RUN /workspace/.venv/bin/python /workspace/docker/write_image_provenance.py && \
-    chmod -R a+rX /opt/tao /workspace/.venv /workspace/cosmos_framework
+    chmod a+rx /workspace /workspace/docker /workspace/docker/entrypoint.sh && \
+    chmod -R a+rX /opt/tao /workspace/.venv /workspace/cosmos_framework && \
+    test -x /workspace/docker/entrypoint.sh && \
+    test -x /workspace/.venv/bin/python
 
 # Triton bundled ptxas doesn't support latest GPU architectures
 ENV TRITON_PTXAS_PATH="/usr/local/cuda/bin/ptxas"
diff --git a/docker/dockerfile_test.py b/docker/dockerfile_test.py
new file mode 100644
index 00000000..9aa65e7c
--- /dev/null
+++ b/docker/dockerfile_test.py
@@ -0,0 +1,14 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+from pathlib import Path
+
+
+def test_runtime_tree_is_accessible_to_arbitrary_container_users() -> None:
+    """Pyxis commonly starts the container with the submitting host UID."""
+
+    dockerfile = (Path(__file__).parents[1] / "Dockerfile").read_text()
+    assert "chmod a+rx /workspace /workspace/docker /workspace/docker/entrypoint.sh" in dockerfile
+    assert "chmod -R a+rX /opt/tao /workspace/.venv /workspace/cosmos_framework" in dockerfile
+    assert "test -x /workspace/docker/entrypoint.sh" in dockerfile
+    assert "test -x /workspace/.venv/bin/python" in dockerfile

From 8e5c8d115c928540360767bb86832277765e7dd1 Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Wed, 5 Aug 2026 02:41:25 +0000
Subject: [PATCH 07/13] fix: make framework runtime image immutable

---
 docker/dockerfile_test.py | 6 ++++++
 docker/entrypoint.sh      | 6 +-----
 2 files changed, 7 insertions(+), 5 deletions(-)

diff --git a/docker/dockerfile_test.py b/docker/dockerfile_test.py
index 9aa65e7c..f0979e7b 100644
--- a/docker/dockerfile_test.py
+++ b/docker/dockerfile_test.py
@@ -12,3 +12,9 @@ def test_runtime_tree_is_accessible_to_arbitrary_container_users() -> None:
     assert "chmod -R a+rX /opt/tao /workspace/.venv /workspace/cosmos_framework" in dockerfile
     assert "test -x /workspace/docker/entrypoint.sh" in dockerfile
     assert "test -x /workspace/.venv/bin/python" in dockerfile
+
+
+def test_entrypoint_never_mutates_the_packaged_python_environment() -> None:
+    entrypoint = (Path(__file__).parent / "entrypoint.sh").read_text()
+    assert "pip install" not in entrypoint
+    assert 'exec "$@"' in entrypoint
diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh
index 8c0f7bbb..2c048ed8 100755
--- a/docker/entrypoint.sh
+++ b/docker/entrypoint.sh
@@ -5,10 +5,6 @@
 
 # Docker entrypoint script.
 
-set -e
-
-if [[ -w /workspace/.venv ]]; then
-    uv pip install --no-deps -e . || true
-fi
+set -euo pipefail
 
 exec "$@"

From 8ee5eeac8f31c841d394d2a82944aec25c91db8b Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Wed, 5 Aug 2026 18:45:37 +0000
Subject: [PATCH 08/13] feat(reasoner): configure Edge video profile at runtime

---
 .../base/reasoner/experiment/wts_vlm.py       | 36 +++++++++++++++++
 .../base/reasoner/experiment/wts_vlm_test.py  | 40 +++++++++++++++++++
 2 files changed, 76 insertions(+)

diff --git a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py
index fc2dec1e..c93f0c80 100644
--- a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py
+++ b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py
@@ -89,6 +89,7 @@ def __init__(
         video_cache_size: int = 8,
         video_device: str = "cuda",
         video_num_threads: int = 1,
+        video_max_pixels: int | str | None = None,
         system_prompt: str = "",
         use_daft_chat_template: bool = False,
     ) -> None:
@@ -104,6 +105,23 @@ def __init__(
         self.video_cache_size = video_cache_size
         self.video_device = video_device
         self.video_num_threads = video_num_threads
+        self.video_max_pixels: int | None = None
+        if video_max_pixels not in (None, "", 0, "0"):
+            parsed_video_max_pixels = int(video_max_pixels)
+            if parsed_video_max_pixels < 1:
+                raise ValueError("video_max_pixels must be >= 1")
+            hf_processor = getattr(processor, "processor", processor)
+            video_processor = getattr(hf_processor, "video_processor", None)
+            size = getattr(video_processor, "size", None)
+            if not isinstance(size, dict):
+                raise ValueError("video_max_pixels requires a processor.video_processor.size mapping")
+            shortest_edge = size.get("shortest_edge")
+            if shortest_edge is not None and parsed_video_max_pixels < int(shortest_edge):
+                raise ValueError(
+                    f"video_max_pixels ({parsed_video_max_pixels}) must be >= shortest_edge ({shortest_edge})"
+                )
+            size["longest_edge"] = parsed_video_max_pixels
+            self.video_max_pixels = parsed_video_max_pixels
         self.system_prompt = system_prompt
         self.use_daft_chat_template = use_daft_chat_template
         if self.use_daft_chat_template:
@@ -212,6 +230,7 @@ def _wts_dataloader(
             video_cache_size="${oc.env:WTS_VIDEO_CACHE_SIZE,8}",
             video_device="${oc.env:TAO_VIDEO_DECODER_DEVICE,cuda}",
             video_num_threads="${oc.env:TAO_VIDEO_DECODER_THREADS,1}",
+            video_max_pixels="${oc.env:WTS_VIDEO_MAX_PIXELS,''}",
             system_prompt="${oc.env:WTS_SYSTEM_PROMPT,''}",
         ),
         batcher=L(PoolPackingBatcher)(
@@ -256,6 +275,7 @@ def _aetc_daft_dataloader(*, split: str, shuffle: bool) -> LazyDict:
             video_cache_size="${oc.env:AETC_VIDEO_CACHE_SIZE,8}",
             video_device="${oc.env:TAO_VIDEO_DECODER_DEVICE,cuda}",
             video_num_threads="${oc.env:TAO_VIDEO_DECODER_THREADS,1}",
+            video_max_pixels="${oc.env:AETC_VIDEO_MAX_PIXELS,''}",
             system_prompt="",
             use_daft_chat_template=True,
         ),
@@ -390,3 +410,19 @@ def _aetc_daft_dataloader(*, split: str, shuffle: bool) -> LazyDict:
     name="wts_vlm_edge",
     node=wts_vlm_edge,
 )
+
+
+# Edge AETC uses the native Edge policy and the same DAFT dataset contract as
+# the Nano AETC recipe. Runtime processor limits are supplied by TAO rather
+# than encoded by modifying the public model checkpoint.
+aetc_daft_vlm_edge = deepcopy(wts_vlm_edge)
+aetc_daft_vlm_edge["job"]["group"] = "aetc_daft_edge_sft"
+aetc_daft_vlm_edge["dataloader_train"] = _aetc_daft_dataloader(split="train", shuffle=True)
+aetc_daft_vlm_edge["dataloader_val"] = _aetc_daft_dataloader(split="val", shuffle=False)
+
+ConfigStore.instance().store(
+    group="experiment",
+    package="_global_",
+    name="aetc_daft_vlm_edge",
+    node=aetc_daft_vlm_edge,
+)
diff --git a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm_test.py b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm_test.py
index 13d44175..65175a0f 100644
--- a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm_test.py
+++ b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm_test.py
@@ -9,9 +9,18 @@
 
 import pytest
 
+# These unit tests exercise dataset/config construction, not video decoding.
+# Keep collection deterministic on CPU hosts that do not have FFmpeg/CUDA
+# libraries; decoder availability is covered by the image/compute preflight.
+torchcodec_video = types.ModuleType("cosmos_framework.utils.generator.torchcodec_video")
+torchcodec_video.TorchCodecVideoReader = object
+sys.modules.setdefault("cosmos_framework.utils.generator.torchcodec_video", torchcodec_video)
+
 from cosmos_framework.configs.base.reasoner.experiment.wts_vlm import (
+    WTSProcessor,
     WTSLlavaDataset,
     aetc_daft_vlm,
+    aetc_daft_vlm_edge,
     wts_vlm_edge,
 )
 from cosmos_framework.data.generator.local_datasets.tao_vl_reason import (
@@ -92,6 +101,37 @@ def test_wts_edge_uses_native_edge_policy() -> None:
     assert wts_vlm_edge["model"]["config"]["policy"]["model_max_length"] == 16000
 
 
+def test_video_max_pixels_is_a_runtime_processor_setting() -> None:
+    tokenizer = types.SimpleNamespace(pad_token_id=0)
+    video_processor = types.SimpleNamespace(size={"shortest_edge": 4096, "longest_edge": 25165824})
+    hf_processor = types.SimpleNamespace(tokenizer=tokenizer, video_processor=video_processor)
+    wrapped_processor = types.SimpleNamespace(tokenizer=tokenizer, processor=hf_processor)
+
+    processor = WTSProcessor(wrapped_processor, video_max_pixels="16384")
+
+    assert processor.video_max_pixels == 16384
+    assert video_processor.size == {"shortest_edge": 4096, "longest_edge": 16384}
+
+
+def test_video_max_pixels_rejects_incompatible_processor_budget() -> None:
+    tokenizer = types.SimpleNamespace(pad_token_id=0)
+    video_processor = types.SimpleNamespace(size={"shortest_edge": 4096, "longest_edge": 25165824})
+    hf_processor = types.SimpleNamespace(tokenizer=tokenizer, video_processor=video_processor)
+    wrapped_processor = types.SimpleNamespace(tokenizer=tokenizer, processor=hf_processor)
+
+    with pytest.raises(ValueError, match="shortest_edge"):
+        WTSProcessor(wrapped_processor, video_max_pixels=1024)
+
+
+def test_edge_aetc_recipe_uses_edge_policy_and_runtime_video_profile() -> None:
+    assert aetc_daft_vlm_edge["defaults"][4] == {"override /vlm_policy": "cosmos3_edge_reasoner"}
+    assert "video_max_pixels" in aetc_daft_vlm_edge["dataloader_train"]["processor"]
+    assert "video_max_pixels" in wts_vlm_edge["dataloader_train"]["processor"]
+    source = __import__("inspect").getsource(sys.modules[WTSProcessor.__module__])
+    assert "AETC_VIDEO_MAX_PIXELS" in source
+    assert "WTS_VIDEO_MAX_PIXELS" in source
+
+
 def test_daft_dataset_matches_internal_hybrid_index_order(monkeypatch) -> None:
     calls = _install_fake_daft(monkeypatch)
     dataset = TaoVlReasonDaftDataset(

From c9a934ec34830631c64e884175057630f1af1433 Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Wed, 5 Aug 2026 19:08:00 +0000
Subject: [PATCH 09/13] feat(reasoner): add dataset-neutral TAO video recipes

---
 .../base/reasoner/experiment/tao_video_sft.py |  22 +++
 ...{wts_vlm_test.py => tao_video_sft_test.py} |  68 ++++---
 .../base/reasoner/experiment/wts_vlm.py       | 177 ++++++++++++++----
 .../generator/local_datasets/tao_vl_reason.py |   4 +-
 4 files changed, 204 insertions(+), 67 deletions(-)
 create mode 100644 cosmos_framework/configs/base/reasoner/experiment/tao_video_sft.py
 rename cosmos_framework/configs/base/reasoner/experiment/{wts_vlm_test.py => tao_video_sft_test.py} (66%)

diff --git a/cosmos_framework/configs/base/reasoner/experiment/tao_video_sft.py b/cosmos_framework/configs/base/reasoner/experiment/tao_video_sft.py
new file mode 100644
index 00000000..d0dd956e
--- /dev/null
+++ b/cosmos_framework/configs/base/reasoner/experiment/tao_video_sft.py
@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: OpenMDW-1.1
+
+"""Dataset-neutral public imports for TAO video-supervision SFT recipes."""
+
+from cosmos_framework.configs.base.reasoner.experiment.wts_vlm import (
+    VideoConversationDataset,
+    VideoSFTProcessor,
+    tao_task_aware_video_reasoning,
+    tao_task_aware_video_reasoning_edge,
+    tao_video_conversation,
+    tao_video_conversation_edge,
+)
+
+__all__ = [
+    "VideoConversationDataset",
+    "VideoSFTProcessor",
+    "tao_task_aware_video_reasoning",
+    "tao_task_aware_video_reasoning_edge",
+    "tao_video_conversation",
+    "tao_video_conversation_edge",
+]
diff --git a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm_test.py b/cosmos_framework/configs/base/reasoner/experiment/tao_video_sft_test.py
similarity index 66%
rename from cosmos_framework/configs/base/reasoner/experiment/wts_vlm_test.py
rename to cosmos_framework/configs/base/reasoner/experiment/tao_video_sft_test.py
index 65175a0f..432fc7fa 100644
--- a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm_test.py
+++ b/cosmos_framework/configs/base/reasoner/experiment/tao_video_sft_test.py
@@ -16,12 +16,12 @@
 torchcodec_video.TorchCodecVideoReader = object
 sys.modules.setdefault("cosmos_framework.utils.generator.torchcodec_video", torchcodec_video)
 
-from cosmos_framework.configs.base.reasoner.experiment.wts_vlm import (
-    WTSProcessor,
-    WTSLlavaDataset,
-    aetc_daft_vlm,
-    aetc_daft_vlm_edge,
-    wts_vlm_edge,
+from cosmos_framework.configs.base.reasoner.experiment.tao_video_sft import (
+    VideoConversationDataset,
+    VideoSFTProcessor,
+    tao_task_aware_video_reasoning,
+    tao_task_aware_video_reasoning_edge,
+    tao_video_conversation_edge,
 )
 from cosmos_framework.data.generator.local_datasets.tao_vl_reason import (
     TaoVlReasonDaftDataset,
@@ -55,7 +55,7 @@ def fake_template(processor) -> None:
     return calls
 
 
-def test_wts_dataset_resolves_video_paths_and_limit(tmp_path) -> None:
+def test_video_conversation_dataset_resolves_media_paths_and_limit(tmp_path) -> None:
     media = tmp_path / "videos"
     media.mkdir()
     (media / "clip.mp4").write_bytes(b"video")
@@ -82,23 +82,36 @@ def test_wts_dataset_resolves_video_paths_and_limit(tmp_path) -> None:
         encoding="utf-8",
     )
 
-    dataset = WTSLlavaDataset(str(annotations), str(media), limit="1")
+    dataset = VideoConversationDataset(str(annotations), str(media), limit="1")
     assert len(dataset) == 1
     assert dataset[0]["video"] == str(media / "clip.mp4")
 
 
-def test_wts_dataset_rejects_invalid_record(tmp_path) -> None:
+def test_video_conversation_dataset_rejects_invalid_record(tmp_path) -> None:
     annotations = tmp_path / "annotations.json"
     annotations.write_text(json.dumps([{"video": "clip.mp4"}]), encoding="utf-8")
 
     with pytest.raises(ValueError, match="conversation"):
-        WTSLlavaDataset(str(annotations), str(tmp_path))
+        VideoConversationDataset(str(annotations), str(tmp_path))
 
 
-def test_wts_edge_uses_native_edge_policy() -> None:
-    assert wts_vlm_edge["defaults"][4] == {"override /vlm_policy": "cosmos3_edge_reasoner"}
-    assert "lr_multipliers" not in wts_vlm_edge["optimizer"]
-    assert wts_vlm_edge["model"]["config"]["policy"]["model_max_length"] == 16000
+def test_video_conversation_dataset_accepts_generic_media_and_messages(tmp_path) -> None:
+    media = tmp_path / "media"
+    media.mkdir()
+    (media / "clip.mp4").write_bytes(b"video")
+    annotations = tmp_path / "annotations.json"
+    annotations.write_text(json.dumps([{
+        "media_path": "clip.mp4",
+        "messages": [{"role": "user", "content": "question"}, {"role": "assistant", "content": "answer"}],
+    }]))
+    dataset = VideoConversationDataset(str(annotations), str(media))
+    assert dataset[0]["video"] == str(media / "clip.mp4")
+
+
+def test_generic_edge_recipe_uses_native_edge_policy() -> None:
+    assert tao_video_conversation_edge["defaults"][4] == {"override /vlm_policy": "cosmos3_edge_reasoner"}
+    assert "lr_multipliers" not in tao_video_conversation_edge["optimizer"]
+    assert tao_video_conversation_edge["model"]["config"]["policy"]["model_max_length"] == 16000
 
 
 def test_video_max_pixels_is_a_runtime_processor_setting() -> None:
@@ -107,7 +120,7 @@ def test_video_max_pixels_is_a_runtime_processor_setting() -> None:
     hf_processor = types.SimpleNamespace(tokenizer=tokenizer, video_processor=video_processor)
     wrapped_processor = types.SimpleNamespace(tokenizer=tokenizer, processor=hf_processor)
 
-    processor = WTSProcessor(wrapped_processor, video_max_pixels="16384")
+    processor = VideoSFTProcessor(wrapped_processor, video_max_pixels="16384")
 
     assert processor.video_max_pixels == 16384
     assert video_processor.size == {"shortest_edge": 4096, "longest_edge": 16384}
@@ -120,23 +133,22 @@ def test_video_max_pixels_rejects_incompatible_processor_budget() -> None:
     wrapped_processor = types.SimpleNamespace(tokenizer=tokenizer, processor=hf_processor)
 
     with pytest.raises(ValueError, match="shortest_edge"):
-        WTSProcessor(wrapped_processor, video_max_pixels=1024)
+        VideoSFTProcessor(wrapped_processor, video_max_pixels=1024)
 
 
-def test_edge_aetc_recipe_uses_edge_policy_and_runtime_video_profile() -> None:
-    assert aetc_daft_vlm_edge["defaults"][4] == {"override /vlm_policy": "cosmos3_edge_reasoner"}
-    assert "video_max_pixels" in aetc_daft_vlm_edge["dataloader_train"]["processor"]
-    assert "video_max_pixels" in wts_vlm_edge["dataloader_train"]["processor"]
-    source = __import__("inspect").getsource(sys.modules[WTSProcessor.__module__])
-    assert "AETC_VIDEO_MAX_PIXELS" in source
-    assert "WTS_VIDEO_MAX_PIXELS" in source
+def test_task_aware_edge_recipe_uses_runtime_video_profile() -> None:
+    assert tao_task_aware_video_reasoning_edge["defaults"][4] == {"override /vlm_policy": "cosmos3_edge_reasoner"}
+    assert "video_max_pixels" in tao_task_aware_video_reasoning_edge["dataloader_train"]["processor"]
+    assert "video_max_pixels" in tao_video_conversation_edge["dataloader_train"]["processor"]
+    source = __import__("inspect").getsource(sys.modules[VideoSFTProcessor.__module__])
+    assert "TAO_VIDEO_MAX_PIXELS" in source
 
 
 def test_daft_dataset_matches_internal_hybrid_index_order(monkeypatch) -> None:
     calls = _install_fake_daft(monkeypatch)
     dataset = TaoVlReasonDaftDataset(
         annotation_paths='["bcq.json", "mcq.json"]',
-        media_root="/data/aetc",
+        media_root="/data/customer-media",
         response_mode="hybrid",
         system_prompt="system",
     )
@@ -149,7 +161,7 @@ def test_daft_dataset_matches_internal_hybrid_index_order(monkeypatch) -> None:
         "daft-3",
     ]
     assert calls[0]["annotation_paths"] == ["bcq.json", "mcq.json"]
-    assert calls[0]["media_roots"] == "/data/aetc"
+    assert calls[0]["media_roots"] == "/data/customer-media"
 
 
 def test_daft_chat_template_targets_wrapped_hf_processor(monkeypatch) -> None:
@@ -159,7 +171,7 @@ def test_daft_chat_template_targets_wrapped_hf_processor(monkeypatch) -> None:
     assert calls == [wrapped.processor]
 
 
-def test_daft_recipe_and_json_path_parser() -> None:
+def test_task_aware_recipe_and_json_path_parser() -> None:
     assert parse_path_list('["a.json", "b.json"]') == ["a.json", "b.json"]
-    assert aetc_daft_vlm["job"]["group"] == "aetc_daft_sft"
-    assert aetc_daft_vlm["dataloader_train"]["processor"]["use_daft_chat_template"]
+    assert tao_task_aware_video_reasoning["job"]["group"] == "tao_task_aware_video_reasoning_sft"
+    assert tao_task_aware_video_reasoning["dataloader_train"]["processor"]["use_daft_chat_template"]
diff --git a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py
index c93f0c80..c163cafc 100644
--- a/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py
+++ b/cosmos_framework/configs/base/reasoner/experiment/wts_vlm.py
@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: OpenMDW-1.1
 
-"""Woven Traffic Safety (WTS) video-QA SFT on the Cosmos3-Nano VLM path."""
+"""Generic conversation and task-aware video SFT recipes for Cosmos3."""
 
 from __future__ import annotations
 
@@ -31,8 +31,8 @@
 from cosmos_framework.utils.reasoner.constant import IGNORE_INDEX
 
 
-class WTSLlavaDataset(Dataset):
-    """Map-style loader for WTS LLaVA JSON annotations."""
+class VideoConversationDataset(Dataset):
+    """Map-style loader for ShareGPT/LLaVA-style video conversations."""
 
     def __init__(
         self,
@@ -52,34 +52,41 @@ def __init__(
         with open(self.annotation_path, encoding="utf-8") as annotation_file:
             records = json.load(annotation_file)
         if not isinstance(records, list):
-            raise TypeError(f"WTS annotations must be a JSON array, got {type(records).__name__}")
+            raise TypeError(f"video-conversation annotations must be a JSON array, got {type(records).__name__}")
         self.records = records[:parsed_limit] if parsed_limit is not None else records
         if not self.records:
-            raise ValueError(f"WTS annotation file contains no usable records: {self.annotation_path}")
+            raise ValueError(f"video-conversation annotation file contains no usable records: {self.annotation_path}")
 
         for index, record in enumerate(self.records):
-            if not isinstance(record, dict) or not isinstance(record.get("video"), str):
-                raise ValueError(f"WTS record {index} must contain a string 'video' field")
-            conversations = record.get("conversations")
+            media_value = next(
+                (record.get(field) for field in ("video", "video_id", "media", "media_path") if isinstance(record.get(field), str)),
+                None,
+            ) if isinstance(record, dict) else None
+            if media_value is None:
+                raise ValueError(f"video-conversation record {index} must contain a string media field")
+            conversations = record.get("conversations") or record.get("messages")
             if not isinstance(conversations, list) or len(conversations) < 2:
-                raise ValueError(f"WTS record {index} must contain at least two conversation turns")
+                raise ValueError(f"video-conversation record {index} must contain at least two conversation turns")
 
     def __len__(self) -> int:
         return len(self.records)
 
     def __getitem__(self, index: int) -> dict[str, Any]:
         record = dict(self.records[index])
-        video_path = record["video"]
+        video_path = next(
+            record[field] for field in ("video", "video_id", "media", "media_path")
+            if isinstance(record.get(field), str)
+        )
         if not os.path.isabs(video_path):
             video_path = os.path.join(self.media_path, video_path)
         if not os.path.isfile(video_path):
-            raise FileNotFoundError(f"WTS video does not exist: {video_path}")
+            raise FileNotFoundError(f"video-conversation media does not exist: {video_path}")
         record["video"] = video_path
         return record
 
 
-class WTSProcessor(VLMProcessor):
-    """Convert WTS ShareGPT records and uniformly sample each video to PIL frames."""
+class VideoSFTProcessor(VLMProcessor):
+    """Convert video-supervision records and uniformly sample media to PIL frames."""
 
     def __init__(
         self,
@@ -141,7 +148,7 @@ def _decode_video(self, video_path: str) -> tuple[list[Image.Image], float]:
         )
         total_frames = len(reader)
         if total_frames < 1:
-            raise ValueError(f"WTS video has zero frames: {video_path}")
+            raise ValueError(f"video-supervision media has zero frames: {video_path}")
         sample_count = min(self.num_video_frames, total_frames)
         if sample_count == 1:
             indices = [0]
@@ -164,6 +171,7 @@ def _decode_video(self, video_path: str) -> tuple[list[Image.Image], float]:
     def _sharegpt_to_openai(self, item: dict) -> list[dict]:
         if "messages" in item:
             messages = deepcopy(item["messages"])
+            video_inserted = False
             for message in messages:
                 content = message.get("content")
                 if not isinstance(content, list):
@@ -173,10 +181,22 @@ def _sharegpt_to_openai(self, item: dict) -> list[dict]:
                         continue
                     video_path = part.get("video")
                     if not isinstance(video_path, str):
-                        raise TypeError("DAFT video content must contain a string path")
+                        raise TypeError("task-aware video content must contain a string path")
                     frames, fps = self._decode_video(video_path)
                     part["video"] = frames
                     part["fps"] = fps
+                    video_inserted = True
+            if not video_inserted and isinstance(item.get("video"), str):
+                frames, fps = self._decode_video(item["video"])
+                for message in messages:
+                    if message.get("role") != "user":
+                        continue
+                    content = message.get("content", "")
+                    message["content"] = [
+                        {"type": "video", "video": frames, "fps": fps},
+                        {"type": "text", "text": content if isinstance(content, str) else ""},
+                    ]
+                    break
             return messages
 
         conversations = item.get("conversations", [])
@@ -202,16 +222,20 @@ def _sharegpt_to_openai(self, item: dict) -> list[dict]:
         return messages
 
 
-def _wts_dataloader(
+def _video_conversation_dataloader(
     *,
     annotation_env: str,
     media_env: str,
     limit_env: str,
     shuffle: bool,
+    frame_env: str = "WTS_NUM_VIDEO_FRAMES",
+    cache_env: str = "WTS_VIDEO_CACHE_SIZE",
+    max_pixels_env: str = "WTS_VIDEO_MAX_PIXELS",
+    system_prompt_env: str = "WTS_SYSTEM_PROMPT",
 ) -> LazyDict:
     return L(CosmosDataLoader)(
         distributor=L(MapDistributor)(
-            dataset=L(WTSLlavaDataset)(
+            dataset=L(VideoConversationDataset)(
                 annotation_path=f"${{oc.env:{annotation_env}}}",
                 media_path=f"${{oc.env:{media_env}}}",
                 limit=f"${{oc.env:{limit_env},''}}",
@@ -220,18 +244,18 @@ def _wts_dataloader(
             seed="${oc.env:TAO_DATALOADER_SEED,42}",
             name="train" if shuffle else "val",
         ),
-        processor=L(WTSProcessor)(
+        processor=L(VideoSFTProcessor)(
             processor=L(build_processor)(
                 tokenizer_type="${model.config.policy.backbone.model_name}",
                 config_variant="hf",
             ),
             ignore_index=IGNORE_INDEX,
-            num_video_frames="${oc.env:WTS_NUM_VIDEO_FRAMES,8}",
-            video_cache_size="${oc.env:WTS_VIDEO_CACHE_SIZE,8}",
+            num_video_frames=f"${{oc.env:{frame_env},8}}",
+            video_cache_size=f"${{oc.env:{cache_env},8}}",
             video_device="${oc.env:TAO_VIDEO_DECODER_DEVICE,cuda}",
             video_num_threads="${oc.env:TAO_VIDEO_DECODER_THREADS,1}",
-            video_max_pixels="${oc.env:WTS_VIDEO_MAX_PIXELS,''}",
-            system_prompt="${oc.env:WTS_SYSTEM_PROMPT,''}",
+            video_max_pixels=f"${{oc.env:{max_pixels_env},''}}",
+            system_prompt=f"${{oc.env:{system_prompt_env},''}}",
         ),
         batcher=L(PoolPackingBatcher)(
             max_tokens="${data_setting.max_tokens}",
@@ -247,17 +271,28 @@ def _wts_dataloader(
     )
 
 
-def _aetc_daft_dataloader(*, split: str, shuffle: bool) -> LazyDict:
-    annotation_env = f"AETC_{split.upper()}_ANNOTATIONS"
-    media_env = f"AETC_{split.upper()}_MEDIA"
-    limit_env = f"AETC_{split.upper()}_LIMIT"
+def _task_aware_video_dataloader(
+    *,
+    split: str,
+    shuffle: bool,
+    annotation_env: str | None = None,
+    media_env: str | None = None,
+    limit_env: str | None = None,
+    frame_env: str = "AETC_NUM_VIDEO_FRAMES",
+    cache_env: str = "AETC_VIDEO_CACHE_SIZE",
+    max_pixels_env: str = "AETC_VIDEO_MAX_PIXELS",
+    system_prompt_env: str = "AETC_SYSTEM_PROMPT",
+) -> LazyDict:
+    annotation_env = annotation_env or f"AETC_{split.upper()}_ANNOTATIONS"
+    media_env = media_env or f"AETC_{split.upper()}_MEDIA"
+    limit_env = limit_env or f"AETC_{split.upper()}_LIMIT"
     return L(CosmosDataLoader)(
         distributor=L(MapDistributor)(
             dataset=L(TaoVlReasonDaftDataset)(
                 annotation_paths=f"${{oc.env:{annotation_env}}}",
                 media_root=f"${{oc.env:{media_env}}}",
                 response_mode="hybrid" if split == "train" else "answer",
-                system_prompt="${oc.env:AETC_SYSTEM_PROMPT,''}",
+                system_prompt=f"${{oc.env:{system_prompt_env},''}}",
                 vision_kwargs={},
                 max_samples=f"${{oc.env:{limit_env},''}}",
             ),
@@ -265,17 +300,17 @@ def _aetc_daft_dataloader(*, split: str, shuffle: bool) -> LazyDict:
             seed="${oc.env:TAO_DATALOADER_SEED,42}",
             name=split,
         ),
-        processor=L(WTSProcessor)(
+        processor=L(VideoSFTProcessor)(
             processor=L(build_processor)(
                 tokenizer_type="${model.config.policy.backbone.model_name}",
                 config_variant="hf",
             ),
             ignore_index=IGNORE_INDEX,
-            num_video_frames="${oc.env:AETC_NUM_VIDEO_FRAMES,8}",
-            video_cache_size="${oc.env:AETC_VIDEO_CACHE_SIZE,8}",
+            num_video_frames=f"${{oc.env:{frame_env},8}}",
+            video_cache_size=f"${{oc.env:{cache_env},8}}",
             video_device="${oc.env:TAO_VIDEO_DECODER_DEVICE,cuda}",
             video_num_threads="${oc.env:TAO_VIDEO_DECODER_THREADS,1}",
-            video_max_pixels="${oc.env:AETC_VIDEO_MAX_PIXELS,''}",
+            video_max_pixels=f"${{oc.env:{max_pixels_env},''}}",
             system_prompt="",
             use_daft_chat_template=True,
         ),
@@ -355,13 +390,13 @@ def _aetc_daft_dataloader(*, split: str, shuffle: bool) -> LazyDict:
             load_from_object_store=dict(enabled=False, credentials="", bucket=""),
             save_to_object_store=dict(enabled=False, credentials="", bucket=""),
         ),
-        dataloader_train=_wts_dataloader(
+        dataloader_train=_video_conversation_dataloader(
             annotation_env="WTS_TRAIN_ANNOTATION",
             media_env="WTS_TRAIN_MEDIA",
             limit_env="WTS_TRAIN_LIMIT",
             shuffle=True,
         ),
-        dataloader_val=_wts_dataloader(
+        dataloader_val=_video_conversation_dataloader(
             annotation_env="WTS_VAL_ANNOTATION",
             media_env="WTS_VAL_MEDIA",
             limit_env="WTS_VAL_LIMIT",
@@ -384,8 +419,8 @@ def _aetc_daft_dataloader(*, split: str, shuffle: bool) -> LazyDict:
 # consuming the same DAFT dataset and Qwen chat-template contract as Cosmos-RL.
 aetc_daft_vlm = deepcopy(wts_vlm)
 aetc_daft_vlm["job"]["group"] = "aetc_daft_sft"
-aetc_daft_vlm["dataloader_train"] = _aetc_daft_dataloader(split="train", shuffle=True)
-aetc_daft_vlm["dataloader_val"] = _aetc_daft_dataloader(split="val", shuffle=False)
+aetc_daft_vlm["dataloader_train"] = _task_aware_video_dataloader(split="train", shuffle=True)
+aetc_daft_vlm["dataloader_val"] = _task_aware_video_dataloader(split="val", shuffle=False)
 
 ConfigStore.instance().store(
     group="experiment",
@@ -417,8 +452,8 @@ def _aetc_daft_dataloader(*, split: str, shuffle: bool) -> LazyDict:
 # than encoded by modifying the public model checkpoint.
 aetc_daft_vlm_edge = deepcopy(wts_vlm_edge)
 aetc_daft_vlm_edge["job"]["group"] = "aetc_daft_edge_sft"
-aetc_daft_vlm_edge["dataloader_train"] = _aetc_daft_dataloader(split="train", shuffle=True)
-aetc_daft_vlm_edge["dataloader_val"] = _aetc_daft_dataloader(split="val", shuffle=False)
+aetc_daft_vlm_edge["dataloader_train"] = _task_aware_video_dataloader(split="train", shuffle=True)
+aetc_daft_vlm_edge["dataloader_val"] = _task_aware_video_dataloader(split="val", shuffle=False)
 
 ConfigStore.instance().store(
     group="experiment",
@@ -426,3 +461,71 @@ def _aetc_daft_dataloader(*, split: str, shuffle: bool) -> LazyDict:
     name="aetc_daft_vlm_edge",
     node=aetc_daft_vlm_edge,
 )
+
+
+# Dataset-neutral TAO contracts. The older experiment registrations above are
+# retained only so existing result configs remain loadable.
+tao_video_conversation = deepcopy(wts_vlm)
+tao_video_conversation["job"]["group"] = "tao_video_conversation_sft"
+tao_video_conversation["dataloader_train"] = _video_conversation_dataloader(
+    annotation_env="TAO_VIDEO_TRAIN_ANNOTATION",
+    media_env="TAO_VIDEO_TRAIN_MEDIA",
+    limit_env="TAO_VIDEO_TRAIN_LIMIT",
+    shuffle=True,
+    frame_env="TAO_VIDEO_NUM_FRAMES",
+    cache_env="TAO_VIDEO_CACHE_SIZE",
+    max_pixels_env="TAO_VIDEO_MAX_PIXELS",
+    system_prompt_env="TAO_VIDEO_SYSTEM_PROMPT",
+)
+tao_video_conversation["dataloader_val"] = _video_conversation_dataloader(
+    annotation_env="TAO_VIDEO_VAL_ANNOTATION",
+    media_env="TAO_VIDEO_VAL_MEDIA",
+    limit_env="TAO_VIDEO_VAL_LIMIT",
+    shuffle=False,
+    frame_env="TAO_VIDEO_NUM_FRAMES",
+    cache_env="TAO_VIDEO_CACHE_SIZE",
+    max_pixels_env="TAO_VIDEO_MAX_PIXELS",
+    system_prompt_env="TAO_VIDEO_SYSTEM_PROMPT",
+)
+
+tao_task_aware_video_reasoning = deepcopy(tao_video_conversation)
+tao_task_aware_video_reasoning["job"]["group"] = "tao_task_aware_video_reasoning_sft"
+tao_task_aware_video_reasoning["dataloader_train"] = _task_aware_video_dataloader(
+    split="train", shuffle=True,
+    annotation_env="TAO_VIDEO_TRAIN_ANNOTATIONS", media_env="TAO_VIDEO_TRAIN_MEDIA_ROOTS",
+    limit_env="TAO_VIDEO_TRAIN_LIMIT", frame_env="TAO_VIDEO_NUM_FRAMES",
+    cache_env="TAO_VIDEO_CACHE_SIZE", max_pixels_env="TAO_VIDEO_MAX_PIXELS",
+    system_prompt_env="TAO_VIDEO_SYSTEM_PROMPT",
+)
+tao_task_aware_video_reasoning["dataloader_val"] = _task_aware_video_dataloader(
+    split="val", shuffle=False,
+    annotation_env="TAO_VIDEO_VAL_ANNOTATIONS", media_env="TAO_VIDEO_VAL_MEDIA_ROOTS",
+    limit_env="TAO_VIDEO_VAL_LIMIT", frame_env="TAO_VIDEO_NUM_FRAMES",
+    cache_env="TAO_VIDEO_CACHE_SIZE", max_pixels_env="TAO_VIDEO_MAX_PIXELS",
+    system_prompt_env="TAO_VIDEO_SYSTEM_PROMPT",
+)
+
+tao_video_conversation_edge = deepcopy(tao_video_conversation)
+tao_video_conversation_edge["defaults"][4] = {"override /vlm_policy": "cosmos3_edge_reasoner"}
+tao_video_conversation_edge["job"]["group"] = "tao_video_conversation_edge_sft"
+tao_video_conversation_edge["optimizer"].pop("lr_multipliers", None)
+tao_video_conversation_edge["model"]["config"]["policy"]["model_max_length"] = 16000
+
+tao_task_aware_video_reasoning_edge = deepcopy(tao_task_aware_video_reasoning)
+tao_task_aware_video_reasoning_edge["defaults"][4] = {"override /vlm_policy": "cosmos3_edge_reasoner"}
+tao_task_aware_video_reasoning_edge["job"]["group"] = "tao_task_aware_video_reasoning_edge_sft"
+tao_task_aware_video_reasoning_edge["optimizer"].pop("lr_multipliers", None)
+tao_task_aware_video_reasoning_edge["model"]["config"]["policy"]["model_max_length"] = 16000
+
+for name, node in (
+    ("tao_video_conversation", tao_video_conversation),
+    ("tao_task_aware_video_reasoning", tao_task_aware_video_reasoning),
+    ("tao_video_conversation_edge", tao_video_conversation_edge),
+    ("tao_task_aware_video_reasoning_edge", tao_task_aware_video_reasoning_edge),
+):
+    ConfigStore.instance().store(group="experiment", package="_global_", name=name, node=node)
+
+
+# Import compatibility for callers that used the development-dataset symbols.
+WTSLlavaDataset = VideoConversationDataset
+WTSProcessor = VideoSFTProcessor
diff --git a/cosmos_framework/data/generator/local_datasets/tao_vl_reason.py b/cosmos_framework/data/generator/local_datasets/tao_vl_reason.py
index 398d1d1d..f8d2386e 100644
--- a/cosmos_framework/data/generator/local_datasets/tao_vl_reason.py
+++ b/cosmos_framework/data/generator/local_datasets/tao_vl_reason.py
@@ -62,7 +62,7 @@ def __init__(
                 TaoVlReasonV1_0CosmosRLConversationDataset,
             )
         except ImportError as exc:
-            raise ImportError("nvidia-tao-daft>=2.9.1 is required for the AETC DAFT recipe") from exc
+            raise ImportError("nvidia-tao-daft>=2.9.1 is required for task-aware video annotations") from exc
 
         paths = parse_path_list(annotation_paths)
         if isinstance(media_root, str) and media_root.strip().startswith("["):
@@ -118,7 +118,7 @@ def apply_daft_chat_template(processor: Any) -> None:
             apply_chat_template_override,
         )
     except ImportError as exc:
-        raise ImportError("nvidia-tao-daft>=2.9.1 is required for the AETC DAFT recipe") from exc
+        raise ImportError("nvidia-tao-daft>=2.9.1 is required for task-aware video annotations") from exc
 
     huggingface_processor = getattr(processor, "processor", processor)
     apply_chat_template_override(huggingface_processor)

From f24a1bb3e5b7519d3fbe3fdae23b3fd1d176e9ae Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Thu, 6 Aug 2026 01:47:25 +0000
Subject: [PATCH 10/13] Respect process cpuset for GPU affinity

---
 cosmos_framework/utils/distributed.py      | 14 +++++++++++--
 cosmos_framework/utils/distributed_test.py | 23 ++++++++++++++++++++++
 2 files changed, 35 insertions(+), 2 deletions(-)

diff --git a/cosmos_framework/utils/distributed.py b/cosmos_framework/utils/distributed.py
index 87cdd39e..bb9eac1f 100644
--- a/cosmos_framework/utils/distributed.py
+++ b/cosmos_framework/utils/distributed.py
@@ -42,8 +42,18 @@ def init() -> int | None:
     local_rank = int(os.getenv("LOCAL_RANK", 0))
     try:
         device = Device(local_rank)
-        os.sched_setaffinity(0, device.get_cpu_affinity())
-    except pynvml.NVMLError as e:
+        allowed_affinity = set(os.sched_getaffinity(0))
+        device_affinity = set(device.get_cpu_affinity())
+        compatible_affinity = allowed_affinity.intersection(device_affinity)
+        if compatible_affinity:
+            os.sched_setaffinity(0, compatible_affinity)
+        else:
+            log.warning(
+                "Skipping GPU CPU affinity because NVML affinity %s does not intersect the process cpuset %s",
+                sorted(device_affinity),
+                sorted(allowed_affinity),
+            )
+    except (pynvml.NVMLError, OSError) as e:
         log.warning(f"Failed to set device affinity: {e}")
     # Set up distributed communication. CPU checkpoint conversion needs Gloo
     # because NCCL cannot synchronize CPU-resident tokenizer or model tensors.
diff --git a/cosmos_framework/utils/distributed_test.py b/cosmos_framework/utils/distributed_test.py
index 56cd7575..e8b7d049 100644
--- a/cosmos_framework/utils/distributed_test.py
+++ b/cosmos_framework/utils/distributed_test.py
@@ -28,6 +28,7 @@ def test_init_uses_gloo_backend_for_cpu(monkeypatch: pytest.MonkeyPatch) -> None
         "Device",
         lambda _local_rank: SimpleNamespace(get_cpu_affinity=lambda: [0]),
     )
+    monkeypatch.setattr(distributed.os, "sched_getaffinity", lambda _pid: {0, 1})
     monkeypatch.setattr(distributed.os, "sched_setaffinity", lambda _pid, _affinity: None)
     monkeypatch.setattr(distributed.torch.cuda, "set_device", lambda _local_rank: None)
 
@@ -36,6 +37,28 @@ def test_init_uses_gloo_backend_for_cpu(monkeypatch: pytest.MonkeyPatch) -> None
     assert init_process_group.call_args.kwargs["backend"] == "gloo"
 
 
+def test_init_intersects_nvml_affinity_with_slurm_cpuset(monkeypatch: pytest.MonkeyPatch) -> None:
+    set_affinity = Mock()
+    monkeypatch.setenv("COSMOS_DEVICE", "cpu")
+    monkeypatch.setenv("TORCH_NCCL_BLOCKING_WAIT", "0")
+    monkeypatch.setenv("TORCH_NCCL_ASYNC_ERROR_HANDLING", "1")
+    monkeypatch.setattr(distributed, "INTERNAL", False)
+    monkeypatch.setattr(distributed.dist, "is_initialized", lambda: False)
+    monkeypatch.setattr(distributed.dist, "is_available", lambda: False)
+    monkeypatch.setattr(distributed.pynvml, "nvmlInit", lambda: None)
+    monkeypatch.setattr(
+        distributed,
+        "Device",
+        lambda _local_rank: SimpleNamespace(get_cpu_affinity=lambda: [0, 32, 33]),
+    )
+    monkeypatch.setattr(distributed.os, "sched_getaffinity", lambda _pid: {32, 33, 40})
+    monkeypatch.setattr(distributed.os, "sched_setaffinity", set_affinity)
+
+    distributed.init()
+
+    set_affinity.assert_called_once_with(0, {32, 33})
+
+
 def _contains_tensor(value: Any) -> bool:
     if isinstance(value, torch.Tensor):
         return True

From 983200cfbdd2bea524cd8d9b211e4e77190ddfac Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Thu, 6 Aug 2026 02:51:04 +0000
Subject: [PATCH 11/13] fix: normalize SLURM affinity logging and data seed

---
 .../data/generator/dataflow/distributors.py   |  5 +++-
 .../data/generator/dataflow/resume_test.py    | 14 ++++++++++
 cosmos_framework/utils/distributed.py         |  6 ++---
 cosmos_framework/utils/distributed_test.py    | 26 +++++++++++++++++++
 4 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/cosmos_framework/data/generator/dataflow/distributors.py b/cosmos_framework/data/generator/dataflow/distributors.py
index 6c1a5467..4210035c 100644
--- a/cosmos_framework/data/generator/dataflow/distributors.py
+++ b/cosmos_framework/data/generator/dataflow/distributors.py
@@ -48,7 +48,10 @@ def __init__(
         name: str = "",
     ):
         self._dataset = dataset
-        self._seed = seed
+        # OmegaConf environment interpolation produces a string even when the
+        # environment value is numeric. Normalize at this typed boundary so
+        # epoch arithmetic and torch.Generator seeding remain deterministic.
+        self._seed = int(seed)
         self._shuffle = shuffle
         self._name = name
 
diff --git a/cosmos_framework/data/generator/dataflow/resume_test.py b/cosmos_framework/data/generator/dataflow/resume_test.py
index 2f3a67a8..f704da63 100644
--- a/cosmos_framework/data/generator/dataflow/resume_test.py
+++ b/cosmos_framework/data/generator/dataflow/resume_test.py
@@ -6,6 +6,8 @@
 
 from __future__ import annotations
 
+from itertools import islice
+
 import torch
 
 from cosmos_framework.callbacks.cosmos_dataloader_state import CosmosDataLoaderStateCallback
@@ -56,3 +58,15 @@ def test_resume_continues_without_dup_or_skip():
     it2 = iter(loader2)  # one iterator: env-var fast-forward happens once, then continues
     resumed = [next(it2)["id"].item() for _ in range(3)]
     assert resumed == [5, 6, 7]
+
+
+def test_map_distributor_normalizes_environment_style_seed_string():
+    distributor = MapDistributor(_IdDS(4), shuffle=True, seed="42")
+
+    assert distributor._seed == 42
+    first = [item["id"].item() for item in islice(distributor.stream(0, 1, 0, 1), 4)]
+    second = [
+        item["id"].item()
+        for item in islice(MapDistributor(_IdDS(4), shuffle=True, seed=42).stream(0, 1, 0, 1), 4)
+    ]
+    assert first == second
diff --git a/cosmos_framework/utils/distributed.py b/cosmos_framework/utils/distributed.py
index bb9eac1f..ac2b1ea3 100644
--- a/cosmos_framework/utils/distributed.py
+++ b/cosmos_framework/utils/distributed.py
@@ -49,9 +49,9 @@ def init() -> int | None:
             os.sched_setaffinity(0, compatible_affinity)
         else:
             log.warning(
-                "Skipping GPU CPU affinity because NVML affinity %s does not intersect the process cpuset %s",
-                sorted(device_affinity),
-                sorted(allowed_affinity),
+                "Skipping GPU CPU affinity because NVML affinity "
+                f"{sorted(device_affinity)} does not intersect the process cpuset "
+                f"{sorted(allowed_affinity)}"
             )
     except (pynvml.NVMLError, OSError) as e:
         log.warning(f"Failed to set device affinity: {e}")
diff --git a/cosmos_framework/utils/distributed_test.py b/cosmos_framework/utils/distributed_test.py
index e8b7d049..6a646dcb 100644
--- a/cosmos_framework/utils/distributed_test.py
+++ b/cosmos_framework/utils/distributed_test.py
@@ -59,6 +59,32 @@ def test_init_intersects_nvml_affinity_with_slurm_cpuset(monkeypatch: pytest.Mon
     set_affinity.assert_called_once_with(0, {32, 33})
 
 
+def test_init_logs_empty_nvml_slurm_affinity_intersection(monkeypatch: pytest.MonkeyPatch) -> None:
+    warning = Mock()
+    monkeypatch.setenv("COSMOS_DEVICE", "cpu")
+    monkeypatch.setenv("TORCH_NCCL_BLOCKING_WAIT", "0")
+    monkeypatch.setenv("TORCH_NCCL_ASYNC_ERROR_HANDLING", "1")
+    monkeypatch.setattr(distributed, "INTERNAL", False)
+    monkeypatch.setattr(distributed.dist, "is_initialized", lambda: False)
+    monkeypatch.setattr(distributed.dist, "is_available", lambda: False)
+    monkeypatch.setattr(distributed.pynvml, "nvmlInit", lambda: None)
+    monkeypatch.setattr(
+        distributed,
+        "Device",
+        lambda _local_rank: SimpleNamespace(get_cpu_affinity=lambda: [0, 1]),
+    )
+    monkeypatch.setattr(distributed.os, "sched_getaffinity", lambda _pid: {32, 33})
+    monkeypatch.setattr(distributed.os, "sched_setaffinity", Mock())
+    monkeypatch.setattr(distributed.log, "warning", warning)
+
+    distributed.init()
+
+    warning.assert_called_once_with(
+        "Skipping GPU CPU affinity because NVML affinity [0, 1] does not "
+        "intersect the process cpuset [32, 33]"
+    )
+
+
 def _contains_tensor(value: Any) -> bool:
     if isinstance(value, torch.Tensor):
         return True

From 986568148e5c6a54111b9b0fd84b16f89ac1907a Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Thu, 6 Aug 2026 04:02:57 +0000
Subject: [PATCH 12/13] fix(vlm): report logical PEFT parameter scope

---
 cosmos_framework/model/generator/vlm_model.py |  5 +++--
 cosmos_framework/utils/generator/lora.py      | 21 +++++++++++++++++++
 cosmos_framework/utils/generator/lora_test.py |  9 ++++++++
 3 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/cosmos_framework/model/generator/vlm_model.py b/cosmos_framework/model/generator/vlm_model.py
index 51bdf37d..c8a085a5 100644
--- a/cosmos_framework/model/generator/vlm_model.py
+++ b/cosmos_framework/model/generator/vlm_model.py
@@ -303,8 +303,9 @@ def __init__(self, config: VLMModelConfig, checkpoint):
             assert n_trainable > 0, "audio freeze policy left 0 trainable parameters — check freeze patterns"
         trainable_parameters = sum(parameter.numel() for parameter in self.model.parameters() if parameter.requires_grad)
         total_parameters = sum(parameter.numel() for parameter in self.model.parameters())
-        self.parameter_summary = {
-            "training_mode": "peft" if self.config.policy.lora_enabled else "dense_sft",
+        peft_summary = getattr(self.model.model, "_tao_peft_parameter_summary", None)
+        self.parameter_summary = peft_summary or {
+            "training_mode": "dense_sft",
             "trainable_parameters": trainable_parameters,
             "total_parameters": total_parameters,
             "frozen_parameters": total_parameters - trainable_parameters,
diff --git a/cosmos_framework/utils/generator/lora.py b/cosmos_framework/utils/generator/lora.py
index 9bc95e1b..8eac7519 100644
--- a/cosmos_framework/utils/generator/lora.py
+++ b/cosmos_framework/utils/generator/lora.py
@@ -267,6 +267,27 @@ def _target_exists(t: str) -> bool:
             param.requires_grad_(False)
             frozen_params += param.numel()
 
+    # Preserve the logical, pre-FSDP adapter scope. Some composable FSDP
+    # transforms recreate parameters with ``requires_grad=True`` even though
+    # the optimizer's ``keys_to_select=['lora_']`` filter still selects only
+    # adapters. Reporting the post-transform flags therefore mislabels every
+    # base parameter as trainable. This captured summary is the authoritative
+    # PEFT scope for status/provenance and cross-backend parity checks.
+    adapter_modules = sorted(
+        name for name, module in network.named_modules() if isinstance(module, LoraInjectedLinear)
+    )
+    network._tao_peft_parameter_summary = {
+        "training_mode": "peft",
+        "trainable_parameters": lora_params,
+        "total_parameters": lora_params + frozen_params,
+        "frozen_parameters": frozen_params,
+        "trainable_parameter_tensors": sum(
+            1 for parameter in network.parameters() if parameter.requires_grad
+        ),
+        "adapter_module_count": len(adapter_modules),
+        "adapter_modules": adapter_modules,
+    }
+
     log.info(
         f"LoRA injection successful: {replaced} modules wrapped, "
         f"{lora_params:,} trainable LoRA params, "
diff --git a/cosmos_framework/utils/generator/lora_test.py b/cosmos_framework/utils/generator/lora_test.py
index 3632f794..4262af43 100644
--- a/cosmos_framework/utils/generator/lora_test.py
+++ b/cosmos_framework/utils/generator/lora_test.py
@@ -38,6 +38,15 @@ def test_vlm_lora_preserves_base_keys_and_reports_trainable_scope() -> None:
     assert model.q_proj.bias.requires_grad is True
     assert model.lm_head.weight.requires_grad is True
     assert model.q_proj.lora_A.weight.requires_grad is True
+    assert model._tao_peft_parameter_summary == {
+        "training_mode": "peft",
+        "trainable_parameters": 23,
+        "total_parameters": 35,
+        "frozen_parameters": 12,
+        "trainable_parameter_tensors": 4,
+        "adapter_module_count": 1,
+        "adapter_modules": ["q_proj"],
+    }
 
 
 def test_rslora_merge_matches_unmerged_forward() -> None:

From ba1226f55fe113cf75fcf0988b281afb278ae9d6 Mon Sep 17 00:00:00 2001
From: Ramanathan Arunachalam <rarunachalam@nvidia.com>
Date: Thu, 6 Aug 2026 04:18:35 +0000
Subject: [PATCH 13/13] fix(vlm): restore adapter-only gradients after
 materialization

---
 cosmos_framework/model/generator/vlm_model.py |  13 +-
 cosmos_framework/utils/generator/lora.py      | 116 +++++++++++-------
 cosmos_framework/utils/generator/lora_test.py |  14 +++
 3 files changed, 97 insertions(+), 46 deletions(-)

diff --git a/cosmos_framework/model/generator/vlm_model.py b/cosmos_framework/model/generator/vlm_model.py
index c8a085a5..39fd310f 100644
--- a/cosmos_framework/model/generator/vlm_model.py
+++ b/cosmos_framework/model/generator/vlm_model.py
@@ -289,7 +289,18 @@ def __init__(self, config: VLMModelConfig, checkpoint):
 
         # Apply freeze before the optimizer is built — ``build_optimizer`` reads
         # ``requires_grad`` off ``named_parameters``.
-        n_trainable = _apply_freeze_config(self.model.model, self.hf_config.model_type, self.config.freeze)
+        if self.config.policy.lora_enabled:
+            from cosmos_framework.utils.generator.lora import apply_lora_trainable_scope
+
+            peft_summary = apply_lora_trainable_scope(
+                self.model.model,
+                lora_target_modules=self.config.policy.lora_target_modules,
+                lora_bias=self.config.policy.lora_bias,
+                lora_modules_to_save=self.config.policy.lora_modules_to_save,
+            )
+            n_trainable = int(peft_summary["trainable_parameter_tensors"])
+        else:
+            n_trainable = _apply_freeze_config(self.model.model, self.hf_config.model_type, self.config.freeze)
         if config.sound_und:
             # The standalone artifact is the sole source of encoder weights.
             # Keep it immutable even when a broad trainable_params expression
diff --git a/cosmos_framework/utils/generator/lora.py b/cosmos_framework/utils/generator/lora.py
index 8eac7519..554f2a9b 100644
--- a/cosmos_framework/utils/generator/lora.py
+++ b/cosmos_framework/utils/generator/lora.py
@@ -161,6 +161,66 @@ def _inject_lora_inplace(
     return replaced
 
 
+def apply_lora_trainable_scope(
+    network: torch.nn.Module,
+    *,
+    lora_target_modules: str,
+    lora_bias: str = "none",
+    lora_modules_to_save: str = "",
+) -> dict[str, object]:
+    """Reapply and report the logical LoRA trainable scope.
+
+    The scope must be restored after model materialization/FSDP transforms,
+    which may recreate parameters with ``requires_grad=True``.
+    """
+    target_modules = [item.strip() for item in lora_target_modules.split(",") if item.strip()]
+    modules_to_save = {item.strip() for item in lora_modules_to_save.split(",") if item.strip()}
+    trainable_parameters = 0
+    frozen_parameters = 0
+    trainable_parameter_tensors = 0
+    for name, parameter in network.named_parameters():
+        module_name = name.rsplit(".", 1)[0]
+        parent_name = module_name.rsplit(".", 1)[0]
+        is_module_to_save = any(
+            module_name == suffix
+            or module_name.endswith(f".{suffix}")
+            or parent_name == suffix
+            or parent_name.endswith(f".{suffix}")
+            for suffix in modules_to_save
+        )
+        is_target_bias = name.endswith(".bias") and any(
+            f".{target}." in f".{name}" for target in target_modules
+        )
+        is_trainable_bias = lora_bias == "all" and name.endswith(".bias")
+        trainable = (
+            "lora_" in name
+            or is_module_to_save
+            or is_trainable_bias
+            or (lora_bias == "lora_only" and is_target_bias)
+        )
+        parameter.requires_grad_(trainable)
+        if trainable:
+            trainable_parameters += parameter.numel()
+            trainable_parameter_tensors += 1
+        else:
+            frozen_parameters += parameter.numel()
+
+    adapter_modules = sorted(
+        name for name, module in network.named_modules() if isinstance(module, LoraInjectedLinear)
+    )
+    summary: dict[str, object] = {
+        "training_mode": "peft",
+        "trainable_parameters": trainable_parameters,
+        "total_parameters": trainable_parameters + frozen_parameters,
+        "frozen_parameters": frozen_parameters,
+        "trainable_parameter_tensors": trainable_parameter_tensors,
+        "adapter_module_count": len(adapter_modules),
+        "adapter_modules": adapter_modules,
+    }
+    network._tao_peft_parameter_summary = summary
+    return summary
+
+
 def inject_lora_pre_fsdp(
     network: torch.nn.Module,
     *,
@@ -244,55 +304,21 @@ def _target_exists(t: str) -> bool:
     if replaced == 0:
         raise ValueError(f"LoRA injection replaced 0 modules; check lora_target_modules={lora_target_modules!r}")
 
-    lora_params = 0
-    frozen_params = 0
-    for name, param in network.named_parameters():
-        module_name = name.rsplit(".", 1)[0]
-        parent_name = module_name.rsplit(".", 1)[0]
-        is_module_to_save = any(
-            module_name == suffix
-            or module_name.endswith(f".{suffix}")
-            or parent_name == suffix
-            or parent_name.endswith(f".{suffix}")
-            for suffix in modules_to_save
-        )
-        is_target_bias = name.endswith(".bias") and any(
-            f".{target}." in f".{name}" for target in target_modules_list
-        )
-        is_trainable_bias = lora_bias == "all" and name.endswith(".bias")
-        if "lora_" in name or is_module_to_save or is_trainable_bias or (lora_bias == "lora_only" and is_target_bias):
-            param.requires_grad_(True)
-            lora_params += param.numel()
-        else:
-            param.requires_grad_(False)
-            frozen_params += param.numel()
-
-    # Preserve the logical, pre-FSDP adapter scope. Some composable FSDP
-    # transforms recreate parameters with ``requires_grad=True`` even though
-    # the optimizer's ``keys_to_select=['lora_']`` filter still selects only
-    # adapters. Reporting the post-transform flags therefore mislabels every
-    # base parameter as trainable. This captured summary is the authoritative
-    # PEFT scope for status/provenance and cross-backend parity checks.
-    adapter_modules = sorted(
-        name for name, module in network.named_modules() if isinstance(module, LoraInjectedLinear)
+    summary = apply_lora_trainable_scope(
+        network,
+        lora_target_modules=lora_target_modules,
+        lora_bias=lora_bias,
+        lora_modules_to_save=lora_modules_to_save,
     )
-    network._tao_peft_parameter_summary = {
-        "training_mode": "peft",
-        "trainable_parameters": lora_params,
-        "total_parameters": lora_params + frozen_params,
-        "frozen_parameters": frozen_params,
-        "trainable_parameter_tensors": sum(
-            1 for parameter in network.parameters() if parameter.requires_grad
-        ),
-        "adapter_module_count": len(adapter_modules),
-        "adapter_modules": adapter_modules,
-    }
+    trainable_parameters = int(summary["trainable_parameters"])
+    frozen_parameters = int(summary["frozen_parameters"])
+    total_parameters = int(summary["total_parameters"])
 
     log.info(
         f"LoRA injection successful: {replaced} modules wrapped, "
-        f"{lora_params:,} trainable LoRA params, "
-        f"{frozen_params:,} frozen base params "
-        f"({100 * lora_params / max(1, lora_params + frozen_params):.3f}% trainable)"
+        f"{trainable_parameters:,} trainable LoRA params, "
+        f"{frozen_parameters:,} frozen base params "
+        f"({100 * trainable_parameters / max(1, total_parameters):.3f}% trainable)"
     )
     return network
 
diff --git a/cosmos_framework/utils/generator/lora_test.py b/cosmos_framework/utils/generator/lora_test.py
index 4262af43..28d49080 100644
--- a/cosmos_framework/utils/generator/lora_test.py
+++ b/cosmos_framework/utils/generator/lora_test.py
@@ -7,6 +7,7 @@
 
 from cosmos_framework.utils.generator.lora import (
     LoraInjectedLinear,
+    apply_lora_trainable_scope,
     inject_lora_pre_fsdp,
     merge_lora_adapters_,
 )
@@ -48,6 +49,19 @@ def test_vlm_lora_preserves_base_keys_and_reports_trainable_scope() -> None:
         "adapter_modules": ["q_proj"],
     }
 
+    model.requires_grad_(True)
+    summary = apply_lora_trainable_scope(
+        model,
+        lora_target_modules="q_proj",
+        lora_bias="lora_only",
+        lora_modules_to_save="lm_head",
+    )
+    assert summary == model._tao_peft_parameter_summary
+    assert model.q_proj.weight.requires_grad is False
+    assert model.q_proj.bias.requires_grad is True
+    assert model.lm_head.weight.requires_grad is True
+    assert model.q_proj.lora_A.weight.requires_grad is True
+
 
 def test_rslora_merge_matches_unmerged_forward() -> None:
     model = _Tiny()