Change dataloader

fastvideo/v1/dataset/parquet_datasets.py

@@ -1,21 +1,29 @@
 import argparse
+import json
 import os
 import random
 import time
+from collections import defaultdict
 
 import numpy as np
 import pyarrow.parquet as pq
 import torch
 import tqdm
+from einops import rearrange
 from torch import distributed as dist
-from torch.utils.data import IterableDataset, get_worker_info
+from torch.utils.data import Dataset
 from torchdata.stateful_dataloader import StatefulDataLoader
 
+from fastvideo.v1.distributed import (get_sequence_model_parallel_rank,
+                                      get_sp_group)
+from fastvideo.v1.logger import init_logger
+
 # Path to your dataset
 dataset_path = "/mnt/sharefs/users/hao.zhang/Vchitect-2M/Vchitect-2M-laten-93x512x512/train/"
+logger = init_logger(__name__)
 
 
-class ParquetVideoTextDataset(IterableDataset):
+class ParquetVideoTextDataset(Dataset):
     """Efficient loader for video-text data from a directory of Parquet files."""
 
     def __init__(self,
@@ -24,237 +32,185 @@ class ParquetVideoTextDataset(IterableDataset):
                  rank: int = 0,
                  world_size: int = 1,
                  cfg_rate: float = 0.0,
-                 num_latent_t: int = 2):
+                 num_latent_t: int = 2,
+                 seed: int = 0):
         super().__init__()
         self.path = str(path)
         self.batch_size = batch_size
         self.rank = rank
-        self.world_size = world_size
+        self.local_rank = get_sequence_model_parallel_rank()
+        self.sp_world_size = world_size
+        self.world_size = int(os.getenv("WORLD_SIZE", 1))
         self.cfg_rate = cfg_rate
         self.num_latent_t = num_latent_t
+        self.local_indices = None
+        self.plan_output_dir = os.path.join(self.path, "data_plan.json")
 
-        # Find all parquet files recursively
-        print(f"Scanning for parquet files in {self.path}")
-        self.parquet_files = []
-        for root, _, files in os.walk(self.path):
-            for file in files:
-                if file.endswith('.parquet'):
-                    self.parquet_files.append(os.path.join(root, file))
-        # Sort files for consistent ordering
-        self.parquet_files.sort()
+        ranks = get_sp_group().ranks
+        group_ranks = [None for _ in range(self.world_size)]
+        torch.distributed.all_gather_object(group_ranks, ranks)
 
-        # Distribute files among workers
-        # drop last unenven files
-        print(f"Total files: {len(self.parquet_files)}")
-        total_files = len(self.parquet_files)
-        base_count = total_files // world_size
-        extra_files = total_files % world_size
+        if rank == 0:
+            # If a plan already exists, then skip creating a new plan
+            # This will be useful when resume training
+            if os.path.exists(self.plan_output_dir):
+                print(f"Using existing plan from {self.plan_output_dir}")
+                return
 
-        if rank < extra_files:
-            start_idx = rank * (base_count + 1)
-            end_idx = start_idx + base_count + 1
-        else:
-            start_idx = rank * base_count + extra_files
-            end_idx = start_idx + base_count
+            # Find all parquet files recursively, and record num_rows for each file
+            print(f"Scanning for parquet files in {self.path}")
+            metadatas = []
+            for root, _, files in os.walk(self.path):
+                for file in sorted(files):
+                    if file.endswith('.parquet'):
+                        file_path = os.path.join(root, file)
+                        num_rows = pq.ParquetFile(file_path).metadata.num_rows
+                        for row_idx in range(num_rows):
+                            metadatas.append((file_path, row_idx))
 
-        self.parquet_files = self.parquet_files[start_idx:end_idx]
+            # Generate the plan that distribute rows among workers
+            random.seed(seed)
+            random.shuffle(metadatas)
 
-        print(f"Files assigned to rank {rank}: {len(self.parquet_files)}")
-        if len(self.parquet_files) > 0:
-            print(f"First file: {self.parquet_files[0]}")
-            print(f"Last file: {self.parquet_files[-1]}")
+            # Get all sp groups
+            # e.g. if num_gpus = 4, sp_size = 2
+            # group_ranks = [(0, 1), (2, 3)]
+            # We will assign the same batches of data to ranks in the same sp group, and we'll assign different batches to ranks in different sp groups
+            # e.g. plan = {0: [row 1, row 4], 1: [row 1, row 4], 2: [row 2, row 3], 3: [row 2, row 3]}
+            group_ranks = list(set(tuple(r) for r in group_ranks))
+            num_sp_groups = len(group_ranks)
+            plan = defaultdict(list)
+            for idx, metadata in enumerate(metadatas):
+                sp_group_idx = idx % num_sp_groups
+                for global_rank in group_ranks[sp_group_idx]:
+                    plan[global_rank].append(metadata)
 
