diff --git a/fastvideo/data_preprocess/preprocess_vae_latents.py b/fastvideo/data_preprocess/preprocess_vae_latents.py index 23072610e..9b2f990a2 100644 --- a/fastvideo/data_preprocess/preprocess_vae_latents.py +++ b/fastvideo/data_preprocess/preprocess_vae_latents.py @@ -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) diff --git a/fastvideo/data_preprocess/v1_preprocess.py b/fastvideo/data_preprocess/v1_preprocess.py index bf372b34f..d7da14865 100644 --- a/fastvideo/data_preprocess/v1_preprocess.py +++ b/fastvideo/data_preprocess/v1_preprocess.py @@ -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) \ No newline at end of file diff --git a/fastvideo/distill.py b/fastvideo/distill.py index bbae02ded..ac776dfc2 100644 --- a/fastvideo/distill.py +++ b/fastvideo/distill.py @@ -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) diff --git a/fastvideo/distill_adv.py b/fastvideo/distill_adv.py index f506fb623..244907e3c 100644 --- a/fastvideo/distill_adv.py +++ b/fastvideo/distill_adv.py @@ -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) diff --git a/fastvideo/train.py b/fastvideo/train.py index 721b603e9..ac0b93a33 100644 --- a/fastvideo/train.py +++ b/fastvideo/train.py @@ -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) diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py index d9e0f9e04..c763d0ddb 100644 --- a/fastvideo/v1/fastvideo_args.py +++ b/fastvideo/v1/fastvideo_args.py @@ -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: diff --git a/fastvideo/v1/models/loader/component_loader.py b/fastvideo/v1/models/loader/component_loader.py index dc4cb9702..110b55875 100644 --- a/fastvideo/v1/models/loader/component_loader.py +++ b/fastvideo/v1/models/loader/component_loader.py @@ -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, diff --git a/fastvideo/v1/models/loader/fsdp_load.py b/fastvideo/v1/models/loader/fsdp_load.py index 982141b37..7189e4483 100644 --- a/fastvideo/v1/models/loader/fsdp_load.py +++ b/fastvideo/v1/models/loader/fsdp_load.py @@ -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, diff --git a/scripts/finetune/finetune_v1.sh b/scripts/finetune/finetune_v1.sh index fa0b1228a..b5fb9ec58 100644 --- a/scripts/finetune/finetune_v1.sh +++ b/scripts/finetune/finetune_v1.sh @@ -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\