cosmos-framework/examples/integration/trainer_level_training.py

# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

"""
Minimal training demo — drive Cosmos's OmniMoTModel from a plain PyTorch loop.

⚠  THIS IS A WIRING DEMO. It shows the smallest possible call sequence to drive
   `model.training_step` from your own loop — it is NOT a fine-tuning recipe.
   Production SFT uses FSDP across ≥ 8 GPUs (AdamW), real datasets (not the
   random tensors used here), and a curriculum / callbacks / EMA. The TOML
   recipes in `examples/toml/*.toml` are the real entry points.

⚠  THE MAIN TRANSFORMER IS RANDOM-INITIALIZED — the demo never loads the
   ~30 GB Cosmos3-Nano DCP shards. Loss values are therefore meaningless;
   the point is to show the call sequence and tensor shapes. For real weight
   loading see `cosmos_framework.inference.model.Cosmos3OmniModel.from_pretrained_dcp`
   and the production trainer in `cosmos_framework.scripts.train`.

================================================================================
SCOPE
================================================================================
This is NOT "extracting the model into another framework". The cosmos_framework package
must be installed (`pip install -e .` from the repo root). OmniMoTModel has deep
imports across cosmos_framework (sequence packing, MoE network, VAE, …) — physically
excising it isn't realistic.

What this demo SHOWS is the integration contract:
    - what to import,
    - what the input batch dict must contain,
    - which model methods to call,
so that you can plug OmniMoTModel into your own training framework as a
black-box `nn.Module` whose `training_step` returns a scalar loss.

What we USE from cosmos_framework:
    cosmos_framework.inference.model.Cosmos3OmniModel          → model class (random-init in this demo;
                                                       use `.from_pretrained_dcp(...)` for real weights)
    cosmos_framework.inference.common.init.init_script         → 1-line torch.distributed init
    cosmos_framework.model.vfm.vlm.qwen3_vl.utils.tokenize_caption
                                                     → text tokenizer (modelling pkg)
    model.training_step(batch, iteration)            → THE training step (flow-matching loss)
    model.config.{action_gen,sound_gen,vision_gen,…} → modality flags

What we DO NOT use:
    cosmos_framework.scripts.train, cosmos_framework.trainer.*           → CLI + Trainer class
    cosmos_framework.data.vfm.joint_dataloader.*               → iterative joint dataloader
    cosmos_framework.data.vfm.augmentor_provider.*             → text/video augmentor pipeline
    cosmos_framework.inference.inference.OmniInference          → inference pipeline

================================================================================
WHY init_script() IS NEEDED
================================================================================
OmniMoTModel uses torch.distributed primitives even on a single GPU
(ParallelDims, DTensor helpers, FSDP composables). `init_script()` runs
`torch.distributed.init_process_group("nccl")` in 1-rank mode and registers DCP
config wrappers. Drop it and the loader crashes with cryptic "default process
group not initialized" errors.

================================================================================
DATA BATCH CONTRACT (single-modality vision branch)
================================================================================
The dict passed to `model.training_step(batch, iteration)` must contain:

    Key                            Type                       Shape / Notes
    ────────────────────────────────────────────────────────────────────────
    model.input_video_key          list[Tensor]  (len=B)      [1, C=3, T, H, W] in [-1, 1]
        (default: "video")                                    For T>1, video; for T=1, image.
    model.input_image_key          list[Tensor]  (len=B)      [1, C=3, 1, H, W] in [-1, 1]
        (default: "images")                                   Alternative image-only entry point.
    model.input_caption_key        list[str]     (len=B)      raw text (NOT re-tokenized by model)
        (default: "ai_caption")
    "text_token_ids"               list[Tensor]  (len=B)      [1, N_tok] long tensor — pre-tokenized
    "image_size"                   list[Tensor]  (len=B)      [1, 4] float — (H, W, H, W)
    "fps"                          Tensor                     [B]  float
    "conditioning_fps"             Tensor                     [B]  float
    "num_frames"                   Tensor                     [B]  int
    "is_preprocessed"              bool                       True ⇒ video already normalized

For ACTION training (forward dynamics / policy) the batch also needs `action`,
`domain_id`, `raw_action_dim`, `mode`, and a hand-built `sequence_plan` — see
`make_action_fdm_batch` below for a worked example, or
`cosmos_framework/inference/action.py: build_action_batch` for the canonical impl.

GOTCHA — video shape differs between training and inference batches:
    Training (this file, is_preprocessed=True) expects a FLAT list of tensors:
        batch[model.input_video_key] = [video]              # [1, C, T, H, W]
    Inference (`cosmos_framework.inference.action.build_action_batch`) uses NESTED:
        batch[model.input_video_key] = [[video]]            # one extra []
    Copying an inference batch into a training loop produces a confusing
    `_normalize_video_databatch_inplace` error. Use the flat convention here.

================================================================================
MEMORY (READ THIS BEFORE RUNNING)
================================================================================
Full-fine-tuning the 8B Cosmos3-Nano on a single 80 GB GPU does NOT fit with
AdamW (param + grad + Adam moments ≈ 96 GB). For a single-GPU demo we use SGD
(no optimizer state) and small inputs; full SFT in production uses FSDP across
≥ 8 GPUs and/or LoRA — see `cosmos_framework.scripts.train` and `examples/toml/*.toml`.

To make full-fine-tuning fit on real hardware, you would either:
    - shard with FSDP (`cosmos_framework.utils.vfm.parallelism.ParallelDims` + FSDP wrap),
    - inject LoRA (`model.add_lora(...)`), or
    - swap the optimizer for one with lower state (Adafactor, 8-bit AdamW).

================================================================================
RUN
================================================================================
    PYTHONPATH=. python examples/integration/trainer_level_training.py
    PYTHONPATH=. python examples/integration/trainer_level_training.py --config-dir /path/to/dir/with/config.json
"""

