[Fix] drop file when num_parquet % num_workers != 0

fastvideo/data_preprocess/preprocess.py

@@ -0,0 +1,111 @@
+import argparse
+import json
+import os
+
+import torch
+import torch.distributed as dist
+
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.utils import maybe_download_model, shallow_asdict
+from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.configs.models.vaes import WanVAEConfig
+from fastvideo import PipelineConfig
+from fastvideo.v1.pipelines.data_preprocess import PreprocessPipeline
+
+logger = init_logger(__name__)
+
+BASE_MODEL_PATH = "/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
+                                  local_dir=os.path.join(
+                                      'data', BASE_MODEL_PATH))
+
+def main(args):
+    # Assume using torchrun
+    local_rank = int(os.getenv("RANK", 0))
+    rank = int(os.environ.get("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    init_distributed_environment(world_size=world_size, rank=rank, local_rank=local_rank)
+    initialize_model_parallel(tensor_model_parallel_size=world_size, sequence_model_parallel_size=world_size)
+    torch.cuda.set_device(local_rank)
+    if not dist.is_initialized():
+        dist.init_process_group(backend="nccl", init_method="env://", world_size=world_size, rank=local_rank)
+
+    pipeline_config = PipelineConfig.from_pretrained(MODEL_PATH)
+    kwargs = {
+        "use_cpu_offload": False,
+        "vae_precision": "fp32",
+        "vae_config": WanVAEConfig(load_encoder=True, load_decoder=False),
+    }
+    pipeline_config_args = shallow_asdict(pipeline_config)
+    pipeline_config_args.update(kwargs)
+    fastvideo_args = FastVideoArgs(model_path=MODEL_PATH,
+                                   num_gpus=world_size,
+                                   device_str="cuda",
+                                   **pipeline_config_args,
+                                   )
+    fastvideo_args.check_fastvideo_args()
+    fastvideo_args.device = torch.device(f"cuda:{local_rank}")
+
+    pipeline = PreprocessPipeline(MODEL_PATH, fastvideo_args)
+    pipeline.forward(batch=None, fastvideo_args=fastvideo_args, args=args)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # dataset & dataloader
+    parser.add_argument("--model_path", type=str, default="data/mochi")
+    parser.add_argument("--model_type", type=str, default="mochi")
+    parser.add_argument("--data_merge_path", type=str, required=True)
+    parser.add_argument("--validation_prompt_txt", type=str)
+    parser.add_argument("--num_frames", type=int, default=163)
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=1,
+        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
+    )
+    parser.add_argument(
+        "--preprocess_video_batch_size",
+        type=int,
+        default=2,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument(
+        "--preprocess_text_batch_size",
+        type=int,
+        default=8,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
+    parser.add_argument("--max_height", type=int, default=480)
+    parser.add_argument("--max_width", type=int, default=848)
+    parser.add_argument("--video_length_tolerance_range", type=int, default=2.0)
+    parser.add_argument("--group_frame", action="store_true")  # TODO
+    parser.add_argument("--group_resolution", action="store_true")  # TODO
+    parser.add_argument("--dataset", default="t2v")
+    parser.add_argument("--train_fps", type=int, default=30)
+    parser.add_argument("--use_image_num", type=int, default=0)
+    parser.add_argument("--text_max_length", type=int, default=256)
+    parser.add_argument("--speed_factor", type=float, default=1.0)
+    parser.add_argument("--drop_short_ratio", type=float, default=1.0)
+    # text encoder & vae & diffusion model
+    parser.add_argument("--text_encoder_name", type=str, default="google/t5-v1_1-xxl")
+    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
+    parser.add_argument("--cfg", type=float, default=0.0)
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
+    )
+
+    args = parser.parse_args()
+    main(args)

fastvideo/distill_wan.py

@@ -0,0 +1,752 @@
+# !/bin/python3
+# isort: skip_file
+import argparse
+import math
+import os
+import time
+from collections import deque
+from copy import deepcopy
+
+import torch
+import torch.distributed as dist
+import wandb
+from accelerate.utils import set_seed
+from diffusers.optimization import get_scheduler
+from diffusers.utils import check_min_version
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import ShardingStrategy
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+from tqdm.auto import tqdm
+
+from fastvideo.dataset.latent_datasets import (LatentDataset, latent_collate_function)
+from fastvideo.distill.solver import EulerSolver, extract_into_tensor
+from fastvideo.utils.checkpoint import (save_checkpoint, save_lora_checkpoint)
+from fastvideo.utils.communications import (broadcast, sp_parallel_dataloader_wrapper)
+from fastvideo.utils.dataset_utils import LengthGroupedSampler
+from fastvideo.utils.parallel_states import (destroy_sequence_parallel_group, get_sequence_parallel_state)
+from fastvideo.utils.validation import log_validation
+
+from fastvideo.v1.utils import maybe_download_model
+from fastvideo.v1.configs.models.dits import WanVideoConfig
+from fastvideo.v1.models.loader.component_loader import TransformerLoader, SchedulerLoader
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel
+from fastvideo.v1.forward_context import set_forward_context
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+
+# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
+check_min_version("0.31.0")
+
+logger = init_logger(__name__)
+
+BASE_MODEL_PATH = "/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+MODEL_PATH = maybe_download_model(BASE_MODEL_PATH,
+                                  local_dir=os.path.join(
+                                      'data', BASE_MODEL_PATH))
+TRANSFORMER_PATH = os.path.join(MODEL_PATH, "transformer")
+SCHEDULER_PATH = os.path.join(MODEL_PATH, "scheduler")
+
+
+def reshard_fsdp(model):
+    for m in FSDP.fsdp_modules(model):
+        if m._has_params and m.sharding_strategy is not ShardingStrategy.NO_SHARD:
+            torch.distributed.fsdp._runtime_utils._reshard(m, m._handle, True)
+
+
+def get_norm(model_pred, norms, gradient_accumulation_steps):
+    fro_norm = (
+        torch.linalg.matrix_norm(model_pred, ord="fro") /  # codespell:ignore
+        gradient_accumulation_steps)
+    largest_singular_value = (torch.linalg.matrix_norm(model_pred, ord=2) / gradient_accumulation_steps)
+    absolute_mean = torch.mean(torch.abs(model_pred)) / gradient_accumulation_steps
+    absolute_max = torch.max(torch.abs(model_pred)) / gradient_accumulation_steps
+    dist.all_reduce(fro_norm, op=dist.ReduceOp.AVG)
+    dist.all_reduce(largest_singular_value, op=dist.ReduceOp.AVG)
+    dist.all_reduce(absolute_mean, op=dist.ReduceOp.AVG)
+    norms["fro"] += torch.mean(fro_norm).item()  # codespell:ignore
+    norms["largest singular value"] += torch.mean(largest_singular_value).item()
+    norms["absolute mean"] += absolute_mean.item()
+    norms["absolute max"] += absolute_max.item()
+
+
+def distill_one_step(
+    transformer,
+    model_type,
+    teacher_transformer,
+    ema_transformer,
+    optimizer,
+    lr_scheduler,
+    loader,
+    noise_scheduler,
+    solver,
+    noise_random_generator,
+    gradient_accumulation_steps,
+    sp_size,
+    max_grad_norm,
+    uncond_prompt_embed,
+    uncond_prompt_mask,
+    num_euler_timesteps,
+    multiphase,
+    not_apply_cfg_solver,
+    distill_cfg,
+    ema_decay,
+    pred_decay_weight,
+    pred_decay_type,
+    hunyuan_teacher_disable_cfg,
+):
+    total_loss = 0.0
+    optimizer.zero_grad()
+    model_pred_norm = {
+        "fro": 0.0,  # codespell:ignore
+        "largest singular value": 0.0,
+        "absolute mean": 0.0,
+        "absolute max": 0.0,
+    }
+    for _ in range(gradient_accumulation_steps):
+        (
+            latents,
+            encoder_hidden_states,
+            latents_attention_mask,
+            encoder_attention_mask,
+        ) = next(loader)
+        # model_input = normalize_dit_input(model_type, latents)
+        model_input = latents
+        noise = torch.randn_like(model_input)
+        bsz = model_input.shape[0]
+        index = torch.randint(0, num_euler_timesteps, (bsz, ), device=model_input.device).long()
+        if sp_size > 1:
+            broadcast(index)
+        # Add noise according to flow matching.
+        # sigmas = get_sigmas(start_timesteps, n_dim=model_input.ndim, dtype=model_input.dtype)
+        sigmas = extract_into_tensor(solver.sigmas, index, model_input.shape)
+        sigmas_prev = extract_into_tensor(solver.sigmas_prev, index, model_input.shape)
+
+        timesteps = (sigmas * noise_scheduler.config.num_train_timesteps).view(-1)
+        # if squeeze to [], unsqueeze to [1]
+
+        timesteps_prev = (sigmas_prev * noise_scheduler.config.num_train_timesteps).view(-1)
+        noisy_model_input = sigmas * noise + (1.0 - sigmas) * model_input
+        noisy_model_input = noisy_model_input.to(torch.bfloat16)
+
+        forward_batch = ForwardBatch(data_type="video", enable_teacache=False)
+        # Predict the noise residual
+        with torch.autocast("cuda", dtype=torch.bfloat16):
+            teacher_kwargs = {
+                "hidden_states": noisy_model_input,
+                "encoder_hidden_states": encoder_hidden_states,
+                "timestep": timesteps,
+                "encoder_attention_mask": encoder_attention_mask,  # B, L
+                "return_dict": False,
+            }
+            if hunyuan_teacher_disable_cfg:
+                teacher_kwargs["guidance"] = torch.tensor([1000.0],
+                                                          device=noisy_model_input.device,
+                                                          dtype=torch.bfloat16)
+            with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch):
+                with torch.autograd.graph.save_on_cpu(pin_memory=True):
+                    model_pred = transformer(**teacher_kwargs)
+
+        # if accelerator.is_main_process:
+        model_pred, end_index = solver.euler_style_multiphase_pred(noisy_model_input, model_pred, index, multiphase)
+        with torch.no_grad():
+            w = distill_cfg
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch):
+                    cond_teacher_output = teacher_transformer(
+                        noisy_model_input,
+                        encoder_hidden_states,
+                        timesteps,
+                        encoder_attention_mask,  # B, L
+                        return_dict=False,
+                    ).float()
+            if not_apply_cfg_solver:
+                uncond_teacher_output = cond_teacher_output
+            else:
+                # Get teacher model prediction on noisy_latents and unconditional embedding
+                with torch.autocast("cuda", dtype=torch.bfloat16):
+                    with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch):
+                        uncond_teacher_output = teacher_transformer(
+                            noisy_model_input,
+                            uncond_prompt_embed.unsqueeze(0).expand(bsz, -1, -1),
+                            timesteps,
+                            uncond_prompt_mask.unsqueeze(0).expand(bsz, -1),
+                            return_dict=False,
+                        ).float()
+            teacher_output = uncond_teacher_output + w * (cond_teacher_output - uncond_teacher_output)
+            x_prev = solver.euler_step(noisy_model_input, teacher_output, index).to(torch.bfloat16)
+
+        # 20.4.12. Get target LCM prediction on x_prev, w, c, t_n
+        with torch.no_grad():
+            with torch.autocast("cuda", dtype=torch.bfloat16):
+                if ema_transformer is not None:
+                    target_pred = ema_transformer(
+                        x_prev.float(),
+                        encoder_hidden_states,
+                        timesteps_prev,
+                        encoder_attention_mask,  # B, L
+                        return_dict=False,
+                    )[0]
+                else:
+                    with set_forward_context(current_timestep=0, attn_metadata=None, forward_batch=forward_batch):
+                        with torch.autograd.graph.save_on_cpu(pin_memory=True):
+                            target_pred = transformer(
+                                x_prev,
+                                encoder_hidden_states,
+                                timesteps_prev,
+                                encoder_attention_mask,  # B, L
+                                return_dict=False,
+                            )
+
+            target, end_index = solver.euler_style_multiphase_pred(x_prev, target_pred, index, multiphase, True)
+
+        huber_c = 0.001
+        # loss = loss.mean()
+        loss = (torch.mean(torch.sqrt((model_pred.float() - target.float())**2 + huber_c**2) - huber_c) /
+                gradient_accumulation_steps)
+        if pred_decay_weight > 0:
+            if pred_decay_type == "l1":
+                pred_decay_loss = (torch.mean(torch.sqrt(model_pred.float()**2)) * pred_decay_weight /
+                                   gradient_accumulation_steps)
+                loss += pred_decay_loss
+            elif pred_decay_type == "l2":
+                # essnetially k2?
+                pred_decay_loss = (torch.mean(model_pred.float()**2) * pred_decay_weight / gradient_accumulation_steps)
+                loss += pred_decay_loss
+            else:
+                assert NotImplementedError("pred_decay_type is not implemented")
+
+        # calculate model_pred norm and mean
+        get_norm(model_pred.detach().float(), model_pred_norm, gradient_accumulation_steps)
+        loss.backward()
+
+        avg_loss = loss.detach().clone()
+        dist.all_reduce(avg_loss, op=dist.ReduceOp.AVG)
+        total_loss += avg_loss.item()
+
+    # update ema
+    if ema_transformer is not None:
+        reshard_fsdp(ema_transformer)
+        for p_averaged, p_model in zip(ema_transformer.parameters(), transformer.parameters()):
+            with torch.no_grad():
+                p_averaged.copy_(torch.lerp(p_averaged.detach(), p_model.detach(), 1 - ema_decay))
+
+    # grad_norm = transformer.clip_grad_norm_(max_grad_norm)
+    grad_norm = torch.nn.utils.clip_grad_norm_(transformer.parameters(), max_norm=max_grad_norm)
+    optimizer.step()
+    lr_scheduler.step()
+
+    return total_loss, grad_norm.item(), model_pred_norm
+
+
+def main(args):
+    torch.backends.cuda.matmul.allow_tf32 = True
+
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    rank = int(os.environ.get("RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    torch.cuda.set_device(rank)
+    init_distributed_environment(world_size=world_size, rank=rank, local_rank=local_rank)
+    initialize_model_parallel(tensor_model_parallel_size=args.sp_size, sequence_model_parallel_size=args.sp_size)
+    fastvideo_args = FastVideoArgs(model_path=MODEL_PATH,
+                                   num_gpus=world_size,
+                                   use_cpu_offload=False,
+                                   precision=args.master_weight_type,
+                                   dit_config=WanVideoConfig(),
+                                   device_str="cuda",
+                                   )
+    fastvideo_args.check_fastvideo_args()
+
+    device_str = f"cuda:{rank}"
+    device = torch.device(device_str)
+    fastvideo_args.device = device
+
+    # If passed along, set the training seed now. On GPU...
+    if args.seed is not None:
+        # TODO: t within the same seq parallel group should be the same. Noise should be different.
+        set_seed(args.seed + rank)
+    # We use different seeds for the noise generation in each process to ensure that the noise is different in a batch.
+    noise_random_generator = None
+
+    # Handle the repository creation
+    if rank <= 0 and args.output_dir is not None:
+        os.makedirs(args.output_dir, exist_ok=True)
+
+    # For mixed precision training we cast all non-trainable weights to half-precision
+    # as these weights are only used for inference, keeping weights in full precision is not required.
+
+    # Create model:
+
+    logger.info("--> loading model from %s", TRANSFORMER_PATH)
+
+    fastvideo_args.device = device
+    transformer_loader = TransformerLoader()
+    transformer = transformer_loader.load(TRANSFORMER_PATH, "", fastvideo_args)
+    transformer = transformer.train()
+    transformer.requires_grad_(True)
+
+    teacher_loader = TransformerLoader()
+    teacher_transformer = teacher_loader.load(TRANSFORMER_PATH, "", fastvideo_args)
+    if args.use_ema:
+        ema_transformer = teacher_loader.load(TRANSFORMER_PATH, "", fastvideo_args)
+    else:
+        ema_transformer = None
+
+    logger.info(
+        "  Total training parameters = %s M", sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e6)
+    logger.info("--> model loaded")
+
+    teacher_transformer.requires_grad_(False)
+    if args.use_ema:
+        ema_transformer.requires_grad_(False)
+
+    # scheduler
+    noise_scheduler_loader = SchedulerLoader()
+    noise_scheduler = noise_scheduler_loader.load(SCHEDULER_PATH, "", fastvideo_args)
+    solver = EulerSolver(
+        noise_scheduler.sigmas.numpy()[::-1],
+        noise_scheduler.config.num_train_timesteps,
+        euler_timesteps=args.num_euler_timesteps,
+    )
+    solver.to(device)
+    params_to_optimize = transformer.parameters()
+    params_to_optimize = list(filter(lambda p: p.requires_grad, params_to_optimize))
+
+    optimizer = torch.optim.AdamW(
+        params_to_optimize,
+        lr=args.learning_rate,
+        betas=(0.9, 0.999),
+        weight_decay=args.weight_decay,
+        eps=1e-8,
+    )
+
+    init_steps = 0
+    logger.info("optimizer: %s", optimizer)
+
+    # todo add lr scheduler
+    lr_scheduler = get_scheduler(
+        args.lr_scheduler,
+        optimizer=optimizer,
+        num_warmup_steps=args.lr_warmup_steps * world_size,
+        num_training_steps=args.max_train_steps * world_size,
+        num_cycles=args.lr_num_cycles,
+        power=args.lr_power,
+        last_epoch=init_steps - 1,
+    )
+
+    train_dataset = LatentDataset(args.data_json_path, args.num_latent_t, args.cfg)
+    uncond_prompt_embed = train_dataset.uncond_prompt_embed
+    uncond_prompt_mask = train_dataset.uncond_prompt_mask
+    sampler = (LengthGroupedSampler(
+        args.train_batch_size,
+        rank=rank,
+        world_size=world_size,
+        lengths=train_dataset.lengths,
+        group_frame=args.group_frame,
+        group_resolution=args.group_resolution,
+    ) if (args.group_frame or args.group_resolution) else DistributedSampler(
+        train_dataset, rank=rank, num_replicas=world_size, shuffle=False))
+
+    train_dataloader = DataLoader(
+        train_dataset,
+        sampler=sampler,
+        collate_fn=latent_collate_function,
+        pin_memory=True,
+        batch_size=args.train_batch_size,
+        num_workers=args.dataloader_num_workers,
+        drop_last=True,
+    )
+
+    num_update_steps_per_epoch = math.ceil(
+        len(train_dataloader) / args.gradient_accumulation_steps * args.sp_size / args.train_sp_batch_size)
+    args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
+
+    if rank <= 0:
+        project = args.tracker_project_name or "fastvideo"
+        wandb.init(project=project, config=args)
+
+    # Train!
+    total_batch_size = (world_size * args.gradient_accumulation_steps / args.sp_size * args.train_sp_batch_size)
+    logger.info("***** Running training *****")
+    logger.info("  Num examples = %s", len(train_dataset))
+    logger.info("  Dataloader size = %s", len(train_dataloader))
+    logger.info("  Num Epochs = %s", args.num_train_epochs)
+    logger.info("  Resume training from step %s", init_steps)
+    logger.info("  Instantaneous batch size per device = %s", args.train_batch_size)
+    logger.info("  Total train batch size (w. data & sequence parallel, accumulation) = %s", total_batch_size)
+    logger.info("  Gradient Accumulation steps = %s", args.gradient_accumulation_steps)
+    logger.info("  Total optimization steps = %s", args.max_train_steps)
+    logger.info(
+        "  Total training parameters per FSDP shard = %s B", sum(p.numel() for p in transformer.parameters() if p.requires_grad) / 1e9
+    )
+    # print dtype
+    logger.info("  Master weight dtype: %s", transformer.parameters().__next__().dtype)
+
+    # Potentially load in the weights and states from a previous save
+    if args.resume_from_checkpoint:
+        assert NotImplementedError("resume_from_checkpoint is not supported now.")
+        # TODO
+
+    progress_bar = tqdm(
+        range(0, args.max_train_steps),
+        initial=init_steps,
+        desc="Steps",
+        # Only show the progress bar once on each machine.
+        disable=local_rank > 0,
+    )
+
+    loader = sp_parallel_dataloader_wrapper(
+        train_dataloader,
+        device,
+        args.train_batch_size,
+        args.sp_size,
+        args.train_sp_batch_size,
+    )
+
+    step_times = deque(maxlen=100)
+
+    # todo future
+    for i in range(init_steps):
+        next(loader)
+
+    # log_validation(args, transformer, device,
+    #             torch.bfloat16, 0, scheduler_type=args.scheduler_type, shift=args.shift, num_euler_timesteps=args.num_euler_timesteps, linear_quadratic_threshold=args.linear_quadratic_threshold,ema=False)
+    def get_num_phases(multi_phased_distill_schedule, step):
+        # step-phase,step-phase
+        multi_phases = multi_phased_distill_schedule.split(",")
+        phase = multi_phases[-1].split("-")[-1]
+        for step_phases in multi_phases:
+            phase_step, phase = step_phases.split("-")
+            if step <= int(phase_step):
+                return int(phase)
+        return phase
+
+    for step in range(init_steps + 1, args.max_train_steps + 1):
+        start_time = time.time()
+        assert args.multi_phased_distill_schedule is not None
+        num_phases = get_num_phases(args.multi_phased_distill_schedule, step)
+
+        loss, grad_norm, pred_norm = distill_one_step(
+            transformer,
+            args.model_type,
+            teacher_transformer,
+            ema_transformer,
+            optimizer,
+            lr_scheduler,
+            loader,
+            noise_scheduler,
+            solver,
+            noise_random_generator,
+            args.gradient_accumulation_steps,
+            args.sp_size,
+            args.max_grad_norm,
+            uncond_prompt_embed,
+            uncond_prompt_mask,
+            args.num_euler_timesteps,
+            num_phases,
+            args.not_apply_cfg_solver,
+            args.distill_cfg,
+            args.ema_decay,
+            args.pred_decay_weight,
+            args.pred_decay_type,
+            args.hunyuan_teacher_disable_cfg,
+        )
+
+        step_time = time.time() - start_time
+        step_times.append(step_time)
+        avg_step_time = sum(step_times) / len(step_times)
+
+        progress_bar.set_postfix({
+            "loss": f"{loss:.4f}",
+            "step_time": f"{step_time:.2f}s",
+            "grad_norm": grad_norm,
+            "phases": num_phases,
+        })
+        progress_bar.update(1)
+        if rank <= 0:
+            wandb.log(
+                {
+                    "train_loss": loss,
+                    "learning_rate": lr_scheduler.get_last_lr()[0],
+                    "step_time": step_time,
+                    "avg_step_time": avg_step_time,
+                    "grad_norm": grad_norm,
+                    "pred_fro_norm": pred_norm["fro"],  # codespell:ignore
+                    "pred_largest_singular_value": pred_norm["largest singular value"],
+                    "pred_absolute_mean": pred_norm["absolute mean"],
+                    "pred_absolute_max": pred_norm["absolute max"],
+                },
+                step=step,
+            )
+        if step % args.checkpointing_steps == 0:
+            if args.use_lora:
+                # Save LoRA weights
+                save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, step)
+            else:
+                # Your existing checkpoint saving code
+                if args.use_ema:
+                    save_checkpoint(ema_transformer, rank, args.output_dir, step)
+                else:
+                    save_checkpoint(transformer, rank, args.output_dir, step)
+            dist.barrier()
+        if args.log_validation and step % args.validation_steps == 0:
+            log_validation(
+                args,
+                transformer,
+                device,
+                torch.bfloat16,
+                step,
+                scheduler_type=args.scheduler_type,
+                shift=args.shift,
+                num_euler_timesteps=args.num_euler_timesteps,
+                linear_quadratic_threshold=args.linear_quadratic_threshold,
+                linear_range=args.linear_range,
+                ema=False,
+            )
+            if args.use_ema:
+                log_validation(
+                    args,
+                    ema_transformer,
+                    device,
+                    torch.bfloat16,
+                    step,
+                    scheduler_type=args.scheduler_type,
+                    shift=args.shift,
+                    num_euler_timesteps=args.num_euler_timesteps,
+                    linear_quadratic_threshold=args.linear_quadratic_threshold,
+                    linear_range=args.linear_range,
+                    ema=True,
+                )
+
+    if args.use_lora:
+        save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps)
+    else:
+        save_checkpoint(transformer, rank, args.output_dir, args.max_train_steps)
+
+    if get_sequence_parallel_state():
+        destroy_sequence_parallel_group()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--model_type", type=str, default="mochi", help="The type of model to train.")
+
+    # dataset & dataloader
+    parser.add_argument("--data_json_path", type=str, required=True)
+    parser.add_argument("--num_height", type=int, default=480)
+    parser.add_argument("--num_width", type=int, default=848)
+    parser.add_argument("--num_frames", type=int, default=163)
+    parser.add_argument(
+        "--dataloader_num_workers",
+        type=int,
+        default=10,
+        help="Number of subprocesses to use for data loading. 0 means that the data will be loaded in the main process.",
+    )
+    parser.add_argument(
+        "--train_batch_size",
+        type=int,
+        default=16,
+        help="Batch size (per device) for the training dataloader.",
+    )
+    parser.add_argument("--num_latent_t", type=int, default=28, help="Number of latent timesteps.")
+    parser.add_argument("--group_frame", action="store_true")  # TODO
+    parser.add_argument("--group_resolution", action="store_true")  # TODO
+
+    # text encoder & vae & diffusion model
+    parser.add_argument("--pretrained_model_name_or_path", type=str)
+    parser.add_argument("--dit_model_name_or_path", type=str)
+    parser.add_argument("--cache_dir", type=str, default="./cache_dir")
+
+    # diffusion setting
+    parser.add_argument("--ema_decay", type=float, default=0.95)
+    parser.add_argument("--ema_start_step", type=int, default=0)
+    parser.add_argument("--cfg", type=float, default=0.1)
+
+    # validation & logs
+    parser.add_argument("--validation_prompt_dir", type=str)
+    parser.add_argument("--validation_sampling_steps", type=str, default="64")
+    parser.add_argument("--validation_guidance_scale", type=str, default="4.5")
+
+    parser.add_argument("--validation_steps", type=float, default=64)
+    parser.add_argument("--log_validation", action="store_true")
+    parser.add_argument("--tracker_project_name", type=str, default=None)
+    parser.add_argument("--seed", type=int, default=None, help="A seed for reproducible training.")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+    parser.add_argument(
+        "--checkpoints_total_limit",
+        type=int,
+        default=None,
+        help=("Max number of checkpoints to store."),
+    )
+    parser.add_argument(
+        "--checkpointing_steps",
+        type=int,
+        default=500,
+        help=("Save a checkpoint of the training state every X updates. These checkpoints can be used both as final"
+              " checkpoints in case they are better than the last checkpoint, and are also suitable for resuming"
+              " training using `--resume_from_checkpoint`."),
+    )
+    parser.add_argument("--shift", type=float, default=1.0)
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help=("Whether training should be resumed from a previous checkpoint. Use a path saved by"
+              ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
+    )
+    parser.add_argument(
+        "--resume_from_lora_checkpoint",
+        type=str,
+        default=None,
+        help=("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
+              ' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
+    )
+    parser.add_argument(
+        "--logging_dir",
+        type=str,
+        default="logs",
+        help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
+              " *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
+    )
+
+    # optimizer & scheduler & Training
+    parser.add_argument("--num_train_epochs", type=int, default=100)
+    parser.add_argument(
+        "--max_train_steps",
+        type=int,
+        default=None,
+        help="Total number of training steps to perform.  If provided, overrides num_train_epochs.",
+    )
+    parser.add_argument(
+        "--gradient_accumulation_steps",
+        type=int,
+        default=1,
+        help="Number of updates steps to accumulate before performing a backward/update pass.",
+    )
+    parser.add_argument(
+        "--learning_rate",
+        type=float,
+        default=1e-4,
+        help="Initial learning rate (after the potential warmup period) to use.",
+    )
+    parser.add_argument(
+        "--scale_lr",
+        action="store_true",
+        default=False,
+        help="Scale the learning rate by the number of GPUs, gradient accumulation steps, and batch size.",
+    )
+    parser.add_argument(
+        "--lr_warmup_steps",
+        type=int,
+        default=10,
+        help="Number of steps for the warmup in the lr scheduler.",
+    )
+    parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
+    parser.add_argument(
+        "--gradient_checkpointing",
+        action="store_true",
+        help="Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.",
+    )
+    parser.add_argument("--selective_checkpointing", type=float, default=1.0)
+    parser.add_argument(
+        "--allow_tf32",
+        action="store_true",
+        help=("Whether or not to allow TF32 on Ampere GPUs. Can be used to speed up training. For more information, see"
+              " https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices"),
+    )
+    parser.add_argument(
+        "--mixed_precision",
+        type=str,
+        default=None,
+        choices=["no", "fp16", "bf16"],
+        help=(
+            "Whether to use mixed precision. Choose between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >="
+            " 1.10.and an Nvidia Ampere GPU.  Default to the value of accelerate config of the current system or the"
+            " flag passed with the `accelerate.launch` command. Use this argument to override the accelerate config."),
+    )
+    parser.add_argument(
+        "--use_cpu_offload",
+        action="store_true",
+        help="Whether to use CPU offload for param & gradient & optimizer states.",
+    )
+
+    parser.add_argument("--sp_size", type=int, default=1, help="For sequence parallel")
+    parser.add_argument(
+        "--train_sp_batch_size",
+        type=int,
+        default=1,
+        help="Batch size for sequence parallel training",
+    )
+
+    parser.add_argument(
+        "--use_lora",
+        action="store_true",
+        default=False,
+        help="Whether to use LoRA for finetuning.",
+    )
+    parser.add_argument("--lora_alpha", type=int, default=256, help="Alpha parameter for LoRA.")
+    parser.add_argument("--lora_rank", type=int, default=128, help="LoRA rank parameter. ")
+    parser.add_argument("--fsdp_sharding_startegy", default="full")
+
+    # lr_scheduler
+    parser.add_argument(
+        "--lr_scheduler",
+        type=str,
+        default="constant",
+        help=('The scheduler type to use. Choose between ["linear", "cosine", "cosine_with_restarts", "polynomial",'
+              ' "constant", "constant_with_warmup"]'),
+    )
+    parser.add_argument("--num_euler_timesteps", type=int, default=100)
+    parser.add_argument(
+        "--lr_num_cycles",
+        type=int,
+        default=1,
+        help="Number of cycles in the learning rate scheduler.",
+    )
+    parser.add_argument(
+        "--lr_power",
+        type=float,
+        default=1.0,
+        help="Power factor of the polynomial scheduler.",
+    )
+    parser.add_argument(
+        "--not_apply_cfg_solver",
+        action="store_true",
+        help="Whether to apply the cfg_solver.",
+    )
+    parser.add_argument("--distill_cfg", type=float, default=3.0, help="Distillation coefficient.")
+    # ["euler_linear_quadratic", "pcm", "pcm_linear_qudratic"]
+    parser.add_argument("--scheduler_type", type=str, default="pcm", help="The scheduler type to use.")
+    parser.add_argument(
+        "--linear_quadratic_threshold",
+        type=float,
+        default=0.025,
+        help="Threshold for linear quadratic scheduler.",
+    )
+    parser.add_argument(
+        "--linear_range",
+        type=float,
+        default=0.5,
+        help="Range for linear quadratic scheduler.",
+    )
+    parser.add_argument("--weight_decay", type=float, default=0.001, help="Weight decay to apply.")
+    parser.add_argument("--use_ema", action="store_true", help="Whether to use EMA.")
+    parser.add_argument("--multi_phased_distill_schedule", type=str, default=None)
+    parser.add_argument("--pred_decay_weight", type=float, default=0.0)
+    parser.add_argument("--pred_decay_type", default="l1")
+    parser.add_argument("--hunyuan_teacher_disable_cfg", action="store_true")
+    parser.add_argument(
+        "--master_weight_type",
+        type=str,
+        default="fp32",
+        help="Weight type to use - fp32 or bf16.",
+    )
+    args = parser.parse_args()
+    main(args)

fastvideo/v1/attention/layer.py

@@ -85,8 +85,8 @@ class DistributedAttention(nn.Module):
         assert q.dim() == 4 and k.dim() == 4 and v.dim(
         ) == 4, "Expected 4D tensors"
         # assert bs = 1
-        assert q.shape[
-            0] == 1, "Batch size must be 1, and there should be no padding tokens"
+        # assert q.shape[
+        #     0] == 1, "Batch size must be 1, and there should be no padding tokens"
         batch_size, seq_len, num_heads, head_dim = q.shape
         local_rank = get_sequence_model_parallel_rank()
         world_size = get_sequence_model_parallel_world_size()

fastvideo/v1/dataset/__init__.py

@@ -2,8 +2,8 @@ from torchvision import transforms
 from torchvision.transforms import Lambda
 from transformers import AutoTokenizer
 
-from fastvideo.dataset.t2v_datasets import T2V_dataset
-from fastvideo.dataset.transform import CenterCropResizeVideo, Normalize255, TemporalRandomCrop
+from fastvideo.v1.dataset.t2v_datasets import T2V_dataset
+from fastvideo.v1.dataset.transform import CenterCropResizeVideo, Normalize255, TemporalRandomCrop
 
 
 def getdataset(args):
@@ -44,7 +44,7 @@ if __name__ == "__main__":
     from accelerate import Accelerator
     from tqdm import tqdm
 
-    from fastvideo.dataset.t2v_datasets import dataset_prog
+    from fastvideo.v1.dataset.t2v_datasets import dataset_prog
 
     args = type(
         "args",

fastvideo/v1/dataset/dataloader/schema.py

@@ -0,0 +1,46 @@
+# schema.py
+"""
+Unified data schema and format for saving and loading image/video data after
+preprocessing.
+
+It uses apache arrow in-memory format that can be consumed by modern data
+frameworks that can handle parquet or lance file.
+"""
+
+import pyarrow as pa
+
+pyarrow_schema = pa.schema(
+    [
+        pa.field("id", pa.string()),
+        # --- Image/Video VAE latents ---
+        # Tensors are stored as raw bytes with shape and dtype info for loading
+        pa.field("vae_latent_bytes", pa.binary()),
+        # e.g., [C, T, H, W] or [C, H, W]
+        pa.field("vae_latent_shape", pa.list_(pa.int64())),
+        # e.g., 'float32'
+        pa.field("vae_latent_dtype", pa.string()),
+        # --- Text encoder output tensor ---
+        # Tensors are stored as raw bytes with shape and dtype info for loading
+        pa.field("text_embedding_bytes", pa.binary()),
+        # e.g., [SeqLen, Dim]
+        pa.field("text_embedding_shape", pa.list_(pa.int64())),
+        # e.g., 'bfloat16' or 'float32'
+        pa.field("text_embedding_dtype", pa.string()),
+        pa.field("text_attention_mask_bytes", pa.binary()),
+        # e.g., [SeqLen]
+        pa.field("text_attention_mask_shape", pa.list_(pa.int64())),
+        # e.g., 'bool' or 'int8'
+        pa.field("text_attention_mask_dtype", pa.string()),
+        # --- Metadata ---
+        pa.field("file_name", pa.string()),
+        pa.field("caption", pa.string()),
+        pa.field("media_type", pa.string()),  # 'image' or 'video'
+        pa.field("width", pa.int64()),
+        pa.field("height", pa.int64()),
+        # -- Video-specific (can be null/default for images) ---
+        # Number of frames processed (e.g., 1 for image, N for video)
+        pa.field("num_frames", pa.int64()),
+        pa.field("duration_sec", pa.float64()),
+        pa.field("fps", pa.float64()),
+    ]
+)

fastvideo/v1/dataset/parquet_datasets.py

@@ -0,0 +1,421 @@
+import argparse
+import os
+import random
+import time
+from typing import Any, Dict, Iterator, List, Optional, Tuple
+
+import numpy as np
+import pyarrow.parquet as pq
+import torch
+import tqdm
+from torch.utils.data import IterableDataset, get_worker_info
+from torchdata.stateful_dataloader import StatefulDataLoader
+
+# Path to your dataset
+dataset_path = "/mnt/sharefs/users/hao.zhang/Vchitect-2M/Vchitect-2M-laten-93x512x512/train/"
+
+
+class ParquetVideoTextDataset(IterableDataset):
+    """Efficient loader for video-text data from a directory of Parquet files."""
+
+    def __init__(self,
+                 path: str,
+                 batch_size: int = 1024,
+                 rank: int = 0,
+                 world_size: int = 1,
+                 cfg_rate: float = 0.0,
+                 num_workers: int = 1,
+                 row_per_parquet: int = 32,
+                 split: str = "train"):
+        super().__init__()
+        self.path = str(path)
+        self.batch_size = batch_size
+        self.rank = rank
+        self.world_size = world_size
+        self.cfg_rate = cfg_rate
+
+        assert split in ["train", "validation"]
+        # Find all parquet files recursively
+        print(f"Scanning for parquet files in {self.path}")
+        self.parquet_files: List[str] = []
+        for root, _, files in os.walk(self.path):
+            for file in files:
+                if file.endswith('.parquet') and split in file:
+                    self.parquet_files.append(os.path.join(root, file))
+        # Sort files for consistent ordering
+        self.parquet_files.sort()
+
+        # 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 < 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
+
+        self.parquet_files = self.parquet_files[start_idx:end_idx]
+
+        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]}")
+
+        # Initialize current file index
+        self.current_file_idx = 0
+        self.current_reader: Optional[pq.ParquetFile] = None
+        self.current_batches: Optional[Iterator[Any]] = None
+        self.total_samples = 0
+
+    def _open_next_file(self) -> bool:
+        """Open the next parquet file for reading."""
+        worker_info = get_worker_info()
+        if worker_info is None:
+            return False
+
+        num_workers = worker_info.num_workers
+        worker_id = 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)
+            if self.current_reader is not None:
+                self.current_batches = self.current_reader.iter_batches(
+                    batch_size=self.batch_size)
+                self.current_file_idx += 1
+                return True
+            return False
+        except Exception as e:
+            print(f"Error opening file {file_path}: {str(e)}")
+            return False
+
+    def __iter__(
+        self
+    ) -> Iterator[Tuple[torch.Tensor, torch.Tensor, torch.Tensor, Dict[str,
+                                                                       Any]]]:
+        """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:
+            try:
+                # Get next batch from current file
+                if self.current_batches is None:
+                    raise StopIteration
+
+                batch = next(self.current_batches)
+                batch_dict = batch.to_pydict()
+                processed = self._process_batch(batch_dict)
+
+                # Update sample count
+                batch_size = len(processed["latents"])
+                self.total_samples += batch_size
+
+                # Print progress
+                if self.total_samples % 1000 == 0:
+                    print(
+                        f"Rank {self.rank}: Processed {self.total_samples} samples"
+                    )
+
+                # Yield each item in the batch
+                for lat, emb, mask, info in zip(processed["latents"],
+                                                processed["embeddings"],
+                                                processed["masks"],
+                                                processed["info"]):
+                    yield from (lat, emb, mask, info)
+
+            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
+
+        # Clean up
+        if self.current_reader is not None:
+            self.current_reader.close()
+
+    def _process_batch(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        """Process a PyArrow batch into tensors."""
+        out: Dict[str, List[torch.Tensor]] = {
+            "lat": [],
+            "emb": [],
+            "msk": [],
+            "info": []
+        }
+
+        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]
+
+            # 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, :]
+
+            # 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.int64).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_)
+            # 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": str(batch["file_name"][i]),
+                "caption": batch["caption"][i],
+            }
+
+            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"]
+        }
+
+
+def bind_cpu_cores(local_rank: int, cpu_per_process: int = 16) -> None:
+    """根据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)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description='Benchmark Parquet dataset loading speed')
+    parser.add_argument('--path',
+                        type=str,
+                        default=dataset_path,
+                        help='Path to Parquet dataset')
+    parser.add_argument('--batch_size',
+                        type=int,
+                        default=4,
+                        help='Batch size for DataLoader')
+    parser.add_argument('--num_batches',
+                        type=int,
+                        default=100,
+                        help='Number of batches to benchmark')
+    parser.add_argument('--vae_debug', action="store_true")
+    args = parser.parse_args()
+
+    # Initialize distributed training
+    local_rank = int(os.environ.get("LOCAL_RANK", 0))
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    rank = int(os.environ.get("RANK", 0))
+
+    # Initialize CUDA device first
+    if torch.cuda.is_available():
+        torch.cuda.set_device(local_rank)
+        device = torch.device(f"cuda:{local_rank}")
+    else:
+        device = torch.device("cpu")
+
+    # Initialize distributed training
+    if world_size > 1:
+        from torch import distributed as dist
+        dist.init_process_group(backend="nccl",
+                                init_method="env://",
+                                world_size=world_size,
+                                rank=rank)
+        print(
+            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,
+        batch_size=args.batch_size,
+        rank=rank,
+        world_size=world_size,
+        num_workers=1,
+        split="train",
+    )
+
+    # Create DataLoader with proper settings
+    dataloader = StatefulDataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        num_workers=1,  # Reduce number of workers to avoid memory issues
+        prefetch_factor=2,
+        shuffle=False,
+        pin_memory=True,
+        drop_last=True)
+
+    # Warm-up with synchronization
+    if rank == 0:
+        print("Warming up...")
+    for i, (latents, embeddings, masks, infos) in enumerate(dataloader):
+        assert torch.sum(masks[0]).item() == torch.count_nonzero(
+            embeddings[0]).item() // 4096
+        if args.vae_debug:
+            from diffusers.utils import export_to_video
+            from diffusers.video_processor import VideoProcessor
+
+            from fastvideo.v1.configs.models.vaes import WanVAEConfig
+            from fastvideo.v1.fastvideo_args import FastVideoArgs
+            from fastvideo.v1.models.loader.component_loader import VAELoader
+            VAE_PATH = "/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers/vae"
+            fastvideo_args = FastVideoArgs(
+                model_path=VAE_PATH,
+                vae_config=WanVAEConfig(load_encoder=False),
+                vae_precision="fp32")
+            fastvideo_args.device = device
+            vae_loader = VAELoader()
+            vae = vae_loader.load(model_path=VAE_PATH,
+                                  architecture="",
+                                  fastvideo_args=fastvideo_args)
+
+            videoprocessor = VideoProcessor(vae_scale_factor=8)
+
+            with torch.inference_mode():
+                video = vae.decode(latents[0].unsqueeze(0).to(device))
+                video = videoprocessor.postprocess_video(video)
+                video_path = os.path.join("/workspace/FastVideo/debug_videos",
+                                          infos["caption"][0][:50] + ".mp4")
+                export_to_video(video[0], video_path, fps=16)
+        if i >= 5:
+            break
+        # Move data to device
+        latents = latents.to(device)
+        embeddings = embeddings.to(device)
+    if world_size > 1:
+        dist.barrier()
+
+    # Benchmark
+    if rank == 0:
+        print(f"Benchmarking with batch_size={args.batch_size}")
+    start_time = time.time()
+    total_samples = 0
+    for i, (latents, embeddings, masks,
+            infos) in enumerate(tqdm.tqdm(dataloader, total=args.num_batches)):
+        if i >= args.num_batches:
+            break
+
+        # Move data to device
+        latents = latents.to(device)
+        embeddings = embeddings.to(device)
+
+        # Calculate actual batch size
+        batch_size = latents.size(0)
+        total_samples += batch_size
+
+        # Print progress only from rank 0
+        if rank == 0 and (i + 1) % 10 == 0:
+            elapsed = time.time() - start_time
+            samples_per_sec = total_samples / elapsed
+            print(
+                f"Batch {i+1}/{args.num_batches}, Speed: {samples_per_sec:.2f} samples/sec"
+            )
+
+    # Final statistics
+    if world_size > 1:
+        dist.barrier()
+
+    if rank == 0:
+        elapsed = time.time() - start_time
+        samples_per_sec = total_samples / elapsed
+
+        print("\nBenchmark Results:")
+        print(f"Total time: {elapsed:.2f} seconds")
+        print(f"Total samples: {total_samples}")
+        print(f"Average speed: {samples_per_sec:.2f} samples/sec")
+        print(f"Time per batch: {elapsed/args.num_batches*1000:.2f} ms")
+
+    if world_size > 1:
+        dist.destroy_process_group()

