[diffusion] refactor: remove training-related code (#13860)

python/sglang/multimodal_gen/runtime/layers/lora/linear.py

@@ -3,7 +3,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Code adapted from SGLang https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/lora/layers.py
 
-import math
 
 import torch
 from torch import nn
@@ -44,7 +43,6 @@ class BaseLayerWithLoRA(nn.Module):
         base_layer: nn.Module,
         lora_rank: int | None = None,
         lora_alpha: int | None = None,
-        training_mode: bool = False,
     ):
         super().__init__()
         self.base_layer: nn.Module = base_layer
@@ -56,39 +54,10 @@ class BaseLayerWithLoRA(nn.Module):
         self.disable_lora: bool = False
         self.lora_rank = lora_rank
         self.lora_alpha = lora_alpha
-        self.training_mode = training_mode
         self.lora_path: str | None = None
 
-        if training_mode:
+        self.lora_A = None
-            assert (
+        self.lora_B = None
-                self.lora_rank is not None
-            ), "LoRA rank  must be set for training mode"
-            if self.lora_rank is None or self.lora_alpha is None:
-                self.lora_alpha = lora_rank
-            self.base_layer.requires_grad_(False)
-            in_dim = self.base_layer.weight.shape[1]
-            out_dim = self.base_layer.weight.shape[0]
-            self.lora_A = nn.Parameter(
-                torch.zeros(
-                    self.lora_rank,
-                    in_dim,
-                    device=self.base_layer.weight.device,
-                    dtype=self.base_layer.weight.dtype,
-                )
-            )
-            self.lora_B = nn.Parameter(
-                torch.zeros(
-                    out_dim,
-                    self.lora_rank,
-                    device=self.base_layer.weight.device,
-                    dtype=self.base_layer.weight.dtype,
-                )
-            )
-            torch.nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5))
-            torch.nn.init.zeros_(self.lora_B)
-        else:
-            self.lora_A = None
-            self.lora_B = None
 
     @torch.compile()
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -122,7 +91,6 @@ class BaseLayerWithLoRA(nn.Module):
         self,
         A: torch.Tensor,
         B: torch.Tensor,
-        training_mode: bool = False,
         lora_path: str | None = None,
     ) -> None:
         self.lora_A = torch.nn.Parameter(
@@ -130,8 +98,7 @@ class BaseLayerWithLoRA(nn.Module):
         )  # share storage with weights in the pipeline
         self.lora_B = torch.nn.Parameter(B)
         self.disable_lora = False
-        if not training_mode:
+        self.merge_lora_weights()
-            self.merge_lora_weights()
         self.lora_path = lora_path
 
     @torch.no_grad()
@@ -245,9 +212,8 @@ class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA):
         base_layer: ColumnParallelLinear,
         lora_rank: int | None = None,
         lora_alpha: int | None = None,
-        training_mode: bool = False,
     ) -> None:
-        super().__init__(base_layer, lora_rank, lora_alpha, training_mode)
+        super().__init__(base_layer, lora_rank, lora_alpha)
 
     def forward(self, input_: torch.Tensor) -> torch.Tensor:
         # duplicate the logic in ColumnParallelLinear
@@ -281,9 +247,8 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
         base_layer: MergedColumnParallelLinear,
         lora_rank: int | None = None,
         lora_alpha: int | None = None,
-        training_mode: bool = False,
     ) -> None:
-        super().__init__(base_layer, lora_rank, lora_alpha, training_mode)
+        super().__init__(base_layer, lora_rank, lora_alpha)
 
     def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
         return A.to(self.base_layer.weight)
@@ -304,9 +269,8 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA):
         base_layer: QKVParallelLinear,
         lora_rank: int | None = None,
         lora_alpha: int | None = None,
-        training_mode: bool = False,
     ) -> None:
-        super().__init__(base_layer, lora_rank, lora_alpha, training_mode)
+        super().__init__(base_layer, lora_rank, lora_alpha)
 
     def slice_lora_a_weights(self, A: torch.Tensor) -> torch.Tensor:
         return A
@@ -338,9 +302,8 @@ class RowParallelLinearWithLoRA(BaseLayerWithLoRA):
         base_layer: RowParallelLinear,
         lora_rank: int | None = None,
         lora_alpha: int | None = None,
