tao-finetune-cosmos-reason

name	tao-finetune-cosmos-reason
description	Shared Cosmos3 frontend that explicitly routes Cosmos Framework and Cosmos-RL, validates runtime model/video-dataset/SLURM inputs, builds clean repository-derived images, prepares checkpoints, gates full training on a smoke run, and returns token-weighted losses and task-aware accuracy.
license	Apache-2.0
compatibility	Docker with NVIDIA Container Toolkit, or SLURM with Pyxis/Enroot and a user-supplied shared-storage configuration.
metadata	{"author":"NVIDIA Corporation","version":"0.3.1"}
allowed-tools	Read Bash
tags	["cosmos","vlm","sft","peft","video","reasoning","slurm"]

Cosmos3 TAO training

Keep this as one shared model-facing frontend. Shared concepts live here; backend-native runtime contracts remain separate in references/cosmos-framework-backend.yaml and references/cosmos-rl-backend.yaml. Never translate one backend's TOML into the other backend's schema.

Mandatory runtime intake

Before planning training, collect all of the following. Do not infer a path from history, another user, a prior job, an image, or a developer checkout.

• base_model_path_or_uri; require base_model_revision for a URI/model ID.
• base_model_format; Nano URI inputs use qwen3_vl or cosmos3_omni. Cosmos3-Edge is inferred as cosmos3_edge from the resolved model ID.
• optional prepared_checkpoint_path; validate it instead of silently replacing it.
• training/validation annotation paths and media roots for conversation-style or task-aware video supervision, plus optional task selection.
• explicit backend for a comparison; cosmos-framework or cosmos-rl.
• training_mode; dense or peft. PEFT also requires rank, alpha, dropout, target modules, bias, RS-LoRA, modules-to-save, and adapter precision.
• user-owned results_dir, checkpoint_dir, cache_dir, and, for SLURM, sqsh_cache_dir, sqsh_path, ssh_key_path, container mounts, and all scheduler settings.
• clean repository paths and exact commits for the selected native backend, TAO integration, DAFT, and TAO Core; image tags/base images/build context are runtime inputs.

The planner preserves each original path and reports an accessible realpath. Missing paths fail. No fallback dataset, checkpoint, cache, result directory, image, partition, account, mount, SSH key, or shared-storage root is allowed.

Public Cosmos3-Edge checkpoint contract

For Cosmos3-Edge, accept the public Hugging Face model ID plus an immutable revision, or a complete local snapshot of that same public model. Do not ask the user for a second workload-specific checkpoint and do not copy or edit processor_config.json or video_preprocessor_config.json to create one. Keep the base-model fingerprint separate from the processor-profile fingerprint.

The model-aware TAO profile supplies Edge runtime defaults: 6 sampled video frames, a 1280 x 720 per-frame reference budget (5,529,600 aggregate video pixels), sequence length 16,000, and flash_attention_2. These are runtime settings, not checkpoint contents. Record whether each value came from the skill profile or an explicit user override, include it in parity and cache keys, and require the normal compute-node smoke gate before full training. Framework receives the pixel budget through TAO_VIDEO_MAX_PIXELS; Nano keeps its native processor limit unless the user explicitly overrides it.

Backend selection

Run scripts/cosmos_workflow.py resolve first.

Request	Automatic selection
Cosmos3-Nano plain train	Cosmos-RL (compatibility default)
Cosmos3-Nano AutoML/HPO	Cosmos-RL
Nano Framework-DCP export	Cosmos Framework
Nano evaluate/inference/microservice with no explicit backend	Cosmos-RL
Nano quantize	Cosmos-RL
Cosmos3-Edge train/export/evaluate/inference/microservice	Cosmos Framework

An explicit supported backend wins, so users can select Cosmos Framework for Nano training without changing model ownership. Comparative runs reject auto, so both sides of an experiment are deliberately forced. Framework-trained checkpoints use the native exact-key exporter, then the repository-backed TAO evaluation adapter. That does not make Framework a Cosmos-RL version.

Framework checkpoint pre-action

Before every Framework evaluate, inference, or inference_microservice action, invoke scripts/framework_checkpoint_action.py plan. Do not ask the user to find or run an exporter manually. The helper preserves the supplied checkpoint path, detects a complete HF safetensors directory versus a native Framework DCP, infers the saved Framework config only from the checkpoint's standard run layout, and otherwise requires an explicit config_file.

For DCP input, run the helper's prepare verb in the newly built Framework TAO action image before starting the requested action. Its runner may be local Docker or an srun Pyxis command supplied through --command-prefix. On SLURM, stage this checked-in helper and cosmos_common.py with checksums in the job input directory; mount only that job directory, not a source checkout. The helper invokes the repository-owned exact-key exporter, verifies the DCP metadata/config/base-model/revision/exported weights and manifest, writes .tao_export_complete, and returns action_model_path. Put that verified path into model.model_name for evaluation or model_path for inference and the microservice. Run verify once more from the target compute frame before the action child starts.

An already verified export is reused without conversion. A stale or partial export is never silently accepted: prepare creates and verifies a sibling temporary export, preserves the invalid directory under an .invalid-* name, then atomically installs the replacement. Capture the pre-action JSON, child exit code, export manifest checksum, source checkpoint/config fingerprints, and final action child code in the job metadata. A failed export blocks the evaluate or inference allocation/action and emits terminal FAILURE.

Required gates

Execute these stages in order and persist their outputs.

