update (#5)

training/args.py

@@ -96,6 +96,11 @@ def _get_dataset_args(parser: argparse.ArgumentParser) -> None:
         default=0,
         help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
     )
+    parser.add_argument(
+        "--pin_memory",
+        action="store_true",
+        help="Whether or not to use the pinned memory setting in pytorch dataloader.",
+    )
 
 
 def _get_validation_args(parser: argparse.ArgumentParser) -> None:

training/cogvideox_text_to_video_lora.py

@@ -25,6 +25,7 @@ from typing import Any, Dict
 import diffusers
 import torch
 import transformers
+import wandb
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import (
@@ -51,8 +52,6 @@ from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 from transformers import AutoTokenizer, T5EncoderModel
 
-import wandb
-
 
 from args import get_args  # isort:skip
 from dataset import BucketSampler, VideoDatasetWithResizing  # isort:skip
@@ -475,32 +474,14 @@ def main(args):
         random_flip=args.random_flip,
     )
 
-    def collate_fn_without_pre_encoding(data):
+    def collate_fn(data):
         prompts = [x["prompt"] for x in data[0]]
 
-        videos = [x["video"] for x in data[0]]
-        videos = torch.stack(videos)
-        videos = videos.to(accelerator.device, dtype=weight_dtype)
-        videos = videos.permute(0, 2, 1, 3, 4)  # [B, C, F, H, W]
-        latent_dist = vae.encode(videos).latent_dist
-        videos = latent_dist.sample() * VAE_SCALING_FACTOR
-        videos = videos.permute(0, 2, 1, 3, 4)  # [B, F, C, H, W]
-        videos = videos.to(memory_format=torch.contiguous_format).float()
-
-        return {
-            "videos": videos,
-            "prompts": prompts,
-        }
-
-    def collate_fn_with_pre_encoding(data):
-        prompts = [x["prompt"] for x in data[0]]
-        prompts = torch.stack(prompts).to(accelerator.device, dtype=weight_dtype)
+        if args.load_tensors:
+            prompts = torch.stack(prompts).to(dtype=weight_dtype, non_blocking=True)
 
         videos = [x["video"] for x in data[0]]
-        videos = torch.stack(videos).to(accelerator.device, dtype=weight_dtype)
-        videos = DiagonalGaussianDistribution(videos).sample() * VAE_SCALING_FACTOR
-        videos = videos.permute(0, 2, 1, 3, 4)  # [B, F, C, H, W]
-        videos = videos.to(memory_format=torch.contiguous_format).float()
+        videos = torch.stack(videos).to(dtype=weight_dtype, non_blocking=True)
 
         return {
             "videos": videos,
@@ -511,8 +492,9 @@ def main(args):
         train_dataset,
         batch_size=1,
         sampler=BucketSampler(train_dataset, batch_size=args.train_batch_size, shuffle=True),
-        collate_fn=collate_fn_with_pre_encoding if args.load_tensors else collate_fn_without_pre_encoding,
+        collate_fn=collate_fn,
         num_workers=args.dataloader_num_workers,
+        pin_memory=args.pin_memory,
     )
 
     # Scheduler and math around the number of training steps.
@@ -637,9 +619,21 @@ def main(args):
             models_to_accumulate = [transformer]
 
             with accelerator.accumulate(models_to_accumulate):
-                model_input = batch["videos"]
+                videos = batch["videos"].to(accelerator.device, non_blocking=True)
                 prompts = batch["prompts"]
 
