diff --git a/fastvideo/v1/dataset/parquet_dataset_iterable_style.py b/fastvideo/v1/dataset/parquet_dataset_iterable_style.py
index 20082fc22..4469d9cd1 100644
--- a/fastvideo/v1/dataset/parquet_dataset_iterable_style.py
+++ b/fastvideo/v1/dataset/parquet_dataset_iterable_style.py
@@ -239,7 +239,7 @@ def shard_parquet_files_across_sp_groups_and_workers(
         logger.info("Saved sharding info to %s", save_dir)
 
     # wait for all ranks to finish
-    torch.distributed.barrier()
+    torch.distributed.barrier(device_ids=[get_world_rank()])
     # recursive call
     return shard_parquet_files_across_sp_groups_and_workers(
         path, num_sp_groups, num_workers, seed)
diff --git a/fastvideo/v1/dataset/parquet_dataset_map_style.py b/fastvideo/v1/dataset/parquet_dataset_map_style.py
index 067f04340..8458049c0 100644
--- a/fastvideo/v1/dataset/parquet_dataset_map_style.py
+++ b/fastvideo/v1/dataset/parquet_dataset_map_style.py
@@ -132,7 +132,7 @@ def get_parquet_files_and_length(path: str):
 
     # Wait for rank 0 to finish saving
     if get_world_size() > 1:
-        torch.distributed.barrier()
+        torch.distributed.barrier(device_ids=[get_world_rank()])
 
     return get_parquet_files_and_length(path)
 
diff --git a/fastvideo/v1/fastvideo_args.py b/fastvideo/v1/fastvideo_args.py
index 40bf88815..d789eb8a3 100644
--- a/fastvideo/v1/fastvideo_args.py
+++ b/fastvideo/v1/fastvideo_args.py
@@ -446,6 +446,13 @@ class TrainingArgs(FastVideoArgs):
     # VSA training decay parameters
     VSA_decay_rate: float = 0.01  # decay rate -> 0.02
     VSA_decay_interval_steps: int = 1  # decay interval steps -> 50
+    
+    # distillation args
+    student_critic_update_ratio: int = 5
+    denoising_step_list: List[int] = field(default_factory=lambda: [1000, 757, 522])
+    min_step_ratio: float = 0.2
+    max_step_ratio: float = 0.98
+    teacher_guidance_scale: float = 3.5
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs":
@@ -710,69 +717,33 @@ class TrainingArgs(FastVideoArgs):
             type=int,
             default=TrainingArgs.VSA_decay_interval_steps,
             help="VSA decay interval steps")
+        
+        # Distillation arguments
+        parser.add_argument("--student-critic-update-ratio",
+            type=int,
+            default=TrainingArgs.student_critic_update_ratio,
+            help="Ratio of student updates to critic updates.")
+        parser.add_argument("--denoising-step-list",
+            type=parse_int_list,
+            default=[1000, 757, 522],
+            help="Comma-separated list of denoising steps (e.g., '1000,757,522')")
+        parser.add_argument("--min-step-ratio",
+            type=float,
+            default=TrainingArgs.min_step_ratio,
+            help="Minimum step ratio")
+        parser.add_argument("--max-step-ratio",
+            type=float,
+            default=TrainingArgs.max_step_ratio,
+            help="Maximum step ratio")
+        parser.add_argument("--teacher-guidance-scale",
+            type=float,
+            default=TrainingArgs.teacher_guidance_scale,
+            help="Teacher guidance scale")
 
         return parser
 
-
 def parse_int_list(value: str) -> List[int]:
     """Parse a comma-separated string of integers into a list."""
     if not value:
         return []
