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:

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:

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