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

text preprocessing ready

Browse Source
This commit is contained in:
RandNMR73
2025-09-20 05:00:53 +00:00
committed by JerryZhou54
parent 1f78d8114e
commit ddf5e76ceb
3 changed files with 64 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/distill/SFWan2.1-T2V/distill_dmd_t2v_1.3B.sh
+3 -3
View File
@@ -18,20 +18,20 @@ export TORCH_NCCL_ENABLE_MONITORING=0
export TRITON_CACHE_DIR=/tmp/triton_cache_${SLURM_PROCID}
export MASTER_PORT=29503
export TOKENIZERS_PARALLELISM=false
export WANDB_API_KEY="50632ebd88ffd970521cec9ab4a1a2d7e85bfc45"
export WANDB_API_KEY="2f25ad37933894dbf0966c838c0b8494987f9f2f"
export WANDB_BASE_URL="https://api.wandb.ai"
export WANDB_MODE=online
export FASTVIDEO_ATTENTION_BACKEND=FLASH_ATTN
# Configs
NUM_GPUS=4
NUM_GPUS=1
# Model paths for Self-Forcing DMD distillation:
GENERATOR_MODEL_PATH="wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers"
REAL_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-14B-Diffusers" # Teacher model
FAKE_SCORE_MODEL_PATH="Wan-AI/Wan2.1-T2V-1.3B-Diffusers" # Critic model
DATA_DIR="data/crush-smol_processed_t2v/combined_parquet_dataset/"
DATA_DIR="data/test-text-preprocessing/Node_0_GPU_1_File_1/combined_parquet_dataset/"
VALIDATION_DATASET_FILE="data/crush-smol-single_processed_t2v/validation.json"
# export CUDA_VISIBLE_DEVICES=4,5
# IP=[MASTER NODE IP]
fastvideo/training/self_forcing_distillation_pipeline.py
+37 -1
View File
@@ -23,7 +23,7 @@ from fastvideo.training.training_utils import (
)
from fastvideo.models.utils import pred_noise_to_pred_video
from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import SelfForcingFlowMatchScheduler
from fastvideo.distributed import get_world_group
from fastvideo.distributed import get_world_group, get_local_torch_device
import torch.distributed as dist
import numpy as np
from fastvideo.utils import set_random_seed, is_vsa_available
@@ -579,6 +579,42 @@ class SelfForcingDistillationPipeline(DistillationPipeline):
assert cache_dict["k"].shape[2] == num_attention_heads, f"Attention heads mismatch in crossattn_cache block {i}"
assert cache_dict["k"].shape[3] == attention_head_dim, f"Attention head dim mismatch in crossattn_cache block {i}"
def _get_next_batch(self, training_batch: TrainingBatch) -> TrainingBatch:
batch = next(self.train_loader_iter, None) # type: ignore
if batch is None:
self.current_epoch += 1
logger.info("Starting epoch %s", self.current_epoch)
# Reset iterator for next epoch
self.train_loader_iter = iter(self.train_dataloader)
# Get first batch of new epoch
batch = next(self.train_loader_iter)
# latents, encoder_hidden_states, encoder_attention_mask, infos = batch
encoder_hidden_states = batch['text_embedding']
encoder_attention_mask = batch['text_attention_mask']
infos = batch['info_list']
batch_size = encoder_hidden_states.shape[0]
vae_config = self.training_args.pipeline_config.vae_config.arch_config
num_channels = vae_config.z_dim
spatial_compression_ratio = vae_config.spatial_compression_ratio
latent_height = self.training_args.num_height // spatial_compression_ratio
latent_width = self.training_args.num_width // spatial_compression_ratio
latents = torch.randn(batch_size, num_channels, self.training_args.num_latent_t,
latent_height, latent_width).to(get_local_torch_device(), dtype=torch.bfloat16)
training_batch.latents = latents.to(get_local_torch_device(),
dtype=torch.bfloat16)
training_batch.encoder_hidden_states = encoder_hidden_states.to(
get_local_torch_device(), dtype=torch.bfloat16)
training_batch.encoder_attention_mask = encoder_attention_mask.to(
get_local_torch_device(), dtype=torch.bfloat16)
training_batch.infos = infos
return training_batch
def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch:
"""
Self-forcing training step that alternates between generator and critic training.
fastvideo/training/training_pipeline.py
+24 -21
View File
@@ -22,7 +22,7 @@ from tqdm.auto import tqdm
import fastvideo.envs as envs
from fastvideo.attention.backends.video_sparse_attn import (
VideoSparseAttentionMetadataBuilder)
from fastvideo.attention.backends.vmoba import VideoMobaAttentionMetadataBuilder
# from fastvideo.attention.backends.vmoba import VideoMobaAttentionMetadataBuilder
from fastvideo.configs.sample import SamplingParam
from fastvideo.dataset import build_parquet_map_style_dataloader
from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v
@@ -42,13 +42,15 @@ from fastvideo.training.training_utils import (
compute_density_for_timestep_sampling, get_scheduler, get_sigmas,
load_checkpoint, normalize_dit_input, save_checkpoint,
shard_latents_across_sp)
from fastvideo.utils import (is_vmoba_available, is_vsa_available,
# from fastvideo.utils import (is_vmoba_available, is_vsa_available,
# set_random_seed, shallow_asdict)
from fastvideo.utils import (is_vsa_available,
set_random_seed, shallow_asdict)
import wandb # isort: skip
vsa_available = is_vsa_available()
vmoba_available = is_vmoba_available()
# vmoba_available = is_vmoba_available()
logger = init_logger(__name__)
@@ -279,20 +281,20 @@ class TrainingPipeline(LoRAPipeline, ABC):
patch_size=patch_size,
VSA_sparsity=current_vsa_sparsity,
device=get_local_torch_device())
elif vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN":
moba_params = self.training_args.moba_config.copy()
moba_params.update({
"current_timestep":
training_batch.timesteps,
"raw_latent_shape":
training_batch.raw_latent_shape[2:5],
"patch_size":
self.training_args.pipeline_config.dit_config.patch_size,
"device":
get_local_torch_device(),
})
training_batch.attn_metadata = VideoMobaAttentionMetadataBuilder(
).build(**moba_params)
# elif vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN":
# moba_params = self.training_args.moba_config.copy()
# moba_params.update({
# "current_timestep":
# training_batch.timesteps,
# "raw_latent_shape":
# training_batch.raw_latent_shape[2:5],
# "patch_size":
# self.training_args.pipeline_config.dit_config.patch_size,
# "device":
# get_local_torch_device(),
# })
# training_batch.attn_metadata = VideoMobaAttentionMetadataBuilder(
# ).build(**moba_params)
else:
training_batch.attn_metadata = None
@@ -317,7 +319,8 @@ class TrainingPipeline(LoRAPipeline, ABC):
def _transformer_forward_and_compute_loss(
self, training_batch: TrainingBatch) -> TrainingBatch:
if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN" or vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN":
# if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN" or vmoba_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VMOBA_ATTN":
if vsa_available and envs.FASTVIDEO_ATTENTION_BACKEND == "VIDEO_SPARSE_ATTN":
assert training_batch.attn_metadata is not None
else:
assert training_batch.attn_metadata is None
@@ -483,9 +486,9 @@ class TrainingPipeline(LoRAPipeline, ABC):
current_decay_times = min(step // vsa_decay_interval_steps,
vsa_sparsity // vsa_decay_rate)
current_vsa_sparsity = current_decay_times * vsa_decay_rate
elif vmoba_available:
# TODO: add vmoba sparsity scheduling here
current_vsa_sparsity = 0.0
# elif vmoba_available:
# # TODO: add vmoba sparsity scheduling here
# pass
else:
current_vsa_sparsity = 0.0