-        # Initialize current file index
-        self.current_file_idx = 0
-        self.current_reader = None
-        self.current_batches = None
-        self.total_samples = 0
+            with open(self.plan_output_dir, "w") as f:
+                json.dump(plan, f)
 
-    def _open_next_file(self):
-        """Open the next parquet file for reading."""
-        num_workers = get_worker_info().num_workers
-        worker_id = get_worker_info().id
-        total_files = len(self.parquet_files)
-        base_count = total_files // num_workers
-        extra_files = total_files % num_workers
-
-        if worker_id < extra_files:
-            start_idx = worker_id * (base_count + 1)
-            end_idx = start_idx + base_count + 1
-        else:
-            start_idx = worker_id * base_count + extra_files
-            end_idx = start_idx + base_count
-
-        worker_parquet_files = self.parquet_files[start_idx:end_idx]
-        if self.current_file_idx >= len(worker_parquet_files):
-            print(
-                f"Rank {self.rank}, Worker {worker_id}: No more files to open (current_idx={self.current_file_idx}, total_files={len(worker_parquet_files)})"
-            )
-            return False
-
-        if self.current_reader is not None:
-            self.current_reader.close()
-
-        file_path = worker_parquet_files[self.current_file_idx]
-        print(
-            f"Rank {self.rank}, Worker {worker_id}: Opening file {self.current_file_idx + 1}/{len(worker_parquet_files)}: {file_path}"
-        )
-
-        try:
-            self.current_reader = pq.ParquetFile(file_path)
-            self.current_batches = self.current_reader.iter_batches(
-                batch_size=self.batch_size)
-            self.current_file_idx += 1
-            return True
-        except Exception as e:
-            print(f"Error opening file {file_path}: {str(e)}")
-            return False
-
-    def __iter__(self):
-        """Iterate over the dataset in a streaming fashion."""
-        print(f"Rank {self.rank}: Starting iteration")
-
-        # First try to open a file
-        if not self._open_next_file():
-            print(f"Rank {self.rank}: Failed to open first file")
-            return
-
-        while True:
+    def __len__(self):
+        if self.local_indices is None:
             try:
-                # Get next batch from current file
-                batch = next(self.current_batches)
-                batch_dict = batch.to_pydict()
-                processed = self._process_batch(batch_dict)
+                with open(self.plan_output_dir) as f:
+                    plan = json.load(f)
+                self.local_indices = plan[str(self.rank)]
+            except:
+                raise Exception("The data plan hasn't been created yet")
+        return len(self.local_indices)
 
-                # Update sample count
-                batch_size = len(processed["latents"])
-                self.total_samples += batch_size
+    def __getitem__(self, idx):
+        if self.local_indices is None:
+            try:
+                with open(self.plan_output_dir) as f:
+                    plan = json.load(f)
+                self.local_indices = plan[self.rank]
+            except:
+                raise Exception("The data plan hasn't been created yet")
+        file_path, row_idx = self.local_indices[idx]
+        parquet_file = pq.ParquetFile(file_path)
 
-                # Print progress
-                if self.total_samples % 1000 == 0:
-                    print(
-                        f"Rank {self.rank}: Processed {self.total_samples} samples"
-                    )
+        # Calculate the row group to read into memory and the local idx
+        # This way we can avoid reading in the entire parquet file
+        cumulative = 0
+        for i in range(parquet_file.num_row_groups):
+            num_rows = parquet_file.metadata.row_group(i).num_rows
+            if cumulative + num_rows > idx:
+                row_group_index = i
+                local_index = idx - cumulative
+                break
+            cumulative += num_rows
 
-                # Yield each item in the batch
-                for lat, emb, mask, info in zip(processed["latents"],
-                                                processed["embeddings"],
-                                                processed["masks"],
-                                                processed["info"]):
-                    if lat.numel() == 0:  # Split is validation
-                        yield lat, emb, mask, info
-                    else:
-                        yield lat[:, -self.num_latent_t:], emb, mask, info
+        row_group = parquet_file.read_row_group(row_group_index).to_pydict()
+        row_dict = {k: v[local_index] for k, v in row_group.items()}
+        del row_group
 
