Generated banner for Tao Finetune Cosmos Reason
Generated logo for Tao Finetune Cosmos Reason

Tao Finetune Cosmos Reason

NVIDIAOfficialCreated by NVIDIA Corporation

  • Skill
  • Workflow
  • Data and AI
  • Apache-2.0
  • Medium risk
  • Has scripts

From the author

Shared Cosmos3 frontend that explicitly routes Cosmos Framework and Cosmos-RL, validates runtime model/video-dataset/SLURM inputs, consumes an SQSH or packaged backend image, optionally plans explicit clean source builds, prepares checkpoints, validates the first update in-process, and returns token-weighted losses and task-aware accuracy.
— the description by NVIDIA Corporation

AutoGineer summary

Written by AutoGineer from the skill’s own files, not by the author

The tao-finetune-cosmos-reason skill provides a structured, unified interface for managing the lifecycle of Cosmos3 multimodal model training and evaluation. It handles complex routing between Cosmos Framework and Cosmos-RL backends, enforces strict input validation for datasets and model checkpoints, and manages the orchestration of training jobs on platforms like SLURM. The skill ensures reproducibility by sealing plans, validating runtime environments, and enforcing rigorous data and model contracts, while automating necessary conversions for Nano-tier architectures.

What it does

  • Routes training requests to the appropriate Cosmos Framework or Cosmos-RL backend.
  • Validates runtime inputs, including model paths, video datasets, and SLURM configurations.
  • Automates the conversion of Nano-tier checkpoints to required formats.
  • Executes training jobs with integrated validation of gradients and loss metrics.
  • Manages evaluation workflows, inheriting training configurations for consistency.
  • Ensures reproducibility through sealed plans and immutable reference resolution.

When to use it

  • When training or fine-tuning Cosmos3-Nano or Cosmos3-Edge models.
  • When you need to perform task-aware video reasoning or conversation-style training.
  • When orchestrating training jobs on SLURM clusters with specific container requirements.
  • When you require a reproducible, validated pipeline for model evaluation and metric extraction.

Advantages

  • Enforces strict reproducibility through sealed, immutable plan artifacts.
  • Simplifies backend selection by automatically routing based on the task type.
  • Provides robust validation of inputs and environment configurations before job submission.
  • Ensures consistency between training and evaluation phases by inheriting configurations.
  • Supports both dense and PEFT training modes with standardized parameter contracts.

Usage examples

“Train a Cosmos3-Nano model using the Cosmos-RL backend on my video conversation dataset.”
The agent will validate the dataset structure, resolve the base model, prepare the Nano checkpoint, and generate a sealed plan for the training job.
“Evaluate my fine-tuned Cosmos3-Edge model on the validation set.”
The agent will inherit the training configuration, run the evaluation workflow, and extract final metrics such as token-weighted loss and task accuracy.

Install

  1. Download the skill as a ZIP:

    https://mp-staging.autogineer.com/agent-skills/nvidia/tao-finetune-cosmos-reason/download
  2. Unzip it into the cross-client skills folder of your project:

    mkdir -p .agents/skills && unzip tao-finetune-cosmos-reason.zip -d .agents/skills

Documentation by NVIDIA Corporation

Cosmos3 TAO training

