nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[feat] cosmos3: native DiT + framework parity (Cosmos3VFMTransformer)

Browse Source
This commit is contained in:
SolitaryThinker
2026-07-13 03:22:33 -07:00
parent b48eafe515
commit 6771f4ec2e
3 changed files with 1099 additions and 375 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/models/dits/cosmos3.py
+699 -284
View File
File diff suppressed because it is too large Load Diff
tests/local_tests/cosmos3/test_cosmos3_dit_parity.py
+256
View File
@@ -0,0 +1,256 @@
# SPDX-License-Identifier: Apache-2.0
"""Numerical-parity test: FastVideo Cosmos3 DiT vs official ``Cosmos3VFMNetwork``.
Builds a tiny official-framework ``Cosmos3VFMNetwork`` AND a tiny FastVideo
``Cosmos3VFMTransformer`` from the SAME tiny config, copies the framework
weights into the FastVideo DiT via an explicit framework->fastvideo name map,
runs BOTH forwards on identical deterministic inputs (CPU / float32), and
asserts ``torch.allclose`` on the vision prediction output (``preds_vision``)
and the per-token ``last_hidden_state``.
The official model is the parity ORACLE. It runs on CPU / float32 via the SDPA
monkey-patch in ``test_cosmos3_reference_forward`` (flash2/flash3/natten are
CUDA-only). The FastVideo DiT runs natively on CPU with plain SDPA.
Run:
cd <worktree> && <fv-cosmos3 python> -m pytest \
tests/local_tests/cosmos3/test_cosmos3_dit_parity.py -q
"""
from __future__ import annotations
import pytest
import torch
# The official framework provides the parity oracle.
cosmos_framework = pytest.importorskip(
"cosmos_framework",
reason="cosmos_framework not installed; run in fv-cosmos3 env.",
)
# Reuse the reference harness's tiny-model builder + SDPA monkey-patch.
from .test_cosmos3_reference_forward import ( # noqa: E402
_apply_sdpa_patches,
_build_tiny_cosmos3,
_build_tiny_packed_seq,
)
pytestmark = [pytest.mark.local]
# Ensure the CPU/float32 SDPA patches are installed (idempotent).
_apply_sdpa_patches()
# ---------------------------------------------------------------------------
# Tiny config shared by both models (must match _build_tiny_cosmos3).
# ---------------------------------------------------------------------------
def _build_tiny_fastvideo_dit() -> "Cosmos3VFMTransformer": # noqa: F821
from fastvideo.configs.models.dits.cosmos3 import (
Cosmos3ArchConfig,
Cosmos3VideoConfig,
)
from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer
arch = Cosmos3ArchConfig(
hidden_size=16,
num_hidden_layers=1,
num_attention_heads=2,
num_key_value_heads=2,
head_dim=8,
intermediate_size=32,
vocab_size=64,
rms_norm_eps=1e-6,
attention_bias=False,
latent_patch_size=2,
latent_channel=16,
rope_theta=5_000_000.0,
mrope_section=[24, 20, 20],
position_embedding_type="3d_rope",
base_fps=24.0,
temporal_compression_factor=4,
enable_fps_modulation=False,
# Dormant heads present in the checkpoint surface (constructed for
# strict-load parity; not exercised by this video-path forward).
action_gen=True,
action_dim=64,
max_action_dim=64,
num_embodiment_domains=32,
sound_gen=True,
sound_dim=64,
)
cfg = Cosmos3VideoConfig(arch_config=arch)
model = Cosmos3VFMTransformer(cfg, hf_config={})
return model.to(torch.float32).eval()
# ---------------------------------------------------------------------------
# Framework -> FastVideo weight name map.
# ---------------------------------------------------------------------------
def _framework_to_fastvideo_state_dict(vfm, num_layers: int) -> dict[str, torch.Tensor]:
"""Translate framework param names into the FastVideo DiT param names.
Framework (Cosmos3VFMNetwork):
language_model.model.{embed_tokens,norm,norm_moe_gen}
language_model.lm_head
language_model.model.layers.{i}.self_attn.{q,k,v,o}_proj(+ _moe_gen)
language_model.model.layers.{i}.self_attn.{q,k}_norm(+ _moe_gen)
language_model.model.layers.{i}.{mlp,mlp_moe_gen}.{gate,up,down}_proj
language_model.model.layers.{i}.{input,post_attention}_layernorm(+ _moe_gen)
vae2llm / llm2vae / time_embedder.mlp.{0,2}
FastVideo (Cosmos3VFMTransformer):
embed_tokens / norm / norm_moe_gen / lm_head
layers.{i}.self_attn.{to_q,to_k,to_v,to_out} (und)
layers.{i}.self_attn.{add_q,add_k,add_v}_proj / to_add_out (gen)
layers.{i}.self_attn.{norm_q,norm_k,norm_added_q,norm_added_k}
layers.{i}.{mlp,mlp_moe_gen}.{gate,up,down}_proj
layers.{i}.{input,post_attention}_layernorm(+ _moe_gen)
proj_in / proj_out / time_embedder.linear_{1,2}
"""
src = dict(vfm.named_parameters())
out: dict[str, torch.Tensor] = {}
def take(name: str) -> torch.Tensor:
return src[name].detach().clone()
# ---- Top-level backbone ----
out["embed_tokens.weight"] = take("language_model.model.embed_tokens.weight")
out["norm.weight"] = take("language_model.model.norm.weight")
out["norm_moe_gen.weight"] = take("language_model.model.norm_moe_gen.weight")
out["lm_head.weight"] = take("language_model.lm_head.weight")
# ---- Vision adapters ----
out["proj_in.weight"] = take("vae2llm.weight")
out["proj_in.bias"] = take("vae2llm.bias")
out["proj_out.weight"] = take("llm2vae.weight")
out["proj_out.bias"] = take("llm2vae.bias")
# ---- Timestep embedder (mlp.0/mlp.2 -> linear_1/linear_2) ----
out["time_embedder.linear_1.weight"] = take("time_embedder.mlp.0.weight")
out["time_embedder.linear_1.bias"] = take("time_embedder.mlp.0.bias")
out["time_embedder.linear_2.weight"] = take("time_embedder.mlp.2.weight")
out["time_embedder.linear_2.bias"] = take("time_embedder.mlp.2.bias")
# ---- Per layer ----
und_attn = {"q_proj": "to_q", "k_proj": "to_k", "v_proj": "to_v", "o_proj": "to_out"}
gen_attn = {
"q_proj_moe_gen": "add_q_proj",
"k_proj_moe_gen": "add_k_proj",
"v_proj_moe_gen": "add_v_proj",
"o_proj_moe_gen": "to_add_out",
}
und_norm = {"q_norm": "norm_q", "k_norm": "norm_k"}
gen_norm = {"q_norm_moe_gen": "norm_added_q", "k_norm_moe_gen": "norm_added_k"}
for i in range(num_layers):
fw = f"language_model.model.layers.{i}"
fv = f"layers.{i}"
for s, d in und_attn.items():
out[f"{fv}.self_attn.{d}.weight"] = take(f"{fw}.self_attn.{s}.weight")
for s, d in gen_attn.items():
out[f"{fv}.self_attn.{d}.weight"] = take(f"{fw}.self_attn.{s}.weight")
for s, d in und_norm.items():
out[f"{fv}.self_attn.{d}.weight"] = take(f"{fw}.self_attn.{s}.weight")
for s, d in gen_norm.items():
out[f"{fv}.self_attn.{d}.weight"] = take(f"{fw}.self_attn.{s}.weight")
for mlp in ("mlp", "mlp_moe_gen"):
for proj in ("gate_proj", "up_proj", "down_proj"):
out[f"{fv}.{mlp}.{proj}.weight"] = take(f"{fw}.{mlp}.{proj}.weight")
for ln in ("input_layernorm", "input_layernorm_moe_gen", "post_attention_layernorm",
"post_attention_layernorm_moe_gen"):
out[f"{fv}.{ln}.weight"] = take(f"{fw}.{ln}.weight")
return out
def _copy_weights(vfm, dit) -> None:
"""Copy framework weights into the FastVideo DiT (shape-checked)."""
mapped = _framework_to_fastvideo_state_dict(vfm, num_layers=dit.num_hidden_layers)
dst = dict(dit.named_parameters())
# Every mapped tensor must land on an existing FastVideo param with a matching shape.
for name, tensor in mapped.items():
assert name in dst, f"FastVideo DiT missing param for mapped key {name!r}"
assert dst[name].shape == tensor.shape, (f"shape mismatch for {name}: "
f"dit={tuple(dst[name].shape)} fw={tuple(tensor.shape)}")
with torch.no_grad():
for name, tensor in mapped.items():
dst[name].copy_(tensor.to(dst[name].dtype))
def _fastvideo_inputs_from_packed_seq(ps) -> dict:
"""Build the FastVideo DiT forward kwargs from a framework PackedSequence."""
v = ps.vision
return dict(
text_ids=ps.text_ids,
text_indexes=ps.text_indexes,
position_ids=ps.position_ids,
sequence_length=int(ps.sequence_length),
split_lens=list(ps.split_lens),
attn_modes=list(ps.attn_modes),
vision_tokens=list(v.tokens),
vision_token_shapes=list(v.token_shapes),
vision_sequence_indexes=v.sequence_indexes,
vision_timesteps=v.timesteps,
vision_mse_loss_indexes=v.mse_loss_indexes,
vision_noisy_frame_indexes=list(v.noisy_frame_indexes),
fps_vision=None,
)
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
class TestCosmos3DiTParity:
def _run_both(self, seed_model: int = 42, seed_data: int = 7):
vfm = _build_tiny_cosmos3(seed=seed_model)
dit = _build_tiny_fastvideo_dit()
_copy_weights(vfm, dit)
ps = _build_tiny_packed_seq(n_text=4, seed=seed_data)
with torch.no_grad():
fw_out = vfm(packed_seq=ps)
fv_out = dit(**_fastvideo_inputs_from_packed_seq(ps))
return fw_out, fv_out
def test_weight_map_is_complete(self):
"""The framework->fastvideo map must cover EVERY FastVideo DiT parameter
that is exercised by the video path (i.e. all non-dormant params).
Dormant action/audio heads have no framework counterpart in this tiny
vision-only setup, so they are excluded from the copy; everything else
must be covered.
"""
vfm = _build_tiny_cosmos3(seed=42)
dit = _build_tiny_fastvideo_dit()
mapped = set(_framework_to_fastvideo_state_dict(vfm, num_layers=dit.num_hidden_layers))
dit_params = set(n for n, _ in dit.named_parameters())
dormant = {
n
for n in dit_params
if n.startswith(("action_", "audio_"))
}
uncovered = dit_params - mapped - dormant
assert not uncovered, f"FastVideo DiT params not covered by weight map: {sorted(uncovered)}"
def test_preds_vision_parity(self):
fw_out, fv_out = self._run_both()
fw_pv = fw_out["preds_vision"][0] # [1, C, T, H, W]
fv_pv = fv_out["preds_vision"][0]
assert fw_pv.shape == fv_pv.shape, f"shape mismatch: fw={fw_pv.shape} fv={fv_pv.shape}"
max_abs = (fw_pv - fv_pv).abs().max().item()
print(f"\n[preds_vision] max abs diff = {max_abs:.3e}")
torch.testing.assert_close(fv_pv, fw_pv, atol=1e-4, rtol=1e-3)
def test_last_hidden_state_parity(self):
fw_out, fv_out = self._run_both()
fw_lhs = fw_out["last_hidden_state"] # [N, hidden]
fv_lhs = fv_out["last_hidden_state"]
assert fw_lhs.shape == fv_lhs.shape, f"shape mismatch: fw={fw_lhs.shape} fv={fv_lhs.shape}"
max_abs = (fw_lhs - fv_lhs).abs().max().item()
print(f"\n[last_hidden_state] max abs diff = {max_abs:.3e}")
torch.testing.assert_close(fv_lhs, fw_lhs, atol=1e-4, rtol=1e-3)
def test_parity_holds_across_seeds(self):
"""Re-running with a different random init still matches (not a fluke)."""
fw_out, fv_out = self._run_both(seed_model=99, seed_data=13)
torch.testing.assert_close(fv_out["preds_vision"][0], fw_out["preds_vision"][0], atol=1e-4, rtol=1e-3)
tests/local_tests/cosmos3/test_cosmos3_state_dict_keys.py
+144 -91
View File
@@ -1,28 +1,38 @@
# SPDX-License-Identifier: Apache-2.0
"""Cosmos3 state-dict-key contract for checkpoint conversion (Tier A scaffold).
"""Cosmos3 DiT param-name / checkpoint-key-surface contract.
Reference: ``vllm_omni/diffusion/models/cosmos3/pipeline_cosmos3.py:319-409``
(``Cosmos3OmniDiffusersPipeline._remap_ckpt_key``). The vllm-omni loader's
remap is the canonical source of truth: it defines which Diffusers
checkpoint keys become which target module-tree parameter names.
The FastVideo native Cosmos3 DiT (``fastvideo.models.dits.cosmos3.
Cosmos3VFMTransformer``) deliberately mirrors the *published diffusers
checkpoint* transformer key surface so the converter at
``scripts/checkpoint_conversion/cosmos3_convert.py`` can strict-load with a
near-identity ``param_names_mapping``.
The FastVideo conversion script at
``scripts/checkpoint_conversion/cosmos3_convert.py`` (to be authored in
Phase 5) must produce a state-dict that loads cleanly into the FastVideo
Cosmos3 DiT module tree. This test pins the target side of that contract:
This pins the FastVideo side of that contract: a tiny 1-layer DiT must expose
exactly the published checkpoint's parameter names. The checkpoint key surface
(validated against ``nvidia/Cosmos3-Nano``; ``{i}`` ranges over the layers):
* ``embed_tokens``/``norm`` live under ``language_model.``
* ``norm_moe_gen`` is a top-level (not per-layer) module
* Per-layer keys split into UND (``language_model.layers.{i}``) and
GEN (``gen_layers.{i}``) sub-trees:
- UND attention: ``self_attn.{q,k,v,o}_proj``, ``self_attn.{q,k}_norm``
- GEN cross-attention: ``cross_attention.{q,k,v,o}_proj``,
``cross_attention.{q,k}_norm`` (from the ``*_moe_gen`` source keys)
- UND/GEN norms: ``input_layernorm``, ``post_attention_layernorm``
- UND/GEN MLPs: ``mlp.{gate,up,down}_proj``
* Top-level adapters ``vae2llm``, ``llm2vae``, ``time_embedder`` live
under ``transformer.`` (no remapping inside the transformer namespace).
* ``lm_head.weight`` is skipped (mapped to ``None``).
Top level:
embed_tokens.weight, norm.weight, norm_moe_gen.weight, lm_head.weight,
proj_in.{weight,bias}, proj_out.{weight,bias},
time_embedder.linear_1.{weight,bias}, time_embedder.linear_2.{weight,bias}
Dormant heads (present for strict-load):
action_proj_in.fc.weight, action_proj_in.bias.weight,
action_proj_out.fc.weight, action_proj_out.bias.weight,
action_modality_embed,
audio_proj_in.{weight,bias}, audio_proj_out.{weight,bias},
audio_modality_embed
Per layer ``layers.{i}``:
self_attn.{to_q,to_k,to_v,to_out,add_q_proj,add_k_proj,add_v_proj,
to_add_out,norm_q,norm_k,norm_added_q,norm_added_k}.weight,
mlp.{gate_proj,up_proj,down_proj}.weight,
mlp_moe_gen.{gate_proj,up_proj,down_proj}.weight,
{input_layernorm,input_layernorm_moe_gen,post_attention_layernorm,
post_attention_layernorm_moe_gen}.weight
The earlier scaffold pinned the dead vllm-omni layout
(``language_model.layers`` / ``gen_layers`` / ``cross_attention`` /
``vae2llm`` / ``llm2vae``); that structure is gone — the native DiT is a single
dual-pathway ``layers`` ModuleList matching the diffusers checkpoint.
"""
from __future__ import annotations
@@ -31,86 +41,129 @@ import pytest
pytestmark = [pytest.mark.local]
EXPECTED_REMAPS: dict[str, str | None] = {
"model.embed_tokens.weight": "transformer.language_model.embed_tokens.weight",
"model.norm.weight": "transformer.language_model.norm.weight",
"model.norm_moe_gen.weight": "transformer.norm_moe_gen.weight",
"model.layers.3.self_attn.q_proj.weight": "transformer.language_model.layers.3.self_attn.q_proj.weight",
"model.layers.3.self_attn.q_proj_moe_gen.weight": "transformer.gen_layers.3.cross_attention.q_proj.weight",
"model.layers.3.self_attn.k_norm_moe_gen.weight": "transformer.gen_layers.3.cross_attention.k_norm.weight",
"model.layers.3.input_layernorm.weight": "transformer.language_model.layers.3.input_layernorm.weight",
"model.layers.3.input_layernorm_moe_gen.weight": "transformer.gen_layers.3.input_layernorm.weight",
"model.layers.3.mlp.gate_proj.weight": "transformer.language_model.layers.3.mlp.gate_proj.weight",
"model.layers.3.mlp_moe_gen.up_proj.weight": "transformer.gen_layers.3.mlp.up_proj.weight",
"vae2llm.weight": "transformer.vae2llm.weight",
"llm2vae.weight": "transformer.llm2vae.weight",
"time_embedder.linear_1.weight": "transformer.time_embedder.linear_1.weight",
"lm_head.weight": None,
}
def test_fastvideo_cosmos3_remap_matches_reference() -> None:
"""Asserts the FastVideo Cosmos3 checkpoint remap table matches the
canonical vllm-omni reference at pipeline_cosmos3.py:319-409.
Once the FastVideo Cosmos3 pipeline (or a standalone remap utility
in ``scripts/checkpoint_conversion/cosmos3_convert.py``) exposes a
``_remap_ckpt_key`` callable, this test verifies every key in
``EXPECTED_REMAPS`` maps to the expected target name (or ``None``
for skipped keys like ``lm_head.weight``).
"""
def _build_tiny_dit():
"""Construct a tiny 1-layer FastVideo Cosmos3 DiT, or skip if unavailable."""
try:
from fastvideo.pipelines.basic.cosmos3.cosmos3_pipeline import ( # type: ignore
Cosmos3OmniDiffusersPipeline,
from fastvideo.configs.models.dits.cosmos3 import (
Cosmos3ArchConfig,
Cosmos3VideoConfig,
)
except ImportError:
pytest.skip("FastVideo Cosmos3 pipeline not yet implemented (Phase 2b)")
from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer
except ImportError: # pragma: no cover - import guard
pytest.skip("FastVideo Cosmos3 DiT not importable in this environment")
remap = getattr(Cosmos3OmniDiffusersPipeline, "_remap_ckpt_key", None)
if remap is None:
pytest.skip("Cosmos3OmniDiffusersPipeline._remap_ckpt_key not yet defined")
import torch
actual = {key: remap(key) for key in EXPECTED_REMAPS}
assert actual == EXPECTED_REMAPS
arch = Cosmos3ArchConfig(
hidden_size=16,
num_hidden_layers=1,
num_attention_heads=2,
num_key_value_heads=2,
head_dim=8,
intermediate_size=32,
vocab_size=64,
latent_patch_size=2,
latent_channel=16,
position_embedding_type="3d_rope",
enable_fps_modulation=False,
action_gen=True,
action_dim=64,
max_action_dim=64,
num_embodiment_domains=32,
sound_gen=True,
sound_dim=64,
)
cfg = Cosmos3VideoConfig(arch_config=arch)
return Cosmos3VFMTransformer(cfg, hf_config={}).to(torch.float32)
def test_fastvideo_cosmos3_dit_module_tree_param_names() -> None:
"""Asserts the FastVideo Cosmos3 DiT module tree produces the param
names that the checkpoint converter is expected to write.
"""The native DiT param-name set must equal the published checkpoint surface."""
model = _build_tiny_dit()
names = {name for name, _ in model.named_parameters()}
Once ``fastvideo.models.dits.cosmos3.Cosmos3VFMTransformer`` is
constructible with a tiny config (see ``conftest._tiny_cosmos3_config``),
this test instantiates a 1-layer model and verifies that the
parameter-name set contains the expected UND/GEN/adapter prefixes:
* ``language_model.layers.0.self_attn.q_proj.weight``
* ``gen_layers.0.cross_attention.q_proj.weight``
* ``norm_moe_gen.weight``
* ``vae2llm.*`` and ``llm2vae.*``
"""
try:
from fastvideo.models.dits.cosmos3 import Cosmos3VFMTransformer # type: ignore
except ImportError:
pytest.skip("FastVideo Cosmos3 not yet implemented (Phase 2b)")
expected_top = {
"embed_tokens.weight",
"norm.weight",
"norm_moe_gen.weight",
"lm_head.weight",
"proj_in.weight",
"proj_in.bias",
"proj_out.weight",
"proj_out.bias",
"time_embedder.linear_1.weight",
"time_embedder.linear_1.bias",
"time_embedder.linear_2.weight",
"time_embedder.linear_2.bias",
}
expected_dormant = {
"action_proj_in.fc.weight",
"action_proj_in.bias.weight",
"action_proj_out.fc.weight",
"action_proj_out.bias.weight",
"action_modality_embed",
"audio_proj_in.weight",
"audio_proj_in.bias",
"audio_proj_out.weight",
"audio_proj_out.bias",
"audio_modality_embed",
}
layer_suffixes = {
"self_attn.to_q.weight",
"self_attn.to_k.weight",
"self_attn.to_v.weight",
"self_attn.to_out.weight",
"self_attn.add_q_proj.weight",
"self_attn.add_k_proj.weight",
"self_attn.add_v_proj.weight",
"self_attn.to_add_out.weight",
"self_attn.norm_q.weight",
"self_attn.norm_k.weight",
"self_attn.norm_added_q.weight",
"self_attn.norm_added_k.weight",
"mlp.gate_proj.weight",
"mlp.up_proj.weight",
"mlp.down_proj.weight",
"mlp_moe_gen.gate_proj.weight",
"mlp_moe_gen.up_proj.weight",
"mlp_moe_gen.down_proj.weight",
"input_layernorm.weight",
"input_layernorm_moe_gen.weight",
"post_attention_layernorm.weight",
"post_attention_layernorm_moe_gen.weight",
}
expected_layer = {f"layers.0.{s}" for s in layer_suffixes}
expected = expected_top | expected_dormant | expected_layer
from types import SimpleNamespace
import torch
assert names == expected, (f"Cosmos3 DiT param surface mismatch.\n"
f" missing: {sorted(expected - names)}\n"
f" unexpected: {sorted(names - expected)}")
from .conftest import _tiny_cosmos3_config
config = _tiny_cosmos3_config(num_hidden_layers=1)
model = Cosmos3VFMTransformer(SimpleNamespace(tf_model_config=config, dtype=torch.float32))
names = set(name for name, _ in model.named_parameters())
required_prefixes = [
"language_model.layers.0.self_attn.q_proj",
"gen_layers.0.cross_attention.q_proj",
"norm_moe_gen",
def test_fastvideo_cosmos3_dit_no_dead_vllm_omni_layout() -> None:
"""The dead vllm-omni layout (split language_model/gen_layers/cross_attention,
vae2llm/llm2vae) must NOT appear in the native DiT param tree."""
model = _build_tiny_dit()
names = {name for name, _ in model.named_parameters()}
dead_fragments = (
"language_model.",
"gen_layers.",
"cross_attention.",
"vae2llm",
"llm2vae",
]
for prefix in required_prefixes:
assert any(name.startswith(prefix) for name in names), (
f"Cosmos3 DiT module tree missing expected prefix {prefix!r}; "
f"converter at scripts/checkpoint_conversion/cosmos3_convert.py "
f"will fail to load weights for this subtree."
)
)
offenders = sorted(n for n in names if any(frag in n for frag in dead_fragments))
assert not offenders, f"native DiT still exposes dead vllm-omni param names: {offenders}"
def test_fastvideo_cosmos3_dit_per_layer_counts() -> None:
"""Per-layer block count: 22 weights/layer; full 1-layer model = 44 params.
(12 attention + 6 dual-MLP + 4 layernorm per layer; + 12 top-level
+ 10 dormant-head params.)
"""
model = _build_tiny_dit()
names = [name for name, _ in model.named_parameters()]
per_layer = [n for n in names if n.startswith("layers.0.")]
assert len(per_layer) == 22, f"expected 22 per-layer params, got {len(per_layer)}"
assert len(names) == 44, f"expected 44 total params for tiny 1-layer DiT, got {len(names)}"