update unit test

Author: Huy Vu2

tests/unit_tests/automodel/test_flow_matching_pipeline.py

@@ -333,12 +333,14 @@ def test_loss_weighting_enabled(self, simple_adapter):
         model_pred = torch.randn(2, 16, 4, 8, 8)
         target = torch.randn_like(model_pred)
         sigma = torch.tensor([0.3, 0.7])
+        batch = {}
 
-        weighted_loss, unweighted_loss, loss_weight = pipeline.compute_loss(model_pred, target, sigma)
+        # Returns: weighted_loss, average_weighted_loss, unweighted_loss, average_unweighted_loss, loss_weight, loss_mask
+        _, scalar_weighted_loss, _, scalar_unweighted_loss, loss_weight, _ = pipeline.compute_loss(model_pred, target, sigma, batch)
 
         # Verify shapes
-        assert weighted_loss.ndim == 0, "Weighted loss should be scalar"
-        assert unweighted_loss.ndim == 0, "Unweighted loss should be scalar"
+        assert scalar_weighted_loss.ndim == 0, "Weighted loss should be scalar"
+        assert scalar_unweighted_loss.ndim == 0, "Unweighted loss should be scalar"
 
         # Verify weight formula: w = 1 + shift * σ
         expected_weights = 1.0 + 3.0 * sigma
@@ -357,11 +359,12 @@ def test_loss_weighting_disabled(self, simple_adapter):
         model_pred = torch.randn(2, 16, 4, 8, 8)
         target = torch.randn_like(model_pred)
         sigma = torch.tensor([0.3, 0.7])
+        batch = {}
 
-        weighted_loss, unweighted_loss, loss_weight = pipeline.compute_loss(model_pred, target, sigma)
+        _, scalar_weighted_loss, _, scalar_unweighted_loss, loss_weight, _ = pipeline.compute_loss(model_pred, target, sigma, batch)
 
         # Without weighting, weighted loss should equal unweighted loss
-        assert torch.allclose(weighted_loss, unweighted_loss, atol=1e-6)
+        assert torch.allclose(scalar_weighted_loss, scalar_unweighted_loss, atol=1e-6)
 
         # All weights should be 1
         assert torch.allclose(loss_weight, torch.ones_like(loss_weight))
@@ -377,10 +380,11 @@ def test_loss_weight_formula(self, simple_adapter):
 
         model_pred = torch.zeros(4, 16, 4, 8, 8)
         target = torch.ones_like(model_pred)
+        batch = {}
 
         for sigma_val in [0.0, 0.25, 0.5, 0.75, 1.0]:
             sigma = torch.full((4,), sigma_val)
-            _, _, loss_weight = pipeline.compute_loss(model_pred, target, sigma)
+            _, _, _, _, loss_weight, _ = pipeline.compute_loss(model_pred, target, sigma, batch)
 
             expected_weight = 1.0 + flow_shift * sigma_val
             actual_weight = loss_weight[0, 0, 0, 0, 0].item()
@@ -396,11 +400,12 @@ def test_loss_is_non_negative(self, simple_adapter):
         model_pred = torch.randn(2, 16, 4, 8, 8)
         target = torch.randn_like(model_pred)
         sigma = torch.rand(2)
+        batch = {}
 
-        weighted_loss, unweighted_loss, _ = pipeline.compute_loss(model_pred, target, sigma)
+        _, scalar_weighted_loss, _, scalar_unweighted_loss, _, _ = pipeline.compute_loss(model_pred, target, sigma, batch)
 
-        assert weighted_loss >= 0, "Weighted loss should be non-negative"
-        assert unweighted_loss >= 0, "Unweighted loss should be non-negative"
+        assert scalar_weighted_loss >= 0, "Weighted loss should be non-negative"
+        assert scalar_unweighted_loss >= 0, "Unweighted loss should be non-negative"
 
     def test_loss_is_finite(self, simple_adapter):
         """Test that computed loss is finite."""
@@ -409,11 +414,12 @@ def test_loss_is_finite(self, simple_adapter):
         model_pred = torch.randn(2, 16, 4, 8, 8)
         target = torch.randn_like(model_pred)
         sigma = torch.rand(2)
+        batch = {}
 
-        weighted_loss, unweighted_loss, _ = pipeline.compute_loss(model_pred, target, sigma)
+        _, scalar_weighted_loss, _, scalar_unweighted_loss, _, _ = pipeline.compute_loss(model_pred, target, sigma, batch)
 