Keep one shared model-facing frontend. Backend image fields and contract paths live under backend_contracts in references/skill_info.yaml; image literals are stamped from versions.yaml, while referenced backend YAMLs define native runtime schemas. Never translate between them.

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. For a Hugging Face model ID or URL, accept an optional friendly base_model_revision such as a branch or tag. If omitted, resolve main; do not ask the user for a commit SHA. Resolve the selected ref read-only through the Hub API to its immutable commit and seal the model ID, requested ref, and resolved SHA in the plan. A complete local snapshot needs no revision and is sealed by its file fingerprints.

  • For Cosmos3-Nano, an explicit input-checkpoint model_type: qwen3_vl or cosmos3_omni. If the user did not supply it, ask once before planning; never infer the choice from config.json, a model ID, a path name, or a previous run. Explain the two choices in plain language: qwen3_vl uses a compatible Hugging Face checkpoint directly, while cosmos3_omni requires an immutable conversion to exact Qwen3-VL safetensors before training. Record the answer as base_model_format. Cosmos3-Edge is inferred as cosmos3_edge from the resolved model ID and does not present this Nano-only choice.

  • Do not expose Omni preparation implementation fields during normal intake. For Nano, use Qwen/Qwen3-VL-8B-Instruct, resolve both Hub models to immutable commits, and run the selected backend's packaged converter. Framework uses cosmos_framework.scripts.convert_model_to_vlm_safetensors; Cosmos-RL uses cosmos_rl.model_preparation.vlm_safetensors with its pinned isolated converter runtime. Validate an explicitly supplied prepared_checkpoint_path or donor, but never ask for one by default.

  • Accept hf_model://nvidia/Cosmos3-Nano directly. If a gated/private model cannot be resolved, ask the user only to set HF_TOKEN in the session environment; never ask them to discover a SHA or provide the token value in chat.

  • explicit video sampling mode: either uniform nframes or fps. FPS mode may also set min_frames and max_frames; both modes may set clip-time, resize, and pixel-budget fields supported by the selected backend.

  • 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, ssh_key_path, mounts, and scheduler settings.

  • Runtime order: compute-readable sqsh_path, explicit image, then the selected backend image in references/skill_info.yaml. On SLURM reuse or convert it once under sqsh_cache_dir. Never compare an SQSH filename with an image tag or request source provenance/SHA.

  • Repository paths, commits/trees, branch, base image, build context, and timestamp are advanced inputs required only for explicit source-build. See references/cosmos-backend-operations.md; never infer a build from runtime selection.

The planner preserves each original path and reports an accessible realpath. Missing required paths fail. A missing supplied SQSH fails; an omitted SQSH selects the packaged image. No historical fallback path or image is allowed.

Nano checkpoint model-type choice

Treat the selected base_model_format as a user decision and verify that it matches the supplied local checkpoint's config.json.model_type. A mismatch fails; it is not permission to relabel or rewrite the source checkpoint.

For qwen3_vl, fingerprint and use a complete compatible Hugging Face Nano checkpoint directly. For cosmos3_omni, tell the user in the launch review that Cosmos-RL will prepare its compatible checkpoint automatically, name the planned output path, and preserve the source checkpoint unchanged. Do not ask the user for an architecture donor, preparation image, or preparation SQSH. The backend-owned preparation step emits a verified qwen3_vl checkpoint under the selected platform's user-owned checkpoint_dir:

  • Docker: the local Docker host's checkpoint_dir.

  • SLURM: the compute-node-verified shared checkpoint_dir, covered by an explicit container mount. Run conversion through the SLURM/Pyxis contract; do not write the converted checkpoint to controller-local storage.

Use the packaged Nano architecture mapping and selected backend runtime unless an advanced override was explicitly supplied. Fingerprint the source, architecture mapping, converted config/tokenizer/processor/index/shards, and conversion provenance. Before conversion, bind the sealed plan's backend-native training fields and VLM_SAFETENSORS_PATH to the planned converted checkpoint inside the selected container. Validate that exact output before training without mutating the sealed plan. The original Omni/Hugging Face path remains provenance only and must never remain as the runtime model path. Reuse is allowed only when the exact target has complete matching conversion provenance and passes the same validation.

Public Cosmos3-Edge checkpoint contract

Accept the public model at its resolved immutable revision or a local snapshot; never request a second checkpoint. Apply the model-aware runtime defaults from references/skill_info.yaml and references/cosmos-framework-backend.yaml, preserving separate model and processor-profile fingerprints and each default/override origin.

Backend selection

Run scripts/cosmos_workflow.py resolve first.

RequestAutomatic selection
Cosmos3-Nano plain trainCosmos-RL (compatibility default)
Cosmos3-Nano AutoML/HPOCosmos-RL
Nano Framework-DCP exportCosmos Framework
Nano evaluate/inference/microservice with no explicit backendCosmos-RL
Nano quantizeCosmos-RL
Cosmos3-Edge train/export/evaluate/inference/microserviceCosmos 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.

