[SP Sharding] Fix SP loss sharding on token axis (thw) with padding; add distributed correctness tests (#1045)

Fixes the sequence parallel sharding on t, now SP shards on t*h*w

fastvideo/tests/distributed/test_hsdp_sp_mismatch_sp2_hsdp_shard4.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Mismatch compatibility test (direction B):
+  - WORLD_SIZE=4
+  - SP: sp_size=2
+  - HSDP/FSDP2 mesh: replicate_dim=1, shard_dim=4 (mesh_shape=(1,4))
+
+We verify gradients from:
+  (A) full MSE: mean((pred-target)^2)
+  (B) SP-sharded loss: sp_size * local_sse / total_numel
+match under this mismatch.
+
+Run:
+  source /home/hao_lab/miniconda3/etc/profile.d/conda.sh
+  conda activate alexfv
+  torchrun --standalone --nproc_per_node=4 -m pytest -q fastvideo/tests/distributed/test_hsdp_sp_mismatch_sp2_hsdp_shard4.py
+"""
+
+import os
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from fastvideo.distributed import cleanup_dist_env_and_memory
+from fastvideo.distributed.parallel_state import (
+    get_sp_group,
+    maybe_init_distributed_environment_and_model_parallel,
+)
+from fastvideo.training.training_utils import shard_latents_across_sp
+
+
+def _world_size() -> int:
+    return int(os.environ.get("WORLD_SIZE", "1"))
+
+
+def _rank() -> int:
+    return int(os.environ.get("RANK", "0"))
+
+
+@pytest.fixture(scope="module")
+def dist_setup():
+    if _world_size() != 4:
+        pytest.skip("Designed for torchrun WORLD_SIZE=4")
+    if not torch.cuda.is_available():
+        pytest.skip("Requires CUDA")
+
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    torch.cuda.set_device(local_rank)
+
+    maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=2)
+    assert get_sp_group().world_size == 2
+    yield
+    if dist.is_available() and dist.is_initialized():
+        dist.barrier()
+    cleanup_dist_env_and_memory()
+
+
+class ScaleModule(torch.nn.Module):
+    def __init__(self, c: int, init: torch.Tensor):
+        super().__init__()
+        assert init.shape == (c,)
+        self.scale = torch.nn.Parameter(init.clone())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.scale.view(1, -1, 1, 1, 1)
+
+
+def _broadcast_(t: torch.Tensor, src: int = 0) -> None:
+    dist.broadcast(t, src=src)
+
+
+def _gather_full_vec_allranks(local_shard: torch.Tensor) -> torch.Tensor:
+    gathered = [torch.empty_like(local_shard) for _ in range(dist.get_world_size())]
+    dist.all_gather(gathered, local_shard)
+    return torch.cat(gathered, dim=0)
+
+
+def test_sp_sharded_loss_matches_full_mse_under_mismatch(dist_setup):
+    from torch.distributed.device_mesh import init_device_mesh
+    from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
+
+    device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}")
+    torch.manual_seed(2027)
+
+    b, c, t, h, w = 2, 8, 1, 1, 8  # thw=8 divisible by sp=2
+
+    x = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    if _rank() == 0:
+        x.normal_()
+    _broadcast_(x)
+    target = (0.25 * x).contiguous()
+
+    init = torch.empty((c,), device=device, dtype=torch.float32)
+    if _rank() == 0:
+        init.normal_()
+    _broadcast_(init)
+
+    mesh = init_device_mesh(
+        "cuda",
+        mesh_shape=(1, 4),
+        mesh_dim_names=("replicate", "shard"),
+    )
+
+    mp = MixedPrecisionPolicy(
+        param_dtype=torch.float32,
+        reduce_dtype=torch.float32,
+        output_dtype=torch.float32,
+        cast_forward_inputs=False,
+    )
+
+    # (A) Full loss.
+    m_full = ScaleModule(c=c, init=init).to(device)
+    fully_shard(m_full, mesh=mesh, reshard_after_forward=True, mp_policy=mp)
+    pred_full = m_full(x)
+    loss_full = ((pred_full - target) ** 2).mean()
+    loss_full.backward()
+    g_full = m_full.scale.grad
+    g_full_local = g_full.to_local() if hasattr(g_full, "to_local") else g_full
+    g_full_vec = _gather_full_vec_allranks(g_full_local.flatten())
+
+    # (B) SP-sharded loss.
+    m_sp = ScaleModule(c=c, init=init).to(device)
+    fully_shard(m_sp, mesh=mesh, reshard_after_forward=True, mp_policy=mp)
+    pred = m_sp(x)
+    sp = get_sp_group().world_size  # 2
+    sharded_pred = shard_latents_across_sp(pred)
+    sharded_target = shard_latents_across_sp(target)
+    local_sse = ((sharded_pred - sharded_target) ** 2).sum()
+    loss_sp = (sp * local_sse) / pred.numel()
+    loss_sp.backward()
+    g_sp = m_sp.scale.grad
+    g_sp_local = g_sp.to_local() if hasattr(g_sp, "to_local") else g_sp
+    g_sp_vec = _gather_full_vec_allranks(g_sp_local.flatten())
+
+    torch.testing.assert_close(g_sp_vec, g_full_vec, rtol=1e-4, atol=1e-5)
+
+

fastvideo/tests/distributed/test_hsdp_sp_mismatch_sp4_hsdp_shard2.py

@@ -0,0 +1,152 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+Mismatch compatibility test (direction A):
+  - WORLD_SIZE=4
+  - SP: sp_size=4
+  - HSDP/FSDP2 mesh: replicate_dim=2, shard_dim=2 (mesh_shape=(2,2))
+
+We verify gradients from:
+  (A) full MSE: mean((pred-target)^2)
+  (B) SP-sharded loss: sp_size * local_sse / total_numel
+match under this mismatch.
+
+Run:
+  source /home/hao_lab/miniconda3/etc/profile.d/conda.sh
+  conda activate alexfv
+  torchrun --standalone --nproc_per_node=4 -m pytest -q fastvideo/tests/distributed/test_hsdp_sp_mismatch_sp4_hsdp_shard2.py
+"""
+
+import os
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from fastvideo.distributed import cleanup_dist_env_and_memory
+from fastvideo.distributed.parallel_state import (
+    get_sp_group,
+    maybe_init_distributed_environment_and_model_parallel,
+)
+from fastvideo.training.training_utils import shard_latents_across_sp
+
+
+def _world_size() -> int:
+    return int(os.environ.get("WORLD_SIZE", "1"))
+
+
+def _rank() -> int:
+    return int(os.environ.get("RANK", "0"))
+
+
+@pytest.fixture(scope="module")
+def dist_setup():
+    if _world_size() != 4:
+        pytest.skip("Designed for torchrun WORLD_SIZE=4")
+    if not torch.cuda.is_available():
+        pytest.skip("Requires CUDA")
+
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    torch.cuda.set_device(local_rank)
+
+    maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=4)
+    assert get_sp_group().world_size == 4
+    yield
+    if dist.is_available() and dist.is_initialized():
+        dist.barrier()
+    cleanup_dist_env_and_memory()
+
+
+class ScaleModule(torch.nn.Module):
+    def __init__(self, c: int, init: torch.Tensor):
+        super().__init__()
+        assert init.shape == (c,)
+        self.scale = torch.nn.Parameter(init.clone())
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.scale.view(1, -1, 1, 1, 1)
+
+
+def _broadcast_(t: torch.Tensor, src: int = 0) -> None:
+    dist.broadcast(t, src=src)
+
+
+def _reconstruct_full_vec_from_replicate0(local_shard: torch.Tensor) -> torch.Tensor:
+    """
+    Mesh layout: [[0,2],[1,3]] with dim_names=("replicate","shard").
+    replicate0 group is ranks [0,2], and shard order inside replicate0 is [0,2].
+    Reconstruct a single full vector from replicate0 only, then broadcast to all ranks.
+    """
+    group_ranks = [0, 2]
+    pg = dist.new_group(ranks=group_ranks)
+    if dist.get_rank() in group_ranks:
+        gathered = [torch.empty_like(local_shard) for _ in range(len(group_ranks))]
+        dist.all_gather(gathered, local_shard, group=pg)
+        full = torch.cat(gathered, dim=0)
+    else:
+        full = torch.empty((local_shard.numel() * len(group_ranks),),
+                           device=local_shard.device,
+                           dtype=local_shard.dtype)
+    dist.broadcast(full, src=0)
+    return full
+
+
+def test_sp_sharded_loss_matches_full_mse_under_mismatch(dist_setup):
+    from torch.distributed.device_mesh import DeviceMesh
+    from torch.distributed.fsdp import MixedPrecisionPolicy, fully_shard
+
+    device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}")
+    torch.manual_seed(2026)
+
+    b, c, t, h, w = 2, 8, 1, 1, 8  # thw=8 divisible by sp=4
+
+    x = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    if _rank() == 0:
+        x.normal_()
+    _broadcast_(x)
+    target = (0.5 * x).contiguous()
+
+    init = torch.empty((c,), device=device, dtype=torch.float32)
+    if _rank() == 0:
+        init.normal_()
+    _broadcast_(init)
+
+    mesh = DeviceMesh(
+        "cuda",
+        [[0, 2], [1, 3]],
+        mesh_dim_names=("replicate", "shard"),
+    )
+
+    mp = MixedPrecisionPolicy(
+        param_dtype=torch.float32,
+        reduce_dtype=torch.float32,
+        output_dtype=torch.float32,
+        cast_forward_inputs=False,
+    )
+
+    # (A) Full loss.
+    m_full = ScaleModule(c=c, init=init).to(device)
+    fully_shard(m_full, mesh=mesh, reshard_after_forward=True, mp_policy=mp)
+    pred_full = m_full(x)
+    loss_full = ((pred_full - target) ** 2).mean()
+    loss_full.backward()
+    g_full = m_full.scale.grad
+    g_full_local = g_full.to_local() if hasattr(g_full, "to_local") else g_full
+    g_full_vec = _reconstruct_full_vec_from_replicate0(g_full_local.flatten())
+
+    # (B) SP-sharded loss.
+    m_sp = ScaleModule(c=c, init=init).to(device)
+    fully_shard(m_sp, mesh=mesh, reshard_after_forward=True, mp_policy=mp)
+    pred = m_sp(x)
+    sp = get_sp_group().world_size  # 4
+    sharded_pred = shard_latents_across_sp(pred)
+    sharded_target = shard_latents_across_sp(target)
+    local_sse = ((sharded_pred - sharded_target) ** 2).sum()
+    loss_sp = (sp * local_sse) / pred.numel()
+    loss_sp.backward()
+    g_sp = m_sp.scale.grad
+    g_sp_local = g_sp.to_local() if hasattr(g_sp, "to_local") else g_sp
+    g_sp_vec = _reconstruct_full_vec_from_replicate0(g_sp_local.flatten())
+
+    torch.testing.assert_close(g_sp_vec, g_full_vec, rtol=1e-4, atol=1e-5)
+
+

