---
library_name: torch
tags:
  - cosmos
  - robotics
  - torch-dcp
---

# Cosmos-H2R

Raw PyTorch **Distributed Checkpoint (DCP)** shards for a Cosmos3 `OmniMoTModel`
(vision SFT, nano, curriculum stage 3), taken at training iteration **1068**.

## Contents

```
config.yaml          # full training config for the run that produced this checkpoint
model/
  .metadata          # DCP metadata — required for loading
  __0_0.distcp ... __31_0.distcp   # 32 rank shards (~86 GB total)
```

`config.yaml` is the Hydra config dumped by the training job (8 nodes x 4 GPUs =
32 ranks, `job.project=cosmos3`, `job.group=sft`,
`job.name=vision_sft_nano_curr_stage3`). The model is built from its `model:`
block (`cosmos_framework.model.vfm.omni_mot_model.OmniMoTModel`); the rest
records the dataloader, optimizer and scheduler settings used for the run.

Only the model state is published here; optimizer, scheduler, and trainer state
are not included, so this checkpoint is for **inference / fine-tuning init**, not
for resuming the original training run.

## Loading

This is **not** an HF `from_pretrained` checkpoint. Load it with
`torch.distributed.checkpoint`:

```python
import torch.distributed.checkpoint as dcp
from torch.distributed.checkpoint.filesystem import FileSystemReader
from torch.distributed.checkpoint.state_dict import get_model_state_dict

state_dict = get_model_state_dict(model)          # your OmniMoTModel
dcp.load(state_dict=state_dict, storage_reader=FileSystemReader("model"))
```

Within the `cosmos_framework` codebase, the checkpoint directory can be passed
directly to the DCP checkpointer, or converted to an HF-native
`Cosmos3OmniModel` directory with:

```bash
python -m cosmos_framework.scripts.export_model \
    --config-file config.yaml \
    --checkpoint-path model \
    -o <out dir>
```

The state dict holds 2138 tensors: 1069 under the `net.` prefix (regular weights)
and 1069 under `net_ema.` (EMA weights, as EMA was enabled during training).