1. Resolve model/backend/action and load the selected backend contract.
2. Check credentials by presence only. Never read or persist credential values. Require a token only for the operation that needs it.
3. Validate host tools, clean repository commits/trees, build context, free storage, every original/resolved path, and image build inputs.
4. Build the selected native image and TAO action image from clean commits. Inspect /opt/tao/image-provenance.json; reject dirty, missing, or mismatched source. Never mount a host source checkout into training.
5. If needed, prepare the model with the converter packaged in the clean Framework image. URI downloads require immutable revisions. Validate exact tensor/config keys and fingerprint model, tokenizer, and processor files.
6. Validate every annotation and referenced media file, record counts, duplicates, train/validation overlap, task selection, and fingerprints. Verify the resolved inputs again from an allocated compute node. When SLURM storage is not mounted on the launch host, let cosmos_workflow.py stream its checked-in cosmos_common.py inspector to a login host over SSH. It runs from stdin, preserves remote realpath values, and creates no remote script or source overlay. Do not require local Lustre, sbatch, or srun on an SSH-based launch host. Run this expensive input inspection exactly once with the plan verb and pass a local --plan-artifact <path> so the resolved request and inspection results are sealed for the remaining launch verbs.
7. Prepare the decoder input selected by the structural dataset contract. Cosmos-RL defaults to direct on-demand sample processing so training starts without a dataset-cache prewarm phase. Conversation-style runs prewarm only when the user explicitly selects --rl-dataset-cache-mode prewarm; that opt-in uses separate deterministic train and validation keys. Task-aware runs decode directly on GPU and require a fingerprinted override artifact from the packaged cosmos_rl.utils.video_override_artifacts builder. Pair every annotation with its actual media root, scan the full input set for the NVDEC macroblock limit, force every validation annotation with --force-annotation, and add only independently diagnosed train streams with --force-video. Validate the artifact with the packaged cosmos_rl.utils.validate_video_override_artifacts command, including full validation-media coverage and GPU random-access decoding, before smoke. The JSON plan emits exact decoder_artifact.preparation_command and decoder_artifact.validation_command values for the selected clean image; re-plan once with their map, manifest, and artifact fingerprint outputs, seal that plan, and never reuse another run's cache or override artifact.

Dataset contracts

Resolve datasets by structure, not by project, benchmark, directory, or file name. The supported families are:

• video_conversation: a JSON array with media and at least two ShareGPT, LLaVA, or OpenAI-style conversation turns;
• task_aware_video_reasoning: one or more item-envelope or array annotation files with media, task identity, and conversation/response targets.

Default dataset_family to auto, inspect every annotation, and require train and validation to resolve to the same family. Capture record count, unique media count, media reuse, extensions, byte-size distribution, task/metric metadata, and any declared width, height, FPS, and duration. Select processor, cache, smoke-size, and resource profiles from those characteristics and the model tier. Never branch on a customer dataset name.

Tasks declaring accuracy participate in deterministic accuracy; common binary and multiple-choice task types are recognized. Generative tasks report their declared metrics and are excluded from aggregate accuracy with a reason. Aggregate accuracy is example-weighted over records with an accuracy definition.

Dense and PEFT contracts

Dense SFT must have no active LoRA block and must report trainable, frozen, and total parameter counts. PEFT must represent the same rank, alpha, dropout, target modules, bias, RS-LoRA, modules-to-save, precision, and trainable count on both backends. Reject a paired PEFT run when semantics cannot be matched.

For fair comparisons, force the same logical model, train/validation records, media, prompt, frames, sequence length, precision, seed, epochs, effective global batch, optimizer, learning rate, schedule, warmup, weight decay, clipping, loss masking, validation/checkpoint cadence, evaluated checkpoint, and generation/normalization settings. Classify differences as equivalent syntax, unavoidable implementation difference, or invalid mismatch; an invalid mismatch blocks the full pair.

Metrics and completion

The required primary metrics are:

• complete-run globally reduced token-weighted training loss, with numerator and valid-label denominator;
• final-validation globally reduced token-weighted loss, with numerator and valid-label denominator;
• repository-evaluator validation accuracy, with correct/total, coverage, per-task metrics, aggregation definition, exclusions, and evaluator version.

Do not average console lines or rank means. A step loss is not the average training loss. A validation heartbeat is not final validation loss. A generative exact match is not accuracy unless the task defines it.

The native Framework callback and native Cosmos-RL logger own early failure, checkpoint, progress, metric, and terminal events. Do not stage a status bridge or patch status at container startup. COMPLETED from SLURM is failure when the child exit code is nonzero or the terminal TAO state is not successful.

SLURM invariants

Generated jobs use #!/usr/bin/env bash and are syntax-checked by Bash. They use SQSH via Pyxis, disable requeue by default, run one launcher task per node, write stdout/stderr to runtime-supplied paths, and exit with the training child code. Framework topology is shard=gpus_per_node, replica=nodes. Cosmos-RL uses one controller on node zero and its policy-worker topology. Asynchronous distributed checkpointing is rejected for multi-node runs.

Every job metadata record must validate against schemas/cosmos-job-metadata.schema.json. It records supplied/resolved paths, model/data/image fingerprints, source commits, config and SQSH checksums, requested/allocated topology, scheduler and child states, logs/results, timestamps, and terminal TAO status without credentials.

Source-affecting recovery

If a run exposes a code or image defect, stop the affected path, change the owning repository, add a test, commit it, rebuild both image and SQSH from a clean checkout, rerun smoke, and rerun every affected full job. Never edit a running container, patch an existing image, reuse an old SQSH after a source change, or rely on a temporary launch script as the implementation.

Use references/cosmos-reproducibility-gates.md as the source-owner and test map before proposing a workaround in a fresh session.