Evaluation intake and inheritance

Run scripts/evaluation_workflow.py for every evaluate action and follow references/cosmos-reason-evaluate.md completely. A parent training plan owns all inheritable model, dataset, prompt, preprocessing, precision, checkpoint, and scoring fields; never ask the user to repeat them. Ask once only for the helper's required_user_inputs, execute its automated_actions, and launch only a checksum-valid ready=true plan. Cosmos-RL policy checkpoints require the emitted cosmos_rl_checkpoint_pre_action; Framework DCP inputs require their emitted export pre-action. The ready plan includes a validated spec_bundle.execution; pass it unchanged to the selected platform. Cosmos owns runtime attestation and evaluator configuration; the platform owns launch.

Framework checkpoint pre-action

Before Framework evaluate, inference, or microservice actions, run scripts/framework_checkpoint_action.py plan and its emitted prepare and verify steps. Follow references/cosmos-backend-operations.md; never ask the user to export DCP manually. On SLURM stage only the helper dependency set declared by workflow_contract.action_helper_dependencies; the platform verifies the closed bundle. Use only the verified action_model_path, reuse only matching complete exports, and record the pre-action, manifest, fingerprints, and independent child result.

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 tools, storage, paths, and runtime selection in the mandatory intake order; existing-SQSH and packaged-image modes skip source gates.

  4. Only for explicit source-build, validate/build clean sources and verify /opt/tao/image-provenance.json. Never mount host source into training.

  5. Enforce the explicit Nano checkpoint model-type choice. If it is cosmos3_omni, show the conversion and platform-owned output path in the launch review, then prepare the model through the shared TAO integration entrypoint packaged in the selected clean backend image after approval. Resolve URI/model-ID refs to immutable Hub commits automatically. Validate exact tensor/config keys and fingerprint model, tokenizer, processor, weight index, every shard, and provenance. Assign the verified converted path—not the original source path—to training.

  6. Validate inputs, counts, duplicates, overlap, tasks, and identities. A no-byte-hashing request selects both fast fingerprint flags; metadata stays hashed while weight/media payloads use path+size. 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. Resolve the video runtime from the structural dataset contract and enforce every profile gate in references/cosmos-reproducibility-gates.md. Cosmos-RL auto selects source-baked pynv-device-rgbp for video_conversation and system-pyav for task_aware_video_reasoning; Framework selects torchcodec-cuda-on-demand. All defaults use bounded rank-local, on-demand memory with no disk prewarm. Repeated-media validation may use the backend-native grouped sharder and validation-only caches only while preserving records, explicit batch, weighting, prompts, and preprocessing. Throughput settings that relax batch-1 parity require user authorization.

  8. Generate backend-native TOML, environment, topology, preflight commands, parity data, runtime profile, and metadata. Full specs contain no sample limit. Reuse one sealed --plan-artifact across read-only preflight, post-review materialize, and render-slurm; do not repeat original inputs. Materialize atomically in the verified compute frame, derive container paths only from explicit mounts, and never copy source patches to the cluster. Diagnostic subsets are explicit opt-ins, never launch prerequisites.

  9. On SLURM reuse the supplied/derived SQSH or convert the selected image once before GPU submit. SQSH SHA and source provenance are not runtime gates. When Omni preparation is required, inspect the SQSH filesystem and reject it unless it contains the shared TAO launcher, the native Framework converter, and /opt/tao/framework-converter-runtime.json. The runtime artifact must report validation_mode=imported_converter_module, proving the isolated converter's transitive dependency graph imported during the image build; file presence alone is not sufficient. Verify Pyxis/Enroot, mounts, Python/packages, decoder, GPU, CUDA, NCCL, and storage in the training allocation.

  10. Launch the requested training job directly; do not submit a separate smoke job unless the user explicitly requests one. For Cosmos-RL VLM training, require packaged source that emits a padding-aware attention_mask and runs the visual-gradient contract after the first backward pass and before the first optimizer update. Persist total/trainable/frozen counts and gradient norms for the vision encoder, visual projector, language model, and language head. Fail immediately when a trainable visual component has no gradient, a non-finite norm, or a zero norm. Report an explicitly frozen visual component as not applicable rather than as a failure.

  11. Materialize the full spec once and verify its SHA256 in the compute frame before rendering the job from the same plan artifact. Monitor scheduler and structured TAO state to a terminal result, and preserve the child exit code independently of scheduler state. Require child exit zero, structured SUCCESS, finite global train/validation loss, checkpoint completion, and a final evaluator metric before reporting completion.

  12. Resolve evaluation with scripts/evaluation_workflow.py. Inherit exact fine-tuning artifacts, collect only its remaining user inputs, run its backend-owned automated checkpoint pre-actions, and require ready=true. On Cosmos-RL, verify the selected HF export with cosmos_rl_checkpoint_action.py, rerun resolution with its manifest, and submit the emitted spec-bundle through the chosen platform; never select the native policy directory, copy the lifecycle into an application-owned launcher, or improvise result/status variables. Evaluate the selected checkpoint with identical prompt, preprocessing, generation, normalization, and task scoring. Extract final metrics with scripts/extract_cosmos_metrics.py.

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 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 has no LoRA block and reports trainable, frozen, and total counts. PEFT preserves rank, alpha, dropout, targets, bias, RS-LoRA, saved modules, precision, and trainable count across backends; mismatched semantics block a pair.

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;

  • the repository evaluator's final validation metric and any supporting values it emits. Do not add a second post-evaluation gate.

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 Bash, SQSH via Pyxis, no requeue by default, one launcher task per node, runtime-supplied logs, and the training child exit code. A single-node exclusive allocation preserves the requested CPU count in the SBATCH contract but passes every CPU actually granted by SLURM to the training step and records requested, allocated, and step counts. Before the training child, the same allocation runs the planner-owned packaged-runtime gate as the training job's first Pyxis step; a failed gate writes the independent child exit artifact and blocks training. Generated jobs set SLURM_EXPORT_ENV=ALL, and the platform consumer still submits with sbatch --export=ALL. Every Pyxis step also receives the planner-owned runtime variable names through --container-env, so a baked image ENV value cannot override the selected decoder, frame-transfer, cache, or worker profile. 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 and contain paths, fingerprints, runtime identity, config, resources, states, outputs, and timing without credentials. Source provenance exists only for source-build; SQSH SHA is optional and never a runtime gate.

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, then restart every affected training job from its clean sealed plan. 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/test map.

