[feat] cosmos3: strict-load verified (identity; needs_conversion=no)
This commit is contained in:
SolitaryThinker
@@ -1,246 +1,87 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Convert a Diffusers-format Cosmos3 checkpoint to FastVideo state-dict layout.
|
||||
"""Cosmos3 checkpoint strict-load verifier (no weight conversion required).
|
||||
|
||||
Cosmos3 weights (when published as ``nvidia/Cosmos3-Nano``) ship in a
|
||||
Diffusers-style namespace with ``model.embed_tokens`` /
|
||||
``model.layers.{i}.self_attn.*_moe_gen`` / ``vae2llm`` / ``llm2vae`` /
|
||||
``time_embedder`` layout. FastVideo's ``Cosmos3VFMTransformer`` splits the
|
||||
UND/GEN paths into separate ``language_model.layers.{i}`` and
|
||||
``gen_layers.{i}`` subtrees. This script applies the 14-rule remap from
|
||||
``Cosmos3OmniDiffusersPipeline._remap_ckpt_key`` to produce a state-dict
|
||||
that loads cleanly into
|
||||
``fastvideo.models.dits.cosmos3.Cosmos3VFMTransformer``.
|
||||
The published ``nvidia/Cosmos3-Nano`` checkpoint is diffusers-format and its
|
||||
transformer weight keys map 1:1 (identity) onto FastVideo's native
|
||||
``Cosmos3VFMTransformer`` parameters -- ``needs_conversion=no``. There is no
|
||||
remap to apply; the checkpoint loads directly.
|
||||
|
||||
Reference: vllm-omni PR #3454 (https://github.com/vllm-project/vllm-omni/pull/3454),
|
||||
HEAD ``8536f5b1421f``; upstream ``pipeline_cosmos3.py:319-411``.
|
||||
This utility verifies strict-load completeness (every checkpoint key has a
|
||||
matching DiT parameter of the right shape, and every DiT parameter is provided
|
||||
by the checkpoint) without allocating the full ~30 GB model, by reading
|
||||
safetensors headers and instantiating the DiT on the ``meta`` device.
|
||||
|
||||
Usage:
|
||||
python scripts/checkpoint_conversion/cosmos3_convert.py \\
|
||||
--src /path/to/nvidia--Cosmos3-Nano \\
|
||||
--dst /path/to/cosmos3_nano_fastvideo
|
||||
python scripts/checkpoint_conversion/cosmos3_convert.py --smoke-test
|
||||
python scripts/checkpoint_conversion/cosmos3_convert.py \
|
||||
--transformer official_weights/cosmos3/transformer
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
import glob
|
||||
import os
|
||||
import re
|
||||
|
||||
import safetensors.torch
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
|
||||
from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import (
|
||||
Cosmos3OmniDiffusersPipeline,
|
||||
)
|
||||
from fastvideo.configs.models.dits.cosmos3 import Cosmos3VideoConfig
|
||||
from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer
|
||||
|
||||
|
||||
def convert_state_dict(
|
||||
src_state: dict[str, torch.Tensor],
|
||||
) -> tuple[dict[str, torch.Tensor], list[str], list[str]]:
|
||||
"""Apply the Cosmos3 remap to a single state dict.
|
||||
|
||||
Returns ``(new_state, skipped_keys, unmapped_keys)`` where:
|
||||
|
||||
* ``new_state`` is the FastVideo-namespaced state dict.
|
||||
* ``skipped_keys`` lists source keys deliberately dropped by the remap
|
||||
(currently only ``lm_head.*``; the language-model head is not part of
|
||||
the diffusion transformer).
|
||||
* ``unmapped_keys`` lists source keys that did not match any remap rule.
|
||||
Should be empty for a clean Diffusers Cosmos3 checkpoint; non-empty
|
||||
results signal either an unexpected upstream key or a missing rule.
|
||||
"""
|
||||
new_state: dict[str, torch.Tensor] = {}
|
||||
skipped: list[str] = []
|
||||
unmapped: list[str] = []
|
||||
for src_key, tensor in src_state.items():
|
||||
new_key = Cosmos3OmniDiffusersPipeline._remap_ckpt_key(src_key)
|
||||
if new_key is None:
|
||||
# Either a deliberate skip (lm_head.*) or an unmapped key.
|
||||
if src_key.startswith("lm_head.") or src_key.startswith(
|
||||
"transformer.lm_head."):
|
||||
skipped.append(src_key)
|
||||
else:
|
||||
unmapped.append(src_key)
|
||||
continue
|
||||
if new_key in new_state:
|
||||
raise KeyError(
|
||||
f"Duplicate target key {new_key!r} produced by remap; "
|
||||
f"colliding source keys include {src_key!r}")
|
||||
new_state[new_key] = tensor
|
||||
return new_state, skipped, unmapped
|
||||
def checkpoint_key_shapes(transformer_dir: str) -> dict[str, tuple[int, ...]]:
|
||||
"""Read ``{key: shape}`` from a sharded safetensors transformer dir."""
|
||||
shards = sorted(glob.glob(os.path.join(transformer_dir, "*.safetensors")))
|
||||
if not shards:
|
||||
raise FileNotFoundError(f"no .safetensors found in {transformer_dir}")
|
||||
shapes: dict[str, tuple[int, ...]] = {}
|
||||
for shard in shards:
|
||||
with safe_open(shard, framework="pt") as handle:
|
||||
for key in handle.keys():
|
||||
shapes[key] = tuple(handle.get_slice(key).get_shape())
|
||||
return shapes
|
||||
|
||||
|
||||
def convert_checkpoint_dir(src: Path, dst: Path) -> None:
|
||||
"""Convert all ``*.safetensors`` shards in ``src`` and write to ``dst``.
|
||||
def verify_strict_load(transformer_dir: str) -> None:
|
||||
"""Raise SystemExit if the checkpoint does not strict-load into the DiT."""
|
||||
ckpt = checkpoint_key_shapes(transformer_dir)
|
||||
cfg = Cosmos3VideoConfig()
|
||||
with torch.device("meta"):
|
||||
dit = Cosmos3VFMTransformer(cfg, hf_config={})
|
||||
params = {name: tuple(p.shape) for name, p in dit.named_parameters()}
|
||||
buffers = {name for name, _ in dit.named_buffers()}
|
||||
name_map: dict[str, str] = cfg.arch_config.param_names_mapping
|
||||
|
||||
Reads each safetensors shard, applies the remap, and writes a single
|
||||
consolidated safetensors file at ``dst/model.safetensors``. This keeps
|
||||
the converter dependency surface small (single file out) for early
|
||||
Phase 5 work; once shard splitting becomes necessary, switch to
|
||||
per-shard outputs and emit ``model.safetensors.index.json``.
|
||||
def remap(key: str) -> str:
|
||||
for pattern, replacement in name_map.items():
|
||||
if re.match(pattern, key):
|
||||
return re.sub(pattern, replacement, key)
|
||||
return key
|
||||
|
||||
TODO: Emit a sharded layout with ``model.safetensors.index.json`` when
|
||||
the real ``nvidia/Cosmos3-Nano`` weights land and exceed the
|
||||
single-file safetensors size budget.
|
||||
"""
|
||||
src = Path(src)
|
||||
dst = Path(dst)
|
||||
if not src.exists():
|
||||
raise FileNotFoundError(f"Source checkpoint directory not found: {src}")
|
||||
if not src.is_dir():
|
||||
raise NotADirectoryError(
|
||||
f"Source path is not a directory: {src}")
|
||||
mapped = {remap(key): shape for key, shape in ckpt.items()}
|
||||
unexpected = sorted(set(mapped) - set(params) - buffers)
|
||||
missing = sorted(set(params) - set(mapped))
|
||||
mismatched = [(k, mapped[k], params[k]) for k in (set(mapped) & set(params)) if mapped[k] != params[k]]
|
||||
|
||||
shard_paths = sorted(src.glob("*.safetensors"))
|
||||
if not shard_paths:
|
||||
raise FileNotFoundError(
|
||||
f"No *.safetensors files found in {src}. Expected a Diffusers-format "
|
||||
"checkpoint directory.")
|
||||
|
||||
dst.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"Loading {len(shard_paths)} shard(s) from {src}")
|
||||
src_state: dict[str, torch.Tensor] = {}
|
||||
for shard_path in shard_paths:
|
||||
print(f" Reading {shard_path.name}")
|
||||
shard = safetensors.torch.load_file(str(shard_path))
|
||||
for key, tensor in shard.items():
|
||||
if key in src_state:
|
||||
raise KeyError(
|
||||
f"Duplicate key {key!r} found across shards "
|
||||
f"(latest in {shard_path.name}); aborting.")
|
||||
src_state[key] = tensor
|
||||
|
||||
print(f"Loaded {len(src_state)} parameters; applying remap")
|
||||
new_state, skipped, unmapped = convert_state_dict(src_state)
|
||||
print(f" Remapped {len(new_state)} parameters")
|
||||
print(f" Skipped {len(skipped)} parameters (lm_head.*)")
|
||||
if unmapped:
|
||||
print(f" WARNING: {len(unmapped)} unmapped key(s):")
|
||||
for k in unmapped[:10]:
|
||||
print(f" {k}")
|
||||
if len(unmapped) > 10:
|
||||
print(f" ... and {len(unmapped) - 10} more")
|
||||
|
||||
out_path = dst / "model.safetensors"
|
||||
print(f"Writing {out_path}")
|
||||
safetensors.torch.save_file(new_state, str(out_path))
|
||||
size_gb = out_path.stat().st_size / (1024**3)
|
||||
print(f" Wrote {size_gb:.2f} GB to {out_path}")
|
||||
if unexpected or missing or mismatched:
|
||||
raise SystemExit("strict-load FAILED: "
|
||||
f"unexpected={unexpected[:10]} missing={missing[:10]} "
|
||||
f"shape_mismatch={mismatched[:10]}")
|
||||
print(f"strict-load OK: {len(ckpt)} checkpoint keys map 1:1 onto "
|
||||
f"{len(params)} DiT params (identity; needs_conversion=no)")
|
||||
|
||||
|
||||
def smoke_test() -> int:
|
||||
"""Exercise the remap on a synthetic state dict covering all rule branches.
|
||||
|
||||
Runs WITHOUT real weights and acts as a CI guard against drift between
|
||||
``_remap_ckpt_key`` and this converter. Returns 0 on success, 1 on
|
||||
failure.
|
||||
|
||||
The synthetic input contains one representative key per remap branch:
|
||||
top-level (``vae2llm`` / ``llm2vae`` / ``time_embedder``), language-model
|
||||
embeddings/norms, the standalone ``norm_moe_gen``, both UND and GEN
|
||||
attention/norm/MLP variants inside ``model.layers.{i}.*``, and the
|
||||
deliberately-skipped ``lm_head.weight``.
|
||||
"""
|
||||
expected: dict[str, str | None] = {
|
||||
# Language-model trunk (UND path).
|
||||
"model.embed_tokens.weight":
|
||||
"transformer.language_model.embed_tokens.weight",
|
||||
"model.norm.weight":
|
||||
"transformer.language_model.norm.weight",
|
||||
# Standalone GEN norm at the top of the trunk.
|
||||
"model.norm_moe_gen.weight": "transformer.norm_moe_gen.weight",
|
||||
# UND self-attention (representative: q_proj).
|
||||
"model.layers.3.self_attn.q_proj.weight":
|
||||
"transformer.language_model.layers.3.self_attn.q_proj.weight",
|
||||
# GEN cross-attention via *_moe_gen suffix (representative: q_proj).
|
||||
"model.layers.3.self_attn.q_proj_moe_gen.weight":
|
||||
"transformer.gen_layers.3.cross_attention.q_proj.weight",
|
||||
# GEN cross-attention norm (representative: k_norm_moe_gen).
|
||||
"model.layers.3.self_attn.k_norm_moe_gen.weight":
|
||||
"transformer.gen_layers.3.cross_attention.k_norm.weight",
|
||||
# UND block norms.
|
||||
"model.layers.3.input_layernorm.weight":
|
||||
"transformer.language_model.layers.3.input_layernorm.weight",
|
||||
# GEN block norms.
|
||||
"model.layers.3.input_layernorm_moe_gen.weight":
|
||||
"transformer.gen_layers.3.input_layernorm.weight",
|
||||
# UND MLP (representative: gate_proj).
|
||||
"model.layers.3.mlp.gate_proj.weight":
|
||||
"transformer.language_model.layers.3.mlp.gate_proj.weight",
|
||||
# GEN MLP (representative: up_proj).
|
||||
"model.layers.3.mlp_moe_gen.up_proj.weight":
|
||||
"transformer.gen_layers.3.mlp.up_proj.weight",
|
||||
# Top-level generation projections.
|
||||
"vae2llm.weight": "transformer.vae2llm.weight",
|
||||
"llm2vae.weight": "transformer.llm2vae.weight",
|
||||
"time_embedder.linear_1.weight":
|
||||
"transformer.time_embedder.linear_1.weight",
|
||||
# Deliberately skipped.
|
||||
"lm_head.weight": None,
|
||||
}
|
||||
|
||||
synthetic = {k: torch.zeros(1) for k in expected}
|
||||
new_state, skipped, unmapped = convert_state_dict(synthetic)
|
||||
|
||||
failures: list[str] = []
|
||||
for src_key, expected_tgt in expected.items():
|
||||
if expected_tgt is None:
|
||||
if src_key not in skipped:
|
||||
failures.append(
|
||||
f" expected {src_key!r} in skipped; "
|
||||
f"new_state contains: {sorted(new_state)}; "
|
||||
f"unmapped contains: {unmapped}")
|
||||
continue
|
||||
if expected_tgt not in new_state:
|
||||
failures.append(
|
||||
f" expected target key {expected_tgt!r} (from src "
|
||||
f"{src_key!r}) in new_state; got keys {sorted(new_state)}")
|
||||
if unmapped:
|
||||
failures.append(
|
||||
f" unmapped keys (should be empty): {unmapped}")
|
||||
|
||||
if failures:
|
||||
print("smoke_test FAILED:")
|
||||
for line in failures:
|
||||
print(line)
|
||||
return 1
|
||||
print(
|
||||
f"smoke_test PASSED: {len(new_state)} new keys, "
|
||||
f"{len(skipped)} skipped, {len(unmapped)} unmapped")
|
||||
return 0
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Convert a Diffusers-format Cosmos3 checkpoint to "
|
||||
"FastVideo state-dict layout.")
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--src",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Diffusers-format source checkpoint directory (e.g. an HF "
|
||||
"snapshot of nvidia/Cosmos3-Nano).")
|
||||
parser.add_argument(
|
||||
"--dst",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Output directory; the converted state dict is written to "
|
||||
"<dst>/model.safetensors.")
|
||||
parser.add_argument(
|
||||
"--smoke-test",
|
||||
action="store_true",
|
||||
help="Run a no-weights smoke test that exercises every remap branch "
|
||||
"on a synthetic state dict. Useful as a CI guard.")
|
||||
"--transformer",
|
||||
default=os.path.join("official_weights", "cosmos3", "transformer"),
|
||||
help="path to the checkpoint transformer/ directory",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.smoke_test:
|
||||
return smoke_test()
|
||||
if args.src is None or args.dst is None:
|
||||
parser.error("--src and --dst are required when not running --smoke-test")
|
||||
convert_checkpoint_dir(args.src, args.dst)
|
||||
return 0
|
||||
verify_strict_load(args.transformer)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
main()
|
||||
|
||||
@@ -0,0 +1,80 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Strict-load completeness: real Cosmos3-Nano transformer <-> FastVideo DiT.
|
||||
|
||||
The published ``nvidia/Cosmos3-Nano`` checkpoint is diffusers-format
|
||||
(``needs_conversion=no``). This test verifies that EVERY transformer weight key
|
||||
in the checkpoint maps 1:1 (via the DiT's ``param_names_mapping``) onto a
|
||||
FastVideo ``Cosmos3VFMTransformer`` parameter of matching shape, and that no DiT
|
||||
parameter is left unfilled -- i.e. a ``strict=True`` load will succeed.
|
||||
|
||||
It runs on the ``meta`` device (no 30 GB allocation, no GPU) by reading only the
|
||||
safetensors headers. Skips cleanly if the checkpoint is not present locally.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import glob
|
||||
import os
|
||||
import re
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = [pytest.mark.local]
|
||||
|
||||
_CKPT_DIR = os.path.join("official_weights", "cosmos3", "transformer")
|
||||
|
||||
|
||||
def _checkpoint_key_shapes() -> dict[str, tuple[int, ...]]:
|
||||
from safetensors import safe_open
|
||||
shards = sorted(glob.glob(os.path.join(_CKPT_DIR, "*.safetensors")))
|
||||
if not shards:
|
||||
pytest.skip(f"Cosmos3 transformer checkpoint not present: {_CKPT_DIR}")
|
||||
out: dict[str, tuple[int, ...]] = {}
|
||||
for shard in shards:
|
||||
with safe_open(shard, framework="pt") as f:
|
||||
for k in f.keys():
|
||||
out[k] = tuple(f.get_slice(k).get_shape())
|
||||
return out
|
||||
|
||||
|
||||
def _meta_dit():
|
||||
from fastvideo.configs.models.dits.cosmos3 import Cosmos3VideoConfig
|
||||
from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer
|
||||
cfg = Cosmos3VideoConfig()
|
||||
with torch.device("meta"):
|
||||
dit = Cosmos3VFMTransformer(cfg, hf_config={})
|
||||
return dit, cfg
|
||||
|
||||
|
||||
def _apply_mapping(key: str, pmap: dict[str, str]) -> str:
|
||||
for pat, repl in pmap.items():
|
||||
if re.match(pat, key):
|
||||
return re.sub(pat, repl, key)
|
||||
return key
|
||||
|
||||
|
||||
def test_strict_load_completeness():
|
||||
ckpt = _checkpoint_key_shapes()
|
||||
dit, cfg = _meta_dit()
|
||||
dit_params = {n: tuple(p.shape) for n, p in dit.named_parameters()}
|
||||
dit_buffers = {n: tuple(b.shape) for n, b in dit.named_buffers()}
|
||||
pmap = cfg.arch_config.param_names_mapping
|
||||
|
||||
mapped = {_apply_mapping(k, pmap): v for k, v in ckpt.items()}
|
||||
ckpt_names = set(mapped)
|
||||
param_names = set(dit_params)
|
||||
buffer_names = set(dit_buffers)
|
||||
|
||||
# Every checkpoint key must land on a DiT parameter (non-persistent buffers excepted).
|
||||
unexpected = sorted(ckpt_names - param_names - buffer_names)
|
||||
assert not unexpected, f"checkpoint keys with no DiT param: {unexpected[:20]}"
|
||||
|
||||
# Every DiT parameter must be provided by the checkpoint (true strict load).
|
||||
missing = sorted(param_names - ckpt_names)
|
||||
assert not missing, f"DiT params not provided by checkpoint: {missing[:20]}"
|
||||
|
||||
# Shapes must match exactly.
|
||||
mism = [(k, mapped[k], dit_params[k]) for k in (ckpt_names & param_names) if mapped[k] != dit_params[k]]
|
||||
assert not mism, f"shape mismatches: {mism[:10]}"
|
||||
|
||||
assert len(ckpt) == len(dit_params), f"key count mismatch: ckpt={len(ckpt)} dit={len(dit_params)}"
|
||||
Reference in New Issue
Block a user