diff --git a/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml b/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml index 73783b023..3965f8a72 100644 --- a/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml +++ b/examples/train/configs/genrl_wan2.1_t2v_1.3B_longcat.yaml @@ -115,8 +115,7 @@ training: enable_gradient_checkpointing_type: full callbacks: - grad_clip: - max_grad_norm: 0.0 # Disabled; GenRLMethod clips internally. + # Gradnorm call back Disabled; GenRLMethod clips internally. ema: decay: 0.9 start_iter: 0 diff --git a/fastvideo/train/methods/rl/genrl.py b/fastvideo/train/methods/rl/genrl.py index 73eba7c43..24cc80243 100644 --- a/fastvideo/train/methods/rl/genrl.py +++ b/fastvideo/train/methods/rl/genrl.py @@ -31,27 +31,27 @@ from fastvideo.train.methods.base import ( LogScalar, TrainingMethod, ) -from fastvideo.train.methods.rl.advantages import ( +from fastvideo.train.methods.rl.utils.advantages import ( compute_advantages, ) -from fastvideo.train.methods.rl.data import ( +from fastvideo.train.methods.rl.utils.data import ( build_prompt_dataloaders, ) -from fastvideo.train.methods.rl.diffusion import ( +from fastvideo.train.methods.rl.utils.diffusion import ( compute_log_prob, ) -from fastvideo.train.methods.rl.embeddings import ( +from fastvideo.train.methods.rl.utils.embeddings import ( compute_text_embeddings, ) -from fastvideo.train.methods.rl.rewards import ( +from fastvideo.train.methods.rl.utils.rewards import ( move_reward_models, multi_score, reward_models_on_device, ) -from fastvideo.train.methods.rl.sampling import ( +from fastvideo.train.methods.rl.utils.sampling import ( sample_epoch, ) -from fastvideo.train.methods.rl.stat_tracking import ( +from fastvideo.train.methods.rl.utils.stat_tracking import ( PerPromptStatTracker, ) from fastvideo.train.models.base import ModelBase diff --git a/fastvideo/train/methods/rl/utils/__init__.py b/fastvideo/train/methods/rl/utils/__init__.py new file mode 100644 index 000000000..e4282003a --- /dev/null +++ b/fastvideo/train/methods/rl/utils/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Utility modules for RL training.""" diff --git a/fastvideo/train/methods/rl/advantages.py b/fastvideo/train/methods/rl/utils/advantages.py similarity index 99% rename from fastvideo/train/methods/rl/advantages.py rename to fastvideo/train/methods/rl/utils/advantages.py index f8fb09c56..b4d935af9 100644 --- a/fastvideo/train/methods/rl/advantages.py +++ b/fastvideo/train/methods/rl/utils/advantages.py @@ -8,7 +8,7 @@ from typing import Any import numpy as np from fastvideo.logger import init_logger -from fastvideo.train.methods.rl.stat_tracking import ( +from fastvideo.train.methods.rl.utils.stat_tracking import ( EPSILON, PerPromptStatTracker, ) diff --git a/fastvideo/train/methods/rl/data.py b/fastvideo/train/methods/rl/utils/data.py similarity index 100% rename from fastvideo/train/methods/rl/data.py rename to fastvideo/train/methods/rl/utils/data.py diff --git a/fastvideo/train/methods/rl/diffusion.py b/fastvideo/train/methods/rl/utils/diffusion.py similarity index 98% rename from fastvideo/train/methods/rl/diffusion.py rename to fastvideo/train/methods/rl/utils/diffusion.py index 3fc47639c..e43a8e438 100644 --- a/fastvideo/train/methods/rl/diffusion.py +++ b/fastvideo/train/methods/rl/utils/diffusion.py @@ -5,7 +5,7 @@ from __future__ import annotations import torch -from fastvideo.train.methods.rl.sde import ( +from fastvideo.train.methods.rl.utils.sde import ( sde_step_with_logprob, ) diff --git a/fastvideo/train/methods/rl/embeddings.py b/fastvideo/train/methods/rl/utils/embeddings.py similarity index 100% rename from fastvideo/train/methods/rl/embeddings.py rename to fastvideo/train/methods/rl/utils/embeddings.py diff --git a/fastvideo/train/methods/rl/evaluation.py b/fastvideo/train/methods/rl/utils/evaluation.py similarity index 96% rename from fastvideo/train/methods/rl/evaluation.py rename to fastvideo/train/methods/rl/utils/evaluation.py index d019c591b..1c10bf64b 100644 --- a/fastvideo/train/methods/rl/evaluation.py +++ b/fastvideo/train/methods/rl/utils/evaluation.py @@ -9,10 +9,10 @@ from typing import Any import torch from fastvideo.logger import init_logger -from fastvideo.train.methods.rl.embeddings import ( +from fastvideo.train.methods.rl.utils.embeddings import ( compute_text_embeddings, ) -from fastvideo.train.methods.rl.pipeline import ( +from fastvideo.train.methods.rl.utils.pipeline import ( wan_denoising_with_logprob, ) diff --git a/fastvideo/train/methods/rl/pipeline.py b/fastvideo/train/methods/rl/utils/pipeline.py similarity index 99% rename from fastvideo/train/methods/rl/pipeline.py rename to fastvideo/train/methods/rl/utils/pipeline.py index 59ee42ca7..197ec7989 100644 --- a/fastvideo/train/methods/rl/pipeline.py +++ b/fastvideo/train/methods/rl/utils/pipeline.py @@ -19,7 +19,7 @@ from fastvideo.logger import init_logger _pipeline_logger = init_logger(__name__) -from fastvideo.train.methods.rl.sde import ( +from fastvideo.train.methods.rl.utils.sde import ( sde_step_with_logprob, ) diff --git a/fastvideo/train/methods/rl/rewards.py b/fastvideo/train/methods/rl/utils/rewards.py similarity index 100% rename from fastvideo/train/methods/rl/rewards.py rename to fastvideo/train/methods/rl/utils/rewards.py diff --git a/fastvideo/train/methods/rl/sampling.py b/fastvideo/train/methods/rl/utils/sampling.py similarity index 98% rename from fastvideo/train/methods/rl/sampling.py rename to fastvideo/train/methods/rl/utils/sampling.py index cc663d0ef..8f68730b2 100644 --- a/fastvideo/train/methods/rl/sampling.py +++ b/fastvideo/train/methods/rl/utils/sampling.py @@ -11,10 +11,10 @@ from typing import Any import torch from fastvideo.logger import init_logger -from fastvideo.train.methods.rl.embeddings import ( +from fastvideo.train.methods.rl.utils.embeddings import ( compute_text_embeddings, ) -from fastvideo.train.methods.rl.pipeline import ( +from fastvideo.train.methods.rl.utils.pipeline import ( wan_denoising_with_logprob, ) diff --git a/fastvideo/train/methods/rl/sde.py b/fastvideo/train/methods/rl/utils/sde.py similarity index 100% rename from fastvideo/train/methods/rl/sde.py rename to fastvideo/train/methods/rl/utils/sde.py diff --git a/fastvideo/train/methods/rl/stat_tracking.py b/fastvideo/train/methods/rl/utils/stat_tracking.py similarity index 100% rename from fastvideo/train/methods/rl/stat_tracking.py rename to fastvideo/train/methods/rl/utils/stat_tracking.py