For infrastructure retries, the launch skill classifies the failure and opens the new --retry-of record; SLURM supplies validated node inventory and exclusions. Run cosmos_workflow.py retry-plan with the new record's <action-root>/config/train.toml; it rebases all writable paths and reseals the Cosmos request. Render that plan; never patch SBATCH.

Files

references/cosmos-reason-parameters.md

Cosmos-RL Parameters, Hardware, Errors, and Inference Mapping

Load this only when SKILL.md points here. If this conflicts with SKILL.md, skill_info.yaml, schemas, or platform/model skills, the current/compact source wins.

Important Parameters

Training Loop
  • train.epoch: Number of training epochs. Default 2. Keep it fixed for default Cosmos-RL AutoML searches unless the user explicitly asks to tune it. Use at least 2 for local AutoML or validation runs that need a host-visible best checkpoint for evaluate/inference; one-epoch runs can leave only a broken best symlink after checkpoint cleanup.

  • train.train_batch_per_replica: Global batch size per training step. Ideally >= 32 for stability. CRITICAL: must be divisible by train.train_policy.mini_batch (default 1 in the packaged template). Recommended production value: 32.

  • train.compile: Set to true for potential speedup on newer GPUs (H100), else false.

  • train.output_dir: Output directory for checkpoints and logs.

