Compare commits
8
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 SolitaryThinker
|
445aaac585 | ||
|
SolitaryThinker
|
2276ad7d51 | ||
|
SolitaryThinker
|
9f0eacf35f | ||
|
SolitaryThinker
|
a9fe0b48c4 | ||
|
SolitaryThinker
|
661cac1a4a | ||
|
SolitaryThinker
|
dd94fe6139 | ||
|
SolitaryThinker
|
075bc69d5e | ||
|
William Lin
|
ae1d112c6a |
@@ -5,14 +5,13 @@ import random
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
import torch
|
||||
import tqdm
|
||||
from torch.utils.data import IterableDataset, get_worker_info
|
||||
from torchdata.stateful_dataloader import StatefulDataLoader
|
||||
|
||||
from fastvideo.v1.dataset.utils import collate_latents_embs_masks
|
||||
from fastvideo.v1.distributed import (get_sp_world_size, get_world_rank,
|
||||
get_world_size)
|
||||
from fastvideo.v1.distributed import (get_sp_world_size, get_world_group,
|
||||
get_world_rank, get_world_size)
|
||||
from fastvideo.v1.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -157,95 +156,123 @@ def shard_parquet_files_across_sp_groups_and_workers(
|
||||
# Check if sharding plan already exists
|
||||
sharding_info_dir = os.path.join(
|
||||
path, f"sharding_info_{num_sp_groups}_sp_groups_{num_workers}_workers")
|
||||
if os.path.exists(sharding_info_dir):
|
||||
logger.info("Sharding plan already exists")
|
||||
logger.info("Loading sharding plan from %s", sharding_info_dir)
|
||||
try:
|
||||
|
||||
# Only rank 0 handles cache checking and file scanning
|
||||
if get_world_rank() == 0:
|
||||
cache_loaded = False
|
||||
shard_parquet_files = None
|
||||
shard_total_samples = None
|
||||
shard_parquet_lengths = None
|
||||
|
||||
# First try to load existing sharding plan
|
||||
if os.path.exists(sharding_info_dir):
|
||||
logger.info("Loading sharding plan from %s", sharding_info_dir)
|
||||
try:
|
||||
with open(
|
||||
os.path.join(sharding_info_dir,
|
||||
"shard_parquet_files.pkl"), "rb") as f:
|
||||
shard_parquet_files = pickle.load(f)
|
||||
with open(
|
||||
os.path.join(sharding_info_dir,
|
||||
"shard_total_samples.pkl"), "rb") as f:
|
||||
shard_total_samples = pickle.load(f)
|
||||
with open(
|
||||
os.path.join(sharding_info_dir,
|
||||
"shard_parquet_lengths.pkl"), "rb") as f:
|
||||
shard_parquet_lengths = pickle.load(f)
|
||||
cache_loaded = True
|
||||
logger.info("Successfully loaded sharding plan")
|
||||
except Exception as e:
|
||||
logger.error("Error loading sharding plan: %s", str(e))
|
||||
logger.info("Falling back to creating new sharding plan")
|
||||
cache_loaded = False
|
||||
|
||||
# If cache not loaded (either doesn't exist or failed to load), create sharding plan
|
||||
if not cache_loaded:
|
||||
logger.info("Creating new sharding plan")
|
||||
logger.info("Scanning for parquet files in %s", path)
|
||||
|
||||
# Find all parquet files
|
||||
parquet_files = []
|
||||
|
||||
for root, _, files in os.walk(path):
|
||||
for file in files:
|
||||
if file.endswith('.parquet'):
|
||||
parquet_files.append(os.path.join(root, file))
|
||||
|
||||
if not parquet_files:
|
||||
raise ValueError("No parquet files found in %s", path)
|
||||
|
||||
# Calculate file lengths efficiently using a single pass
|
||||
logger.info("Calculating file lengths...")
|
||||
lengths = []
|
||||
for file in tqdm.tqdm(parquet_files, desc="Reading parquet files"):
|
||||
lengths.append(pq.ParquetFile(file).metadata.num_rows)
|
||||
|
||||
total_samples = sum(lengths)
|
||||
logger.info("Found %d files with %d total samples",
|
||||
len(parquet_files), total_samples)
|
||||
|
||||
# Sort files by length for better balancing
|
||||
sorted_indices = np.argsort(lengths)
|
||||
sorted_files = [parquet_files[i] for i in sorted_indices]
|
||||
sorted_lengths = [lengths[i] for i in sorted_indices]
|
||||
|
||||
# Create shards
|
||||
num_shards = num_sp_groups * num_workers
|
||||
shard_parquet_files = [[] for _ in range(num_shards)]
|
||||
shard_total_samples = [0] * num_shards
|
||||
shard_parquet_lengths = [{} for _ in range(num_shards)]
|
||||
|
||||
# Distribute files to shards using a greedy approach
|
||||
logger.info("Distributing files to shards...")
|
||||
for file, length in zip(reversed(sorted_files),
|
||||
reversed(sorted_lengths),
|
||||
strict=True):
|
||||
# Find shard with minimum current length
|
||||
target_shard = np.argmin(shard_total_samples)
|
||||
shard_parquet_files[target_shard].append(file)
|
||||
shard_total_samples[target_shard] += length
|
||||
shard_parquet_lengths[target_shard][file] = length
|
||||
#randomize each shard
|
||||
for shard in shard_parquet_files:
|
||||
random.seed(seed)
|
||||
random.shuffle(shard)
|
||||
|
||||
# Save the sharding plan
|
||||
os.makedirs(sharding_info_dir, exist_ok=True)
|
||||
with open(
|
||||
os.path.join(sharding_info_dir, "shard_parquet_files.pkl"),
|
||||
"rb") as f:
|
||||
shard_parquet_files = pickle.load(f)
|
||||
"wb") as f:
|
||||
pickle.dump(shard_parquet_files, f)
|
||||
with open(
|
||||
os.path.join(sharding_info_dir, "shard_total_samples.pkl"),
|
||||
"rb") as f:
|
||||
shard_total_samples = pickle.load(f)
|
||||
"wb") as f:
|
||||
pickle.dump(shard_total_samples, f)
|
||||
with open(
|
||||
os.path.join(sharding_info_dir,
|
||||
"shard_parquet_lengths.pkl"), "rb") as f:
|
||||
shard_parquet_lengths = pickle.load(f)
|
||||
return shard_parquet_files, shard_total_samples, shard_parquet_lengths
|
||||
except Exception as e:
|
||||
logger.error("Error loading sharding plan: %s", str(e))
|
||||
logger.info("Falling back to creating new sharding plan")
|
||||
"shard_parquet_lengths.pkl"), "wb") as f:
|
||||
pickle.dump(shard_parquet_lengths, f)
|
||||
logger.info("Saved sharding info to %s", sharding_info_dir)
|
||||
|
||||
if get_world_rank() == 0:
|
||||
logger.info("Scanning for parquet files in %s", path)
|
||||
# Wait for rank 0 to finish creating/loading sharding plan
|
||||
world_group = get_world_group()
|
||||
world_group.barrier()
|
||||
|
||||
# Find all parquet files
|
||||
parquet_files = []
|
||||
# Now all ranks load the sharding plan (it should exist and be valid now)
|
||||
logger.info("Loading sharding plan from %s after barrier",
|
||||
sharding_info_dir)
|
||||
with open(os.path.join(sharding_info_dir, "shard_parquet_files.pkl"),
|
||||
"rb") as f:
|
||||
shard_parquet_files = pickle.load(f)
|
||||
with open(os.path.join(sharding_info_dir, "shard_total_samples.pkl"),
|
||||
"rb") as f:
|
||||
shard_total_samples = pickle.load(f)
|
||||
with open(os.path.join(sharding_info_dir, "shard_parquet_lengths.pkl"),
|
||||
"rb") as f:
|
||||
shard_parquet_lengths = pickle.load(f)
|
||||
|
||||
for root, _, files in os.walk(path):
|
||||
for file in files:
|
||||
if file.endswith('.parquet'):
|
||||
parquet_files.append(os.path.join(root, file))
|
||||
|
||||
if not parquet_files:
|
||||
raise ValueError("No parquet files found in %s", path)
|
||||
|
||||
# Calculate file lengths efficiently using a single pass
|
||||
logger.info("Calculating file lengths...")
|
||||
lengths = []
|
||||
for file in tqdm.tqdm(parquet_files, desc="Reading parquet files"):
|
||||
lengths.append(pq.ParquetFile(file).metadata.num_rows)
|
||||
|
||||
total_samples = sum(lengths)
|
||||
logger.info("Found %d files with %d total samples", len(parquet_files),
|
||||
total_samples)
|
||||
|
||||
# Sort files by length for better balancing
|
||||
sorted_indices = np.argsort(lengths)
|
||||
sorted_files = [parquet_files[i] for i in sorted_indices]
|
||||
sorted_lengths = [lengths[i] for i in sorted_indices]
|
||||
|
||||
# Create shards
|
||||
num_shards = num_sp_groups * num_workers
|
||||
shard_parquet_files = [[] for _ in range(num_shards)]
|
||||
shard_total_samples = [0] * num_shards
|
||||
shard_parquet_lengths = [{} for _ in range(num_shards)]
|
||||
|
||||
# Distribute files to shards using a greedy approach
|
||||
logger.info("Distributing files to shards...")
|
||||
for file, length in zip(reversed(sorted_files),
|
||||
reversed(sorted_lengths),
|
||||
strict=True):
|
||||
# Find shard with minimum current length
|
||||
target_shard = np.argmin(shard_total_samples)
|
||||
shard_parquet_files[target_shard].append(file)
|
||||
shard_total_samples[target_shard] += length
|
||||
shard_parquet_lengths[target_shard][file] = length
|
||||
#randomize each shard
|
||||
for shard in shard_parquet_files:
|
||||
random.seed(seed)
|
||||
random.shuffle(shard)
|
||||
|
||||
save_dir = os.path.join(
|
||||
path,
|
||||
f"sharding_info_{num_sp_groups}_sp_groups_{num_workers}_workers")
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
with open(os.path.join(save_dir, "shard_parquet_files.pkl"), "wb") as f:
|
||||
pickle.dump(shard_parquet_files, f)
|
||||
with open(os.path.join(save_dir, "shard_total_samples.pkl"), "wb") as f:
|
||||
pickle.dump(shard_total_samples, f)
|
||||
with open(os.path.join(save_dir, "shard_parquet_lengths.pkl"),
|
||||
"wb") as f:
|
||||
pickle.dump(shard_parquet_lengths, f)
|
||||
logger.info("Saved sharding info to %s", save_dir)
|
||||
|
||||
# wait for all ranks to finish
|
||||
torch.distributed.barrier()
|
||||
# recursive call
|
||||
return shard_parquet_files_across_sp_groups_and_workers(
|
||||
path, num_sp_groups, num_workers, seed)
|
||||
return shard_parquet_files, shard_total_samples, shard_parquet_lengths
|
||||
|
||||
|
||||
def build_parquet_iterable_style_dataloader(
|
||||
|
||||
@@ -13,8 +13,8 @@ from torch.utils.data import Dataset, Sampler
|
||||
from torchdata.stateful_dataloader import StatefulDataLoader
|
||||
|
||||
from fastvideo.v1.dataset.utils import collate_rows_from_parquet_schema
|
||||
from fastvideo.v1.distributed import (get_sp_world_size, get_world_rank,
|
||||
get_world_size)
|
||||
from fastvideo.v1.distributed import (get_sp_world_size, get_world_group,
|
||||
get_world_rank, get_world_size)
|
||||
from fastvideo.v1.logger import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
@@ -97,48 +97,64 @@ def get_parquet_files_and_length(path: str):
|
||||
cache_dir = os.path.join(path, "map_style_cache")
|
||||
cache_file = os.path.join(cache_dir, "file_info.pkl")
|
||||
|
||||
if os.path.exists(cache_file):
|
||||
logger.info("Loading cached file info from %s", cache_file)
|
||||
try:
|
||||
with open(cache_file, "rb") as f:
|
||||
file_names_sorted, lengths_sorted = pickle.load(f)
|
||||
return file_names_sorted, lengths_sorted
|
||||
except Exception as e:
|
||||
logger.error("Error loading cached file info: %s", str(e))
|
||||
logger.info("Falling back to scanning files")
|
||||
|
||||
# If no cache exists or loading failed, scan files
|
||||
# Only rank 0 checks for cache and scans files if needed
|
||||
if get_world_rank() == 0:
|
||||
lengths = []
|
||||
file_names = []
|
||||
for root, _, files in os.walk(path):
|
||||
for file in sorted(files):
|
||||
if file.endswith('.parquet'):
|
||||
file_path = os.path.join(root, file)
|
||||
file_names.append(file_path)
|
||||
for file_path in tqdm.tqdm(file_names,
|
||||
desc="Reading parquet files to get lengths"):
|
||||
num_rows = pq.ParquetFile(file_path).metadata.num_rows
|
||||
lengths.append(num_rows)
|
||||
# sort according to file name to ensure all rank has the same order (in case os.walk is not sorted)
|
||||
file_names_sorted, lengths_sorted = zip(*sorted(zip(file_names,
|
||||
lengths,
|
||||
strict=True),
|
||||
key=lambda x: x[0]),
|
||||
strict=True)
|
||||
assert len(
|
||||
file_names_sorted) != 0, "No parquet files found in the dataset"
|
||||
cache_loaded = False
|
||||
file_names_sorted = None
|
||||
lengths_sorted = None
|
||||
|
||||
os.makedirs(cache_dir, exist_ok=True)
|
||||
with open(cache_file, "wb") as f:
|
||||
pickle.dump((file_names_sorted, lengths_sorted), f)
|
||||
logger.info("Saved file info to %s", cache_file)
|
||||
# First try to load existing cache
|
||||
if os.path.exists(cache_file):
|
||||
logger.info("Loading cached file info from %s", cache_file)
|
||||
try:
|
||||
with open(cache_file, "rb") as f:
|
||||
file_names_sorted, lengths_sorted = pickle.load(f)
|
||||
cache_loaded = True
|
||||
logger.info("Successfully loaded cached file info")
|
||||
except Exception as e:
|
||||
logger.error("Error loading cached file info: %s", str(e))
|
||||
logger.info("Falling back to scanning files")
|
||||
cache_loaded = False
|
||||
|
||||
# Wait for rank 0 to finish saving
|
||||
if get_world_size() > 1:
|
||||
torch.distributed.barrier()
|
||||
# If cache not loaded (either doesn't exist or failed to load), scan files
|
||||
if not cache_loaded:
|
||||
logger.info("Scanning parquet files to get lengths")
|
||||
lengths = []
|
||||
file_names = []
|
||||
for root, _, files in os.walk(path):
|
||||
for file in sorted(files):
|
||||
if file.endswith('.parquet'):
|
||||
file_path = os.path.join(root, file)
|
||||
file_names.append(file_path)
|
||||
for file_path in tqdm.tqdm(
|
||||
file_names, desc="Reading parquet files to get lengths"):
|
||||
num_rows = pq.ParquetFile(file_path).metadata.num_rows
|
||||
lengths.append(num_rows)
|
||||
# sort according to file name to ensure all rank has the same order
|
||||
file_names_sorted, lengths_sorted = zip(*sorted(zip(file_names,
|
||||
lengths,
|
||||
strict=True),
|
||||
key=lambda x: x[0]),
|
||||
strict=True)
|
||||
assert len(
|
||||
file_names_sorted) != 0, "No parquet files found in the dataset"
|
||||
|
||||
return get_parquet_files_and_length(path)
|
||||
# Save the cache
|
||||
os.makedirs(cache_dir, exist_ok=True)
|
||||
with open(cache_file, "wb") as f:
|
||||
pickle.dump((file_names_sorted, lengths_sorted), f)
|
||||
logger.info("Saved file info to %s", cache_file)
|
||||
|
||||
# Wait for rank 0 to finish creating/loading cache
|
||||
world_group = get_world_group()
|
||||
world_group.barrier()
|
||||
|
||||
# Now all ranks load the cache (it should exist and be valid now)
|
||||
logger.info("Loading cached file info from %s after barrier", cache_file)
|
||||
with open(cache_file, "rb") as f:
|
||||
file_names_sorted, lengths_sorted = pickle.load(f)
|
||||
|
||||
return file_names_sorted, lengths_sorted
|
||||
|
||||
|
||||
def read_row_from_parquet_file(parquet_files: list[str], global_row_idx: int,
|
||||
@@ -153,17 +169,27 @@ def read_row_from_parquet_file(parquet_files: list[str], global_row_idx: int,
|
||||
'''
|
||||
# find the parquet file and local row index
|
||||
cumulative = 0
|
||||
file_index = 0
|
||||
local_row_idx = 0
|
||||
|
||||
for file_index in range(len(lengths)):
|
||||
if cumulative + lengths[file_index] > global_row_idx:
|
||||
local_row_idx = global_row_idx - cumulative
|
||||
break
|
||||
cumulative += lengths[file_index]
|
||||
else:
|
||||
# If we reach here, global_row_idx is out of bounds
|
||||
raise IndexError(
|
||||
f"global_row_idx {global_row_idx} is out of bounds for dataset")
|
||||
|
||||
parquet_file = pq.ParquetFile(parquet_files[file_index])
|
||||
|
||||
# Calculate the row group to read into memory and the local idx
|
||||
# This way we can avoid reading in the entire parquet file
|
||||
cumulative = 0
|
||||
row_group_index = 0
|
||||
local_index = 0
|
||||
|
||||
for i in range(parquet_file.num_row_groups):
|
||||
num_rows = parquet_file.metadata.row_group(i).num_rows
|
||||
if cumulative + num_rows > local_row_idx:
|
||||
@@ -171,6 +197,11 @@ def read_row_from_parquet_file(parquet_files: list[str], global_row_idx: int,
|
||||
local_index = local_row_idx - cumulative
|
||||
break
|
||||
cumulative += num_rows
|
||||
else:
|
||||
# If we reach here, local_row_idx is out of bounds for this parquet file
|
||||
raise IndexError(
|
||||
f"local_row_idx {local_row_idx} is out of bounds for parquet file {parquet_files[file_index]}"
|
||||
)
|
||||
|
||||
row_group = parquet_file.read_row_group(row_group_index).to_pydict()
|
||||
row_dict = {k: v[local_index] for k, v in row_group.items()}
|
||||
|
||||
@@ -323,6 +323,23 @@ class GroupCoordinator:
|
||||
return input_
|
||||
return self.device_communicator.gather(input_, dst, dim)
|
||||
|
||||
def gather_object(self, obj: Any, dst: int = 0) -> list[Any] | None:
|
||||
"""Gather the input object.
|
||||
NOTE: `dst` is the global rank of the destination rank.
|
||||
"""
|
||||
world_size = self.world_size
|
||||
if self.world_size == 1:
|
||||
return [obj]
|
||||
gather_list = None
|
||||
if dst == self.rank:
|
||||
gather_list = [None] * world_size
|
||||
|
||||
torch.distributed.gather_object(obj,
|
||||
gather_list,
|
||||
dst,
|
||||
group=self.cpu_group)
|
||||
return gather_list
|
||||
|
||||
def all_to_all_4D(self,
|
||||
input_: torch.Tensor,
|
||||
scatter_dim: int = 2,
|
||||
@@ -806,6 +823,7 @@ def get_dp_group() -> GroupCoordinator:
|
||||
def initialize_model_parallel(
|
||||
tensor_model_parallel_size: int = 1,
|
||||
sequence_model_parallel_size: int = 1,
|
||||
data_parallel_size: int = 1,
|
||||
backend: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
@@ -925,7 +943,6 @@ def maybe_init_distributed_environment_and_model_parallel(
|
||||
world_size = int(os.environ.get("WORLD_SIZE", 1))
|
||||
rank = int(os.environ.get("RANK", 0))
|
||||
device = torch.device(f"cuda:{local_rank}")
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
init_distributed_environment(
|
||||
world_size=world_size,
|
||||
@@ -935,6 +952,7 @@ def maybe_init_distributed_environment_and_model_parallel(
|
||||
device_id=device)
|
||||
initialize_model_parallel(tensor_model_parallel_size=tp_size,
|
||||
sequence_model_parallel_size=sp_size)
|
||||
torch.cuda.set_device(device)
|
||||
|
||||
|
||||
def model_parallel_is_initialized() -> bool:
|
||||
|
||||
@@ -327,7 +327,8 @@ class ImageProcessorLoader(ComponentLoader):
|
||||
"""Load the image processor based on the model path, and inference args."""
|
||||
logger.info("Loading image processor from %s", model_path)
|
||||
|
||||
image_processor = AutoImageProcessor.from_pretrained(model_path, )
|
||||
image_processor = AutoImageProcessor.from_pretrained(model_path,
|
||||
use_fast=True)
|
||||
logger.info("Loaded image processor: %s",
|
||||
image_processor.__class__.__name__)
|
||||
return image_processor
|
||||
|
||||
@@ -108,8 +108,7 @@ class DecodingStage(PipelineStage):
|
||||
# Normalize image to [0, 1] range
|
||||
image = (image / 2 + 0.5).clamp(0, 1)
|
||||
|
||||
# Convert to CPU float32 for compatibility
|
||||
image = image.cpu().float()
|
||||
image = image.float()
|
||||
|
||||
# Update batch with decoded image
|
||||
batch.output = image
|
||||
|
||||
@@ -1 +1 @@
|
||||
{"step_time":0.6983645600266755,"_wandb":{"runtime":107},"grad_norm":0.50390625,"avg_step_time":1.002151239803061,"_step":5,"validation_videos_50_steps":{"captions":false,"_type":"videos","count":8,"videos":[{"size":159131,"path":"media/videos/validation_videos_50_steps_0_dc447599dbe48350e9c9.mp4","_type":"video-file","sha256":"dc447599dbe48350e9c920f4971e1786bde580dd20d52b4aa147ae8d3dc564d6"},{"_type":"video-file","sha256":"4e283876ddfbf5a2cb6f8aca07a39f832b5f806fbefde8854bc73ed904ff20ee","size":160315,"path":"media/videos/validation_videos_50_steps_0_4e283876ddfbf5a2cb6f.mp4"},{"size":135225,"path":"media/videos/validation_videos_50_steps_0_78185c41e1935306e93c.mp4","_type":"video-file","sha256":"78185c41e1935306e93c2d416ee40b31d038abffb029cb5bfb11c2a634eb2fcf"},{"_type":"video-file","sha256":"27e9819d002d3f63c8918bbdc5bf2857b5effe0caf5d5b8374b9c590fc6432eb","size":197873,"path":"media/videos/validation_videos_50_steps_0_27e9819d002d3f63c891.mp4"},{"_type":"video-file","sha256":"46fe548e86144ca60a9396fcd15e8788d3a05c066c6819ccdd6a9041feaaec8f","size":170601,"path":"media/videos/validation_videos_50_steps_0_46fe548e86144ca60a93.mp4"},{"sha256":"91ec338774bec870b9c5c81be4330a8f0f2535124cb372e881c3703e6b65ed77","size":164462,"path":"media/videos/validation_videos_50_steps_0_91ec338774bec870b9c5.mp4","_type":"video-file"},{"_type":"video-file","sha256":"ee4e811080a619215fd7541b39203dfb69c2db4cdadbfdd7a234a2cff684f6f3","size":139435,"path":"media/videos/validation_videos_50_steps_0_ee4e811080a619215fd7.mp4"},{"sha256":"22e31e048ba5e5b9d6587d7306904d6edc6c618b8f77c3ceb05ba2d602309274","size":147072,"path":"media/videos/validation_videos_50_steps_0_22e31e048ba5e5b9d658.mp4","_type":"video-file"}]},"_timestamp":1.75118195270901e+09,"vsa_sparsity":0.05,"learning_rate":1e-05,"train_loss":0.19960195198655128,"_runtime":107.325113071}
|
||||
{"step_time":0.6983645600266755,"_wandb":{"runtime":107},"grad_norm":1.39390625,"avg_step_time":1.002151239803061,"_step":5,"validation_videos_50_steps":{"captions":false,"_type":"videos","count":8,"videos":[{"size":159131,"path":"media/videos/validation_videos_50_steps_0_dc447599dbe48350e9c9.mp4","_type":"video-file","sha256":"dc447599dbe48350e9c920f4971e1786bde580dd20d52b4aa147ae8d3dc564d6"},{"_type":"video-file","sha256":"4e283876ddfbf5a2cb6f8aca07a39f832b5f806fbefde8854bc73ed904ff20ee","size":160315,"path":"media/videos/validation_videos_50_steps_0_4e283876ddfbf5a2cb6f.mp4"},{"size":135225,"path":"media/videos/validation_videos_50_steps_0_78185c41e1935306e93c.mp4","_type":"video-file","sha256":"78185c41e1935306e93c2d416ee40b31d038abffb029cb5bfb11c2a634eb2fcf"},{"_type":"video-file","sha256":"27e9819d002d3f63c8918bbdc5bf2857b5effe0caf5d5b8374b9c590fc6432eb","size":197873,"path":"media/videos/validation_videos_50_steps_0_27e9819d002d3f63c891.mp4"},{"_type":"video-file","sha256":"46fe548e86144ca60a9396fcd15e8788d3a05c066c6819ccdd6a9041feaaec8f","size":170601,"path":"media/videos/validation_videos_50_steps_0_46fe548e86144ca60a93.mp4"},{"sha256":"91ec338774bec870b9c5c81be4330a8f0f2535124cb372e881c3703e6b65ed77","size":164462,"path":"media/videos/validation_videos_50_steps_0_91ec338774bec870b9c5.mp4","_type":"video-file"},{"_type":"video-file","sha256":"ee4e811080a619215fd7541b39203dfb69c2db4cdadbfdd7a234a2cff684f6f3","size":139435,"path":"media/videos/validation_videos_50_steps_0_ee4e811080a619215fd7.mp4"},{"sha256":"22e31e048ba5e5b9d6587d7306904d6edc6c618b8f77c3ceb05ba2d602309274","size":147072,"path":"media/videos/validation_videos_50_steps_0_22e31e048ba5e5b9d658.mp4","_type":"video-file"}]},"_timestamp":1.75118195270901e+09,"vsa_sparsity":0.05,"learning_rate":1e-05,"train_loss":0.16960195198655128,"_runtime":107.325113071}
|
||||
@@ -111,7 +111,7 @@ def test_distributed_training():
|
||||
'avg_step_time': 1.0,
|
||||
'grad_norm': 0.1,
|
||||
'step_time': 1.0,
|
||||
'train_loss': 0.001
|
||||
'train_loss': 0.01
|
||||
}
|
||||
|
||||
failures = []
|
||||
|
||||
@@ -1 +1 @@
|
||||
{"step_time":5.501357046999999,"grad_norm":0.384765625,"train_loss":0.07890288904309273,"avg_step_time":5.831571423200001}
|
||||
{"step_time":3.501357046999999,"grad_norm":0.384765625,"train_loss":0.05890288904309273,"avg_step_time":3.831571423200001}
|
||||
@@ -121,10 +121,10 @@ def test_distributed_training():
|
||||
wandb_summary = json.load(open(summary_file))
|
||||
|
||||
fields_and_thresholds = {
|
||||
'avg_step_time': 6.0,
|
||||
'avg_step_time': 3.0,
|
||||
'grad_norm': 0.3,
|
||||
'step_time': 6.0,
|
||||
'train_loss': 0.0025
|
||||
'step_time': 3.0,
|
||||
'train_loss': 0.01
|
||||
}
|
||||
|
||||
failures = []
|
||||
|
||||
@@ -11,7 +11,6 @@ from typing import Any
|
||||
import imageio
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision
|
||||
from diffusers import FlowMatchEulerDiscreteScheduler
|
||||
from diffusers.optimization import get_scheduler
|
||||
from einops import rearrange
|
||||
@@ -430,7 +429,8 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
|
||||
self.seed)
|
||||
logger.info("Initialized random seeds with seed: %s", self.seed)
|
||||
|
||||
self.noise_scheduler = FlowMatchEulerDiscreteScheduler()
|
||||
self.noise_scheduler = FlowMatchEulerDiscreteScheduler(
|
||||
shift=self.training_args.pipeline_config.flow_shift, )
|
||||
|
||||
if self.training_args.resume_from_checkpoint:
|
||||
self._resume_from_checkpoint()
|
||||
@@ -595,6 +595,47 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
|
||||
|
||||
logger.info("Starting validation")
|
||||
|
||||
# Setup validation
|
||||
sampling_param, validation_dataloader, validation_steps = self._setup_validation(
|
||||
training_args)
|
||||
transformer.eval()
|
||||
world_group = get_world_group()
|
||||
|
||||
# Process each validation step
|
||||
for num_inference_steps in validation_steps:
|
||||
logger.info("rank: %s: num_inference_steps: %s",
|
||||
self.global_rank,
|
||||
num_inference_steps,
|
||||
local_main_process_only=False)
|
||||
|
||||
# Run inference for this step
|
||||
local_videos, local_captions, final_video_shape = self._run_validation_step(
|
||||
sampling_param, training_args, validation_dataloader,
|
||||
num_inference_steps)
|
||||
|
||||
# Gather results from all ranks
|
||||
all_videos_gathered = world_group.gather(local_videos, dst=0, dim=0)
|
||||
# all_videos_gathered: [num_validation_videos * world_size, num_frames, height, width, 3]
|
||||
all_captions_gathered = world_group.gather_object(local_captions,
|
||||
dst=0)
|
||||
|
||||
# Log results (only on rank 0)
|
||||
if self.global_rank == 0:
|
||||
self._log_gathered_results(all_videos_gathered,
|
||||
all_captions_gathered,
|
||||
num_inference_steps, global_step,
|
||||
training_args, sampling_param)
|
||||
|
||||
world_group.barrier()
|
||||
# Re-enable gradients for training
|
||||
training_args.inference_mode = False
|
||||
transformer.train()
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def _setup_validation(
|
||||
self, training_args) -> tuple[SamplingParam, DataLoader, list[int]]:
|
||||
"""Setup validation parameters and data."""
|
||||
# Create sampling parameters if not provided
|
||||
sampling_param = SamplingParam.from_pretrained(training_args.model_path)
|
||||
|
||||
@@ -613,100 +654,135 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
|
||||
batch_size=None,
|
||||
num_workers=0)
|
||||
|
||||
transformer.eval()
|
||||
|
||||
validation_steps = training_args.validation_sampling_steps.split(",")
|
||||
validation_steps = [int(step) for step in validation_steps]
|
||||
validation_steps = [step for step in validation_steps if step > 0]
|
||||
# Log validation results for this step
|
||||
world_group = get_world_group()
|
||||
num_sp_groups = world_group.world_size // self.sp_group.world_size
|
||||
|
||||
# Process each validation prompt for each validation step
|
||||
for num_inference_steps in validation_steps:
|
||||
logger.info("rank: %s: num_inference_steps: %s",
|
||||
return sampling_param, validation_dataloader, validation_steps
|
||||
|
||||
def _run_validation_step(
|
||||
self, sampling_param, training_args, validation_dataloader,
|
||||
num_inference_steps
|
||||
) -> tuple[torch.Tensor, list[str], tuple[int, int, int, int, int]]:
|
||||
"""Run validation inference for one step."""
|
||||
step_video_tensors: list[torch.Tensor] = []
|
||||
step_captions: list[str] = []
|
||||
batch = None
|
||||
|
||||
for validation_batch in validation_dataloader:
|
||||
batch = self._prepare_validation_batch(sampling_param,
|
||||
training_args,
|
||||
validation_batch,
|
||||
num_inference_steps)
|
||||
logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
|
||||
self.global_rank,
|
||||
num_inference_steps,
|
||||
self.rank_in_sp_group,
|
||||
batch.prompt,
|
||||
local_main_process_only=False)
|
||||
step_videos: list[np.ndarray] = []
|
||||
step_captions: list[str] = []
|
||||
|
||||
for validation_batch in validation_dataloader:
|
||||
batch = self._prepare_validation_batch(sampling_param,
|
||||
training_args,
|
||||
validation_batch,
|
||||
num_inference_steps)
|
||||
logger.info("rank: %s: rank_in_sp_group: %s, batch.prompt: %s",
|
||||
self.global_rank,
|
||||
self.rank_in_sp_group,
|
||||
batch.prompt,
|
||||
local_main_process_only=False)
|
||||
assert batch.prompt is not None and isinstance(batch.prompt, str)
|
||||
step_captions.append(batch.prompt)
|
||||
|
||||
assert batch.prompt is not None and isinstance(
|
||||
batch.prompt, str)
|
||||
step_captions.append(batch.prompt)
|
||||
# Run validation inference
|
||||
output_batch = self.validation_pipeline.forward(
|
||||
batch, training_args)
|
||||
samples = output_batch.output
|
||||
logger.info("Samples device: %s", samples.device)
|
||||
|
||||
# Run validation inference
|
||||
output_batch = self.validation_pipeline.forward(
|
||||
batch, training_args)
|
||||
samples = output_batch.output
|
||||
if self.rank_in_sp_group != 0:
|
||||
continue
|
||||
|
||||
if self.rank_in_sp_group != 0:
|
||||
continue
|
||||
# Process outputs
|
||||
assert samples.shape[
|
||||
0] == 1, "validation samples should have batch size 1"
|
||||
video = rearrange(samples, "b c t h w -> b t h w c")
|
||||
video = video * 255
|
||||
step_video_tensors.append(video)
|
||||
|
||||
# Process outputs
|
||||
video = rearrange(samples, "b c t h w -> t b c h w")
|
||||
frames = []
|
||||
for x in video:
|
||||
x = torchvision.utils.make_grid(x, nrow=6)
|
||||
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
|
||||
frames.append((x * 255).numpy().astype(np.uint8))
|
||||
step_videos.append(frames)
|
||||
# ValidationDataset will always pad the dataset so that the number
|
||||
# of videos is a multiple of the number of sp groups. Each sp group
|
||||
# will have the same number of videos
|
||||
num_validation_videos = len(step_captions)
|
||||
assert batch is not None
|
||||
assert batch.height is not None
|
||||
assert batch.width is not None
|
||||
final_video_shape = (num_validation_videos, batch.num_frames,
|
||||
batch.height, batch.width, 3)
|
||||
logger.info("Final video shape: %s",
|
||||
final_video_shape,
|
||||
local_main_process_only=False)
|
||||
|
||||
torch.distributed.barrier()
|
||||
# Only sp_group leaders (rank_in_sp_group == 0) need to send their
|
||||
# results to global rank 0
|
||||
if self.rank_in_sp_group == 0:
|
||||
if self.global_rank == 0:
|
||||
# Global rank 0 collects results from all sp_group leaders
|
||||
all_videos = step_videos # Start with own results
|
||||
all_captions = step_captions
|
||||
# Collect validation results from all SP group leaders using
|
||||
# all_gather_object.
|
||||
# Prepare data for gathering - only SP group leaders have valid
|
||||
# data, other ranks have duplicate data and so we send empty data.
|
||||
if self.rank_in_sp_group == 0:
|
||||
# SP group leaders contribute their data
|
||||
local_videos = torch.cat(step_video_tensors, dim=0)
|
||||
local_captions = step_captions
|
||||
else:
|
||||
# Other ranks contribute empty data
|
||||
local_videos = torch.zeros(final_video_shape, device=self.device)
|
||||
local_captions = []
|
||||
|
||||
# Receive from other sp_group leaders
|
||||
for sp_group_idx in range(1, num_sp_groups):
|
||||
src_rank = sp_group_idx * self.sp_world_size # Global rank of other sp_group leaders
|
||||
recv_videos = world_group.recv_object(src=src_rank)
|
||||
recv_captions = world_group.recv_object(src=src_rank)
|
||||
all_videos.extend(recv_videos)
|
||||
all_captions.extend(recv_captions)
|
||||
return local_videos, local_captions, final_video_shape
|
||||
|
||||
video_filenames = []
|
||||
for i, (video, caption) in enumerate(
|
||||
zip(all_videos, all_captions, strict=True)):
|
||||
os.makedirs(training_args.output_dir, exist_ok=True)
|
||||
filename = os.path.join(
|
||||
training_args.output_dir,
|
||||
f"validation_step_{global_step}_inference_steps_{num_inference_steps}_video_{i}.mp4"
|
||||
)
|
||||
imageio.mimsave(filename, video, fps=sampling_param.fps)
|
||||
video_filenames.append(filename)
|
||||
def _log_gathered_results(self, all_videos_gathered, all_captions_gathered,
|
||||
num_inference_steps, global_step, training_args,
|
||||
sampling_param) -> None:
|
||||
"""Process and log gathered validation results."""
|
||||
assert all_videos_gathered is not None
|
||||
assert all_captions_gathered is not None
|
||||
assert len(all_captions_gathered) == get_world_group().world_size
|
||||
|
||||
logs = {
|
||||
f"validation_videos_{num_inference_steps}_steps": [
|
||||
wandb.Video(filename, caption=caption)
|
||||
for filename, caption in zip(
|
||||
video_filenames, all_captions, strict=True)
|
||||
]
|
||||
}
|
||||
wandb.log(logs, step=global_step)
|
||||
else:
|
||||
# Other sp_group leaders send their results to global rank 0
|
||||
world_group.send_object(step_videos, dst=0)
|
||||
world_group.send_object(step_captions, dst=0)
|
||||
torch.distributed.barrier()
|
||||
all_videos_chunked_by_rank = all_videos_gathered.chunk(
|
||||
get_world_group().world_size, dim=0)
|
||||
num_validation_videos = all_videos_chunked_by_rank[0].shape[0]
|
||||
assert num_validation_videos > 0, "mismatch in num_validation_videos and how many videos were gathered"
|
||||
|
||||
# Re-enable gradients for training
|
||||
training_args.inference_mode = False
|
||||
transformer.train()
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
# Flatten the gathered data (filter out empty contributions)
|
||||
all_sp_rank_0_videos = []
|
||||
all_sp_rank_0_captions = []
|
||||
for idx in range(0,
|
||||
get_world_group().world_size,
|
||||
self.sp_group.world_size):
|
||||
all_sp_rank_0_videos.append(all_videos_chunked_by_rank[idx])
|
||||
all_sp_rank_0_captions.extend(all_captions_gathered[idx])
|
||||
|
||||
all_videos_tensor = torch.cat(all_sp_rank_0_videos, dim=0)
|
||||
# all_videos_tensor: [num_validation_videos * num_sp_groups, num_frames, height, width, 3]
|
||||
assert len(all_videos_tensor.shape) == 5
|
||||
all_videos_tensor = all_videos_tensor.cpu()
|
||||
|
||||
all_videos_processed = []
|
||||
for video in all_videos_tensor:
|
||||
assert len(video.shape) == 4
|
||||
frames = []
|
||||
for frame in video:
|
||||
frames.append(frame.numpy().astype(np.uint8))
|
||||
all_videos_processed.append(frames)
|
||||
|
||||
all_captions = all_sp_rank_0_captions
|
||||
assert len(all_videos_processed) == len(all_captions), (
|
||||
f"mismatch in number of videos and captions: "
|
||||
f"{len(all_videos_processed)} != {len(all_captions)}")
|
||||
|
||||
# Save videos and log to wandb
|
||||
video_filenames = []
|
||||
for i, (video, caption) in enumerate(
|
||||
zip(all_videos_processed, all_captions, strict=True)):
|
||||
os.makedirs(training_args.output_dir, exist_ok=True)
|
||||
filename = os.path.join(
|
||||
training_args.output_dir,
|
||||
f"validation_step_{global_step}_inference_steps_{num_inference_steps}_video_{i}.mp4"
|
||||
)
|
||||
imageio.mimsave(filename, video, fps=sampling_param.fps)
|
||||
video_filenames.append(filename)
|
||||
|
||||
logs = {
|
||||
f"validation_videos_{num_inference_steps}_steps": [
|
||||
wandb.Video(filename, caption=caption) for filename, caption in
|
||||
zip(video_filenames, all_captions, strict=True)
|
||||
]
|
||||
}
|
||||
wandb.log(logs, step=global_step)
|
||||
|
||||
Reference in New Issue
Block a user