nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Preprocessing Stage: 1. Save to parquets periodically 2. Allow resuming from the middle 3. Doesn't support multi-gpu for now. Data Loader Stage: 1. Allow multi-gpu dataloader 2. Doesn't drop files even if number of parquet files is not divisible by num_gpus 3. Able to resume training

Browse Source
This commit is contained in:
JerryZhou54
2025-05-23 00:57:39 +00:00
parent 982bfcfdc8
commit 42a2f272d5
6 changed files with 183 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/data_preprocess/preprocess.py
+11
View File
@@ -77,6 +77,17 @@ if __name__ == "__main__":
default=8,
help="Batch size (per device) for the training dataloader.",
)
parser.add_argument(
"--samples_per_file",
type=int,
default=64
)
parser.add_argument(
"--flush_frequency",
type=int,
default=256,
help="how often to save to parquet files"
)
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)
fastvideo/v1/dataset/__init__.py
+2 -1
View File
@@ -6,7 +6,7 @@ from fastvideo.v1.dataset.t2v_datasets import T2V_dataset
from fastvideo.v1.dataset.transform import CenterCropResizeVideo, Normalize255, TemporalRandomCrop
def getdataset(args):
def getdataset(args, start_idx=0):
temporal_sample = TemporalRandomCrop(args.num_frames) # 16 x
norm_fun = Lambda(lambda x: 2.0 * x - 1.0)
resize_topcrop = [
@@ -33,6 +33,7 @@ def getdataset(args):
temporal_sample=temporal_sample,
tokenizer=tokenizer,
transform_topcrop=transform_topcrop,
start_idx=start_idx
)
raise NotImplementedError(args.dataset)
fastvideo/v1/dataset/parquet_datasets.py
+49 -19
View File
@@ -16,6 +16,7 @@ import pyarrow.parquet as pq
import torch
import numpy as np
from torch.utils.data import IterableDataset, get_worker_info
from torch import distributed as dist
# Path to your dataset
dataset_path = "/mnt/sharefs/users/hao.zhang/Vchitect-2M/Vchitect-2M-laten-93x512x512/train/"
@@ -38,7 +39,7 @@ class ParquetVideoTextDataset(IterableDataset):
self.parquet_files = []
for root, _, files in os.walk(self.path):
for file in files:
if file.endswith('.parquet') and split in file:
if file.endswith('.parquet'):
self.parquet_files.append(os.path.join(root, file))
# Sort files for consistent ordering
self.parquet_files.sort()
@@ -46,14 +47,20 @@ class ParquetVideoTextDataset(IterableDataset):
# Distribute files among workers
# drop last unenven files
print(f"Total files: {len(self.parquet_files)}")
if len(self.parquet_files) % (world_size * num_workers) != 0:
total_length = len(self.parquet_files)
length_after_drop = total_length - total_length % (world_size * num_workers)
self.parquet_files = self.parquet_files[:length_after_drop]
print(f"Total files after dropping: {len(self.parquet_files)}")
self.parquet_files = self.parquet_files[rank::world_size]
print(f"Found {len(self.parquet_files)} parquet files for rank {rank}")
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]}")
@@ -68,7 +75,18 @@ class ParquetVideoTextDataset(IterableDataset):
"""Open the next parquet file for reading."""
num_workers = get_worker_info().num_workers
worker_id = get_worker_info().id
worker_parquet_files = self.parquet_files[worker_id::num_workers]
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
@@ -170,7 +188,7 @@ class ParquetVideoTextDataset(IterableDataset):
# 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 = np.frombuffer(text_attention_mask_bytes, dtype=np.uint8).astype(np.bool_)
msk = msk.reshape(1, -1)
# Make array writable
msk = np.copy(msk)
@@ -238,7 +256,6 @@ if __name__ == "__main__":
# Initialize distributed training
if world_size > 1:
from torch import distributed as dist
dist.init_process_group(
backend="nccl",
init_method="env://",
@@ -270,11 +287,27 @@ if __name__ == "__main__":
pin_memory=True,
drop_last=True
)
# Example of how to load dataloader state
# if os.path.exists("/workspace/FastVideo/dataloader_state.pt"):
# dataloader_state = torch.load("/workspace/FastVideo/dataloader_state.pt")
# dataloader.load_state_dict(dataloader_state[rank])
# Warm-up with synchronization
if rank == 0:
print("Warming up...")
for i, (latents, embeddings, masks, infos) in enumerate(dataloader):
# Example of how to save dataloader state
# if i == 30:
# dist.barrier()
# local_data = {rank: dataloader.state_dict()}
# gathered_data = [None] * world_size
# dist.all_gather_object(gathered_data, local_data)
# if rank == 0:
# global_state_dict = {}
# for d in gathered_data:
# global_state_dict.update(d)
# torch.save(global_state_dict, "dataloader_state.pt")
assert torch.sum(masks[0]).item() == torch.count_nonzero(embeddings[0]).item() // 4096
if args.vae_debug:
from fastvideo.v1.fastvideo_args import FastVideoArgs
@@ -299,16 +332,13 @@ if __name__ == "__main__":
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)
# latents = latents.to(device)
# embeddings = embeddings.to(device)
if world_size > 1:
dist.barrier()
# Benchmark
if rank == 0:
fastvideo/v1/dataset/t2v_datasets.py
+3 -2
View File
@@ -73,7 +73,8 @@ def filter_resolution(h, w, max_h_div_w_ratio=17 / 16, min_h_div_w_ratio=8 / 16)
class T2V_dataset(Dataset):
def __init__(self, args, transform, temporal_sample, tokenizer, transform_topcrop):
def __init__(self, args, transform, temporal_sample, tokenizer, transform_topcrop, start_idx=0):
self.start_idx = start_idx
self.data = args.data_merge_path
self.num_frames = args.num_frames
self.train_fps = args.train_fps
@@ -322,5 +323,5 @@ class T2V_dataset(Dataset):
return cap_lists
def get_cap_list(self):
cap_lists = self.read_jsons(self.data)
cap_lists = self.read_jsons(self.data)[self.start_idx:]
return cap_lists
fastvideo/v1/pipelines/data_preprocess.py
+115 -157
View File
@@ -58,15 +58,27 @@ class PreprocessPipeline(ComposedPipelineBase):
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)
self.preprocess_video_and_text(fastvideo_args, args)
def preprocess_video(self, fastvideo_args: FastVideoArgs, args):
def preprocess_video_and_text(self, fastvideo_args: FastVideoArgs, args):
os.makedirs(args.output_dir, exist_ok=True)
# Create directory for combined data
combined_parquet_dir = os.path.join(args.output_dir, f"combined_parquet_dataset")
os.makedirs(combined_parquet_dir, exist_ok=True)
local_rank = int(os.getenv("RANK", 0))
world_size = int(os.getenv("WORLD_SIZE", 1))
# Get how many samples have already been processed
start_idx = 0
for root, _, files in os.walk(combined_parquet_dir):
for file in files:
if file.endswith('.parquet'):
table = pq.read_table(os.path.join(root, file))
start_idx += table.num_rows
# Loading dataset
train_dataset = getdataset(args)
sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=True)
train_dataset = getdataset(args, start_idx=start_idx)
sampler = DistributedSampler(train_dataset, rank=local_rank, num_replicas=world_size, shuffle=False)
train_dataloader = DataLoader(
train_dataset,
sampler=sampler,
@@ -74,11 +86,10 @@ class PreprocessPipeline(ComposedPipelineBase):
num_workers=args.dataloader_num_workers,
)
num_processed_samples = 0
# 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
@@ -88,6 +99,7 @@ class PreprocessPipeline(ComposedPipelineBase):
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)
num_processed_samples += len(valid_indices)
if not valid_indices:
continue
@@ -103,47 +115,11 @@ class PreprocessPipeline(ComposedPipelineBase):
# 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():
# Get corresponding captions for this batch
batch_captions = valid_data["text"]
batch = ForwardBatch(
data_type="video",
prompt=batch_captions,
@@ -154,38 +130,39 @@ class PreprocessPipeline(ComposedPipelineBase):
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
# 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:
save_pbar = tqdm(enumerate(valid_data["path"]), desc="Saving outputs", unit="item", leave=False)
for idx, video_path in save_pbar:
# Get the corresponding latent and info using video name
latent = self.latent_data[video_name]
info = self.video_data[video_name]
latent = latents[idx].cpu()
video_name = os.path.basename(video_path).split(".")[0]
height, width = valid_data["pixel_values"][idx].shape[-2:]
# 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()
text_attention_mask = prompt_attention_mask[idx].cpu().numpy().astype(np.uint8)
# Create record for Parquet dataset
record = {
@@ -200,17 +177,16 @@ class PreprocessPipeline(ComposedPipelineBase):
"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"],
"caption": valid_data["text"][idx],
"media_type": "video",
"width": info["width"],
"height": info["height"],
"num_frames": info["num_frames"],
"duration_sec": info["duration_sec"],
"fps": info["fps"],
"width": width,
"height": height,
"num_frames": latents[idx].shape[1],
"duration_sec": float(valid_data["duration"][idx]),
"fps": float(valid_data["fps"][idx]),
}
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")
@@ -245,89 +221,64 @@ class PreprocessPipeline(ComposedPipelineBase):
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
if num_processed_samples >= args.flush_frequency:
assert hasattr(self, 'all_tables') and self.all_tables
print(f"Combining {len(self.all_tables)} batches...")
combined_table = pa.concat_tables(self.all_tables)
assert len(combined_table) == num_processed_samples
print(f"Total samples collected: {len(combined_table)}")
# Calculate total number of chunks needed, discarding remainder
total_chunks = max(num_processed_samples // args.samples_per_file, 1)
print(f"Fixed samples per parquet file: {args.samples_per_file}")
print(f"Total number of parquet files: {total_chunks}")
print(f"Total samples to be processed: {total_chunks * args.samples_per_file} (discarding {num_processed_samples % args.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
print(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, args.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}")
num_processed_samples = 0
self.all_tables = []
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")
@@ -463,6 +414,13 @@ class PreprocessPipeline(ComposedPipelineBase):
# Create worker-specific subdirectory
worker_dir = os.path.join(output_dir, f"worker_{worker_id}")
os.makedirs(worker_dir, exist_ok=True)
# Check how many files there are already in the dir, and update i accordingly
num_parquets = 0
for root, _, files in os.walk(worker_dir):
for file in files:
if file.endswith('.parquet'):
num_parquets += 1
for i in range(start_idx, end_idx):
start_sample = i * samples_per_file
@@ -470,7 +428,7 @@ class PreprocessPipeline(ComposedPipelineBase):
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")
chunk_path = os.path.join(worker_dir, f"data_chunk_{i + num_parquets}.parquet")
temp_path = chunk_path + '.tmp'
try:
scripts/preprocess/preprocess_wan_data.sh
+3 -2
View File
@@ -12,7 +12,6 @@ torchrun --nproc_per_node=$GPU_NUM \
--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 \
@@ -21,4 +20,6 @@ torchrun --nproc_per_node=$GPU_NUM \
--model_type $MODEL_TYPE \
--text_encoder_name $TEXT_ENCODER_PATH \
--train_fps 16 \
--validation_prompt_txt $VALIDATION_PATH
--validation_prompt_txt $VALIDATION_PATH \
--samples_per_file 108 \
--flush_frequency 108