+                # Encode videos
+                if not args.load_tensors:
+                    videos = videos.permute(0, 2, 1, 3, 4)  # [B, C, F, H, W]
+                    latent_dist = vae.encode(videos).latent_dist
+                else:
+                    latent_dist = DiagonalGaussianDistribution(videos)
+
+                videos = latent_dist.sample() * VAE_SCALING_FACTOR
+                videos = videos.permute(0, 2, 1, 3, 4)  # [B, F, C, H, W]
+                videos = videos.to(memory_format=torch.contiguous_format).float()
+                model_input = videos
+
                 # Encode prompts
                 if not args.load_tensors:
                     prompt_embeds = compute_prompt_embeddings(

training/cogvideox_text_to_video_sft.py

@@ -25,6 +25,7 @@ from typing import Any, Dict
 import diffusers
 import torch
 import transformers
+import wandb
 from accelerate import Accelerator
 from accelerate.logging import get_logger
 from accelerate.utils import (
@@ -50,8 +51,6 @@ from torch.utils.data import DataLoader
 from tqdm.auto import tqdm
 from transformers import AutoTokenizer, T5EncoderModel
 
-import wandb
-
 
 from args import get_args  # isort:skip
 from dataset import BucketSampler, VideoDatasetWithResizing  # isort:skip
@@ -406,32 +405,14 @@ def main(args):
         random_flip=args.random_flip,
     )
 
-    def collate_fn_without_pre_encoding(data):
+    def collate_fn(data):
         prompts = [x["prompt"] for x in data[0]]
 
-        videos = [x["video"] for x in data[0]]
-        videos = torch.stack(videos)
-        videos = videos.to(accelerator.device, dtype=weight_dtype)
-        videos = videos.permute(0, 2, 1, 3, 4)  # [B, C, F, H, W]
-        latent_dist = vae.encode(videos).latent_dist
-        videos = latent_dist.sample() * VAE_SCALING_FACTOR
-        videos = videos.permute(0, 2, 1, 3, 4)  # [B, F, C, H, W]
-        videos = videos.to(memory_format=torch.contiguous_format).float()
-
-        return {
-            "videos": videos,
-            "prompts": prompts,
-        }
-
-    def collate_fn_with_pre_encoding(data):
-        prompts = [x["prompt"] for x in data[0]]
-        prompts = torch.stack(prompts).to(accelerator.device, dtype=weight_dtype)
+        if args.load_tensors:
+            prompts = torch.stack(prompts).to(dtype=weight_dtype, non_blocking=True)
 
         videos = [x["video"] for x in data[0]]
-        videos = torch.stack(videos).to(accelerator.device, dtype=weight_dtype)
-        videos = DiagonalGaussianDistribution(videos).sample() * VAE_SCALING_FACTOR
-        videos = videos.permute(0, 2, 1, 3, 4)  # [B, F, C, H, W]
-        videos = videos.to(memory_format=torch.contiguous_format).float()
+        videos = torch.stack(videos).to(dtype=weight_dtype, non_blocking=True)
 
         return {
             "videos": videos,
@@ -442,8 +423,9 @@ def main(args):
         train_dataset,
         batch_size=1,
         sampler=BucketSampler(train_dataset, batch_size=args.train_batch_size, shuffle=True),
-        collate_fn=collate_fn_with_pre_encoding if args.load_tensors else collate_fn_without_pre_encoding,
+        collate_fn=collate_fn,
         num_workers=args.dataloader_num_workers,
+        pin_memory=args.pin_memory,
     )
 
     # Scheduler and math around the number of training steps.
@@ -568,9 +550,21 @@ def main(args):
             models_to_accumulate = [transformer]
 
             with accelerator.accumulate(models_to_accumulate):
-                model_input = batch["videos"]
+                videos = batch["videos"].to(accelerator.device, non_blocking=True)
                 prompts = batch["prompts"]
 
+                # Encode videos
+                if not args.load_tensors:
+                    videos = videos.permute(0, 2, 1, 3, 4)  # [B, C, F, H, W]
+                    latent_dist = vae.encode(videos).latent_dist
+                else:
+                    latent_dist = DiagonalGaussianDistribution(videos)
+
+                videos = latent_dist.sample() * VAE_SCALING_FACTOR
+                videos = videos.permute(0, 2, 1, 3, 4)  # [B, F, C, H, W]
+                videos = videos.to(memory_format=torch.contiguous_format).float()
+                model_input = videos
+
                 # Encode prompts
                 if not args.load_tensors:
                     prompt_embeds = compute_prompt_embeddings(