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

[feat] GenRL: add runtime and memory stability helpers (#1402)

Browse Source 
Co-authored-by: Davids048 <jundasu@ucsd.edu>
This commit is contained in:
Adam
2026-06-05 01:14:43 +00:00
committed by Davids048
co-authored by Davids048
parent cd69575095
commit e1be068c46
6 changed files with 76 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/train/run_slurm.sh
+1 -1
View File
@@ -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[*]:-}
fastvideo/train/callbacks/ema.py
+6 -1
View File
@@ -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 "
fastvideo/train/methods/rl/utils/evaluation.py
+11 -1
View File
@@ -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",
)
fastvideo/train/methods/rl/utils/pipeline.py
+7 -8
View File
@@ -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)
)
(
_,
fastvideo/train/methods/rl/utils/sampling.py
+37 -16
View File
@@ -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()
fastvideo/train/methods/rl/utils/sde.py
+14 -8
View File
@@ -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(