-        training_mode: bool = False,
     ) -> None:
-        super().__init__(base_layer, lora_rank, lora_alpha, training_mode)
+        super().__init__(base_layer, lora_rank, lora_alpha)
 
     def forward(self, input_: torch.Tensor):
         # duplicate the logic in RowParallelLinear
@@ -392,7 +355,6 @@ def get_lora_layer(
     layer: nn.Module,
     lora_rank: int | None = None,
     lora_alpha: int | None = None,
-    training_mode: bool = False,
 ) -> BaseLayerWithLoRA | None:
     supported_layer_types: dict[type[LinearBase], type[BaseLayerWithLoRA]] = {
         # the order matters
@@ -409,7 +371,6 @@ def get_lora_layer(
                 layer,
                 lora_rank=lora_rank,
                 lora_alpha=lora_alpha,
-                training_mode=training_mode,
             )
             return ret
     return None

python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py

@@ -389,13 +389,6 @@ class Qwen2_5_VLTextModel(nn.Module):
                 "You must specify exactly one of input_ids or inputs_embeds"
             )
 
-        if self.gradient_checkpointing and self.training:
-            if use_cache:
-                logger.warn(
-                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
-                )
-                use_cache = False
-
         # torch.jit.trace() doesn't support cache objects in the output
         if use_cache and past_key_values is None and not torch.jit.is_tracing():
             past_key_values = DynamicCache(config=self.config)

python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_self_forcing_flow_match.py

@@ -26,7 +26,6 @@ class SelfForcingFlowMatchSchedulerOutput(BaseOutput):
 
 
 class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
-
     config_name = "scheduler_config.json"
     order = 1
 
@@ -41,7 +40,8 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
         inverse_timesteps=False,
         extra_one_step=False,
         reverse_sigmas=False,
-        training=False,
+        *args,
+        **kwargs,
     ):
         self.num_train_timesteps = num_train_timesteps
         self.shift = shift
@@ -50,13 +50,12 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
         self.inverse_timesteps = inverse_timesteps
         self.extra_one_step = extra_one_step
         self.reverse_sigmas = reverse_sigmas
-        self.set_timesteps(num_inference_steps, training=training)
+        self.set_timesteps(num_inference_steps)
 
     def set_timesteps(
         self,
         num_inference_steps=100,
         denoising_strength=1.0,
-        training=False,
         return_dict=False,
         **kwargs,
     ):
@@ -77,14 +76,6 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
         if self.reverse_sigmas:
             self.sigmas = 1 - self.sigmas
         self.timesteps = self.sigmas * self.num_train_timesteps
-        if training:
-            x = self.timesteps
-            y = torch.exp(
-                -2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2
-            )
-            y_shifted = y - y.min()
-            bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum())
-            self.linear_timesteps_weights = bsmntw_weighing
 
     def step(
         self,
@@ -139,27 +130,6 @@ class SelfForcingFlowMatchScheduler(BaseScheduler, ConfigMixin, SchedulerMixin):
         sample = (1 - sigma) * original_samples + sigma * noise
         return sample.type_as(noise)
 
-    def training_target(self, sample, noise, timestep):
-        target = noise - sample
-        return target
-
-    def training_weight(self, timestep):
-        """
-        Input:
-            - timestep: the timestep with shape [B*T]
-        Output: the corresponding weighting [B*T]
-        """
-        if timestep.ndim == 2:
-            timestep = timestep.flatten(0, 1)
-        self.linear_timesteps_weights = self.linear_timesteps_weights.to(
-            timestep.device
-        )
-        timestep_id = torch.argmin(
-            (self.timesteps.unsqueeze(1) - timestep.unsqueeze(0)).abs(), dim=0
-        )
-        weights = self.linear_timesteps_weights[timestep_id]
-        return weights
-
     def scale_model_input(
         self, sample: torch.Tensor, timestep: int | None = None
     ) -> torch.Tensor:

python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py

@@ -29,7 +29,6 @@ logger = init_logger(__name__)
 class LoRAPipeline(ComposedPipelineBase):
     """
     Pipeline that supports injecting LoRA adapters into the diffusion transformer.
-    TODO: support training.
     """
 
     lora_adapters: dict[str, dict[str, torch.Tensor]] = defaultdict(