File size: 3,878 Bytes
4e7b91c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """r0b_probe.py — ANIMA R0b: the fork-integration gate (runs INSIDE the pod
clone of diffusion-pipe @ feat/aleph-adapter). Certifies, without a dataset:
G1 anima weights load through CosmosPredict2Pipeline (dit_config printed);
G2 aleph relays attach post-materialization (site count, param count);
G3 DTYPE LAW: relay dtype == block dtype (bf16);
G4 param groups: aleph_relay bucket present, trunk freezable by lr=0;
G5 zero-init: all-off/enabled-at-init forward delta == 0 through one block.
Full exp004 training goes to the evening slot per the day-3 plan.
Run: cd /workspace/geolip2/diffusion-pipe && python3 ../pod2/r0b_probe.py
"""
from __future__ import annotations
import sys
import torch
sys.path.insert(0, ".")
import utils.common # bind the fork utils BEFORE ComfyUI (regular-pkg shadowing)
sys.path.append("submodules/ComfyUI")
CFG = {
"model": {
"type": "anima",
"transformer_path":
"/workspace/models/anima/split_files/diffusion_models/"
"anima-base-v1.0.safetensors",
"llm_path": "/workspace/models/anima/split_files/text_encoders/"
"qwen_3_06b_base.safetensors",
"vae_path": "/workspace/models/anima/split_files/vae/"
"qwen_image_vae.safetensors",
"dtype": torch.bfloat16,
"aleph_relay": True,
"aleph_relay_every": 1,
"aleph_relay_lr": 1e-3,
"self_attn_lr": 0, "cross_attn_lr": 0, "mlp_lr": 0, "mod_lr": 0,
"llm_adapter_lr": 0,
},
"optimizer": {"lr": 0},
"reentrant_activation_checkpointing": False,
}
def main():
import os
for k, v in (("MASTER_ADDR", "127.0.0.1"), ("MASTER_PORT", "29571"),
("RANK", "0"), ("WORLD_SIZE", "1"), ("LOCAL_RANK", "0")):
os.environ.setdefault(k, v)
import deepspeed
deepspeed.init_distributed()
from models.cosmos_predict2 import CosmosPredict2Pipeline
pipe = CosmosPredict2Pipeline(CFG)
print("[G1] pipeline constructed (text encoder loaded, name="
f"{pipe.name})", flush=True)
pipe.load_diffusion_model()
tr = pipe.transformer
relays = [getattr(b, "aleph_relay", None) for b in tr.blocks]
n_sites = sum(r is not None for r in relays)
n_params = sum(p.numel() for r in relays if r is not None
for p in r.parameters())
print(f"[G2] {n_sites}/{len(tr.blocks)} blocks carry relays, "
f"{n_params:,} adapter params", flush=True)
assert n_sites == len(tr.blocks) > 0
r0 = next(r for r in relays if r is not None)
b0 = tr.blocks[0]
bdt = next(p for n, p in b0.named_parameters()
if "aleph_relay" not in n).dtype
rdt = next(r0.parameters()).dtype
print(f"[G3] block dtype {bdt} | relay dtype {rdt}", flush=True)
assert rdt == bdt == torch.bfloat16, "dtype law violated"
params = [p for p in tr.parameters()]
for p, (n, _) in zip(params, tr.named_parameters()):
pass
groups = pipe.get_param_groups(
[p for p in tr.parameters() if hasattr(p, "original_name")])
n_trainable = sum(p.numel() for g in groups for p in g["params"])
print(f"[G4] {len(groups)} param groups, trainable {n_trainable:,} "
f"(should == adapter count {n_params:,})", flush=True)
assert n_trainable == n_params, "freeze-by-lr-0 leaked trunk params"
r0.assert_zero_init()
x = torch.randn(1, 2, 4, 4, tr.model_channels, dtype=bdt)
with torch.no_grad():
out_on = r0(x)
r0.enabled = False
out_off = r0(x)
r0.enabled = True
assert torch.equal(out_on, x) and out_off is x, "zero-init/toggle broken"
print("[G5] zero-init + toggle exact on the DiT relay", flush=True)
print("R0B ALL GATES GREEN — fork integration certified; exp004 train "
"is evening-slot ready", flush=True)
if __name__ == "__main__":
main()
|