fastvideo/tests/distributed/test_sp_shard_latents_and_loss.py

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+import os
+
+import pytest
+import torch
+import torch.distributed as dist
+
+from fastvideo.distributed import (cleanup_dist_env_and_memory,
+                                  get_sp_group,
+                                  maybe_init_distributed_environment_and_model_parallel)
+from fastvideo.training.training_utils import shard_latents_across_sp
+
+
+def _world_size() -> int:
+    return int(os.environ.get("WORLD_SIZE", "1"))
+
+
+def _rank() -> int:
+    return int(os.environ.get("RANK", "0"))
+
+
+@pytest.fixture(scope="module")
+def sp_dist():
+    """
+    These tests are intended to be run under torchrun with multiple GPUs, e.g.
+      torchrun --nproc_per_node=2 -m pytest -q fastvideo/tests/distributed/test_sp_shard_latents_and_loss.py
+      torchrun --nproc_per_node=4 -m pytest -q fastvideo/tests/distributed/test_sp_shard_latents_and_loss.py
+      torchrun --nproc_per_node=8 -m pytest -q fastvideo/tests/distributed/test_sp_shard_latents_and_loss.py
+    """
+    ws = _world_size()
+    if ws <= 1:
+        pytest.skip("Requires torchrun with WORLD_SIZE>1")
+    if not torch.cuda.is_available():
+        pytest.skip("Requires CUDA")
+
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    torch.cuda.set_device(local_rank)
+
+    torch.manual_seed(1234)
+    # Initialize FastVideo dist + model-parallel groups.
+    maybe_init_distributed_environment_and_model_parallel(tp_size=1, sp_size=ws)
+    assert get_sp_group().world_size == ws
+
+    yield
+
+    if dist.is_available() and dist.is_initialized():
+        dist.barrier()
+    cleanup_dist_env_and_memory()
+
+
+def _broadcast_tensor_(t: torch.Tensor, src: int = 0) -> torch.Tensor:
+    if dist.is_available() and dist.is_initialized():
+        dist.broadcast(t, src=src)
+    return t
+
+
+def _all_gather_cat(t: torch.Tensor, dim: int) -> torch.Tensor:
+    ws = _world_size()
+    gathered = [torch.empty_like(t) for _ in range(ws)]
+    dist.all_gather(gathered, t)
+    return torch.cat(gathered, dim=dim)
+
+
+def test_shard_latents_t_not_divisible_but_thw_divisible_no_error(sp_dist):
+    """
+    Regression test for the commit "fix splitting on t":
+    - t is NOT divisible by sp_size
+    - (t*h*w) IS divisible by sp_size
+    - sharding should NOT raise, and shards should round-trip to the original.
+    """
+    ws = _world_size()
+    device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}")
+
+    b, c, t, h, w = 2, 3, 3, 2, 4  # t=3 (not divisible by 8), thw=24 (divisible by 8)
+    assert (t * h * w) % ws == 0
+    assert t % ws != 0
+
+    latents = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    if _rank() == 0:
+        latents.normal_()
+    _broadcast_tensor_(latents)
+
+    shard = shard_latents_across_sp(latents)
+    assert shard.shape == (b, c, (t * h * w) // ws)
+
+    gathered = _all_gather_cat(shard, dim=2)
+    torch.testing.assert_close(gathered, latents.reshape(b, c, t * h * w))
+
+    dist.barrier()
+
+
+@pytest.mark.parametrize(
+    "shape",
+    [
+        # No padding needed: thw divisible by 8, but t is not.
+        (2, 1, 3, 2, 4),  # thw=24
+        # Padding needed: thw not divisible by 8.
+        (2, 1, 3, 2, 5),  # thw=30 -> pad to 32
+    ],
+)
+def test_sharded_loss_gradient_matches_full_mse_after_rank_avg(sp_dist, shape):
+    """
+    Validates the math behind the new loss computation:
+
+    If each rank computes:
+      loss_rank = sp_world_size * local_sse / total_numel
+    and the training stack averages gradients across ranks, then the resulting
+    gradient should match the full MSE gradient.
+    """
+    ws = _world_size()
+    device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}")
+    b, c, t, h, w = shape
+
+    # Make inputs identical across ranks for deterministic comparison.
+    init = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    x = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    y = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    if _rank() == 0:
+        init.normal_()
+        x.normal_()
+        y.normal_()
+    _broadcast_tensor_(init)
+    _broadcast_tensor_(x)
+    _broadcast_tensor_(y)
+
+    # Full loss baseline (redundant compute; correct reference).
+    w_full = torch.nn.Parameter(init.clone())
+    pred_full = w_full * x
+    full_loss = ((pred_full - y)**2).mean()
+    full_loss.backward()
+    grad_full = w_full.grad.detach().clone()
+
+    # Sharded loss (compute on a shard only).
+    w_shard = torch.nn.Parameter(init.clone())
+    pred = w_shard * x
+    sharded_pred = shard_latents_across_sp(pred)
+    sharded_target = shard_latents_across_sp(y)
+    local_sse = ((sharded_pred - sharded_target)**2).sum()
+    shard_loss = ws * local_sse / pred.numel()
+    shard_loss.backward()
+    grad_shard = w_shard.grad.detach().clone()
+
+    # Simulate "gradient averaging across ranks" (DDP/FSDP replicate-dim behavior).
+    dist.all_reduce(grad_shard, op=dist.ReduceOp.SUM)
+    grad_shard /= ws
+
+    # (Optional) also average grad_full, for symmetry.
+    dist.all_reduce(grad_full, op=dist.ReduceOp.SUM)
+    grad_full /= ws
+
+    torch.testing.assert_close(grad_shard, grad_full, rtol=1e-5, atol=1e-6)
+
+    dist.barrier()
+
+
+def test_padding_does_not_change_global_sse(sp_dist):
+    """
+    Padding-specific correctness test.
+
+    When (t*h*w) is NOT divisible by sp_size, shard_latents_across_sp pads with
+    zeros on the flattened axis. This test validates that:
+    1) Summed SSE across SP ranks equals the full (unpadded) SSE.
+    2) Any padded tokens that land on a rank are exactly zero for both pred/target.
+    """
+    from fastvideo.distributed.utils import compute_padding_for_sp
+
+    ws = _world_size()
+    device = torch.device(f"cuda:{int(os.environ.get('LOCAL_RANK', '0'))}")
+
+    b, c, t, h = 2, 2, 1, 1
+    # Force padding for any ws>1: seq_len = 2*ws + 1  -> remainder 1
+    w = 2 * ws + 1
+    seq_len = t * h * w
+    assert seq_len % ws != 0
+
+    pred = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    target = torch.empty((b, c, t, h, w), device=device, dtype=torch.float32)
+    if _rank() == 0:
+        pred.normal_()
+        target.normal_()
+    _broadcast_tensor_(pred)
+    _broadcast_tensor_(target)
+
+    # Full (unpadded) SSE reference.
+    full_sse = ((pred - target)**2).sum()
+
+    # Sharded local SSE (includes padded region, which should contribute 0).
+    sharded_pred = shard_latents_across_sp(pred)
+    sharded_target = shard_latents_across_sp(target)
+    local_sse = ((sharded_pred - sharded_target)**2).sum()
+
+    global_sse = local_sse.clone()
+    dist.all_reduce(global_sse, op=dist.ReduceOp.SUM)
+    torch.testing.assert_close(global_sse, full_sse, rtol=1e-5, atol=1e-6)
+
+    # Explicitly validate padded tokens are zeros on ranks that cover them.
+    padded_seq_len, padding_amount = compute_padding_for_sp(seq_len, ws)
+    assert padding_amount > 0
+    elements_per_rank = padded_seq_len // ws
+    start = _rank() * elements_per_rank
+    end = (_rank() + 1) * elements_per_rank
+
+    if end > seq_len:
+        # There is padding on this rank: positions [max(seq_len, start), end)
+        pad_start_global = max(seq_len, start)
+        pad_end_global = end
+        pad_start_local = pad_start_global - start
+        pad_end_local = pad_end_global - start
+        assert pad_start_local < pad_end_local
+
+        pad_slice_pred = sharded_pred[:, :, pad_start_local:pad_end_local]
+        pad_slice_target = sharded_target[:, :, pad_start_local:pad_end_local]
+        torch.testing.assert_close(pad_slice_pred, torch.zeros_like(pad_slice_pred))
+        torch.testing.assert_close(pad_slice_target, torch.zeros_like(pad_slice_target))
+
+    dist.barrier()
+
+