-        assert torch.isfinite(weighted_loss), "Weighted loss should be finite"
-        assert torch.isfinite(unweighted_loss), "Unweighted loss should be finite"
+        assert torch.isfinite(scalar_weighted_loss), "Weighted loss should be finite"
+        assert torch.isfinite(scalar_unweighted_loss), "Unweighted loss should be finite"
 
     def test_loss_mse_correctness(self, simple_adapter):
         """Test that base loss is MSE."""
@@ -425,13 +431,14 @@ def test_loss_mse_correctness(self, simple_adapter):
         model_pred = torch.randn(2, 16, 4, 8, 8)
         target = torch.randn_like(model_pred)
         sigma = torch.rand(2)
+        batch = {}
 
-        _, unweighted_loss, _ = pipeline.compute_loss(model_pred, target, sigma)
+        _, _, _, scalar_unweighted_loss, _, _ = pipeline.compute_loss(model_pred, target, sigma, batch)
 
         # Manual MSE calculation
         expected_mse = nn.functional.mse_loss(model_pred.float(), target.float())
 
-        assert torch.allclose(unweighted_loss, expected_mse, atol=1e-6)
+        assert torch.allclose(scalar_unweighted_loss, expected_mse, atol=1e-6)
 
 
 class TestFullTrainingStep:
@@ -442,7 +449,8 @@ def test_basic_training_step(self, pipeline, mock_model, sample_batch):
         device = torch.device("cpu")
         dtype = torch.bfloat16
 
-        loss, metrics = pipeline.step(mock_model, sample_batch, device, dtype, global_step=0)
+        # Returns: weighted_loss, average_weighted_loss, loss_mask, metrics
+        _, loss, _, metrics = pipeline.step(mock_model, sample_batch, device, dtype, global_step=0)
 
         # Verify loss
         assert isinstance(loss, torch.Tensor), "Loss should be a tensor"
@@ -478,7 +486,7 @@ def test_step_with_different_batch_sizes(self, simple_adapter, mock_model):
                 "text_embeddings": torch.randn(batch_size, 77, 4096),
             }
 
-            loss, metrics = pipeline.step(mock_model, batch, device, dtype, global_step=0)
+            _, loss, _, metrics = pipeline.step(mock_model, batch, device, dtype, global_step=0)
 
             assert isinstance(loss, torch.Tensor), f"Loss should be tensor for batch_size={batch_size}"
             assert not torch.isnan(loss), f"Loss should not be NaN for batch_size={batch_size}"
@@ -493,7 +501,7 @@ def test_step_with_4d_video_latents(self, pipeline, mock_model):
             "text_embeddings": torch.randn(77, 4096),  # 2D instead of 3D
         }
 
-        loss, metrics = pipeline.step(mock_model, batch, device, dtype, global_step=0)
+        _, loss, _, metrics = pipeline.step(mock_model, batch, device, dtype, global_step=0)
 
         assert isinstance(loss, torch.Tensor)
         assert not torch.isnan(loss)
@@ -503,7 +511,7 @@ def test_step_metrics_collection(self, pipeline, mock_model, sample_batch):
         device = torch.device("cpu")
         dtype = torch.bfloat16
 
