unify parallel entrypoint

Author: Musisoul

lightx2v_train/lightx2v_train/runtime/ddp.py

@@ -100,14 +100,15 @@ def apply_ddp(model, config):
         logger.info("DP(DDP) skipped for {} because the denoiser has no trainable parameters.", model.__class__.__name__)
         return model
 
-    wrapped = LightX2VDistributedDataParallel(denoiser, **_ddp_kwargs(config))
+    ddp_kwargs = _ddp_kwargs(config)
+    wrapped = LightX2VDistributedDataParallel(denoiser, **ddp_kwargs)
     if getattr(model, "transformer", None) is not denoiser:
         raise RuntimeError(f"{model.__class__.__name__} must store its trainable denoiser in self.transformer to use DP(DDP).")
     model.transformer = wrapped
     logger.info(
         "DP(DDP) transformer wrapped: broadcast_buffers={} find_unused_parameters={} static_graph={}",
-        wrapped.broadcast_buffers,
-        wrapped.find_unused_parameters,
-        wrapped.static_graph,
+        ddp_kwargs["broadcast_buffers"],
+        ddp_kwargs["find_unused_parameters"],
+        ddp_kwargs["static_graph"],
     )
     return model

lightx2v_train/lightx2v_train/runtime/fsdp.py

@@ -2,14 +2,11 @@
 from loguru import logger
 from torch.distributed.fsdp import FSDPModule, MixedPrecisionPolicy, fully_shard
 
-from lightx2v_train.runtime.ddp import ddp_enabled
 from lightx2v_train.runtime.distributed import get_device_mesh, is_distributed
 from lightx2v_train.utils.utils import get_running_dtype
 
 
 def fsdp2_enabled(config):
-    if ddp_enabled(config):
-        return False
     fsdp_config = config.get("distributed", {}).get("fsdp2", {})
     return is_distributed() and fsdp_config.get("enabled", True)
 

lightx2v_train/lightx2v_train/runtime/parallel.py

@@ -0,0 +1,26 @@
+from loguru import logger
+
+from lightx2v_train.runtime.ddp import apply_ddp, ddp_enabled, set_ddp_gradient_sync
+from lightx2v_train.runtime.distributed import is_distributed
+from lightx2v_train.runtime.fsdp import apply_fsdp2, fsdp2_enabled
+
+
+def apply_parallel(model, config):
+    """Apply the configured distributed parallel strategy exactly once."""
+
+    if not is_distributed():
+        return model
+
+    if ddp_enabled(config):
+        return apply_ddp(model, config)
+
+    if fsdp2_enabled(config):
+        return apply_fsdp2(model, config)
+
+    logger.warning("Distributed training is initialized, but neither DP(DDP) nor FSDP2 is enabled. The model will run without distributed wrapping.")
+    return model
+
+
+def set_parallel_gradient_sync(model, enabled):
+    model.set_fsdp2_gradient_sync(enabled)
+    set_ddp_gradient_sync(model.denoiser_module(), enabled)

lightx2v_train/lightx2v_train/trainers/base.py

@@ -9,9 +9,8 @@
 
 from lightx2v_train.infer import build_inferencer
 from lightx2v_train.runtime.checkpoint import find_latest_checkpoint, parse_checkpoint_iteration, prune_checkpoints
-from lightx2v_train.runtime.ddp import apply_ddp, set_ddp_gradient_sync
 from lightx2v_train.runtime.distributed import barrier, get_world_size, is_main_process
-from lightx2v_train.runtime.fsdp import apply_fsdp2
+from lightx2v_train.runtime.parallel import apply_parallel, set_parallel_gradient_sync
 from lightx2v_train.schedulers.flow_matching import RectifiedFlowMatchingScheduler
 from lightx2v_train.utils.utils import get_running_dtype
 
