# Code Structure

<!--TOC-->

______________________________________________________________________

**Table of Contents**

- [Repository Layout](#repository-layout)
- [The `cosmos_framework/` Package](#the-cosmos_framework-package)
  - [`cosmos_framework/algorithm/`](#cosmos_frameworkalgorithm)
  - [`cosmos_framework/callbacks/`](#cosmos_frameworkcallbacks)
  - [`cosmos_framework/checkpoint/`](#cosmos_frameworkcheckpoint)
  - [`cosmos_framework/communicator/`](#cosmos_frameworkcommunicator)
  - [`cosmos_framework/configs/`](#cosmos_frameworkconfigs)
  - [`cosmos_framework/controller/`](#cosmos_frameworkcontroller)
  - [`cosmos_framework/data/`](#cosmos_frameworkdata)
  - [`cosmos_framework/evaluation/`](#cosmos_frameworkevaluation)
  - [`cosmos_framework/inference/`](#cosmos_frameworkinference)
  - [`cosmos_framework/launcher/`](#cosmos_frameworklauncher)
  - [`cosmos_framework/model/`](#cosmos_frameworkmodel)
  - [`cosmos_framework/tools/`](#cosmos_frameworktools)
  - [`cosmos_framework/trainer/`](#cosmos_frameworktrainer)
  - [`cosmos_framework/utils/`](#cosmos_frameworkutils)
  - [`cosmos_framework/workers/`](#cosmos_frameworkworkers)
- [Supporting Directories](#supporting-directories)
- [Where to Add New Code](#where-to-add-new-code)

______________________________________________________________________

<!--TOC-->

## Repository Layout

```text
Cosmos/
├── cosmos_framework/             # Main package (training infra + inference subpackage)
│   ├── inference/      # Inference subpackage (model, args, defaults, Ray serving, common helpers, SFT experiment configs)
│   └── ...             # Training-infra subpackages: data, model, trainer, callbacks, checkpoint, …
│   └── scripts/        # CLI entry-point scripts: train.py, _train.py, inference.py, eval.py, export_model.py, …
├── packages/           # Backend shim packages: diffusers-cosmos3, transformers-cosmos3, vllm-cosmos3
├── docs/               # User documentation (you are here)
├── docker/             # Dockerfiles for reproducible environments
├── examples/           # Runnable training / fine-tuning / inference examples
├── tests/              # Unit and integration tests
├── tools/              # Standalone CLI utilities (e.g. checkpoint conversion)
├── pyproject.toml      # uv-managed dependency manifest
├── uv.lock             # Pinned dependency graph (do not edit by hand)
└── .python-version     # Python version pin (used by uv)
```

`cosmos_framework/` is the single Python package. Training infrastructure (data, model, trainer, callbacks, checkpoint, utils, …) lives in top-level subpackages; inference (Diffusers / Transformers / vLLM-friendly inference core, Ray serving, per-modality defaults, training-side experiment YAMLs) lives under `cosmos_framework/inference/`. The library-style backend shims that load Cosmos3 checkpoints into upstream ecosystems live under `packages/{diffusers,transformers,vllm}-cosmos3/`.

## The `cosmos_framework/` Package

The `cosmos_framework/` package is organized around the workflow of a large-scale, distributed training run — particularly post-training and reinforcement-learning regimes — with each subpackage owning one concern.

```text
cosmos_framework/
├── algorithm/      # Loss functions, reward models, RL algorithms        [planned]
│   ├── loss/
│   ├── reward/
│   └── rl/
├── callbacks/      # Lifecycle hooks (logging, profiling, eval triggers, checkpoint cadence)
├── checkpoint/     # Saving, loading, conversion (DCP ↔ HF safetensors)
├── communicator/   # Inter-process / inter-worker communication primitives  [planned]
├── configs/        # Pydantic-validated TOML schema + LazyConfig experiment SKUs
│   ├── base/       # config.py, experiment/{action,sft,posttrain_video}/, vlm/, defaults/
│   └── toml_config/# sft_config.py (pydantic schema), toml_config_helper.py
├── controller/     # Top-level orchestration of multi-worker training jobs  [planned]
├── data/           # Dataset loading, batching, augmentation, sharding
├── evaluation/     # Eval harness for trained checkpoints                [planned]
├── inference/      # Inference engine + modality defaults + Ray serving + shared helpers
├── launcher/       # Job launching (Slurm, torchrun, k8s)                [planned]
├── model/          # Model definitions and parallelism wrappers
├── tools/          # In-package CLI utilities (flops/, visualize/)
├── trainer/        # Training loop, optimizer step, gradient accumulation
├── utils/          # Shared low-level utilities (logging, config, distributed helpers, vfm/, vlm/)
└── workers/        # Specialized roles in a distributed RL job           [planned]
    ├── reference/  # Reference / frozen-policy worker (KL anchor)
    ├── reward/     # Reward-model worker
    ├── rollout/    # On-policy rollout generation worker
    └── simulations/# Simulator-driven environment worker
```

> **Note — some subpackages above don't exist yet.** Entries tagged **`[planned]`** (`algorithm/`, `communicator/`, `controller/`, `evaluation/`, `launcher/`, `workers/`) describe the intended home for functionality that has not yet landed in this release. The directories are intentionally absent on disk — they will be created when their respective feature areas (RL training, multi-worker orchestration, the distributed evaluation harness, job launchers, RL worker roles) come online. The "Where to Add New Code" table at the bottom of this document still applies: when you build one of these features, create the matching subpackage and put the code there. Subpackages without the tag (`callbacks/`, `checkpoint/`, `configs/`, `data/`, `inference/`, `model/`, `tools/`, `trainer/`, `utils/`, plus `scripts/` covered separately) exist today.

### `cosmos_framework/algorithm/`

*Planned — not yet present in this release; the subpackage will be created when RL / loss work lands.*

Algorithmic primitives that are independent of the model and trainer.

- `loss/` — supervised and distillation losses (cross-entropy, flow-matching, KL, etc.).
- `reward/` — reward functions and learned reward heads.
- `rl/` — RL update rules (PPO, GRPO, DPO-family) that consume losses and rewards.

Add new objectives here, not inside the trainer.

### `cosmos_framework/callbacks/`

Pluggable lifecycle hooks invoked by the trainer at well-defined points (step begin/end, epoch boundary, eval, save, exception). Use callbacks for cross-cutting concerns such as wandb/W&B logging, gradient clipping, MoE stability monitoring, dataloader-state checkpointing, and learning-rate logging.

### `cosmos_framework/checkpoint/`

All checkpoint I/O lives here:

- DCP (PyTorch Distributed Checkpoint) save/load
- HuggingFace `safetensors` import/export
- Schema migration and resume-from-step logic

### `cosmos_framework/communicator/`

*Planned — not yet present in this release; the subpackage will be created when multi-worker comms land.*

Communication primitives between processes — point-to-point send/recv, broadcast helpers, and any RPC-style channels used between the controller and workers. Keep raw `torch.distributed` / NCCL calls out of business logic; route them through this layer.

### `cosmos_framework/configs/`

Configuration system for training runs:

- `configs/toml_config/` — the user-facing pydantic schema for the structured TOML interface consumed by `cosmos_framework.scripts.train --sft-toml=…`. `sft_config.py` defines `SFTExperimentConfig` (with `extra="forbid"` on every sub-model). `toml_config_helper.py` handles VFM↔VLM path remaps and OmegaConf env interpolation.
- `configs/base/` — internal LazyConfig-based experiment SKUs. `base/experiment/sft/*.py` registers Cosmos3 SFT experiments (e.g. `vision_sft_nano.py`, `vision_sft_super.py`). `base/vlm/` is the VLM-side analogue with its own `config.py`, `experiment/`, `defaults/`, and `freeze_config.py`.

See [`docs/sft_config.md`](./sft_config.md) for the full field-by-field TOML reference.

### `cosmos_framework/controller/`

*Planned — not yet present in this release; the subpackage will be created when multi-worker orchestration lands.*

The orchestrator for a multi-worker job. The controller drives the training loop, hands batches to rollout/reward workers, collects gradients, and decides when to checkpoint or evaluate. Think "head node logic" — there is one controller per job.

### `cosmos_framework/data/`

Datasets, samplers, collators, augmentations, and data-side parallelism (e.g. sequence packing, multi-aspect batching). New dataset formats and new augmentations both live here.

### `cosmos_framework/evaluation/`

*Planned — not yet present in this release; the subpackage will be created when the offline eval harness lands.*

Evaluation harnesses run against trained checkpoints — metrics, dataset-driven eval loops, and reporting. Distinct from `inference/`: evaluation is offline and metric-oriented.

### `cosmos_framework/inference/`

The full inference subpackage:

- `args.py` — sampling/setup args (`SamplingArgs`, `SamplingOverrides`, `OmniSetupArgs`, `OmniSetupOverrides`, `OmniSampleOverrides`), plus the modality-defaults loader and the `_RESOLUTION_SHIFT_DEFAULTS` table.
- `model.py`, `inference.py` — model + inference engine entry points used by `cosmos_framework/scripts/inference.py`.
- `common/` — shared helpers for args, init, config, checkpoints (used by both training and inference scripts).
- `defaults/<mode>/sample_args.json` — per-modality default sample arguments (text2image, text2video, image2video, image2image, video2video, forward_dynamics, inverse_dynamics, policy, reasoner) plus `prompt_upsampler.txt` and `video_captioner.txt` system prompts.
- `ray/` — Ray Serve / Submit / Gradio entry points (`cosmos_framework.inference.ray.serve`, `cosmos_framework.inference.ray.submit`, `cosmos_framework.inference.ray.gradio`) and their YAML configs under `ray/configs/`.
- `configs/{checkpoint,model}/` — per-checkpoint and per-model inference configs.
- Modality entry points: `vision.py`, `action.py`, `sound.py`, `transfer.py`, `interactive.py`, `prompt_upsampling.py`, `dataset.py`.

Training-side experiment SKUs live separately at `cosmos_framework/configs/base/experiment/sft/*.py` (see [`cosmos_framework/configs/`](#cosmos_frameworkconfigs)) — not under `inference/`.

Library-style backend shims that adapt Cosmos3 checkpoints to the Diffusers / Transformers / vLLM ecosystems live separately under `packages/{diffusers,transformers,vllm}-cosmos3/`.

### `cosmos_framework/launcher/`

*Planned — not yet present in this release; the subpackage will be created when launcher back-ends land.*

Job launching back-ends: Slurm, `torchrun`, and Kubernetes adapters. Selects the launch path based on the environment and forwards process rank/world-size to the controller.

### `cosmos_framework/model/`

Model architectures and the parallelism wrappers around them (FSDP, tensor parallel, context parallel, pipeline parallel). The trainer is model-agnostic; everything the trainer touches goes through this layer.

### `cosmos_framework/tools/`

CLI entry points surfaced from the package (as opposed to standalone scripts in the top-level `tools/`). Use this for utilities that need to import `cosmos_framework.*` internals.

### `cosmos_framework/trainer/`

The training loop itself — gradient accumulation, optimizer step, scheduler step, mixed-precision policy, and the dispatcher that fires callbacks. Stays narrow on purpose: model details live in `model/`, algorithm details in `algorithm/`.

### `cosmos_framework/utils/`

Shared low-level helpers (logging, config loading, distributed setup, profiling). Keep this folder *thin* — anything substantial should grow into its own subpackage.

### `cosmos_framework/workers/`

*Planned — not yet present in this release; the subpackage will be created when RL worker roles land.*

Specialized worker roles for distributed RL jobs. Each worker is a long-running process the controller talks to:

- `reference/` — frozen reference policy (for KL anchoring in PPO/GRPO/DPO).
- `reward/` — reward-model worker; computes scalar rewards for rollouts.
- `rollout/` — on-policy generation worker; samples trajectories from the current policy.
- `simulations/` — simulator-backed environment worker (used when reward comes from a sim rather than a learned model).

Add new worker types as sibling subpackages — each owns its own startup, message loop, and shutdown.

## Supporting Directories

- `tests/` — pytest tests, mirroring the `cosmos_framework/` package layout.
- `examples/` — runnable end-to-end examples; see `examples/README.md`.
- `docker/` — Dockerfiles and image build helpers; see `docker/README.md`.
- `cosmos_framework/scripts/` — CLI entry-point scripts (`train.py`, `inference.py`, `eval.py`, `export_model.py`, …); invoke as `python -m cosmos_framework.scripts.<name>`. Primary training entry point: `cosmos_framework.scripts.train` driven by a structured, pydantic-validated TOML interface (`--sft-toml=<recipe-toml>`); recipe TOMLs live under [`examples/toml/sft_config/`](../examples/toml/sft_config/) and the schema is defined in [`cosmos_framework/configs/toml_config/sft_config.py`](../cosmos_framework/configs/toml_config/sft_config.py) — see [`examples/README.md`](../examples/README.md) and [`docs/training.md`](./training.md).
- `packages/` — library-style backend shims: `packages/{diffusers,transformers,vllm}-cosmos3/`, each installable independently.

## Where to Add New Code

| You want to add…                            | Put it in…                                                                       |
| ------------------------------------------- | -------------------------------------------------------------------------------- |
| A new loss function                         | `cosmos_framework/algorithm/loss/`                                               |
| A new RL update rule                        | `cosmos_framework/algorithm/rl/`                                                 |
| A new reward function or head               | `cosmos_framework/algorithm/reward/`                                             |
| A new model architecture                    | `cosmos_framework/model/`                                                        |
| A new dataset format / augmentation         | `cosmos_framework/data/`                                                         |
| A new training callback                     | `cosmos_framework/callbacks/`                                                    |
| A new checkpoint format or converter        | `cosmos_framework/checkpoint/`                                                   |
| A new launcher back-end (Slurm flavor, k8s) | `cosmos_framework/launcher/`                                                     |
| A new RL worker role                        | `cosmos_framework/workers/<new_role>/`                                           |
| A new evaluation suite                      | `cosmos_framework/evaluation/`                                                   |
| A new runnable example                      | `examples/`                                                                      |
| A new standalone CLI tool                   | `tools/` (repo root) for non-cosmos imports, otherwise `cosmos_framework/tools/` |
| A new test                                  | `tests/` mirroring the package path                                              |