From 9384fa272903b2610d1e8a34fd09ab7e9a40cd78 Mon Sep 17 00:00:00 2001 From: Mick Date: Tue, 25 Nov 2025 11:38:50 +0800 Subject: [PATCH] [diffusion] refactor: remove training-related code (#13860) --- .../runtime/layers/lora/linear.py | 53 +++---------------- .../runtime/models/encoders/qwen2_5vl.py | 7 --- .../scheduling_self_forcing_flow_match.py | 36 ++----------- .../runtime/pipelines_core/lora_pipeline.py | 1 - 4 files changed, 10 insertions(+), 87 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py index e21e4dd6e..fbe6a4495 100644 --- a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py +++ b/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: - assert ( - 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 + 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 diff --git a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py b/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py index 65bb83280..c354f9237 100644 --- a/python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_self_forcing_flow_match.py b/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_self_forcing_flow_match.py index 08fc4d8bb..9a4749a6d 100644 --- a/python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_self_forcing_flow_match.py +++ b/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: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py index 6f5c52a04..3af42dee4 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py +++ b/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(