-        loss, metrics = pipeline.step(mock_model, sample_batch, device, dtype, global_step=100)
+        _, loss, _, metrics = pipeline.step(mock_model, sample_batch, device, dtype, global_step=100)
 
         expected_keys = [
             "loss",
@@ -530,7 +538,7 @@ def test_step_sigma_in_valid_range(self, pipeline, mock_model, sample_batch):
         device = torch.device("cpu")
         dtype = torch.bfloat16
 
-        loss, metrics = pipeline.step(mock_model, sample_batch, device, dtype, global_step=0)
+        _, loss, _, metrics = pipeline.step(mock_model, sample_batch, device, dtype, global_step=0)
 
         assert 0.0 <= metrics["sigma_min"] <= 1.0, "Sigma min should be in [0, 1]"
         assert 0.0 <= metrics["sigma_max"] <= 1.0, "Sigma max should be in [0, 1]"
@@ -548,7 +556,7 @@ def test_step_timesteps_in_valid_range(self, simple_adapter, mock_model, sample_
         device = torch.device("cpu")
         dtype = torch.bfloat16
 
-        loss, metrics = pipeline.step(mock_model, sample_batch, device, dtype, global_step=0)
+        _, loss, _, metrics = pipeline.step(mock_model, sample_batch, device, dtype, global_step=0)
 
         assert 0.0 <= metrics["timestep_min"] <= num_timesteps
         assert 0.0 <= metrics["timestep_max"] <= num_timesteps
@@ -558,7 +566,7 @@ def test_step_noisy_latents_are_finite(self, pipeline, mock_model, sample_batch)
         device = torch.device("cpu")
         dtype = torch.bfloat16
 
-        loss, metrics = pipeline.step(mock_model, sample_batch, device, dtype, global_step=0)
+        _, loss, _, metrics = pipeline.step(mock_model, sample_batch, device, dtype, global_step=0)
 
         assert torch.isfinite(torch.tensor(metrics["noisy_min"])), "Noisy min should be finite"
         assert torch.isfinite(torch.tensor(metrics["noisy_max"])), "Noisy max should be finite"
@@ -568,7 +576,7 @@ def test_step_with_image_batch(self, pipeline, mock_model, image_batch):
         device = torch.device("cpu")
         dtype = torch.bfloat16
 
-        loss, metrics = pipeline.step(mock_model, image_batch, device, dtype, global_step=0)
+        _, loss, _, metrics = pipeline.step(mock_model, image_batch, device, dtype, global_step=0)
 
         assert isinstance(loss, torch.Tensor)
         assert not torch.isnan(loss)
@@ -588,11 +596,11 @@ def test_deterministic_with_seed(self, simple_adapter, mock_model, sample_batch)
 
         # First run
         torch.manual_seed(42)
-        _, metrics1 = pipeline.step(mock_model, sample_batch, device, dtype, global_step=0)
+        _, _, _, metrics1 = pipeline.step(mock_model, sample_batch, device, dtype, global_step=0)
 
         # Second run with same seed
         torch.manual_seed(42)
-        _, metrics2 = pipeline.step(mock_model, sample_batch, device, dtype, global_step=0)
+        _, _, _, metrics2 = pipeline.step(mock_model, sample_batch, device, dtype, global_step=0)
 
         # Sigma values should be identical
         assert abs(metrics1["sigma_min"] - metrics2["sigma_min"]) < 1e-6
@@ -675,7 +683,7 @@ def test_empty_batch_handling(self, simple_adapter):
         }
 
         mock_model = MockModel()
-        loss, metrics = pipeline.step(mock_model, batch, torch.device("cpu"), torch.float32, global_step=0)
+        _, loss, _, metrics = pipeline.step(mock_model, batch, torch.device("cpu"), torch.float32, global_step=0)
 
         assert not torch.isnan(loss)
 
@@ -693,7 +701,7 @@ def test_large_batch_handling(self, simple_adapter):
         }
 
         mock_model = MockModel()
-        loss, metrics = pipeline.step(mock_model, batch, torch.device("cpu"), torch.float32, global_step=0)
+        _, loss, _, metrics = pipeline.step(mock_model, batch, torch.device("cpu"), torch.float32, global_step=0)
 
         assert not torch.isnan(loss)
 
@@ -716,7 +724,7 @@ def test_extreme_flow_shift_values(self, simple_adapter):
             }
 
             mock_model = MockModel()
-            loss, metrics = pipeline.step(mock_model, batch, torch.device("cpu"), torch.float32, global_step=0)
+            _, loss, _, metrics = pipeline.step(mock_model, batch, torch.device("cpu"), torch.float32, global_step=0)
 
             assert torch.isfinite(loss), f"Loss should be finite for shift={shift}"
 
@@ -755,7 +763,7 @@ def test_multiple_training_steps(self, simple_adapter):
                 "text_embeddings": torch.randn(2, 77, 4096),
             }
 
-            loss, metrics = pipeline.step(mock_model, batch, torch.device("cpu"), torch.float32, global_step=step)
+            _, loss, _, metrics = pipeline.step(mock_model, batch, torch.device("cpu"), torch.float32, global_step=step)
             losses.append(loss.item())
 
             assert not torch.isnan(loss), f"Loss became NaN at step {step}"
@@ -779,7 +787,7 @@ def test_pipeline_with_all_sampling_methods(self, simple_adapter):
                 "text_embeddings": torch.randn(2, 77, 4096),
             }
 
-            loss, metrics = pipeline.step(mock_model, batch, torch.device("cpu"), torch.float32, global_step=0)
+            _, loss, _, metrics = pipeline.step(mock_model, batch, torch.device("cpu"), torch.float32, global_step=0)
 
             assert not torch.isnan(loss), f"Loss should not be NaN for method={method}"