Supporting Wan model

dfm/src/megatron/data/common/diffusion_sample.py

@@ -34,9 +34,11 @@ class DiffusionSample(Sample):
         num_frames (Optional[torch.Tensor]): Number of frames in the video.
         padding_mask (Optional[torch.Tensor]): Mask indicating padding positions.
         seq_len_q (Optional[torch.Tensor]): Sequence length for query embeddings.
+        seq_len_q_padded (Optional[torch.Tensor]): Sequence length for query embeddings after padding.
         seq_len_kv (Optional[torch.Tensor]): Sequence length for key/value embeddings.
         pos_ids (Optional[torch.Tensor]): Positional IDs.
         latent_shape (Optional[torch.Tensor]): Shape of the latent tensor.
+        video_metadata (Optional[dict]): Metadata of the video.
     """
 
     video: torch.Tensor  # video latents (C T H W)
@@ -48,9 +50,12 @@ class DiffusionSample(Sample):
     num_frames: Optional[torch.Tensor] = None
     padding_mask: Optional[torch.Tensor] = None
     seq_len_q: Optional[torch.Tensor] = None
+    seq_len_q_padded: Optional[torch.Tensor] = None
     seq_len_kv: Optional[torch.Tensor] = None
+    seq_len_kv_padded: Optional[torch.Tensor] = None
     pos_ids: Optional[torch.Tensor] = None
     latent_shape: Optional[torch.Tensor] = None
+    video_metadata: Optional[dict] = None
 
     def to_dict(self) -> dict:
         """Converts the sample to a dictionary."""
@@ -64,9 +69,12 @@ def to_dict(self) -> dict:
             num_frames=self.num_frames,
             padding_mask=self.padding_mask,
             seq_len_q=self.seq_len_q,
+            seq_len_q_padded=self.seq_len_q_padded,
             seq_len_kv=self.seq_len_kv,
+            seq_len_kv_padded=self.seq_len_kv_padded,
             pos_ids=self.pos_ids,
             latent_shape=self.latent_shape,
+            video_metadata=self.video_metadata,
         )
 
     def __add__(self, other: Any) -> int:

dfm/src/megatron/data/common/diffusion_task_encoder_with_sp.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 import random
 from abc import ABC, abstractmethod
 from typing import List
@@ -103,9 +102,12 @@ def cat(attr):
             context_embeddings=cat("context_embeddings"),
             loss_mask=cat("loss_mask"),
             seq_len_q=cat("seq_len_q"),
+            seq_len_q_padded=cat("seq_len_q_padded"),
             seq_len_kv=cat("seq_len_kv"),
+            seq_len_kv_padded=cat("seq_len_kv_padded"),
             pos_ids=cat("pos_ids"),
             latent_shape=stack("latent_shape"),
+            video_metadata=[sample.video_metadata for sample in samples],
         )
 
     @stateless

dfm/src/megatron/data/wan/wan_energon_datamodule.py

@@ -0,0 +1,49 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=C0115,C0116,C0301
+
+from dataclasses import dataclass
+
+from megatron.bridge.data.utils import DatasetBuildContext
+from torch import int_repr
+
+from dfm.src.megatron.data.common.diffusion_energon_datamodule import DiffusionDataModule, DiffusionDataModuleConfig
+from dfm.src.megatron.data.wan.wan_taskencoder import WanTaskEncoder
+
+
+@dataclass(kw_only=True)
+class WanDataModuleConfig(DiffusionDataModuleConfig):
+    path: str
+    seq_length: int
+    packing_buffer_size: int
+    micro_batch_size: int
+    global_batch_size: int
+    num_workers: int_repr
+    dataloader_type: str = "external"
+
+    def __post_init__(self):
+        self.dataset = DiffusionDataModule(
+            path=self.path,
+            seq_length=self.seq_length,
+            packing_buffer_size=self.packing_buffer_size,
+            task_encoder=WanTaskEncoder(seq_length=self.seq_length, packing_buffer_size=self.packing_buffer_size),
+            micro_batch_size=self.micro_batch_size,
+            global_batch_size=self.global_batch_size,
+            num_workers=self.num_workers,
+        )
+        self.sequence_length = self.dataset.seq_length
+
+    def build_datasets(self, context: DatasetBuildContext):
+        return self.dataset.train_dataloader(), self.dataset.train_dataloader(), self.dataset.train_dataloader()

dfm/src/megatron/data/wan/wan_mock_datamodule.py

@@ -0,0 +1,144 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=C0115,C0116,C0301
+
+from dataclasses import dataclass
+
+import torch
+from megatron.bridge.data.utils import DatasetBuildContext, DatasetProvider
+from torch.utils.data import DataLoader, Dataset
+
+from dfm.src.megatron.model.wan.utils import patchify
+
+
+class _MockDataset(Dataset):
+    def __init__(self, length: int):
+        self.length = max(int(length), 1)
+
+    def __len__(self) -> int:
+        return self.length
+
+    def __getitem__(self, idx: int) -> dict:
+        return {}
+
+
+def mock_batch(
+    F_latents: int,
+    H_latents: int,
+    W_latents: int,
+    patch_temporal: int,
+    patch_spatial: int,
+    number_packed_samples: int,
+    context_seq_len: int,
+    context_embeddings_dim: int,
+) -> dict:
+    # set mock values for one video sample
+    video_latent = torch.randn(16, F_latents, H_latents, W_latents, dtype=torch.float32)
+    grid_size = torch.tensor(
+        [
+            video_latent.shape[1] // patch_temporal,
+            video_latent.shape[2] // patch_spatial,
+            video_latent.shape[3] // patch_spatial,
+        ],
+        dtype=torch.int32,
+    )
+    video_latent = patchify([video_latent], (patch_temporal, patch_spatial, patch_spatial))[0]
+    video_latent = torch.as_tensor(video_latent, dtype=torch.float32)
+    seq_len_q = video_latent.shape[0]
+    seq_len_q_padded = seq_len_q
+    loss_mask = torch.ones(seq_len_q, dtype=torch.bfloat16)
+    context_embeddings = torch.randn(context_seq_len, context_embeddings_dim, dtype=torch.float32)
+    seq_len_kv = context_embeddings.shape[0]
+    seq_len_kv_padded = seq_len_kv
+    video_metadata = {}
+
+    # set mock values for packed video samples
+    video_latents_packed = [video_latent for _ in range(number_packed_samples)]
+    video_latents_packed = torch.cat(video_latents_packed, dim=0)
+    loss_masks_packed = [loss_mask for _ in range(number_packed_samples)]
+    loss_masks_packed = torch.cat(loss_masks_packed, dim=0)
+    seq_len_q_packed = torch.tensor([seq_len_q for _ in range(number_packed_samples)], dtype=torch.int32)
+    seq_len_q_padded_packed = torch.tensor([seq_len_q_padded for _ in range(number_packed_samples)], dtype=torch.int32)
+    seq_len_kv_packed = torch.tensor([seq_len_kv for _ in range(number_packed_samples)], dtype=torch.int32)
+    seq_len_kv_padded_packed = torch.tensor(
+        [seq_len_kv_padded for _ in range(number_packed_samples)], dtype=torch.int32
+    )
+    grid_sizes_packed = torch.stack([grid_size for _ in range(number_packed_samples)], dim=0)
+    context_embeddings_packed = [context_embeddings for _ in range(number_packed_samples)]
+    context_embeddings_packed = torch.cat(context_embeddings_packed, dim=0)
+
+    ### Note: shape of sample's values
+    # video_latent: [num_patches, latents_channels * pF * pH * pW]
+    # grid_size: [F_patches, W_patches, H_patches]
+    # context_embeddings: [context_seq_len, text_embedding_dim]
+
+    batch = dict(
+        video_latents=video_latents_packed.unsqueeze(1),
+        context_embeddings=context_embeddings_packed.unsqueeze(1),
+        loss_mask=loss_masks_packed.unsqueeze(1),
+        seq_len_q=seq_len_q_packed,
+        seq_len_q_padded=seq_len_q_padded_packed,
+        seq_len_kv=seq_len_kv_packed,
+        seq_len_kv_padded=seq_len_kv_padded_packed,
+        grid_sizes=grid_sizes_packed,
+        video_metadata=video_metadata,
+    )
+
+    return batch
+
+
+@dataclass(kw_only=True)
+class WanMockDataModuleConfig(DatasetProvider):
+    path: str = ""
+    seq_length: int
+    packing_buffer_size: int
+    micro_batch_size: int
+    global_batch_size: int
+    num_workers: int
+    dataloader_type: str = "external"
+    F_latents: int = 24
+    H_latents: int = 104
+    W_latents: int = 60
+    patch_spatial: int = 2
+    patch_temporal: int = 1
+    number_packed_samples: int = 3
+    context_seq_len: int = 512
+    context_embeddings_dim: int = 4096
+
+    def __post_init__(self):
+        mock_ds = _MockDataset(length=1024)
+        self._train_dl = DataLoader(
+            mock_ds,
+            batch_size=self.micro_batch_size,
+            num_workers=self.num_workers,
+            collate_fn=lambda samples: mock_batch(
+                F_latents=self.F_latents,
+                H_latents=self.H_latents,
+                W_latents=self.W_latents,
+                patch_temporal=self.patch_temporal,
+                patch_spatial=self.patch_spatial,
+                number_packed_samples=self.number_packed_samples,
+                context_seq_len=self.context_seq_len,
+                context_embeddings_dim=self.context_embeddings_dim,
+            ),
+            shuffle=False,
+            drop_last=False,
+        )
+        self.sequence_length = self.seq_length
+
+    def build_datasets(self, _context: DatasetBuildContext):
+        if hasattr(self, "dataset"):
+            return self.dataset.train_dataloader(), self.dataset.train_dataloader(), self.dataset.train_dataloader()
+        return self._train_dl, self._train_dl, self._train_dl