@@ -121,13 +120,11 @@ def _build_lr_scheduler(self, optimizer, num_training_steps=None, num_warmup_ste
     def setup(self, resume_ckpt_path=None):
         self._setup_trainable_model(self.model)
 
-        apply_fsdp2(self.model, self.config)
+        apply_parallel(self.model, self.config)
 
         if self.gradient_checkpointing:
             self.model.enable_gradient_checkpointing()
 
-        apply_ddp(self.model, self.config)
-
         if self.infer_every_iters:
             self.inferencer = build_inferencer(self.config)
             self.inferencer.set_model(self.model)
@@ -225,8 +222,7 @@ def _resolve_resume(self):
         return ckpt_path, current_iter
 
     def _set_gradient_sync(self, enabled):
-        self.model.set_fsdp2_gradient_sync(enabled)
-        set_ddp_gradient_sync(self.model.denoiser_module(), enabled)
+        set_parallel_gradient_sync(self.model, enabled)
 
     def run_inference(self, current_iter):
         base_output_dir = self.infer_config.get("output_dir", "./output_infer")

lightx2v_train/lightx2v_train/trainers/dmd.py

@@ -12,9 +12,8 @@
 
 from lightx2v_train.model_zoo import build_model
 from lightx2v_train.runtime.checkpoint import prune_checkpoints
-from lightx2v_train.runtime.ddp import apply_ddp, set_ddp_gradient_sync
 from lightx2v_train.runtime.distributed import barrier, get_world_size, is_distributed, is_main_process, reduce_mean
-from lightx2v_train.runtime.fsdp import apply_fsdp2
+from lightx2v_train.runtime.parallel import apply_parallel, set_parallel_gradient_sync
 from lightx2v_train.schedulers import DMDFlowMatchingScheduler
 from lightx2v_train.schedulers.flow_matching import CausalForcingFlowMatchScheduler
 from lightx2v_train.utils.registry import TRAINER_REGISTER
@@ -71,10 +70,9 @@ def setup(self, resume_ckpt_path=None):
         self.fake_model = build_model(fake_model_config)
         self.fake_model.load_components(transformer_only=True, reference_model=self.model)
         self._setup_trainable_model(self.fake_model)
-        apply_fsdp2(self.fake_model, self.config)
+        apply_parallel(self.fake_model, self.config)
         if self.gradient_checkpointing:
             self.fake_model.enable_gradient_checkpointing()
-        apply_ddp(self.fake_model, self.config)
 
         teacher_model_config = copy.deepcopy(self.config)
         teacher_model_config["model"] = copy.deepcopy(base_model_config)
@@ -83,7 +81,7 @@ def setup(self, resume_ckpt_path=None):
         self.teacher_model.load_components(transformer_only=True, reference_model=self.model)
         self.teacher_model.transformer.requires_grad_(False)
         self.teacher_model.transformer.eval()
-        apply_fsdp2(self.teacher_model, self.config)
+        apply_parallel(self.teacher_model, self.config)
         self.teacher_model.transformer.eval()
 
         self.fake_trainable_params = list(self.fake_model.trainable_parameters())
@@ -445,12 +443,10 @@ def train(self):
         logger.info("[train] finished iter={}/{}", current_iter, max_train_iters)
 
     def _set_student_gradient_sync(self, enabled):
-        self.model.set_fsdp2_gradient_sync(enabled)
-        set_ddp_gradient_sync(self.model.denoiser_module(), enabled)
+        set_parallel_gradient_sync(self.model, enabled)
 
     def _set_fake_gradient_sync(self, enabled):
-        self.fake_model.set_fsdp2_gradient_sync(enabled)
-        set_ddp_gradient_sync(self.fake_model.denoiser_module(), enabled)
+        set_parallel_gradient_sync(self.fake_model, enabled)
 
     def _set_gradient_sync(self, enabled):
         self._set_student_gradient_sync(enabled)

lightx2v_train/lightx2v_train/trainers/dopsd.py

@@ -16,7 +16,7 @@
 from lightx2v_train.infer.dopsd_trajectory_viz import save_student_teacher_trajectory_grid
 from lightx2v_train.runtime.checkpoint import find_latest_checkpoint, parse_checkpoint_iteration, prune_checkpoints
 from lightx2v_train.runtime.distributed import barrier, get_rank, get_world_size, is_distributed, is_main_process, reduce_mean
-from lightx2v_train.runtime.fsdp import apply_fsdp2
+from lightx2v_train.runtime.parallel import apply_parallel
 from lightx2v_train.utils.registry import TRAINER_REGISTER
 from lightx2v_train.utils.utils import get_running_dtype
 
@@ -102,7 +102,7 @@ def setup(self, resume_ckpt_path=None):
         )
         self.model.set_dual_lora_trainable(self.student_adapter, self.teacher_adapter)
 
-        apply_fsdp2(self.model, self.config)
+        apply_parallel(self.model, self.config)
 
         if self.gradient_checkpointing:
             self.model.enable_gradient_checkpointing()