from cosmos_framework.inference.common.init import init_script

init_script(training=True)  # ← see docstring above

import argparse
import json
from pathlib import Path

import attrs
import torch

from cosmos_framework.configs.base.defaults.compile import CompileConfig
from cosmos_framework.configs.base.defaults.parallelism import ParallelismConfig
from cosmos_framework.data.vfm.action.domain_utils import get_domain_id
from cosmos_framework.data.vfm.action.transforms import build_sequence_plan_from_mode
from cosmos_framework.data.vfm.sequence_packing import SequencePlan
from cosmos_framework.inference.args import DEFAULT_CHECKPOINT
from cosmos_framework.inference.model import Cosmos3OmniConfig, Cosmos3OmniModel
from cosmos_framework.model.vfm.vlm.qwen3_vl.utils import tokenize_caption


def _load_omni_model(*, config_dir_arg: str | None):
    """Build OmniMoTModel with RANDOM main-transformer weights — wiring demo only.

    This helper exists so the demo can run without downloading the ~30 GB transformer
    DCP. Only ``config.json`` is fetched (single ~5 KB file) and the main net is
    instantiated via ``hydra.utils.instantiate`` with random parameters. Auxiliary
    sub-models (Qwen3-VL tokenizer, Wan2.2 VAE, AVAE) still load from the HF cache
    during ``Cosmos3OmniModel.__init__`` — they are not stubbed out.

    For REAL weight loading, see
    :func:`cosmos_framework.inference.model.Cosmos3OmniModel.from_pretrained_dcp`
    and the production trainer in :mod:`cosmos_framework.scripts.train`.
    """
    if config_dir_arg is None:
        from huggingface_hub import hf_hub_download
        config_dir = Path(hf_hub_download(
            repo_id=DEFAULT_CHECKPOINT.hf.repository,
            filename="config.json",
            revision=DEFAULT_CHECKPOINT.hf.revision,
        )).parent
    else:
        config_dir = Path(config_dir_arg)
    # Shipped DCPs nest config.json one level deeper under model/.
    if not (config_dir / "config.json").exists() and (config_dir / "model" / "config.json").exists():
        config_dir = config_dir / "model"
    print(f"Loading config from: {config_dir / 'config.json'}")

    # Shipped configs carry stale `cosmos3._src.*` dotted module strings in `_type` / `_target_`
    # fields. cosmos_framework's CONFIG_REPLACEMENTS_INVERSE only rewrites the slash-form
    # paths, so we rewrite the dotted form here before constructing the config.
    config_text = (config_dir / "config.json").read_text()
    for _old, _new in [
        ("cosmos3._src.vfm.configs.base.", "cosmos_framework.configs.base."),
        ("cosmos3._src.vfm.models.", "cosmos_framework.model.vfm."),
        ("cosmos3._src.vfm.tokenizers.", "cosmos_framework.model.vfm.tokenizers."),
        ("cosmos3._src.imaginaire.", "cosmos_framework."),
    ]:
        config_text = config_text.replace(_old, _new)
    config = Cosmos3OmniConfig(model=json.loads(config_text)["model"])
    config.parallelism = attrs.asdict(ParallelismConfig())
    config.compile = attrs.asdict(CompileConfig(enabled=False))
    return Cosmos3OmniModel(config).model