-    return [int(x.strip()) for x in value.split(",")]
-
-@dataclasses.dataclass
-class DistillationArgs(TrainingArgs):
-    """
-    Distillation arguments. Inherits from TrainingArgs and adds distillation-specific
-    arguments.
-    """
-    student_critic_update_ratio: int = 5
-    denoising_step_list: List[int] = [1000, 757, 522]
-    min_step_ratio: float = 0.2
-    max_step_ratio: float = 0.98
-    teacher_guidance_scale: float = 3.5
-
-    @staticmethod
-    def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:        
-        # Then add distillation-specific arguments
-        parser.add_argument("--student-critic-update-ratio",
-                            type=int,
-                            default=DistillationArgs.student_critic_update_ratio,
-                            help="Ratio of student updates to critic updates.")
-        parser.add_argument("--denoising-step-list",
-                            type=parse_int_list,
-                            default=DistillationArgs.denoising_step_list,
-                            help="Comma-separated list of denoising steps (e.g., '1000,757,522')")
-        parser.add_argument("--min-step-ratio",
-                            type=float,
-                            default=DistillationArgs.min_step_ratio,
-                            help="Minimum step ratio")
-        parser.add_argument("--max-step-ratio",
-                            type=float,
-                            default=DistillationArgs.max_step_ratio,
-                            help="Maximum step ratio")
-        parser.add_argument("--teacher-guidance-scale",
-                            type=float,
-                            default=DistillationArgs.teacher_guidance_scale,
-                            help="Teacher guidance scale")
-        return parser
-
-    @classmethod
-    def from_cli_args(cls, args: argparse.Namespace) -> "DistillationArgs":
-        provided_args = clean_cli_args(args)
-        # Get all fields from the dataclass
-        attrs = [attr.name for attr in dataclasses.fields(cls)]
-        logger.info(provided_args)
-        # Create a dictionary of attribute values, with defaults for missing attributes
-        kwargs = {}
-        for attr in attrs:
-            if attr == 'pipeline_config':
-                pipeline_config = PipelineConfig.from_kwargs(provided_args)
-                kwargs[attr] = pipeline_config
-            # Use getattr with default value from the dataclass for potentially missing attributes
-            else:
-                default_value = getattr(cls, attr, None)
-                value = getattr(args, attr, default_value)
-                kwargs[attr] = value  # type: ignore
-
-        return cls(**kwargs)  # type: ignore
+    return [int(x.strip()) for x in value.split(",")]
\ No newline at end of file
diff --git a/fastvideo/v1/models/loader/component_loader.py b/fastvideo/v1/models/loader/component_loader.py
index 270bc8387..0e47b6b9c 100644
--- a/fastvideo/v1/models/loader/component_loader.py
+++ b/fastvideo/v1/models/loader/component_loader.py
@@ -70,6 +70,8 @@ class ComponentLoader(ABC):
         module_loaders = {
             "scheduler": (SchedulerLoader, "diffusers"),
             "transformer": (TransformerLoader, "diffusers"),
+            "teacher_transformer": (TransformerLoader, "diffusers"),
+            "critic_transformer": (TransformerLoader, "diffusers"),
             "vae": (VAELoader, "diffusers"),
             "text_encoder": (TextEncoderLoader, "transformers"),
             "text_encoder_2": (TextEncoderLoader, "transformers"),
diff --git a/fastvideo/v1/models/schedulers/scheduling_flow_unipc_multistep.py b/fastvideo/v1/models/schedulers/scheduling_flow_unipc_multistep.py
index 42df10b85..85391a6cd 100644
--- a/fastvideo/v1/models/schedulers/scheduling_flow_unipc_multistep.py
+++ b/fastvideo/v1/models/schedulers/scheduling_flow_unipc_multistep.py
@@ -779,40 +779,21 @@ class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin, BaseScheduler):
         noise: torch.Tensor,
         timesteps: torch.IntTensor,
     ) -> torch.Tensor:
-        # Make sure sigmas and timesteps have the same device and dtype as original_samples
-        sigmas = self.sigmas.to(device=original_samples.device,
-                                dtype=original_samples.dtype)
-        if original_samples.device.type == "mps" and torch.is_floating_point(
-                timesteps):
-            # mps does not support float64
-            schedule_timesteps = self.timesteps.to(original_samples.device,
-                                                   dtype=torch.float32)
-            timesteps = timesteps.to(original_samples.device,
-                                     dtype=torch.float32)
-        else:
-            schedule_timesteps = self.timesteps.to(original_samples.device)
-            timesteps = timesteps.to(original_samples.device)
-
-        # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index
-        if self.begin_index is None:
-            step_indices = [
-                self.index_for_timestep(t, schedule_timesteps)
-                for t in timesteps
-            ]
-        elif self.step_index is not None:
-            # add_noise is called after first denoising step (for inpainting)
-            step_indices = [self.step_index] * timesteps.shape[0]
-        else:
-            # add noise is called before first denoising step to create initial latent(img2img)
-            step_indices = [self.begin_index] * timesteps.shape[0]
-
-        sigma = sigmas[step_indices].flatten()
-        while len(sigma.shape) < len(original_samples.shape):
-            sigma = sigma.unsqueeze(-1)
-
-        alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
-        noisy_samples = alpha_t * original_samples + sigma_t * noise
-        return noisy_samples
+        """
+        Diffusion forward corruption process.
+        Input:
+            - clean_latent: the clean latent with shape [B, C, H, W]
+            - noise: the noise with shape [B, C, H, W]
+            - timestep: the timestep with shape [B]
+        Output: the corrupted latent with shape [B, C, H, W]
+        """
+        self.sigmas = self.sigmas.to(noise.device)
+        self.timesteps = self.timesteps.to(noise.device)
+        timestep_id = torch.argmin(
+            (self.timesteps.unsqueeze(0) - timesteps.unsqueeze(1)).abs(), dim=1)
+        sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) # [21, 1, 1, 1]
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample.type_as(noise)
 
     def __len__(self):
         return self.config.num_train_timesteps