fastvideo/v1/dataset/t2v_datasets.py

@@ -166,6 +166,8 @@ class T2V_dataset(Dataset):
             input_ids=input_ids,
             cond_mask=cond_mask,
             path=video_path,
+            fps=dataset_prog.cap_list[idx]["fps"],
+            duration=dataset_prog.cap_list[idx]["duration"]
         )
 
     def get_image(self, idx):

fastvideo/v1/fastvideo_args.py

@@ -70,7 +70,7 @@ class FastVideoArgs:
     # Text encoder configuration
     DEFAULT_TEXT_ENCODER_PRECISIONS = (
         "fp16",
-        "fp16",
+        # "fp16",
     )
     text_encoder_precisions: Tuple[str, ...] = field(
         default_factory=lambda: FastVideoArgs.DEFAULT_TEXT_ENCODER_PRECISIONS)

fastvideo/v1/layers/layernorm.py

@@ -146,7 +146,7 @@ class ScaleResidualLayerNormScaleShift(nn.Module):
         # Apply normalization
         normalized = self.norm(residual_output)
         # Apply scale and shift
-        modulated = normalized * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        modulated = normalized * (1.0 + scale) + shift
         return modulated, residual_output
 
 
@@ -182,4 +182,4 @@ class LayerNormScaleShift(nn.Module):
                 scale: torch.Tensor) -> torch.Tensor:
         """Apply ln followed by scale and shift in a single fused operation."""
         normalized = self.norm(x)