Model & Policy
  • policy.model_name_or_path: Required runtime Hugging Face model ID/URL or cluster-local snapshot path. The planner resolves an omitted revision as main, or a supplied branch/tag, to an immutable Hub commit automatically. Local snapshots use content fingerprints and need no revision.

  • policy.model_max_length: Context window size. Must be 40960 for video SFT. Affected by FPS, resolution, and prompt length.

  • policy.model_gradient_checkpointing: Save VRAM by recomputing activations. Keep true for large models.

Parallelism (Multi-GPU / Multi-Node)
  • policy.parallelism.dp_shard_size: Data-parallel shard size. For single-node SFT it equals visible GPUs. The validated custom policy-only multi-node launcher sets it to total policy ranks (nodes × GPUs per node).

  • policy.parallelism.dp_replicate_size: Keep 1 for the validated single-node and custom policy-only multi-node SFT launchers.

  • policy.parallelism.tp_size: Tensor parallelism. Default 1.

  • policy.parallelism.cp_size: Context parallelism. Default 1.

  • policy.parallelism.pp_size: Pipeline parallelism. Default 1.

For multi-node policy-only SFT, do not submit the single-node CLI unchanged. Start one controller on node zero and one policy replica on every node, pass --ngpus, --nnodes, and the SLURM rendezvous endpoint, and preserve the policy process exit code. Other Cosmos-RL role topologies must use their own declared contract; do not infer them from this SFT profile.

Training and evaluation can use different Slurm shapes. If the user requests multi-node training and single-node evaluation, preserve that distinction: submit train/AutoML recommendations with the requested multi-node shape and run Cosmos-RL evaluation on the smaller eval shape, usually 1 node with the requested GPUs per node.

For platform-side multi-node setup (sbatch flags on SLURM, Indexed Job + Service on Kubernetes), see the platform skill's "Multi-node training" section: skills/platform/tao-run-on-slurm, skills/platform/tao-run-on-kubernetes. Brev and local Docker are single-host only.

Optimization & Data Loading
  • train.optm_lr: Learning rate. Default 1e-6.

  • train.train_policy.type: Training policy. Default sft.

  • train.train_policy.mini_batch: Micro-batch size per GPU. If OOM, reduce this. Constraint: train_batch_per_replica % mini_batch == 0.

  • train.train_policy.dataset.name: Unique ID for dataset cache. IMPORTANT: change this if you modify fps or total_pixels to force cache regeneration.

  • train.train_policy.dataset.test_size: Validation split. Float (0.0–1.0) = ratio; Int = absolute number.

Vision Encoders
  • custom.vision.fps or custom.vision.nframes — mutually exclusive, set exactly one.

    • nframes (default in template): extract this many frames evenly across the clip. This is the safest default for 1-GPU AutoML validation runs.

    • fps: extract frames at this rate. High motion: 3. Low motion/static: 1–2. Select this only when the selected videos, policy.model_max_length, and GPU memory can absorb the expanded token count.

    • Setting both makes qwen-vl-utils' decord backend error out (Only accept either fps or nframes) and silently fall back to torchvision, which deadlocks under multi-worker dataloading (BlockingIOError [Errno 11] swscaler errors). If you switch from fps to nframes, also delete fps from your spec.

    • Default Cosmos AutoML search must not tune custom.vision.fps while the packaged template uses custom.vision.nframes. Only include custom.vision.fps when the user explicitly requests FPS tuning and the assembled spec deletes custom.vision.nframes.

  • custom.vision.min_frames / custom.vision.max_frames — optional lower and upper sampled-frame bounds used only with fps. Qwen rounds the resolved count to its frame factor and clamps it to the selected clip length.

  • custom.vision.video_start / custom.vision.video_end — optional nonnegative clip boundaries in seconds; when both are present, start must be less than end.

  • custom.vision.resized_height / custom.vision.resized_width — explicit frame dimensions. Set both together; they replace automatic pixel-budget resizing.

  • custom.vision.min_pixels / custom.vision.max_pixels / custom.vision.total_pixels — per-frame minimum, per-frame maximum, and aggregate video pixel budgets forwarded to qwen-vl-utils.

  • custom.system_prompt: Instructions prepended to every prompt.