# ────────────────────────────────────────────────────────────────────────────
# Per-modality batch builders. Each returns a B=1 dict in the shape that
# model.training_step expects. Plug your own dataset in by producing the
# same keys per sample and collating into list-valued entries.
# ────────────────────────────────────────────────────────────────────────────

def _tokenize(model, caption: str, device) -> torch.Tensor:
    """Tokenize a caption using the model's own VLM tokenizer."""
    ids = tokenize_caption(
        caption,
        model.vlm_tokenizer,
        is_video=False,
        use_system_prompt=model.vlm_config.use_system_prompt,
    )
    # Shape [1, N_tok]. The collate format in cosmos_framework.data.vfm.joint_dataloader
    # keeps text_token_ids as a list of [1, N] tensors (one per sample) because
    # token counts vary across the batch.
    return torch.tensor(ids, dtype=torch.long, device=device).unsqueeze(0)


def make_text_to_image_batch(model, *, caption: str, h: int = 128, w: int = 128, device="cuda") -> dict:
    """Text-to-image: vision branch with T=1."""
    image = (torch.randn(1, 3, 1, h, w, device=device) * 0.3).clamp(-1, 1)  # must be in [-1, 1]
    return {
        model.input_image_key:   [image],                                                       # T=1 → image branch
        model.input_caption_key: [caption],
        "text_token_ids":        [_tokenize(model, caption, device)],
        "image_size":            [torch.tensor([[h, w, h, w]], dtype=torch.float32, device=device)],
        "fps":              torch.tensor([16.0], device=device),
        "conditioning_fps": torch.tensor([16.0], device=device),
        "num_frames":       torch.tensor([1], device=device),
        "is_preprocessed":  True,
    }


def make_text_to_video_batch(model, *, caption: str, num_frames: int = 17,
                             h: int = 128, w: int = 128, device="cuda") -> dict:
    """Text-to-video: vision branch with T>1. Same model, same loss — only T differs."""
    video = (torch.randn(1, 3, num_frames, h, w, device=device) * 0.3).clamp(-1, 1)
    return {
        model.input_video_key:   [video],
        model.input_caption_key: [caption],
        "text_token_ids":        [_tokenize(model, caption, device)],
        "image_size":            [torch.tensor([[h, w, h, w]], dtype=torch.float32, device=device)],
        "fps":              torch.tensor([16.0], device=device),
        "conditioning_fps": torch.tensor([16.0], device=device),
        "num_frames":       torch.tensor([num_frames], device=device),
        "is_preprocessed":  True,
    }


