diff --git a/fastvideo/models/dits/cosmos3.py b/fastvideo/models/dits/cosmos3.py index 391e4ac0b..c9166d829 100644 --- a/fastvideo/models/dits/cosmos3.py +++ b/fastvideo/models/dits/cosmos3.py @@ -322,13 +322,15 @@ class Cosmos3TimestepEmbedder(nn.Module): self.linear_2 = ReplicatedLinear(hidden_size, hidden_size, bias=True) def forward(self, t: torch.Tensor) -> torch.Tensor: - t_freq = timestep_embedding(t, self.frequency_embedding_size) - # Sinusoidal runs in fp32 (timestep_scale makes inputs tiny); cast to the - # MLP weight dtype (bf16 at inference; no-op in the fp32 parity tests). - t_freq = t_freq.to(self.linear_1.weight.dtype) - h, _ = self.linear_1(t_freq) - h = self.act(h) - out, _ = self.linear_2(h) + # The official Cosmos3 pipeline keeps the timestep MLP in fp32 even + # when the rest of the transformer runs in bf16. + with torch.autocast("cuda", enabled=t.is_cuda, dtype=torch.float32): + t_freq = timestep_embedding(t, self.frequency_embedding_size) + if not t.is_cuda: + t_freq = t_freq.to(self.linear_1.weight.dtype) + h, _ = self.linear_1(t_freq) + h = self.act(h) + out, _ = self.linear_2(h) return out diff --git a/tests/local_tests/cosmos3/test_cosmos3_dit_parity.py b/tests/local_tests/cosmos3/test_cosmos3_dit_parity.py index 835212493..1e4eb7bbd 100644 --- a/tests/local_tests/cosmos3/test_cosmos3_dit_parity.py +++ b/tests/local_tests/cosmos3/test_cosmos3_dit_parity.py @@ -254,3 +254,25 @@ class TestCosmos3DiTParity: """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) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Cosmos3 timestep precision parity requires CUDA autocast") +def test_timestep_embedder_bf16_matches_official_fp32_policy() -> None: + from cosmos_framework.model.generator.mot.modeling_utils import TimestepEmbedder + + from fastvideo.models.dits.cosmos3 import Cosmos3TimestepEmbedder + + official = TimestepEmbedder(hidden_size=32).cuda().bfloat16().eval() + fastvideo = Cosmos3TimestepEmbedder(hidden_size=32).cuda().bfloat16().eval() + with torch.no_grad(): + fastvideo.linear_1.weight.copy_(official.mlp[0].weight) + fastvideo.linear_1.bias.copy_(official.mlp[0].bias) + fastvideo.linear_2.weight.copy_(official.mlp[2].weight) + fastvideo.linear_2.bias.copy_(official.mlp[2].bias) + + timesteps = torch.linspace(0, 1, 7, device="cuda", dtype=torch.float32) + with torch.autocast("cuda", enabled=True, dtype=torch.float32): + expected = official(timesteps) + actual = fastvideo(timesteps) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0)