rebase fixes

fastvideo/utils/checkpoint.py

@@ -11,7 +11,6 @@ from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_
 from torch.distributed.fsdp import FullOptimStateDictConfig, FullStateDictConfig
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp import StateDictType
-import dataclasses
 
 from fastvideo.utils.logging_ import main_print
 
@@ -45,50 +44,13 @@ def save_checkpoint_optimizer(model, optimizer, rank, output_dir, step, discrimi
         optimizer_path = os.path.join(save_dir, "optimizer.pt")
         torch.save(optim_state, optimizer_path)
     else:
-        weight_path = os.path.join(save_dstate_dictir, "discriminator_pytorch_model.safetensors")
+        weight_path = os.path.join(save_dir, "discriminator_pytorch_model.safetensors")
         save_file(cpu_state, weight_path)
         optimizer_path = os.path.join(save_dir, "discriminator_optimizer.pt")
         torch.save(optim_state, optimizer_path)
     main_print(f"--> checkpoint saved at step {step}")
 
 
-def save_checkpoint_v1(transformer, rank, output_dir, step):
-    # from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
-# from torch.distributed.fsdp import StateDictType, FullStateDictConfig
-
-    # Configure FSDP to save full state dict
-    FSDP.set_state_dict_type(
-        transformer,
-        state_dict_type=StateDictType.FULL_STATE_DICT,
-        state_dict_config=FullStateDictConfig(offload_to_cpu=True, rank0_only=True),
-    )
-
-    # Now get the state dict
-    cpu_state = transformer.state_dict()
-
-    # Save it (only on rank 0 since we used rank0_only=True)
-    # if torch.distributed.get_rank() == 0:
-    #     torch.save(state_dict, "model_checkpoint.pt")
-    if rank <= 0:
-        save_dir = os.path.join(output_dir, f"checkpoint-{step}")
-        os.makedirs(save_dir, exist_ok=True)
-        # save using safetensors
-        # weight_path = os.path.join(save_dir, "diffusion_pytorch_model.safetensors")
-        weight_path = os.path.join(save_dir, "diffusion_pytorch_model.pt")
-        print(weight_path)
-        # save_file(cpu_state, weight_path)
-        torch.save(cpu_state, weight_path)
-        config_dict = transformer.hf_config
-        if "dtype" in config_dict:
-            del config_dict["dtype"]  # TODO
-        config_path = os.path.join(save_dir, "config.json")
-        # save dict as json
-        with open(config_path, "w") as f:
-            json.dump(config_dict, f, indent=4)
-    main_print(f"--> checkpoint saved at step {step}")
-
-
-
 def save_checkpoint(transformer, rank, output_dir, step):
     main_print(f"--> saving checkpoint at step {step}")
     with FSDP.state_dict_type(

fastvideo/v1/pipelines/training_pipeline.py

@@ -19,7 +19,6 @@ from tqdm.auto import tqdm
 # import torch.distributed as dist
 import wandb
 from fastvideo.models.mochi_hf.mochi_latents_utils import normalize_dit_input
-from fastvideo.utils.checkpoint import save_checkpoint_v1
 from fastvideo.v1.configs.sample import SamplingParam
 from fastvideo.v1.dataset.parquet_datasets import ParquetVideoTextDataset
 from fastvideo.v1.distributed import cleanup_dist_env_and_memory, get_sp_group
@@ -30,7 +29,7 @@ from fastvideo.v1.pipelines import ComposedPipelineBase
 from fastvideo.v1.pipelines.pipeline_batch_info import ForwardBatch
 from fastvideo.v1.pipelines.training_utils import (
     _clip_grad_norm_while_handling_failing_dtensor_cases,
-    compute_density_for_timestep_sampling, get_sigmas)
+    compute_density_for_timestep_sampling, get_sigmas, save_checkpoint)
 from fastvideo.v1.pipelines.wan.wan_pipeline import WanValidationPipeline
 
 logger = init_logger(__name__)
@@ -558,12 +557,14 @@ class WanTrainingPipeline(TrainingPipeline):
         total_loss = 0.0
         optimizer.zero_grad()
         for _ in range(gradient_accumulation_steps):
+            logger.info(f"Rank {self.rank}: Training step {_}")
             (
                 latents,
                 encoder_hidden_states,
                 encoder_attention_mask,
                 infos,
             ) = next(loader_iter)
+            logger.info(f"Rank {self.rank}: Training step {_} loaded data")
             latents = latents.to(self.fastvideo_args.device,
                                  dtype=torch.bfloat16)
             encoder_hidden_states = encoder_hidden_states.to(
@@ -772,26 +773,16 @@ class WanTrainingPipeline(TrainingPipeline):
                     step=step,
                 )
             if step % args.checkpointing_steps == 0:
-                if args.use_lora:
-                    raise NotImplementedError("LoRA is not supported now")
-                    # Save LoRA weights
-                    save_lora_checkpoint(transformer, optimizer, rank,
-                                         args.output_dir, step, pipe)
-                else:
-                    # Your existing checkpoint saving code
-                    save_checkpoint_v1(self.transformer, self.rank,
-                                       args.output_dir, step)
-                    self.transformer.train()
+                # Your existing checkpoint saving code
+                save_checkpoint(self.transformer, self.rank, args.output_dir,
+                                step)
+                self.transformer.train()
                 self.sp_group.barrier()
             if args.log_validation and step % args.validation_steps == 0:
                 self.log_validation(self.transformer, args, step)
 
-        if args.use_lora:
-            raise NotImplementedError("LoRA is not supported now")
-            # save_lora_checkpoint(transformer, optimizer, rank, args.output_dir, args.max_train_steps, pipe)
-        else:
-            save_checkpoint_v1(self.transformer, self.rank, args.output_dir,
-                               args.max_train_steps)
+        save_checkpoint(self.transformer, self.rank, args.output_dir,
+                        args.max_train_steps)
 
         if get_sp_group():
             cleanup_dist_env_and_memory()

fastvideo/v1/pipelines/training_utils.py

@@ -1,16 +1,14 @@
 import json
-import os
 import math
+import os
 from typing import List, Optional, Tuple, Union
 
 import torch
 import torch.distributed as dist
-
-import torch
+import torch.distributed.tensor
 from torch.distributed.fsdp import FullStateDictConfig
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.distributed.fsdp import StateDictType
-import torch.distributed.tensor
 
 from fastvideo.v1.logger import init_logger
 
@@ -94,9 +92,7 @@ def save_checkpoint(transformer, rank, output_dir, step):
         # save dict as json
         with open(config_path, "w") as f:
             json.dump(config_dict, f, indent=4)
-    logger.info("--> checkpoint saved at step {step} to {weight_path}",
-                step=step,
-                weight_path=weight_path)
+    logger.info("--> checkpoint saved at step %s to %s", step, weight_path)
 
 
 def _clip_grad_norm_while_handling_failing_dtensor_cases(

train_v1.sh

@@ -25,13 +25,13 @@ torchrun --nnodes 1 --nproc_per_node $NUM_GPUS\
     --sp_size $NUM_GPUS \
     --tp_size $NUM_GPUS \
     --train_sp_batch_size 1\
-    --dataloader_num_workers 1\
+    --dataloader_num_workers 5\
     --gradient_accumulation_steps=1\
     --max_train_steps=120 \
     --learning_rate=1e-6\
     --mixed_precision="bf16"\
-    --checkpointing_steps=60 \
-    --validation_steps 20\
+    --checkpointing_steps=5 \
+    --validation_steps 10\
     --validation_sampling_steps "2,4,8" \
     --log_validation \
     --checkpoints_total_limit 3\