Spaces:
Running on L40S
Running on L40S
File size: 15,765 Bytes
9f818c5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 | # 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 |
|