Checkpointing
  • train.ckpt.save_freq_in_epoch: Save every N epochs. Default 1. Preserve this epoch-based default unless the user explicitly requests step-based checkpointing; dataset and hardware guidance must not override it.

  • train.ckpt.save_freq: Save every N optimizer steps. Use this only when the user explicitly requests step-based checkpointing. In that case, omit save_freq_in_epoch or set it to 0; enabling both can cause the current Cosmos-RL image to honor only the epoch cadence.

  • train.ckpt.max_keep: Keep N most recent checkpoints. Default 2 for AutoML/minimal runs so the best LoRA adapter remains available even when the container records the best validation step before later epoch cleanup.

  • train.ckpt.export_safetensors: Export in safetensors format. Default true.

When verifying downstream handoff, prefer train_output_dir/best/safetensors only if it resolves inside the results mount. In the current local Docker image, epoch-based saving writes concrete artifacts under train_output_dir/<timestamp>/safetensors/epoch_N and train_output_dir/<timestamp>/checkpoints/epoch_N/policy, but best/best_score.json and the best/{safetensors,checkpoints} symlinks can record step_N targets that do not exist. If a best symlink points at a missing step_* directory, resolve the best validation step back to the corresponding retained epoch_* directory and use that exact folder. Do not fall back to "latest" silently.

For evaluate, pass the resolved LoRA folder directly: model.model_name=<train_output_dir>/<timestamp>/safetensors/epoch_N, model.enable_lora=true, and model.base_model_path=<same base model used for training> (default the user-supplied immutable model URI, or the local base-model snapshot path). For resume/retrain, pass the exact Cosmos checkpoint policy folder as a string: train.resume=<train_output_dir>/<timestamp>/checkpoints/epoch_N/policy. Avoid train.resume=true for local Docker epoch-based checkpoints because the current resolver scans step_* checkpoint directories and can miss the epoch_* folders. Do not count downloaded base-model shards under ptm/, launcher staging files under inputs/, or the broken best symlink itself as fine-tuned checkpoints for handoff.

Validation
  • validation.freq_in_epoch: Run validation every N epochs. Default 1. Preserve this epoch-based default unless the user explicitly requests a step-based cadence.

  • validation.freq: Run a complete validation pass every N optimizer steps. Use only when explicitly requested and account for the full dataset cost; on a large validation split, frequent complete passes can dominate total runtime.

Logging
  • logging.logger: Options: console, wandb.

  • logging.project_name / logging.experiment_name: W&B experiment tracking.

Hardware

Cosmos-RL models are 8B parameters and use FSDP sharding. Dense SFT requires at least 256 GB of cumulative visible GPU memory, with no fixed device count or per-device capacity. PEFT (training_mode=peft) requires at least 72 GB instead. Set dp_shard_size to the actual visible GPU count and dp_replicate_size=1 for a single node. The production recommendation remains 8x A100 or H100 (80 GB each). A single high-memory GB300 is also viable with a compatible image, dp_shard_size=1, and the single-GPU video guards documented in cosmos-reason-single-gpu-video.md. Every visible architecture must be supported by the selected image and pass the allocated-node CUDA-stack gate.

Read runtime_requirements.gpu_host from references/skill_info.yaml and pass it as minimum-version overrides to tao-setup-nvidia-gpu-host; do not copy the values into a platform-wide rule or convert them into exact pins.

Error Patterns

CUDA out of memory (train): Reduce train.train_policy.mini_batch or increase dp_shard_size. Enable fsdp_offload if GPU memory is limited. Also check custom.vision.total_pixels — high resolution increases memory significantly.