fastvideo/training/training_pipeline.py

@@ -443,11 +443,26 @@ class TrainingPipeline(LoRAPipeline, ABC):
             # make sure no implicit broadcasting happens
             assert model_pred.shape == target.shape, f"model_pred.shape: {model_pred.shape}, target.shape: {target.shape}"
 
-            sharded_pred = shard_latents_across_sp(model_pred)
-            sharded_target = shard_latents_across_sp(target)
-            loss = (torch.mean(
-                (sharded_pred.float() - sharded_target.float())**2) /
-                    self.training_args.gradient_accumulation_steps)
+            # Defensive: avoid NaNs/div0 if an upstream bug ever produces empty tensors.
+            # mean(empty) -> NaN; division by 0 -> inf/NaN. Keep a 0 loss with grad history.
+            if model_pred.numel() == 0:
+                loss = (model_pred.sum() *
+                        0.0) / self.training_args.gradient_accumulation_steps
+            else:
+                # Compute MSE on one SP shard. We shard along the flattened token axis
+                # (t*h*w) with optional padding so (t*h*w) need not be divisible by sp_size.
+                sp_world_size = get_sp_group().world_size
+                if sp_world_size > 1:
+                    sharded_pred = shard_latents_across_sp(model_pred)
+                    sharded_target = shard_latents_across_sp(target)
+                    local_sse = ((sharded_pred.float() -
+                                  sharded_target.float())**2).sum()
+                    loss = (sp_world_size * local_sse / model_pred.numel()
+                            ) / self.training_args.gradient_accumulation_steps
+                else:
+                    loss = (torch.mean(
+                        (model_pred.float() - target.float())**2) /
+                            self.training_args.gradient_accumulation_steps)
 
             loss.backward()
 

