Merge pull request #3095 from bghira/agent/dataset-batch-distributed-runtime

Support variable dataset batch sizes across distributed ranks

Author: bghira

simpletuner/helpers/data_backend/runtime/context_parallel_sync.py

@@ -22,6 +22,7 @@
 import os
 import random
 from contextlib import contextmanager
+from dataclasses import dataclass
 from typing import Any, Optional, Tuple
 
 import torch
@@ -40,6 +41,18 @@
 CP_SKIP_SAMPLING_SENTINEL = "__CP_SKIP_SAMPLING__"
 
 
+@dataclass(frozen=True)
+class DistributedBatchLayout:
+    local_batch_size: int
+    global_batch_size: int
+    local_batch_offset: int
+    data_rank: int
+    data_parallel_size: int
+    model_replica_size: int
+    world_size: int
+    data_replica_batch_sizes: tuple[int, ...]
+
+
 def _normalize_parallel_size(value: Any, name: str) -> int:
     if value is None:
         return 1
@@ -219,6 +232,122 @@ def get_model_replica_data_info(
     return True, data_rank, data_local_rank, data_group_size, dp_replicate_size
 
 
+def resolve_distributed_batch_layout(accelerator, local_batch_size: int) -> DistributedBatchLayout:
+    """Resolve sample counts and this data replica's offset using fixed-shape collectives."""
+    local_batch_size = _normalize_parallel_size(local_batch_size, "local_batch_size")
+    if local_batch_size < 1:
+        raise ValueError("local_batch_size must be greater than 0.")
+
+    if accelerator is None:
+        return DistributedBatchLayout(local_batch_size, local_batch_size, 0, 0, 1, 1, 1, (local_batch_size,))
+
+    world_size = _normalize_parallel_size(getattr(accelerator, "num_processes", 1), "num_processes")
+    process_index = _normalize_parallel_size(getattr(accelerator, "process_index", 0), "process_index")
+    if world_size == 1:
+        return DistributedBatchLayout(local_batch_size, local_batch_size, 0, 0, 1, 1, 1, (local_batch_size,))
+
+    count = torch.tensor([local_batch_size], device=accelerator.device, dtype=torch.long)
+    gathered_counts = accelerator.gather(count).reshape(-1)
+    if gathered_counts.numel() != world_size:
+        raise RuntimeError(
+            f"Distributed batch count gather returned {gathered_counts.numel()} values for world size {world_size}."
+        )
+    world_batch_sizes = tuple(int(value) for value in gathered_counts.cpu().tolist())
+
+    data_enabled, data_rank, _data_local_rank, model_replica_size, data_parallel_size = get_model_replica_data_info(
+        accelerator
+    )
+    if data_enabled:
+        replica_batch_sizes = []
+        for replica_rank in range(data_parallel_size):
+            start = replica_rank * model_replica_size
+            replica_counts = world_batch_sizes[start : start + model_replica_size]
+            if len(set(replica_counts)) != 1:
+                raise RuntimeError(
+                    "Model-parallel ranks received different local batch sizes: "
+                    f"data replica {replica_rank} reported {replica_counts}."
+                )
+            replica_batch_sizes.append(replica_counts[0])
+        data_replica_batch_sizes = tuple(replica_batch_sizes)
+    else:
+        data_rank = process_index
+        data_parallel_size = world_size
+        model_replica_size = 1
+        data_replica_batch_sizes = world_batch_sizes
+
+    return DistributedBatchLayout(
+        local_batch_size=local_batch_size,
+        global_batch_size=sum(data_replica_batch_sizes),
+        local_batch_offset=sum(data_replica_batch_sizes[:data_rank]),
+        data_rank=data_rank,
+        data_parallel_size=data_parallel_size,
+        model_replica_size=model_replica_size,
+        world_size=world_size,
+        data_replica_batch_sizes=data_replica_batch_sizes,
+    )
+
+
+def gather_variable_batch_tensor(
+    tensor: torch.Tensor,
+    accelerator,
+    layout: Optional[DistributedBatchLayout] = None,
+) -> torch.Tensor:
+    """Gather every rank's per-sample tensor when data replicas have different batch sizes."""
+    if tensor.ndim < 1:
+        raise ValueError("Variable batch tensor gather requires a batch dimension.")
+    if layout is None:
+        layout = resolve_distributed_batch_layout(accelerator, tensor.shape[0])
+    if tensor.shape[0] != layout.local_batch_size:
+        raise ValueError(f"Tensor batch dimension {tensor.shape[0]} does not match layout size {layout.local_batch_size}.")
+    if layout.world_size == 1:
+        return tensor.detach()
+
+    maximum_batch_size = max(layout.data_replica_batch_sizes)
+    padded = tensor.detach()
+    if padded.shape[0] < maximum_batch_size:
+        padding = padded.new_zeros((maximum_batch_size - padded.shape[0], *padded.shape[1:]))
+        padded = torch.cat((padded, padding), dim=0)
+
+    gathered = accelerator.gather(padded)
+    expected_first_dimension = layout.world_size * maximum_batch_size
+    if gathered.shape[0] != expected_first_dimension:
+        raise RuntimeError(
+            "Distributed tensor gather returned an unexpected first dimension: "
+            f"expected {expected_first_dimension}, got {gathered.shape[0]}."
+        )
+    gathered = gathered.reshape(layout.world_size, maximum_batch_size, *padded.shape[1:])
+    rank_tensors = []
+    for world_rank in range(layout.world_size):
+        replica_rank = world_rank // layout.model_replica_size
+        batch_size = layout.data_replica_batch_sizes[replica_rank]
+        rank_tensors.append(gathered[world_rank, :batch_size])
+    return torch.cat(rank_tensors, dim=0)
+
+
+def gather_sample_weighted_scalar(value: torch.Tensor, local_batch_size: int, accelerator) -> torch.Tensor:
+    """Return a sample-weighted world mean from one fixed-shape contribution per rank."""
+    local_batch_size = _normalize_parallel_size(local_batch_size, "local_batch_size")
+    if local_batch_size < 1:
+        raise ValueError("local_batch_size must be greater than 0.")
+    if value.numel() != 1:
+        raise ValueError("Sample-weighted scalar gather requires a scalar tensor.")
+
+    value = value.detach().float().reshape(())
+    world_size = _normalize_parallel_size(getattr(accelerator, "num_processes", 1), "num_processes")
+    if world_size == 1:
+        return value
+
+    local_count = value.new_tensor(float(local_batch_size))
+    contribution = torch.stack((value * local_count, local_count))
+    gathered = accelerator.gather(contribution).reshape(-1, 2)
+    if gathered.shape[0] != world_size:
+        raise RuntimeError(
+            f"Distributed loss gather returned {gathered.shape[0]} contributions for world size {world_size}."
+        )
+    totals = gathered.sum(dim=0)
+    return totals[0] / totals[1]
+
+
 class ContextParallelBatchSynchronizer:
     """
     Caches context-parallel information for efficient batch synchronization.

simpletuner/helpers/distillation/anyflow/distiller.py

@@ -10,7 +10,10 @@
 from safetensors.torch import load_file, save_file
 
 from simpletuner.helpers.data_backend.dataset_types import DatasetType
-from simpletuner.helpers.data_backend.runtime.context_parallel_sync import get_model_replica_data_info
+from simpletuner.helpers.data_backend.runtime.context_parallel_sync import (
+    gather_variable_batch_tensor,
+    resolve_distributed_batch_layout,
+)
 from simpletuner.helpers.distillation.anyflow.scheduler import AnyFlowValidationScheduler
 from simpletuner.helpers.distillation.common import DistillationBase
 from simpletuner.helpers.distillation.registry import DistillationRegistry
@@ -585,20 +588,17 @@ def _prepare_meanflow_pair(self, prepared_batch: Dict[str, Any], model) -> tuple
         r_base = torch.minimum(first, second)
 
         accelerator = getattr(model, "accelerator", getattr(self.teacher_model, "accelerator", None))
-        process_index = int(getattr(accelerator, "process_index", 0) or 0)
-        num_processes = int(getattr(accelerator, "num_processes", 1) or 1)
-        data_parallel_enabled, data_rank, _, _, data_parallel_size = get_model_replica_data_info(accelerator)
-        if data_parallel_enabled:
-            process_index = data_rank
-            num_processes = data_parallel_size
-        global_batch_size = batch_size * num_processes
-        global_indices = process_index * batch_size + torch.arange(batch_size, device=device)
+        batch_layout = resolve_distributed_batch_layout(accelerator, batch_size)
+        self._distributed_batch_accelerator = accelerator
+        self._distributed_batch_layout = batch_layout
+        global_batch_size = batch_layout.global_batch_size
+        global_indices = batch_layout.local_batch_offset + torch.arange(batch_size, device=device)
         diffusion_count = round(float(self.config["diffusion_ratio"]) * global_batch_size)
         consistency_count = round(float(self.config["consistency_ratio"]) * global_batch_size)
         effective_diffusion_count = min(diffusion_count, global_batch_size)
         effective_consistency_count = min(consistency_count, global_batch_size - effective_diffusion_count)
         arbitrary_count = global_batch_size - effective_diffusion_count - effective_consistency_count
-        if process_index == 0 and not getattr(self, "_meanflow_branch_mix_logged", False):
+        if batch_layout.data_rank == 0 and not getattr(self, "_meanflow_branch_mix_logged", False):
             self.logger.info(
                 "AnyFlow interval mixture at global batch %d: diffusion=%d, consistency=%d, arbitrary=%d.",
                 global_batch_size,
@@ -1253,13 +1253,16 @@ def _clear_prediction_buffers(output: Dict[str, Any]) -> None:
         if isinstance(hidden_states_buffer, dict):
             hidden_states_buffer.clear()
 
-    @staticmethod
-    def _gather_detached(tensor: torch.Tensor) -> torch.Tensor:
-        if not torch.distributed.is_available() or not torch.distributed.is_initialized():
-            return tensor.detach()
-        gathered = [torch.empty_like(tensor) for _ in range(torch.distributed.get_world_size())]
-        torch.distributed.all_gather(gathered, tensor.detach())
-        return torch.cat(gathered, dim=0)
+    def _gather_detached(self, tensor: torch.Tensor) -> torch.Tensor:
+        accelerator = getattr(
+            self,
+            "_distributed_batch_accelerator",
+            getattr(self.teacher_model, "accelerator", None),
+        )
+        layout = getattr(self, "_distributed_batch_layout", None)
+        if layout is None or layout.local_batch_size != tensor.shape[0]:
+            layout = resolve_distributed_batch_layout(accelerator, tensor.shape[0])
+        return gather_variable_batch_tensor(tensor, accelerator, layout)
 
     def _discriminator_state_dict(self) -> Dict[str, torch.Tensor]:
         adapter_name = str(self._discriminator_adapter_name)

simpletuner/helpers/models/common.py

@@ -33,6 +33,7 @@
 
 from simpletuner.diff2flow import DiffusionToFlowBridge
 from simpletuner.helpers.assistant_lora import build_adapter_stack, set_adapter_stack
+from simpletuner.helpers.data_backend.runtime.context_parallel_sync import resolve_distributed_batch_layout
 from simpletuner.helpers.models.foundation_mixins import (
     AudioTransformMixin,
     PipelineSupportMixin,
@@ -4714,31 +4715,30 @@ def sample_flow_sigmas(self, batch: dict, state: dict) -> tuple[torch.Tensor, to
                 timesteps = base_timesteps.expand(bsz)
             else:
                 if timestep_mode == "round-robin":
-                    world_size = max(1, int(getattr(self.accelerator, "num_processes", 1) or 1))
-                    process_index = int(getattr(self.accelerator, "process_index", 0) or 0)
-                    if base_timesteps.numel() < bsz * world_size and not getattr(
+                    batch_layout = resolve_distributed_batch_layout(self.accelerator, bsz)
+                    global_batch_size = batch_layout.global_batch_size
+                    if base_timesteps.numel() < global_batch_size and not getattr(
                         self, "_flow_custom_timestep_overlap_warning_logged", False
                     ):
                         logger.warning(
                             "flow_timesteps_mode=round-robin has %s custom timestep(s), but the global batch "
                             "consumes %s sample(s) per step. Different ranks may reuse timestep entries on the same step; "
-                            "provide at least train_batch_size * num_processes values for non-overlapping per-step "
-                            "coverage.",
+                            "provide at least the global per-step sample count for non-overlapping coverage.",
                             base_timesteps.numel(),
-                            bsz * world_size,
+                            global_batch_size,
                         )
                         self._flow_custom_timestep_overlap_warning_logged = True
                     if not hasattr(self, "_flow_custom_timestep_cursor"):
                         resume_step = getattr(self, "_flow_custom_timestep_resume_step", None)
                         completed_steps = int(resume_step if resume_step is not None else state.get("global_step", 0) or 0)
-                        self._flow_custom_timestep_cursor = (
-                            completed_steps * bsz * world_size + process_index * bsz
-                        ) % base_timesteps.numel()
+                        self._flow_custom_timestep_cursor = (completed_steps * global_batch_size) % base_timesteps.numel()
                         if resume_step is not None:
                             delattr(self, "_flow_custom_timestep_resume_step")
                     cursor = int(getattr(self, "_flow_custom_timestep_cursor", 0))
-                    indices = (torch.arange(bsz, device=self.accelerator.device) + cursor) % base_timesteps.numel()
-                    self._flow_custom_timestep_cursor = (cursor + bsz * world_size) % base_timesteps.numel()
+                    indices = (
+                        torch.arange(bsz, device=self.accelerator.device) + cursor + batch_layout.local_batch_offset
+                    ) % base_timesteps.numel()
+                    self._flow_custom_timestep_cursor = (cursor + global_batch_size) % base_timesteps.numel()
                 else:
                     indices = torch.randint(0, base_timesteps.numel(), (bsz,), device=self.accelerator.device)
                 sigmas = base_sigmas[indices]

simpletuner/helpers/training/trainer.py

@@ -48,7 +48,10 @@
     run_distillation_cache_generation,
 )
 from simpletuner.helpers.data_backend.runtime import random_dataloader_iterator
-from simpletuner.helpers.data_backend.runtime.context_parallel_sync import ContextParallelBatchSynchronizer
+from simpletuner.helpers.data_backend.runtime.context_parallel_sync import (
+    ContextParallelBatchSynchronizer,
+    gather_sample_weighted_scalar,
+)
 from simpletuner.helpers.data_backend.runtime.schedule import normalize_start_epoch, normalize_start_step
 from simpletuner.helpers.distillation.composition import resolve_configured_distiller_requirement_profile
 from simpletuner.helpers.distillation.requirements import EMPTY_PROFILE, DistillerRequirementProfile
@@ -7031,11 +7034,11 @@ def train(self):
                             f"filepaths={batch_filepaths}, loss_logs={loss_logs_context})."
                         )
 
-                    # Gather the losses across all processes for logging (if using distributed training)
-                    avg_loss = self.accelerator.gather(loss.repeat(int(bsz))).mean()
+                    # Keep metric collectives fixed-shape when ranks use different dataset batch sizes.
+                    avg_loss = gather_sample_weighted_scalar(loss, int(bsz), self.accelerator)
                     self.train_loss += avg_loss.item() / self.config.gradient_accumulation_steps
                     if aux_loss_logs is not None:
-                        avg_diffusion_loss = self.accelerator.gather(diffusion_loss.repeat(int(bsz))).mean()
+                        avg_diffusion_loss = gather_sample_weighted_scalar(diffusion_loss, int(bsz), self.accelerator)
                         self.train_diffusion_loss += avg_diffusion_loss.item() / self.config.gradient_accumulation_steps
                     # Backpropagate
                     self.grad_norm = None

tests/helpers/distillation/test_anyflow_distiller.py

@@ -11,6 +11,7 @@
 from safetensors.torch import save_file
 
 import tests.test_stubs  # noqa: F401
+from simpletuner.helpers.data_backend.runtime.context_parallel_sync import DistributedBatchLayout
 from simpletuner.helpers.distillation.anyflow.distiller import AnyFlowDistiller
 from simpletuner.helpers.distillation.anyflow.scheduler import AnyFlowValidationScheduler
 from simpletuner.helpers.distillation.factory import DistillerFactory
@@ -529,8 +530,17 @@ def test_meanflow_branch_assignment_uses_data_replica_rank_with_context_parallel
         with (
             patch("torch.rand", side_effect=draws),
             patch(
-                "simpletuner.helpers.distillation.anyflow.distiller.get_model_replica_data_info",
-                return_value=(True, 1, 0, 2, 4),
+                "simpletuner.helpers.distillation.anyflow.distiller.resolve_distributed_batch_layout",
+                return_value=DistributedBatchLayout(
+                    local_batch_size=2,
+                    global_batch_size=8,
+                    local_batch_offset=2,
+                    data_rank=1,
+                    data_parallel_size=4,
+                    model_replica_size=2,
+                    world_size=8,
+                    data_replica_batch_sizes=(2, 2, 2, 2),
+                ),
             ),
         ):
             batch = _prepared_batch()
@@ -539,6 +549,53 @@ def test_meanflow_branch_assignment_uses_data_replica_rank_with_context_parallel
         self.assertTrue(torch.equal(batch["anyflow_diffusion_mask"], torch.tensor([True, True])))
         self.assertTrue(torch.equal(batch["anyflow_consistency_mask"], torch.tensor([False, False])))
 
+    def test_meanflow_branch_assignment_uses_rank_varying_batch_offsets(self):
+        model = _FlowModel()
+        distiller = AnyFlowDistiller(
+            teacher_model=model,
+            noise_scheduler=None,
+            config={
+                "model_type": "lora",
+                "diffusion_ratio": 0.5,
+                "consistency_ratio": 0.25,
+            },
+        )
+        latents = torch.zeros(3, 1, 2, 2)
+        noise = torch.ones_like(latents)
+        sigmas = torch.tensor([0.9, 0.6, 0.3]).view(3, 1, 1, 1)
+        batch = {
+            "latents": latents,
+            "noise": noise,
+            "input_noise": noise.clone(),
+            "sigmas": sigmas,
+            "timesteps": torch.tensor([900.0, 600.0, 300.0]),
+            "noisy_latents": (1 - sigmas) * latents + sigmas * noise,
+        }
+        draws = [torch.tensor([0.8, 0.7, 0.6]), torch.tensor([0.2, 0.4, 0.3])]
+        layout = DistributedBatchLayout(
+            local_batch_size=3,
+            global_batch_size=4,
+            local_batch_offset=1,
+            data_rank=1,
+            data_parallel_size=2,
+            model_replica_size=1,
+            world_size=2,
+            data_replica_batch_sizes=(1, 3),
+        )
+
+        with (
+            patch("torch.rand", side_effect=draws),
+            patch(
+                "simpletuner.helpers.distillation.anyflow.distiller.resolve_distributed_batch_layout",
+                return_value=layout,
+            ),
+        ):
+            distiller._prepare_meanflow_pair(batch, model)
+
+        self.assertTrue(torch.equal(batch["anyflow_diffusion_mask"], torch.tensor([True, False, False])))
+        self.assertTrue(torch.equal(batch["anyflow_consistency_mask"], torch.tensor([False, True, False])))
+        self.assertTrue(torch.equal(batch["anyflow_arbitrary_mask"], torch.tensor([False, False, True])))
+
     def test_meanflow_warns_once_when_global_batch_omits_enabled_branch(self):
         model = _FlowModel()
         distiller = AnyFlowDistiller(