diff --git a/CMakeLists.txt b/CMakeLists.txt index 7fd244c..33f7861 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -76,6 +76,11 @@ if(EXISTS "${KIMODO_GGML_SOURCE_DIR}/CMakeLists.txt") target_include_directories(kimodo-multiprompt-fixture-parity PRIVATE src) target_link_libraries(kimodo-multiprompt-fixture-parity PRIVATE ggml ggml-vulkan) target_compile_definitions(kimodo-multiprompt-fixture-parity PRIVATE KIMODO_HAVE_GGML=1 KIMODO_HAVE_GGML_VULKAN=1) + add_executable(kimodo-stage-fixture-parity tests/stage_fixture_parity.cpp) + target_sources(kimodo-stage-fixture-parity PRIVATE src/gguf.cpp src/ggml_weights.cpp src/denoiser.cpp src/motion_rep.cpp src/diffusion.cpp) + target_include_directories(kimodo-stage-fixture-parity PRIVATE src) + target_link_libraries(kimodo-stage-fixture-parity PRIVATE ggml ggml-vulkan) + target_compile_definitions(kimodo-stage-fixture-parity PRIVATE KIMODO_HAVE_GGML=1 KIMODO_HAVE_GGML_VULKAN=1) add_executable(kimodo-llm-text-session-parity tests/llm_text_session_parity.cpp src/llm_text_encoder.cpp src/llm_tokenizer.cpp) target_include_directories(kimodo-llm-text-session-parity PRIVATE src) diff --git a/reference/dump_kimodo_multiprompt_reference.py b/reference/dump_kimodo_multiprompt_reference.py index 1a99926..f2c5351 100644 --- a/reference/dump_kimodo_multiprompt_reference.py +++ b/reference/dump_kimodo_multiprompt_reference.py @@ -85,12 +85,29 @@ def main() -> None: sys.path.insert(0, str(upstream)) from kimodo import load_model # pylint: disable=import-outside-toplevel + # The NVIDIA PyTorch container enables TF32 globally. Kimodo GGML uses + # F32 accumulation for reference parity, so a CUDA capture must disable + # Tensor Core TF32 before any model module is materialised. + if args.device.startswith("cuda"): + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + torch.set_float32_matmul_precision("highest") + # TransformerEncoder otherwise selects PyTorch's CUDA fast path, + # whose fused attention reductions have a different F32 accumulation + # order from Kimodo's explicit GGML attention graph. + torch.backends.mha.set_fastpath_enabled(False) + torch.backends.cuda.enable_flash_sdp(False) + torch.backends.cuda.enable_mem_efficient_sdp(False) + torch.backends.cuda.enable_math_sdp(True) + torch.manual_seed(args.seed) encoder = CapturedEmbeddings(args.prompt, args.embedding) model, resolved = load_model("kimodo-smplx-rp", device=args.device, text_encoder=encoder, return_resolved_name=True) calls: list[dict[str, list[torch.Tensor] | torch.Tensor]] = [] inverse_inputs: list[torch.Tensor] = [] + root_calls: list[tuple[tuple[torch.Tensor, ...], torch.Tensor]] = [] + body_calls: list[tuple[tuple[torch.Tensor, ...], torch.Tensor]] = [] original_step = model.denoising_step def capture_step(*values: Any, **kwargs: Any) -> torch.Tensor: @@ -110,6 +127,18 @@ def main() -> None: call["output"].append(result.detach().clone()) # type: ignore[index] return result + def capture_stage(calls_for_stage: list[tuple[tuple[torch.Tensor, ...], torch.Tensor]]): + def hook(_module: torch.nn.Module, values: tuple[Any, ...], output: torch.Tensor) -> None: + if not calls_for_stage: + calls_for_stage.append(( + tuple(value.detach().clone() for value in values if isinstance(value, torch.Tensor)), + output.detach().clone(), + )) + return hook + + root_hook = model.denoiser.model.root_model.register_forward_hook(capture_stage(root_calls)) + body_hook = model.denoiser.model.body_model.register_forward_hook(capture_stage(body_calls)) + original_inverse = model.motion_rep.inverse def capture_inverse(motion: torch.Tensor, *values: Any, **kwargs: Any): @@ -129,10 +158,17 @@ def main() -> None: finally: model.denoising_step = original_step model.motion_rep.inverse = original_inverse - if len(calls) != len(args.prompt) or not inverse_inputs: + root_hook.remove() + body_hook.remove() + if len(calls) != len(args.prompt) or not inverse_inputs or not root_calls or not body_calls: raise RuntimeError(f"expected {len(args.prompt)} segment trajectories, got {len(calls)}") arrays: dict[str, np.ndarray] = {"stitched_motion_rep": as_f32(inverse_inputs[-1])} + for stage, stage_calls in (("root", root_calls), ("body", body_calls)): + values, stage_output = stage_calls[0] + arrays[f"{stage}_output"] = as_f32(stage_output) + for index, value in enumerate(values): + arrays[f"{stage}_input_{index}"] = as_f32(value) for index, call in enumerate(calls): prefix = f"segment_{index:02d}_" for name in ("pad_mask", "text_features", "text_pad_mask", "first_heading_angle", "motion_mask", "observed_motion"): @@ -158,6 +194,8 @@ def main() -> None: "transition_frames": args.transition_frames, "diffusion_steps": args.steps, "seed": args.seed, "cfg_type": "separated", "cfg_weight": [2.0, 2.0], "post_processing": False, "device": args.device, "torch": torch.__version__, + "cuda_tf32": torch.backends.cuda.matmul.allow_tf32 if args.device.startswith("cuda") else None, + "cuda_mha_fastpath": torch.backends.mha.get_fastpath_enabled() if args.device.startswith("cuda") else None, "python": platform.python_version(), "embedding_sources": [ {"path": str(path), "sha256": sha256(path), "note": note} for path, note in zip(args.embedding, notes) diff --git a/tests/stage_fixture_parity.cpp b/tests/stage_fixture_parity.cpp new file mode 100644 index 0000000..9f47f81 --- /dev/null +++ b/tests/stage_fixture_parity.cpp @@ -0,0 +1,36 @@ +#include "denoiser.hpp" +#include "ggml_weights.hpp" +#include +#include +#include +#include +#include +#include +#include + +static std::vector read(const std::string &path) { + std::ifstream input(path, std::ios::binary | std::ios::ate); + if (!input || input.tellg() < 0 || input.tellg() % 4) throw std::runtime_error("bad fixture: " + path); + std::vector value(static_cast(input.tellg()) / 4); input.seekg(0); + input.read(reinterpret_cast(value.data()), static_cast(value.size() * 4)); + if (!input) throw std::runtime_error("short fixture"); return value; +} +int main(int argc, char **argv) try { + if (argc != 5) return 2; + const std::string stage(argv[3]), directory = std::string(argv[2]) + "/"; + const auto frames = static_cast(std::stoul(argv[4])); + if (stage != "root" && stage != "body") throw std::runtime_error("stage must be root or body"); + auto weights = kimodo::detail::ggml_motion_weights::load(argv[1]); if (!weights) throw std::runtime_error(weights.error()); + const size_t dimension = stage == "root" ? 546 : 545; + auto output = kimodo::detail::run_motion_transformer( + **weights, stage + "_model.", read(directory + stage + "_input_0.f32"), dimension, + read(directory + stage + "_input_2.f32"), read(directory + stage + "_input_4.f32"), + read(directory + stage + "_input_5.f32"), 3, frames); + if (!output) throw std::runtime_error(output.error()); + const auto expected = read(directory + stage + "_output.f32"); + if (output->size() != expected.size()) throw std::runtime_error("shape mismatch"); + float maximum = 0; double error = 0, reference = 0; + for (size_t i = 0; i < expected.size(); ++i) { const float d = (*output)[i] - expected[i]; maximum = std::max(maximum, std::abs(d)); error += double(d)*d; reference += double(expected[i])*expected[i]; } + std::printf("%s max_abs=%g rel_l2=%g\n", stage.c_str(), maximum, std::sqrt(error/reference)); + return maximum < 2.e-3f ? 0 : 1; +} catch (const std::exception &error) { std::fprintf(stderr, "%s\n", error.what()); return 1; }