def make_sound_video_batch(model, *, caption: str, num_video_frames: int = 5,
                           audio_hop_count: int = 8, h: int = 128, w: int = 128,
                           device="cuda") -> dict:
    """Joint text→video+sound batch (t2vs mode).

    Requires `model.config.sound_gen=True`. The model's AVAE expects stereo
    audio at 48 kHz with hop_size=1920 (Cosmos3-Nano defaults), so we round
    `num_audio_samples = audio_hop_count * 1920`. Audio and video duration
    don't have to match exactly; cosmos_framework handles temporal alignment via RoPE
    fps modulation in `_get_sound_fps_for_rope`.
    """
    # Stereo (AVAE expects 2 channels). 8 hops × 1920 = 15360 samples = 0.32 s @ 48 kHz.
    audio_channels = 2
    num_audio_samples = audio_hop_count * 1920
    waveform = (torch.randn(audio_channels, num_audio_samples, device=device) * 0.1).clamp(-1, 1)

    video = (torch.randn(1, 3, num_video_frames, h, w, device=device) * 0.3).clamp(-1, 1)

    # Sequence plan has both vision and sound; default condition indexes ([]) mean
    # all frames / all sound latent steps are noised and supervised.
    sequence_plan = SequencePlan(
        has_text=True,
        has_vision=True,
        has_sound=True,
    )

    return {
        model.input_video_key:   [video],
        "sound":                 [waveform],
        model.input_caption_key: [caption],
        "text_token_ids":        [_tokenize(model, caption, device)],
        "image_size":            [torch.tensor([[h, w, h, w]], dtype=torch.float32, device=device)],
        "fps":              torch.tensor([16.0], device=device),
        "conditioning_fps": torch.tensor([16.0], device=device),
        "num_frames":       torch.tensor([num_video_frames], device=device),
        "sequence_plan":    [sequence_plan],
        "is_preprocessed":  True,
    }


def make_action_fdm_batch(model, *, caption: str, num_video_frames: int = 5,
                          action_chunk: int = 4, raw_action_dim: int = 7,
                          h: int = 128, w: int = 128,
                          domain_name: str = "bridge_orig_lerobot", device="cuda") -> dict:
    """Action forward-dynamics: predict future video given 1st frame + action sequence.

    Requires `model.config.action_gen=True`. The batch contract is a superset of
    the vision batch: the same `video` / text fields plus an `action` tensor, a
    `domain_id` (cross-embodiment routing), `raw_action_dim` (un-padded dim;
    cosmos_framework pads to `max_action_dim`), `mode`, and a hand-built `sequence_plan`.
    See `cosmos_framework/inference/action.py: build_action_batch` for the canonical impl.

    `domain_name` selects the cross-embodiment routing; see
    `cosmos_framework/data/vfm/action/domain_utils.py` for the full list of supported
    embodiments.
    """
    # First frame is the conditioning anchor; remaining frames are predicted.
    video = (torch.randn(1, 3, num_video_frames, h, w, device=device) * 0.3).clamp(-1, 1)  # [1, C, T, H, W]

    # Pad raw action (e.g. 7-DoF: xyz + rpy + gripper) to max_action_dim.
    action = torch.zeros(action_chunk, model.config.max_action_dim, device=device)
    action[:, :raw_action_dim] = torch.randn(action_chunk, raw_action_dim, device=device) * 0.1

    # Hand-built sequence plan tells the packer which frames are conditioning.
    sequence_plan = build_sequence_plan_from_mode(
        mode="forward_dynamics",
        video_length=num_video_frames,
        action_length=action_chunk,
        has_text=True,
    )

    # Note: the inference-side `build_action_batch` uses `[[video]]` (nested) but
    # the training-side _normalize_video_databatch_inplace expects a flat list of
    # tensors when is_preprocessed=True. Use the flat-list convention here.
    return {
        model.input_video_key:   [video],
        "action":                [action],
        "raw_action_dim":        [torch.tensor(raw_action_dim, dtype=torch.long, device=device)],
        "mode":                  ["forward_dynamics"],
        model.input_caption_key: [caption],
        "text_token_ids":        [_tokenize(model, caption, device)],
        "image_size":            [torch.tensor([[h, w, h, w]], dtype=torch.float32, device=device)],
        "fps":              torch.tensor([16.0], device=device),
        "conditioning_fps": torch.tensor([16.0], device=device),
        "num_frames":       torch.tensor([num_video_frames], device=device),
        "domain_id":        [torch.tensor(get_domain_id(domain_name), dtype=torch.long, device=device)],
        "sequence_plan":    [sequence_plan],
        "is_preprocessed":  True,
    }


