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

[misc] rename dp_size to hdsp_replicate_dim (#491)

Browse Source
This commit is contained in:
Zhang Peiyuan
2025-06-10 16:36:56 -07:00
committed by GitHub
parent 2a46902ecb
commit 0f2bbe71ac
9 changed files with 31 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/data_preprocess/preprocess_vae_latents.py
+1 -7
View File
@@ -104,13 +104,7 @@ if __name__ == "__main__":
default=None,
help="The output directory where the model predictions and checkpoints will be written.",
)
parser.add_argument(
"--logging_dir",
type=str,
default="logs",
help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
" *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
)
args = parser.parse_args()
main(args)
fastvideo/data_preprocess/v1_preprocess.py
-7
View File
@@ -97,13 +97,6 @@ if __name__ == "__main__":
default=None,
help="The output directory where the model predictions and checkpoints will be written.",
)
parser.add_argument(
"--logging_dir",
type=str,
default="logs",
help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
" *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
)
args = parser.parse_args()
main(args)
fastvideo/distill.py
-7
View File
@@ -671,13 +671,6 @@ if __name__ == "__main__":
help=("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
)
parser.add_argument(
"--logging_dir",
type=str,
default="logs",
help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
" *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
)
# optimizer & scheduler & Training
parser.add_argument("--num_train_epochs", type=int, default=100)
fastvideo/distill_adv.py
-7
View File
@@ -693,13 +693,6 @@ if __name__ == "__main__":
help=("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
)
parser.add_argument(
"--logging_dir",
type=str,
default="logs",
help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
" *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
)
# optimizer & scheduler & Training
parser.add_argument("--num_train_epochs", type=int, default=100)
fastvideo/train.py
+1 -7
View File
@@ -520,13 +520,7 @@ if __name__ == "__main__":
help=("Whether training should be resumed from a previous lora checkpoint. Use a path saved by"
' `--checkpointing_steps`, or `"latest"` to automatically select the last available checkpoint.'),
)
parser.add_argument(
"--logging_dir",
type=str,
default="logs",
help=("[TensorBoard](https://www.tensorflow.org/tensorboard) log directory. Will default to"
" *output_dir/runs/**CURRENT_DATETIME_HOSTNAME***."),
)
# optimizer & scheduler & Training
parser.add_argument("--num_train_epochs", type=int, default=100)
fastvideo/v1/fastvideo_args.py
+16 -27
View File
@@ -44,8 +44,8 @@ class FastVideoArgs:
num_gpus: int = 1
tp_size: int = -1
sp_size: int = -1
dp_size: int = 1
dp_shards: int = -1
hsdp_replicate_dim: int = 1
hsdp_shard_dim: int = -1
dist_timeout: Optional[int] = None # timeout for torch.distributed
# Video generation parameters
@@ -189,17 +189,15 @@ class FastVideoArgs:
help="The sequence parallelism size.",
)
parser.add_argument(
"--data-parallel-size",
"--dp-size",
"--hsdp-replicate-dim",
type=int,
default=FastVideoArgs.dp_size,
default=FastVideoArgs.hsdp_replicate_dim,
help="The data parallelism size.",
)
parser.add_argument(
"--data-parallel-shards",
"--dp-shards",
"--hsdp-shard-dim",
type=int,
default=FastVideoArgs.dp_shards,
default=FastVideoArgs.hsdp_shard_dim,
help="The data parallelism shards.",
)
parser.add_argument(
@@ -379,10 +377,6 @@ class FastVideoArgs:
kwargs[attr] = args.tensor_parallel_size
elif attr == 'sp_size' and hasattr(args, 'sequence_parallel_size'):
kwargs[attr] = args.sequence_parallel_size
elif attr == 'dp_size' and hasattr(args, 'data_parallel_size'):
kwargs[attr] = args.data_parallel_size
elif attr == 'dp_shards' and hasattr(args, 'data_parallel_shards'):
kwargs[attr] = args.data_parallel_shards
elif attr == 'flow_shift' and hasattr(args, 'shift'):
kwargs[attr] = args.shift
# Use getattr with default value from the dataclass for potentially missing attributes
@@ -397,19 +391,20 @@ class FastVideoArgs:
def check_fastvideo_args(self) -> None:
"""Validate inference arguments for consistency"""
if not self.inference_mode:
assert self.dp_size is not -1, "dp_size must be set for training"
assert self.dp_shards is not -1, "dp_shards must be set for training"
assert self.sp_size is not -1, "sp_size must be set for training"
assert self.hsdp_replicate_dim != -1, "hsdp_replicate_dim must be set for training"
assert self.hsdp_shard_dim != -1, "hsdp_shard_dim must be set for training"
assert self.sp_size != -1, "sp_size must be set for training"
if self.tp_size is -1:
if self.tp_size == -1:
self.tp_size = self.num_gpus
if self.sp_size is -1:
if self.sp_size == -1:
self.sp_size = self.num_gpus
if self.dp_shards is -1:
self.dp_shards = self.num_gpus
if self.hsdp_shard_dim == -1:
self.hsdp_shard_dim = self.num_gpus
assert self.sp_size <= self.num_gpus and self.num_gpus % self.sp_size == 0, "num_gpus must >= and be divisible by sp_size"
assert self.dp_size <= self.num_gpus and self.num_gpus % self.dp_size == 0, "num_gpus must >= and be divisible by dp_size"
assert self.dp_shards <= self.num_gpus and self.num_gpus % self.dp_shards == 0, "num_gpus must >= and be divisible by dp_shards"
assert self.hsdp_replicate_dim <= self.num_gpus and self.num_gpus % self.hsdp_replicate_dim == 0, "num_gpus must >= and be divisible by hsdp_replicate_dim"
assert self.hsdp_shard_dim <= self.num_gpus and self.num_gpus % self.hsdp_shard_dim == 0, "num_gpus must >= and be divisible by hsdp_shard_dim"
if self.num_gpus < max(self.tp_size, self.sp_size):
self.num_gpus = max(self.tp_size, self.sp_size)
@@ -541,7 +536,6 @@ class TrainingArgs(FastVideoArgs):
checkpoints_total_limit: int = 0
checkpointing_steps: int = 0
resume_from_checkpoint: bool = False
logging_dir: str = ""
# optimizer & scheduler
num_train_epochs: int = 0
@@ -600,11 +594,6 @@ class TrainingArgs(FastVideoArgs):
kwargs[attr] = args.sequence_parallel_size
elif attr == 'flow_shift' and hasattr(args, 'shift'):
kwargs[attr] = args.shift
elif attr == 'dp_size' and hasattr(args, 'data_parallel_size'):
kwargs[attr] = args.data_parallel_size
elif attr == 'dp_shards' and hasattr(args, 'data_parallel_shards'):
kwargs[attr] = args.data_parallel_shards
# Use getattr with default value from the dataclass for potentially missing attributes
else:
default_value = getattr(cls, attr, None)
if getattr(args, attr, default_value) is not None:
fastvideo/v1/models/loader/component_loader.py
+3 -3
View File
@@ -403,7 +403,7 @@ class TransformerLoader(ComponentLoader):
# Load the model using FSDP loader
logger.info("Loading model from %s, default_dtype: %s", cls_name,
default_dtype)
assert fastvideo_args.dp_shards is not None
assert fastvideo_args.hsdp_shard_dim is not None
model = maybe_load_fsdp_model(
model_cls=model_cls,
init_params={
@@ -412,8 +412,8 @@ class TransformerLoader(ComponentLoader):
},
weight_dir_list=safetensors_list,
device=get_torch_device(),
data_parallel_size=fastvideo_args.dp_size,
data_parallel_shards=fastvideo_args.dp_shards,
hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim,
hsdp_shard_dim=fastvideo_args.hsdp_shard_dim,
cpu_offload=fastvideo_args.use_cpu_offload,
fsdp_inference=fastvideo_args.use_fsdp_inference,
default_dtype=default_dtype,
fastvideo/v1/models/loader/fsdp_load.py
+8 -6
View File
@@ -60,8 +60,8 @@ def maybe_load_fsdp_model(
init_params: Dict[str, Any],
weight_dir_list: List[str],
device: torch.device,
data_parallel_size: int,
data_parallel_shards: int,
hsdp_replicate_dim: int,
hsdp_shard_dim: int,
default_dtype: torch.dtype,
param_dtype: torch.dtype,
reduce_dtype: torch.dtype,
@@ -87,13 +87,15 @@ def maybe_load_fsdp_model(
with set_default_dtype(default_dtype), torch.device("meta"):
model = model_cls(**init_params)
dp_size = data_parallel_size if fsdp_inference or training_mode else 1
world_size = hsdp_replicate_dim * hsdp_shard_dim
if not training_mode and not fsdp_inference:
hsdp_replicate_dim = world_size
hsdp_shard_dim = 1
device_mesh = init_device_mesh(
"cuda",
# (Replicate(), Shard(dim=0))
mesh_shape=(dp_size, data_parallel_shards),
mesh_dim_names=("dp_replicate", "dp_shards"),
mesh_shape=(hsdp_replicate_dim, hsdp_shard_dim),
mesh_dim_names=("replicate", "shard"),
)
shard_model(model,
cpu_offload=cpu_offload,
scripts/finetune/finetune_v1.sh
+2 -2
View File
@@ -21,8 +21,8 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
--num_latent_t 20 \
--sp_size 4 \
--tp_size 4 \
--dp_size 1 \
--dp_shards 4 \
--hsdp_replicate_dim 1 \
--hsdp_shards 4 \
--num_gpus $NUM_GPUS \
--train_sp_batch_size 1\
--dataloader_num_workers 10\