fastvideo/training/training_utils.py

@@ -20,9 +20,10 @@ from fastvideo.training.checkpointing_utils import (ModelWrapper,
                                                     RandomStateWrapper,
                                                     SchedulerWrapper)
 
-from einops import rearrange
 from fastvideo.distributed.parallel_state import (get_sp_parallel_rank,
                                                   get_sp_world_size)
+from fastvideo.distributed.utils import (compute_padding_for_sp,
+                                         pad_sequence_tensor)
 
 logger = init_logger(__name__)
 
@@ -867,10 +868,25 @@ def shard_latents_across_sp(latents: torch.Tensor) -> torch.Tensor:
     sp_world_size = get_sp_world_size()
     rank_in_sp_group = get_sp_parallel_rank()
     if sp_world_size > 1:
-        latents = rearrange(latents,
-                            "b c (n s) h w -> b c n s h w",
-                            n=sp_world_size).contiguous()
-        latents = latents[:, :, rank_in_sp_group, :, :, :]
+        # Shard on the flattened token axis (t*h*w) rather than raw `t`, so we
+        # don't require t % sp_world_size == 0. Pad on the flattened axis if needed.
+        assert latents.ndim == 5, f"Expected latents [b,c,t,h,w], got {latents.shape}"
+        b, c, t, h, w = latents.shape
+        latents = latents.reshape(b, c, t * h * w)
+
+        original_seq_len = latents.shape[2]
+        padded_seq_len, padding_amount = compute_padding_for_sp(
+            original_seq_len, sp_world_size)
+        if padding_amount > 0:
+            latents = pad_sequence_tensor(latents,
+                                          padded_seq_len,
+                                          seq_dim=2,
+                                          pad_value=0.0)
+
+        elements_per_rank = padded_seq_len // sp_world_size
+        start = rank_in_sp_group * elements_per_rank
+        end = (rank_in_sp_group + 1) * elements_per_rank
+        latents = latents[:, :, start:end].contiguous()
     return latents