Adding dataset files for Wan

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

Author: pthombre

dfm/examples/Automodel/finetune/wan2_1_t2v_flow.yaml

@@ -14,7 +14,7 @@ model:
 
 data:
   dataloader:
-    _target_: nemo_automodel.components.datasets.diffusion.build_wan21_dataloader
+    _target_: Automodel.datasets.build_wan21_dataloader
     meta_folder: /lustre/fsw/portfolios/coreai/users/linnanw/hdvilla_sample/pika/wan21_codes/1.3B_meta/
     batch_size: 1
     num_workers: 2
@@ -50,12 +50,12 @@ fsdp:
   dp_size: 8
 
 logging:
-  save_every: 50
+  save_every: 1000
   log_every: 2
 
 checkpoint:
   enabled: true
-  checkpoint_dir: /opt/DFM/wan_t2v_flow_outputs_base_recipe_dfm/
+  checkpoint_dir: /opt/DFM/wan_t2v_flow_outputs_base_recipe_dfm_test_2/
   model_save_format: torch_save
   save_consolidated: false
   restore_from: null

dfm/src/Automodel/_diffusers/auto_diffusion_pipeline.py

@@ -116,7 +116,7 @@ def from_pretrained(
         load_for_training: bool = False,
         components_to_load: Optional[Iterable[str]] = None,
         **kwargs,
-    ) -> DiffusionPipeline:
+    ) -> tuple[DiffusionPipeline, Dict[str, FSDP2Manager]]:
         pipe: DiffusionPipeline = DiffusionPipeline.from_pretrained(
             pretrained_model_name_or_path,
             *model_args,

dfm/src/Automodel/datasets/__init__.py

@@ -0,0 +1,30 @@
+# Copyright (c) 2025, 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.
+
+from Automodel.datasets.wan21 import (
+    MetaFilesDataset,
+    build_node_parallel_sampler,
+    build_wan21_dataloader,
+    collate_fn,
+    create_dataloader,
+)
+
+
+__all__ = [
+    "MetaFilesDataset",
+    "build_node_parallel_sampler",
+    "build_wan21_dataloader",
+    "collate_fn",
+    "create_dataloader",
+]

dfm/src/Automodel/datasets/wan21.py

@@ -0,0 +1,217 @@
+# Copyright (c) 2025, 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.
+
+from __future__ import annotations
+
+import logging
+import os
+import pickle
+from pathlib import Path
+from typing import Callable, Dict, List, Optional, Tuple
+
+import torch
+import torch.distributed as dist
+from torch.utils.data import DataLoader, Dataset, DistributedSampler
+
+
+logger = logging.getLogger(__name__)
+
+
+class MetaFilesDataset(Dataset):
+    """PyTorch dataset for WAN2.1 `.meta` files."""
+
+    def __init__(
+        self,
+        meta_folder: str,
+        transform_text: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+        transform_video: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+        filter_fn: Optional[Callable[[Dict], bool]] = None,
+        device: str = "cpu",
+        max_files: Optional[int] = None,
+    ) -> None:
+        self.meta_folder = Path(meta_folder)
+        self.transform_text = transform_text
+        self.transform_video = transform_video
+        self.filter_fn = filter_fn
+        self.device = device
+
+        self.meta_files = sorted(self.meta_folder.glob("*.meta"))
+        if max_files is None:
+            max_files_env = os.environ.get("MAX_META_FILES")
+            if max_files_env is not None:
+                try:
+                    max_files = int(max_files_env)
+                except ValueError:
+                    logger.warning("Invalid MAX_META_FILES=%s", max_files_env)
+
+        if max_files is not None and max_files > 0:
+            self.meta_files = self.meta_files[:max_files]
+            logger.info("Limited to first %d meta files", len(self.meta_files))
+
+        if not self.meta_files:
+            raise ValueError(f"No .meta files found in {meta_folder}")
+
+        if self.filter_fn:
+            filtered = []
+            for path in self.meta_files:
+                try:
+                    with open(path, "rb") as f:
+                        data = pickle.load(f)
+                except Exception as exc:  # pragma: no cover - best effort logging
+                    logger.warning("Failed to load %s during filtering: %s", path, exc)
+                    continue
+                if self.filter_fn(data.get("metadata", {})):
+                    filtered.append(path)
+            self.meta_files = filtered
+            logger.info("Filtered meta files count: %d", len(self.meta_files))
+
+        self._log_dataset_stats()
+
+    def _log_dataset_stats(self) -> None:
+        sample_paths = self.meta_files[: min(5, len(self.meta_files))]
+        stats: List[Tuple[torch.Size, torch.Size, str]] = []
+        for path in sample_paths:
+            try:
+                with open(path, "rb") as f:
+                    data = pickle.load(f)
+                stats.append(
+                    (
+                        data["text_embeddings"].shape,
+                        data["video_latents"].shape,
+                        str(data.get("deterministic_latents", "unknown")),
+                    )
+                )
+            except Exception as exc:  # pragma: no cover - stats only
+                logger.debug("Failed to sample %s: %s", path, exc)
+
+        if stats:
+            text_shapes, video_shapes, modes = zip(*stats, strict=False)
+            logger.info("Sample text embeddings: %s", text_shapes)
+            logger.info("Sample video latents: %s", video_shapes)
+            logger.info("Sample encoding modes: %s", set(modes))
+
+    def __len__(self) -> int:
+        return len(self.meta_files)
+
+    def __getitem__(self, index: int) -> Dict[str, torch.Tensor]:  # type: ignore[override]
+        path = self.meta_files[index]
+        with open(path, "rb") as f:
+            data = pickle.load(f)
+
+        text_embeddings: torch.Tensor = data["text_embeddings"].to(self.device)
+        video_latents: torch.Tensor = data["video_latents"].to(self.device)
+
+        if self.transform_text is not None:
+            text_embeddings = self.transform_text(text_embeddings)
+        if self.transform_video is not None:
+            video_latents = self.transform_video(video_latents)
+
+        file_info = {
+            "meta_filename": Path(path).name,
+            "original_filename": data.get("original_filename", "unknown"),
+            "original_video_path": data.get("original_video_path", "unknown"),
+            "deterministic_latents": data.get("deterministic_latents", "unknown"),
+            "memory_optimization": data.get("memory_optimization", "unknown"),
+            "num_frames": data.get("num_frames", "unknown"),
+        }
+
+        return {
+            "text_embeddings": text_embeddings,
+            "video_latents": video_latents,
+            "metadata": data.get("metadata", {}),
+            "file_info": file_info,
+        }
+
+
+def collate_fn(batch: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
+    text_embeddings = torch.stack([item["text_embeddings"] for item in batch])
+    video_latents = torch.stack([item["video_latents"] for item in batch])
+    return {
+        "text_embeddings": text_embeddings,
+        "video_latents": video_latents,
+        "metadata": [item["metadata"] for item in batch],
+        "file_info": [item["file_info"] for item in batch],
+    }
+
+
+def build_node_parallel_sampler(
+    dataset: Dataset,
+    num_nodes: Optional[int] = None,
+    shuffle: bool = True,
+) -> Optional[DistributedSampler]:
+    if not dist.is_initialized():
+        return None
+
+    world_size = dist.get_world_size()
+    local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", world_size))
+    local_world_size = max(local_world_size, 1)
+    if num_nodes is None:
+        num_nodes = max(1, world_size // local_world_size)
+
+    node_rank = dist.get_rank() // local_world_size
+    replicas = num_nodes
+
+    return DistributedSampler(
+        dataset,
+        num_replicas=replicas,
+        rank=node_rank,
+        shuffle=shuffle,
+        drop_last=False,
+    )
+
+
+def build_wan21_dataloader(
+    *,
+    meta_folder: str,
+    batch_size: int,
+    shuffle: bool = True,
+    num_workers: int = 2,
+    device: str = "cpu",
+    transform_text: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+    transform_video: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
+    filter_fn: Optional[Callable[[Dict], bool]] = None,
+    max_files: Optional[int] = None,
+    num_nodes: Optional[int] = None,
+) -> Tuple[DataLoader, Optional[DistributedSampler]]:
+    dataset = MetaFilesDataset(
+        meta_folder=meta_folder,
+        transform_text=transform_text,
+        transform_video=transform_video,
+        filter_fn=filter_fn,
+        device=device,
+        max_files=max_files,
+    )
+
+    sampler = build_node_parallel_sampler(dataset, num_nodes, shuffle=shuffle)
+
+    use_pin_memory = device == "cpu"
+    dataloader = DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=(sampler is None and shuffle),
+        sampler=sampler,
+        num_workers=num_workers,
+        collate_fn=collate_fn,
+        pin_memory=use_pin_memory,
+    )
+
+    return dataloader, sampler
+
+
+def create_dataloader(
+    meta_folder: str,
+    batch_size: int,
+    num_nodes: int,
+) -> Tuple[DataLoader, Optional[DistributedSampler]]:
+    return build_wan21_dataloader(meta_folder=meta_folder, batch_size=batch_size, num_nodes=num_nodes)

dfm/src/Automodel/recipes/finetune.py

@@ -21,7 +21,6 @@
 
 import torch
 import torch.distributed as dist
-import wandb
 from Automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline
 from Automodel.flow_matching.training_step_t2v import (
     step_fsdp_transformer_t2v,
@@ -36,6 +35,8 @@
 from torch.distributed.fsdp import MixedPrecisionPolicy
 from transformers.utils.hub import TRANSFORMERS_CACHE
 
+import wandb
+
 
 def build_model_and_optimizer(
     *,
@@ -462,10 +463,6 @@ def run_train_validation_loop(self):
             if is_main_process() and wandb.run is not None:
                 wandb.log({"epoch/avg_loss": avg_loss, "epoch/num": epoch + 1}, step=global_step)
 
-        logging.info("[INFO] Training complete, saving final checkpoint...")
-
-        self.save_checkpoint(epoch=self.step_scheduler.epoch, step=global_step)
-
         if is_main_process():
             logging.info(f"[INFO] Saved final checkpoint at step {global_step}")
             if wandb.run is not None: