Transfer-Auto post-training¶
This is the bidirectional RGB/LiDAR curriculum, not Sim-Auto's causal TF → TF-dCM → SF-DMD workflow. Start the curriculum from Cosmos3-Nano. Starting from released Transfer-Auto instead is a fresh fine-tune using phase 4.
Use Framework setup. Training needs the train extra, the CUDA group
matching the allocated GPU driver, and working native TorchCodec/FFmpeg decoding.
Run all commands from the Framework repository root. The workflow below uses
your local paired recordings, controls and captions.
Prepare checkpoints¶
Download a complete, authorized nvidia/Cosmos3-Nano snapshot at an immutable
revision. Convert the main model with:
python -m cosmos_framework.scripts.convert_model_to_dcp \
--checkpoint-path "$NANO_HF" \
-o "$NANO_DCP"
Set COSMOS_TRANSFER_AUTO_NANO_DCP to the resulting iteration-root directory
containing model/.metadata. This is model-only initialization, not a resumable
training checkpoint. The first phase initializes only new LiDAR projections and
the physical sensor embedding table, and begins a new EMA history.
Download the matching lidar_vae/ from an authorized Transfer-Auto HF snapshot
with hf download --include 'lidar_vae/*'. Keep its config and safetensors together.
Set COSMOS_TRANSFER_AUTO_LIDAR_VAE to the local safetensors file. The tokenizer
itself is not converted to DCP. Set COSMOS_TRANSFER_AUTO_VIDEO_VAE to the matching
Wan2.2 VAE .pth file, available from Wan-AI/Wan2.2-TI2V-5B.
Prepare local inputs¶
The user supplies synchronized real RGB, rendered WSM, target-grounded per-view
captions, and measured/HD-map LiDAR ranges. The loader does not produce WSM or
invent measured targets. Camera keys retain the physical IDs in
camera_attributes.py.
RGB and WSM use the same temporal selection and resize/center crop. LiDAR uses
the matched 128×1808 projection, metric range, normalized measured intensity,
and validity; HD-map controls use unit intensity on occupied rays.
For a new Auto bundle, create a JSONL index with one scene interval per row:
{"id":"drive1-clip1","stratum":"480","assets":{"video_camera_front_wide_120fov":"rgb/front.mp4","wsm_camera_front_wide_120fov":"wsm/front.mp4","caption_camera_front_wide_120fov":"captions/front.json","lidar/lidar_rangemap":"ranges/real.tar","lidar/hdmap_rangemap":"ranges/hdmap.tar"}}
Paths are relative to the index. Add the same three camera members for every
selected view. Supply the prepared range-video format accepted by
load_prepared_lidar_video_3ch,
not arbitrary images reshaped to the same dimensions. These are sparse sweep
archives containing the source-frame-keyed row/column/range NPZ members and,
for measured targets, intensity members; they are not single dense NPZ files.
Existing encoded bundles
can be used directly. The front-wide caption metadata is also needed for
LiDAR-only samples even if that camera is not a selected RGB view.
A camera caption file contains caption_structured: a JSON object or serialized
JSON object of chunks. Each chunk has start_frame, exclusive end_frame, and a
caption describing the recorded target interval. Keep the same clock across
cameras and range sweep IDs. Use disjoint drives/scenes for train and validation.
python -m cosmos_framework.data.generator.local_simulation.auto_bundle \
--index "$AUTO_INDEX" \
--tokenizer "$LOCAL_QWEN_TOKENIZER" \
--output "$AUTO_BUNDLE" \
--camera-keys camera_front_wide_120fov
The tokenizer directory is the complete local Qwen tokenizer/processor asset
directory matching the model. For the Nano HF export above, use
LOCAL_QWEN_TOKENIZER="$NANO_HF/text_tokenizer". Keep each index id matched
to the clip key embedded in its LiDAR archives. The builder creates a new directory only. Its
completion marker establishes packaging, not sensor alignment or training
compatibility. Decode an actual sample/batch before training. Optional --hardlink
avoids a media copy on the same filesystem but requires immutable input files.
Joint phases retain the bundle's rig-selection policy. With all eleven reference
cameras, the builder uses the reference 480p 60% seven-view / 30% eleven-view /
10% other-rig mix and the 720p 90% six-of-seven / 10% other-rig mix. With fewer
cameras it removes unavailable rigs and renormalizes, which is a documented data
adaptation, not exact rig-exposure reproduction.
Set COSMOS_TRANSFER_AUTO_LOCAL_BUNDLE to the paired camera/range Auto bundle.
All four phases use this source; no separate general-video or driving-prediction
dataset is required. All phases use selective activation checkpointing; joint
phases also save only marked operations.
Resolve and run each phase¶
The explicit registration module provides these selectors:
| Selector | Task and logical-sample mixture | Default stop | State from predecessor |
|---|---|---|---|
transfer_auto_phase1 |
Separate sensors: 30% single-camera RGB, up to 397 frames / 70% LiDAR, 21 sweeps | 70k | Nano weights only, fresh optimizer/schedule/EMA |
transfer_auto_phase2 |
Short joint multiview, 61 RGB frames / 21 sweeps | 10k new | Model/EMA/Adam retained, counter/scheduler/data reset |
transfer_auto_phase3 |
Long joint multiview, up to 397 RGB frames / 133 sweeps | Joint counter 22k | Retain model/EMA/Adam/counter/scheduler, restart data |
transfer_auto_phase4 |
Continue long joint training with the same attention and dropout policy | Joint counter 50k | Retain model/EMA/Adam/counter/scheduler, restart data |
Phase 1 requests 400 RGB frames, trimmed to the tokenizer-compatible maximum of 397; its separate 21-sweep LiDAR samples do not need to match the RGB duration. All joint phases use deduplicated maskless attention with a preceding 0.4-second cross-view window, including the current time. Phase 1 uses single-view/two-way attention. Phases 3 and 4 both drop the complete LiDAR target and control with 10% probability; phase 2 does not drop LiDAR.
Caption dropout is 10% sample-wide in every phase. Joint phases additionally drop each camera's description independently with 10% probability, preserving the camera headers. Sample-wide dropout blanks every view's caption together; the task system prompt remains. RGB views use the rig/subset sampling described above, not independent whole-RGB modality dropout. View-completion, control dropout and LiDAR-prefix dropout are disabled. T2V/I2V conditioning is 50/50.
Joint streams use 60/40 resolution exposure. The reference context/FSDP degree is 32.
Adapt topology using both model_parallel.context_parallel_size and
model.config.parallelism.context_parallel_shard_degree, plus the FSDP shard
degree. Check global batch and memory on the actual hardware.
Resolve the selected configuration before allocating training GPUs:
python -m cosmos_framework.scripts.train_experiment \
--config cosmos_framework/configs/simulation/config.py \
--experiment-module cosmos_framework.configs.simulation.experiment.auto.transfer_auto \
--dryrun \
-- \
experiment=transfer_auto_phase1
Use the same entry point for training under the allocation's launcher:
Set NNODES, GPUS_PER_NODE, NODE_RANK, and MASTER_ADDR from the allocation.
Choose CONTEXT_PARALLEL and FSDP_SHARDS compatible with that world size and the
model. The reference uses 32 for both. Keep all three topology settings consistent.
torchrun \
--nnodes="$NNODES" \
--nproc-per-node="$GPUS_PER_NODE" \
--node-rank="$NODE_RANK" \
--rdzv-backend=c10d \
--rdzv-endpoint="$MASTER_ADDR:29500" \
-m cosmos_framework.scripts.train_experiment \
--config cosmos_framework/configs/simulation/config.py \
--experiment-module cosmos_framework.configs.simulation.experiment.auto.transfer_auto \
-- \
experiment=transfer_auto_phase1 \
model_parallel.context_parallel_size="$CONTEXT_PARALLEL" \
model.config.parallelism.context_parallel_shard_degree="$CONTEXT_PARALLEL" \
model.config.parallelism.data_parallel_shard_degree="$FSDP_SHARDS"
The launcher must provide the stated resources and a reachable rendezvous.
Choose a fresh IMAGINAIRE_OUTPUT_ROOT and job.name for each new phase. After a
complete phase-1 save, set COSMOS_TRANSFER_AUTO_PHASE1_CHECKPOINT to that exact
output and select phase 2. Then bind COSMOS_TRANSFER_AUTO_PHASE2_CHECKPOINT
for phase 3, and COSMOS_TRANSFER_AUTO_PHASE3_CHECKPOINT for phase 4.
These variables name complete iteration-root checkpoints, not model-only files.
Shortened smoke and recovery¶
For a concrete single-node starting point, use the
reference-scene short check:
four H200 GPUs, CP=2/FSDP=2, 21 aligned RGB frames and synchronous checkpointing.
CP=4/FSDP=1 does not shard model/optimizer state. These short-check settings do
not establish memory requirements for the full curriculum.
See checkpoint memory
for host RAM and /dev/shm requirements before enabling async saves.
Shorten native trainer.max_iter and checkpoint.save_iter overrides. Keep
the canonical scheduler cycles and attention policies. A shorter input/view
selection tests execution, not the original maximum horizon or rig exposure.
Check every task, finite modality losses, intended parameter updates, completed
checkpoint saves, and actual state restoration.
For a bounded smoke, align save/handoff boundaries to complete CP data windows
(iteration * grad_accum_iter divisible by the context degree). This exercises
full-window restoration without claiming bit-exact replay of discarded partial windows.
trainer.max_iter is an absolute endpoint. A same-phase checkpoint at iteration
N takes one more update with trainer.max_iter=N+1. Reuse its job identity and
output root so native latest-checkpoint resume restores all state. Phase-transition
resets do not apply again to a same-job resume.
When shortening phase 1, set model.config.ema.iteration_shift to its actual
completed update count in all joint phases. The default is 70,000. Phase 2
resets the trainer/scheduler counter; phases 3 and 4 retain that joint counter
and EMA age offset. Loading Adam while resetting a schedule applies that new schedule's
initial LR before the first update, without discarding its moments.
For a fresh direct fine-tune from released Transfer-Auto, use
transfer_auto_phase4, override checkpoint.load_path
with the converted release model weights, set checkpoint.load_training_state=False,
checkpoint.keys_not_to_resume=[] and model.config.ema.iteration_shift=0, and
choose a fresh job identity, LR, schedule and duration. For model-only DCP without
EMA, use checkpoint.keys_to_skip_loading=[net_ema.]; retain strict regular-weight
loading and do not skip the trained LiDAR heads or sensor embedding table. Begin
fresh optimizer/scheduler/counter/EMA history, not a full-state resume. HF exports
do not contain optimizer or data-cursor state. Export the chosen regular/EMA
result with matched config and tokenizers only after a successful train/save/reload
check. A curriculum smoke does not separately validate release-HF direct fine-tuning.
Export¶
Export a phase checkpoint to Hugging Face format with
cosmos_framework.scripts.export_model, passing the phase experiment
(--config-file cosmos_framework/configs/simulation/config.py --experiment transfer_auto_phase2,
for example) and its checkpoint. The export is self-contained for inference:
- LiDAR VAE (optional). With
--bundle-lidar-vae, the VAE named byCOSMOS_TRANSFER_AUTO_LIDAR_VAEis copied intolidar_vae/(diffusion_pytorch_model.safetensorsplusconfig.json), the same layout as the released Transfer-Auto snapshot, and the exportedconfig.jsonpointslidar_tokenizer.vae_pathat it. The loader binds the bundled file, so the training-time VAE path does not need to exist at inference time. Without the flag, the export keeps the training-time location. - Per-view caption layout. Phases 2–4 tokenize one caption per camera.
export_modelderives this from the training dataloaders and recordsper_view_captions: trueandper_view_caption_task: "av_joint_camera_lidar_transfer"in the exportedconfig.json; phase 1 recordsper_view_captions: false. Inference reads these keys to pack prompts the way the weights were trained. For a config without dataloaders, set them with--per-view-captions on|offand--per-view-caption-task.
convert_model_to_diffusers carries both keys into the Diffusers layout; pass
--carry-lidar-vae to also copy a bundled lidar_vae/.