diag-nan.py 3.56 KB
"""定位 NaN 的产生位置:对 dit / text_encoder 挂 forward hook,
找出第一个『输入有限、输出非有限』的模块。"""
import os, sys, time, torch
sys.argv = [sys.argv[0]]
BASE = "/data/models/minimax-h3"
os.environ.setdefault("DIFFSYNTH_ATTENTION_IMPLEMENTATION", "torch")
os.environ.setdefault("DIFFSYNTH_MODEL_BASE_PATH", f"{BASE}/models")
os.environ.setdefault("MODELSCOPE_CACHE", f"{BASE}/.mscache")

from diffsynth.pipelines.minimax_h3_audio_video import MiniMaxH3Pipeline, ModelConfig

DTYPE = torch.float16
free_gb = torch.cuda.mem_get_info("cuda")[0] / 1024**3
vram_limit = max(6.0, free_gb - 2.0)
print(f"[env] free={free_gb:.1f}G vram_limit={vram_limit:.1f}G", flush=True)

vc = dict(offload_dtype=DTYPE, offload_device="cpu", onload_dtype=DTYPE, onload_device="cpu",
          preparing_dtype=DTYPE, preparing_device="cuda",
          computation_dtype=DTYPE, computation_device="cuda")
NF4 = "DiffSynth-Studio/MiniMax-H3-NF4"
t0 = time.time()
pipe = MiniMaxH3Pipeline.from_pretrained(
    torch_dtype=DTYPE, device="cuda",
    model_configs=[
        ModelConfig(model_id=NF4, origin_file_pattern="minimax-h3-fl2va-nf4.safetensors", **vc),
        ModelConfig(model_id=NF4, origin_file_pattern="minimax-h3-text-encoder-nf4.safetensors", **vc),
        ModelConfig(model_id=NF4, origin_file_pattern="video_vae_nf4.safetensors", **vc),
        ModelConfig(model_id=NF4, origin_file_pattern="audio_vae_nf4.safetensors", **vc),
    ],
    processor_config=ModelConfig(model_id="MiniMax/MiniMax-H3", origin_file_pattern="FL2VA/processor/"),
    vram_limit=vram_limit)
print(f"[load] {time.time()-t0:.0f}s", flush=True)

pipe.load_lora(pipe.dit, ModelConfig(model_id="lightx2v/Minimax-h3-Turbo",
    origin_file_pattern="minimax_h3_fl2v_turbo_4step_v1.0_768p_bf16.safetensors"))

def finite(x):
    if torch.is_tensor(x):
        return bool(torch.isfinite(x).all()), (float(x.abs().max()) if x.numel() and x.is_floating_point() else None)
    return True, None

hits = []      # 第一个坏掉的模块
seen = set()
def mk(name):
    def hook(mod, inp, out):
        if name in seen: return
        ok_in = all(finite(t)[0] for t in inp if torch.is_tensor(t))
        outs = out if isinstance(out, (tuple, list)) else (out,)
        for o in outs:
            ok, mx = finite(o)
            if not ok:
                seen.add(name)
                hits.append((name, type(mod).__name__, ok_in, mx))
                return
    return hook

handles = []
for tag, mod in (("dit", pipe.dit), ("txt", pipe.text_encoder)):
    for n, m in mod.named_modules():
        if n:  # 跳过根
            handles.append(m.register_forward_hook(mk(f"{tag}.{n}")))
print(f"[hook] {len(handles)} 个模块已挂钩", flush=True)

video, audio = pipe(prompt="一只橘猫在窗台上打哈欠,阳光洒进来",
                    height=480, width=832, num_frames=124,
                    num_inference_steps=4, seed=42, flow_shift=6.0)
for h in handles: h.remove()

print("\n===== 结果 =====", flush=True)
for t, nm in (("video", video), ("audio", audio)):
    if torch.is_tensor(nm):
        print(f"{t}: shape={tuple(nm.shape)} dtype={nm.dtype} finite={bool(torch.isfinite(nm).all())} "
              f"absmax={float(nm.abs().max()) if nm.is_floating_point() else 'n/a'}")
    else:
        print(f"{t}: type={type(nm)}")
print(f"\n第一批非有限输出的模块(共 {len(hits)},只列前 15):")
for n, cls, ok_in, mx in hits[:15]:
    print(f"  {n}  [{cls}]  输入有限={ok_in}  输出absmax={mx}")
if not hits:
    print("  无 —— dit/text_encoder 内部全程有限,问题在 VAE 解码或后处理")