rooty2020

README: document config.yaml and export command

2e927b2 verified 2 months ago

1.99 kB

---
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).

	---
	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).