OOM during evaluation with LoRA: Loading the base model + LoRA adapter uses more memory than zero-shot eval. If zero-shot eval passes but post-training eval OOMs, reduce evaluation.batch_size (e.g., from 10 to 1) or lower vision.total_pixels. The OOM typically manifests as the node killing the process mid-run (no Python traceback — just ERR_PROGRAM with a node-level OOM event). This is especially likely in DEFT workflows where the same eval spec is used for both zero-shot and post-training evaluation.

NaN loss: Learning rate may be too high. Reduce optm_lr and increase optm_warmup_epochs.

vision_embeds.shape[0] must be equal to n_tokens: model_max_length is too small for the video input at the current FPS and resolution. Increase policy.model_max_length to 40960.

Quantize image/video token mismatch: Mismatch in image token count between text and input_ids during calibration means quantize.max_sequence_length is too small for the sampled media tokens. The packaged template uses 4096; do not lower it to tiny values such as 128 for video calibration.

train_batch_per_replica not divisible by mini_batch: The default train_batch_per_replica=1 from the TAO Core schema is invalid because mini_batch defaults to 4. Immediate AssertionError on all ranks. Fix: set train_batch_per_replica to a multiple of mini_batch (recommended: 32 for large datasets, 4 for small datasets).

train_batch_per_replica larger than samples per rank: With FSDP, each rank sees total_samples / dp_shard_size samples. If train_batch_per_replica exceeds this, the trainer completes 0 training steps and attempts to save a checkpoint before the optimizer/scheduler is initialized, crashing with 'NoneType' object has no attribute 'state_dict'. Fix: ensure train_batch_per_replica <= total_samples / dp_shard_size. For small datasets (e.g., 31 DEFT-generated samples on 8 GPUs = ~4 per rank), set train_batch_per_replica to 4.

Stale dataset cache after changing fps/total_pixels: Change train.train_policy.dataset.name to a new unique identifier to force cache regeneration.

Forked CUDA video decoder fails or returns zero frames: A positive train or validation dataloader_num_workers must use a clean spawn multiprocessing context. Forking after the policy process initializes CUDA/NCCL can make the FFmpeg NVDEC codec fail to open or produce total_frames=0. Use a Cosmos-RL revision with the spawned-worker fix and a spawn-picklable dataset cache. When the worker count is 0, omit the multiprocessing context and prefetch factor.

Checkpoint save failure (scheduler is None): The cosmos-rl trainer crashes with 'NoneType' object has no attribute 'state_dict' when saving a checkpoint before any training step has executed. This happens when the dataset is too small for the batch size (0 steps per epoch). See the batch size error above.

You are trying to access a gated repo: The HuggingFace model nvidia/Cosmos3-Nano requires authentication. All ranks will retry in a loop until they time out. Fix: ensure HF_TOKEN is set in your environment (export HF_TOKEN=... in your shell, or set -a; source /path/to/.env; set +a) and passed into the container with -e HF_TOKEN. The user must also accept the model agreement at https://huggingface.co/nvidia/Cosmos3-Nano.

Cosmos-RL GPU resource and architecture gate: For dense SFT, the actionable launch gate is at least 256 GB of cumulative visible GPU memory, a GPU architecture supported by the selected Cosmos-RL image, and normal platform, container, S3, and credential preflight. Set dp_shard_size to the actual GPU count and do not require a fixed policy/rollout topology for SFT. Run scripts/check_tao_launch_preflight.py --gpu-min-total-memory-gb 256 --gpu-arch-allowlist cosmos_rl=sm_80,sm_90,sm_100,sm_103,sm_103a,sm_120 before launching. For PEFT, run the same command with --gpu-min-total-memory-gb 72; every other gate still applies. The PEFT floor was measured for the LoRA profile (rank 16, batch 4 per replica); a higher rank or a larger batch needs more. If the target architecture is known but cannot be detected from the launch host, pass --gpu-arch sm_XX explicitly. Architecture-specific suffixes such as a and f match the same base SM family. sm_121 is not launchable with this image unless direct runtime validation confirms sm_121 support or a newer compatible image is selected. If a resource-qualified platform still fails with a kernel JIT error such as nvrtc: invalid --gpu-architecture, classify it as an image/toolchain defect to fix with a compatible image, not as a platform resource incompatibility.

