mp training

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())):

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)
@@ -130,6 +147,7 @@ def shard_model(
     *,
     cpu_offload: bool,
     reshard_after_forward: bool = True,
+    mp_policy: Optional[MixedPrecisionPolicy] = None,
     dp_mesh: Optional[DeviceMesh] = None,
 ) -> None:
     """
@@ -157,14 +175,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)
@@ -211,6 +232,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)

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)

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

fastvideo/v1/pipelines/training_pipeline.py

@@ -191,14 +191,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)
 
             # Prepare batch for validation
             # print('shape of embeddings', prompt_embeds.shape)
@@ -215,7 +216,7 @@ class TrainingPipeline(ComposedPipelineBase, ABC):
                 width=args.num_width,
                 num_frames=args.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,
@@ -225,10 +226,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")
@@ -529,6 +531,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)
 

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__)
 

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 
+    --master_weight_type "fp32" \
+    --max_grad_norm 1.0