[diffusion] chore: make torch.compile opt-in for speed mode (#34173)
This commit is contained in:
@@ -80,7 +80,7 @@ Use `sglang generate --help` and `sglang serve --help` for the full argument lis
|
||||
- `--lora-path {PATH}` and `--lora-nickname {NAME}`: load a LoRA adapter
|
||||
- `--lora-merge-mode {auto|merge|dynamic}`: choose how LoRA is applied. `auto` statically merges regular weights and uses dynamic LoRA for FSDP-sharded weights to avoid full-gather peaks.
|
||||
- `--num-gpus {N}`: number of GPUs to use
|
||||
- `--performance-mode {manual|auto|speed|memory}` / `--mode`: preset for latency/throughput and memory defaults. `auto` is the default and keeps safe offload defaults, using FSDP only for validated DiT-offload replacement paths; `speed` also enables `--enable-torch-compile` unless the model-specific deployment config opts out or you explicitly disable it. Use `manual` to keep performance-related server args under explicit user control. Explicit offload, FSDP, and parallelism flags take precedence in all modes.
|
||||
- `--performance-mode {manual|auto|speed|memory}` / `--mode`: preset for latency/throughput and memory defaults. `auto` is the default and keeps safe offload defaults, using FSDP only for validated DiT-offload replacement paths. `speed` keeps `torch.compile` disabled unless a model-specific deployment config opts in after validation; pass `--enable-torch-compile true` to enable it explicitly. Use `manual` to keep performance-related server args under explicit user control. Explicit offload, FSDP, and parallelism flags take precedence in all modes.
|
||||
- `--tp-size {N}`: tensor parallelism size. Depending on the pipeline, it can shard the DiT, one or more encoders, or both.
|
||||
- `--sp-degree {N}`: sequence parallelism size
|
||||
- `--dp-size {N}` (alias `--data-parallel-size`): number of data-parallel replicas. Each replica is a full copy of the engine on `num_gpus / N` GPUs with its own ingress; generation requests round-robin across replicas, realtime sessions stick to the replica holding their state, and control operations (weights, LoRA, memory occupation, shutdown) apply to every replica. Combines with the other parallelism axes (`num_gpus = dp × cfg × tp × sp`); monolithic serving only.
|
||||
|
||||
@@ -103,7 +103,7 @@ status-code contract and warmup-mode behavior.
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`speed`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Favors GPU-resident execution for lower latency and higher throughput. Disables CPU offload when unset and enables `torch.compile` by default unless explicitly disabled; may OOM.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Favors GPU-resident execution for lower latency and higher throughput. Disables CPU offload when unset. `torch.compile` stays off unless the model has a validated default or it is enabled explicitly; may OOM.</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`memory`</td>
|
||||
@@ -112,7 +112,7 @@ status-code contract and warmup-mode behavior.
|
||||
</tbody>
|
||||
</table>
|
||||
|
||||
`auto` checks selected GPU memory before applying FSDP. In multi-GPU runs it uses the least available memory across selected GPUs, and only turns on FSDP automatically when doing so can replace DiT offload. Text encoder, image encoder, and other component residency still follow the offload policy unless the model marks a high-memory resident path as safe. When the model default uses CFG and the user did not set a parallelism policy, `auto` may also enable CFG parallelism. `speed` intentionally does not check memory; it is the mode for users who prefer latency/throughput and accept OOM risk, and it will turn on `torch.compile` automatically unless `--enable-torch-compile false` is set explicitly.
|
||||
`auto` checks selected GPU memory before applying FSDP. In multi-GPU runs it uses the least available memory across selected GPUs, and only turns on FSDP automatically when doing so can replace DiT offload. Text encoder, image encoder, and other component residency still follow the offload policy unless the model marks a high-memory resident path as safe. When the model default uses CFG and the user did not set a parallelism policy, `auto` may also enable CFG parallelism. `speed` intentionally does not check memory; it is the mode for users who prefer latency/throughput and accept OOM risk. It keeps `torch.compile` disabled by default because its effect varies by model and workload. A model-specific deployment config may enable a validated compile path, and `--enable-torch-compile true` always opts in explicitly.
|
||||
|
||||
The modes tune residency for native pipeline components declared to the component residency manager. Today this covers the major DiT, text/image encoder, VAE, vocoder, and upsampler components; DiT can use layerwise offload when supported, while text encoders use either resident execution or component CPU offload. Do not assume text-encoder layerwise offload unless a model implements and validates it.
|
||||
|
||||
|
||||
@@ -25,9 +25,9 @@ class ModelDeploymentConfig:
|
||||
auto_enable_cfg_parallel: bool = True
|
||||
# degree 1 keeps CFG parallel disabled and leaves GPUs available for SP
|
||||
auto_cfg_parallel_degree_by_num_gpus: tuple[tuple[int, int], ...] = ()
|
||||
# Let performance_mode=speed opt into torch.compile unless the model has
|
||||
# established that the compiled path changes its numerical contract.
|
||||
speed_mode_enable_torch_compile_by_default: bool = True
|
||||
# torch.compile is model opt-in because it can be slower than eager for
|
||||
# diffusion workloads dominated by already-optimized kernels
|
||||
speed_mode_enable_torch_compile_by_default: bool = False
|
||||
supports_cfg_parallel: bool = True
|
||||
|
||||
def get_auto_cfg_parallel_degree(self, num_gpus: int) -> int:
|
||||
|
||||
@@ -78,9 +78,7 @@ class ServerArgsAutoTuner:
|
||||
and not args.enable_torch_compile
|
||||
and not args.is_arg_explicitly_set("enable_torch_compile")
|
||||
):
|
||||
# speed means fastest: compile by default. An explicit
|
||||
# --enable-torch-compile false still wins (e.g. models where
|
||||
# compile is slower or changes the numerical contract).
|
||||
# only models with a validated compile win opt in by default
|
||||
args.enable_torch_compile = True
|
||||
logger.info(
|
||||
"performance_mode=speed enables torch.compile "
|
||||
|
||||
@@ -24,6 +24,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import (
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import (
|
||||
MiniMaxH3PipelineConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import (
|
||||
ModelDeploymentConfig,
|
||||
)
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig
|
||||
from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
|
||||
QwenImagePipelineConfig,
|
||||
@@ -1590,7 +1593,7 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
self.assertFalse(args.text_encoder_cpu_offload)
|
||||
self.assertFalse(args.image_encoder_cpu_offload)
|
||||
|
||||
def test_speed_mode_enables_torch_compile_by_default(self):
|
||||
def test_speed_mode_keeps_torch_compile_off_by_default(self):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
QwenImagePipelineConfig(),
|
||||
kwargs={
|
||||
@@ -1599,20 +1602,40 @@ class TestOffloadDefaults(unittest.TestCase):
|
||||
},
|
||||
)
|
||||
|
||||
self.assertTrue(args.enable_torch_compile)
|
||||
|
||||
def test_speed_mode_preserves_explicit_torch_compile_off(self):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
QwenImagePipelineConfig(),
|
||||
kwargs={
|
||||
"model_path": "Qwen/Qwen-Image",
|
||||
"performance_mode": "speed",
|
||||
"enable_torch_compile": False,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertFalse(args.enable_torch_compile)
|
||||
|
||||
def test_speed_mode_preserves_explicit_torch_compile_setting(self):
|
||||
for enabled in (False, True):
|
||||
with self.subTest(enabled=enabled):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
QwenImagePipelineConfig(),
|
||||
kwargs={
|
||||
"model_path": "Qwen/Qwen-Image",
|
||||
"performance_mode": "speed",
|
||||
"enable_torch_compile": enabled,
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(args.enable_torch_compile, enabled)
|
||||
|
||||
def test_speed_mode_honors_model_torch_compile_opt_in(self):
|
||||
with patch.object(
|
||||
QwenImagePipelineConfig,
|
||||
"get_model_deployment_config",
|
||||
return_value=ModelDeploymentConfig(
|
||||
speed_mode_enable_torch_compile_by_default=True
|
||||
),
|
||||
):
|
||||
args = self._from_dict_with_pipeline_config(
|
||||
QwenImagePipelineConfig(),
|
||||
kwargs={
|
||||
"model_path": "Qwen/Qwen-Image",
|
||||
"performance_mode": "speed",
|
||||
},
|
||||
)
|
||||
|
||||
self.assertTrue(args.enable_torch_compile)
|
||||
|
||||
def test_speed_mode_uses_minimax_h3_compile_policy(self):
|
||||
for explicit, expected in ((None, False), (True, True)):
|
||||
kwargs = {
|
||||
|
||||
Reference in New Issue
Block a user