chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mick
2025-11-25 11:38:50 +08:00
committed by GitHub
parent 173e73fa1e
commit 9384fa2729
4 changed files with 10 additions and 87 deletions
Download Patch File Download Diff File Expand all files Collapse all files

53
python/sglang/multimodal_gen/runtime/layers/lora/linear.py
View File

@@ -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

7
python/sglang/multimodal_gen/runtime/models/encoders/qwen2_5vl.py
View File

@@ -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)

36
python/sglang/multimodal_gen/runtime/models/schedulers/scheduling_self_forcing_flow_match.py
View File

@@ -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:

1
python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py
View File

@@ -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(