-            except StopIteration:
-                # Current file is exhausted, try next file
-                print(
-                    f"Rank {self.rank}: Current file exhausted, trying next file"
-                )
-                self.current_batches = None
-                if not self._open_next_file():
-                    print(
-                        f"Rank {self.rank}: No more files to process. Total samples: {self.total_samples}"
-                    )
-                    break
-            except Exception as e:
-                print(f"Error processing batch: {str(e)}")
-                self.current_batches = None
-                if not self._open_next_file():
-                    print(
-                        f"Rank {self.rank}: Failed to open next file after error"
-                    )
-                    break
+        processed = self._process_row(row_dict)
+        lat, emb, mask, info = processed["latents"], processed[
+            "embeddings"], processed["masks"], processed["info"]
+        if lat.numel() == 0:  # Validation parquet
+            return lat, emb, mask, info
+        else:
+            lat = lat[:, -self.num_latent_t:]
+            if self.sp_world_size > 1:
+                lat = rearrange(lat,
+                                "t (n s) h w -> t n s h w",
+                                n=self.sp_world_size).contiguous()
+                lat = lat[:, self.local_rank, :, :, :]
+            return lat, emb, mask, info
 
-        # Clean up
-        if self.current_reader is not None:
-            self.current_reader.close()
-
-    def _process_batch(self, batch):
+    def _process_row(self, row):
         """Process a PyArrow batch into tensors."""
-        out = {"lat": [], "emb": [], "msk": [], "info": []}
+        out = {"lat": None, "emb": None, "msk": None, "info": None}
 
-        for i in range(len(batch["vae_latent_bytes"])):
-            vae_latent_bytes = batch["vae_latent_bytes"][i]
-            vae_latent_shape = batch["vae_latent_shape"][i]
-            text_embedding_bytes = batch["text_embedding_bytes"][i]
-            text_embedding_shape = batch["text_embedding_shape"][i]
-            text_attention_mask_bytes = batch["text_attention_mask_bytes"][i]
-            text_attention_mask_shape = batch["text_attention_mask_shape"][i]
+        vae_latent_bytes = row["vae_latent_bytes"]
+        vae_latent_shape = row["vae_latent_shape"]
+        text_embedding_bytes = row["text_embedding_bytes"]
+        text_embedding_shape = row["text_embedding_shape"]
+        text_attention_mask_bytes = row["text_attention_mask_bytes"]
+        text_attention_mask_shape = row["text_attention_mask_shape"]
 
-            # Process latent
-            if not vae_latent_shape:  # No VAE latent is stored. Split is validation
-                lat = np.array([])
-            else:
-                lat = np.frombuffer(vae_latent_bytes,
-                                    dtype=np.float32).reshape(vae_latent_shape)
-                # Make array writable
-                lat = np.copy(lat)
+        # Process latent
+        if not vae_latent_shape:  # No VAE latent is stored. Split is validation
+            lat = np.array([])
+        else:
+            lat = np.frombuffer(vae_latent_bytes,
+                                dtype=np.float32).reshape(vae_latent_shape)
+            # Make array writable
+            lat = np.copy(lat)
 
-            if random.random() < self.cfg_rate:
-                emb = np.zeros((512, 4096), dtype=np.float32)
-            else:
-                emb = np.frombuffer(
-                    text_embedding_bytes,
-                    dtype=np.float32).reshape(text_embedding_shape)
-                # Make array writable
-                emb = np.copy(emb)
-            if emb.shape[0] < 512:
-                padded_emb = np.zeros((512, emb.shape[1]), dtype=np.float32)
-                padded_emb[:emb.shape[0], :] = emb
-                emb = padded_emb
-            elif emb.shape[0] > 512:
-                emb = emb[:512, :]
+        if random.random() < self.cfg_rate:
+            emb = np.zeros((512, 4096), dtype=np.float32)
+        else:
+            emb = np.frombuffer(text_embedding_bytes,
+                                dtype=np.float32).reshape(text_embedding_shape)
+            # Make array writable
+            emb = np.copy(emb)
+        if emb.shape[0] < 512:
+            padded_emb = np.zeros((512, emb.shape[1]), dtype=np.float32)
+            padded_emb[:emb.shape[0], :] = emb
+            emb = padded_emb
+        elif emb.shape[0] > 512:
+            emb = emb[:512, :]
 