# ────────────────────────────────────────────────────────────────────────────
# Main loop. Three things only: build batch → training_step → backward+step.
# ────────────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--config-dir",
        type=str,
        default=None,
        help="Local directory containing config.json (architecture only — weights are "
             "randomly initialized). If omitted, fetches Cosmos3-Nano's config.json from HF.",
    )
    parser.add_argument("--num-iters", type=int, default=4)
    args = parser.parse_args()

    output_dir = Path("outputs/trainer_level_training").absolute()
    output_dir.mkdir(parents=True, exist_ok=True)

    # 1) Build the bare OmniMoTModel (random weights — see module docstring) ---
    model = _load_omni_model(config_dir_arg=args.config_dir)
    model.train()

    print(f"Modality flags: vision_gen={model.config.vision_gen}, "
          f"action_gen={model.config.action_gen}, sound_gen={model.config.sound_gen}")

    # 2) Optimizer — SGD (zero state) so the demo fits on a single 80GB GPU.
    #    Production cosmos_framework training uses AdamW with FSDP across ≥ 8 GPUs.
    optimizer = torch.optim.SGD(
        [p for p in model.parameters() if p.requires_grad],
        lr=1e-5,
    )

    # 3) Build an alternating multi-modality stream -----------------------
    caption_img = "A neon city street at night, rain reflecting the signs."
    caption_vid = "A camera dollies through a forest of giant glowing mushrooms."
    caption_act = "A robot arm picks up a red block from the table."
    caption_snd = "Wind howling through pine trees, distant thunder."

    def next_batch(it: int):
        # Round-robin through 4 modalities. Replace with your real dataloader.
        kind = ["T2I", "T2V", "ACTION_FDM", "T2VS"][it % 4]
        if kind == "T2I":
            return (kind, make_text_to_image_batch(model, caption=caption_img))
        if kind == "T2V":
            return (kind, make_text_to_video_batch(model, caption=caption_vid))
        if kind == "ACTION_FDM":
            return (kind, make_action_fdm_batch(model, caption=caption_act))
        return (kind, make_sound_video_batch(model, caption=caption_snd))

    # 4) Training loop ----------------------------------------------------
    # model.training_step does, end-to-end:
    #   tokenize text → VAE-encode video → sample t & noise (rectified flow)
    #   → pack tokens → run MoT network → flow-matching velocity loss.
    # We just call it.
    for it in range(args.num_iters):
        kind, batch = next_batch(it)

        aux, loss = model.training_step(batch, iteration=it)

        loss.backward()
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

        print(f"iter {it:>3d}  [{kind}]  loss={loss.item():.4f}")

    # 5) Save weights — plain torch.save ----------------------------------
    # NOTE: production cosmos_framework writes sharded DCP via cosmos_framework.utils.checkpoint
    # (FSDP-aware, resumable). torch.save is fine for this single-GPU demo
    # but won't capture FSDP shards or optimizer state.
    save_path = output_dir / "model.pt"
    torch.save(model.state_dict(), save_path)
    print(f"Saved weights: {save_path}")


if __name__ == "__main__":
    main()