diff --git a/fastvideo/v1/pipelines/__init__.py b/fastvideo/v1/pipelines/__init__.py
index 0840fabea..4a78dec0f 100644
--- a/fastvideo/v1/pipelines/__init__.py
+++ b/fastvideo/v1/pipelines/__init__.py
@@ -12,8 +12,7 @@ from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines.composed_pipeline_base import ComposedPipelineBase
 from fastvideo.v1.pipelines.lora_pipeline import LoRAPipeline
 from fastvideo.v1.pipelines.pipeline_batch_info import (ForwardBatch,
-                                                        TrainingBatch,
-                                                        DistillBatch)
+                                                        TrainingBatch)
 from fastvideo.v1.pipelines.pipeline_registry import PipelineRegistry
 from fastvideo.v1.utils import (maybe_download_model,
                                 verify_model_config_and_directory)
@@ -66,5 +65,4 @@ __all__ = [
     "ForwardBatch",
     "LoRAPipeline",
     "TrainingBatch",
-    "DistillBatch",
 ]
diff --git a/fastvideo/v1/pipelines/composed_pipeline_base.py b/fastvideo/v1/pipelines/composed_pipeline_base.py
index 55dae4ccf..988cd683f 100644
--- a/fastvideo/v1/pipelines/composed_pipeline_base.py
+++ b/fastvideo/v1/pipelines/composed_pipeline_base.py
@@ -224,12 +224,13 @@ class ComposedPipelineBase(ABC):
         assert len(
             model_index
         ) > 1, "model_index.json must contain at least one pipeline module"
-
+        
         for module_name in self.required_config_modules:
             if module_name not in model_index:
-                raise ValueError(
-                    f"model_index.json must contain a {module_name} module")
-
+                logger.warning(
+                    f"model_index.json does not contain a {module_name} module, adding {module_name} to model_index")
+                if 'transformer' in module_name:
+                    model_index[module_name] = model_index['transformer']
         # all the component models used by the pipeline
         required_modules = self.required_config_modules
         logger.info("Loading required modules: %s", required_modules)
@@ -244,7 +245,9 @@ class ComposedPipelineBase(ABC):
                 logger.info("Using module %s already provided", module_name)
                 modules[module_name] = loaded_modules[module_name]
                 continue
