[Diffusion][SenseNova] support SenseNova-U1.5-8B-MoT (#36606)
Co-authored-by: wuyuefeng <wuyuefeng@noreply.gitcode.com> Co-authored-by: Xiaoyu Zhang <1182563586@qq.com> Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
Yuefeng Wu
co-authored by
wuyuefeng
Xiaoyu Zhang
ronnie_zheng
wuyuefeng
Xiaoyu Zhang
ronnie_zheng
parent
78da625190
commit
daf66f6670
@@ -0,0 +1,140 @@
|
||||
---
|
||||
title: SenseNova-U1.5-8B-MoT
|
||||
metatags:
|
||||
description: "Run SenseNova-U1.5-8B-MoT text-to-image generation with SGLang's native multimodal generation runtime."
|
||||
---
|
||||
|
||||
import { DiffusionModelTags } from '/src/snippets/diffusion/model-tags.jsx';
|
||||
|
||||
<DiffusionModelTags tags={["image", "text-to-image", "Chinese typography", "poster layouts", "native pipeline"]} />
|
||||
|
||||
## 1. Model Introduction
|
||||
|
||||
[SenseNova-U1.5-8B-MoT](https://huggingface.co/sensenova/SenseNova-U1.5-8B-MoT) is a text-to-image model in the SenseNova-U1 family. It is useful for high-resolution poster, infographic, and enterprise visual-generation workloads where prompt following, layout structure, and Chinese text rendering are important.
|
||||
|
||||
Built on [NEO-Unify](https://huggingface.co/blog/sensenova/neo-unify), this release improves composition, color harmony, material rendering, and local visual detail for native high-resolution generation. It also focuses on clearer Chinese and English text rendering, stronger infographic layout structure, and more reliable execution of complex prompts with object counts, spatial relationships, styles, and multiple constraints.
|
||||
|
||||
## 2. SGLang-diffusion Installation
|
||||
|
||||
SGLang-diffusion offers multiple installation methods. You can choose the most suitable installation method based on your hardware platform and requirements.
|
||||
|
||||
Please refer to the [official SGLang-diffusion installation guide](../../../docs/sglang-diffusion/installation) for installation instructions.
|
||||
|
||||
## 3. Model Deployment
|
||||
|
||||
### 3.1 Offline Image Generation
|
||||
|
||||
SenseNova-U1.5-8B-MoT is a text-to-image model. You can run a single offline generation job with `sglang generate`:
|
||||
|
||||
```bash Command
|
||||
sglang generate \
|
||||
--model-path sensenova/SenseNova-U1.5-8B-MoT \
|
||||
--prompt "A cinematic mountain lake at sunrise, realistic photography." \
|
||||
--width 2048 \
|
||||
--height 2048 \
|
||||
--guidance-scale 4.0 \
|
||||
--num-inference-steps 50 \
|
||||
--seed 42 \
|
||||
--output-file-path output_sensenova_u15.png
|
||||
```
|
||||
|
||||
### 3.2 Recommended Resolutions
|
||||
|
||||
SenseNova-U1 was trained around the 2K-pixel resolution tier. The SGLang pipeline only supports positive `--width` and `--height` values that are divisible by 32, and generation quality may degrade for untrained size combinations.
|
||||
|
||||
Recommended resolution tiers:
|
||||
|
||||
| Aspect ratio | Resolution |
|
||||
| --- | --- |
|
||||
| 1:1 | `2048 x 2048` |
|
||||
| 16:9 / 9:16 | `2720 x 1536` / `1536 x 2720` |
|
||||
| 3:2 / 2:3 | `2496 x 1664` / `1664 x 2496` |
|
||||
| 4:3 / 3:4 | `2368 x 1760` / `1760 x 2368` |
|
||||
| 2:1 / 1:2 | `2880 x 1440` / `1440 x 2880` |
|
||||
| 3:1 / 1:3 | `3456 x 1152` / `1152 x 3456` |
|
||||
| 9:16 4K | `2304 x 4096` |
|
||||
|
||||
## 4. API Usage
|
||||
|
||||
For complete API documentation, please refer to the [official API usage guide](../../../docs/sglang-diffusion/api/openai_api).
|
||||
|
||||
### 4.1 Generate an Image
|
||||
|
||||
```python Example
|
||||
import base64
|
||||
from openai import OpenAI
|
||||
|
||||
client = OpenAI(api_key="EMPTY", base_url="http://localhost:30000/v1")
|
||||
|
||||
response = client.images.generate(
|
||||
model="sensenova/SenseNova-U1.5-8B-MoT",
|
||||
prompt="A clean technology poster with structured layout and clear typography.",
|
||||
n=1,
|
||||
response_format="b64_json",
|
||||
)
|
||||
|
||||
image_bytes = base64.b64decode(response.data[0].b64_json)
|
||||
with open("output_sensenova_u15.png", "wb") as f:
|
||||
f.write(image_bytes)
|
||||
```
|
||||
|
||||
## 5. Benchmark
|
||||
|
||||
### 5.1 Speedup Benchmark
|
||||
|
||||
#### 5.1.1 Single-run Profile
|
||||
|
||||
<Tabs>
|
||||
<Tab title="NVIDIA A800">
|
||||
|
||||
This is a single-run profile for smoke validation rather than a formal benchmark. Environment: NVIDIA A800; workload: vbench text-to-image, 1 prompt, concurrency 1, 1 output per prompt; warmup: none; repetitions: 1; dispersion: not reported; baseline: none; revisions: PR checkout and resolved model snapshot at run time; output validation: request completed and returned one image.
|
||||
|
||||
**Server Command**:
|
||||
|
||||
```shell Command
|
||||
sglang serve \
|
||||
--model-path sensenova/SenseNova-U1.5-8B-MoT \
|
||||
--port 30000
|
||||
```
|
||||
|
||||
**Benchmark Command**:
|
||||
|
||||
```shell Command
|
||||
python3 -m sglang.multimodal_gen.benchmarks.bench_serving \
|
||||
--dataset vbench \
|
||||
--task text-to-image \
|
||||
--num-prompts 1 \
|
||||
--max-concurrency 1 \
|
||||
--port 30000
|
||||
```
|
||||
|
||||
**Result**:
|
||||
|
||||
```text Output
|
||||
================= Serving Benchmark Result =================
|
||||
Task: text-to-image
|
||||
Model: /workspace/models/ModelScope/SenseNova/SenseNova-U1.5-8B-MoT
|
||||
Dataset: vbench
|
||||
--------------------------------------------------
|
||||
Benchmark duration (s): 13.10
|
||||
Request rate: inf
|
||||
Max request concurrency: 1
|
||||
Successful requests: 1/1
|
||||
Completed outputs: 1
|
||||
Outputs per prompt: 1
|
||||
--------------------------------------------------
|
||||
Request throughput (req/s): 0.08
|
||||
Output throughput (outputs/s): 0.08
|
||||
Latency Mean (s): 13.10
|
||||
Latency Median (s): 13.10
|
||||
Latency P90 (s): 13.10
|
||||
Latency P95 (s): 13.10
|
||||
Latency P99 (s): 13.10
|
||||
--------------------------------------------------
|
||||
Peak Memory Max (MB): 34326.00
|
||||
Peak Memory Mean (MB): 34326.00
|
||||
Peak Memory Median (MB): 34326.00
|
||||
------------------------------------------------------------
|
||||
```
|
||||
</Tab>
|
||||
</Tabs>
|
||||
@@ -31,6 +31,12 @@ Image models generate one image request as a bounded denoising job, usually with
|
||||
href="/cookbook/diffusion/Qwen-Image/Qwen-Image"
|
||||
img="/cards/logos/qwen.png"
|
||||
/>
|
||||
<Card
|
||||
title="SenseNova-U1.5"
|
||||
mode="card"
|
||||
href="/cookbook/diffusion/SenseNova/SenseNova-U1.5-8B-MoT"
|
||||
img="/cards/logos/internvl.png"
|
||||
/>
|
||||
<Card
|
||||
title="LongCat-Image"
|
||||
mode="card"
|
||||
|
||||
@@ -1512,6 +1512,13 @@
|
||||
"cookbook/diffusion/Qwen-Image/Qwen-Image-Edit"
|
||||
]
|
||||
},
|
||||
{
|
||||
"group": "SenseNova",
|
||||
"tag": "NEW",
|
||||
"pages": [
|
||||
"cookbook/diffusion/SenseNova/SenseNova-U1.5-8B-MoT"
|
||||
]
|
||||
},
|
||||
{
|
||||
"group": "LongCat-Image",
|
||||
"tag": "NEW",
|
||||
|
||||
@@ -103,6 +103,11 @@ export const DiffusionModelCatalog = ({ category }) => {
|
||||
name: "JoyAI-Image",
|
||||
modelIds: ["jdopensource/JoyAI-Image-Edit-Diffusers"],
|
||||
},
|
||||
{
|
||||
name: "SenseNova-U1.5",
|
||||
modelIds: ["sensenova/SenseNova-U1.5-8B-MoT"],
|
||||
cookbook: "/cookbook/diffusion/SenseNova/SenseNova-U1.5-8B-MoT",
|
||||
},
|
||||
{
|
||||
name: "GLM-Image",
|
||||
modelIds: ["zai-org/GLM-Image"],
|
||||
|
||||
@@ -53,6 +53,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.sana import SanaPipelineConf
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.sana_video import (
|
||||
SanaVideoPipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.sensenova_u1 import (
|
||||
SenseNovaU1PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.stablediffusion3 import (
|
||||
StableDiffusion3PipelineConfig,
|
||||
)
|
||||
@@ -83,6 +86,7 @@ __all__ = [
|
||||
"PipelineConfig",
|
||||
"SanaPipelineConfig",
|
||||
"SanaVideoPipelineConfig",
|
||||
"SenseNovaU1PipelineConfig",
|
||||
"SlidingTileAttnConfig",
|
||||
"MOVAPipelineConfig",
|
||||
"Pi05PipelineConfig",
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from dataclasses import dataclass
|
||||
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.base import (
|
||||
ModelTaskType,
|
||||
PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
|
||||
ModelDeploymentConfig,
|
||||
)
|
||||
|
||||
|
||||
def _is_runtime_option_requested(value) -> bool:
|
||||
if value is None:
|
||||
return False
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
if isinstance(value, (dict, list, tuple, set)):
|
||||
return bool(value)
|
||||
return True
|
||||
|
||||
|
||||
def _is_arg_explicitly_set(server_args, option: str) -> bool:
|
||||
is_explicit = getattr(server_args, "is_arg_explicitly_set", None)
|
||||
if callable(is_explicit):
|
||||
return is_explicit(option)
|
||||
return _is_runtime_option_requested(getattr(server_args, option, None))
|
||||
|
||||
|
||||
def _component_residency_requests_offload(value) -> bool:
|
||||
if not _is_runtime_option_requested(value):
|
||||
return False
|
||||
if isinstance(value, dict):
|
||||
values = value.values()
|
||||
elif isinstance(value, str):
|
||||
values = value.split(",")
|
||||
else:
|
||||
values = value
|
||||
|
||||
for raw_value in values:
|
||||
mode = str(raw_value).split("=", 1)[-1].strip().replace("_", "-").lower()
|
||||
if mode in ("component-offload", "layerwise-offload"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _set_compatible_runtime_defaults(server_args) -> None:
|
||||
compatible_defaults = {
|
||||
"component_residency": None,
|
||||
"cpu_offload_components": None,
|
||||
"dit_cpu_offload": False,
|
||||
"text_encoder_cpu_offload": False,
|
||||
"image_encoder_cpu_offload": False,
|
||||
"vae_cpu_offload": False,
|
||||
"dit_layerwise_offload": False,
|
||||
"layerwise_offload_components": None,
|
||||
"quantization": None,
|
||||
"quantization_ignored_layers": None,
|
||||
"transformer_weights_path": None,
|
||||
"component_paths": {},
|
||||
"component_weights_paths": {},
|
||||
"component_quantizations": {},
|
||||
"component_quantization_ignored_layers": {},
|
||||
"component_precisions": {},
|
||||
"attention_backend": None,
|
||||
"component_attention_backends": {},
|
||||
"attention_backend_config": None,
|
||||
}
|
||||
for option, value in compatible_defaults.items():
|
||||
if not _is_arg_explicitly_set(server_args, option):
|
||||
setattr(server_args, option, value)
|
||||
|
||||
|
||||
@dataclass
|
||||
class SenseNovaU1PipelineConfig(PipelineConfig):
|
||||
"""Native SenseNova-U1 text-to-image pipeline configuration."""
|
||||
|
||||
task_type: ModelTaskType = ModelTaskType.T2I
|
||||
model_precision: str = "bf16"
|
||||
should_use_guidance: bool = True
|
||||
supports_cfg_parallel: bool = False
|
||||
|
||||
def supports_dynamic_batching(self):
|
||||
return False
|
||||
|
||||
def supports_disaggregation(self) -> bool:
|
||||
return False
|
||||
|
||||
def supports_sequential_multi_output_inference(self):
|
||||
return True
|
||||
|
||||
def validate_server_args(self, server_args) -> None:
|
||||
if server_args.num_gpus != 1:
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline currently supports num_gpus=1. "
|
||||
"Native tensor/pipeline parallelism is not implemented yet."
|
||||
)
|
||||
if getattr(server_args, "enable_torch_compile", False):
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline does not support torch.compile yet. "
|
||||
"Please omit --enable-torch-compile."
|
||||
)
|
||||
if getattr(server_args, "lora_path", None):
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline does not support LoRA adapters yet. "
|
||||
"Please omit --lora-path."
|
||||
)
|
||||
_set_compatible_runtime_defaults(server_args)
|
||||
if _is_arg_explicitly_set(
|
||||
server_args, "component_residency"
|
||||
) and _component_residency_requests_offload(
|
||||
getattr(server_args, "component_residency", None)
|
||||
):
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline does not support component residency "
|
||||
"offload modes yet. Please omit --component-residency."
|
||||
)
|
||||
unsupported_runtime_options = {
|
||||
"cpu_offload_components": "CPU offload",
|
||||
"dit_cpu_offload": "DiT CPU offload",
|
||||
"text_encoder_cpu_offload": "text encoder CPU offload",
|
||||
"image_encoder_cpu_offload": "image encoder CPU offload",
|
||||
"vae_cpu_offload": "VAE CPU offload",
|
||||
"dit_layerwise_offload": "DiT layerwise offload",
|
||||
"layerwise_offload_components": "layerwise offload",
|
||||
"quantization": "quantization",
|
||||
"quantization_ignored_layers": "quantization ignored layers",
|
||||
"transformer_weights_path": "pre-quantized transformer weights",
|
||||
"component_paths": "component path overrides",
|
||||
"component_weights_paths": "component weight path overrides",
|
||||
"component_quantizations": "component quantization",
|
||||
"component_quantization_ignored_layers": (
|
||||
"component quantization ignored layers"
|
||||
),
|
||||
"component_precisions": "component precision overrides",
|
||||
}
|
||||
for option, description in unsupported_runtime_options.items():
|
||||
if _is_arg_explicitly_set(
|
||||
server_args, option
|
||||
) and _is_runtime_option_requested(getattr(server_args, option, None)):
|
||||
raise ValueError(
|
||||
f"SenseNovaU1Pipeline does not support {description} yet. "
|
||||
f"Please omit --{option.replace('_', '-')}."
|
||||
)
|
||||
if _is_arg_explicitly_set(
|
||||
server_args, "attention_backend"
|
||||
) and _is_runtime_option_requested(
|
||||
getattr(server_args, "attention_backend", None)
|
||||
):
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline does not support custom attention backends yet. "
|
||||
"Please omit --attention-backend."
|
||||
)
|
||||
if _is_arg_explicitly_set(
|
||||
server_args, "component_attention_backends"
|
||||
) and _is_runtime_option_requested(
|
||||
getattr(server_args, "component_attention_backends", None)
|
||||
):
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline does not support component attention backends yet. "
|
||||
"Please omit --component-attention-backends."
|
||||
)
|
||||
if _is_arg_explicitly_set(
|
||||
server_args, "attention_backend_config"
|
||||
) and _is_runtime_option_requested(
|
||||
getattr(server_args, "attention_backend_config", None)
|
||||
):
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline does not support attention backend config yet. "
|
||||
"Please omit --attention-backend-config."
|
||||
)
|
||||
|
||||
def get_model_deployment_config(self) -> ModelDeploymentConfig:
|
||||
return ModelDeploymentConfig(
|
||||
speed_mode_enable_torch_compile_by_default=False,
|
||||
keep_resident_min_available_gb=80,
|
||||
auto_enable_cfg_parallel=False,
|
||||
supports_cfg_parallel=False,
|
||||
)
|
||||
@@ -10,6 +10,9 @@ from sglang.multimodal_gen.configs.sample.lingbot_video_moe import (
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.pi05 import Pi05SamplingParams
|
||||
from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams
|
||||
from sglang.multimodal_gen.configs.sample.sensenova_u1 import (
|
||||
SenseNovaU1SamplingParams,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"SamplingParams",
|
||||
@@ -18,4 +21,5 @@ __all__ = [
|
||||
"Ideogram4SamplingParams",
|
||||
"Pi05SamplingParams",
|
||||
"LingBotVideoMoESamplingParams",
|
||||
"SenseNovaU1SamplingParams",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,112 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
from sglang.multimodal_gen.configs.sample.sampling_params import (
|
||||
DataType,
|
||||
SamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sensenova_u1 import (
|
||||
DEFAULT_CFG_INTERVAL,
|
||||
DEFAULT_CFG_NORM,
|
||||
DEFAULT_ENABLE_TIMESTEP_SHIFT,
|
||||
DEFAULT_T_EPS,
|
||||
DEFAULT_THINK_MODE,
|
||||
DEFAULT_TIMESTEP_SHIFT,
|
||||
SENSENOVA_U1_CFG_NORM_CHOICES,
|
||||
SENSENOVA_U1_REQUEST_EXTRA_KEY,
|
||||
SENSENOVA_U1_RESOLUTION_ALIGNMENT,
|
||||
)
|
||||
|
||||
_PUBLIC_OVERRIDE_FIELDS = {
|
||||
"prompt",
|
||||
"prompt_path",
|
||||
"height",
|
||||
"width",
|
||||
"num_inference_steps",
|
||||
"guidance_scale",
|
||||
"num_outputs_per_prompt",
|
||||
"seed",
|
||||
"save_output",
|
||||
"output_path",
|
||||
"output_file_name",
|
||||
"output_quality",
|
||||
"output_compression",
|
||||
"quality",
|
||||
}
|
||||
|
||||
|
||||
@dataclass
|
||||
class SenseNovaU1SamplingParams(SamplingParams):
|
||||
data_type: DataType = field(default=DataType.IMAGE, init=False)
|
||||
height: int = 2048
|
||||
width: int = 2048
|
||||
num_frames: int = 1
|
||||
fps: int = 1
|
||||
num_inference_steps: int = 50
|
||||
guidance_scale: float = 4.0
|
||||
cfg_norm: str = DEFAULT_CFG_NORM
|
||||
timestep_shift: float = DEFAULT_TIMESTEP_SHIFT
|
||||
enable_timestep_shift: bool = DEFAULT_ENABLE_TIMESTEP_SHIFT
|
||||
cfg_interval: tuple[float, float] = DEFAULT_CFG_INTERVAL
|
||||
t_eps: float = DEFAULT_T_EPS
|
||||
think_mode: bool = DEFAULT_THINK_MODE
|
||||
negative_prompt: None = field(default=None, init=False)
|
||||
|
||||
@classmethod
|
||||
def supported_override_fields(cls) -> set[str]:
|
||||
return set(_PUBLIC_OVERRIDE_FIELDS)
|
||||
|
||||
@classmethod
|
||||
def get_cli_args(cls, args):
|
||||
cli_args = super().get_cli_args(args)
|
||||
return {
|
||||
key: value
|
||||
for key, value in cli_args.items()
|
||||
if key in _PUBLIC_OVERRIDE_FIELDS
|
||||
}
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if isinstance(self.cfg_interval, list):
|
||||
self.cfg_interval = tuple(float(x) for x in self.cfg_interval)
|
||||
super().__post_init__()
|
||||
|
||||
def _validate(self) -> None:
|
||||
super()._validate()
|
||||
if (
|
||||
self.width % SENSENOVA_U1_RESOLUTION_ALIGNMENT != 0
|
||||
or self.height % SENSENOVA_U1_RESOLUTION_ALIGNMENT != 0
|
||||
):
|
||||
raise ValueError(
|
||||
"SenseNova-U1 requires width and height to be divisible by "
|
||||
f"{SENSENOVA_U1_RESOLUTION_ALIGNMENT}, got "
|
||||
f"{self.width}x{self.height}."
|
||||
)
|
||||
if self.num_frames != 1:
|
||||
raise ValueError(
|
||||
f"SenseNova-U1 is an image model and requires num_frames=1, got {self.num_frames}."
|
||||
)
|
||||
if self.cfg_norm not in SENSENOVA_U1_CFG_NORM_CHOICES:
|
||||
raise ValueError(
|
||||
f"cfg_norm must be one of {SENSENOVA_U1_CFG_NORM_CHOICES}, "
|
||||
f"got {self.cfg_norm!r}"
|
||||
)
|
||||
if len(self.cfg_interval) != 2:
|
||||
raise ValueError("cfg_interval must contain exactly two values")
|
||||
start, end = self.cfg_interval
|
||||
if not 0.0 <= float(start) <= float(end) <= 1.0:
|
||||
raise ValueError(
|
||||
f"cfg_interval must satisfy 0 <= start <= end <= 1, got {self.cfg_interval!r}"
|
||||
)
|
||||
|
||||
def build_request_extra(self) -> dict[str, Any]:
|
||||
extra = super().build_request_extra()
|
||||
extra[SENSENOVA_U1_REQUEST_EXTRA_KEY] = {
|
||||
"cfg_norm": self.cfg_norm,
|
||||
"timestep_shift": self.timestep_shift,
|
||||
"enable_timestep_shift": self.enable_timestep_shift,
|
||||
"cfg_interval": tuple(self.cfg_interval),
|
||||
"t_eps": self.t_eps,
|
||||
"think_mode": self.think_mode,
|
||||
}
|
||||
return extra
|
||||
@@ -0,0 +1,55 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Shared constants for the native SenseNova-U1 integration."""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
SENSENOVA_U1_REQUEST_EXTRA_KEY = "sensenova_u1"
|
||||
|
||||
SENSENOVA_U1_MODEL_IDS = {
|
||||
"sensenova/sensenova-u1.5-8b-mot",
|
||||
}
|
||||
SENSENOVA_U1_ADAPTER_ONLY_MODEL_IDS = {
|
||||
"sensenova/sensenova-u1.5-8b-mot-loras",
|
||||
}
|
||||
|
||||
SENSENOVA_U1_CFG_NORM_CHOICES = (
|
||||
"none",
|
||||
"global",
|
||||
"channel",
|
||||
"cfg_zero_star",
|
||||
)
|
||||
SENSENOVA_U1_RESOLUTION_ALIGNMENT = 32
|
||||
|
||||
DEFAULT_CFG_NORM = "none"
|
||||
DEFAULT_TIMESTEP_SHIFT = 3.0
|
||||
DEFAULT_ENABLE_TIMESTEP_SHIFT = True
|
||||
DEFAULT_CFG_INTERVAL = (0.0, 1.0)
|
||||
DEFAULT_T_EPS = 0.02
|
||||
DEFAULT_THINK_MODE = False
|
||||
|
||||
|
||||
def is_sensenova_u1_model(model_path: str) -> bool:
|
||||
"""Identify SenseNova-U1 Hub IDs and local base checkpoints."""
|
||||
if os.path.isdir(model_path):
|
||||
config_path = os.path.join(model_path, "config.json")
|
||||
try:
|
||||
with open(config_path) as config_file:
|
||||
config = json.load(config_file)
|
||||
except (OSError, json.JSONDecodeError):
|
||||
return False
|
||||
|
||||
if not isinstance(config, dict):
|
||||
return False
|
||||
architectures = config.get("architectures", [])
|
||||
return (
|
||||
config.get("model_type") == "neo_chat"
|
||||
and isinstance(architectures, list)
|
||||
and "NEOChatModel" in architectures
|
||||
)
|
||||
|
||||
return model_path.rstrip("/").lower() in SENSENOVA_U1_MODEL_IDS
|
||||
|
||||
|
||||
def is_sensenova_u1_adapter_only_model(model_path: str) -> bool:
|
||||
return model_path.rstrip("/").lower() in SENSENOVA_U1_ADAPTER_ONLY_MODEL_IDS
|
||||
@@ -101,6 +101,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.sana_video import (
|
||||
SanaVideoPipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import SanaWMPipelineConfig
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.sensenova_u1 import (
|
||||
SenseNovaU1PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.stablediffusion3 import (
|
||||
StableDiffusion3PipelineConfig,
|
||||
)
|
||||
@@ -181,6 +184,9 @@ from sglang.multimodal_gen.configs.sample.qwenimage import (
|
||||
from sglang.multimodal_gen.configs.sample.sana import SanaSamplingParams
|
||||
from sglang.multimodal_gen.configs.sample.sana_video import SanaVideoSamplingParams
|
||||
from sglang.multimodal_gen.configs.sample.sana_wm import SanaWMSamplingParams
|
||||
from sglang.multimodal_gen.configs.sample.sensenova_u1 import (
|
||||
SenseNovaU1SamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.stablediffusion3 import (
|
||||
StableDiffusion3SamplingParams,
|
||||
)
|
||||
@@ -200,6 +206,11 @@ from sglang.multimodal_gen.configs.sample.zimage import (
|
||||
ZImageSamplingParams,
|
||||
ZImageTurboSamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sensenova_u1 import (
|
||||
SENSENOVA_U1_MODEL_IDS,
|
||||
is_sensenova_u1_adapter_only_model,
|
||||
is_sensenova_u1_model,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
|
||||
ComposedPipelineBase,
|
||||
)
|
||||
@@ -458,17 +469,24 @@ def has_registered_diffusion_model_path(model_path: str) -> bool:
|
||||
_ensure_registry_initialized()
|
||||
all_model_hf_paths = sorted(_MODEL_HF_PATH_TO_NAME.keys(), key=len, reverse=True)
|
||||
|
||||
if is_sensenova_u1_model(model_path):
|
||||
return True
|
||||
|
||||
if model_path in _MODEL_HF_PATH_TO_NAME:
|
||||
return True
|
||||
|
||||
model_short_name = get_model_short_name(model_path.lower())
|
||||
for registered_model_hf_id in all_model_hf_paths:
|
||||
if registered_model_hf_id.lower() in SENSENOVA_U1_MODEL_IDS:
|
||||
continue
|
||||
registered_model_name = get_model_short_name(registered_model_hf_id.lower())
|
||||
if registered_model_name in model_short_name:
|
||||
return True
|
||||
|
||||
normalized_model_path = _normalize_hf_cache_path(model_path)
|
||||
for registered_model_hf_id in all_model_hf_paths:
|
||||
if registered_model_hf_id.lower() in SENSENOVA_U1_MODEL_IDS:
|
||||
continue
|
||||
cache_repo_fragment = (
|
||||
f"models--{registered_model_hf_id.lower().replace('/', '--')}"
|
||||
)
|
||||
@@ -502,6 +520,13 @@ def _get_config_info(
|
||||
"falling back to automatic detection."
|
||||
)
|
||||
|
||||
# SenseNova Hub IDs require an exact match, while local checkpoints are
|
||||
# identified from their config metadata rather than their directory name.
|
||||
if is_sensenova_u1_model(model_path):
|
||||
for registered_hf_id in all_model_hf_paths:
|
||||
if registered_hf_id.lower() in SENSENOVA_U1_MODEL_IDS:
|
||||
return _CONFIG_REGISTRY.get(_MODEL_HF_PATH_TO_NAME[registered_hf_id])
|
||||
|
||||
# 1. Exact match
|
||||
if model_path in _MODEL_HF_PATH_TO_NAME:
|
||||
model_id = _MODEL_HF_PATH_TO_NAME[model_path]
|
||||
@@ -511,6 +536,8 @@ def _get_config_info(
|
||||
# 2. Partial match: find the best (longest) match against all registered model hf paths.
|
||||
model_short_name = get_model_short_name(model_path.lower())
|
||||
for registered_model_hf_id in all_model_hf_paths:
|
||||
if registered_model_hf_id.lower() in SENSENOVA_U1_MODEL_IDS:
|
||||
continue
|
||||
registered_model_name = get_model_short_name(registered_model_hf_id.lower())
|
||||
|
||||
if registered_model_name in model_short_name:
|
||||
@@ -529,6 +556,8 @@ def _get_config_info(
|
||||
# -> models--black-forest-labs--flux.2-dev-nvfp4 (to match with cache_repo_fragment)
|
||||
normalized_model_path = _normalize_hf_cache_path(model_path)
|
||||
for registered_model_hf_id in all_model_hf_paths:
|
||||
if registered_model_hf_id.lower() in SENSENOVA_U1_MODEL_IDS:
|
||||
continue
|
||||
cache_repo_fragment = (
|
||||
f"models--{registered_model_hf_id.lower().replace('/', '--')}"
|
||||
)
|
||||
@@ -661,6 +690,16 @@ def get_model_info(
|
||||
elif isinstance(backend, str):
|
||||
backend = Backend.from_string(backend)
|
||||
|
||||
if is_sensenova_u1_adapter_only_model(model_path):
|
||||
logger.error(
|
||||
"SenseNova-U1 adapter-only checkpoint '%s' does not contain base "
|
||||
"model weights or config. SenseNova-U1 adapters are not supported "
|
||||
"yet; use the base checkpoint 'sensenova/SenseNova-U1.5-8B-MoT' "
|
||||
"directly.",
|
||||
model_path,
|
||||
)
|
||||
return None
|
||||
|
||||
# Handle explicit diffusers backend
|
||||
if backend == Backend.DIFFUSERS:
|
||||
logger.info(
|
||||
@@ -980,6 +1019,13 @@ def _register_configs():
|
||||
)
|
||||
],
|
||||
)
|
||||
register_configs(
|
||||
sampling_param_cls=SenseNovaU1SamplingParams,
|
||||
pipeline_config_cls=SenseNovaU1PipelineConfig,
|
||||
hf_model_paths=[
|
||||
"sensenova/SenseNova-U1.5-8B-MoT",
|
||||
],
|
||||
)
|
||||
register_configs(
|
||||
sampling_param_cls=FastH3SamplingParams,
|
||||
pipeline_config_cls=FastH3PipelineConfig,
|
||||
@@ -1392,6 +1438,9 @@ def is_known_non_diffusers_multimodal_model(model_path: str) -> bool:
|
||||
|
||||
def get_non_diffusers_pipeline_name(model_path: str) -> Optional[str]:
|
||||
"""Get the pipeline name for a known non-diffusers model."""
|
||||
if is_sensenova_u1_model(model_path):
|
||||
return "SenseNovaU1Pipeline"
|
||||
|
||||
normalized_model_path = _normalize_hf_cache_path(model_path)
|
||||
model_short_name = get_model_short_name(normalized_model_path)
|
||||
for pattern, pipeline_name in KNOWN_NON_DIFFUSERS_DIFFUSION_MODEL_PATTERNS.items():
|
||||
|
||||
@@ -163,9 +163,12 @@ def generate_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None
|
||||
# respect config file by overriding args with args parsed from it
|
||||
if config_file:
|
||||
config_args = ServerArgs.load_config_file(config_file) or {}
|
||||
sampling_param_fields = {
|
||||
field.name for field in dataclasses.fields(sampling_params_cls)
|
||||
}
|
||||
if hasattr(sampling_params_cls, "supported_override_fields"):
|
||||
sampling_param_fields = sampling_params_cls.supported_override_fields()
|
||||
else:
|
||||
sampling_param_fields = {
|
||||
field.name for field in dataclasses.fields(sampling_params_cls)
|
||||
}
|
||||
sampling_params_kwargs.update(
|
||||
{
|
||||
key: value
|
||||
|
||||
@@ -18,7 +18,6 @@ import tempfile
|
||||
import threading
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import contextmanager
|
||||
from copy import copy
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, List, Optional, Sequence, Union
|
||||
|
||||
@@ -41,6 +40,12 @@ from sglang.multimodal_gen.configs.sample.sampling_params import (
|
||||
DataType,
|
||||
SamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.request_utils import (
|
||||
expand_request_outputs as expand_request_outputs,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.request_utils import (
|
||||
normalize_output_seeds as normalize_output_seeds,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import CYAN, RESET, init_logger
|
||||
@@ -214,129 +219,6 @@ class MaterializedOutput:
|
||||
fps: int = 0
|
||||
|
||||
|
||||
def normalize_output_seeds(
|
||||
seed: int | list[int],
|
||||
*,
|
||||
num_outputs_per_prompt: int,
|
||||
num_prompts: int = 1,
|
||||
prompt_index: int = 0,
|
||||
) -> list[int]:
|
||||
"""
|
||||
return a list of seed with size equal to `num_outputs_per_prompt`
|
||||
"""
|
||||
if num_outputs_per_prompt <= 0:
|
||||
raise ValueError(
|
||||
f"num_outputs_per_prompt must be positive, got {num_outputs_per_prompt}"
|
||||
)
|
||||
|
||||
if isinstance(seed, list):
|
||||
seeds = [int(item) for item in seed]
|
||||
total_outputs = num_outputs_per_prompt * num_prompts
|
||||
if len(seeds) == num_outputs_per_prompt:
|
||||
return seeds
|
||||
if len(seeds) == total_outputs:
|
||||
start = prompt_index * num_outputs_per_prompt
|
||||
return seeds[start : start + num_outputs_per_prompt]
|
||||
raise ValueError(
|
||||
"seed list length must match num_outputs_per_prompt "
|
||||
f"({num_outputs_per_prompt}) or total outputs ({total_outputs}), "
|
||||
f"got {len(seeds)}"
|
||||
)
|
||||
|
||||
base_seed = int(seed)
|
||||
return [base_seed + i for i in range(num_outputs_per_prompt)]
|
||||
|
||||
|
||||
def _with_output_index_suffix(output_file_name: str, output_index: int) -> str:
|
||||
base, ext = os.path.splitext(output_file_name)
|
||||
return f"{base}_{output_index}{ext}"
|
||||
|
||||
|
||||
def _copy_trace_ctx_for_output(req: Req, request_id: str | None, output_index: int):
|
||||
trace_ctx = req.trace_ctx
|
||||
if output_index == 0 or not trace_ctx.tracing_enable:
|
||||
return trace_ctx
|
||||
|
||||
output_trace_ctx = TraceReqContext(
|
||||
rid=request_id,
|
||||
module_name=trace_ctx.module_name,
|
||||
external_trace_header=trace_ctx.external_trace_header,
|
||||
)
|
||||
output_trace_ctx.trace_req_start()
|
||||
return output_trace_ctx
|
||||
|
||||
|
||||
def _copy_req_for_output(
|
||||
req: Req,
|
||||
*,
|
||||
request_id: str | None,
|
||||
output_index: int,
|
||||
) -> Req:
|
||||
"""Create a lightweight per-output ``Req`` without deep-copying tensors."""
|
||||
output_req = copy(req)
|
||||
output_req.sampling_params = copy(req.sampling_params)
|
||||
output_req.extra = dict(req.extra)
|
||||
output_req.condition_inputs = dict(req.condition_inputs)
|
||||
output_req.trace_ctx = _copy_trace_ctx_for_output(req, request_id, output_index)
|
||||
return output_req
|
||||
|
||||
|
||||
def expand_request_outputs(
|
||||
req: Req,
|
||||
*,
|
||||
num_prompts: int = 1,
|
||||
prompt_index: int = 0,
|
||||
) -> list[Req]:
|
||||
"""
|
||||
Expand a req to a list with size equal to `num_prompts`
|
||||
"""
|
||||
num_outputs = int(req.num_outputs_per_prompt)
|
||||
# each req must has different seed
|
||||
seeds = normalize_output_seeds(
|
||||
req.seed,
|
||||
num_outputs_per_prompt=num_outputs,
|
||||
num_prompts=num_prompts,
|
||||
prompt_index=prompt_index,
|
||||
)
|
||||
|
||||
if num_outputs == 1:
|
||||
req.seed = seeds[0]
|
||||
req.seeds = None
|
||||
req.generator = None
|
||||
req.sampling_params.refresh_request_extra_after_output_expansion(req)
|
||||
return [req]
|
||||
|
||||
expanded: list[Req] = []
|
||||
for output_index, seed in enumerate(seeds):
|
||||
output_request_id = (
|
||||
f"{req.request_id}:{output_index}" if req.request_id is not None else None
|
||||
)
|
||||
output_req = _copy_req_for_output(
|
||||
req, request_id=output_request_id, output_index=output_index
|
||||
)
|
||||
output_req.seed = seed
|
||||
output_req.num_outputs_per_prompt = 1
|
||||
output_req.seeds = None
|
||||
output_req.generator = None
|
||||
output_req.extra["parent_request_id"] = req.request_id
|
||||
output_req.extra["output_index"] = output_index
|
||||
|
||||
if output_request_id is not None:
|
||||
output_req.request_id = output_request_id
|
||||
|
||||
if req.output_file_name:
|
||||
output_req.output_file_name = _with_output_index_suffix(
|
||||
req.output_file_name, output_index
|
||||
)
|
||||
output_req.sampling_params.refresh_request_extra_after_output_expansion(
|
||||
output_req
|
||||
)
|
||||
output_req.validate()
|
||||
expanded.append(output_req)
|
||||
|
||||
return expanded
|
||||
|
||||
|
||||
def _normalize_audio_to_numpy(audio: Any) -> np.ndarray | None:
|
||||
"""Convert audio (torch / numpy) into a float32 numpy array in [-1, 1], best-effort."""
|
||||
if audio is None:
|
||||
|
||||
@@ -1371,13 +1371,15 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
|
||||
"""
|
||||
merge batched output
|
||||
"""
|
||||
if parts.output_file_paths:
|
||||
merged.output_file_paths = parts.output_file_paths
|
||||
if any(metrics is not None for metrics in parts.metrics_list):
|
||||
merged.metrics_list = parts.metrics_list
|
||||
merged.metrics = next(
|
||||
metrics for metrics in parts.metrics_list if metrics is not None
|
||||
)
|
||||
if merged.error is not None:
|
||||
return
|
||||
if parts.output_file_paths:
|
||||
merged.output_file_paths = parts.output_file_paths
|
||||
if parts.tensor_outputs:
|
||||
merged.output = torch.cat(parts.tensor_outputs, dim=0)
|
||||
elif parts.list_outputs:
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""SenseNova-U1 native model registration for multimodal generation."""
|
||||
|
||||
from sglang.multimodal_gen.runtime.models.sensenova_u1.neo_unify import ( # noqa: F401
|
||||
NEOChatConfig,
|
||||
NEOChatModel,
|
||||
NEOLLMConfig,
|
||||
NEOMoELLMConfig,
|
||||
NEOVisionConfig,
|
||||
NEOVisionModel,
|
||||
register,
|
||||
)
|
||||
|
||||
register()
|
||||
|
||||
__all__ = [
|
||||
"NEOChatConfig",
|
||||
"NEOChatModel",
|
||||
"NEOLLMConfig",
|
||||
"NEOMoELLMConfig",
|
||||
"NEOVisionConfig",
|
||||
"NEOVisionModel",
|
||||
"register",
|
||||
]
|
||||
@@ -0,0 +1,42 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from transformers import AutoModel, AutoTokenizer
|
||||
|
||||
from sglang.multimodal_gen.runtime.distributed import get_local_torch_device
|
||||
from sglang.multimodal_gen.runtime.models import ( # noqa: F401
|
||||
sensenova_u1 as _sensenova_u1,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.utils import PRECISION_TO_TYPE
|
||||
|
||||
|
||||
def load_model_and_tokenizer(
|
||||
model_path: str,
|
||||
server_args: ServerArgs,
|
||||
) -> dict[str, Any]:
|
||||
dtype = PRECISION_TO_TYPE.get(
|
||||
server_args.pipeline_config.model_precision, torch.bfloat16
|
||||
)
|
||||
model_kwargs: dict[str, Any] = {"torch_dtype": dtype}
|
||||
if server_args.trust_remote_code:
|
||||
model_kwargs["trust_remote_code"] = True
|
||||
if server_args.revision is not None:
|
||||
model_kwargs["revision"] = server_args.revision
|
||||
|
||||
tokenizer_kwargs: dict[str, Any] = {}
|
||||
if server_args.trust_remote_code:
|
||||
tokenizer_kwargs["trust_remote_code"] = True
|
||||
if server_args.revision is not None:
|
||||
tokenizer_kwargs["revision"] = server_args.revision
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(model_path, **tokenizer_kwargs)
|
||||
model = AutoModel.from_pretrained(model_path, **model_kwargs).eval()
|
||||
device = get_local_torch_device()
|
||||
current_platform.set_device(device)
|
||||
model = model.to(device)
|
||||
return {"model": model, "tokenizer": tokenizer}
|
||||
@@ -0,0 +1,8 @@
|
||||
# SenseNova-U1 Vendored Model Files
|
||||
|
||||
The Python files in this directory are adapted from:
|
||||
|
||||
- Repository: https://github.com/OpenSenseNova/SenseNova-U1
|
||||
- Commit: 2f42002f9b819506c9deb44599f0809e30252aba
|
||||
|
||||
They were modified for SGLang multimodal_gen native integration.
|
||||
@@ -0,0 +1,57 @@
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from .configuration_neo_chat import NEOChatConfig, NEOLLMConfig, NEOMoELLMConfig
|
||||
from .configuration_neo_vit import NEOVisionConfig
|
||||
from .modeling_neo_chat import NEOChatModel
|
||||
from .modeling_neo_vit import NEOVisionModel
|
||||
from .modeling_qwen3 import _HAS_FLASH_ATTN as has_flash_attn
|
||||
from .modeling_qwen3 import (
|
||||
Qwen3ForCausalLM,
|
||||
effective_attn_backend,
|
||||
get_attn_backend,
|
||||
set_attn_backend,
|
||||
)
|
||||
from .modeling_qwen3_moe import Qwen3MoeForCausalLM
|
||||
|
||||
__all__ = [
|
||||
"NEOChatConfig",
|
||||
"NEOLLMConfig",
|
||||
"NEOMoELLMConfig",
|
||||
"NEOVisionConfig",
|
||||
"NEOChatModel",
|
||||
"NEOVisionModel",
|
||||
"Qwen3ForCausalLM",
|
||||
"Qwen3MoeForCausalLM",
|
||||
"register",
|
||||
"set_attn_backend",
|
||||
"get_attn_backend",
|
||||
"effective_attn_backend",
|
||||
"has_flash_attn",
|
||||
]
|
||||
|
||||
|
||||
_REGISTERED = False
|
||||
|
||||
|
||||
def register() -> None:
|
||||
"""Register NEO-Unify types with ``transformers.Auto*``.
|
||||
|
||||
After calling this (or simply ``import sensenova_u1``), users can load a
|
||||
SenseNova-U1 checkpoint via plain ``AutoConfig.from_pretrained`` /
|
||||
``AutoModel.from_pretrained``.
|
||||
"""
|
||||
global _REGISTERED
|
||||
if _REGISTERED:
|
||||
return
|
||||
|
||||
from transformers import AutoConfig, AutoModel
|
||||
|
||||
AutoConfig.register("neo_vision", NEOVisionConfig, exist_ok=True)
|
||||
AutoConfig.register("neo_chat", NEOChatConfig, exist_ok=True)
|
||||
|
||||
AutoModel.register(NEOVisionConfig, NEOVisionModel, exist_ok=True)
|
||||
AutoModel.register(NEOChatConfig, NEOChatModel, exist_ok=True)
|
||||
|
||||
_REGISTERED = True
|
||||
+217
@@ -0,0 +1,217 @@
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
import copy
|
||||
|
||||
from transformers import Qwen3Config, Qwen3MoeConfig
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.utils import logging
|
||||
|
||||
from .configuration_neo_vit import NEOVisionConfig
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
def _restore_legacy_rope_theta(config) -> None:
|
||||
"""Expose the v4 rope attribute expected by the vendored model code."""
|
||||
if hasattr(config, "rope_theta"):
|
||||
return
|
||||
rope_parameters = getattr(config, "rope_parameters", None) or {}
|
||||
config.rope_theta = float(rope_parameters.get("rope_theta", 10000.0))
|
||||
|
||||
|
||||
class NEOLLMConfig(Qwen3Config):
|
||||
"""Config for the dense Qwen3 backbone used by NEO-Unify.
|
||||
|
||||
Extends ``Qwen3Config`` with two extra rope knobs used by the spatial
|
||||
(height/width) rotary axes that are layered on top of the temporal one.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, rope_theta_hw=10000.0, max_position_embeddings_hw=10000, **kwargs
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
_restore_legacy_rope_theta(self)
|
||||
self.rope_theta_hw = rope_theta_hw
|
||||
self.max_position_embeddings_hw = max_position_embeddings_hw
|
||||
|
||||
|
||||
class NEOMoELLMConfig(Qwen3MoeConfig):
|
||||
"""Config for the Qwen3-MoE backbone used by NEO-Unify.
|
||||
|
||||
Extends ``Qwen3MoeConfig`` with the same ``rope_theta_hw`` /
|
||||
``max_position_embeddings_hw`` extras as :class:`NEOLLMConfig`, and adds a
|
||||
*generation-path* MoE branch alongside the standard understanding-path one.
|
||||
In the A3B unified model every decoder layer carries two parallel sparse
|
||||
MoE blocks routed by the per-token ``image_gen_indicators`` mask:
|
||||
|
||||
* ``mlp`` - sparse MoE for the understanding path
|
||||
(``num_experts`` experts, ``num_experts_per_tok`` active,
|
||||
expert width ``moe_intermediate_size``).
|
||||
* ``mlp_mot_gen`` - sparse MoE for the image generation path
|
||||
(``gen_num_experts`` experts, ``gen_num_experts_per_tok``
|
||||
active, expert width ``gen_moe_intermediate_size``).
|
||||
|
||||
Each gen-path knob falls back to its understanding-path counterpart when
|
||||
unset, so vanilla single-MoE configs keep working without changes.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
rope_theta_hw=10000.0,
|
||||
max_position_embeddings_hw=10000,
|
||||
gen_num_experts=None,
|
||||
gen_num_experts_per_tok=None,
|
||||
gen_moe_intermediate_size=None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
_restore_legacy_rope_theta(self)
|
||||
self.rope_theta_hw = rope_theta_hw
|
||||
self.max_position_embeddings_hw = max_position_embeddings_hw
|
||||
|
||||
# Generation-path MoE knobs default to the understanding-path values
|
||||
# so legacy single-MoE configs (where both branches share the same
|
||||
# router width / expert count) keep working unchanged.
|
||||
self.gen_num_experts = (
|
||||
int(gen_num_experts)
|
||||
if gen_num_experts is not None
|
||||
else int(self.num_experts)
|
||||
)
|
||||
self.gen_num_experts_per_tok = (
|
||||
int(gen_num_experts_per_tok)
|
||||
if gen_num_experts_per_tok is not None
|
||||
else int(self.num_experts_per_tok)
|
||||
)
|
||||
self.gen_moe_intermediate_size = (
|
||||
int(gen_moe_intermediate_size)
|
||||
if gen_moe_intermediate_size is not None
|
||||
else int(self.moe_intermediate_size)
|
||||
)
|
||||
|
||||
# ``Qwen3Attention`` (used by NEO-Unify MoE layers) reads
|
||||
# ``config.layer_types[layer_idx]`` to decide between ``"full_attention"``
|
||||
# and ``"sliding_attention"``. Older / vanilla ``Qwen3MoeConfig`` does
|
||||
# not populate that field, so we backfill it here mirroring the dense
|
||||
# ``Qwen3Config`` behaviour: sliding-attention layers start at
|
||||
# ``max_window_layers`` when ``use_sliding_window`` is enabled.
|
||||
existing = getattr(self, "layer_types", None)
|
||||
if not existing or len(existing) != self.num_hidden_layers:
|
||||
use_swa = (
|
||||
bool(getattr(self, "use_sliding_window", False))
|
||||
and getattr(self, "sliding_window", None) is not None
|
||||
)
|
||||
max_window_layers = int(getattr(self, "max_window_layers", 0) or 0)
|
||||
self.layer_types = [
|
||||
(
|
||||
"sliding_attention"
|
||||
if (use_swa and i >= max_window_layers)
|
||||
else "full_attention"
|
||||
)
|
||||
for i in range(self.num_hidden_layers)
|
||||
]
|
||||
|
||||
|
||||
def _is_moe_llm_config(llm_config) -> bool:
|
||||
"""Detect whether an ``llm_config`` (dict or object) targets a MoE backbone.
|
||||
|
||||
Order of checks: explicit ``model_type``, ``architectures`` entry that
|
||||
contains ``MoE/MoeForCausalLM``, or presence of MoE-specific keys
|
||||
(``num_experts``).
|
||||
"""
|
||||
if isinstance(llm_config, dict):
|
||||
model_type = llm_config.get("model_type", "")
|
||||
archs = llm_config.get("architectures") or []
|
||||
has_num_experts = "num_experts" in llm_config
|
||||
else:
|
||||
model_type = getattr(llm_config, "model_type", "")
|
||||
archs = getattr(llm_config, "architectures", None) or []
|
||||
has_num_experts = hasattr(llm_config, "num_experts")
|
||||
|
||||
if isinstance(model_type, str) and "moe" in model_type.lower():
|
||||
return True
|
||||
for arch in archs:
|
||||
arch_str = str(arch)
|
||||
if "Moe" in arch_str or "MoE" in arch_str:
|
||||
return True
|
||||
return (
|
||||
bool(has_num_experts)
|
||||
and getattr(llm_config, "num_experts", 0)
|
||||
and int(getattr(llm_config, "num_experts", 0)) > 1
|
||||
)
|
||||
|
||||
|
||||
def _build_llm_config(llm_config):
|
||||
"""Instantiate the right LLM config object from a dict or pre-built config."""
|
||||
if isinstance(llm_config, dict):
|
||||
if _is_moe_llm_config(llm_config):
|
||||
return NEOMoELLMConfig(**llm_config)
|
||||
return NEOLLMConfig(**llm_config)
|
||||
return llm_config
|
||||
|
||||
|
||||
class NEOChatConfig(PretrainedConfig):
|
||||
model_type = "neo_chat"
|
||||
is_composition = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vision_config=None,
|
||||
llm_config=None,
|
||||
use_backbone_lora=0,
|
||||
use_llm_lora=0,
|
||||
downsample_ratio=0.5,
|
||||
template=None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
if vision_config is None:
|
||||
vision_config = {"architectures": ["NEOVisionModel"]}
|
||||
logger.info(
|
||||
"vision_config is None. Initializing the NEOVisionConfig with default values."
|
||||
)
|
||||
|
||||
if llm_config is None:
|
||||
llm_config = {"architectures": ["Qwen3ForCausalLM"]}
|
||||
logger.info(
|
||||
"llm_config is None. Initializing the LlamaConfig config with default values (`LlamaConfig`)."
|
||||
)
|
||||
assert "architectures" in llm_config, (
|
||||
"Should specify architecture in llm_config"
|
||||
)
|
||||
|
||||
if isinstance(vision_config, dict):
|
||||
self.vision_config = NEOVisionConfig(**vision_config)
|
||||
else:
|
||||
self.vision_config = vision_config
|
||||
|
||||
self.llm_config = _build_llm_config(llm_config)
|
||||
|
||||
self.use_backbone_lora = use_backbone_lora
|
||||
self.use_llm_lora = use_llm_lora
|
||||
self.downsample_ratio = downsample_ratio
|
||||
self.template = template
|
||||
self.tie_word_embeddings = self.llm_config.tie_word_embeddings
|
||||
|
||||
@property
|
||||
def is_moe_llm(self) -> bool:
|
||||
"""Convenience flag so callers can switch between dense / MoE LLM."""
|
||||
return isinstance(self.llm_config, NEOMoELLMConfig)
|
||||
|
||||
def to_dict(self):
|
||||
"""
|
||||
Serializes this instance to a Python dictionary. Override the default [`~PretrainedConfig.to_dict`].
|
||||
|
||||
Returns:
|
||||
`Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
|
||||
"""
|
||||
output = copy.deepcopy(self.__dict__)
|
||||
output["vision_config"] = self.vision_config.to_dict()
|
||||
output["llm_config"] = self.llm_config.to_dict()
|
||||
output["model_type"] = self.__class__.model_type
|
||||
output["use_backbone_lora"] = self.use_backbone_lora
|
||||
output["use_llm_lora"] = self.use_llm_lora
|
||||
output["downsample_ratio"] = self.downsample_ratio
|
||||
output["template"] = self.template
|
||||
|
||||
return output
|
||||
+70
@@ -0,0 +1,70 @@
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
import os
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, Union
|
||||
|
||||
from transformers.configuration_utils import PretrainedConfig
|
||||
from transformers.utils import logging
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
|
||||
def _as_singleton_tuple(value: Any) -> tuple[Any, ...]:
|
||||
if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
|
||||
if len(value) == 1:
|
||||
return _as_singleton_tuple(value[0])
|
||||
return tuple(value)
|
||||
return (value,)
|
||||
|
||||
|
||||
class NEOVisionConfig(PretrainedConfig):
|
||||
model_type = "neo_vision"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_channels=3,
|
||||
patch_size=16,
|
||||
hidden_size=1024,
|
||||
llm_hidden_size=2048,
|
||||
downsample_ratio=0.5,
|
||||
rope_theta_vision=10000.0,
|
||||
max_position_embeddings_vision=10000,
|
||||
min_pixels=65536,
|
||||
max_pixels=4194304,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(**kwargs)
|
||||
|
||||
self.hidden_size = hidden_size
|
||||
self.llm_hidden_size = _as_singleton_tuple(llm_hidden_size)
|
||||
self.downsample_ratio = _as_singleton_tuple(downsample_ratio)
|
||||
self.rope_theta_vision = rope_theta_vision
|
||||
self.max_position_embeddings_vision = max_position_embeddings_vision
|
||||
self.num_channels = num_channels
|
||||
self.patch_size = patch_size
|
||||
self.min_pixels = min_pixels
|
||||
self.max_pixels = max_pixels
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls, pretrained_model_name_or_path: Union[str, os.PathLike], **kwargs
|
||||
) -> "PretrainedConfig":
|
||||
config_dict, kwargs = cls.get_config_dict(
|
||||
pretrained_model_name_or_path, **kwargs
|
||||
)
|
||||
|
||||
if "vision_config" in config_dict:
|
||||
config_dict = config_dict["vision_config"]
|
||||
|
||||
if (
|
||||
"model_type" in config_dict
|
||||
and hasattr(cls, "model_type")
|
||||
and config_dict["model_type"] != cls.model_type
|
||||
):
|
||||
logger.warning(
|
||||
f"You are using a model of type {config_dict['model_type']} to instantiate a model of type "
|
||||
f"{cls.model_type}. This is not supported for all configurations of models and can yield errors."
|
||||
)
|
||||
|
||||
return cls.from_dict(config_dict, **kwargs)
|
||||
@@ -0,0 +1,420 @@
|
||||
"""
|
||||
Conversation prompt templates.
|
||||
|
||||
We kindly request that you import fastchat instead of copying this file if you wish to use it.
|
||||
If you have changes in mind, please contribute back so the community can benefit collectively and continue to maintain these valuable templates.
|
||||
|
||||
Modified from https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py
|
||||
"""
|
||||
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
import dataclasses
|
||||
from enum import IntEnum, auto
|
||||
from typing import Dict, List, Tuple, Union
|
||||
|
||||
|
||||
class SeparatorStyle(IntEnum):
|
||||
"""Separator styles."""
|
||||
|
||||
ADD_COLON_SINGLE = auto()
|
||||
ADD_COLON_TWO = auto()
|
||||
ADD_COLON_SPACE_SINGLE = auto()
|
||||
NO_COLON_SINGLE = auto()
|
||||
NO_COLON_TWO = auto()
|
||||
ADD_NEW_LINE_SINGLE = auto()
|
||||
LLAMA2 = auto()
|
||||
CHATGLM = auto()
|
||||
CHATML = auto()
|
||||
CHATINTERN = auto()
|
||||
DOLLY = auto()
|
||||
RWKV = auto()
|
||||
PHOENIX = auto()
|
||||
ROBIN = auto()
|
||||
FALCON_CHAT = auto()
|
||||
CHATGLM3 = auto()
|
||||
INTERNVL_ZH = auto()
|
||||
MPT = auto()
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
class Conversation:
|
||||
"""A class that manages prompt templates and keeps all conversation history."""
|
||||
|
||||
# The name of this template
|
||||
name: str
|
||||
# The template of the system prompt
|
||||
system_template: str = "{system_message}"
|
||||
# The system message
|
||||
system_message: str = ""
|
||||
# The names of two roles
|
||||
roles: Tuple[str] = ("USER", "ASSISTANT")
|
||||
# All messages. Each item is (role, message).
|
||||
messages: List[List[str]] = ()
|
||||
# The number of few shot examples
|
||||
offset: int = 0
|
||||
# The separator style and configurations
|
||||
sep_style: SeparatorStyle = SeparatorStyle.ADD_COLON_SINGLE
|
||||
sep: str = "\n"
|
||||
sep2: str = None
|
||||
# Stop criteria (the default one is EOS token)
|
||||
stop_str: Union[str, List[str]] = None
|
||||
# Stops generation if meeting any token in this list
|
||||
stop_token_ids: List[int] = None
|
||||
|
||||
def get_prompt(self) -> str:
|
||||
"""Get the prompt for generation."""
|
||||
if self.system_message is not None and self.system_message != "":
|
||||
system_prompt = self.system_template.format(
|
||||
system_message=self.system_message
|
||||
)
|
||||
else:
|
||||
system_prompt = ""
|
||||
|
||||
if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + ": " + message + self.sep
|
||||
else:
|
||||
ret += role + ":"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.ADD_COLON_TWO:
|
||||
seps = [self.sep, self.sep2]
|
||||
ret = "" if system_prompt == "" else system_prompt + seps[0]
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
if message:
|
||||
ret += role + ": " + message + seps[i % 2]
|
||||
else:
|
||||
ret += role + ":"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.ADD_COLON_SPACE_SINGLE:
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + ": " + message + self.sep
|
||||
else:
|
||||
ret += role + ": " # must be end with a space
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.ADD_NEW_LINE_SINGLE:
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + "\n" + message + self.sep
|
||||
else:
|
||||
ret += role + "\n"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
|
||||
ret = system_prompt
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + message + self.sep
|
||||
else:
|
||||
ret += role
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.NO_COLON_TWO:
|
||||
seps = [self.sep, self.sep2]
|
||||
ret = system_prompt
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
if message:
|
||||
ret += role + message + seps[i % 2]
|
||||
else:
|
||||
ret += role
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.RWKV:
|
||||
ret = system_prompt
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
if message:
|
||||
ret += (
|
||||
role
|
||||
+ ": "
|
||||
+ message.replace("\r\n", "\n").replace("\n\n", "\n")
|
||||
)
|
||||
ret += "\n\n"
|
||||
else:
|
||||
ret += role + ":"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.LLAMA2:
|
||||
seps = [self.sep, self.sep2]
|
||||
ret = system_prompt if system_prompt != "" else "[INST] "
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
tag = self.roles[i % 2]
|
||||
if message:
|
||||
if i == 0:
|
||||
ret += message + " "
|
||||
else:
|
||||
ret += tag + " " + message + seps[i % 2]
|
||||
else:
|
||||
ret += tag
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.CHATGLM:
|
||||
# source: https://huggingface.co/THUDM/chatglm-6b/blob/1d240ba371910e9282298d4592532d7f0f3e9f3e/modeling_chatglm.py#L1302-L1308
|
||||
# source2: https://huggingface.co/THUDM/chatglm2-6b/blob/e186c891cf64310ac66ef10a87e6635fa6c2a579/modeling_chatglm.py#L926
|
||||
round_add_n = 1 if self.name == "chatglm2" else 0
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep
|
||||
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
if i % 2 == 0:
|
||||
ret += f"[Round {i // 2 + round_add_n}]{self.sep}"
|
||||
|
||||
if message:
|
||||
ret += f"{role}：{message}{self.sep}"
|
||||
else:
|
||||
ret += f"{role}："
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.CHATML:
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep + "\n"
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + "\n" + message + self.sep + "\n"
|
||||
else:
|
||||
ret += role + "\n"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.CHATGLM3:
|
||||
ret = system_prompt
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + "\n" + " " + message
|
||||
else:
|
||||
ret += role
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.CHATINTERN:
|
||||
# source: https://huggingface.co/internlm/internlm-chat-7b-8k/blob/bd546fa984b4b0b86958f56bf37f94aa75ab8831/modeling_internlm.py#L771
|
||||
seps = [self.sep, self.sep2]
|
||||
ret = system_prompt
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
# if i % 2 == 0:
|
||||
# ret += "<s>"
|
||||
if message:
|
||||
ret += role + ":" + message + seps[i % 2] + "\n"
|
||||
else:
|
||||
ret += role + ":"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.DOLLY:
|
||||
seps = [self.sep, self.sep2]
|
||||
ret = system_prompt
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
if message:
|
||||
ret += role + ":\n" + message + seps[i % 2]
|
||||
if i % 2 == 1:
|
||||
ret += "\n\n"
|
||||
else:
|
||||
ret += role + ":\n"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.PHOENIX:
|
||||
ret = system_prompt
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + ": " + "<s>" + message + "</s>"
|
||||
else:
|
||||
ret += role + ": " + "<s>"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.ROBIN:
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + ":\n" + message + self.sep
|
||||
else:
|
||||
ret += role + ":\n"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.FALCON_CHAT:
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep
|
||||
for role, message in self.messages:
|
||||
if message:
|
||||
ret += role + ": " + message + self.sep
|
||||
else:
|
||||
ret += role + ":"
|
||||
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.INTERNVL_ZH:
|
||||
seps = [self.sep, self.sep2]
|
||||
ret = "" if system_prompt == "" else self.system_message + seps[0]
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
if message:
|
||||
ret += role + ": " + message + seps[i % 2]
|
||||
else:
|
||||
ret += role + ":"
|
||||
return ret
|
||||
elif self.sep_style == SeparatorStyle.MPT:
|
||||
ret = "" if system_prompt == "" else system_prompt + self.sep
|
||||
for i, (role, message) in enumerate(self.messages):
|
||||
if message:
|
||||
if type(message) is tuple:
|
||||
message, _, _ = message
|
||||
ret += role + message + self.sep
|
||||
else:
|
||||
if i != len(self.messages) and message is not None:
|
||||
ret += role + self.sep
|
||||
else:
|
||||
ret += role
|
||||
return ret
|
||||
else:
|
||||
raise ValueError(f"Invalid style: {self.sep_style}")
|
||||
|
||||
def set_system_message(self, system_message: str):
|
||||
"""Set the system message."""
|
||||
self.system_message = system_message
|
||||
|
||||
def append_message(self, role: str, message: str):
|
||||
"""Append a new message."""
|
||||
self.messages.append([role, message])
|
||||
|
||||
def update_last_message(self, message: str):
|
||||
"""Update the last output.
|
||||
|
||||
The last message is typically set to be None when constructing the prompt,
|
||||
so we need to update it in-place after getting the response from a model.
|
||||
"""
|
||||
self.messages[-1][1] = message
|
||||
|
||||
def to_gradio_chatbot(self):
|
||||
"""Convert the conversation to gradio chatbot format."""
|
||||
ret = []
|
||||
for i, (role, msg) in enumerate(self.messages[self.offset :]):
|
||||
if i % 2 == 0:
|
||||
ret.append([msg, None])
|
||||
else:
|
||||
ret[-1][-1] = msg
|
||||
return ret
|
||||
|
||||
def to_openai_api_messages(self):
|
||||
"""Convert the conversation to OpenAI chat completion format."""
|
||||
ret = [{"role": "system", "content": self.system_message}]
|
||||
|
||||
for i, (_, msg) in enumerate(self.messages[self.offset :]):
|
||||
if i % 2 == 0:
|
||||
ret.append({"role": "user", "content": msg})
|
||||
else:
|
||||
if msg is not None:
|
||||
ret.append({"role": "assistant", "content": msg})
|
||||
return ret
|
||||
|
||||
def copy(self):
|
||||
return Conversation(
|
||||
name=self.name,
|
||||
system_template=self.system_template,
|
||||
system_message=self.system_message,
|
||||
roles=self.roles,
|
||||
messages=[[x, y] for x, y in self.messages],
|
||||
offset=self.offset,
|
||||
sep_style=self.sep_style,
|
||||
sep=self.sep,
|
||||
sep2=self.sep2,
|
||||
stop_str=self.stop_str,
|
||||
stop_token_ids=self.stop_token_ids,
|
||||
)
|
||||
|
||||
def dict(self):
|
||||
return {
|
||||
"template_name": self.name,
|
||||
"system_message": self.system_message,
|
||||
"roles": self.roles,
|
||||
"messages": self.messages,
|
||||
"offset": self.offset,
|
||||
}
|
||||
|
||||
|
||||
# A global registry for all conversation templates
|
||||
conv_templates: Dict[str, Conversation] = {}
|
||||
|
||||
|
||||
def register_conv_template(template: Conversation, override: bool = False):
|
||||
"""Register a new conversation template."""
|
||||
if not override:
|
||||
assert template.name not in conv_templates, (
|
||||
f"{template.name} has been registered."
|
||||
)
|
||||
|
||||
conv_templates[template.name] = template
|
||||
|
||||
|
||||
def get_conv_template(name: str) -> Conversation:
|
||||
"""Get a conversation template."""
|
||||
return conv_templates[name].copy()
|
||||
|
||||
|
||||
# Both Hermes-2 and neo1_0-chat are chatml-format conversation templates. The difference
|
||||
# is that during training, the preprocessing function for the Hermes-2 template doesn't add
|
||||
# <s> at the beginning of the tokenized sequence, while the neo1_0-chat template does.
|
||||
# Therefore, they are completely equivalent during inference.
|
||||
|
||||
# These Unicode-escaped strings preserve the exact Chinese prompts from the
|
||||
# pinned upstream source while complying with the multimodal_gen source lint.
|
||||
# Upstream prompt retained for benchmark compatibility:
|
||||
_INTERNVL_SYSTEM_MESSAGE = (
|
||||
"\u4f60\u662f\u7531\u4e0a\u6d77\u4eba\u5de5\u667a\u80fd\u5b9e\u9a8c\u5ba4"
|
||||
"\u8054\u5408\u5546\u6c64\u79d1\u6280\u5f00\u53d1\u7684\u4e66\u751f\u591a"
|
||||
"\u6a21\u6001\u5927\u6a21\u578b\uff0c\u82f1\u6587\u540d\u53ebInternVL, "
|
||||
"\u662f\u4e00\u4e2a\u6709\u7528\u65e0\u5bb3\u7684\u4eba\u5de5\u667a\u80fd"
|
||||
"\u52a9\u624b\u3002"
|
||||
)
|
||||
# Newer upstream prompt, intentionally not used by the first three templates:
|
||||
_INTERNVL2_5_SYSTEM_MESSAGE = (
|
||||
"\u4f60\u662f\u4e66\u751f\xb7\u4e07\u8c61\uff0c\u82f1\u6587\u540d\u662f"
|
||||
"InternVL\uff0c\u662f\u7531\u4e0a\u6d77\u4eba\u5de5\u667a\u80fd\u5b9e\u9a8c"
|
||||
"\u5ba4\u3001\u6e05\u534e\u5927\u5b66\u53ca\u591a\u5bb6\u5408\u4f5c\u5355"
|
||||
"\u4f4d\u8054\u5408\u5f00\u53d1\u7684\u591a\u6a21\u6001\u5927\u8bed\u8a00"
|
||||
"\u6a21\u578b\u3002"
|
||||
)
|
||||
|
||||
register_conv_template(
|
||||
Conversation(
|
||||
name="Hermes-2",
|
||||
system_template="<|im_start|>system\n{system_message}",
|
||||
# note: The new system prompt was not used here to avoid changes in benchmark performance.
|
||||
# system_message=_INTERNVL2_5_SYSTEM_MESSAGE
|
||||
system_message=_INTERNVL_SYSTEM_MESSAGE,
|
||||
roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
|
||||
sep_style=SeparatorStyle.MPT,
|
||||
sep="<|im_end|>",
|
||||
stop_str="<|endoftext|>",
|
||||
)
|
||||
)
|
||||
|
||||
register_conv_template(
|
||||
Conversation(
|
||||
name="internlm2-chat",
|
||||
system_template="<|im_start|>system\n{system_message}",
|
||||
# note: The new system prompt was not used here to avoid changes in benchmark performance.
|
||||
# system_message=_INTERNVL2_5_SYSTEM_MESSAGE
|
||||
system_message=_INTERNVL_SYSTEM_MESSAGE,
|
||||
roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
|
||||
sep_style=SeparatorStyle.MPT,
|
||||
sep="<|im_end|>",
|
||||
)
|
||||
)
|
||||
|
||||
register_conv_template(
|
||||
Conversation(
|
||||
name="phi3-chat",
|
||||
system_template="<|system|>\n{system_message}",
|
||||
# note: The new system prompt was not used here to avoid changes in benchmark performance.
|
||||
# system_message=_INTERNVL2_5_SYSTEM_MESSAGE
|
||||
system_message=_INTERNVL_SYSTEM_MESSAGE,
|
||||
roles=("<|user|>\n", "<|assistant|>\n"),
|
||||
sep_style=SeparatorStyle.MPT,
|
||||
sep="<|end|>",
|
||||
)
|
||||
)
|
||||
|
||||
register_conv_template(
|
||||
Conversation(
|
||||
name="internvl2_5",
|
||||
system_template="<|im_start|>system\n{system_message}",
|
||||
system_message=_INTERNVL2_5_SYSTEM_MESSAGE,
|
||||
roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
|
||||
sep_style=SeparatorStyle.MPT,
|
||||
sep="<|im_end|>\n",
|
||||
)
|
||||
)
|
||||
|
||||
register_conv_template(
|
||||
Conversation(
|
||||
name="neo1_0",
|
||||
system_template="<|im_start|>system\n{system_message}",
|
||||
system_message="",
|
||||
roles=("<|im_start|>user\n", "<|im_start|>assistant\n"),
|
||||
sep_style=SeparatorStyle.MPT,
|
||||
sep="<|im_end|>\n",
|
||||
)
|
||||
)
|
||||
+648
@@ -0,0 +1,648 @@
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
import logging
|
||||
import math
|
||||
from functools import lru_cache
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def modulate(x, shift, scale=None):
|
||||
if shift is None:
|
||||
return x * (1 + scale)
|
||||
return x * (1 + scale) + shift
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-5):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
output = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
|
||||
return output * self.weight
|
||||
|
||||
|
||||
class TimestepEmbedder(nn.Module):
|
||||
"""
|
||||
Embeds scalar timesteps into vector representations.
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size, frequency_embedding_size=256):
|
||||
super().__init__()
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_size, hidden_size, bias=True),
|
||||
)
|
||||
self.frequency_embedding_size = frequency_embedding_size
|
||||
|
||||
@staticmethod
|
||||
def timestep_embedding(t: torch.Tensor, dim: int, max_period: float = 10000.0):
|
||||
"""
|
||||
Create sinusoidal timestep embeddings.
|
||||
:param t: a 1-D Tensor of N indices, one per batch element. These may be fractional.
|
||||
:param dim: the dimension of the output.
|
||||
:param max_period: controls the minimum frequency of the embeddings.
|
||||
:return: an (N, D) Tensor of positional embeddings.
|
||||
"""
|
||||
# https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
|
||||
half = dim // 2
|
||||
freqs = torch.exp(
|
||||
-math.log(max_period)
|
||||
* torch.arange(start=0, end=half, dtype=torch.float32)
|
||||
/ half
|
||||
).to(device=t.device)
|
||||
args = t[:, None].float() * freqs[None]
|
||||
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
if dim % 2:
|
||||
embedding = torch.cat(
|
||||
[embedding, torch.zeros_like(embedding[:, :1])], dim=-1
|
||||
)
|
||||
return embedding
|
||||
|
||||
def forward(self, t):
|
||||
t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
|
||||
t_emb = self.mlp(t_freq.to(self.mlp[0].weight.dtype))
|
||||
return t_emb
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
def __init__(self, channels, mlp_ratio=1.0):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.intermediate_size = int(channels * mlp_ratio)
|
||||
|
||||
self.in_ln = nn.LayerNorm(self.channels, eps=1e-6)
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(self.channels, self.intermediate_size),
|
||||
nn.SiLU(),
|
||||
nn.Linear(self.intermediate_size, self.channels),
|
||||
)
|
||||
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(), nn.Linear(channels, 3 * channels, bias=True)
|
||||
)
|
||||
|
||||
def forward(self, x, y):
|
||||
shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(y).chunk(3, dim=-1)
|
||||
h = modulate(self.in_ln(x), shift_mlp, scale_mlp)
|
||||
h = self.mlp(h)
|
||||
return x + gate_mlp * h
|
||||
|
||||
|
||||
# class FinalLayer(nn.Module):
|
||||
|
||||
# def __init__(self, model_channels, out_channels):
|
||||
# super().__init__()
|
||||
# self.norm_final = nn.LayerNorm(model_channels, elementwise_affine=False, eps=1e-6)
|
||||
# self.linear = nn.Linear(model_channels, out_channels, bias=True)
|
||||
# self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(model_channels, 2 * model_channels, bias=True))
|
||||
|
||||
# def forward(self, x, c):
|
||||
# shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
|
||||
# x = modulate(self.norm_final(x), shift, scale)
|
||||
# x = self.linear(x)
|
||||
# return x
|
||||
|
||||
# class SimpleMLPAdaLN(nn.Module):
|
||||
|
||||
# def __init__(self, input_dim, out_dim, dim=1536, layers=12, mlp_ratio=1.0):
|
||||
# super().__init__()
|
||||
# self.input_dim = input_dim
|
||||
# self.out_dim = out_dim
|
||||
# self.dim = dim
|
||||
# self.layers = layers
|
||||
# self.mlp_ratio = mlp_ratio
|
||||
|
||||
# self.time_embed = TimestepEmbedder(dim)
|
||||
# self.input_proj = nn.Linear(input_dim, dim)
|
||||
|
||||
# res_blocks = []
|
||||
# for _ in range(layers):
|
||||
# res_blocks.append(ResBlock(dim, mlp_ratio))
|
||||
# self.res_blocks = nn.ModuleList(res_blocks)
|
||||
|
||||
# self.final_layer = FinalLayer(dim, out_dim)
|
||||
|
||||
# self.grad_checkpointing = False
|
||||
|
||||
# self.initialize_weights()
|
||||
|
||||
# def initialize_weights(self):
|
||||
# def _basic_init(module):
|
||||
# if isinstance(module, nn.Linear):
|
||||
# torch.nn.init.xavier_uniform_(module.weight)
|
||||
# if module.bias is not None:
|
||||
# nn.init.constant_(module.bias, 0)
|
||||
|
||||
# self.apply(_basic_init)
|
||||
|
||||
# # Initialize timestep embedding MLP
|
||||
# nn.init.normal_(self.time_embed.mlp[0].weight, std=0.02)
|
||||
# nn.init.normal_(self.time_embed.mlp[2].weight, std=0.02)
|
||||
|
||||
# # Zero-out adaLN modulation layers
|
||||
# for block in self.res_blocks:
|
||||
# nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
|
||||
# nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
|
||||
|
||||
# # Zero-out output layers
|
||||
# nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
|
||||
# nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
|
||||
# nn.init.constant_(self.final_layer.linear.weight, 0)
|
||||
# nn.init.constant_(self.final_layer.linear.bias, 0)
|
||||
|
||||
# def forward(self, x, t):
|
||||
# """
|
||||
# x.shape = (bsz, input_dim)
|
||||
# t.shape = (bsz,)
|
||||
# """
|
||||
|
||||
# x = self.input_proj(x)
|
||||
# t = self.time_embed(t)
|
||||
|
||||
# y = t
|
||||
|
||||
# for block in self.res_blocks:
|
||||
# if self.grad_checkpointing and self.training:
|
||||
# x = checkpoint(block, x, y, use_reentrant=True)
|
||||
# else:
|
||||
# x = block(x, y)
|
||||
|
||||
# return self.final_layer(x, y)
|
||||
|
||||
|
||||
class FlowMatchingHead(nn.Module):
|
||||
def __init__(self, input_dim, out_dim, dim=1536, layers=12, mlp_ratio=1.0):
|
||||
super(FlowMatchingHead, self).__init__()
|
||||
self.net = SimpleMLPAdaLN(
|
||||
input_dim=input_dim,
|
||||
out_dim=out_dim,
|
||||
dim=dim,
|
||||
layers=layers,
|
||||
mlp_ratio=mlp_ratio,
|
||||
)
|
||||
|
||||
@property
|
||||
def dtype(self):
|
||||
return self.net.input_proj.weight.dtype
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
return self.net.input_proj.weight.device
|
||||
|
||||
def forward(self, x, t):
|
||||
x = self.net(x, t)
|
||||
return x
|
||||
|
||||
|
||||
def precompute_freqs_cis_2d(
|
||||
dim: int, height: int, width: int, theta: float = 10000.0, scale=16.0
|
||||
):
|
||||
# assert H * H == end
|
||||
# flat_patch_pos = torch.linspace(-1, 1, end) # N = end
|
||||
x_pos = torch.linspace(0, scale, width)
|
||||
y_pos = torch.linspace(0, scale, height)
|
||||
y_pos, x_pos = torch.meshgrid(y_pos, x_pos, indexing="ij")
|
||||
y_pos = y_pos.reshape(-1)
|
||||
x_pos = x_pos.reshape(-1)
|
||||
freqs = 1.0 / (
|
||||
theta ** (torch.arange(0, dim, 4)[: (dim // 4)].float() / dim)
|
||||
) # Hc/4
|
||||
x_freqs = torch.outer(x_pos, freqs).float() # N Hc/4
|
||||
y_freqs = torch.outer(y_pos, freqs).float() # N Hc/4
|
||||
x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs)
|
||||
y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs)
|
||||
freqs_cis = torch.cat(
|
||||
[x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1
|
||||
) # N,Hc/4,2
|
||||
freqs_cis = freqs_cis.reshape(height * width, -1)
|
||||
return freqs_cis
|
||||
|
||||
|
||||
class NerfEmbedder(nn.Module):
|
||||
def __init__(self, in_channels, hidden_size_input, max_freqs):
|
||||
super().__init__()
|
||||
self.max_freqs = max_freqs
|
||||
self.hidden_size_input = hidden_size_input
|
||||
self.embedder = nn.Sequential(
|
||||
nn.Linear(in_channels + max_freqs**2, hidden_size_input, bias=True),
|
||||
)
|
||||
|
||||
@lru_cache
|
||||
def fetch_pos(self, patch_size, device, dtype):
|
||||
pos = precompute_freqs_cis_2d(
|
||||
self.max_freqs**2 * 2, patch_size, patch_size
|
||||
).real
|
||||
pos = pos[None, :, :].to(device=device, dtype=dtype)
|
||||
return pos
|
||||
|
||||
def forward(self, inputs):
|
||||
B, P2, C = inputs.shape
|
||||
patch_size = int(P2**0.5)
|
||||
device = inputs.device
|
||||
dtype = inputs.dtype
|
||||
dct = self.fetch_pos(patch_size, device, dtype)
|
||||
dct = dct.repeat(B, 1, 1)
|
||||
inputs = torch.cat([inputs, dct], dim=-1)
|
||||
inputs = self.embedder(inputs)
|
||||
return inputs
|
||||
|
||||
|
||||
class SimpleMLPAdaLN(nn.Module):
|
||||
"""
|
||||
The MLP for Diffusion Loss.
|
||||
:param in_channels: channels in the input Tensor.
|
||||
:param model_channels: base channel count for the model.
|
||||
:param out_channels: channels in the output Tensor.
|
||||
:param z_channels: channels in the condition.
|
||||
:param num_res_blocks: number of residual blocks per downsample.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
model_channels,
|
||||
out_channels,
|
||||
z_channels,
|
||||
num_res_blocks,
|
||||
patch_size,
|
||||
grad_checkpointing=False,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.model_channels = model_channels
|
||||
self.out_channels = out_channels
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.grad_checkpointing = grad_checkpointing
|
||||
self.patch_size = patch_size
|
||||
|
||||
self.cond_embed = nn.Linear(z_channels, patch_size**2 * model_channels)
|
||||
|
||||
self.input_proj = nn.Linear(in_channels, model_channels)
|
||||
|
||||
res_blocks = []
|
||||
for i in range(num_res_blocks):
|
||||
res_blocks.append(
|
||||
ResBlock(
|
||||
model_channels,
|
||||
)
|
||||
)
|
||||
|
||||
self.res_blocks = nn.ModuleList(res_blocks)
|
||||
self.final_layer = FinalLayer(model_channels, out_channels)
|
||||
|
||||
self.initialize_weights()
|
||||
|
||||
def initialize_weights(self):
|
||||
def _basic_init(module):
|
||||
if isinstance(module, nn.Linear):
|
||||
torch.nn.init.xavier_uniform_(module.weight)
|
||||
if module.bias is not None:
|
||||
nn.init.constant_(module.bias, 0)
|
||||
|
||||
self.apply(_basic_init)
|
||||
|
||||
# Zero-out adaLN modulation layers
|
||||
for block in self.res_blocks:
|
||||
nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
|
||||
nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
|
||||
|
||||
# Zero-out output layers
|
||||
nn.init.constant_(self.final_layer.linear.weight, 0)
|
||||
nn.init.constant_(self.final_layer.linear.bias, 0)
|
||||
|
||||
def forward(self, x, c):
|
||||
"""
|
||||
Apply the model to an input batch.
|
||||
:param x: an [N x C] Tensor of inputs.
|
||||
:param t: a 1-D batch of timesteps.
|
||||
:param c: conditioning from AR transformer.
|
||||
:return: an [N x C] Tensor of outputs.
|
||||
"""
|
||||
x = self.input_proj(x)
|
||||
c = self.cond_embed(c)
|
||||
|
||||
y = c.reshape(-1, self.patch_size**2, self.model_channels)
|
||||
|
||||
for block in self.res_blocks:
|
||||
x = block(x, y)
|
||||
|
||||
return self.final_layer(x)
|
||||
|
||||
|
||||
class FinalLayer(nn.Module):
|
||||
"""
|
||||
The final layer adopted from DiT.
|
||||
"""
|
||||
|
||||
def __init__(self, model_channels, out_channels):
|
||||
super().__init__()
|
||||
self.norm_final = nn.LayerNorm(
|
||||
model_channels, elementwise_affine=False, eps=1e-6
|
||||
)
|
||||
self.linear = nn.Linear(model_channels, out_channels, bias=True)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.norm_final(x)
|
||||
x = self.linear(x)
|
||||
return x
|
||||
|
||||
|
||||
#################################################################################
|
||||
# Sine/Cosine Positional Embedding Functions #
|
||||
#################################################################################
|
||||
# https://github.com/facebookresearch/mae/blob/main/util/pos_embed.py
|
||||
|
||||
|
||||
def get_2d_sincos_pos_embed(
|
||||
embed_dim, grid_size, cls_token=False, extra_tokens=0, pe_interpolation=1.0
|
||||
):
|
||||
"""
|
||||
grid_size: int of the grid height and width
|
||||
return:
|
||||
pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
|
||||
"""
|
||||
grid_h = np.arange(grid_size, dtype=np.float32) / pe_interpolation
|
||||
grid_w = np.arange(grid_size, dtype=np.float32) / pe_interpolation
|
||||
grid = np.meshgrid(grid_w, grid_h) # here w goes first
|
||||
grid = np.stack(grid, axis=0)
|
||||
|
||||
grid = grid.reshape([2, 1, grid_size, grid_size])
|
||||
pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
|
||||
if cls_token and extra_tokens > 0:
|
||||
pos_embed = np.concatenate(
|
||||
[np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0
|
||||
)
|
||||
return pos_embed
|
||||
|
||||
|
||||
def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
|
||||
assert embed_dim % 2 == 0
|
||||
|
||||
# use half of dimensions to encode grid_h
|
||||
emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)
|
||||
emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)
|
||||
|
||||
emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
|
||||
return emb
|
||||
|
||||
|
||||
def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
|
||||
"""
|
||||
embed_dim: output dimension for each position
|
||||
pos: a list of positions to be encoded: size (M,)
|
||||
out: (M, D)
|
||||
"""
|
||||
assert embed_dim % 2 == 0
|
||||
omega = np.arange(embed_dim // 2, dtype=np.float64)
|
||||
omega /= embed_dim / 2.0
|
||||
omega = 1.0 / 10000**omega # (D/2,)
|
||||
|
||||
pos = pos.reshape(-1) # (M,)
|
||||
out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product
|
||||
|
||||
emb_sin = np.sin(out) # (M, D/2)
|
||||
emb_cos = np.cos(out) # (M, D/2)
|
||||
|
||||
emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)
|
||||
return emb
|
||||
|
||||
|
||||
# --------------------------------------------------------
|
||||
# Interpolate position embeddings for high-resolution
|
||||
# References:
|
||||
# DeiT: https://github.com/facebookresearch/deit
|
||||
# --------------------------------------------------------
|
||||
def interpolate_pos_embed(
|
||||
model_path, pe_key: str = "gen_pos_embed", new_len: int = 4096
|
||||
):
|
||||
state_dict = torch.load(model_path, map_location="cpu")
|
||||
|
||||
pos_embed_1d = state_dict[pe_key]
|
||||
_, ori_len, embed_dim = pos_embed_1d.shape
|
||||
|
||||
ori_size = int(ori_len**0.5)
|
||||
new_size = int(new_len**0.5)
|
||||
|
||||
if ori_size != new_size:
|
||||
logger.info(
|
||||
"Position interpolate from %dx%d to %dx%d"
|
||||
% (ori_size, ori_size, new_size, new_size)
|
||||
)
|
||||
pos_embed_2d = pos_embed_1d.reshape(-1, ori_size, ori_size, embed_dim).permute(
|
||||
0, 3, 1, 2
|
||||
)
|
||||
pos_embed_2d = torch.nn.functional.interpolate(
|
||||
pos_embed_2d, size=(new_size, new_size), mode="bicubic", align_corners=False
|
||||
)
|
||||
pos_embed_1d = pos_embed_2d.permute(0, 2, 3, 1).flatten(1, 2)
|
||||
state_dict[pe_key] = pos_embed_1d
|
||||
|
||||
torch.save(state_dict, model_path)
|
||||
|
||||
|
||||
class PositionEmbedding(nn.Module):
|
||||
def __init__(self, max_num_patch_per_side, hidden_size):
|
||||
super().__init__()
|
||||
self.max_num_patch_per_side = max_num_patch_per_side
|
||||
self.hidden_size = hidden_size
|
||||
self.pos_embed = nn.Parameter(
|
||||
torch.zeros(max_num_patch_per_side**2, hidden_size), requires_grad=False
|
||||
)
|
||||
self._init_weights()
|
||||
|
||||
def _init_weights(self):
|
||||
# Initialize (and freeze) pos_embed by sin-cos embedding:
|
||||
pos_embed = get_2d_sincos_pos_embed(
|
||||
self.hidden_size, self.max_num_patch_per_side
|
||||
)
|
||||
self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float())
|
||||
|
||||
def forward(self, position_ids):
|
||||
return self.pos_embed[position_ids]
|
||||
|
||||
|
||||
class ResidualConvBlock(nn.Module):
|
||||
def __init__(self, channels: int):
|
||||
super().__init__()
|
||||
self.block = nn.Sequential(
|
||||
nn.Conv2d(channels, channels, kernel_size=3, padding=1),
|
||||
nn.SiLU(),
|
||||
nn.Conv2d(channels, channels, kernel_size=3, padding=1),
|
||||
)
|
||||
nn.init.zeros_(self.block[2].weight)
|
||||
nn.init.zeros_(self.block[2].bias)
|
||||
|
||||
def forward(self, x):
|
||||
return x + self.block(x)
|
||||
|
||||
|
||||
class PostConvSmoother(nn.Module):
|
||||
def __init__(self, in_channels=3, hidden_channels=64, num_blocks=3):
|
||||
super().__init__()
|
||||
self.in_proj = nn.Conv2d(in_channels, hidden_channels, kernel_size=3, padding=1)
|
||||
self.blocks = nn.Sequential(
|
||||
*[ResidualConvBlock(hidden_channels) for _ in range(num_blocks)]
|
||||
)
|
||||
self.out_proj = nn.Conv2d(hidden_channels, in_channels, kernel_size=1)
|
||||
|
||||
nn.init.zeros_(self.out_proj.weight)
|
||||
nn.init.zeros_(self.out_proj.bias)
|
||||
|
||||
def forward(self, x):
|
||||
h = self.in_proj(x)
|
||||
h = self.blocks(h)
|
||||
return x + self.out_proj(h)
|
||||
|
||||
|
||||
class ProgressiveConvDecoder(nn.Module):
|
||||
def __init__(self, hidden_dim=4096, out_channels=3):
|
||||
super().__init__()
|
||||
|
||||
# self.proj = nn.Linear(hidden_dim, 1024)
|
||||
# self.act = nn.SiLU()
|
||||
|
||||
self.up_blocks = nn.ModuleList(
|
||||
[
|
||||
nn.Sequential(
|
||||
nn.Upsample(scale_factor=2, mode="nearest"),
|
||||
nn.Conv2d(hidden_dim, 512, kernel_size=3, padding=1),
|
||||
nn.GroupNorm(32, 512),
|
||||
nn.SiLU(),
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.Upsample(scale_factor=2, mode="nearest"),
|
||||
nn.Conv2d(512, 256, kernel_size=3, padding=1),
|
||||
nn.GroupNorm(32, 256),
|
||||
nn.SiLU(),
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.Upsample(scale_factor=2, mode="nearest"),
|
||||
nn.Conv2d(256, 64, kernel_size=3, padding=1),
|
||||
nn.GroupNorm(32, 64),
|
||||
nn.SiLU(),
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.Upsample(scale_factor=2, mode="nearest"),
|
||||
nn.Conv2d(64, 32, kernel_size=3, padding=1),
|
||||
nn.GroupNorm(16, 32),
|
||||
nn.SiLU(),
|
||||
),
|
||||
nn.Sequential(
|
||||
nn.Upsample(scale_factor=2, mode="nearest"),
|
||||
nn.Conv2d(32, 16, kernel_size=3, padding=1),
|
||||
nn.SiLU(),
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
self.out_conv = nn.Conv2d(16, out_channels, kernel_size=3, padding=1)
|
||||
|
||||
def forward(self, x_2d):
|
||||
# B, C, H, W = x_2d.shape
|
||||
# x = x_2d.permute(0, 2, 3, 1).contiguous() # (B, H, W, C)
|
||||
# x = self.proj(x)
|
||||
# x = self.act(x)
|
||||
# x = x.permute(0, 3, 1, 2).contiguous() # (B, 512, H, W)
|
||||
x = x_2d
|
||||
for block in self.up_blocks:
|
||||
x = block(x)
|
||||
|
||||
out = self.out_conv(x)
|
||||
return out
|
||||
|
||||
|
||||
class PatchDecoder_postps(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# layer 1: H/32 -> H/8 (4x upscale)
|
||||
|
||||
self.conv1 = nn.Conv2d(4096, 4096, kernel_size=3, padding=1)
|
||||
self.ps1 = nn.PixelShuffle(4)
|
||||
self.act1 = nn.GELU()
|
||||
|
||||
# layer 2: H/8 -> H (8x upscale)
|
||||
self.conv2 = nn.Conv2d(256, 192, kernel_size=3, padding=1)
|
||||
self.ps2 = nn.PixelShuffle(8)
|
||||
|
||||
def forward(self, x):
|
||||
# x shape: [B, 4096, H/32, W/32]
|
||||
x = self.ps1(self.act1(self.conv1(x))) # -> [B, 256, H/8, W/8]
|
||||
x = self.ps2(self.conv2(x)) # -> [B, 3, H, W]
|
||||
return x
|
||||
|
||||
|
||||
class PatchDecoder_preps(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# layer 1: H/32 -> H/16 (2x upscale)
|
||||
self.ps1 = nn.PixelShuffle(2)
|
||||
self.conv1 = nn.Conv2d(1024, 1024, kernel_size=3, padding=1)
|
||||
self.act1 = nn.GELU()
|
||||
|
||||
# layer 2: H/16 -> H/8 (2x upscale)
|
||||
self.ps2 = nn.PixelShuffle(2)
|
||||
self.conv2 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
|
||||
self.act2 = nn.GELU()
|
||||
|
||||
# layer 3: H/8 -> H (8x upscale)
|
||||
self.ps3 = nn.PixelShuffle(8)
|
||||
self.conv3 = nn.Conv2d(4, 3, kernel_size=3, padding=1)
|
||||
|
||||
def forward(self, x):
|
||||
# x shape: [B, 4096, H/32, W/32]
|
||||
x = self.act1(self.conv1(self.ps1((x)))) # -> [B, 256, H/16, W/16]
|
||||
x = self.act2(self.conv2(self.ps2((x)))) # -> [B, 256, H/8, W/8]
|
||||
x = self.conv3(self.ps3((x))) # -> [B, 3, H, W]
|
||||
return x
|
||||
|
||||
|
||||
class PatchDecoder_preps1(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
# layer 1: H/32 -> H/16 (2x upscale)
|
||||
self.ps1 = nn.PixelShuffle(2)
|
||||
self.conv1 = nn.Conv2d(1024, 1024, kernel_size=3, padding=1)
|
||||
self.act1 = nn.GELU()
|
||||
|
||||
# layer 2: H/16 -> H/8 (2x upscale)
|
||||
self.ps2 = nn.PixelShuffle(2)
|
||||
self.conv2 = nn.Conv2d(256, 192, kernel_size=3, padding=1)
|
||||
|
||||
# layer 3: H/8 -> H (8x upscale)
|
||||
self.ps3 = nn.PixelShuffle(8)
|
||||
|
||||
def forward(self, x):
|
||||
# x shape: [B, 4096, H/32, W/32]
|
||||
x = self.act1(self.conv1(self.ps1((x)))) # -> [B, 256, H/16, W/16]
|
||||
x = self.ps3(self.conv2(self.ps2((x)))) # -> [B, 256, H/8, W/8]
|
||||
return x
|
||||
|
||||
|
||||
class ConvDecoder(nn.Module):
|
||||
def __init__(self, input_dim=4096, hidden_dim=1024):
|
||||
super().__init__()
|
||||
# layer 1: H/32 -> H/16 (2x upscale)
|
||||
self.ps1 = nn.PixelShuffle(2)
|
||||
self.conv1 = nn.Conv2d(input_dim // 4, hidden_dim, kernel_size=3, padding=1)
|
||||
self.act1 = nn.GELU()
|
||||
|
||||
# layer 2: H/16 -> H/8 (2x upscale)
|
||||
self.ps2 = nn.PixelShuffle(2)
|
||||
self.conv2 = nn.Conv2d(hidden_dim // 4, 192, kernel_size=3, padding=1)
|
||||
|
||||
# layer 3: H/8 -> H (8x upscale)
|
||||
self.ps3 = nn.PixelShuffle(8)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.act1(self.conv1(self.ps1((x))))
|
||||
x = self.ps3(self.conv2(self.ps2((x))))
|
||||
return x
|
||||
+2739
File diff suppressed because it is too large
Load Diff
+285
@@ -0,0 +1,285 @@
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
from typing import Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.utils.checkpoint
|
||||
from torch import nn
|
||||
from transformers.modeling_outputs import BaseModelOutputWithPooling
|
||||
from transformers.modeling_utils import PreTrainedModel
|
||||
|
||||
from .configuration_neo_vit import NEOVisionConfig
|
||||
|
||||
|
||||
def precompute_rope_freqs_sincos(
|
||||
dim: int, max_position: int, base: float = 10000.0, device=None
|
||||
):
|
||||
"""Precompute 1D RoPE cosine and sine values."""
|
||||
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device).float() / dim))
|
||||
t = torch.arange(max_position, device=device).type_as(inv_freq)
|
||||
freqs = torch.outer(t, inv_freq)
|
||||
return torch.cos(freqs), torch.sin(freqs)
|
||||
|
||||
|
||||
def build_abs_positions_from_grid_hw(grid_hw: torch.Tensor, device=None):
|
||||
"""
|
||||
Compute patch coordinates (x, y)
|
||||
|
||||
Args:
|
||||
grid_hw: (B, 2) tensor representing (H, W) per image
|
||||
"""
|
||||
device = grid_hw.device
|
||||
B = grid_hw.shape[0]
|
||||
|
||||
# Get the number of patches per image
|
||||
H = grid_hw[:, 0]
|
||||
W = grid_hw[:, 1]
|
||||
N = H * W
|
||||
N_total = N.sum()
|
||||
|
||||
# Create the batch index for each patch (B x patch count)
|
||||
patch_to_sample = torch.repeat_interleave(
|
||||
torch.arange(B, device=device), N
|
||||
) # (N_total,)
|
||||
|
||||
# Generate intra-image patch index (row-major order)
|
||||
patch_id_within_image = torch.arange(N_total, device=device)
|
||||
patch_id_within_image = (
|
||||
patch_id_within_image
|
||||
- torch.cumsum(torch.cat([torch.tensor([0], device=device), N[:-1]]), dim=0)[
|
||||
patch_to_sample
|
||||
]
|
||||
)
|
||||
|
||||
# Get H/W for each patch according to its image
|
||||
W_per_patch = W[patch_to_sample]
|
||||
abs_x = patch_id_within_image % W_per_patch
|
||||
abs_y = patch_id_within_image // W_per_patch
|
||||
|
||||
return abs_x, abs_y
|
||||
|
||||
|
||||
def apply_rotary_emb_1d(
|
||||
x: torch.Tensor,
|
||||
cos_cached: torch.Tensor,
|
||||
sin_cached: torch.Tensor,
|
||||
positions: torch.Tensor,
|
||||
):
|
||||
"""Apply 1D RoPE to part of the input tensor."""
|
||||
# x: (..., seq_len, dim_part)
|
||||
# positions: (..., seq_len)
|
||||
# cos_cached: (max_pos, dim_part / 2)
|
||||
|
||||
cos = cos_cached[positions] # Shape: (positions.shape, dim_part / 2)
|
||||
sin = sin_cached[positions] # Shape: (positions.shape, dim_part / 2)
|
||||
|
||||
x1 = x[..., 0::2]
|
||||
x2 = x[..., 1::2]
|
||||
|
||||
rotated_x1 = x1 * cos - x2 * sin
|
||||
rotated_x2 = x1 * sin + x2 * cos
|
||||
|
||||
x_rotated = torch.empty_like(x)
|
||||
x_rotated[..., 0::2] = rotated_x1
|
||||
x_rotated[..., 1::2] = rotated_x2
|
||||
return x_rotated
|
||||
|
||||
|
||||
def apply_2d_rotary_pos_emb(
|
||||
x: torch.Tensor,
|
||||
cos_cached_x: torch.Tensor,
|
||||
sin_cached_x: torch.Tensor,
|
||||
cos_cached_y: torch.Tensor,
|
||||
sin_cached_y: torch.Tensor,
|
||||
abs_positions_x: torch.Tensor,
|
||||
abs_positions_y: torch.Tensor,
|
||||
):
|
||||
"""Apply 2D RoPE to input tensor x."""
|
||||
dim = x.shape[-1]
|
||||
dim_half = dim // 2
|
||||
|
||||
# Use the first half of the embedding for one RoPE direction and the second
|
||||
# half for the other direction. The split order must stay consistent.
|
||||
x_part_1 = x[..., :dim_half]
|
||||
x_part_2 = x[..., dim_half:]
|
||||
|
||||
# Apply rotations associated with abs_positions_x to x_part_1.
|
||||
rotated_part_1 = apply_rotary_emb_1d(
|
||||
x_part_1, cos_cached_x, sin_cached_x, abs_positions_x
|
||||
)
|
||||
# Apply rotations associated with abs_positions_y to x_part_2.
|
||||
rotated_part_2 = apply_rotary_emb_1d(
|
||||
x_part_2, cos_cached_y, sin_cached_y, abs_positions_y
|
||||
)
|
||||
|
||||
# Concatenate them back in the same order used by the split.
|
||||
return torch.cat((rotated_part_1, rotated_part_2), dim=-1)
|
||||
|
||||
|
||||
class NEOVisionEmbeddings(nn.Module):
|
||||
"""
|
||||
Embedding Module for Vision.
|
||||
"""
|
||||
|
||||
def __init__(self, config: NEOVisionConfig):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.embed_dim = config.hidden_size
|
||||
self.llm_embed_dim = config.llm_hidden_size[0]
|
||||
self.downsample_factor = int(1 / config.downsample_ratio[0])
|
||||
self.patch_size = config.patch_size
|
||||
|
||||
self.patch_embedding = nn.Conv2d(
|
||||
in_channels=config.num_channels,
|
||||
out_channels=self.embed_dim,
|
||||
kernel_size=self.patch_size,
|
||||
stride=self.patch_size,
|
||||
)
|
||||
self.dense_embedding = nn.Conv2d(
|
||||
in_channels=self.embed_dim,
|
||||
out_channels=self.llm_embed_dim,
|
||||
kernel_size=self.downsample_factor,
|
||||
stride=self.downsample_factor,
|
||||
)
|
||||
self.gelu = nn.GELU()
|
||||
|
||||
self.rope_dim_part = self.embed_dim // 2
|
||||
self.max_position_embeddings_vision = config.max_position_embeddings_vision
|
||||
self.rope_theta_vision = config.rope_theta_vision
|
||||
|
||||
# These deterministic caches are not checkpoint state. In
|
||||
# Transformers 5, ``from_pretrained`` constructs models on the meta
|
||||
# device; tensors computed here would later be materialized as
|
||||
# uninitialized memory because persistent=False buffers are absent
|
||||
# from the checkpoint. Build them lazily on the first real device.
|
||||
self.register_buffer("cos_cached_x", None, persistent=False)
|
||||
self.register_buffer("sin_cached_x", None, persistent=False)
|
||||
self.register_buffer("cos_cached_y", None, persistent=False)
|
||||
self.register_buffer("sin_cached_y", None, persistent=False)
|
||||
|
||||
def _ensure_rope_cache(self, device: torch.device) -> None:
|
||||
if self.cos_cached_x is not None and self.cos_cached_x.device == device:
|
||||
return
|
||||
|
||||
cos, sin = precompute_rope_freqs_sincos(
|
||||
self.rope_dim_part,
|
||||
self.max_position_embeddings_vision,
|
||||
base=self.rope_theta_vision,
|
||||
device=device,
|
||||
)
|
||||
self.cos_cached_x = cos
|
||||
self.sin_cached_x = sin
|
||||
self.cos_cached_y = cos.clone()
|
||||
self.sin_cached_y = sin.clone()
|
||||
|
||||
def _apply_2d_rotary_pos_emb(self, patch_embeds, grid_hw):
|
||||
"""
|
||||
Apply 2D Rotary Position Embedding to the patch embeddings.
|
||||
"""
|
||||
abs_pos_x, abs_pos_y = build_abs_positions_from_grid_hw(
|
||||
grid_hw, device=patch_embeds.device
|
||||
)
|
||||
embeddings = apply_2d_rotary_pos_emb(
|
||||
patch_embeds.to(
|
||||
torch.float32
|
||||
), # RoPE calculations are often more stable in float32
|
||||
self.cos_cached_x,
|
||||
self.sin_cached_x,
|
||||
self.cos_cached_y,
|
||||
self.sin_cached_y,
|
||||
abs_pos_x,
|
||||
abs_pos_y,
|
||||
).to(self.patch_embedding.weight.dtype)
|
||||
return embeddings
|
||||
|
||||
def forward(self, pixel_values: torch.FloatTensor, grid_hw=None) -> torch.Tensor:
|
||||
|
||||
pixel_values = pixel_values.view( #
|
||||
-1,
|
||||
3,
|
||||
self.patch_size,
|
||||
self.patch_size,
|
||||
) # [28072, 768] -> [28072, 3, 16, 16]
|
||||
patch_embeds = self.gelu(self.patch_embedding(pixel_values)).view(
|
||||
-1, self.embed_dim
|
||||
)
|
||||
self._ensure_rope_cache(patch_embeds.device)
|
||||
patch_embeds = self._apply_2d_rotary_pos_emb(
|
||||
patch_embeds, grid_hw
|
||||
) # [28072, 1024]
|
||||
assert (grid_hw[:, 0] * grid_hw[:, 1]).sum() == patch_embeds.shape[0]
|
||||
|
||||
patches_list = []
|
||||
cur_position = 0
|
||||
for i in range(grid_hw.shape[0]):
|
||||
h, w = grid_hw[i]
|
||||
patches_per_img = (
|
||||
patch_embeds[cur_position : cur_position + h * w]
|
||||
.view(h, w, -1)
|
||||
.unsqueeze(0)
|
||||
)
|
||||
patches_per_img = self.dense_embedding(patches_per_img.permute(0, 3, 1, 2))
|
||||
patches_per_img = patches_per_img.permute(0, 2, 3, 1)
|
||||
patches_list.append(patches_per_img.view(-1, patches_per_img.shape[-1]))
|
||||
cur_position += h * w
|
||||
|
||||
embeddings = torch.cat(
|
||||
patches_list, dim=0
|
||||
) # (N_total // downsample_factor**2, C)
|
||||
|
||||
assert cur_position == patch_embeds.shape[0]
|
||||
assert embeddings.shape[0] == int(
|
||||
patch_embeds.shape[0] / self.downsample_factor**2
|
||||
)
|
||||
|
||||
return embeddings
|
||||
|
||||
|
||||
class NEOVisionModel(PreTrainedModel):
|
||||
main_input_name = "pixel_values"
|
||||
_supports_flash_attn_2 = True
|
||||
supports_gradient_checkpointing = True
|
||||
config_class = NEOVisionConfig
|
||||
# support transformers 4.51.+
|
||||
_tp_plan = ""
|
||||
|
||||
def __init__(self, config: NEOVisionConfig):
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
|
||||
self.embeddings = NEOVisionEmbeddings(config)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
pixel_values: Optional[torch.FloatTensor] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
pixel_embeds: Optional[torch.FloatTensor] = None,
|
||||
grid_hw: Optional[torch.Tensor] = None,
|
||||
) -> Union[Tuple, BaseModelOutputWithPooling]:
|
||||
output_hidden_states = (
|
||||
output_hidden_states
|
||||
if output_hidden_states is not None
|
||||
else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = (
|
||||
return_dict if return_dict is not None else self.config.use_return_dict
|
||||
)
|
||||
|
||||
if pixel_values is None and pixel_embeds is None:
|
||||
raise ValueError("You have to specify pixel_values or pixel_embeds")
|
||||
|
||||
if pixel_embeds is not None:
|
||||
hidden_states = pixel_embeds
|
||||
else:
|
||||
assert pixel_values.dim() == 2, (
|
||||
f"pixel_values must be 2D for native resolution, got: {pixel_values.dim()}"
|
||||
)
|
||||
hidden_states = self.embeddings(pixel_values, grid_hw=grid_hw)
|
||||
|
||||
return BaseModelOutputWithPooling(
|
||||
last_hidden_state=hidden_states,
|
||||
pooler_output=None,
|
||||
hidden_states=None,
|
||||
attentions=None,
|
||||
)
|
||||
+1540
File diff suppressed because it is too large
Load Diff
+619
@@ -0,0 +1,619 @@
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import nn
|
||||
from transformers.cache_utils import Cache, DynamicCache
|
||||
from transformers.generation import GenerationMixin
|
||||
from transformers.masking_utils import create_causal_mask
|
||||
from transformers.modeling_layers import GradientCheckpointingLayer
|
||||
from transformers.modeling_outputs import (
|
||||
BaseModelOutputWithPast,
|
||||
CausalLMOutputWithPast,
|
||||
)
|
||||
from transformers.modeling_utils import PreTrainedModel
|
||||
from transformers.processing_utils import Unpack
|
||||
from transformers.utils import TransformersKwargs, can_return_tuple
|
||||
from transformers.utils.deprecation import deprecate_kwarg
|
||||
|
||||
from .configuration_neo_chat import NEOMoELLMConfig
|
||||
from .modeling_qwen3 import (
|
||||
Qwen3Attention,
|
||||
Qwen3RMSNorm,
|
||||
create_block_causal_mask,
|
||||
)
|
||||
from .transformers_compat import (
|
||||
causal_mask_kwargs,
|
||||
model_input_compat,
|
||||
tied_weights_keys,
|
||||
)
|
||||
|
||||
|
||||
class Qwen3MoeMLP(nn.Module):
|
||||
"""Single expert FFN. Same structure as :class:`Qwen3MLP` but the
|
||||
intermediate size is parameterised so it can be ``moe_intermediate_size``
|
||||
(per-expert) for experts and ``intermediate_size`` for any dense fallback.
|
||||
"""
|
||||
|
||||
def __init__(self, config, intermediate_size: Optional[int] = None):
|
||||
super().__init__()
|
||||
from transformers.activations import ACT2FN
|
||||
|
||||
self.config = config
|
||||
self.hidden_size = config.hidden_size
|
||||
self.intermediate_size = (
|
||||
intermediate_size
|
||||
if intermediate_size is not None
|
||||
else config.intermediate_size
|
||||
)
|
||||
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
|
||||
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
|
||||
self.act_fn = ACT2FN[config.hidden_act]
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
|
||||
|
||||
|
||||
class Qwen3MoeSparseMoeBlock(nn.Module):
|
||||
"""Top-k softmax-routed MoE block matching HuggingFace's Qwen3-MoE layout.
|
||||
|
||||
Parameter names (``gate.weight``, ``experts.{i}.gate_proj/up_proj/down_proj``)
|
||||
are kept identical so converted A3B checkpoints load directly via the
|
||||
``mlp.*`` / ``mlp_mot_gen.*`` keys. The block is parameterised explicitly
|
||||
so the same class can serve both the understanding branch (``num_experts``
|
||||
experts, top-k = ``num_experts_per_tok``, width ``moe_intermediate_size``)
|
||||
and the image-generation branch (``gen_num_experts`` etc.).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: NEOMoELLMConfig,
|
||||
num_experts: Optional[int] = None,
|
||||
num_experts_per_tok: Optional[int] = None,
|
||||
moe_intermediate_size: Optional[int] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.num_experts = (
|
||||
int(num_experts) if num_experts is not None else int(config.num_experts)
|
||||
)
|
||||
self.top_k = int(
|
||||
num_experts_per_tok
|
||||
if num_experts_per_tok is not None
|
||||
else config.num_experts_per_tok
|
||||
)
|
||||
self.norm_topk_prob = bool(getattr(config, "norm_topk_prob", True))
|
||||
self.hidden_size = config.hidden_size
|
||||
|
||||
expert_intermediate_size = int(
|
||||
moe_intermediate_size
|
||||
if moe_intermediate_size is not None
|
||||
else config.moe_intermediate_size
|
||||
)
|
||||
|
||||
self.gate = nn.Linear(config.hidden_size, self.num_experts, bias=False)
|
||||
self.experts = nn.ModuleList(
|
||||
[
|
||||
Qwen3MoeMLP(config, intermediate_size=expert_intermediate_size)
|
||||
for _ in range(self.num_experts)
|
||||
]
|
||||
)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
orig_shape = hidden_states.shape
|
||||
hidden_dim = orig_shape[-1]
|
||||
flat = hidden_states.view(-1, hidden_dim)
|
||||
n_tokens = flat.shape[0]
|
||||
|
||||
router_logits = self.gate(flat)
|
||||
routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float32)
|
||||
routing_weights, selected_experts = torch.topk(
|
||||
routing_weights, self.top_k, dim=-1
|
||||
)
|
||||
if self.norm_topk_prob:
|
||||
routing_weights = routing_weights / routing_weights.sum(
|
||||
dim=-1, keepdim=True
|
||||
)
|
||||
routing_weights = routing_weights.to(flat.dtype)
|
||||
|
||||
output = torch.zeros(
|
||||
(n_tokens, hidden_dim), dtype=flat.dtype, device=flat.device
|
||||
)
|
||||
# (num_experts, top_k, num_tokens)
|
||||
expert_mask = F.one_hot(selected_experts, num_classes=self.num_experts).permute(
|
||||
2, 1, 0
|
||||
)
|
||||
|
||||
for expert_idx in range(self.num_experts):
|
||||
idx, top_x = torch.where(expert_mask[expert_idx])
|
||||
if top_x.numel() == 0:
|
||||
continue
|
||||
expert_layer = self.experts[expert_idx]
|
||||
current_state = flat.index_select(0, top_x)
|
||||
current_out = (
|
||||
expert_layer(current_state) * routing_weights[top_x, idx, None]
|
||||
)
|
||||
output.index_add_(0, top_x, current_out.to(flat.dtype))
|
||||
|
||||
return output.view(*orig_shape)
|
||||
|
||||
|
||||
class Qwen3MoeDecoderLayer(GradientCheckpointingLayer):
|
||||
"""A Qwen3-MoE decoder block with the NEO-Unify two-branch structure.
|
||||
|
||||
Mirrors ``Qwen3DecoderLayer`` from :mod:`modeling_qwen3` but uses sparse
|
||||
MoE blocks on *both* branches:
|
||||
|
||||
* ``self.mlp`` - understanding-path MoE
|
||||
(``num_experts`` / ``num_experts_per_tok`` /
|
||||
``moe_intermediate_size``)
|
||||
* ``self.mlp_mot_gen`` - image-generation-path MoE
|
||||
(``gen_num_experts`` / ``gen_num_experts_per_tok`` /
|
||||
``gen_moe_intermediate_size``)
|
||||
|
||||
Layers listed in ``mlp_only_layers`` or those not aligned with
|
||||
``decoder_sparse_step`` fall back to a dense :class:`Qwen3MoeMLP` on the
|
||||
understanding branch (matching upstream Qwen3-MoE), while the
|
||||
generation branch still uses a sparse MoE.
|
||||
"""
|
||||
|
||||
def __init__(self, config: NEOMoELLMConfig, layer_idx: int):
|
||||
super().__init__()
|
||||
self.hidden_size = config.hidden_size
|
||||
self.self_attn = Qwen3Attention(config=config, layer_idx=layer_idx)
|
||||
|
||||
mlp_only_layers = list(getattr(config, "mlp_only_layers", []) or [])
|
||||
decoder_sparse_step = int(getattr(config, "decoder_sparse_step", 1) or 1)
|
||||
is_sparse = (
|
||||
int(config.num_experts) > 0
|
||||
and layer_idx not in mlp_only_layers
|
||||
and (layer_idx + 1) % decoder_sparse_step == 0
|
||||
)
|
||||
|
||||
if is_sparse:
|
||||
self.mlp = Qwen3MoeSparseMoeBlock(
|
||||
config,
|
||||
num_experts=config.num_experts,
|
||||
num_experts_per_tok=config.num_experts_per_tok,
|
||||
moe_intermediate_size=config.moe_intermediate_size,
|
||||
)
|
||||
else:
|
||||
self.mlp = Qwen3MoeMLP(config, intermediate_size=config.intermediate_size)
|
||||
|
||||
# Image-generation branch: in the A3B checkpoint this is *also* a sparse
|
||||
# MoE block (``gen_num_experts`` experts, typically smaller than the und
|
||||
# branch's ``num_experts``). ``NEOMoELLMConfig`` defaults the gen-path
|
||||
# knobs to their und-path counterparts so legacy single-pool configs
|
||||
# keep working.
|
||||
self.mlp_mot_gen = Qwen3MoeSparseMoeBlock(
|
||||
config,
|
||||
num_experts=getattr(config, "gen_num_experts", config.num_experts),
|
||||
num_experts_per_tok=getattr(
|
||||
config, "gen_num_experts_per_tok", config.num_experts_per_tok
|
||||
),
|
||||
moe_intermediate_size=getattr(
|
||||
config, "gen_moe_intermediate_size", config.moe_intermediate_size
|
||||
),
|
||||
)
|
||||
|
||||
self.input_layernorm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.input_layernorm_mot_gen = Qwen3RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps
|
||||
)
|
||||
self.post_attention_layernorm = Qwen3RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps
|
||||
)
|
||||
self.post_attention_layernorm_mot_gen = Qwen3RMSNorm(
|
||||
config.hidden_size, eps=config.rms_norm_eps
|
||||
)
|
||||
self.attention_type = config.layer_types[layer_idx]
|
||||
|
||||
def forward_und(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
image_gen_indicators: torch.Tensor,
|
||||
exist_non_image_gen_tokens: bool,
|
||||
exist_image_gen_tokens: bool,
|
||||
indexes: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
use_cache: Optional[bool] = False,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> torch.Tensor:
|
||||
residual = hidden_states
|
||||
hidden_states = self.input_layernorm(hidden_states)
|
||||
hidden_states, _ = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
image_gen_indicators=image_gen_indicators,
|
||||
exist_non_image_gen_tokens=exist_non_image_gen_tokens,
|
||||
exist_image_gen_tokens=exist_image_gen_tokens,
|
||||
indexes=indexes,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
residual = hidden_states
|
||||
hidden_states = self.post_attention_layernorm(hidden_states)
|
||||
hidden_states = self.mlp(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
return hidden_states
|
||||
|
||||
def forward_gen(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
image_gen_indicators: torch.Tensor,
|
||||
exist_non_image_gen_tokens: bool,
|
||||
exist_image_gen_tokens: bool,
|
||||
indexes: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
use_cache: Optional[bool] = False,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> torch.Tensor:
|
||||
residual = hidden_states
|
||||
hidden_states = self.input_layernorm_mot_gen(hidden_states)
|
||||
hidden_states, _ = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
image_gen_indicators=image_gen_indicators,
|
||||
exist_non_image_gen_tokens=exist_non_image_gen_tokens,
|
||||
exist_image_gen_tokens=exist_image_gen_tokens,
|
||||
indexes=indexes,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
residual = hidden_states
|
||||
hidden_states = self.post_attention_layernorm_mot_gen(hidden_states)
|
||||
hidden_states = self.mlp_mot_gen(hidden_states)
|
||||
hidden_states = residual + hidden_states
|
||||
return hidden_states
|
||||
|
||||
@deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
image_gen_indicators: torch.Tensor,
|
||||
exist_non_image_gen_tokens: bool,
|
||||
exist_image_gen_tokens: bool,
|
||||
indexes: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
use_cache: Optional[bool] = False,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> torch.Tensor:
|
||||
if exist_non_image_gen_tokens and not exist_image_gen_tokens:
|
||||
return self.forward_und(
|
||||
hidden_states,
|
||||
image_gen_indicators,
|
||||
exist_non_image_gen_tokens,
|
||||
exist_image_gen_tokens,
|
||||
indexes,
|
||||
attention_mask,
|
||||
position_ids,
|
||||
past_key_values,
|
||||
use_cache,
|
||||
cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
if not exist_non_image_gen_tokens and exist_image_gen_tokens:
|
||||
return self.forward_gen(
|
||||
hidden_states,
|
||||
image_gen_indicators,
|
||||
exist_non_image_gen_tokens,
|
||||
exist_image_gen_tokens,
|
||||
indexes,
|
||||
attention_mask,
|
||||
position_ids,
|
||||
past_key_values,
|
||||
use_cache,
|
||||
cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
# Mixed und/gen path — see the NOTE in Qwen3Attention.forward (modeling_qwen3.py).
|
||||
raise NotImplementedError(
|
||||
"Mixed und/gen decoder-layer forward is not yet validated (issue #207). "
|
||||
"Split the sequence at token-type boundaries and use forward_und / forward_gen."
|
||||
)
|
||||
|
||||
# Mixed batch: dispatch tokens per branch then merge back. Matches the
|
||||
# dense ``Qwen3DecoderLayer.forward`` mixed-path implementation.
|
||||
residual = hidden_states
|
||||
|
||||
_hidden_states = hidden_states.new_zeros(hidden_states.shape)
|
||||
if exist_non_image_gen_tokens:
|
||||
_hidden_states[~image_gen_indicators] = self.input_layernorm(
|
||||
hidden_states[~image_gen_indicators]
|
||||
)
|
||||
if exist_image_gen_tokens:
|
||||
_hidden_states[image_gen_indicators] = self.input_layernorm_mot_gen(
|
||||
hidden_states[image_gen_indicators]
|
||||
)
|
||||
hidden_states = _hidden_states
|
||||
|
||||
hidden_states, _ = self.self_attn(
|
||||
hidden_states=hidden_states,
|
||||
image_gen_indicators=image_gen_indicators,
|
||||
exist_non_image_gen_tokens=exist_non_image_gen_tokens,
|
||||
exist_image_gen_tokens=exist_image_gen_tokens,
|
||||
indexes=indexes,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
hidden_states = residual + hidden_states
|
||||
|
||||
residual = hidden_states
|
||||
|
||||
_hidden_states = hidden_states.new_zeros(hidden_states.shape)
|
||||
if exist_non_image_gen_tokens:
|
||||
und_hidden = self.post_attention_layernorm(
|
||||
hidden_states[~image_gen_indicators]
|
||||
)
|
||||
# MoE expects a 3D input (batch, seq, hidden); promote then squeeze.
|
||||
if und_hidden.dim() == 2:
|
||||
und_hidden = und_hidden.unsqueeze(0)
|
||||
_hidden_states[~image_gen_indicators] = self.mlp(und_hidden).squeeze(0)
|
||||
else:
|
||||
_hidden_states[~image_gen_indicators] = self.mlp(und_hidden)
|
||||
if exist_image_gen_tokens:
|
||||
_hidden_states[image_gen_indicators] = self.mlp_mot_gen(
|
||||
self.post_attention_layernorm_mot_gen(
|
||||
hidden_states[image_gen_indicators]
|
||||
)
|
||||
)
|
||||
|
||||
hidden_states = _hidden_states
|
||||
hidden_states = residual + hidden_states
|
||||
return hidden_states
|
||||
|
||||
|
||||
class Qwen3MoePreTrainedModel(PreTrainedModel):
|
||||
config: NEOMoELLMConfig
|
||||
base_model_prefix = "model"
|
||||
supports_gradient_checkpointing = True
|
||||
_no_split_modules = ["Qwen3MoeDecoderLayer"]
|
||||
_skip_keys_device_placement = ["past_key_values"]
|
||||
_supports_flash_attn = True
|
||||
_supports_sdpa = True
|
||||
_supports_flex_attn = True
|
||||
|
||||
_can_compile_fullgraph = False # MoE routing has data-dependent control flow.
|
||||
_supports_attention_backend = True
|
||||
_can_record_outputs = {
|
||||
"hidden_states": Qwen3MoeDecoderLayer,
|
||||
"attentions": Qwen3Attention,
|
||||
}
|
||||
|
||||
|
||||
class Qwen3MoeModel(Qwen3MoePreTrainedModel):
|
||||
def __init__(self, config: NEOMoELLMConfig):
|
||||
super().__init__(config)
|
||||
self.padding_idx = config.pad_token_id
|
||||
self.vocab_size = config.vocab_size
|
||||
|
||||
self.embed_tokens = nn.Embedding(
|
||||
config.vocab_size, config.hidden_size, self.padding_idx
|
||||
)
|
||||
self.layers = nn.ModuleList(
|
||||
[
|
||||
Qwen3MoeDecoderLayer(config, layer_idx)
|
||||
for layer_idx in range(config.num_hidden_layers)
|
||||
]
|
||||
)
|
||||
self.norm = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
self.norm_mot_gen = Qwen3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
self.has_sliding_layers = "sliding_attention" in self.config.layer_types
|
||||
self.current_index = -1
|
||||
|
||||
self.post_init()
|
||||
|
||||
@model_input_compat
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Optional[torch.LongTensor] = None,
|
||||
image_gen_indicators: Optional[torch.Tensor] = None,
|
||||
indexes: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> BaseModelOutputWithPast:
|
||||
if image_gen_indicators is None:
|
||||
exist_non_image_gen_tokens = True
|
||||
exist_image_gen_tokens = False
|
||||
else:
|
||||
# Convert the CUDA reductions once before the decoder loop. If the
|
||||
# scalar tensors reach every layer, each Python branch can force a
|
||||
# host-device synchronization and collapse async weight prefetch.
|
||||
exist_non_image_gen_tokens = bool((~image_gen_indicators).any().item())
|
||||
exist_image_gen_tokens = bool(image_gen_indicators.any().item())
|
||||
|
||||
if (input_ids is None) ^ (inputs_embeds is not None):
|
||||
raise ValueError(
|
||||
"You must specify exactly one of input_ids or inputs_embeds"
|
||||
)
|
||||
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds = self.embed_tokens(input_ids)
|
||||
|
||||
if use_cache and past_key_values is None:
|
||||
past_key_values = DynamicCache(config=self.config)
|
||||
|
||||
if cache_position is None:
|
||||
past_seen_tokens = (
|
||||
past_key_values.get_seq_length() if past_key_values is not None else 0
|
||||
)
|
||||
cache_position = torch.arange(
|
||||
past_seen_tokens,
|
||||
past_seen_tokens + inputs_embeds.shape[1],
|
||||
device=inputs_embeds.device,
|
||||
)
|
||||
|
||||
if position_ids is None:
|
||||
position_ids = cache_position.unsqueeze(0)
|
||||
|
||||
if not isinstance(causal_mask_mapping := attention_mask, dict):
|
||||
if input_ids is not None:
|
||||
mask_kwargs = causal_mask_kwargs(
|
||||
create_causal_mask,
|
||||
config=self.config,
|
||||
inputs_embeds=inputs_embeds,
|
||||
attention_mask=attention_mask,
|
||||
cache_position=cache_position,
|
||||
past_key_values=past_key_values,
|
||||
position_ids=position_ids,
|
||||
)
|
||||
causal_mask_mapping = {
|
||||
"full_attention": create_causal_mask(**mask_kwargs),
|
||||
}
|
||||
self.current_index += 1
|
||||
indexes = torch.LongTensor([[self.current_index], [0], [0]]).to(
|
||||
input_ids.device
|
||||
)
|
||||
else:
|
||||
causal_mask_mapping = {
|
||||
"full_attention": create_block_causal_mask(indexes[0]),
|
||||
}
|
||||
self.current_index = indexes[0].max()
|
||||
else:
|
||||
self.current_index = indexes[0].max()
|
||||
|
||||
hidden_states = inputs_embeds
|
||||
|
||||
for decoder_layer in self.layers[: self.config.num_hidden_layers]:
|
||||
hidden_states = decoder_layer(
|
||||
hidden_states,
|
||||
image_gen_indicators=image_gen_indicators,
|
||||
exist_non_image_gen_tokens=exist_non_image_gen_tokens,
|
||||
exist_image_gen_tokens=exist_image_gen_tokens,
|
||||
indexes=indexes,
|
||||
attention_mask=causal_mask_mapping[decoder_layer.attention_type],
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
if not exist_image_gen_tokens:
|
||||
hidden_states = self.norm(hidden_states)
|
||||
elif not exist_non_image_gen_tokens:
|
||||
hidden_states = self.norm_mot_gen(hidden_states)
|
||||
else:
|
||||
_hidden_states = hidden_states.new_zeros(hidden_states.shape)
|
||||
_hidden_states[~image_gen_indicators] = self.norm(
|
||||
hidden_states[~image_gen_indicators]
|
||||
)
|
||||
_hidden_states[image_gen_indicators] = self.norm_mot_gen(
|
||||
hidden_states[image_gen_indicators]
|
||||
)
|
||||
hidden_states = _hidden_states
|
||||
|
||||
return BaseModelOutputWithPast(
|
||||
last_hidden_state=hidden_states,
|
||||
past_key_values=past_key_values if use_cache else None,
|
||||
)
|
||||
|
||||
|
||||
class Qwen3MoeForCausalLM(Qwen3MoePreTrainedModel, GenerationMixin):
|
||||
_tied_weights_keys = tied_weights_keys(
|
||||
"lm_head.weight", "model.embed_tokens.weight"
|
||||
)
|
||||
_tp_plan = {"lm_head": "colwise_rep"}
|
||||
_pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
|
||||
|
||||
def __init__(self, config: NEOMoELLMConfig):
|
||||
super().__init__(config)
|
||||
self.model = Qwen3MoeModel(config)
|
||||
self.vocab_size = config.vocab_size
|
||||
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
|
||||
self.post_init()
|
||||
|
||||
@can_return_tuple
|
||||
def forward(
|
||||
self,
|
||||
input_ids: Optional[torch.LongTensor] = None,
|
||||
indexes: Optional[torch.LongTensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
position_ids: Optional[torch.LongTensor] = None,
|
||||
past_key_values: Optional[Cache] = None,
|
||||
inputs_embeds: Optional[torch.FloatTensor] = None,
|
||||
labels: Optional[torch.LongTensor] = None,
|
||||
use_cache: Optional[bool] = None,
|
||||
cache_position: Optional[torch.LongTensor] = None,
|
||||
logits_to_keep: Union[int, torch.Tensor] = 0,
|
||||
**kwargs: Unpack[TransformersKwargs],
|
||||
) -> CausalLMOutputWithPast:
|
||||
outputs: BaseModelOutputWithPast = self.model(
|
||||
input_ids=input_ids,
|
||||
indexes=indexes,
|
||||
attention_mask=attention_mask,
|
||||
position_ids=position_ids,
|
||||
past_key_values=past_key_values,
|
||||
inputs_embeds=inputs_embeds,
|
||||
use_cache=use_cache,
|
||||
cache_position=cache_position,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
hidden_states = outputs.last_hidden_state
|
||||
slice_indices = (
|
||||
slice(-logits_to_keep, None)
|
||||
if isinstance(logits_to_keep, int)
|
||||
else logits_to_keep
|
||||
)
|
||||
logits = self.lm_head(hidden_states[:, slice_indices, :])
|
||||
|
||||
loss = None
|
||||
if labels is not None:
|
||||
loss = self.loss_function(
|
||||
logits=logits,
|
||||
labels=labels,
|
||||
vocab_size=self.config.vocab_size,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
return CausalLMOutputWithPast(
|
||||
loss=loss,
|
||||
logits=logits,
|
||||
past_key_values=outputs.past_key_values,
|
||||
hidden_states=outputs.hidden_states,
|
||||
attentions=outputs.attentions,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Qwen3MoeForCausalLM",
|
||||
"Qwen3MoeModel",
|
||||
"Qwen3MoePreTrainedModel",
|
||||
"Qwen3MoeDecoderLayer",
|
||||
"Qwen3MoeSparseMoeBlock",
|
||||
"Qwen3MoeMLP",
|
||||
]
|
||||
+71
@@ -0,0 +1,71 @@
|
||||
"""Small compatibility seams for the supported Transformers 4/5 window."""
|
||||
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
from functools import lru_cache
|
||||
from typing import Any, Callable
|
||||
|
||||
import torch
|
||||
from packaging.version import Version
|
||||
|
||||
try:
|
||||
from transformers.utils.generic import merge_with_config_defaults
|
||||
from transformers.utils.output_capturing import capture_outputs
|
||||
|
||||
def model_input_compat(func: Callable[..., Any]) -> Callable[..., Any]:
|
||||
return merge_with_config_defaults(capture_outputs(func))
|
||||
|
||||
except ImportError: # Transformers 4.x
|
||||
pass
|
||||
|
||||
|
||||
@lru_cache(maxsize=None)
|
||||
def _parameter_names(callable_: Callable[..., Any]) -> frozenset[str]:
|
||||
return frozenset(inspect.signature(callable_).parameters)
|
||||
|
||||
|
||||
def causal_mask_kwargs(
|
||||
mask_factory: Callable[..., Any],
|
||||
*,
|
||||
config: Any,
|
||||
inputs_embeds: torch.Tensor,
|
||||
attention_mask: torch.Tensor | None,
|
||||
cache_position: torch.Tensor,
|
||||
past_key_values: Any,
|
||||
position_ids: torch.Tensor,
|
||||
) -> dict[str, Any]:
|
||||
"""Build arguments accepted by the installed ``create_causal_mask``.
|
||||
|
||||
Transformers 4.57 uses ``input_embeds`` and ``cache_position`` while newer
|
||||
5.x releases use ``inputs_embeds`` and derive the cache position internally.
|
||||
"""
|
||||
parameters = _parameter_names(mask_factory)
|
||||
embedding_parameter = (
|
||||
"inputs_embeds" if "inputs_embeds" in parameters else "input_embeds"
|
||||
)
|
||||
candidates = {
|
||||
"config": config,
|
||||
embedding_parameter: inputs_embeds,
|
||||
"attention_mask": attention_mask,
|
||||
"cache_position": cache_position,
|
||||
"past_key_values": past_key_values,
|
||||
"position_ids": position_ids,
|
||||
}
|
||||
return {name: value for name, value in candidates.items() if name in parameters}
|
||||
|
||||
|
||||
def pretrained_dtype_kwargs(dtype: torch.dtype) -> dict[str, torch.dtype]:
|
||||
"""Use the public dtype keyword supported throughout Transformers 4.57+."""
|
||||
return {"dtype": dtype}
|
||||
|
||||
|
||||
def tied_weights_keys(output_key: str, input_key: str) -> list[str] | dict[str, str]:
|
||||
"""Return the `_tied_weights_keys` shape expected by Transformers 4 or 5."""
|
||||
import transformers
|
||||
|
||||
if Version(transformers.__version__).major >= 5:
|
||||
return {output_key: input_key}
|
||||
return [output_key]
|
||||
@@ -0,0 +1,165 @@
|
||||
# Modified for SGLang; see this directory's README.md for upstream source.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torchvision.transforms as T
|
||||
from PIL import Image
|
||||
|
||||
SYSTEM_MESSAGE_FOR_GEN = (
|
||||
"You are an image generation and editing assistant that accurately understands and executes "
|
||||
"user intent.\n\nYou support two modes:\n\n1. Think Mode:\nIf the task requires reasoning, you "
|
||||
"MUST start with a <think></think> block. Put all reasoning inside the block using plain text. "
|
||||
"DO NOT include any image tags. Keep it reasonable and directly useful for producing the final "
|
||||
"image.\n\n2. Non-Think Mode:\nIf no reasoning is needed, directly produce the final image.\n\n"
|
||||
"Task Types:\n\nA. Text-to-Image Generation:\n"
|
||||
"- Generate a high-quality image based on the user's description.\n"
|
||||
"- Ensure visual clarity, semantic consistency, and completeness.\n"
|
||||
"- DO NOT introduce elements that contradict or override the user's intent.\n\n"
|
||||
"B. Image Editing:\n"
|
||||
"- Use the provided image(s) as input or reference for modification or transformation.\n"
|
||||
"- The result can be an edited image or a new image based on the reference(s).\n"
|
||||
"- Preserve all unspecified attributes unless explicitly changed.\n\n"
|
||||
"General Rules:\n"
|
||||
"- For any visible text in the image, follow the language specified for the rendered text in "
|
||||
"the user's description, not the language of the prompt. If no language is specified, use the "
|
||||
"user's input language."
|
||||
)
|
||||
|
||||
IMAGENET_MEAN = (0.485, 0.456, 0.406)
|
||||
IMAGENET_STD = (0.229, 0.224, 0.225)
|
||||
|
||||
|
||||
def round_by_factor(number: float, factor: int) -> int:
|
||||
"""Returns the closest integer to `number` that is divisible by `factor`."""
|
||||
return round(number / factor) * factor
|
||||
|
||||
|
||||
def ceil_by_factor(number: float, factor: int) -> int:
|
||||
"""Returns the smallest integer >= `number` that is divisible by `factor`."""
|
||||
return math.ceil(number / factor) * factor
|
||||
|
||||
|
||||
def floor_by_factor(number: float, factor: int) -> int:
|
||||
"""Returns the largest integer <= `number` that is divisible by `factor`."""
|
||||
return math.floor(number / factor) * factor
|
||||
|
||||
|
||||
def smart_resize(
|
||||
height: int,
|
||||
width: int,
|
||||
factor: int = 32,
|
||||
min_pixels: int = 65536,
|
||||
max_pixels: int = 4194304,
|
||||
) -> tuple[int, int]:
|
||||
"""Rescale so that H/W are divisible by `factor` and total pixels ∈ [min, max].
|
||||
|
||||
Copied from https://github.com/QwenLM/Qwen2.5-VL/blob/main/qwen-vl-utils/src/qwen_vl_utils/vision_process.py#L60
|
||||
"""
|
||||
if max(height, width) / min(height, width) > 200:
|
||||
raise ValueError(
|
||||
f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
|
||||
)
|
||||
h_bar = max(factor, round_by_factor(height, factor))
|
||||
w_bar = max(factor, round_by_factor(width, factor))
|
||||
if h_bar * w_bar > max_pixels:
|
||||
beta = math.sqrt((height * width) / max_pixels)
|
||||
h_bar = max(factor, floor_by_factor(height / beta, factor))
|
||||
w_bar = max(factor, floor_by_factor(width / beta, factor))
|
||||
elif h_bar * w_bar < min_pixels:
|
||||
beta = math.sqrt(min_pixels / (height * width))
|
||||
h_bar = ceil_by_factor(height * beta, factor)
|
||||
w_bar = ceil_by_factor(width * beta, factor)
|
||||
return h_bar, w_bar
|
||||
|
||||
|
||||
def dynamic_preprocess_native_resolution(
|
||||
image: Image.Image,
|
||||
size_factor: int = 32,
|
||||
min_pixels: int = 65536,
|
||||
max_pixels: int = 4194304,
|
||||
**_kwargs,
|
||||
) -> Image.Image:
|
||||
width, height = image.size
|
||||
resized_height, resized_width = smart_resize(
|
||||
height,
|
||||
width,
|
||||
factor=size_factor,
|
||||
min_pixels=min_pixels,
|
||||
max_pixels=max_pixels,
|
||||
)
|
||||
return image.resize((resized_width, resized_height))
|
||||
|
||||
|
||||
def preprocess_pixel_values(pixel_values: torch.Tensor, patch_size: int = 16):
|
||||
c, h, w = pixel_values.shape
|
||||
grid_h = h // patch_size
|
||||
grid_w = w // patch_size
|
||||
|
||||
flatten_pixel_values = (
|
||||
pixel_values.view(c, grid_h, patch_size, grid_w, patch_size)
|
||||
.permute(1, 3, 0, 2, 4) # [grid_h, grid_w, c, patch_size, patch_size]
|
||||
.reshape(grid_h * grid_w, c * patch_size**2)
|
||||
)
|
||||
|
||||
grid_hw = torch.tensor([[grid_h, grid_w]], device=pixel_values.device)
|
||||
return flatten_pixel_values, grid_hw
|
||||
|
||||
|
||||
def get_contrasting_background(image: Image.Image):
|
||||
"""Return a background color for RGBA->RGB conversion, or ``None`` to use default.
|
||||
|
||||
The original Neo_Unify implementation computed a contrasting background
|
||||
from the alpha channel. For this open-source release we fall back to a
|
||||
plain white background; callers that need the smarter behavior can override
|
||||
this function.
|
||||
"""
|
||||
del image
|
||||
return (255, 255, 255)
|
||||
|
||||
|
||||
def load_image_native(
|
||||
image,
|
||||
patch_size: int = 16,
|
||||
downsample_ratio: float = 0.5,
|
||||
min_pixels: int = 65536,
|
||||
max_pixels: int = 4194304,
|
||||
upscale: bool = False,
|
||||
):
|
||||
"""Load and preprocess an image: RGB convert, smart-resize, normalize, patchify."""
|
||||
if not isinstance(image, Image.Image):
|
||||
image = Image.open(image)
|
||||
if image.mode == "RGBA":
|
||||
bg_color = get_contrasting_background(image)
|
||||
if bg_color:
|
||||
background = Image.new("RGB", image.size, bg_color)
|
||||
background.paste(image, mask=image.split()[3])
|
||||
image = background.convert("RGB")
|
||||
else:
|
||||
image = image.convert("RGB")
|
||||
else:
|
||||
image = image.convert("RGB")
|
||||
|
||||
if upscale:
|
||||
image = image.resize((image.width * 2, image.height * 2), Image.BILINEAR)
|
||||
|
||||
transform = T.Compose(
|
||||
[
|
||||
T.Lambda(lambda img: img.convert("RGB") if img.mode != "RGB" else img),
|
||||
T.ToTensor(),
|
||||
T.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD),
|
||||
]
|
||||
)
|
||||
|
||||
new_image = dynamic_preprocess_native_resolution(
|
||||
image,
|
||||
size_factor=int(patch_size // downsample_ratio),
|
||||
min_pixels=min_pixels,
|
||||
max_pixels=max_pixels,
|
||||
)
|
||||
pixel_values, grid_hw = preprocess_pixel_values(
|
||||
transform(new_image).to(torch.float32), patch_size=patch_size
|
||||
)
|
||||
return pixel_values, grid_hw
|
||||
@@ -0,0 +1,73 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.sensenova_u1 import (
|
||||
SenseNovaU1PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.sensenova_u1 import (
|
||||
SenseNovaU1SamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
|
||||
from sglang.multimodal_gen.runtime.models.sensenova_u1.loader import (
|
||||
load_model_and_tokenizer,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
|
||||
ComposedPipelineBase,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages import InputValidationStage
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sensenova_u1 import (
|
||||
SenseNovaU1GenerationStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class SenseNovaU1Pipeline(ComposedPipelineBase):
|
||||
pipeline_name = "SenseNovaU1Pipeline"
|
||||
pipeline_config_cls = SenseNovaU1PipelineConfig
|
||||
sampling_params_cls = SenseNovaU1SamplingParams
|
||||
_required_config_modules: list[str] = []
|
||||
|
||||
def validate_disagg_role(self, role: RoleType) -> None:
|
||||
if role != RoleType.MONOLITHIC:
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline only supports monolithic deployment; "
|
||||
f"disaggregation role {role.value!r} is not supported"
|
||||
)
|
||||
|
||||
def load_modules(
|
||||
self,
|
||||
server_args: ServerArgs,
|
||||
loaded_modules: dict[str, torch.nn.Module] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
if loaded_modules is not None and {"model", "tokenizer"} <= set(loaded_modules):
|
||||
return loaded_modules
|
||||
|
||||
if server_args.num_gpus != 1:
|
||||
raise ValueError(
|
||||
"SenseNovaU1Pipeline currently supports num_gpus=1. "
|
||||
"Native tensor/pipeline parallelism is not implemented yet."
|
||||
)
|
||||
modules = load_model_and_tokenizer(self.model_path, server_args)
|
||||
logger.info("Loaded SenseNova-U1 model from %s", self.model_path)
|
||||
return modules
|
||||
|
||||
def create_pipeline_stages(self, server_args: ServerArgs) -> None:
|
||||
del server_args
|
||||
self.add_stage(InputValidationStage())
|
||||
self.add_stage(
|
||||
SenseNovaU1GenerationStage(
|
||||
model=self.get_module("model"),
|
||||
tokenizer=self.get_module("tokenizer"),
|
||||
),
|
||||
"sensenova_u1_generation_stage",
|
||||
)
|
||||
|
||||
|
||||
EntryClass = SenseNovaU1Pipeline
|
||||
@@ -0,0 +1,135 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Shared helpers for expanding one request into per-output requests."""
|
||||
|
||||
import os
|
||||
from copy import copy, deepcopy
|
||||
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
|
||||
from sglang.srt.observability.trace import TraceReqContext
|
||||
|
||||
|
||||
def normalize_output_seeds(
|
||||
seed: int | list[int],
|
||||
*,
|
||||
num_outputs_per_prompt: int,
|
||||
num_prompts: int = 1,
|
||||
prompt_index: int = 0,
|
||||
) -> list[int]:
|
||||
"""Return the seeds for one prompt's outputs."""
|
||||
if num_outputs_per_prompt <= 0:
|
||||
raise ValueError(
|
||||
f"num_outputs_per_prompt must be positive, got {num_outputs_per_prompt}"
|
||||
)
|
||||
|
||||
if isinstance(seed, list):
|
||||
seeds = [int(item) for item in seed]
|
||||
total_outputs = num_outputs_per_prompt * num_prompts
|
||||
if len(seeds) == num_outputs_per_prompt:
|
||||
return seeds
|
||||
if len(seeds) == total_outputs:
|
||||
start = prompt_index * num_outputs_per_prompt
|
||||
return seeds[start : start + num_outputs_per_prompt]
|
||||
raise ValueError(
|
||||
"seed list length must match num_outputs_per_prompt "
|
||||
f"({num_outputs_per_prompt}) or total outputs ({total_outputs}), "
|
||||
f"got {len(seeds)}"
|
||||
)
|
||||
|
||||
base_seed = int(seed)
|
||||
return [base_seed + i for i in range(num_outputs_per_prompt)]
|
||||
|
||||
|
||||
def _with_output_index_suffix(output_file_name: str, output_index: int) -> str:
|
||||
base, ext = os.path.splitext(output_file_name)
|
||||
return f"{base}_{output_index}{ext}"
|
||||
|
||||
|
||||
def _trace_ctx_for_output(
|
||||
req: Req,
|
||||
request_id: str | None,
|
||||
output_index: int,
|
||||
*,
|
||||
reuse_parent_trace_ctx: bool,
|
||||
):
|
||||
trace_ctx = req.trace_ctx
|
||||
if reuse_parent_trace_ctx or output_index == 0 or not trace_ctx.tracing_enable:
|
||||
return trace_ctx
|
||||
|
||||
output_trace_ctx = TraceReqContext(
|
||||
rid=request_id,
|
||||
module_name=trace_ctx.module_name,
|
||||
external_trace_header=trace_ctx.external_trace_header,
|
||||
)
|
||||
output_trace_ctx.trace_req_start()
|
||||
return output_trace_ctx
|
||||
|
||||
|
||||
def expand_request_outputs(
|
||||
req: Req,
|
||||
*,
|
||||
num_prompts: int = 1,
|
||||
prompt_index: int = 0,
|
||||
reuse_parent_trace_ctx: bool = False,
|
||||
preserve_parent_metrics: bool = False,
|
||||
) -> list[Req]:
|
||||
"""Expand one request into independent per-output requests.
|
||||
|
||||
Entry points use separate trace roots because they own and finish every
|
||||
expanded request scope. Sequential pipeline execution reuses the parent
|
||||
context because the executor owns only the parent request trace lifecycle.
|
||||
"""
|
||||
num_outputs = int(req.num_outputs_per_prompt)
|
||||
seeds = normalize_output_seeds(
|
||||
req.seed,
|
||||
num_outputs_per_prompt=num_outputs,
|
||||
num_prompts=num_prompts,
|
||||
prompt_index=prompt_index,
|
||||
)
|
||||
|
||||
if num_outputs == 1:
|
||||
req.seed = seeds[0]
|
||||
req.seeds = None
|
||||
req.generator = None
|
||||
req.sampling_params.refresh_request_extra_after_output_expansion(req)
|
||||
return [req]
|
||||
|
||||
expanded: list[Req] = []
|
||||
for output_index, seed in enumerate(seeds):
|
||||
output_request_id = (
|
||||
f"{req.request_id}:{output_index}" if req.request_id is not None else None
|
||||
)
|
||||
output_metrics = deepcopy(req.metrics) if preserve_parent_metrics else None
|
||||
output_req = copy(req)
|
||||
output_req.sampling_params = copy(req.sampling_params)
|
||||
output_req.extra = dict(req.extra)
|
||||
output_req.condition_inputs = dict(req.condition_inputs)
|
||||
output_req.trace_ctx = _trace_ctx_for_output(
|
||||
req,
|
||||
output_request_id,
|
||||
output_index,
|
||||
reuse_parent_trace_ctx=reuse_parent_trace_ctx,
|
||||
)
|
||||
output_req.seed = seed
|
||||
output_req.num_outputs_per_prompt = 1
|
||||
output_req.seeds = None
|
||||
output_req.generator = None
|
||||
output_req.extra["parent_request_id"] = req.request_id
|
||||
output_req.extra["output_index"] = output_index
|
||||
|
||||
if output_request_id is not None:
|
||||
output_req.request_id = output_request_id
|
||||
|
||||
if req.output_file_name:
|
||||
output_req.output_file_name = _with_output_index_suffix(
|
||||
req.output_file_name, output_index
|
||||
)
|
||||
output_req.sampling_params.refresh_request_extra_after_output_expansion(
|
||||
output_req
|
||||
)
|
||||
output_req.validate()
|
||||
if output_metrics is not None:
|
||||
output_req.metrics = output_metrics
|
||||
output_req.metrics.request_id = output_req.request_id
|
||||
expanded.append(output_req)
|
||||
|
||||
return expanded
|
||||
@@ -5,6 +5,8 @@
|
||||
Input validation stage for diffusion pipelines.
|
||||
"""
|
||||
|
||||
from typing import Iterator
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision.transforms.functional as TF
|
||||
@@ -13,6 +15,9 @@ from PIL import Image
|
||||
from sglang.multimodal_gen.configs.pipeline_configs import WanI2V480PConfig
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.request_utils import (
|
||||
expand_request_outputs,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
|
||||
@@ -48,6 +53,24 @@ class InputValidationStage(PipelineStage):
|
||||
super().__init__()
|
||||
self.vae_image_processor = vae_image_processor
|
||||
|
||||
def iter_sequential_requests(
|
||||
self, batch: Req, server_args: ServerArgs
|
||||
) -> Iterator[Req]:
|
||||
if not server_args.pipeline_config.supports_sequential_multi_output_inference():
|
||||
return iter((batch,))
|
||||
|
||||
num_outputs = max(1, int(batch.num_outputs_per_prompt or 1))
|
||||
if num_outputs == 1:
|
||||
return iter((batch,))
|
||||
|
||||
return iter(
|
||||
expand_request_outputs(
|
||||
batch,
|
||||
reuse_parent_trace_ctx=True,
|
||||
preserve_parent_metrics=True,
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _calculate_dimensions_from_area(
|
||||
max_area: float, aspect_ratio: float, mod_value: int
|
||||
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Model-specific helpers and stages for SenseNova-U1."""
|
||||
|
||||
from .stages import SenseNovaU1GenerationStage
|
||||
|
||||
__all__ = ["SenseNovaU1GenerationStage"]
|
||||
+6
@@ -0,0 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Pipeline lifecycle stages for the native SenseNova-U1 implementation."""
|
||||
|
||||
from .generation import SenseNovaU1GenerationStage
|
||||
|
||||
__all__ = ["SenseNovaU1GenerationStage"]
|
||||
+104
@@ -0,0 +1,104 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.configs.sensenova_u1 import (
|
||||
DEFAULT_CFG_INTERVAL,
|
||||
DEFAULT_CFG_NORM,
|
||||
DEFAULT_ENABLE_TIMESTEP_SHIFT,
|
||||
DEFAULT_T_EPS,
|
||||
DEFAULT_THINK_MODE,
|
||||
DEFAULT_TIMESTEP_SHIFT,
|
||||
SENSENOVA_U1_REQUEST_EXTRA_KEY,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.disaggregation.roles import RoleType
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import (
|
||||
OutputBatch,
|
||||
Req,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
|
||||
from sglang.multimodal_gen.runtime.server_args import ServerArgs
|
||||
|
||||
|
||||
def _denorm_sensenova_output(x: torch.Tensor) -> torch.Tensor:
|
||||
"""Convert SenseNova's normalized image tensor from [-1, 1] to [0, 1]."""
|
||||
return ((x.float() + 1.0) * 0.5).clamp(0, 1)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SenseNovaU1GenerationOptions:
|
||||
cfg_norm: str = DEFAULT_CFG_NORM
|
||||
timestep_shift: float = DEFAULT_TIMESTEP_SHIFT
|
||||
enable_timestep_shift: bool = DEFAULT_ENABLE_TIMESTEP_SHIFT
|
||||
cfg_interval: tuple[float, float] = DEFAULT_CFG_INTERVAL
|
||||
t_eps: float = DEFAULT_T_EPS
|
||||
think_mode: bool = DEFAULT_THINK_MODE
|
||||
|
||||
@classmethod
|
||||
def from_batch(cls, batch: Req) -> SenseNovaU1GenerationOptions:
|
||||
extra = batch.extra.get(SENSENOVA_U1_REQUEST_EXTRA_KEY, {})
|
||||
return cls(
|
||||
cfg_norm=extra.get("cfg_norm", DEFAULT_CFG_NORM),
|
||||
timestep_shift=float(extra.get("timestep_shift", DEFAULT_TIMESTEP_SHIFT)),
|
||||
enable_timestep_shift=bool(
|
||||
extra.get("enable_timestep_shift", DEFAULT_ENABLE_TIMESTEP_SHIFT)
|
||||
),
|
||||
cfg_interval=tuple(extra.get("cfg_interval", DEFAULT_CFG_INTERVAL)),
|
||||
t_eps=float(extra.get("t_eps", DEFAULT_T_EPS)),
|
||||
think_mode=bool(extra.get("think_mode", DEFAULT_THINK_MODE)),
|
||||
)
|
||||
|
||||
|
||||
class SenseNovaU1GenerationStage(PipelineStage):
|
||||
def __init__(self, model: torch.nn.Module, tokenizer: Any):
|
||||
super().__init__()
|
||||
self.model = model
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
@property
|
||||
def role_affinity(self) -> RoleType:
|
||||
return RoleType.DENOISER
|
||||
|
||||
def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch:
|
||||
del server_args
|
||||
options = SenseNovaU1GenerationOptions.from_batch(batch)
|
||||
if int(batch.num_outputs_per_prompt) != 1:
|
||||
raise ValueError(
|
||||
"SenseNova-U1 expects output expansion before generation; "
|
||||
f"got num_outputs_per_prompt={batch.num_outputs_per_prompt}."
|
||||
)
|
||||
seed = batch.seed[0] if isinstance(batch.seed, list) else int(batch.seed)
|
||||
|
||||
out = self.model.t2i_generate(
|
||||
self.tokenizer,
|
||||
batch.prompt,
|
||||
image_size=(int(batch.width), int(batch.height)),
|
||||
cfg_scale=float(batch.guidance_scale),
|
||||
cfg_norm=options.cfg_norm,
|
||||
timestep_shift=options.timestep_shift,
|
||||
enable_timestep_shift=options.enable_timestep_shift,
|
||||
cfg_interval=options.cfg_interval,
|
||||
num_steps=int(batch.num_inference_steps),
|
||||
batch_size=1,
|
||||
t_eps=options.t_eps,
|
||||
think_mode=options.think_mode,
|
||||
seed=seed,
|
||||
)
|
||||
think_text = None
|
||||
if options.think_mode:
|
||||
images, think_text = out
|
||||
else:
|
||||
images = out
|
||||
|
||||
images = _denorm_sensenova_output(images)
|
||||
samples = [sample.contiguous() for sample in images]
|
||||
usage = {"think_text": think_text} if think_text is not None else None
|
||||
return OutputBatch(
|
||||
output=samples,
|
||||
metrics=batch.metrics,
|
||||
usage=usage,
|
||||
)
|
||||
@@ -112,6 +112,27 @@ def is_ltx2_two_stage_pipeline_name(pipeline_class_name: str | None) -> bool:
|
||||
return pipeline_class_name in LTX2_TWO_STAGE_PIPELINE_NAMES
|
||||
|
||||
|
||||
def _infer_direct_constructor_explicit_arg_names(server_args) -> set[str]:
|
||||
explicit_arg_names: set[str] = set()
|
||||
for attr in dataclasses.fields(server_args):
|
||||
if not attr.init or attr.name == "_explicit_arg_names":
|
||||
continue
|
||||
|
||||
value = getattr(server_args, attr.name)
|
||||
if attr.default is not dataclasses.MISSING:
|
||||
default = attr.default
|
||||
elif attr.default_factory is not dataclasses.MISSING:
|
||||
default = attr.default_factory()
|
||||
else:
|
||||
explicit_arg_names.add(attr.name)
|
||||
continue
|
||||
|
||||
if value != default:
|
||||
explicit_arg_names.add(attr.name)
|
||||
|
||||
return explicit_arg_names
|
||||
|
||||
|
||||
def _normalize_component_precisions(value: object) -> dict[str, str]:
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError("component_precisions must be a mapping")
|
||||
@@ -1864,11 +1885,17 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
raise ValueError(f"Could not parse attention backend config: {config_str}")
|
||||
|
||||
def __post_init__(self):
|
||||
if not self._explicit_arg_names:
|
||||
self._explicit_arg_names = _infer_direct_constructor_explicit_arg_names(
|
||||
self
|
||||
)
|
||||
|
||||
# configure logger before use
|
||||
configure_logger(server_args=self)
|
||||
|
||||
component_paths: dict[str, str] = {}
|
||||
component_weights_paths = dict(self.component_weights_paths)
|
||||
migrated_component_weight_path = False
|
||||
for component, path in self.component_paths.items():
|
||||
if not is_explicit_weight_file_reference(path):
|
||||
component_paths[component] = path
|
||||
@@ -1880,6 +1907,11 @@ class ServerArgs(DisaggServerArgsMixin):
|
||||
f"{existing!r} and {path!r}"
|
||||
)
|
||||
component_weights_paths[component] = path
|
||||
migrated_component_weight_path = True
|
||||
if migrated_component_weight_path and self.is_arg_explicitly_set(
|
||||
"component_paths"
|
||||
):
|
||||
self._explicit_arg_names.add("component_weights_paths")
|
||||
self.component_paths = component_paths
|
||||
self.component_weights_paths = component_weights_paths
|
||||
self.component_precisions = _normalize_component_precisions(
|
||||
|
||||
@@ -10,6 +10,9 @@ from sglang.multimodal_gen.runtime.entrypoints.utils import (
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import (
|
||||
InputValidationStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import (
|
||||
LatentPreparationStage,
|
||||
)
|
||||
@@ -131,6 +134,55 @@ class TestMultiOutputGrouping(unittest.TestCase):
|
||||
["rid:0", "rid:1"],
|
||||
)
|
||||
|
||||
def test_sequential_stage_matches_entrypoint_output_expansion(self):
|
||||
def make_req():
|
||||
req = Req(
|
||||
sampling_params=SamplingParams(
|
||||
request_id="rid",
|
||||
prompt="p",
|
||||
output_path="/tmp",
|
||||
output_file_name="image.png",
|
||||
num_outputs_per_prompt=2,
|
||||
seed=[100, 101],
|
||||
)
|
||||
)
|
||||
return req
|
||||
|
||||
entrypoint_outputs = expand_request_outputs(make_req())
|
||||
sequential_parent = make_req()
|
||||
sequential_outputs = list(
|
||||
InputValidationStage().iter_sequential_requests(
|
||||
sequential_parent,
|
||||
SimpleNamespace(
|
||||
pipeline_config=SimpleNamespace(
|
||||
supports_sequential_multi_output_inference=lambda: True
|
||||
)
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
def expansion_signature(req):
|
||||
return (
|
||||
req.request_id,
|
||||
req.seed,
|
||||
req.num_outputs_per_prompt,
|
||||
req.output_file_name,
|
||||
req.extra["parent_request_id"],
|
||||
req.extra["output_index"],
|
||||
req.metrics.request_id,
|
||||
)
|
||||
|
||||
self.assertEqual(
|
||||
[expansion_signature(req) for req in sequential_outputs],
|
||||
[expansion_signature(req) for req in entrypoint_outputs],
|
||||
)
|
||||
self.assertTrue(
|
||||
all(
|
||||
req.trace_ctx is sequential_parent.trace_ctx
|
||||
for req in sequential_outputs
|
||||
)
|
||||
)
|
||||
|
||||
def test_split_batched_latents_uses_original_batched_tensor(self):
|
||||
stage = LatentPreparationStage.__new__(LatentPreparationStage)
|
||||
src = Req(sampling_params=SamplingParams(prompt="p"))
|
||||
|
||||
@@ -0,0 +1,851 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import asyncio
|
||||
import json
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.sensenova_u1 import (
|
||||
SenseNovaU1PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sample.sampling_params import DataType
|
||||
from sglang.multimodal_gen.configs.sample.sensenova_u1 import (
|
||||
SenseNovaU1SamplingParams,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.sensenova_u1 import (
|
||||
SENSENOVA_U1_REQUEST_EXTRA_KEY,
|
||||
)
|
||||
from sglang.multimodal_gen.registry import (
|
||||
_get_config_info,
|
||||
get_model_info,
|
||||
get_non_diffusers_pipeline_name,
|
||||
is_registered_diffusion_model_path,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.entrypoints.openai.utils import (
|
||||
process_generation_batch,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.managers.gpu_worker import GPUWorker
|
||||
from sglang.multimodal_gen.runtime.models.sensenova_u1.neo_unify.configuration_neo_vit import (
|
||||
NEOVisionConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.sensenova_u1.neo_unify.conversation import (
|
||||
get_conv_template,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.models.sensenova_u1.neo_unify.modeling_neo_chat import (
|
||||
_randn_with_seed,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor import (
|
||||
PipelineExecutor,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import (
|
||||
InputValidationStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sensenova_u1 import (
|
||||
SenseNovaU1GenerationStage,
|
||||
)
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
from sglang.multimodal_gen.runtime.server_args.server_args import ServerArgs
|
||||
from sglang.multimodal_gen.runtime.utils.perf_logger import MemorySnapshot
|
||||
|
||||
|
||||
class _FakeSenseNovaModel:
|
||||
def __init__(self):
|
||||
self.call_kwargs = None
|
||||
|
||||
def t2i_generate(self, tokenizer, prompt, **kwargs):
|
||||
self.call_kwargs = {"tokenizer": tokenizer, "prompt": prompt, **kwargs}
|
||||
return torch.tensor(
|
||||
[
|
||||
[
|
||||
[[-1.0, 0.0], [0.5, 1.0]],
|
||||
[[-1.0, 0.0], [0.5, 1.0]],
|
||||
[[-1.0, 0.0], [0.5, 1.0]],
|
||||
]
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class _RecordingTraceContext:
|
||||
tracing_enable = True
|
||||
|
||||
def __init__(self):
|
||||
self.finish_count = 0
|
||||
self.started_slices = []
|
||||
self.finished_slices = []
|
||||
|
||||
def trace_req_finish(self):
|
||||
self.finish_count += 1
|
||||
|
||||
def trace_slice_start(self, name, level=0):
|
||||
self.started_slices.append((name, level))
|
||||
|
||||
def trace_slice_end(self, name, level=0, **kwargs):
|
||||
self.finished_slices.append((name, level))
|
||||
|
||||
|
||||
class _SequentialTestExecutor(PipelineExecutor):
|
||||
def __init__(self, server_args, *, fail=False, fail_request_ids=None):
|
||||
super().__init__(server_args)
|
||||
self.fail = fail
|
||||
self.fail_request_ids = set(fail_request_ids or [])
|
||||
self.executed_requests = []
|
||||
|
||||
def execute_group(self, stages, batches, server_args):
|
||||
for batch in batches:
|
||||
batch.metrics.record_stage("InputValidationStage", 0.125)
|
||||
batch.metrics.record_memory_snapshot(
|
||||
"after_validation",
|
||||
MemorySnapshot(
|
||||
allocated_mb=100.0,
|
||||
reserved_mb=200.0,
|
||||
peak_allocated_mb=300.0,
|
||||
peak_reserved_mb=400.0,
|
||||
),
|
||||
)
|
||||
return batches
|
||||
|
||||
def execute(self, stages, batch, server_args):
|
||||
self.executed_requests.append(batch)
|
||||
if self.fail or batch.request_id in self.fail_request_ids:
|
||||
raise RuntimeError(f"generation failed for {batch.request_id}")
|
||||
return OutputBatch(
|
||||
output_file_paths=[batch.output_file_name],
|
||||
metrics=batch.metrics,
|
||||
)
|
||||
|
||||
|
||||
class _SequentialTestPipeline:
|
||||
def __init__(self, server_args, *, fail=False, fail_request_ids=None):
|
||||
self.input_stage = InputValidationStage()
|
||||
self.executor = _SequentialTestExecutor(
|
||||
server_args,
|
||||
fail=fail,
|
||||
fail_request_ids=fail_request_ids,
|
||||
)
|
||||
|
||||
def forward_batch_sequentially(self, batches, server_args):
|
||||
return self.executor.execute_group_sequentially(
|
||||
[self.input_stage, object()], batches, server_args
|
||||
)
|
||||
|
||||
|
||||
class _WorkerBackedSchedulerClient:
|
||||
def __init__(self, worker):
|
||||
self.worker = worker
|
||||
|
||||
async def forward(self, batches):
|
||||
return next(self.worker.execute_forward_sequentially(batches))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("template_name", "expected_system_message"),
|
||||
[
|
||||
(
|
||||
"Hermes-2",
|
||||
"\u4f60\u662f\u7531\u4e0a\u6d77\u4eba\u5de5\u667a\u80fd\u5b9e\u9a8c\u5ba4\u8054\u5408\u5546\u6c64\u79d1\u6280\u5f00\u53d1\u7684\u4e66\u751f\u591a\u6a21\u6001\u5927\u6a21\u578b\uff0c\u82f1\u6587\u540d\u53ebInternVL, \u662f\u4e00\u4e2a\u6709\u7528\u65e0\u5bb3\u7684\u4eba\u5de5\u667a\u80fd\u52a9\u624b\u3002",
|
||||
),
|
||||
(
|
||||
"internlm2-chat",
|
||||
"\u4f60\u662f\u7531\u4e0a\u6d77\u4eba\u5de5\u667a\u80fd\u5b9e\u9a8c\u5ba4\u8054\u5408\u5546\u6c64\u79d1\u6280\u5f00\u53d1\u7684\u4e66\u751f\u591a\u6a21\u6001\u5927\u6a21\u578b\uff0c\u82f1\u6587\u540d\u53ebInternVL, \u662f\u4e00\u4e2a\u6709\u7528\u65e0\u5bb3\u7684\u4eba\u5de5\u667a\u80fd\u52a9\u624b\u3002",
|
||||
),
|
||||
(
|
||||
"phi3-chat",
|
||||
"\u4f60\u662f\u7531\u4e0a\u6d77\u4eba\u5de5\u667a\u80fd\u5b9e\u9a8c\u5ba4\u8054\u5408\u5546\u6c64\u79d1\u6280\u5f00\u53d1\u7684\u4e66\u751f\u591a\u6a21\u6001\u5927\u6a21\u578b\uff0c\u82f1\u6587\u540d\u53ebInternVL, \u662f\u4e00\u4e2a\u6709\u7528\u65e0\u5bb3\u7684\u4eba\u5de5\u667a\u80fd\u52a9\u624b\u3002",
|
||||
),
|
||||
(
|
||||
"internvl2_5",
|
||||
"\u4f60\u662f\u4e66\u751f\xb7\u4e07\u8c61\uff0c\u82f1\u6587\u540d\u662fInternVL\uff0c\u662f\u7531\u4e0a\u6d77\u4eba\u5de5\u667a\u80fd\u5b9e\u9a8c\u5ba4\u3001\u6e05\u534e\u5927\u5b66\u53ca\u591a\u5bb6\u5408\u4f5c\u5355\u4f4d\u8054\u5408\u5f00\u53d1\u7684\u591a\u6a21\u6001\u5927\u8bed\u8a00\u6a21\u578b\u3002",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_sensenova_u1_conversation_preserves_upstream_system_prompt(
|
||||
template_name, expected_system_message
|
||||
):
|
||||
assert get_conv_template(template_name).system_message == expected_system_message
|
||||
|
||||
|
||||
def _force_generator_fallback(monkeypatch, device_type):
|
||||
original_generator = torch.Generator
|
||||
|
||||
def unsupported_device_generator(device="cpu"):
|
||||
if torch.device(device).type == device_type:
|
||||
raise RuntimeError(f"Generator is unsupported on {device_type}")
|
||||
return original_generator(device)
|
||||
|
||||
monkeypatch.setattr(torch, "Generator", unsupported_device_generator)
|
||||
|
||||
|
||||
def test_sensenova_u1_randn_fallback_preserves_cpu_rng(monkeypatch):
|
||||
_force_generator_fallback(monkeypatch, "cpu")
|
||||
rng_state = torch.get_rng_state().clone()
|
||||
|
||||
first = _randn_with_seed((2, 3), device="cpu", dtype=torch.float32, seed=17)
|
||||
second = _randn_with_seed((2, 3), device="cpu", dtype=torch.float32, seed=17)
|
||||
|
||||
assert torch.equal(first, second)
|
||||
assert torch.equal(torch.get_rng_state(), rng_state)
|
||||
|
||||
|
||||
def test_sensenova_u1_randn_fallback_preserves_device_rng(monkeypatch):
|
||||
device_type = current_platform.device_type
|
||||
if not device_type or device_type == "cpu":
|
||||
pytest.skip("No accelerator is available")
|
||||
|
||||
device = torch.device(device_type, 0)
|
||||
device_module = torch.get_device_module(device)
|
||||
if not device_module.is_available():
|
||||
pytest.skip(f"{device_type} is not available")
|
||||
|
||||
_force_generator_fallback(monkeypatch, device_type)
|
||||
cpu_rng_state = torch.get_rng_state().clone()
|
||||
device_rng_state = device_module.get_rng_state(device).clone()
|
||||
|
||||
first = _randn_with_seed((2, 3), device=device, dtype=torch.float32, seed=17)
|
||||
second = _randn_with_seed((2, 3), device=device, dtype=torch.float32, seed=17)
|
||||
|
||||
assert torch.equal(first, second)
|
||||
assert torch.equal(torch.get_rng_state(), cpu_rng_state)
|
||||
assert torch.equal(device_module.get_rng_state(device), device_rng_state)
|
||||
|
||||
|
||||
def test_sensenova_u1_registry_resolves_local_and_hf_paths(tmp_path):
|
||||
_get_config_info.cache_clear()
|
||||
get_model_info.cache_clear()
|
||||
|
||||
local_path = tmp_path / "checkpoint-revision-abc123"
|
||||
local_path.mkdir()
|
||||
(local_path / "config.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
"architectures": ["NEOChatModel"],
|
||||
"model_type": "neo_chat",
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
assert is_registered_diffusion_model_path(str(local_path))
|
||||
assert get_non_diffusers_pipeline_name(str(local_path)) == "SenseNovaU1Pipeline"
|
||||
|
||||
local_model_info = get_model_info(str(local_path))
|
||||
assert local_model_info is not None
|
||||
assert local_model_info.pipeline_config_cls is SenseNovaU1PipelineConfig
|
||||
assert local_model_info.sampling_param_cls is SenseNovaU1SamplingParams
|
||||
|
||||
model_info = get_model_info("sensenova/SenseNova-U1.5-8B-MoT")
|
||||
assert model_info is not None
|
||||
assert model_info.pipeline_config_cls is SenseNovaU1PipelineConfig
|
||||
assert model_info.sampling_param_cls is SenseNovaU1SamplingParams
|
||||
|
||||
modelscope_id = "SenseNova/SenseNova-U1.5-8B-MoT"
|
||||
assert is_registered_diffusion_model_path(modelscope_id)
|
||||
assert get_non_diffusers_pipeline_name(modelscope_id) == "SenseNovaU1Pipeline"
|
||||
assert get_model_info(modelscope_id) is not None
|
||||
get_model_info.cache_clear()
|
||||
|
||||
|
||||
def test_sensenova_u1_registry_requires_exact_hub_id(monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
"sglang.multimodal_gen.registry.maybe_download_model_index",
|
||||
lambda _: {},
|
||||
)
|
||||
_get_config_info.cache_clear()
|
||||
get_model_info.cache_clear()
|
||||
|
||||
unrelated_repo = "acme/SenseNova-U1.5-8B-MoT"
|
||||
assert not is_registered_diffusion_model_path(unrelated_repo)
|
||||
assert get_non_diffusers_pipeline_name(unrelated_repo) is None
|
||||
assert _get_config_info(unrelated_repo) is None
|
||||
|
||||
_get_config_info.cache_clear()
|
||||
get_model_info.cache_clear()
|
||||
|
||||
|
||||
def test_sensenova_u1_registry_does_not_route_lora_only_repositories(tmp_path):
|
||||
lora_repo = "sensenova/SenseNova-U1.5-8B-MoT-LoRA"
|
||||
lora_path = tmp_path / "SenseNova-U1.5-8B-MoT-LoRA"
|
||||
lora_path.mkdir()
|
||||
(lora_path / "adapter_config.json").write_text("{}")
|
||||
|
||||
assert get_non_diffusers_pipeline_name(lora_repo) is None
|
||||
assert get_non_diffusers_pipeline_name(str(lora_path)) is None
|
||||
assert not is_registered_diffusion_model_path(lora_repo)
|
||||
assert not is_registered_diffusion_model_path(str(lora_path))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("backend", ["auto", "sglang", "diffusers"])
|
||||
def test_sensenova_u1_known_adapter_only_repo_rejected_before_backend_resolution(
|
||||
monkeypatch, backend
|
||||
):
|
||||
def fail_model_index_download(_):
|
||||
raise AssertionError("adapter-only repo should not download model_index")
|
||||
|
||||
def fail_diffusers_resolution(**_kwargs):
|
||||
raise AssertionError("adapter-only repo should not resolve diffusers info")
|
||||
|
||||
monkeypatch.setattr(
|
||||
"sglang.multimodal_gen.registry.maybe_download_model_index",
|
||||
fail_model_index_download,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"sglang.multimodal_gen.registry._get_diffusers_model_info",
|
||||
fail_diffusers_resolution,
|
||||
)
|
||||
get_model_info.cache_clear()
|
||||
|
||||
loras_repo = "sensenova/SenseNova-U1.5-8B-MoT-LoRAs"
|
||||
|
||||
assert get_non_diffusers_pipeline_name(loras_repo) is None
|
||||
assert get_model_info(loras_repo, backend=backend) is None
|
||||
get_model_info.cache_clear()
|
||||
|
||||
|
||||
def test_sensenova_u1_sampling_params_keep_private_defaults_internal():
|
||||
params = SenseNovaU1SamplingParams(prompt="hello", width=2304, height=4096)
|
||||
|
||||
assert params.guidance_scale == 4.0
|
||||
assert params.num_inference_steps == 50
|
||||
assert params.num_outputs_per_prompt == 1
|
||||
assert params.cfg_norm == "none"
|
||||
assert params.timestep_shift == 3.0
|
||||
|
||||
extra = params.build_request_extra()[SENSENOVA_U1_REQUEST_EXTRA_KEY]
|
||||
assert extra == {
|
||||
"cfg_norm": "none",
|
||||
"timestep_shift": 3.0,
|
||||
"enable_timestep_shift": True,
|
||||
"cfg_interval": (0.0, 1.0),
|
||||
"t_eps": 0.02,
|
||||
"think_mode": False,
|
||||
}
|
||||
|
||||
|
||||
def test_sensenova_u1_rejects_unaligned_resolution():
|
||||
with pytest.raises(ValueError, match="divisible by 32"):
|
||||
SenseNovaU1SamplingParams(width=2160, height=3840)
|
||||
|
||||
|
||||
def test_sensenova_u1_accepts_openai_image_api_num_frames():
|
||||
params = SenseNovaU1SamplingParams(
|
||||
prompt="hello",
|
||||
width=2048,
|
||||
height=2048,
|
||||
num_frames=1,
|
||||
)
|
||||
|
||||
assert params.num_frames == 1
|
||||
assert params.data_type == DataType.IMAGE
|
||||
|
||||
|
||||
def test_sensenova_u1_scheduler_capabilities():
|
||||
config = SenseNovaU1PipelineConfig()
|
||||
|
||||
assert not config.supports_dynamic_batching()
|
||||
assert config.supports_sequential_multi_output_inference()
|
||||
|
||||
|
||||
def test_sensenova_u1_rejects_multi_gpu_during_arg_validation():
|
||||
config = SenseNovaU1PipelineConfig()
|
||||
|
||||
with pytest.raises(ValueError, match="num_gpus=1"):
|
||||
config.validate_server_args(
|
||||
SimpleNamespace(
|
||||
num_gpus=2,
|
||||
enable_torch_compile=False,
|
||||
lora_path=None,
|
||||
attention_backend=None,
|
||||
component_attention_backends={},
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_sensenova_u1_clears_auto_tuned_runtime_defaults():
|
||||
config = SenseNovaU1PipelineConfig()
|
||||
args = SimpleNamespace(
|
||||
num_gpus=1,
|
||||
enable_torch_compile=False,
|
||||
lora_path=None,
|
||||
component_residency={"transformer": "layerwise-offload"},
|
||||
cpu_offload_components=["transformer"],
|
||||
dit_cpu_offload=True,
|
||||
text_encoder_cpu_offload=True,
|
||||
image_encoder_cpu_offload=True,
|
||||
vae_cpu_offload=True,
|
||||
dit_layerwise_offload=True,
|
||||
layerwise_offload_components=["transformer"],
|
||||
quantization=None,
|
||||
quantization_ignored_layers=None,
|
||||
transformer_weights_path=None,
|
||||
component_paths={"model": "/tmp/component"},
|
||||
component_weights_paths={"model": "/tmp/model.safetensors"},
|
||||
component_quantizations={},
|
||||
component_quantization_ignored_layers={},
|
||||
component_precisions={},
|
||||
attention_backend="aiter",
|
||||
component_attention_backends={"text_encoder": "torch_sdpa"},
|
||||
attention_backend_config={"foo": "bar"},
|
||||
is_arg_explicitly_set=lambda _name: False,
|
||||
)
|
||||
|
||||
config.validate_server_args(args)
|
||||
|
||||
assert args.component_residency is None
|
||||
assert args.cpu_offload_components is None
|
||||
assert args.dit_cpu_offload is False
|
||||
assert args.text_encoder_cpu_offload is False
|
||||
assert args.image_encoder_cpu_offload is False
|
||||
assert args.vae_cpu_offload is False
|
||||
assert args.dit_layerwise_offload is False
|
||||
assert args.layerwise_offload_components is None
|
||||
assert args.component_paths == {}
|
||||
assert args.component_weights_paths == {}
|
||||
assert args.attention_backend is None
|
||||
assert args.component_attention_backends == {}
|
||||
assert args.attention_backend_config is None
|
||||
|
||||
|
||||
def test_sensenova_u1_allows_explicit_resident_component_residency():
|
||||
config = SenseNovaU1PipelineConfig()
|
||||
args = SimpleNamespace(
|
||||
num_gpus=1,
|
||||
enable_torch_compile=False,
|
||||
lora_path=None,
|
||||
component_residency={"transformer": "resident"},
|
||||
cpu_offload_components=None,
|
||||
dit_cpu_offload=False,
|
||||
text_encoder_cpu_offload=False,
|
||||
image_encoder_cpu_offload=False,
|
||||
vae_cpu_offload=False,
|
||||
dit_layerwise_offload=False,
|
||||
layerwise_offload_components=None,
|
||||
quantization=None,
|
||||
quantization_ignored_layers=None,
|
||||
transformer_weights_path=None,
|
||||
component_paths={},
|
||||
component_weights_paths={},
|
||||
component_quantizations={},
|
||||
component_quantization_ignored_layers={},
|
||||
component_precisions={},
|
||||
attention_backend=None,
|
||||
component_attention_backends={},
|
||||
attention_backend_config={},
|
||||
is_arg_explicitly_set=lambda name: name == "component_residency",
|
||||
)
|
||||
|
||||
config.validate_server_args(args)
|
||||
|
||||
assert args.component_residency == {"transformer": "resident"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("override", "expected"),
|
||||
[
|
||||
({"enable_torch_compile": True}, "torch.compile"),
|
||||
({"lora_path": "sensenova/SenseNova-U1.5-8B-MoT-LoRAs"}, "LoRA adapters"),
|
||||
(
|
||||
{"component_residency": {"transformer": "component-offload"}},
|
||||
"component residency offload",
|
||||
),
|
||||
({"cpu_offload_components": ["transformer"]}, "CPU offload"),
|
||||
({"dit_cpu_offload": True}, "DiT CPU offload"),
|
||||
({"text_encoder_cpu_offload": True}, "text encoder CPU offload"),
|
||||
({"image_encoder_cpu_offload": True}, "image encoder CPU offload"),
|
||||
({"vae_cpu_offload": True}, "VAE CPU offload"),
|
||||
({"dit_layerwise_offload": True}, "DiT layerwise offload"),
|
||||
({"layerwise_offload_components": ["transformer"]}, "layerwise offload"),
|
||||
({"quantization": "fp8"}, "quantization"),
|
||||
({"quantization_ignored_layers": ["foo"]}, "quantization ignored layers"),
|
||||
(
|
||||
{"transformer_weights_path": "/tmp/transformer.safetensors"},
|
||||
"pre-quantized transformer weights",
|
||||
),
|
||||
({"component_paths": {"model": "/tmp/component"}}, "component path overrides"),
|
||||
(
|
||||
{"component_weights_paths": {"model": "/tmp/model.safetensors"}},
|
||||
"component weight path overrides",
|
||||
),
|
||||
({"component_quantizations": {"transformer": "fp8"}}, "component quantization"),
|
||||
(
|
||||
{"component_quantization_ignored_layers": {"transformer": ["foo"]}},
|
||||
"component quantization ignored layers",
|
||||
),
|
||||
({"component_precisions": {"transformer": "fp16"}}, "component precision"),
|
||||
({"attention_backend": "fa"}, "custom attention backends"),
|
||||
(
|
||||
{"component_attention_backends": {"text_encoder": "torch_sdpa"}},
|
||||
"component attention backends",
|
||||
),
|
||||
({"attention_backend_config": {"foo": "bar"}}, "attention backend config"),
|
||||
],
|
||||
)
|
||||
def test_sensenova_u1_rejects_unsupported_runtime_modes(override, expected):
|
||||
config = SenseNovaU1PipelineConfig()
|
||||
args = {
|
||||
"num_gpus": 1,
|
||||
"enable_torch_compile": False,
|
||||
"lora_path": None,
|
||||
"component_residency": None,
|
||||
"cpu_offload_components": None,
|
||||
"dit_cpu_offload": None,
|
||||
"text_encoder_cpu_offload": None,
|
||||
"image_encoder_cpu_offload": None,
|
||||
"vae_cpu_offload": False,
|
||||
"dit_layerwise_offload": None,
|
||||
"layerwise_offload_components": None,
|
||||
"quantization": None,
|
||||
"quantization_ignored_layers": None,
|
||||
"transformer_weights_path": None,
|
||||
"component_paths": {},
|
||||
"component_weights_paths": {},
|
||||
"component_quantizations": {},
|
||||
"component_quantization_ignored_layers": {},
|
||||
"component_precisions": {},
|
||||
"attention_backend": None,
|
||||
"component_attention_backends": {},
|
||||
"attention_backend_config": {},
|
||||
}
|
||||
args.update(override)
|
||||
|
||||
with pytest.raises(ValueError, match=expected):
|
||||
config.validate_server_args(SimpleNamespace(**args))
|
||||
|
||||
|
||||
def test_sensenova_u1_rejects_direct_server_args_quantization():
|
||||
config = SenseNovaU1PipelineConfig()
|
||||
|
||||
with pytest.raises(ValueError, match="quantization"):
|
||||
ServerArgs(
|
||||
model_path="sensenova/SenseNova-U1.5-8B-MoT",
|
||||
pipeline_config=config,
|
||||
quantization="fp8",
|
||||
)
|
||||
|
||||
|
||||
def test_sensenova_u1_rejects_file_valued_component_paths(tmp_path):
|
||||
config = SenseNovaU1PipelineConfig()
|
||||
|
||||
with pytest.raises(ValueError, match="component weight path overrides"):
|
||||
ServerArgs(
|
||||
model_path="sensenova/SenseNova-U1.5-8B-MoT",
|
||||
pipeline_config=config,
|
||||
component_paths={"model": str(tmp_path / "model.safetensors")},
|
||||
)
|
||||
|
||||
|
||||
def test_sensenova_u1_vision_config_round_trips_sequence_fields(tmp_path):
|
||||
config = NEOVisionConfig(llm_hidden_size=2048, downsample_ratio=0.5)
|
||||
config.save_pretrained(tmp_path)
|
||||
|
||||
loaded = NEOVisionConfig.from_pretrained(tmp_path)
|
||||
|
||||
assert loaded.llm_hidden_size == (2048,)
|
||||
assert loaded.downsample_ratio == (0.5,)
|
||||
|
||||
|
||||
def test_sensenova_u1_vision_config_normalizes_nested_singletons():
|
||||
config = NEOVisionConfig(llm_hidden_size=[[2048]], downsample_ratio=[[0.5]])
|
||||
|
||||
assert config.llm_hidden_size == (2048,)
|
||||
assert config.downsample_ratio == (0.5,)
|
||||
|
||||
|
||||
def test_sensenova_u1_rejects_video_frame_count():
|
||||
with pytest.raises(ValueError, match="num_frames=1"):
|
||||
SenseNovaU1SamplingParams(width=2048, height=2048, num_frames=2)
|
||||
|
||||
|
||||
def test_sensenova_u1_cli_args_expose_only_sglang_compatible_fields():
|
||||
args = SimpleNamespace(
|
||||
prompt="hello",
|
||||
width=2304,
|
||||
height=4096,
|
||||
guidance_scale=4.5,
|
||||
num_inference_steps=30,
|
||||
num_outputs_per_prompt=2,
|
||||
cfg_norm="global",
|
||||
timestep_shift=9.0,
|
||||
think_mode=True,
|
||||
)
|
||||
|
||||
cli_args = SenseNovaU1SamplingParams.get_cli_args(args)
|
||||
|
||||
assert cli_args["prompt"] == "hello"
|
||||
assert cli_args["width"] == 2304
|
||||
assert cli_args["height"] == 4096
|
||||
assert cli_args["guidance_scale"] == 4.5
|
||||
assert cli_args["num_inference_steps"] == 30
|
||||
assert cli_args["num_outputs_per_prompt"] == 2
|
||||
assert "cfg_norm" not in cli_args
|
||||
assert "timestep_shift" not in cli_args
|
||||
assert "think_mode" not in cli_args
|
||||
|
||||
|
||||
def test_sensenova_u1_generation_stage_uses_sglang_params_and_single_model_batch():
|
||||
sampling = SenseNovaU1SamplingParams(
|
||||
prompt="a mountain lake",
|
||||
width=2304,
|
||||
height=4096,
|
||||
guidance_scale=4.5,
|
||||
num_inference_steps=30,
|
||||
seed=123,
|
||||
)
|
||||
batch = SimpleNamespace(
|
||||
prompt=sampling.prompt,
|
||||
width=sampling.width,
|
||||
height=sampling.height,
|
||||
guidance_scale=sampling.guidance_scale,
|
||||
num_inference_steps=sampling.num_inference_steps,
|
||||
seed=sampling.seed,
|
||||
num_outputs_per_prompt=sampling.num_outputs_per_prompt,
|
||||
extra=sampling.build_request_extra(),
|
||||
metrics=None,
|
||||
)
|
||||
model = _FakeSenseNovaModel()
|
||||
stage = SenseNovaU1GenerationStage(model=model, tokenizer="tok")
|
||||
|
||||
output = stage.forward(batch, server_args=SimpleNamespace())
|
||||
|
||||
assert len(output.output) == 1
|
||||
assert torch.allclose(
|
||||
output.output[0],
|
||||
torch.tensor(
|
||||
[
|
||||
[[0.0, 0.5], [0.75, 1.0]],
|
||||
[[0.0, 0.5], [0.75, 1.0]],
|
||||
[[0.0, 0.5], [0.75, 1.0]],
|
||||
]
|
||||
),
|
||||
)
|
||||
assert model.call_kwargs["tokenizer"] == "tok"
|
||||
assert model.call_kwargs["prompt"] == "a mountain lake"
|
||||
assert model.call_kwargs["image_size"] == (2304, 4096)
|
||||
assert model.call_kwargs["cfg_scale"] == 4.5
|
||||
assert model.call_kwargs["num_steps"] == 30
|
||||
assert model.call_kwargs["batch_size"] == 1
|
||||
assert model.call_kwargs["seed"] == 123
|
||||
|
||||
|
||||
def test_sensenova_u1_multi_output_request_expands_before_generation_stage():
|
||||
sampling = SenseNovaU1SamplingParams(
|
||||
prompt="a mountain lake",
|
||||
width=2304,
|
||||
height=4096,
|
||||
num_outputs_per_prompt=2,
|
||||
)
|
||||
batch = Req(
|
||||
request_id="req-0",
|
||||
prompt=sampling.prompt,
|
||||
width=sampling.width,
|
||||
height=sampling.height,
|
||||
guidance_scale=sampling.guidance_scale,
|
||||
num_inference_steps=sampling.num_inference_steps,
|
||||
seed=42,
|
||||
sampling_params=sampling,
|
||||
extra=sampling.build_request_extra(),
|
||||
output_file_name="sample.png",
|
||||
)
|
||||
server_args = SimpleNamespace(pipeline_config=SenseNovaU1PipelineConfig())
|
||||
input_stage = InputValidationStage()
|
||||
stage = SenseNovaU1GenerationStage(model=_FakeSenseNovaModel(), tokenizer="tok")
|
||||
batch.metrics.record_stage("InputValidationStage", 0.125)
|
||||
batch.metrics.record_memory_snapshot(
|
||||
"after_validation",
|
||||
MemorySnapshot(
|
||||
allocated_mb=100.0,
|
||||
reserved_mb=200.0,
|
||||
peak_allocated_mb=300.0,
|
||||
peak_reserved_mb=400.0,
|
||||
),
|
||||
)
|
||||
|
||||
expanded = list(input_stage.iter_sequential_requests(batch, server_args))
|
||||
|
||||
assert [req.num_outputs_per_prompt for req in expanded] == [1, 1]
|
||||
assert [req.seed for req in expanded] == [42, 43]
|
||||
assert [req.request_id for req in expanded] == ["req-0:0", "req-0:1"]
|
||||
assert [req.output_file_name for req in expanded] == [
|
||||
"sample_0.png",
|
||||
"sample_1.png",
|
||||
]
|
||||
assert [req.metrics.request_id for req in expanded] == ["req-0:0", "req-0:1"]
|
||||
assert all(req.trace_ctx is batch.trace_ctx for req in expanded)
|
||||
assert all(req.metrics is not batch.metrics for req in expanded)
|
||||
assert expanded[0].metrics is not expanded[1].metrics
|
||||
assert all(
|
||||
req.metrics.stages == {"InputValidationStage": 125.0} for req in expanded
|
||||
)
|
||||
assert all(
|
||||
req.metrics.memory_snapshots["after_validation"].peak_reserved_mb == 400.0
|
||||
for req in expanded
|
||||
)
|
||||
assert (
|
||||
expanded[0].metrics.memory_snapshots["after_validation"]
|
||||
is not expanded[1].metrics.memory_snapshots["after_validation"]
|
||||
)
|
||||
|
||||
expanded[0].metrics.record_stage("child-only", 0.5)
|
||||
expanded[0].metrics.memory_snapshots["after_validation"].peak_reserved_mb = 999.0
|
||||
assert "child-only" not in expanded[1].metrics.stages
|
||||
assert "child-only" not in batch.metrics.stages
|
||||
assert (
|
||||
expanded[1].metrics.memory_snapshots["after_validation"].peak_reserved_mb
|
||||
== 400.0
|
||||
)
|
||||
assert batch.metrics.memory_snapshots["after_validation"].peak_reserved_mb == 400.0
|
||||
|
||||
for req in expanded:
|
||||
output = stage.forward(req, server_args=SimpleNamespace())
|
||||
assert len(output.output) == 1
|
||||
|
||||
|
||||
def test_sensenova_u1_multi_output_rejects_short_seed_list():
|
||||
sampling = SenseNovaU1SamplingParams(
|
||||
prompt="a mountain lake",
|
||||
width=2304,
|
||||
height=4096,
|
||||
num_outputs_per_prompt=2,
|
||||
seed=[7],
|
||||
)
|
||||
batch = Req(
|
||||
request_id="req-0",
|
||||
prompt=sampling.prompt,
|
||||
width=sampling.width,
|
||||
height=sampling.height,
|
||||
guidance_scale=sampling.guidance_scale,
|
||||
num_inference_steps=sampling.num_inference_steps,
|
||||
seed=sampling.seed,
|
||||
sampling_params=sampling,
|
||||
extra=sampling.build_request_extra(),
|
||||
output_file_name="sample.png",
|
||||
)
|
||||
server_args = SimpleNamespace(pipeline_config=SenseNovaU1PipelineConfig())
|
||||
|
||||
with pytest.raises(ValueError, match="seed list length"):
|
||||
list(InputValidationStage().iter_sequential_requests(batch, server_args))
|
||||
|
||||
|
||||
def _make_sensenova_u1_sequential_entrypoint(*, fail=False, fail_request_ids=None):
|
||||
sampling = SenseNovaU1SamplingParams(
|
||||
prompt="a mountain lake",
|
||||
width=2304,
|
||||
height=4096,
|
||||
num_outputs_per_prompt=2,
|
||||
save_output=False,
|
||||
suppress_logs=True,
|
||||
)
|
||||
trace_ctx = _RecordingTraceContext()
|
||||
batch = Req(
|
||||
request_id="req-0",
|
||||
prompt=sampling.prompt,
|
||||
width=sampling.width,
|
||||
height=sampling.height,
|
||||
seed=42,
|
||||
sampling_params=sampling,
|
||||
extra=sampling.build_request_extra(),
|
||||
output_file_name="sample.png",
|
||||
trace_ctx=trace_ctx,
|
||||
)
|
||||
server_args = SimpleNamespace(pipeline_config=SenseNovaU1PipelineConfig())
|
||||
pipeline = _SequentialTestPipeline(
|
||||
server_args, fail=fail, fail_request_ids=fail_request_ids
|
||||
)
|
||||
worker = GPUWorker.__new__(GPUWorker)
|
||||
worker.pipeline = pipeline
|
||||
worker.server_args = server_args
|
||||
worker.is_output_rank = True
|
||||
worker._runtime_peak_reserved_mb = 0.0
|
||||
worker._release_warmup_pool_before_serving = False
|
||||
worker._realtime_sessions = SimpleNamespace(attach=lambda _req: None)
|
||||
return batch, trace_ctx, pipeline.executor, _WorkerBackedSchedulerClient(worker)
|
||||
|
||||
|
||||
def _force_cpu_entrypoint(monkeypatch):
|
||||
monkeypatch.setattr(current_platform, "is_cpu", lambda: True)
|
||||
monkeypatch.setattr(current_platform, "is_mps", lambda: False)
|
||||
monkeypatch.setattr(current_platform, "is_npu", lambda: False)
|
||||
monkeypatch.setattr(
|
||||
"sglang.multimodal_gen.runtime.entrypoints.openai.utils.get_global_server_args",
|
||||
lambda: SimpleNamespace(batching_max_size=1),
|
||||
)
|
||||
|
||||
|
||||
def test_sensenova_u1_multi_output_entrypoint_success(monkeypatch):
|
||||
_force_cpu_entrypoint(monkeypatch)
|
||||
batch, trace_ctx, executor, scheduler_client = (
|
||||
_make_sensenova_u1_sequential_entrypoint()
|
||||
)
|
||||
|
||||
paths, result = asyncio.run(process_generation_batch(scheduler_client, batch))
|
||||
|
||||
assert paths == ["sample_0.png", "sample_1.png"]
|
||||
assert result.error is None
|
||||
assert [req.request_id for req in executor.executed_requests] == [
|
||||
"req-0:0",
|
||||
"req-0:1",
|
||||
]
|
||||
assert [req.seed for req in executor.executed_requests] == [42, 43]
|
||||
assert result.metrics_list is not None
|
||||
assert [metrics.request_id for metrics in result.metrics_list] == [
|
||||
"req-0:0",
|
||||
"req-0:1",
|
||||
]
|
||||
assert all(
|
||||
"InputValidationStage" in metrics.stages
|
||||
and "PipelineExecutor.sequential_wait" in metrics.stages
|
||||
and metrics.memory_snapshots["after_validation"].peak_reserved_mb == 400.0
|
||||
for metrics in result.metrics_list
|
||||
)
|
||||
assert all(req.trace_ctx is trace_ctx for req in executor.executed_requests)
|
||||
assert trace_ctx.started_slices == [("gpu_forward", 2)]
|
||||
assert trace_ctx.finished_slices == [("gpu_forward", 2)]
|
||||
assert trace_ctx.finish_count == 1
|
||||
|
||||
|
||||
def test_sensenova_u1_multi_output_entrypoint_failure(monkeypatch):
|
||||
_force_cpu_entrypoint(monkeypatch)
|
||||
batch, trace_ctx, executor, scheduler_client = (
|
||||
_make_sensenova_u1_sequential_entrypoint(fail=True)
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="generation failed for req-0:0"):
|
||||
asyncio.run(process_generation_batch(scheduler_client, batch))
|
||||
|
||||
assert [req.request_id for req in executor.executed_requests] == [
|
||||
"req-0:0",
|
||||
"req-0:1",
|
||||
]
|
||||
assert all(
|
||||
"InputValidationStage" in req.metrics.stages
|
||||
and "PipelineExecutor.sequential_wait" in req.metrics.stages
|
||||
and req.metrics.memory_snapshots["after_validation"].peak_reserved_mb == 400.0
|
||||
for req in executor.executed_requests
|
||||
)
|
||||
assert all(req.trace_ctx is trace_ctx for req in executor.executed_requests)
|
||||
assert trace_ctx.started_slices == [("gpu_forward", 2)]
|
||||
assert trace_ctx.finished_slices == [("gpu_forward", 2)]
|
||||
assert trace_ctx.finish_count == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("failed_request_id", ["req-0:0", "req-0:1"])
|
||||
def test_sensenova_u1_multi_output_entrypoint_mixed_failure_fails_parent(
|
||||
monkeypatch, failed_request_id
|
||||
):
|
||||
_force_cpu_entrypoint(monkeypatch)
|
||||
batch, trace_ctx, executor, scheduler_client = (
|
||||
_make_sensenova_u1_sequential_entrypoint(fail_request_ids={failed_request_id})
|
||||
)
|
||||
|
||||
with pytest.raises(
|
||||
RuntimeError, match=f"generation failed for {failed_request_id}"
|
||||
):
|
||||
asyncio.run(process_generation_batch(scheduler_client, batch))
|
||||
|
||||
assert [req.request_id for req in executor.executed_requests] == [
|
||||
"req-0:0",
|
||||
"req-0:1",
|
||||
]
|
||||
assert all(req.trace_ctx is trace_ctx for req in executor.executed_requests)
|
||||
assert trace_ctx.started_slices == [("gpu_forward", 2)]
|
||||
assert trace_ctx.finished_slices == [("gpu_forward", 2)]
|
||||
assert trace_ctx.finish_count == 1
|
||||
Reference in New Issue
Block a user