-        return normalized * (1.0 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        return normalized * (1.0 + scale) + shift

fastvideo/v1/pipelines/composed_pipeline_base.py

@@ -136,10 +136,7 @@ class ComposedPipelineBase(ABC):
             modules_config
         ) > 1, "model_index.json must contain at least one pipeline module"
 
-        required_modules = [
-            "vae", "text_encoder", "transformer", "scheduler", "tokenizer"
-        ]
-        for module_name in required_modules:
+        for module_name in self.required_config_modules:
             if module_name not in modules_config:
                 raise ValueError(
                     f"model_index.json must contain a {module_name} module")
@@ -149,6 +146,8 @@ class ComposedPipelineBase(ABC):
         modules = {}
         for module_name, (transformers_or_diffusers,
                           architecture) in modules_config.items():
+            if module_name not in self.required_config_modules:
+                continue
             component_model_path = os.path.join(self.model_path, module_name)
             module = PipelineComponentLoader.load_module(
                 module_name=module_name,

fastvideo/v1/pipelines/data_preprocess.py

@@ -0,0 +1,497 @@
+# SPDX-License-Identifier: Apache-2.0
+"""
+T2V Data Preprocessing pipeline implementation.
+
+This module contains an implementation of the T2V Data Preprocessing pipeline
+using the modular pipeline architecture.
+"""
+import os
+from tqdm import tqdm
+
+import torch
+from torch.utils.data import DataLoader
+from torch.utils.data.distributed import DistributedSampler
+import numpy as np
+import pyarrow as pa
+import pyarrow.parquet as pq
+
+import multiprocessing
+from concurrent.futures import ProcessPoolExecutor
+import gc
+
+from fastvideo.v1.dataset import getdataset
+from fastvideo.v1.dataset.dataloader.schema import pyarrow_schema
+from fastvideo.v1.fastvideo_args import FastVideoArgs
+from fastvideo.v1.logger import init_logger
+from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
+from fastvideo.v1.pipelines.stages import TextEncodingStage
+from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
+
+# TODO(will): move PRECISION_TO_TYPE to better place
+
+logger = init_logger(__name__)
+
+
+class PreprocessPipeline(ComposedPipelineBase):
+
+    _required_config_modules = [
+        "text_encoder", "tokenizer", "vae"
+    ]
+
+    def create_pipeline_stages(self, fastvideo_args: FastVideoArgs):
+        """Set up pipeline stages with proper dependency injection."""
+
+        self.add_stage(stage_name="prompt_encoding_stage",
+                       stage=TextEncodingStage(
+                           text_encoders=[self.get_module("text_encoder")],
+                           tokenizers=[self.get_module("tokenizer")],
+                       ))
+    
+    @torch.no_grad()
+    def forward(
+        self,
+        batch: ForwardBatch,
+        fastvideo_args: FastVideoArgs,
+        args,
+    ):
+        # Initialize class variables for data sharing
+        self.video_data = {}  # Store video metadata and paths
+        self.latent_data = {}  # Store latent tensors
+        self.preprocess_validation_text(fastvideo_args, args)
+        self.preprocess_video(fastvideo_args, args)
+        self.preprocess_text(fastvideo_args, args)
+
+    def preprocess_video(self, fastvideo_args: FastVideoArgs, args):
+        local_rank = int(os.getenv("RANK", 0))
+        world_size = int(os.getenv("WORLD_SIZE", 1))
+        # Loading dataset
+        train_dataset = getdataset(args)
+        sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
+        train_dataloader = DataLoader(
+            train_dataset,
+            sampler=sampler,
+            batch_size=args.preprocess_video_batch_size,
+            num_workers=args.dataloader_num_workers,
+        )
+
+        # Add progress bar for video preprocessing
+        pbar = tqdm(train_dataloader, desc="Processing videos", unit="batch", disable=local_rank != 0)
+        for batch_idx, data in enumerate(pbar):
+            if batch_idx == 4:
+                break
+            if data is None:
+                continue            
+
+            with torch.inference_mode():
+                # Filter out invalid samples (those with all zeros)
+                valid_indices = []
+                for i, pixel_values in enumerate(data["pixel_values"]):
+                    if not torch.all(pixel_values == 0):  # Check if all values are zero
+                        valid_indices.append(i)
+                
+                if not valid_indices:
+                    continue
+
+                # Create new batch with only valid samples
+                valid_data = {
+                    "pixel_values": torch.stack([data["pixel_values"][i] for i in valid_indices]),
+                    "text": [data["text"][i] for i in valid_indices],
+                    "path": [data["path"][i] for i in valid_indices],
+                    "fps": [data["fps"][i] for i in valid_indices],
+                    "duration": [data["duration"][i] for i in valid_indices],
+                }
+
+                # VAE
+                with torch.autocast("cuda", dtype=torch.float32):
+                    print(valid_data["pixel_values"].shape)
+                    latents = self.get_module("vae").encode(valid_data["pixel_values"].to(fastvideo_args.device)).mean
+                    print(latents.shape)
+
+                for idx, video_path in enumerate(valid_data["path"]):
+                    video_name = os.path.basename(video_path).split(".")[0]
+                    # Get video dimensions from the pixel values
+                    height, width = valid_data["pixel_values"][idx].shape[-2:]
+                    # Store data in class variables - move tensors to CPU
+                    self.video_data[video_name] = {
+                        "caption": valid_data["text"][idx],
+                        "width": width,
+                        "height": height,
+                        "duration_sec": float(valid_data["duration"][idx]),
+                        "fps": float(valid_data["fps"][idx]),
+                        "num_frames": latents[idx].shape[1],
+                    }
+                    # Move latent tensor to CPU before storing
+                    self.latent_data[video_name] = latents[idx].cpu()
+    
+    def preprocess_text(self, fastvideo_args: FastVideoArgs, args):
+        os.makedirs(args.output_dir, exist_ok=True)
+        # Create directory for combined data
+        job_id = os.environ.get('job_id', str(os.getpid()))
+        combined_parquet_dir = os.path.join(args.output_dir, f"combined_parquet_dataset_{job_id}")
+        os.makedirs(combined_parquet_dir, exist_ok=True)
+
+        # Process all videos in batches
+        video_names = list(self.latent_data.keys())
+        text_batch_size = args.preprocess_text_batch_size
+        
+        # Add progress bar for text preprocessing
+        pbar = tqdm(range(0, len(video_names), text_batch_size), desc="Processing text", unit="batch")
+        for batch_idx in pbar:
+            # Get current batch of video names
+            batch_names = video_names[batch_idx:batch_idx + text_batch_size]
+            
+            # Get corresponding captions for this batch
+            batch_captions = [self.video_data[name]["caption"] for name in batch_names]
+            
+            with torch.inference_mode():
+                batch = ForwardBatch(
+                    data_type="video",
+                    prompt=batch_captions,
+                    prompt_embeds=[],
+                    prompt_attention_mask=[],
+                )
+                result_batch = self.prompt_encoding_stage(batch, fastvideo_args)
+                prompt_embeds, prompt_attention_mask = result_batch.prompt_embeds[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
+            non_padded_embeds = []
+            non_padded_masks = []
+            
+            # Process each item in the batch
+            for i in range(prompt_embeds.size(0)):
+                seq_len = seq_lens[i].item()
+                # Slice the embeddings and masks to keep only non-padding parts
+                non_padded_embeds.append(prompt_embeds[i, :seq_len])
+                non_padded_masks.append(prompt_attention_mask[i, :seq_len])
+            
+            # Update the tensors with non-padded versions
+            prompt_embeds = non_padded_embeds
+            prompt_attention_mask = non_padded_masks
+
+            # Prepare batch data for Parquet dataset
+            batch_data = []
+            
+            # Add progress bar for saving outputs
+            save_pbar = tqdm(enumerate(batch_names), desc="Saving outputs", unit="item", leave=False)
+            for idx, video_name in save_pbar:
+                # Get the corresponding latent and info using video name
+                latent = self.latent_data[video_name]
+                info = self.video_data[video_name]
+                
+                # Convert tensors to numpy arrays
+                vae_latent = latent.cpu().numpy()
+                text_embedding = prompt_embeds[idx].cpu().numpy()
+                text_attention_mask = prompt_attention_mask[idx].cpu().numpy()
+                
+                # Create record for Parquet dataset
+                record = {
+                    "id": video_name,
+                    "vae_latent_bytes": vae_latent.tobytes(),
+                    "vae_latent_shape": list(vae_latent.shape),
+                    "vae_latent_dtype": str(vae_latent.dtype),
+                    "text_embedding_bytes": text_embedding.tobytes(),
+                    "text_embedding_shape": list(text_embedding.shape),
+                    "text_embedding_dtype": str(text_embedding.dtype),
+                    "text_attention_mask_bytes": text_attention_mask.tobytes(),
+                    "text_attention_mask_shape": list(text_attention_mask.shape),
+                    "text_attention_mask_dtype": str(text_attention_mask.dtype),
+                    "file_name": video_name,
+                    "caption": info["caption"],
+                    "media_type": "video",
+                    "width": info["width"],
+                    "height": info["height"],
+                    "num_frames": info["num_frames"],
+                    "duration_sec": info["duration_sec"],
+                    "fps": info["fps"],
+                }
+                batch_data.append(record)
+            
+            # After all batches are processed, combine and write in 1GB chunks
+            if batch_data:
+                # Add progress bar for writing to Parquet dataset
+                write_pbar = tqdm(total=1, desc="Writing to Parquet dataset", unit="batch")
+                # Convert batch data to PyArrow arrays
+                arrays = [
+                    pa.array([record["id"] for record in batch_data]),
+                    pa.array([record["vae_latent_bytes"] for record in batch_data], type=pa.binary()),
+                    pa.array([record["vae_latent_shape"] for record in batch_data], type=pa.list_(pa.int32())),
+                    pa.array([record["vae_latent_dtype"] for record in batch_data]),
+                    pa.array([record["text_embedding_bytes"] for record in batch_data], type=pa.binary()),
+                    pa.array([record["text_embedding_shape"] for record in batch_data], type=pa.list_(pa.int32())),
+                    pa.array([record["text_embedding_dtype"] for record in batch_data]),
+                    pa.array([record["text_attention_mask_bytes"] for record in batch_data], type=pa.binary()),
+                    pa.array([record["text_attention_mask_shape"] for record in batch_data], type=pa.list_(pa.int32())),
+                    pa.array([record["text_attention_mask_dtype"] for record in batch_data]),
+                    pa.array([record["file_name"] for record in batch_data]),
+                    pa.array([record["caption"] for record in batch_data]),
+                    pa.array([record["media_type"] for record in batch_data]),
+                    pa.array([record["width"] for record in batch_data], type=pa.int32()),
+                    pa.array([record["height"] for record in batch_data], type=pa.int32()),
+                    pa.array([record["num_frames"] for record in batch_data], type=pa.int32()),
+                    pa.array([record["duration_sec"] for record in batch_data], type=pa.float32()),
+                    pa.array([record["fps"] for record in batch_data], type=pa.float32()),
+                ]
+                table = pa.Table.from_arrays(arrays, names=[f.name for f in pyarrow_schema])
+                write_pbar.update(1)
+                write_pbar.close()
+                
+                # Store the table in a list for later processing
+                if not hasattr(self, 'all_tables'):
+                    self.all_tables = []
+                self.all_tables.append(table)
+                
+                logger.info(f"Collected batch with {len(table)} samples")
+        
+        logger.info("After text preprocessing loop")
+        
+        # After all batches are processed, combine and write in 1GB chunks
+        if hasattr(self, 'all_tables') and self.all_tables:
+            logger.info(f"Combining {len(self.all_tables)} batches...")
+            combined_table = pa.concat_tables(self.all_tables)
+            logger.info(f"Total samples collected: {len(combined_table)}")
+            
+            # Calculate total number of chunks needed
+            num_samples = len(combined_table)
+            
+            # Calculate samples per file based on actual latent shapes
+            first_latent_shape = combined_table.column("vae_latent_shape")[0].as_py()
+            first_text_shape = combined_table.column("text_embedding_shape")[0].as_py()
+            first_mask_shape = combined_table.column("text_attention_mask_shape")[0].as_py()
+            
+            # Calculate size per sample in bytes
+            latent_size = np.prod(first_latent_shape) * 4  # float32 = 4 bytes
+            text_size = np.prod(first_text_shape) * 4  # float32 = 4 bytes
+            mask_size = np.prod(first_mask_shape)  # uint8 = 1 byte
+            metadata_size = 4 * 5 + 8 * 2  # 5 int32 + 2 float32
+            
+            total_size_per_sample = latent_size + text_size + mask_size + metadata_size
+            
+            # Target file size: 1GB in bytes
+            target_file_size = 512 * 1024 * 1024  # 1GB in bytes
+            samples_per_file = max(1, target_file_size // total_size_per_sample)
+            
+            logger.info(f"Estimated size per sample: {total_size_per_sample/1024/1024:.2f}MB")
+            logger.info(f"Samples per file: {samples_per_file}")
+            
+            # Calculate total number of chunks needed, discarding remainder
+            total_chunks = max(num_samples // samples_per_file, 1)
+            
+            logger.info(f"Fixed samples per parquet file: {samples_per_file}")
+            logger.info(f"Total number of parquet files: {total_chunks}")
+            logger.info(f"Total samples to be processed: {total_chunks * samples_per_file} (discarding {num_samples % samples_per_file} samples)")
+            
+            # Split work among processes
+            num_workers = int(min(multiprocessing.cpu_count(), total_chunks))
+            chunks_per_worker = (total_chunks + num_workers - 1) // num_workers
+            
+            logger.info(f"Using {num_workers} workers to process {total_chunks} chunks")
+            logger.info(f"Chunks per worker: {chunks_per_worker}")
+            
+            # Prepare work ranges
+            work_ranges = []
+            for i in range(num_workers):
+                start_idx = i * chunks_per_worker
+                end_idx = min((i + 1) * chunks_per_worker, total_chunks)
+                if start_idx < total_chunks:
+                    work_ranges.append((start_idx, end_idx, combined_table, i, combined_parquet_dir, samples_per_file))
+            
+            total_written = 0
+            failed_ranges = []
+            with ProcessPoolExecutor(max_workers=num_workers) as executor:
+                futures = {executor.submit(self.process_chunk_range, work_range): work_range for work_range in work_ranges}
+                for future in tqdm(futures, desc="Processing chunks"):
+                    try:
+                        written = future.result()
+                        total_written += written
+                        logger.info(f"Processed chunk with {written} samples")
+                    except Exception as e:
+                        work_range = futures[future]
+                        failed_ranges.append(work_range)
+                        logger.error(f"Failed to process range {work_range[0]}-{work_range[1]}: {str(e)}")
+            
+            # Retry failed ranges sequentially
+            if failed_ranges:
+                logger.warning(f"Retrying {len(failed_ranges)} failed ranges sequentially")
+                for work_range in failed_ranges:
+                    try:
+                        total_written += self.process_chunk_range(work_range)
+                    except Exception as e:
+                        logger.error(f"Failed to process range {work_range[0]}-{work_range[1]} after retry: {str(e)}")
+            
+            logger.info(f"Total samples written: {total_written}")
+            
+            # Clear the collected tables to free memory
+            del self.all_tables
+            gc.collect()  # Force garbage collection            
+                
+    def preprocess_validation_text(self, fastvideo_args: FastVideoArgs, args):
+        # Create Parquet dataset directory for validation
+        validation_parquet_dir = os.path.join(args.output_dir, "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, "r", encoding="utf-8") as file:
+            lines = file.readlines()
+        prompts = [line.strip() for line in lines]
+
+        # Prepare batch data for Parquet dataset
+        batch_data = []
+
+        # Add progress bar for validation text preprocessing
+        pbar = tqdm(enumerate(prompts), desc="Processing validation prompts", unit="prompt")
+        for prompt_idx, prompt in pbar:
+            with torch.inference_mode():
+                # Text Encoder
+                batch = ForwardBatch(
+                    data_type="video",
+                    prompt=prompt,
+                    prompt_embeds=[],
+                    prompt_attention_mask=[],
+                )
+                result_batch = self.prompt_encoding_stage(batch, fastvideo_args)
+            prompt_embeds = result_batch.prompt_embeds[0]
+            prompt_attention_mask = result_batch.prompt_attention_mask[0]
+
+            file_name = prompt.split(".")[0]
+
+            # Get the sequence length from attention mask (number of 1s)
+            seq_len = prompt_attention_mask.sum().item()
+            # Slice the embeddings to keep only the non-padding parts
+            text_embedding = prompt_embeds[0, :seq_len].cpu().numpy()
+            text_attention_mask = prompt_attention_mask[0, :seq_len].cpu().numpy()
+
+            # Log the shapes after removing padding
+            logger.info(f"Shape after removing padding - Embeddings: {text_embedding.shape}, Mask: {text_attention_mask.shape}")
+
+            # Create record for Parquet dataset
+            record = {
+                "id": file_name,
+                "vae_latent_bytes": b"",  # Not available for validation
+                "vae_latent_shape": [],
+                "vae_latent_dtype": "",
+                "text_embedding_bytes": text_embedding.tobytes(),
+                "text_embedding_shape": list(text_embedding.shape),
+                "text_embedding_dtype": str(text_embedding.dtype),
+                "text_attention_mask_bytes": text_attention_mask.tobytes(),
+                "text_attention_mask_shape": list(text_attention_mask.shape),
+                "text_attention_mask_dtype": str(text_attention_mask.dtype),
+                "file_name": file_name,
+                "caption": prompt,
+                "media_type": "video",
+                "width": 0,  # Not available for validation
+                "height": 0,  # Not available for validation
+                "num_frames": 0,  # Not available for validation
+                "duration_sec": 0.0,  # Not available for validation
+                "fps": 0.0,  # Not available for validation
+            }
+            batch_data.append(record)
+                
+            logger.info(f"Saved validation sample: {file_name}")
+
+        if batch_data:
+            # Add progress bar for writing to Parquet dataset
+            write_pbar = tqdm(total=1, desc="Writing to Parquet dataset", unit="batch")
+            # Convert batch data to PyArrow arrays
+            arrays = [
+                pa.array([record["id"] for record in batch_data]),
+                pa.array([record["vae_latent_bytes"] for record in batch_data], type=pa.binary()),
+                pa.array([record["vae_latent_shape"] for record in batch_data], type=pa.list_(pa.int32())),
+                pa.array([record["vae_latent_dtype"] for record in batch_data]),
+                pa.array([record["text_embedding_bytes"] for record in batch_data], type=pa.binary()),
+                pa.array([record["text_embedding_shape"] for record in batch_data], type=pa.list_(pa.int32())),
+                pa.array([record["text_embedding_dtype"] for record in batch_data]),
+                pa.array([record["text_attention_mask_bytes"] for record in batch_data], type=pa.binary()),
+                pa.array([record["text_attention_mask_shape"] for record in batch_data], type=pa.list_(pa.int32())),
+                pa.array([record["text_attention_mask_dtype"] for record in batch_data]),
+                pa.array([record["file_name"] for record in batch_data]),
+                pa.array([record["caption"] for record in batch_data]),
+                pa.array([record["media_type"] for record in batch_data]),
+                pa.array([record["width"] for record in batch_data], type=pa.int32()),
+                pa.array([record["height"] for record in batch_data], type=pa.int32()),
+                pa.array([record["num_frames"] for record in batch_data], type=pa.int32()),
+                pa.array([record["duration_sec"] for record in batch_data], type=pa.float32()),
+                pa.array([record["fps"] for record in batch_data], type=pa.float32()),
+            ]
+            table = pa.Table.from_arrays(arrays, names=[f.name for f in pyarrow_schema])
+            write_pbar.update(1)
+            write_pbar.close()
+            
+            logger.info(f"Total validation samples: {len(table)}")
+            
+            work_range = (0, 1, table, 0, validation_parquet_dir, len(table))
+            
+            total_written = 0
+            failed_ranges = []
+            with ProcessPoolExecutor(max_workers=1) as executor:
+                futures = {executor.submit(self.process_chunk_range, work_range): work_range}
+                for future in tqdm(futures, desc="Processing chunks"):
+                    try:
+                        total_written += future.result()
+                    except Exception as e:
+                        work_range = futures[future]
+                        failed_ranges.append(work_range)
+                        logger.error(f"Failed to process range {work_range[0]}-{work_range[1]}: {str(e)}")
+            
+            # Retry failed ranges sequentially
+            if failed_ranges:
+                logger.warning(f"Retrying {len(failed_ranges)} failed ranges sequentially")
+                for work_range in failed_ranges:
+                    try:
+                        total_written += self.process_chunk_range(work_range)
+                    except Exception as e:
+                        logger.error(f"Failed to process range {work_range[0]}-{work_range[1]} after retry: {str(e)}")
+            
+            logger.info(f"Total validation samples written: {total_written}")
+            
+            # Clear memory
+            del table
+            gc.collect()  # Force garbage collection
+
+    @staticmethod
+    def process_chunk_range(args):
+        start_idx, end_idx, table, worker_id, output_dir, samples_per_file = args
+        try:
+            total_written = 0
+            num_samples = len(table)
+            
+            # Create worker-specific subdirectory
+            worker_dir = os.path.join(output_dir, f"worker_{worker_id}")
+            os.makedirs(worker_dir, exist_ok=True)
+            
+            for i in range(start_idx, end_idx):
+                start_sample = i * samples_per_file
+                end_sample = min((i + 1) * samples_per_file, num_samples)
+                chunk = table.slice(start_sample, end_sample - start_sample)
+                
+                # Create chunk file in worker's directory
+                chunk_path = os.path.join(worker_dir, f"data_chunk_{i}.parquet")
+                temp_path = chunk_path + '.tmp'
+                
+                try:
+                    # Write to temporary file
+                    pq.write_table(chunk, temp_path, compression='zstd')
+                    
+                    # Rename temporary file to final file
+                    if os.path.exists(chunk_path):
+                        os.remove(chunk_path)  # Remove existing file if it exists
+                    os.rename(temp_path, chunk_path)
+                    
+                    total_written += len(chunk)
+                except Exception as e:
+                    # Clean up temporary file if it exists
+                    if os.path.exists(temp_path):
+                        os.remove(temp_path)
+                    raise e
+                    
+            return total_written
+        except Exception as e:
+            logger.error(f"Error processing chunks {start_idx}-{end_idx} for worker {worker_id}: {str(e)}")
+            raise
+
+EntryClass = PreprocessPipeline

fastvideo/v1/pipelines/stages/text_encoding.py

@@ -63,9 +63,13 @@ class TextEncodingStage(PipelineStage):
             if fastvideo_args.use_cpu_offload:
                 text_encoder = text_encoder.to(fastvideo_args.device)
 
-            assert isinstance(batch.prompt, str)
-            text = preprocess_func(batch.prompt)
-            text_inputs = tokenizer(text, **encoder_config.tokenizer_kwargs).to(
+            assert isinstance(batch.prompt, (str, list))
+            if isinstance(batch.prompt, str):
+                batch.prompt = [batch.prompt]
+            texts = []
+            for prompt_str in batch.prompt:
+                texts.append(preprocess_func(prompt_str))
+            text_inputs = tokenizer(texts, **encoder_config.tokenizer_kwargs).to(
                 fastvideo_args.device)
             input_ids = text_inputs["input_ids"]
             attention_mask = text_inputs["attention_mask"]
@@ -78,6 +82,8 @@ class TextEncodingStage(PipelineStage):
             prompt_embeds = postprocess_func(outputs)
 
             batch.prompt_embeds.append(prompt_embeds)
+            if batch.prompt_attention_mask is not None:
+                batch.prompt_attention_mask.append(attention_mask)
 
             if batch.do_classifier_free_guidance:
                 assert isinstance(batch.negative_prompt, str)
@@ -98,6 +104,8 @@ class TextEncodingStage(PipelineStage):
 
                 assert batch.negative_prompt_embeds is not None
                 batch.negative_prompt_embeds.append(negative_prompt_embeds)
+                if batch.negative_attention_mask is not None:
+                    batch.negative_attention_mask.append(negative_attention_mask)
 
             if fastvideo_args.use_cpu_offload:
                 text_encoder.to('cpu')

scripts/distill/distill_wan.sh

@@ -0,0 +1,45 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+
+DATA_DIR=/workspace/data
+num_gpus=1
+IP=127.0.0.1
+
+torchrun --nnodes 1 --nproc_per_node $num_gpus \
+    --node_rank=0 \
+    --rdzv_id=456 \
+    --rdzv_backend=c10d \
+    --rdzv_endpoint=$IP:29500 \
+    fastvideo/distill_wan.py\
+    --seed 42\
+    --cache_dir "$DATA_DIR/.cache"\
+    --data_json_path "$DATA_DIR/HD-Mixkit-Finetune-Hunyuan/videos2caption.json"\
+    --validation_prompt_dir "$DATA_DIR/HD-Mixkit-Finetune-Hunyuan/validation"\
+    --train_batch_size=1 \
+    --num_latent_t 1 \
+    --sp_size $num_gpus \
+    --train_sp_batch_size 1\
+    --dataloader_num_workers 4\
+    --gradient_accumulation_steps=1\
+    --max_train_steps=320\
+    --learning_rate=1e-6\
+    --mixed_precision="bf16"\
+    --master_weight_type="bf16"\
+    --checkpointing_steps=64\
+    --validation_steps 64\
+    --validation_sampling_steps "2,4,8" \
+    --checkpoints_total_limit 3\
+    --allow_tf32\
+    --ema_start_step 0\
+    --cfg 0.0\
+    --log_validation\
+    --output_dir="$DATA_DIR/outputs/hy_phase1_shift17_bs_16_HD"\
+    --tracker_project_name Hunyuan_Distill \
+    --num_height 720 \
+    --num_width 1280 \
+    --num_frames  125 \
+    --shift 17 \
+    --validation_guidance_scale "1.0" \
+    --num_euler_timesteps 50 \
+    --multi_phased_distill_schedule "4000-1" \
+    --not_apply_cfg_solver 

scripts/preprocess/preprocess_wan_data.sh

@@ -0,0 +1,24 @@
+# export WANDB_MODE="offline"
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers"
+TEXT_ENCODER_PATH="/workspace/data/Wan-AI/Wan2.1-T2V-1.3B-Diffusers/tokenizer"
+MODEL_TYPE="wan"
+DATA_MERGE_PATH="/workspace/data/Mixkit-Src/merge.txt"
+OUTPUT_DIR="/workspace/data/HD-Mixkit-Finetune-Wan"
+VALIDATION_PATH="assets/prompt.txt"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/data_preprocess/preprocess.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --preprocess_video_batch_size=4 \
+    --preprocess_text_batch_size=4 \
+    --max_height=480 \
+    --max_width=832 \
+    --num_frames=81 \
+    --dataloader_num_workers 1 \
+    --output_dir=$OUTPUT_DIR \
+    --model_type $MODEL_TYPE \
+    --text_encoder_name $TEXT_ENCODER_PATH \
+    --train_fps 16 \
+    --validation_prompt_txt $VALIDATION_PATH