Merge pull request #3086 from bghira/agent/audio-validation-gather

Gather batch-parallel validation payloads

Author: bghira

simpletuner/helpers/training/validation.py

@@ -18,6 +18,7 @@
 import numpy as np
 import torch
 import wandb
+from accelerate.utils import gather_object
 from tqdm import tqdm
 
 from simpletuner.helpers.caching.memory import reclaim_memory
@@ -76,7 +77,7 @@
 from simpletuner.helpers.utils.checkpoint_manager import CheckpointManager
 
 logger = logging.getLogger("Validation")
-from simpletuner.helpers.training.multi_process import gather_across_processes, should_log, split_across_processes
+from simpletuner.helpers.training.multi_process import should_log, split_across_processes
 
 if should_log():
     logger.setLevel(os.environ.get("SIMPLETUNER_LOG_LEVEL", "INFO"))
@@ -4191,13 +4192,15 @@ def process_prompts(
 
         if use_distributed:
             logger.info(f"[Rank {rank}] Gathering {len(local_payloads)} local payloads")
-            gathered_payloads = gather_across_processes(local_payloads)
+            gathered_payloads = list(gather_object(local_payloads) or [])
+            aggregated_payloads = []
+            for payload_group in gathered_payloads:
+                if isinstance(payload_group, list):
+                    aggregated_payloads.extend(payload_group)
+                else:
+                    aggregated_payloads.append(payload_group)
             if not self.accelerator.is_main_process:
                 return
-            logger.info(
-                f"[Rank {rank}] Gathered {len(gathered_payloads)} payload groups: {[len(p) for p in gathered_payloads]}"
-            )
-            aggregated_payloads = [payload for worker_payloads in gathered_payloads for payload in worker_payloads]
             logger.info(f"[Rank {rank}] Total aggregated payloads: {len(aggregated_payloads)}")
 
             aggregated_payloads.sort(key=lambda payload: payload["index"])

tests/test_validation_context_parallel.py

@@ -1,8 +1,11 @@
 import unittest
 from contextlib import contextmanager
 from types import SimpleNamespace
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
+import torch
+
+from simpletuner.helpers.models.common import AudioModelFoundation
 from simpletuner.helpers.training.validation import Validation, _ValidationWorkItem
 
 
@@ -30,6 +33,23 @@ def _work_items(count: int):
     ]
 
 
+class DummyAudioModel(AudioModelFoundation):
+    def validation_audio_sample_rate(self):
+        return 44100
+
+    def _encode_prompts(self, prompts, is_negative_prompt=False):
+        return {}
+
+    def convert_text_embed_for_pipeline(self, text_embedding):
+        return {}
+
+    def convert_negative_text_embed_for_pipeline(self, text_embedding):
+        return {}
+
+    def model_predict(self, prepared_batch):
+        return None
+
+
 class ValidationContextParallelTests(unittest.TestCase):
     @contextmanager
     def _patch_cp(self, *, data_rank: int, data_local_rank: int, data_parallel_size: int = 4):
@@ -85,6 +105,64 @@ def test_context_parallel_keeps_single_prompt_distributed(self):
         self.assertEqual(worker_count, 4)
         self.assertEqual([item.index for item in local_items], [0])
 
+    def test_batch_parallel_gathers_audio_payloads_from_peer_ranks(self):
+        validation = Validation.__new__(Validation)
+        validation.accelerator = SimpleNamespace(
+            num_processes=2,
+            process_index=0,
+            is_main_process=True,
+        )
+        validation.config = SimpleNamespace(validation_multigpu="batch-parallel")
+        validation.model = DummyAudioModel.__new__(DummyAudioModel)
+        validation.validation_prompt_metadata = {
+            "validation_prompts": ["prompt 0", "prompt 1"],
+            "validation_shortnames": ["song_0", "song_1"],
+        }
+        validation.validation_image_inputs = None
+        validation.validation_prompt_dict = None
+        validation.validation_resolutions = [(0, 0)]
+        validation.save_dir = "validation_images"
+        validation.eval_scores = {}
+        validation.validation_video_paths = {}
+        validation.evaluation_result = None
+        validation._check_abort = MagicMock()
+        validation._use_context_parallel_validation = MagicMock(return_value=False)
+        validation._split_validation_work_items = MagicMock(return_value=([_work_items(2)[0]], True, 2))
+        validation._should_publish_validation_payloads = MagicMock(return_value=True)
+
+        rank0_payload = {
+            "index": 0,
+            "shortname": "song_0",
+            "decorated_shortname": "song_0",
+            "prompt": "prompt 0",
+            "stitched": [],
+            "checkpoint": [],
+            "audio": Validation._serialise_media_list([torch.zeros(1, 4)]),
+        }
+        rank1_payload = {
+            "index": 1,
+            "shortname": "song_1",
+            "decorated_shortname": "song_1",
+            "prompt": "prompt 1",
+            "stitched": [],
+            "checkpoint": [],
+            "audio": Validation._serialise_media_list([torch.ones(1, 4)]),
+        }
+        validation._execute_validation_work_item = MagicMock(return_value=rank0_payload)
+
+        with (
+            patch("simpletuner.helpers.training.validation.gather_object", return_value=[[rank0_payload], [rank1_payload]]),
+            patch("simpletuner.helpers.training.validation.validation_audio.save_audio") as save_audio,
+            patch("simpletuner.helpers.training.validation.validation_audio.log_audio_to_webhook"),
+            patch("simpletuner.helpers.training.validation.validation_audio.log_audio_to_trackers"),
+        ):
+            validation.process_prompts(validation_type="intermediary")
+
+        self.assertEqual(sorted(validation.validation_audios.keys()), ["song_0", "song_1"])
+        torch.testing.assert_close(validation.validation_audios["song_0"][0], torch.zeros(1, 4))
+        torch.testing.assert_close(validation.validation_audios["song_1"][0], torch.ones(1, 4))
+        self.assertEqual(save_audio.call_count, 2)
+
 
 if __name__ == "__main__":
     unittest.main()