-            component_model_path = os.path.join(self.model_path, module_name)
+            loading_module_name = module_name.split("_")[-1]
+            component_model_path = os.path.join(self.model_path, loading_module_name)
+            # torch.distributed.breakpoint()  
             module = PipelineComponentLoader.load_module(
                 module_name=module_name,
                 component_model_path=component_model_path,
diff --git a/fastvideo/v1/pipelines/pipeline_batch_info.py b/fastvideo/v1/pipelines/pipeline_batch_info.py
index 9d4d1df12..b7f7f5a4d 100644
--- a/fastvideo/v1/pipelines/pipeline_batch_info.py
+++ b/fastvideo/v1/pipelines/pipeline_batch_info.py
@@ -21,7 +21,6 @@ from fastvideo.v1.configs.sample.teacache import (TeaCacheParams,
 __all__ = [
     "ForwardBatch",
     "TrainingBatch", 
-    "DistillBatch",
     "AttentionMetadata",
     "VideoSparseAttentionMetadata",
 ]
@@ -174,26 +173,18 @@ class TrainingBatch:
     total_loss: float | None = None
     grad_norm: float | None = None
 
-
-@dataclass
-class DistillBatch(TrainingBatch):
-    """Batch class for distillation training with additional distillation-specific attributes."""
-    
-    # Distillation-specific data
+    # Distillation-specific attributes
     encoder_hidden_states_neg: Optional[torch.Tensor] = None
     encoder_attention_mask_neg: Optional[torch.Tensor] = None
     conditional_dict: Optional[Dict[str, Any]] = None
     unconditional_dict: Optional[Dict[str, Any]] = None
-    clean_latent: Optional[torch.Tensor] = None
     
     # Distillation losses
     student_loss: float = 0.0
     critic_loss: float = 0.0
     
-    # Model predictions
-    model_pred: Optional[torch.Tensor] = None
-    teacher_output: Optional[torch.Tensor] = None
-    
     # Training control
-    dmd_log_dict: Dict[str, Any] = {}
-    critic_log_dict: Dict[str, Any] = {}
\ No newline at end of file
+    dmd_log_dict: Dict[str, Any] = field(default_factory=dict)
+    critic_log_dict: Dict[str, Any] = field(default_factory=dict)
+
+    
diff --git a/fastvideo/v1/training/distillation_pipeline.py b/fastvideo/v1/training/distillation_pipeline.py
index 93ca7b132..2c1f4385b 100644
--- a/fastvideo/v1/training/distillation_pipeline.py
+++ b/fastvideo/v1/training/distillation_pipeline.py
@@ -23,11 +23,11 @@ from fastvideo.v1.configs.sample import SamplingParam
 from fastvideo.v1.dataset import build_parquet_map_style_dataloader
 from fastvideo.v1.distributed import (cleanup_dist_env_and_memory, get_sp_group,
                                       get_torch_device, get_world_group)
-from fastvideo.v1.fastvideo_args import FastVideoArgs, DistillationArgs
+from fastvideo.v1.fastvideo_args import FastVideoArgs,TrainingArgs
 from fastvideo.v1.forward_context import set_forward_context
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.pipelines import (ComposedPipelineBase, ForwardBatch,
-                                    TrainingBatch, DistillBatch)
+                                    TrainingBatch)
 from fastvideo.v1.training.training_pipeline import TrainingPipeline
 from fastvideo.v1.training.training_utils import (
     clip_grad_norm_while_handling_failing_dtensor_cases,
@@ -56,20 +56,18 @@ class DistillationPipeline(TrainingPipeline):
         raise RuntimeError(
             "create_pipeline_stages should not be called for training pipeline")
 
-    def initialize_training_pipeline(self, distillation_args: DistillationArgs):
+    def initialize_training_pipeline(self, training_args: TrainingArgs):
         """Initialize the distillation training pipeline with multiple models."""
         logger.info("Initializing distillation training pipeline...")
         
         # 1. Call parent initialization first
-        super().initialize_training_pipeline(distillation_args)
-        assert isinstance(self.training_args, DistillationArgs)
+        super().initialize_training_pipeline(training_args)
         
+        self.noise_scheduler = self.get_module("scheduler")
         # 2. Distillation-specific initialization
         # The parent class already sets self.transformer as the student model
-        self.teacher_transformer = self._copy_model(self.transformer)
-        assert self.teacher_transformer is not None, "Failed to create teacher model"
-        self.critic_transformer = self._copy_model(self.transformer)
-        assert self.critic_transformer is not None, "Failed to create critic model"
+        self.teacher_transformer = self.get_module("teacher_transformer")
+        self.critic_transformer = self.get_module("critic_transformer")
         self.teacher_transformer.requires_grad_(False)
         self.teacher_transformer.eval()
         self.critic_transformer.requires_grad_(True)
@@ -79,20 +77,21 @@ class DistillationPipeline(TrainingPipeline):
         student_params = list(filter(lambda p: p.requires_grad, self.transformer.parameters()))
         self.student_transformer_optimizer = torch.optim.AdamW(
             student_params,
-            lr=distillation_args.learning_rate,
+            lr=training_args.learning_rate,
             betas=(0.9, 0.999),
-            weight_decay=distillation_args.weight_decay,
+            weight_decay=training_args.weight_decay,
             eps=1e-8,
         )
         critic_params = list(filter(lambda p: p.requires_grad, self.critic_transformer.parameters()))
         self.critic_transformer_optimizer = torch.optim.AdamW(
             critic_params,
-            lr=distillation_args.learning_rate,
+            lr=training_args.learning_rate,
             betas=(0.9, 0.999),
-            weight_decay=distillation_args.weight_decay,
+            weight_decay=training_args.weight_decay,
             eps=1e-8,
         )
         logger.info("Distillation optimizers initialized: student and critic")
+
         self.student_critic_update_ratio = self.training_args.student_critic_update_ratio
         logger.info(f"Distillation pipeline initialized with student_critic_update_ratio={self.student_critic_update_ratio}")
 
@@ -116,7 +115,7 @@ class DistillationPipeline(TrainingPipeline):
 
 
     @abstractmethod
-    def initialize_validation_pipeline(self, distillation_args: DistillationArgs):
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
         """Initialize validation pipeline - must be implemented by subclasses."""
         raise NotImplementedError(
             "Distillation pipelines must implement this method") 
@@ -128,7 +127,7 @@ class DistillationPipeline(TrainingPipeline):
         copied_model.load_state_dict(state_dict)
         return copied_model
 
-    def _prepare_distillation(self, training_batch: DistillBatch) -> DistillBatch:
+    def _prepare_distillation(self, training_batch: TrainingBatch) -> TrainingBatch:
         """Prepare training environment for distillation."""
         self.transformer.requires_grad_(True)
         self.transformer.train()
@@ -171,10 +170,10 @@ class DistillationPipeline(TrainingPipeline):
         else:
             raise NotImplementedError("Unsupported model type {}".format(type))
 
-    def _student_forward(self, conditional_dict: Dict[str, Any], unconditional_dict: Dict[str, Any], clean_latent: torch.Tensor) -> torch.Tensor:
+    def _student_forward(self, conditional_dict: Dict[str, Any], unconditional_dict: Dict[str, Any], latents: torch.Tensor) -> torch.Tensor:
         """Forward pass through student transformer and compute student losses."""
         video_latent_shape = self.video_latent_shape
-        dtype = clean_latent.dtype
+        dtype = latents.dtype
         simulated_noisy_input = []
         for timestep in self.denoising_step_list:
             noise = torch.randn(
@@ -184,13 +183,14 @@ class DistillationPipeline(TrainingPipeline):
                 video_latent_shape[:2], device=self.device, dtype=torch.long)
 
             if timestep != 0:
+                
                 noisy_image = self.noise_scheduler.add_noise(
-                    clean_latent.flatten(0, 1),
+                    latents.flatten(0, 1),
                     noise.flatten(0, 1),
                     noisy_timestep.flatten(0, 1)
                 ).unflatten(0, video_latent_shape[:2])
             else:
-                noisy_image = clean_latent
+                noisy_image = latents
 
             simulated_noisy_input.append(noisy_image)
 
@@ -260,7 +260,7 @@ class DistillationPipeline(TrainingPipeline):
         grad = torch.nan_to_num(grad)
 
         return grad, {
-            "dmdtrain_clean_latent": estimated_clean_video.detach(),
+            "dmdtrain_latents": estimated_clean_video.detach(),
             "dmdtrain_noisy_latent": noisy_video.detach(),
             "dmdtrain_pred_real_video": pred_real_video.detach(),
             "dmdtrain_pred_fake_video": pred_fake_video.detach(),
@@ -313,7 +313,7 @@ class DistillationPipeline(TrainingPipeline):
         
         return dmd_loss, dmd_log_dict
 
-    def _compute_critic_loss(self, training_batch: DistillBatch) -> torch.Tensor:
+    def _compute_critic_loss(self, training_batch: TrainingBatch) -> torch.Tensor:
         """Compute critic loss for adversarial training."""
         assert self.critic_transformer is not None
         assert training_batch.latents is not None
@@ -329,16 +329,16 @@ class DistillationPipeline(TrainingPipeline):
         critic_loss = real_loss + fake_loss
         return critic_loss
 
-    def _student_forward_and_compute_dmd_loss(self, training_batch: DistillBatch) -> Tuple[DistillBatch, torch.Tensor, dict]:
+    def _student_forward_and_compute_dmd_loss(self, training_batch: TrainingBatch) -> Tuple[TrainingBatch, torch.Tensor, dict]:
         """Forward pass through student transformer and compute student losses."""
         assert training_batch.conditional_dict is not None
         assert training_batch.unconditional_dict is not None
-        assert training_batch.clean_latent is not None
+        assert training_batch.latents is not None
         
         pred_video = self._student_forward(
             conditional_dict=training_batch.conditional_dict,
             unconditional_dict=training_batch.unconditional_dict,
-            clean_latent=training_batch.clean_latent
+            latents=training_batch.latents
         )
 
         dmd_loss, dmd_log_dict = self._compute_dmd_loss(
@@ -349,16 +349,16 @@ class DistillationPipeline(TrainingPipeline):
         
         return training_batch, dmd_loss, dmd_log_dict
 
-    def _critic_forward_and_compute_loss(self, training_batch: DistillBatch) -> Tuple[DistillBatch, torch.Tensor, dict]:
+    def _critic_forward_and_compute_loss(self, training_batch: TrainingBatch) -> Tuple[TrainingBatch, torch.Tensor, dict]:
         assert training_batch.conditional_dict is not None
         assert training_batch.unconditional_dict is not None
-        assert training_batch.clean_latent is not None
+        assert training_batch.latents is not None
 
         with torch.no_grad():
             generated_video = self._student_forward(
                 conditional_dict=training_batch.conditional_dict,
                 unconditional_dict=training_batch.unconditional_dict,
-                clean_latent=training_batch.clean_latent
+                latents=training_batch.latents
             )
         video_latent_shape = self.video_latent_shape
         critic_timestep = torch.randint(
@@ -436,7 +436,7 @@ class DistillationPipeline(TrainingPipeline):
         flow_pred = (xt - x0_pred) / sigma_t
         return flow_pred.to(original_dtype)
         
-    def _clip_grad_norm(self, training_batch: DistillBatch) -> DistillBatch:
+    def _clip_grad_norm(self, training_batch: TrainingBatch) -> TrainingBatch:
         assert self.training_args is not None
         max_grad_norm = self.training_args.max_grad_norm
 
@@ -459,7 +459,25 @@ class DistillationPipeline(TrainingPipeline):
         training_batch.grad_norm = grad_norm
         return training_batch
 
-    def train_one_step(self, training_batch: DistillBatch) -> DistillBatch:
+    def _prepare_dit_inputs(self,
+                            training_batch: TrainingBatch) -> TrainingBatch:
+        super()._prepare_dit_inputs(training_batch)
+        conditional_dict = {
+            "encoder_hidden_states": training_batch.encoder_hidden_states,
+            "encoder_attention_mask": training_batch.encoder_attention_mask,
+        }
+        unconditional_dict = {
+            "encoder_hidden_states_neg": self.negative_prompt_embeds,
+            "encoder_attention_mask_neg": self.negative_prompt_attention_mask,
+        }
+
+        training_batch.conditional_dict = conditional_dict
+        training_batch.unconditional_dict = unconditional_dict
+        self.video_latent_shape = training_batch.latents.shape
+
+        return training_batch
+    
+    def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch:
         """Train one step with alternating student and critic updates."""
         assert self.training_args is not None
 
@@ -469,7 +487,6 @@ class DistillationPipeline(TrainingPipeline):
 
         for _ in range(self.training_args.gradient_accumulation_steps):
             training_batch = self._get_next_batch(training_batch)
-            training_batch = DistillBatch(**{k: v for k, v in training_batch.__dict__.items()})
 
             # assert training_batch.latents is not None
             # training_batch.latents = shard_latents_across_sp(
@@ -538,7 +555,7 @@ class DistillationPipeline(TrainingPipeline):
         # Then add distillation-specific information
         logger.info("Distillation-specific settings:")
         logger.info("  Student/Critic update ratio: %s", self.student_critic_update_ratio)
-        assert isinstance(self.training_args, DistillationArgs)
+        assert isinstance(self.training_args, TrainingArgs)
         logger.info("  Max gradient norm: %s", self.training_args.max_grad_norm)
         assert self.teacher_transformer is not None
         logger.info("  Teacher transformer parameters: %s B",
@@ -548,30 +565,30 @@ class DistillationPipeline(TrainingPipeline):
                     sum(p.numel() for p in self.critic_transformer.parameters()) / 1e9)
 
     @torch.no_grad()
-    def _log_validation(self, transformer, distillation_args, global_step) -> None:
+    def _log_validation(self, transformer, training_args, global_step) -> None:
         """Log validation results for distillation training."""
-        assert distillation_args is not None
-        distillation_args.inference_mode = True
-        distillation_args.use_cpu_offload = False
-        if not distillation_args.log_validation:
+        assert training_args is not None
+        training_args.inference_mode = True
+        training_args.use_cpu_offload = False
+        if not training_args.log_validation:
             return
         if self.validation_pipeline is None:
             raise ValueError("Validation pipeline is not set")
 
         logger.info("Starting validation")
 
-        sampling_param = SamplingParam.from_pretrained(distillation_args.model_path)
+        sampling_param = SamplingParam.from_pretrained(training_args.model_path)
 
-        validation_seed = distillation_args.seed if distillation_args.seed is not None else 42
+        validation_seed = training_args.seed if training_args.seed is not None else 42
         torch.manual_seed(validation_seed)
         torch.cuda.manual_seed_all(validation_seed)
 
         logger.info("Using validation seed: %s", validation_seed)
 
         logger.info('fastvideo_args.validation_preprocessed_path: %s',
-                    distillation_args.validation_preprocessed_path)
+                    training_args.validation_preprocessed_path)
         validation_dataset, validation_dataloader = build_parquet_map_style_dataloader(
-            distillation_args.validation_preprocessed_path,
+            training_args.validation_preprocessed_path,
             batch_size=1,
             num_data_workers=0,
             drop_last=False,
@@ -584,7 +601,7 @@ class DistillationPipeline(TrainingPipeline):
         if transformer:
             transformer.eval()
 
-        validation_steps = distillation_args.validation_sampling_steps.split(",")
+        validation_steps = training_args.validation_sampling_steps.split(",")
         validation_steps = [int(step) for step in validation_steps]
         validation_steps = [step for step in validation_steps if step > 0]
 
@@ -602,8 +619,8 @@ class DistillationPipeline(TrainingPipeline):
                                 sampling_param.width // 8]
                 n_tokens = latents_size[0] * latents_size[1] * latents_size[2]
 
-                temporal_compression_factor = distillation_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
-                num_frames = (distillation_args.num_latent_t -
+                temporal_compression_factor = training_args.pipeline_config.vae_config.arch_config.temporal_compression_ratio
+                num_frames = (training_args.num_latent_t -
                               1) * temporal_compression_factor + 1
 
                 batch = ForwardBatch(
@@ -616,8 +633,8 @@ class DistillationPipeline(TrainingPipeline):
                     prompt_attention_mask=[prompt_attention_mask],
                     negative_prompt_embeds=[self.negative_prompt_embeds],
                     negative_attention_mask=[self.negative_prompt_attention_mask],
-                    height=distillation_args.num_height,
-                    width=distillation_args.num_width,
+                    height=training_args.num_height,
+                    width=training_args.num_width,
                     num_frames=num_frames,
                     num_inference_steps=
                     num_inference_steps,
@@ -629,7 +646,7 @@ class DistillationPipeline(TrainingPipeline):
                 with torch.no_grad(), torch.autocast("cuda",
                                                      dtype=torch.bfloat16):
                     output_batch = self.validation_pipeline.forward(
-                        batch, distillation_args)
+                        batch, training_args)
                     samples = output_batch.output
 
                 if self.rank_in_sp_group != 0:
@@ -663,9 +680,9 @@ class DistillationPipeline(TrainingPipeline):
                     for i, (video_frames,
                             caption) in enumerate(zip(all_videos,
                                                       all_captions)):
-                        os.makedirs(distillation_args.output_dir, exist_ok=True)
+                        os.makedirs(training_args.output_dir, exist_ok=True)
                         filename = os.path.join(
-                            distillation_args.output_dir,
+                            training_args.output_dir,
                             f"validation_step_{global_step}_inference_steps_{num_inference_steps}_video_{i}.mp4"
                         )
                         imageio.mimsave(filename, list(video_frames), fps=sampling_param.fps)
@@ -706,15 +723,15 @@ class DistillationPipeline(TrainingPipeline):
         })
 
         # Process DMD training data if available
-        if "dmdtrain_clean_latent" in generator_log_dict:
-            (dmdtrain_clean_latent, dmdtrain_noisy_latent, dmdtrain_pred_real_video, dmdtrain_pred_fake_video) = map(
+        if "dmdtrain_latents" in generator_log_dict:
+            (dmdtrain_latents, dmdtrain_noisy_latent, dmdtrain_pred_real_video, dmdtrain_pred_fake_video) = map(
                 lambda x: self.transformer.vae.decode_to_pixel(x).squeeze(1),
-                [generator_log_dict['dmdtrain_clean_latent'], generator_log_dict['dmdtrain_noisy_latent'],
+                [generator_log_dict['dmdtrain_latents'], generator_log_dict['dmdtrain_noisy_latent'],
                  generator_log_dict['dmdtrain_pred_real_video'], generator_log_dict['dmdtrain_pred_fake_video']]
             )
 
             wandb_loss_dict.update({
-                "dmdtrain_clean_latent": prepare_for_saving(dmdtrain_clean_latent),
+                "dmdtrain_latents": prepare_for_saving(dmdtrain_latents),
                 "dmdtrain_noisy_latent": prepare_for_saving(dmdtrain_noisy_latent),
                 "dmdtrain_pred_real_video": prepare_for_saving(dmdtrain_pred_real_video),
                 "dmdtrain_pred_fake_video": prepare_for_saving(dmdtrain_pred_fake_video)
@@ -737,7 +754,7 @@ class DistillationPipeline(TrainingPipeline):
 
         logger.info("Initialized random seeds with seed: %s", seed)
 
-        self.noise_scheduler = FlowMatchEulerDiscreteScheduler()
+        # self.noise_scheduler = FlowMatchEulerDiscreteScheduler()
 
         if self.training_args.resume_from_checkpoint:
             self._resume_from_checkpoint()
@@ -760,7 +777,7 @@ class DistillationPipeline(TrainingPipeline):
                           self.training_args.max_train_steps + 1):
             start_time = time.perf_counter()
 
-            training_batch = DistillBatch()
+            training_batch = TrainingBatch()
             training_batch.current_timestep = step
             training_batch = self.train_one_step(training_batch)
 
@@ -837,55 +854,6 @@ class DistillationPipeline(TrainingPipeline):
         # if get_sp_group():
         #     cleanup_dist_env_and_memory()
 
-    def _get_next_batch(self, training_batch: DistillBatch) -> DistillBatch:
-        """Override parent method to return DistillBatch."""
-        result = super()._get_next_batch(training_batch)
-        # Convert TrainingBatch to DistillBatch
-        distill_batch = DistillBatch()
-        for key, value in result.__dict__.items():
-            setattr(distill_batch, key, value)
-        return distill_batch
-
-    def _normalize_dit_input(self, training_batch: DistillBatch) -> DistillBatch:
-        """Override parent method to return DistillBatch."""
-        result = super()._normalize_dit_input(training_batch)
-        # Convert TrainingBatch to DistillBatch
-        distill_batch = DistillBatch()
-        for key, value in result.__dict__.items():
-            setattr(distill_batch, key, value)
-        return distill_batch
-
-    def _prepare_dit_inputs(self, training_batch: DistillBatch) -> DistillBatch:
-        """Override parent method to return DistillBatch."""
-        result = super()._prepare_dit_inputs(training_batch)
-        # Convert TrainingBatch to DistillBatch
-        distill_batch = DistillBatch()
-        for key, value in result.__dict__.items():
-            setattr(distill_batch, key, value)
-        conditional_dict = {
-            "encoder_hidden_states": training_batch.encoder_hidden_states,
-            "encoder_attention_mask": training_batch.encoder_attention_mask,
-        }
-        unconditional_dict = {
-            "encoder_hidden_states_neg": self.negative_prompt_embeds,
-            "encoder_attention_mask_neg": self.negative_prompt_attention_mask,
-        }
-
-        distill_batch.conditional_dict = conditional_dict
-        distill_batch.unconditional_dict = unconditional_dict
-        self.video_latent_shape = training_batch.latents.shape
-
-        return distill_batch
-
-    def _build_attention_metadata(self, training_batch: DistillBatch) -> DistillBatch:
-        """Override parent method to return DistillBatch."""
-        result = super()._build_attention_metadata(training_batch)
-        # Convert TrainingBatch to DistillBatch
-        distill_batch = DistillBatch()
-        for key, value in result.__dict__.items():
-            setattr(distill_batch, key, value)
-        return distill_batch
-
 class FlowPredLoss():
     def __call__(
         self, x: torch.Tensor,
diff --git a/fastvideo/v1/training/wan_distillation_pipeline.py b/fastvideo/v1/training/wan_distillation_pipeline.py
index 0314234fd..01e1fd354 100644
--- a/fastvideo/v1/training/wan_distillation_pipeline.py
+++ b/fastvideo/v1/training/wan_distillation_pipeline.py
@@ -2,7 +2,7 @@
 import sys
 from copy import deepcopy
 
-from fastvideo.v1.fastvideo_args import FastVideoArgs, DistillationArgs
+from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs
 from fastvideo.v1.logger import init_logger
 from fastvideo.v1.models.schedulers.scheduling_flow_unipc_multistep import (
     FlowUniPCMultistepScheduler)
@@ -20,20 +20,20 @@ class WanDistillationPipeline(DistillationPipeline):
     A distillation pipeline for Wan that uses a single transformer model.
     The main transformer serves as the student model, and copies are made for teacher and critic.
     """
-    _required_config_modules = ["scheduler", "transformer"]  # Only need scheduler and transformer
-
+    _required_config_modules = ["scheduler", "transformer", "vae", "teacher_transformer", "critic_transformer"]
+    
     def initialize_pipeline(self, fastvideo_args: FastVideoArgs):
         """Initialize Wan-specific scheduler."""
         self.modules["scheduler"] = FlowUniPCMultistepScheduler(
             shift=fastvideo_args.pipeline_config.flow_shift)
 
-    def create_training_stages(self, training_args: DistillationArgs):
+    def create_training_stages(self, training_args: TrainingArgs):
         """
         May be used in future refactors.
         """
         pass
 
-    def initialize_validation_pipeline(self, training_args: DistillationArgs):
+    def initialize_validation_pipeline(self, training_args: TrainingArgs):
         """Initialize Wan validation pipeline."""
         logger.info("Initializing Wan validation pipeline...")
         args_copy = deepcopy(training_args)
@@ -59,9 +59,8 @@ def main(args) -> None:
     pipeline = WanDistillationPipeline.from_pretrained(
         args.pretrained_model_name_or_path, args=args)
     
-    # Convert args to DistillationArgs for training
-    distillation_args = DistillationArgs.from_cli_args(args)
-    
+    # Convert args to TrainingArgs for training
+    distillation_args = TrainingArgs.from_cli_args(args)
     # Initialize the distillation pipeline
     pipeline.initialize_training_pipeline(distillation_args)
     pipeline.initialize_validation_pipeline(distillation_args)
@@ -73,14 +72,11 @@ def main(args) -> None:
 
 if __name__ == "__main__":
     argv = sys.argv
-    from fastvideo.v1.fastvideo_args import DistillationArgs, TrainingArgs
+    from fastvideo.v1.fastvideo_args import TrainingArgs
     from fastvideo.v1.utils import FlexibleArgumentParser
-    from fastvideo.v1.fastvideo_args import FastVideoArgs
     parser = FlexibleArgumentParser()
     parser = TrainingArgs.add_cli_args(parser)
-    parser = DistillationArgs.add_cli_args(parser)
-    parser = FastVideoArgs.add_cli_args(parser)
-    
+    parser = FastVideoArgs.add_cli_args(parser)   
     args = parser.parse_args()
     args.use_cpu_offload = False
     main(args) 
\ No newline at end of file