diff --git a/examples/train/run_slurm.sh b/examples/train/run_slurm.sh index 4e06a0b05..35eb837d6 100755 --- a/examples/train/run_slurm.sh +++ b/examples/train/run_slurm.sh @@ -107,7 +107,7 @@ srun torchrun \\ --node_rank \$SLURM_PROCID \\ --rdzv_backend=c10d \\ --rdzv_endpoint="\$MASTER_ADDR:\$MASTER_PORT" \\ - fastvideo/train/entrypoint/train.py \\ + -m fastvideo.train.entrypoint.train \\ --config ${CONFIG} \\ --training.distributed.num_gpus ${TOTAL_GPUS} \\ ${EXTRA_ARGS[*]:-} diff --git a/fastvideo/train/callbacks/ema.py b/fastvideo/train/callbacks/ema.py index 2bd9f01f0..8452a1995 100644 --- a/fastvideo/train/callbacks/ema.py +++ b/fastvideo/train/callbacks/ema.py @@ -47,9 +47,11 @@ class EMACallback(Callback): *, decay: float = 0.9999, start_iter: int = 0, + update_interval: int = 1, ) -> None: self._decay = float(decay) self._start_iter = int(start_iter) + self._update_interval = max(1, int(update_interval)) self._ema_started = False self.student_ema: EMA_FSDP | None = None @@ -78,9 +80,10 @@ class EMACallback(Callback): ) logger.info( "EMA callback enabled (decay=%s, " - "start_iter=%d).", + "start_iter=%d, update_interval=%d).", self._decay, self._start_iter, + self._update_interval, ) def on_training_step_end( @@ -94,6 +97,8 @@ class EMACallback(Callback): if iteration < self._start_iter: return + if (iteration - self._start_iter) % self._update_interval != 0: + return if not self._ema_started: logger.info( "Starting EMA updates at iteration %d " diff --git a/fastvideo/train/methods/rl/utils/evaluation.py b/fastvideo/train/methods/rl/utils/evaluation.py index 1c10bf64b..6659b99b7 100644 --- a/fastvideo/train/methods/rl/utils/evaluation.py +++ b/fastvideo/train/methods/rl/utils/evaluation.py @@ -40,6 +40,8 @@ def eval_once( rank: int, is_main_process: bool, tracker: Any | None = None, + max_batches: int | None = None, + seed: int = 0, ) -> dict[str, float]: """Run evaluation on test set. @@ -65,11 +67,14 @@ def eval_once( ctx = nullcontext() with ctx: + generator = torch.Generator(device=device) for batch_idx, ( _epoch_tag, prompts, metadata, ) in enumerate(test_dataloader): + if max_batches is not None and batch_idx >= max_batches: + break prompt_embeds = compute_text_embeddings( prompts, text_encoder, @@ -77,8 +82,12 @@ def eval_once( max_sequence_length=512, device=device, ) + neg_prompt_embeds = sample_neg_prompt_embeds[ + : len(prompts) + ] with torch.no_grad(): + generator.manual_seed(seed + batch_idx) ( videos, _latents, @@ -90,13 +99,14 @@ def eval_once( scheduler, prompt_embeds=prompt_embeds, negative_prompt_embeds=( - sample_neg_prompt_embeds + neg_prompt_embeds ), num_inference_steps=eval_num_steps, guidance_scale=eval_guidance_scale, height=height, width=width, num_frames=num_frames, + generator=generator, deterministic=True, sde_type="flow_sde", ) diff --git a/fastvideo/train/methods/rl/utils/pipeline.py b/fastvideo/train/methods/rl/utils/pipeline.py index 197ec7989..a4a79260e 100644 --- a/fastvideo/train/methods/rl/utils/pipeline.py +++ b/fastvideo/train/methods/rl/utils/pipeline.py @@ -280,9 +280,8 @@ def wan_denoising_with_logprob( encoder_hidden_states=prompt_embeds, return_dict=False, ) - ref_noise = ref_noise.to(dtype) - if do_cfg: - with ref_ctx: + ref_noise = ref_noise.to(dtype) + if do_cfg: ref_uncond = ref_model( hidden_states=latents_ori.to( dtype @@ -293,11 +292,11 @@ def wan_denoising_with_logprob( ), return_dict=False, ) - ref_noise = ( - ref_uncond - + guidance_scale - * (ref_noise - ref_uncond) - ) + ref_noise = ( + ref_uncond + + guidance_scale + * (ref_noise - ref_uncond) + ) ( _, diff --git a/fastvideo/train/methods/rl/utils/sampling.py b/fastvideo/train/methods/rl/utils/sampling.py index 8f68730b2..09a635950 100644 --- a/fastvideo/train/methods/rl/utils/sampling.py +++ b/fastvideo/train/methods/rl/utils/sampling.py @@ -4,6 +4,7 @@ compute rewards.""" from __future__ import annotations +import hashlib import time from collections.abc import Callable from typing import Any @@ -31,8 +32,15 @@ def create_generator( """Create deterministic generators seeded by prompt.""" generators = [] for prompt in prompts: + prompt_seed = int.from_bytes( + hashlib.blake2b( + prompt.encode("utf-8"), + digest_size=8, + ).digest(), + "big", + ) g = torch.Generator(device=device) - g.manual_seed(base_seed + hash(prompt) % (2**31)) + g.manual_seed(base_seed + prompt_seed % (2**31)) generators.append(g) return generators @@ -72,13 +80,13 @@ def sample_epoch( ref_transformer: torch.nn.Module | None = None, lora_model: Any | None = None, tracker: Any | None = None, + async_reward_scoring: bool = True, ) -> tuple[ list[dict[str, Any]], list[torch.Tensor], list[list[str]], ]: - """Run one sampling epoch: generate videos, compute - rewards asynchronously. + """Run one sampling epoch: generate videos and compute rewards. Returns: Tuple of (samples, all_videos, all_prompts): @@ -131,7 +139,7 @@ def sample_epoch( if same_latent: gen = create_generator( prompts, - base_seed=epoch * SEED_EPOCH_STRIDE + i, + base_seed=seed + epoch * SEED_EPOCH_STRIDE + i, device=device, ) else: @@ -188,19 +196,25 @@ def sample_epoch( .repeat(sample_batch_size, 1) ) + videos_cpu = videos.detach().cpu() + # Collect decoded videos and prompts for logging. - all_videos.append(videos) + all_videos.append(videos_cpu) all_prompts.append(list(prompts)) - # Async reward computation. - rewards_future = executor.submit( - reward_fn, - videos, - prompts, - prompt_metadata, - True, - ) - time.sleep(0) + if async_reward_scoring: + rewards = executor.submit( + reward_fn, + videos_cpu, + prompts, + prompt_metadata, + True, + ) + time.sleep(0) + else: + rewards = (videos_cpu, list(prompts), prompt_metadata) + + del videos logger.info( "[sample_epoch] batch %d/%d: " @@ -225,7 +239,7 @@ def sample_epoch( "next_latents": latents[:, 1:], "log_probs": log_probs, "kl": kl, - "rewards": rewards_future, + "rewards": rewards, } ) @@ -233,7 +247,14 @@ def sample_epoch( torch.cuda.synchronize() _t_reward_wait = time.perf_counter() for sample in samples: - rewards, _ = sample["rewards"].result() + if async_reward_scoring: + rewards, _ = sample["rewards"].result() + else: + videos_cpu, prompts, prompt_metadata = sample["rewards"] + torch.cuda.empty_cache() + rewards, _ = reward_fn( + videos_cpu, prompts, prompt_metadata, True + ) sample["rewards"] = { key: torch.as_tensor(value, device=device).float() for key, value in rewards.items() diff --git a/fastvideo/train/methods/rl/utils/sde.py b/fastvideo/train/methods/rl/utils/sde.py index 1f676f598..e11cebf55 100644 --- a/fastvideo/train/methods/rl/utils/sde.py +++ b/fastvideo/train/methods/rl/utils/sde.py @@ -115,16 +115,22 @@ def sde_step_with_logprob( if deterministic: prev_sample = sample + dt * model_output - log_prob = ( - -( - (prev_sample.detach() - prev_sample_mean) ** 2 + std_scale = std_dev_t * torch.sqrt(-1 * dt) + if torch.all(std_scale == 0): + log_prob = torch.zeros_like(prev_sample) + else: + std_scale = torch.clamp( + std_scale, + min=math.sqrt(torch.finfo(std_scale.dtype).tiny), ) - / (2 * ((std_dev_t * torch.sqrt(-1 * dt)) ** 2)) - - torch.log(std_dev_t * torch.sqrt(-1 * dt)) - - torch.log( - torch.sqrt(2 * torch.as_tensor(math.pi)) + log_prob = ( + -((prev_sample.detach() - prev_sample_mean) ** 2) + / (2 * (std_scale**2)) + - torch.log(std_scale) + - torch.log( + torch.sqrt(2 * torch.as_tensor(math.pi)) + ) ) - ) elif sde_type == "flow_cps": std_dev_t = sigma_prev * math.sin(