From b8270b080f76db284664eb6281d821730734bb17 Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Sun, 5 Jul 2026 13:12:28 -0700 Subject: [PATCH] [test]: randomize the zero-init head so DreamX AR tiny parity is not vacuous --- .../test_dreamx_world_ar_transformer_parity.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py b/tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py index ff0eb34c9..70c011733 100644 --- a/tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py +++ b/tests/local_tests/dreamx_world/test_dreamx_world_ar_transformer_parity.py @@ -124,6 +124,14 @@ def _make_inputs(): def test_dreamx_world_ar_tiny_forward_matches_official(): official = _load_official_tiny() + # The official init_weights zero-inits the output head (head.head.weight and + # biases), so both models would output exactly zero and the comparison would + # pass vacuously. Randomize the head (deterministically) before copying the + # state dict so the outputs reflect the internal computation. + generator = torch.Generator().manual_seed(7) + with torch.no_grad(): + official.head.head.weight.normal_(std=0.5, generator=generator) + official.head.head.bias.normal_(std=0.5, generator=generator) fastvideo = DreamXWorldARTransformer3DModel(_tiny_config(), {}).eval() fastvideo.load_state_dict(official.state_dict(), strict=True) @@ -133,6 +141,7 @@ def test_dreamx_world_ar_tiny_forward_matches_official(): x, t, context, camera, kv_cache, cross_cache = _make_inputs() fastvideo_out = fastvideo(x=x, t=t, context=context, seq_len=4, y_camera=camera, kv_cache=kv_cache, crossattn_cache=cross_cache).detach() + assert official_out.abs().max() > 0, "official output is all-zero; parity comparison is vacuous" assert_close(fastvideo_out, official_out, atol=1e-5, rtol=1e-5)