[fix]: match Cosmos3 timestep precision

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
 
 

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)