TAO_API_JOB_ID status logging warnings in direct Docker: cosmos-rl-evaluate, cosmos-rl-inference, and cosmos-rl-quantize may log a traceback from tao_status_logger.py when TAO_API_JOB_ID is unset. For direct local-Docker model-skill validation this is nonfatal if the process exits 0 and the action writes its expected result files. Do not hide a real action failure behind this warning, but do not mark an otherwise successful local run failed only because status-file logging was unavailable.

DEFT Support

Cosmos-RL implements the DEFT workflow contract for video QA tasks. Use the packaged model metadata and workflow/deft/deft.md for the pipeline overview; this skill does not package a config.json.

Gap Analysis (scripts/analyze_gaps.py)

Model-specific script that identifies failure cases from cosmos-rl evaluation output.

  • Eval output format: results.json with fields: video_id, response, question, gt

  • Comparison: exact string match after .lower().strip() — requires eval prompts that force short constrained answers (e.g., yes/no)

  • Output: parquet with video_id (full path), question, ground_truth

Limitation: Brittle exact match. If the model responds with full sentences instead of constrained answers, mismatches will be over-reported. The eval prompt design must account for this.

Spec Param / Parent Model Inference

Model-specific inference mappings belong in this MD file. Agents should read this section and apply the mappings before launching a job. This mirrors the old microservices infer_params.py flow.

  • Checkpoint metadata: format: safetensors, folder: true

Inference mappings from TAO Core cosmos-rl.config.json:

ActionSpec FieldInference FunctionMeaning
evaluatemodel.model_nameparent_model_foldermodel folder inferred from the parent job results folder
evaluateresults_diroutput_dircurrent job results directory
inferencemodel_pathparent_model_foldermodel folder inferred from the parent job results folder
inferenceresults_diroutput_dircurrent job results directory
quantizemodel.model_pathparent_model_foldermodel folder inferred from the parent job results folder
quantizeresults_diroutput_dircurrent job results directory
trainresults_diroutput_dircurrent job results directory
traintrain.output_diroutput_dircurrent job results directory
traintrain.resumeresume_modelexact checkpoint policy folder inferred from the current job results folder

For parent_model or parent_model_folder, pass the upstream train/export/AutoML child job id as the parent job id; list the parent result folder, filter checkpoint artifacts, and select the resolved model file or folder. Do not patch generated runner scripts to guess checkpoint paths.

Requirements and compatibility

Works with: Claude Code, Claude apps, OpenAI Codex, GitHub Copilot, Cursor, Gemini CLI, Kiro, OpenCode, Goose, Amp, Roo Code, JetBrains Junie, Factory Droid, Devin (ex-Windsurf), Mistral Vibe, OpenHands, Letta

“Works with” means the tool reads the open Agent Skills format; a skill written for one tool may use features others lack.

Needs: Declares allowed tools, Needs Python, Needs Node.js

The author says: Requires Python 3.11+ with PyYAML and a supported container execution platform; SLURM runs additionally require SSH, sbatch/srun, Pyxis/Enroot, and shared storage.

Security

What our scanner noticed· Medium risk

  • Large base64-encoded blobs
  • Scripts that your agent may run

This skill includes scripts that your agent may run. Read them before you use it.

A static scan is not a guarantee. Skills are third-party content: review a skill before you let an agent use it.

Attribution and licence

Created by
NVIDIA Corporation
Publisher
NVIDIA (Official)
Source
NVIDIA/skills · skills/tao-finetune-cosmos-reason · commit fce3cca
Licence
Apache License 2.0 (redistribution allowed)
  • LICENSE (added by AutoGineer from LICENSE-APACHE)
  • ATTRIBUTION.md (written by AutoGineer)
Synced

AutoGineer is not affiliated with or endorsed by NVIDIA Corporation or NVIDIA.

Download ZIP

The ZIP holds the skill’s files, its licence files and an ATTRIBUTION.md. Download tao-finetune-cosmos-reason.zip