From 3642928abeccc818af33492d3df295f1ea7d5159 Mon Sep 17 00:00:00 2001 From: Will Lin Date: Tue, 27 May 2025 11:38:47 -0700 Subject: [PATCH] mp training --- .../v1/models/loader/component_loader.py | 18 ++++++++-- fastvideo/v1/models/loader/fsdp_load.py | 33 ++++++++++++++++--- .../v1/pipelines/composed_pipeline_base.py | 16 +++++---- fastvideo/v1/pipelines/stages/denoising.py | 5 +-- fastvideo/v1/pipelines/training_pipeline.py | 18 ++++++---- fastvideo/v1/pipelines/wan/wan_pipeline.py | 1 - train_v1.sh | 4 +-- 7 files changed, 71 insertions(+), 24 deletions(-) diff --git a/fastvideo/v1/models/loader/component_loader.py b/fastvideo/v1/models/loader/component_loader.py index ef8626c59..95b565960 100644 --- a/fastvideo/v1/models/loader/component_loader.py +++ b/fastvideo/v1/models/loader/component_loader.py @@ -394,7 +394,16 @@ class TransformerLoader(ComponentLoader): default_dtype = PRECISION_TO_TYPE[fastvideo_args.precision] # Load the model using FSDP loader - logger.info("Loading model from %s", cls_name) + logger.info("Loading model from %s, default_dtype: %s", cls_name, default_dtype) + # model = load_fsdp_model(model_cls=model_cls, + # init_params={ + # "config": dit_config, + # "hf_config": hf_config + # }, + # weight_dir_list=safetensors_list, + # device=fastvideo_args.device, + # cpu_offload=fastvideo_args.use_cpu_offload, + # default_dtype=default_dtype) model = load_fsdp_model(model_cls=model_cls, init_params={ "config": dit_config, @@ -403,7 +412,12 @@ class TransformerLoader(ComponentLoader): weight_dir_list=safetensors_list, device=fastvideo_args.device, cpu_offload=fastvideo_args.use_cpu_offload, - default_dtype=default_dtype) + default_dtype=default_dtype, + # TODO(will): make these configurable + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + output_dtype=None, + ) if fastvideo_args.enable_torch_compile: logger.info("Torch Compile enabled for DiT") for n, m in reversed(list(model.named_modules())): diff --git a/fastvideo/v1/models/loader/fsdp_load.py b/fastvideo/v1/models/loader/fsdp_load.py index 5d10a8738..d2707f163 100644 --- a/fastvideo/v1/models/loader/fsdp_load.py +++ b/fastvideo/v1/models/loader/fsdp_load.py @@ -14,13 +14,16 @@ from typing import (Any, Callable, DefaultDict, Dict, Generator, Hashable, List, import torch from torch import nn from torch.distributed import DeviceMesh, init_device_mesh -from torch.distributed._composable.fsdp import CPUOffloadPolicy, fully_shard +from torch.distributed.fsdp import CPUOffloadPolicy, fully_shard, MixedPrecisionPolicy from torch.distributed._tensor import distribute_tensor from torch.nn.modules.module import _IncompatibleKeys from fastvideo.v1.distributed.parallel_state import ( get_sequence_model_parallel_world_size) from fastvideo.v1.models.loader.weight_utils import safetensors_weights_iterator +from fastvideo.v1.logger import init_logger + +logger = init_logger(__name__) # TODO(PY): move this to utils elsewhere @@ -86,16 +89,29 @@ def get_param_names_mapping( # TODO(PY): add compile option + # param_dtype: torch.dtype, + # reduce_dtype: torch.dtype, + # output_dtype: torch.dtype, + # pp_enabled: bool = False, + # cpu_offload: bool = False, def load_fsdp_model( model_cls: Type[nn.Module], init_params: Dict[str, Any], weight_dir_list: List[str], device: torch.device, + default_dtype: torch.dtype, + param_dtype: torch.dtype, + reduce_dtype: torch.dtype, cpu_offload: bool = False, - default_dtype: Optional[torch.dtype] = torch.bfloat16, + output_dtype: Optional[torch.dtype] = None, ) -> torch.nn.Module: + + mp_policy = MixedPrecisionPolicy(param_dtype, reduce_dtype, output_dtype, cast_forward_inputs=True) + + # with set_default_dtype(default_dtype), torch.device("meta"): with set_default_dtype(default_dtype), torch.device("meta"): model = model_cls(**init_params) + device_mesh = init_device_mesh( "cuda", mesh_shape=(get_sequence_model_parallel_world_size(), ), @@ -104,6 +120,7 @@ def load_fsdp_model( shard_model(model, cpu_offload=cpu_offload, reshard_after_forward=True, + mp_policy=mp_policy, dp_mesh=device_mesh["dp"]) weight_iterator = safetensors_weights_iterator(weight_dir_list) param_names_mapping_fn = get_param_names_mapping(model._param_names_mapping) @@ -129,6 +146,7 @@ def shard_model( *, cpu_offload: bool, reshard_after_forward: bool = True, + mp_policy: Optional[MixedPrecisionPolicy] = None, dp_mesh: Optional[DeviceMesh] = None, ) -> None: """ @@ -156,14 +174,17 @@ def shard_model( """ fsdp_kwargs = { "reshard_after_forward": reshard_after_forward, - "mesh": dp_mesh + "mesh": dp_mesh, + "mp_policy": mp_policy, } if cpu_offload: fsdp_kwargs["offload_policy"] = CPUOffloadPolicy() - # Shard the model with FSDP, iterating in reverse to start with + # iterating in reverse to start with # lowest-level modules first num_layers_sharded = 0 + # TODO(will): don't reshard after forward for the last layer to save on the + # all-gather that will immediately happen Shard the model with FSDP, for n, m in reversed(list(model.named_modules())): if any([ shard_condition(n, m) @@ -210,6 +231,10 @@ def load_fsdp_model_from_full_model_state_dict( NotImplementedError: If got FSDP with more than 1D. """ meta_sharded_sd = model.state_dict() + # s = fully_shard.state(model) + # logger.info(f"type(s): {type(s)}") + # logger.info(f"s: {s}") + # import pdb; pdb.set_trace() sharded_sd = {} to_merge_params: DefaultDict[Hashable, Dict[Any, Any]] = defaultdict(dict) diff --git a/fastvideo/v1/pipelines/composed_pipeline_base.py b/fastvideo/v1/pipelines/composed_pipeline_base.py index d12c58a35..71ffb03bb 100644 --- a/fastvideo/v1/pipelines/composed_pipeline_base.py +++ b/fastvideo/v1/pipelines/composed_pipeline_base.py @@ -155,17 +155,21 @@ class ComposedPipelineBase(ABC): for key, value in config_args.items(): setattr(fastvideo_args, key, value) + # we use cpu offload for training fastvideo_args.use_cpu_offload = False + # make sure we are in training mode fastvideo_args.inference_mode = False + # we hijack the precision to be the master weight type so that the + # model is loaded with the correct precision. Subsequently we will + # use FSDP2's MixedPrecisionPolicy to set the precision for the + # fwd, bwd, and other operations' precision. + fastvideo_args.precision = fastvideo_args.master_weight_type + assert fastvideo_args.precision == 'fp32', 'only fp32 is supported for training' + + fastvideo_args.check_fastvideo_args() logger.info(f"fastvideo_args in from_pretrained: {fastvideo_args}") - # fastvideo_args = FastVideoArgs( - # model_path=model_path, - # device_str=device or "cuda" if torch.cuda.is_available() else "cpu", - # **config_args) - fastvideo_args.check_fastvideo_args() - return cls(model_path, fastvideo_args, required_config_modules=required_config_modules) diff --git a/fastvideo/v1/pipelines/stages/denoising.py b/fastvideo/v1/pipelines/stages/denoising.py index 51c9c7716..f282356f1 100644 --- a/fastvideo/v1/pipelines/stages/denoising.py +++ b/fastvideo/v1/pipelines/stages/denoising.py @@ -74,7 +74,8 @@ class DenoisingStage(PipelineStage): ) # Setup precision and autocast settings - target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision] + # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision] + target_dtype = torch.bfloat16 autocast_enabled = (target_dtype != torch.float32 ) and not fastvideo_args.disable_autocast @@ -189,7 +190,7 @@ class DenoisingStage(PipelineStage): # Predict noise residual with torch.autocast(device_type="cuda", - dtype=target_dtype, + dtype=torch.bfloat16, enabled=autocast_enabled): # TODO(will-refactor): all of this should be in the stage's init diff --git a/fastvideo/v1/pipelines/training_pipeline.py b/fastvideo/v1/pipelines/training_pipeline.py index d6e312887..12f3f9161 100644 --- a/fastvideo/v1/pipelines/training_pipeline.py +++ b/fastvideo/v1/pipelines/training_pipeline.py @@ -190,14 +190,15 @@ class TrainingPipeline(ComposedPipelineBase, ABC): logger.info(f"infos: {infos}") caption = infos['caption'] captions.append(caption) - prompt_embeds = embeddings.to(fastvideo_args.device) - prompt_attention_mask = masks.to(fastvideo_args.device) + prompt_embeds = embeddings.to(fastvideo_args.device).to(torch.bfloat16) + prompt_attention_mask = masks.to(fastvideo_args.device).to(torch.bfloat16) # Calculate sizes latents_size = [(sampling_param.num_frames - 1) // 4 + 1, sampling_param.height // 8, sampling_param.width // 8] n_tokens = latents_size[0] * latents_size[1] * latents_size[2] + logger.info('embed dtype', prompt_embeds.dtype) num_frames = (fastvideo_args.num_latent_t - 1) * 4 + 1 logger.info(f"validation num_frames: {num_frames}") @@ -216,7 +217,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC): width=args.num_width, num_frames=num_frames, # num_inference_steps=fastvideo_args.validation_sampling_steps, - num_inference_steps=10, + num_inference_steps=50, # guidance_scale=fastvideo_args.validation_guidance_scale, guidance_scale=1, n_tokens=n_tokens, @@ -226,10 +227,11 @@ class TrainingPipeline(ComposedPipelineBase, ABC): ) # Run validation inference - with torch.inference_mode(): - output_batch = self.validation_pipeline.forward( - batch, fastvideo_args) - samples = output_batch.output + with torch.autocast("cuda", dtype=torch.bfloat16): + with torch.inference_mode(): + output_batch = self.validation_pipeline.forward( + batch, fastvideo_args) + samples = output_batch.output # Process outputs video = rearrange(samples, "b c t h w -> t b c h w") @@ -555,6 +557,8 @@ class WanTrainingPipeline(TrainingPipeline): args_copy.inference_mode = True args_copy.vae_config.load_encoder = False + # TODO(will): clean this up + args_copy.precision = "bf16" validation_pipeline = WanValidationPipeline.from_pretrained( args.model_path, args=args_copy) diff --git a/fastvideo/v1/pipelines/wan/wan_pipeline.py b/fastvideo/v1/pipelines/wan/wan_pipeline.py index f07af6329..35be7e970 100644 --- a/fastvideo/v1/pipelines/wan/wan_pipeline.py +++ b/fastvideo/v1/pipelines/wan/wan_pipeline.py @@ -15,7 +15,6 @@ from fastvideo.v1.pipelines.stages import (ConditioningStage, DecodingStage, TextEncodingStage, TimestepPreparationStage) -# TODO(will): move PRECISION_TO_TYPE to better place logger = init_logger(__name__) diff --git a/train_v1.sh b/train_v1.sh index 35fc137a5..970d8508a 100644 --- a/train_v1.sh +++ b/train_v1.sh @@ -49,5 +49,5 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\ --multi_phased_distill_schedule "4000-1" \ --weight_decay 0.01 \ --not_apply_cfg_solver \ - --master_weight_type "bf16" \ - --max_grad_norm 1.0 \ No newline at end of file + --master_weight_type "fp32" \ + --max_grad_norm 1.0