-            # Process mask
-            if len(text_attention_mask_bytes) > 0 and len(
-                    text_attention_mask_shape) > 0:
-                msk = np.frombuffer(text_attention_mask_bytes,
-                                    dtype=np.uint8).astype(np.bool_)
-                msk = msk.reshape(1, -1)
-                # Make array writable
-                msk = np.copy(msk)
-                if msk.shape[1] < 512:
-                    padded_msk = np.zeros((1, 512), dtype=np.bool_)
-                    padded_msk[:, :msk.shape[1]] = msk
-                    msk = padded_msk
-                elif msk.shape[1] > 512:
-                    msk = msk[:, :512]
-            else:
-                msk = np.ones((1, 512), dtype=np.bool_)
-            # to string
-            file_name = str(batch["file_name"][i])
-            # Collect metadata
-            info = {
-                "width": batch["width"][i],
-                "height": batch["height"][i],
-                "num_frames": batch["num_frames"][i],
-                "duration_sec": batch["duration_sec"][i],
-                "fps": batch["fps"][i],
-                "file_name": batch["file_name"][i],
-                "caption": batch["caption"][i],
-            }
+        # Process mask
+        if len(text_attention_mask_bytes) > 0 and len(
+                text_attention_mask_shape) > 0:
+            msk = np.frombuffer(text_attention_mask_bytes,
+                                dtype=np.uint8).astype(np.bool_)
+            msk = msk.reshape(1, -1)
+            # Make array writable
+            msk = np.copy(msk)
+            if msk.shape[1] < 512:
+                padded_msk = np.zeros((1, 512), dtype=np.bool_)
+                padded_msk[:, :msk.shape[1]] = msk
+                msk = padded_msk
+            elif msk.shape[1] > 512:
+                msk = msk[:, :512]
+        else:
+            msk = np.ones((1, 512), dtype=np.bool_)
 
-            out["lat"].append(torch.from_numpy(lat))
-            out["emb"].append(torch.from_numpy(emb))
-            out["msk"].append(torch.from_numpy(msk))
-            out["info"].append(info)
-
-        return {
-            "latents": torch.stack(out["lat"]) if out["lat"] else None,
-            "embeddings": torch.stack(out["emb"]) if out["emb"] else None,
-            "masks": torch.stack(out["msk"]) if out["msk"] else None,
-            "info": out["info"]
+        # Collect metadata
+        info = {
+            "width": row["width"],
+            "height": row["height"],
+            "num_frames": row["num_frames"],
+            "duration_sec": row["duration_sec"],
+            "fps": row["fps"],
+            "file_name": row["file_name"],
+            "caption": row["caption"],
         }
 
+        out["lat"] = torch.from_numpy(lat)
+        out["emb"] = torch.from_numpy(emb)
+        out["msk"] = torch.from_numpy(msk)
+        out["info"] = info
 
-def bind_cpu_cores(local_rank, cpu_per_process=16):
-    """根据local_rank绑定固定cpu核。"""
-    start = local_rank * cpu_per_process
-    end = start + cpu_per_process
-    cores = list(range(start, end))
-    print(f"[Rank {local_rank}] Binding to CPU cores: {cores}")
-    os.sched_setaffinity(0, cores)
+        return {
+            "latents": out["lat"],
+            "embeddings": out["emb"],
+            "masks": out["msk"],
+            "info": out["info"]
+        }
 
 
 if __name__ == "__main__":
@@ -297,9 +253,6 @@ if __name__ == "__main__":
             f"Initialized process: rank={rank}, local_rank={local_rank}, world_size={world_size}, device={device}"
         )
 
-    # Bind CPU cores after distributed initialization
-    # bind_cpu_cores(local_rank, cpu_per_process=16)
-
     # Create dataset
     dataset = ParquetVideoTextDataset(
         args.path,

fastvideo/v1/pipelines/preprocess_pipeline.py

@@ -140,7 +140,6 @@ class PreprocessPipeline(ComposedPipelineBase):
                     0], result_batch.prompt_attention_mask[0]
                 assert prompt_embeds.shape[0] == prompt_attention_mask.shape[0]
 
-                # Remove padding from prompt_embeds using attention mask for all batches
                 # Get sequence lengths from attention masks (number of 1s)
                 seq_lens = prompt_attention_mask.sum(dim=1)
                 # Create a list to store non-padded embeddings and masks
@@ -354,9 +353,6 @@ class PreprocessPipeline(ComposedPipelineBase):
                                               "validation_parquet_dataset")
         os.makedirs(validation_parquet_dir, exist_ok=True)
 
-        # Initialize Parquet dataset
-        validation_parquet_path = os.path.join(validation_parquet_dir,
-                                               "data.parquet")
 
         with open(args.validation_prompt_txt, encoding="utf-8") as file:
             lines = file.readlines()