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

[diffusion] optimization: reduce minimax h3 mps memory pressure (#33880)

Browse Source
This commit is contained in:
Mick
2026-08-19 13:35:01 +08:00
committed by GitHub
parent 3391ab3712
commit 593b1a9b8a
26 changed files with 757 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/pipeline_configs/minimax_h3.py
+26
View File
@@ -25,6 +25,9 @@ from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend i
AttentionRequirements,
)
from sglang.multimodal_gen.runtime.layers.attention.selector import get_attn_backend
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
LAYERWISE_OFFLOAD,
)
from sglang.multimodal_gen.runtime.platforms import (
AttentionBackendEnum,
current_platform,
@@ -186,6 +189,29 @@ class MiniMaxH3PipelineConfig(PipelineConfig):
def validate_server_args(self, server_args) -> None:
# Reject known-inexact VAE modes before any large component download.
self.vae_config.resolved_parallel_decode_mode()
if current_platform.is_mps():
required_components = (
"transformer",
"text_encoder",
"video_vae",
"audio_vae",
)
missing_components = [
component
for component in required_components
if server_args.residency_mode(component) != LAYERWISE_OFFLOAD
]
if missing_components:
raise ValueError(
"MiniMax-H3 on MPS requires synchronous layerwise offload for "
f"{missing_components}; pass --layerwise-offload-components "
"transformer text_encoder video_vae audio_vae"
)
if server_args.enable_torch_compile:
raise ValueError(
"MiniMax-H3 MPS execution does not support torch.compile; "
"pass --enable-torch-compile false"
)
component_backends = server_args.component_attention_backends or {}
attention_backend = component_backends.get(
"transformer", self._server_arg_value(server_args.attention_backend)
python/sglang/multimodal_gen/runtime/layers/attention/backends/sdpa.py
+27 -7
View File
@@ -24,6 +24,8 @@ _PYTORCH_DEFAULT_CUDA_SDP_BACKENDS = [
SDPBackend.MATH,
]
_MPS_VARLEN_QUERY_CHUNK_SIZE = 128
class SDPABackend(AttentionBackend):
@@ -128,13 +130,31 @@ class SDPAImpl(AttentionImpl):
for start, stop in zip(bounds[:-1], bounds[1:]):
if start == stop:
continue
segment = self.forward(
query[start:stop].unsqueeze(0),
key[start:stop].unsqueeze(0),
value[start:stop].unsqueeze(0),
None,
)
output[start:stop].copy_(segment[0])
if query.device.type != "mps":
segment = self.forward(
query[start:stop].unsqueeze(0),
key[start:stop].unsqueeze(0),
value[start:stop].unsqueeze(0),
None,
)
output[start:stop].copy_(segment[0])
continue
# mps SDPA materializes a quadratic temporary for a varlen segment
# chunking query rows keeps every row's complete K/V context intact
keys = key[start:stop].unsqueeze(0)
values = value[start:stop].unsqueeze(0)
for query_start in range(start, stop, _MPS_VARLEN_QUERY_CHUNK_SIZE):
query_stop = min(query_start + _MPS_VARLEN_QUERY_CHUNK_SIZE, stop)
segment = self.forward(
query[query_start:query_stop].unsqueeze(0),
keys,
values,
None,
)
output[query_start:query_stop].copy_(segment[0])
torch.mps.synchronize()
torch.mps.empty_cache()
return output
python/sglang/multimodal_gen/runtime/layers/linear.py
+5 -5
View File
@@ -154,11 +154,11 @@ class UnquantizedLinearMethod(LinearMethodBase):
def apply(
self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None
) -> torch.Tensor:
if x.device.type == "mps" and (
x.dtype != torch.float32
or layer.weight.dtype != torch.float32
or (bias is not None and bias.dtype != torch.float32)
):
if x.device.type == "mps":
if x.dtype == layer.weight.dtype and (
bias is None or bias.dtype == x.dtype
):
return F.linear(x, layer.weight, bias)
return F.linear(
x.to(torch.float32),
layer.weight.to(torch.float32),
python/sglang/multimodal_gen/runtime/loader/component_loaders/text_encoder_loader.py
+11 -2
View File
@@ -530,7 +530,7 @@ class TextEncoderLoader(ComponentLoader):
)
component_starts_on_cpu = False
if component_starts_on_cpu and not current_platform.is_mps():
if component_starts_on_cpu:
model_device = torch.device("cpu")
else:
model_device = local_torch_device
@@ -559,6 +559,12 @@ class TextEncoderLoader(ComponentLoader):
)
model.bind_encoder_tp_group(encoder_tp_group)
if current_platform.is_mps() and component_starts_on_cpu:
# the h3 encoder is layered immediately after this loader returns
# compatible CPU safetensors stay mapped instead of copying the
# full Qwen checkpoint into unified memory
model._mps_zero_copy_weight_loading = True
weights_to_load = {name for name, _ in model.named_parameters()}
loaded_weights = model.load_weights(
self._get_all_weights(
@@ -581,7 +587,10 @@ class TextEncoderLoader(ComponentLoader):
if component_starts_on_cpu:
if current_platform.is_mps():
model = model.to(local_torch_device)
logger.info(
"Keeping %s on CPU for MPS layerwise offload",
model.__class__.__name__,
)
else:
model = model.to("cpu")
else:
python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py
+31 -7
View File
@@ -136,6 +136,18 @@ class TransformerLoader(ComponentLoader):
]
expected_library = "diffusers"
def customized_load_kwargs_for_component(
self, server_args: ServerArgs, component_name: str
) -> dict[str, bool]:
if current_platform.is_mps() and self._is_component_set_as_layerwise_load(
server_args, component_name
):
logger.info(
"Loading %s on CPU first for MPS layerwise offload", component_name
)
return {"cpu_offload_flag": True}
return {}
def should_raise_customized_load_error(
self, server_args: ServerArgs, component_name: str
) -> bool:
@@ -151,7 +163,11 @@ class TransformerLoader(ComponentLoader):
)
def load_customized(
self, component_model_path: str, server_args: ServerArgs, component_name: str
self,
component_model_path: str,
server_args: ServerArgs,
component_name: str,
cpu_offload_flag: bool = False,
):
"""Load the transformer based on the model path, and inference args."""
component_server_args = _server_args_for_transformer_component(
@@ -196,8 +212,9 @@ class TransformerLoader(ComponentLoader):
)
# Quantization adapters may require resident weights, so placement must
# be resolved after they have validated the component configuration.
component_starts_on_cpu = server_args.should_start_component_on_cpu(
component_name
component_starts_on_cpu = (
server_args.should_start_component_on_cpu(component_name)
or cpu_offload_flag
)
use_fsdp = server_args.should_use_fsdp_for_component(component_name)
@@ -259,10 +276,14 @@ class TransformerLoader(ComponentLoader):
logger.debug("quantization config: %s", init_params["quant_config"])
local_torch_device = get_local_torch_device()
checkpoint_load_device = _resolve_checkpoint_load_device(
local_torch_device,
component_starts_on_cpu=component_starts_on_cpu,
runtime_quant_config=quant_spec.runtime_quant_config,
checkpoint_load_device = (
torch.device("cpu")
if cpu_offload_flag
else _resolve_checkpoint_load_device(
local_torch_device,
component_starts_on_cpu=component_starts_on_cpu,
runtime_quant_config=quant_spec.runtime_quant_config,
)
)
direct_gpu_weight_loading = bool(
component_server_args.direct_gpu_weight_loading
@@ -276,6 +297,9 @@ class TransformerLoader(ComponentLoader):
needs_device_weight_postprocess=quant_spec.needs_device_weight_postprocess,
component_starts_on_cpu=component_starts_on_cpu,
load_full_state_dict_on_device=direct_gpu_weight_loading,
mps_layerwise_cpu_staging=bool(
cpu_offload_flag and current_platform.is_mps()
),
)
if direct_gpu_weight_loading:
logger.warning(
python/sglang/multimodal_gen/runtime/loader/component_loaders/vae_loader.py
+25 -4
View File
@@ -98,8 +98,24 @@ class VAELoader(ComponentLoader):
component_names = ["vae", "audio_vae", "video_vae"]
expected_library = "diffusers"
def customized_load_kwargs_for_component(
self, server_args: ServerArgs, component_name: str
) -> dict[str, bool]:
if current_platform.is_mps() and self._is_component_set_as_layerwise_load(
server_args, component_name
):
logger.info(
"Loading %s on CPU first for MPS layerwise offload", component_name
)
return {"cpu_offload_flag": True}
return {}
def load_customized(
self, component_model_path: str, server_args: ServerArgs, component_name: str
self,
component_model_path: str,
server_args: ServerArgs,
component_name: str,
cpu_offload_flag: bool = False,
):
"""Load the VAE based on the model path, and inference args."""
config = get_diffusers_component_config(component_path=component_model_path)
@@ -133,8 +149,9 @@ class VAELoader(ComponentLoader):
# NOTE: some post init logics are only available after updated with config
vae_config.post_init()
component_starts_on_cpu = server_args.should_start_component_on_cpu(
component_name
component_starts_on_cpu = (
server_args.should_start_component_on_cpu(component_name)
or cpu_offload_flag
)
target_device = self.target_device(component_starts_on_cpu)
@@ -190,7 +207,11 @@ class VAELoader(ComponentLoader):
loaded.update(safetensors_load_file(sf_path))
_backfill_ltx2_audio_vae_latent_stats(loaded, component_name)
strict_load = native_only
vae.load_state_dict(loaded, strict=strict_load)
vae.load_state_dict(
loaded,
strict=strict_load,
assign=bool(cpu_offload_flag and current_platform.is_mps()),
)
if not strict_load:
state_keys = set(vae.state_dict().keys())
python/sglang/multimodal_gen/runtime/loader/fsdp_load.py
+23 -1
View File
@@ -295,6 +295,15 @@ def maybe_load_fsdp_model(
logger.info("Disabling FSDP for MPS platform as it's not compatible")
weight_load_plan = weight_load_plan or WeightLoadPlan(checkpoint_load_device=device)
mps_zero_copy_weight_loading = bool(
current_platform.is_mps()
and weight_load_plan.mps_layerwise_cpu_staging
and weight_load_plan.checkpoint_load_device.type == "cpu"
)
if mps_zero_copy_weight_loading:
# layerwise offload replaces block parameters with mps placeholders after
# load, so compatible checkpoint tensors stay file-backed on CPU
model._mps_zero_copy_weight_loading = True
defer_cpu_placement = bool(
component_starts_on_cpu
and weight_load_plan.defer_cpu_placement
@@ -408,6 +417,7 @@ def maybe_load_fsdp_model(
strict=strict,
cpu_offload=load_on_cpu,
param_names_mapping=param_names_mapping_fn,
mps_zero_copy_weight_loading=mps_zero_copy_weight_loading,
preconverted_state_dict=preconverted_state_dict,
)
if bnb_quant_states:
@@ -531,6 +541,7 @@ def load_model_from_full_model_state_dict(
strict: bool = False,
cpu_offload: bool = False,
param_names_mapping: Callable[[str], tuple[str, Any, Any]] | None = None,
mps_zero_copy_weight_loading: bool = False,
preconverted_state_dict: (
tuple[
dict[
@@ -555,6 +566,7 @@ def load_model_from_full_model_state_dict(
strict (bool): flag to check if to load the model in strict mode
cpu_offload (bool): flag to check if FSDP offload is enabled
param_names_mapping (Optional[Callable[[str], str]]): a function that maps full param name to sharded param name
mps_zero_copy_weight_loading (bool): retain compatible CPU checkpoint tensors for MPS layerwise offload
Returns:
``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:
* **missing_keys** is a list of str containing the missing keys
@@ -696,7 +708,17 @@ def load_model_from_full_model_state_dict(
if actual_param is not None
else None
)
if weight_loader is not None:
use_checkpoint_tensor_directly = bool(
mps_zero_copy_weight_loading
and actual_param is not None
and not getattr(actual_param, "mps_zero_copy_unsafe", False)
and tuple(meta_sharded_param.shape) == tuple(full_tensor.shape)
and full_tensor.device.type == "cpu"
and full_tensor.dtype == target_dtype
)
if use_checkpoint_tensor_directly:
sharded_tensor = full_tensor
elif weight_loader is not None:
assert actual_param is not None
if _can_assign_cpu_tensor_without_copy(
actual_param,
python/sglang/multimodal_gen/runtime/loader/weight_load_plan.py
+5
View File
@@ -11,6 +11,9 @@ class WeightLoadPlan:
checkpoint_load_device: torch.device
# Device required while running process_weights_after_loading; None means unchanged.
weight_postprocess_device: torch.device | None = None
# mps layerwise loading retains compatible safetensors tensors as CPU-backed
# parameters instead of materializing a second unified-memory copy
mps_layerwise_cpu_staging: bool = False
# Delay final CPU placement until after device-side weight postprocessing.
defer_cpu_placement: bool = False
# keep the complete mapped checkpoint state dict on the load device
@@ -24,6 +27,7 @@ class WeightLoadPlan:
needs_device_weight_postprocess: bool,
component_starts_on_cpu: bool,
load_full_state_dict_on_device: bool = False,
mps_layerwise_cpu_staging: bool = False,
) -> "WeightLoadPlan":
# if on-device weight postprocessing is required, load directly to device to speedup loading
weight_postprocess_device = (
@@ -36,4 +40,5 @@ class WeightLoadPlan:
needs_device_weight_postprocess and component_starts_on_cpu
),
load_full_state_dict_on_device=load_full_state_dict_on_device,
mps_layerwise_cpu_staging=mps_layerwise_cpu_staging,
)
python/sglang/multimodal_gen/runtime/managers/gpu_worker.py
+11 -10
View File
@@ -65,9 +65,7 @@ from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.post_training.gpu_worker_post_training_mixin import (
GPUWorkerPostTrainingMixin,
)
from sglang.multimodal_gen.runtime.realtime.session import (
RealtimeSessionCache,
)
from sglang.multimodal_gen.runtime.realtime.session import RealtimeSessionCache
from sglang.multimodal_gen.runtime.server_args import PortArgs, ServerArgs
from sglang.multimodal_gen.runtime.utils.common import set_cuda_arch, set_musa_arch
from sglang.multimodal_gen.runtime.utils.logging_utils import (
@@ -228,7 +226,8 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
def init_device_and_model(self) -> None:
"""Initialize the device and load the model."""
current_platform.set_device(current_platform.get_device(self.local_rank))
if not current_platform.is_mps():
current_platform.set_device(current_platform.get_device(self.local_rank))
# num_gpus is the total world size across every node; the co-located,
# CPU-contending worker count on THIS host is num_gpus // nnodes.
local_num_gpus = self.server_args.num_gpus // self.server_args.nnodes
@@ -316,9 +315,8 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
if output_batch.metrics:
output_batch.metrics.record_memory_snapshot("mem_analysis", final_snapshot)
# for details on max_memory_reserved: https://docs.pytorch.org/docs/stable/generated/torch.cuda.memory.max_memory_reserved.html
peak_reserved_bytes = torch.get_device_module().max_memory_reserved()
peak_allocated_bytes = torch.get_device_module().max_memory_allocated()
peak_reserved_bytes = final_snapshot.peak_reserved_mb * (1024**2)
peak_allocated_bytes = final_snapshot.peak_allocated_mb * (1024**2)
output_batch.peak_memory_mb = peak_reserved_bytes / (1024**2)
peak_reserved_gb = peak_reserved_bytes / (1024**3)
@@ -466,7 +464,11 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
output_batch = None
forward_failed = False
try:
if self.is_output_rank and not current_platform.is_cpu():
if (
self.is_output_rank
and not current_platform.is_cpu()
and not current_platform.is_mps()
):
torch.get_device_module().reset_peak_memory_stats()
start_time = (
@@ -690,8 +692,7 @@ class GPUWorker(GPUWorkerPostTrainingMixin):
def _record_output_peak_memory(self, output_batch: OutputBatch) -> None:
if not self.is_output_rank or current_platform.is_cpu():
return
peak_reserved_bytes = torch.get_device_module().max_memory_reserved()
output_batch.peak_memory_mb = peak_reserved_bytes / (1024**2)
output_batch.peak_memory_mb = capture_memory_snapshot().peak_reserved_mb
def _forward_group(self, batch: list[Req]) -> OutputBatch:
assert self.pipeline is not None
python/sglang/multimodal_gen/runtime/managers/memory_managers/component_residency_strategies.py
+12
View File
@@ -202,6 +202,14 @@ class LayerwiseOffloadStrategy(ComponentResidencyStrategy):
state: ResidencyState,
) -> None:
if isinstance(module, LayerwiseOffloadableModuleMixin):
# MPS layerwise components retain checkpoint-backed CPU weights and
# synchronously materialize one layer at a time. Moving the whole
# module here would defeat that bounded-residency contract.
if current_platform.is_mps():
if module.mps_stream_non_layer_weights:
return
_module_to_local_device(module, dtype=use.target_dtype)
return
module.prepare_for_next_req()
def finish_use(
@@ -214,6 +222,10 @@ class LayerwiseOffloadStrategy(ComponentResidencyStrategy):
return
for manager in module.layerwise_offload_managers:
manager.release_all()
if current_platform.is_mps():
torch.mps.synchronize()
module.restore_mps_cpu_non_layer_weights()
torch.mps.empty_cache()
def finish_request(
self,
python/sglang/multimodal_gen/runtime/managers/memory_managers/layerwise_offload.py
+237 -16
View File
@@ -1,6 +1,7 @@
import bisect
import re
from collections.abc import Mapping, Sequence
from contextlib import nullcontext
from typing import Any, Dict, List, Set, Tuple
import torch
@@ -84,7 +85,8 @@ class LayerwiseOffloadManager:
"""A lightweight layerwise CPU offload manager.
This utility offloads per-layer parameters/buffers from GPU to CPU, and
supports async H2D prefetch using a dedicated CUDA stream.
supports async H2D prefetch using a dedicated CUDA stream. MPS uses
synchronous per-layer transfers from the checkpoint-backed CPU tensors.
Typical usage:
- Construct the manager with the target model and the list-like module
@@ -104,13 +106,23 @@ class LayerwiseOffloadManager:
pin_cpu_memory: bool = True,
prefetch_size: int = 1,
resident_layers: int = 0,
initialize: bool = True,
residency_policy: str = RESIDENCY_POLICY_LEADING,
) -> None:
self.model = model
self.layers_attr_str = layers_attr_str
self.num_layers = num_layers
self.pin_cpu_memory = pin_cpu_memory
self.prefetch_size = min(max(1, prefetch_size), self.num_layers)
self._synchronous_mps = current_platform.is_mps()
# mps shares physical memory with the CPU and has no pinned host memory
# or CUDA-style copy streams
self.pin_cpu_memory = bool(pin_cpu_memory and not self._synchronous_mps)
# an explicit MPS zero avoids staging the next layer alongside the
# active one; MPS has no transfer overlap to recover from that cost
self.prefetch_size = (
0
if current_platform.is_mps() and prefetch_size == 0
else min(max(1, prefetch_size), self.num_layers)
)
# Layers held on GPU across denoise steps, instead of being re-streamed
# every step. `residency_policy` picks *which* layers those are; see
# compute_streamed_layers for why the choice is not cosmetic.
@@ -134,10 +146,17 @@ class LayerwiseOffloadManager:
self.enabled = bool(enabled and torch.get_device_module().is_available())
if not self.enabled:
return
self.device = torch.device(
current_platform.device_type, torch.get_device_module().current_device()
self.device = (
current_platform.get_local_torch_device()
if current_platform.is_mps()
else torch.device(
current_platform.device_type,
torch.get_device_module().current_device(),
)
)
self.copy_stream = (
None if self._synchronous_mps else torch.get_device_module().Stream()
)
self.copy_stream = torch.get_device_module().Stream()
# ``named_parameters()`` is relative to ``model``, just like the path in
# ``layers_attr_str``. Anchor the match so a manager for top-level
@@ -153,6 +172,9 @@ class LayerwiseOffloadManager:
# layer_idx -> {name: pinned_cpu_tensor_with_original_stride}
# stores tensors whose original non-contiguous stride/layout must be preserved
self._strided_cpu_weights: Dict[int, Dict[str, torch.Tensor]] = {}
# mps keeps the original CPU tensor for each layer instead of building a
# second flattened host copy
self._mps_cpu_weights: Dict[int, Dict[str, torch.Tensor]] = {}
# layer_idx -> {name: {dtype, offset, numel, shape}}
# stores the offset and numel of each weight from a same layer, of same dtype
self._weight_metadata: Dict[int, Dict[str, Dict[str, Any]]] = {}
@@ -168,6 +190,10 @@ class LayerwiseOffloadManager:
# Store forward hooks for removal
self._forward_hooks: List[Any] = []
if initialize:
self._initialize()
def initialize(self) -> None:
self._initialize()
def _match_layer_idx(self, name: str) -> int | None:
@@ -229,6 +255,10 @@ class LayerwiseOffloadManager:
self._named_parameters = dict(self.model.named_parameters())
self._named_buffers = dict(self.model.named_buffers())
if self._synchronous_mps:
self._initialize_mps_cpu_weights()
return
# 1. collect and group layer parameters by dtype. Keep buffers resident:
# shared buffers such as RoPE caches may be referenced by many layers.
layer_groups: Dict[int, Dict[torch.dtype, List[Tuple[str, torch.Tensor]]]] = {}
@@ -361,6 +391,33 @@ class LayerwiseOffloadManager:
return list(range(count))
return self._next_streamed(after=-1, count=count)
@torch.compiler.disable
def _initialize_mps_cpu_weights(self) -> None:
for name, tensor in self._named_parameters.items():
layer_idx = self._match_layer_idx(name)
if layer_idx is None or layer_idx >= self.num_layers:
continue
local_tensor = self._to_local_tensor(tensor).detach()
cpu_tensor = (
local_tensor
if local_tensor.device.type == "cpu"
else local_tensor.to("cpu")
)
self._mps_cpu_weights.setdefault(layer_idx, {})[name] = cpu_tensor
self._weight_metadata.setdefault(layer_idx, {})[name] = {
"dtype": cpu_tensor.dtype,
}
tensor.data = self._get_shared_empty_tensor_for_target(
tensor, cpu_tensor.dtype
)
torch.mps.empty_cache()
self.register_forward_hooks()
self._configured = True
logger.info(
f"Initialized synchronous MPS layerwise offload with {self.num_layers} layers"
)
def prepare_for_next_req(self, non_blocking=True):
"""
Prepare for the next round of denoising loop with prefetching the necessary layers
@@ -435,22 +492,41 @@ class LayerwiseOffloadManager:
"""
idempotent
"""
if not self.enabled or self.device is None or self.copy_stream is None:
if not self.enabled or self.device is None:
return
if layer_idx < 0 or layer_idx >= self.num_layers:
return
if layer_idx in self._gpu_layers:
return
if self._synchronous_mps:
cpu_weights = self._mps_cpu_weights.get(layer_idx)
if not cpu_weights:
return
with torch.inference_mode(False), torch.no_grad():
for name, cpu_tensor in cpu_weights.items():
target = self.get_target_with_name(name)
target.data = self._wrap_for_target(
target,
cpu_tensor.to(device=self.device, non_blocking=False),
)
self._gpu_layers.add(layer_idx)
return
if layer_idx not in self._consolidated_cpu_weights:
return
self.copy_stream.wait_stream(torch.get_device_module().current_stream())
if self.copy_stream is not None:
self.copy_stream.wait_stream(torch.get_device_module().current_stream())
stream_context = torch.get_device_module().stream(self.copy_stream)
else:
# the device has no CUDA-like stream or pinned-memory support
non_blocking = False
stream_context = nullcontext()
# create gpu buffer and load from CPU buffer
gpu_buffers: Dict[torch.dtype, torch.Tensor] = {}
with (
torch.inference_mode(False),
torch.no_grad(),
torch.get_device_module().stream(self.copy_stream),
stream_context,
):
for dtype, cpu_buffer in self._consolidated_cpu_weights[layer_idx].items():
gpu_buffer = torch.empty(
@@ -488,10 +564,11 @@ class LayerwiseOffloadManager:
].view(meta["shape"])
target.data = self._wrap_for_target(target, local_tensor)
# record the prefetch event of this layer after all copies are enqueued
event = torch.get_device_module().Event()
event.record(self.copy_stream)
self._prefetch_events[layer_idx] = event
if self.copy_stream is not None:
# record after all copies so the consumer waits for every weight copy
event = torch.get_device_module().Event()
event.record(self.copy_stream)
self._prefetch_events[layer_idx] = event
self._gpu_layers.add(layer_idx)
@@ -524,6 +601,11 @@ class LayerwiseOffloadManager:
)
self._gpu_layers.discard(layer_idx)
if self._synchronous_mps:
# mps dispatch is asynchronous, so a tensor rebinding alone leaves
# prior layer allocations live until the command buffer drains
torch.mps.synchronize()
torch.mps.empty_cache()
@torch.compiler.disable
def release_all(self) -> None:
@@ -554,6 +636,10 @@ class LayerwiseOffloadManager:
"""Sync a layer's weights from GPU back to CPU."""
if not self.enabled or layer_idx not in self._gpu_layers:
return
if self._synchronous_mps:
# inference does not mutate parameters; retain the original mapped
# tensor instead of materializing a CPU copy after every layer
return
if layer_idx not in self._consolidated_cpu_weights:
return
@@ -616,6 +702,33 @@ class LayerwiseOffloadManager:
return None
updated_names: Set[str] = set()
if self._synchronous_mps:
for name, loaded_weight in weight_dict.items():
layer_idx = self._match_layer_idx(name)
if layer_idx is None:
continue
cpu_tensor = self._mps_cpu_weights.get(layer_idx, {}).get(name)
if cpu_tensor is None:
continue
if tuple(cpu_tensor.shape) != tuple(loaded_weight.shape):
raise ValueError(
f"Shape mismatch for {name}: "
f"expected={tuple(cpu_tensor.shape)}, "
f"loaded={tuple(loaded_weight.shape)}"
)
replacement = loaded_weight.to(
device="cpu", dtype=cpu_tensor.dtype
).detach()
self._mps_cpu_weights[layer_idx][name] = replacement
if layer_idx in self._gpu_layers:
target = self.get_target_with_name(name)
target.data = self._wrap_for_target(
target,
replacement.to(device=target.device, dtype=target.dtype),
)
updated_names.add(name)
return updated_names
for name, loaded_weight in weight_dict.items():
layer_idx = self._match_layer_idx(name)
if layer_idx is None:
@@ -665,6 +778,11 @@ class LayerwiseOffloadManager:
when offload is enabled, this method returns the real weights and
can be used for checksum computation.
"""
if self._synchronous_mps:
for layer_idx in sorted(self._mps_cpu_weights):
yield from self._mps_cpu_weights[layer_idx].items()
return
for layer_idx in sorted(self._weight_metadata):
for name, meta in self._weight_metadata[layer_idx].items():
if meta.get("preserve_strides", False):
@@ -696,7 +814,7 @@ class LayerwiseOffloadManager:
if i not in self._gpu_layers:
# LTX audio VAE traverses decoder.up in reverse order
self.prefetch_layer(i, non_blocking=False)
if i in self._prefetch_events:
if i in self._prefetch_events and self.copy_stream is not None:
torch.get_device_module().current_stream().wait_event(
self._prefetch_events[i]
)
@@ -716,7 +834,7 @@ class LayerwiseOffloadManager:
):
self.prefetch_layer(layer_to_prefetch, non_blocking=True)
# trigger batch prefetch (i + prefetch_size ~ i + 2 * prefetch_size) if needed
elif i % self.prefetch_size == 0:
elif self.prefetch_size and i % self.prefetch_size == 0:
for j in range(i + self.prefetch_size, i + 2 * self.prefetch_size):
layer_to_prefetch = j % self.num_layers
self.prefetch_layer(layer_to_prefetch, non_blocking=True)
@@ -750,11 +868,106 @@ class LayerwiseOffloadableModuleMixin:
# whether the current module is selected by the `dit` group
layerwise_offload_dit_group_enabled: bool = True
# H3 has a large packed-sequence working set on MPS, so its non-block
# weights are materialized only for the subphase that consumes them
mps_stream_non_layer_weights: bool = False
# The list of names of this module's layer/block ModuleList or Sequential attributes.
layer_names: List[str] = []
layerwise_offload_managers: list[LayerwiseOffloadManager] = []
def _capture_mps_cpu_non_layer_weights(self) -> None:
managed_names = {
name
for manager in self.layerwise_offload_managers
for weights in manager._mps_cpu_weights.values()
for name in weights
}
self._mps_cpu_non_layer_parameters = {
name: parameter.detach()
for name, parameter in self.named_parameters()
if name not in managed_names
}
self._mps_cpu_buffers = {
name: buffer.detach() for name, buffer in self.named_buffers()
}
if not self.mps_stream_non_layer_weights:
return
parameters = dict(self.named_parameters())
for name, tensor in self._mps_cpu_non_layer_parameters.items():
parameters[name].data = torch.empty(
(1,),
dtype=tensor.dtype,
device=current_platform.get_local_torch_device(),
)
buffers = dict(self.named_buffers())
for name, tensor in self._mps_cpu_buffers.items():
buffers[name].data = torch.empty(
(1,),
dtype=tensor.dtype,
device=current_platform.get_local_torch_device(),
)
@staticmethod
def _matches_mps_weight_prefix(name: str, prefixes: tuple[str, ...]) -> bool:
return any(
name == prefix or name.startswith(f"{prefix}.") for prefix in prefixes
)
def materialize_mps_non_layer_weights(self, *prefixes: str) -> None:
if not current_platform.is_mps() or not self.mps_stream_non_layer_weights:
return
selected_prefixes = tuple(prefixes)
with torch.inference_mode(False), torch.no_grad():
parameters = dict(self.named_parameters())
for name, tensor in self._mps_cpu_non_layer_parameters.items():
if self._matches_mps_weight_prefix(name, selected_prefixes):
parameters[name].data = tensor.to(
current_platform.get_local_torch_device()
)
buffers = dict(self.named_buffers())
for name, tensor in self._mps_cpu_buffers.items():
if self._matches_mps_weight_prefix(name, selected_prefixes):
buffers[name].data = tensor.to(
current_platform.get_local_torch_device()
)
def release_mps_non_layer_weights(self, *prefixes: str) -> None:
if not current_platform.is_mps() or not self.mps_stream_non_layer_weights:
return
selected_prefixes = tuple(prefixes)
with torch.inference_mode(False), torch.no_grad():
parameters = dict(self.named_parameters())
for name, tensor in self._mps_cpu_non_layer_parameters.items():
if self._matches_mps_weight_prefix(name, selected_prefixes):
parameters[name].data = torch.empty(
(1,),
dtype=tensor.dtype,
device=current_platform.get_local_torch_device(),
)
buffers = dict(self.named_buffers())
for name, tensor in self._mps_cpu_buffers.items():
if self._matches_mps_weight_prefix(name, selected_prefixes):
buffers[name].data = torch.empty(
(1,),
dtype=tensor.dtype,
device=current_platform.get_local_torch_device(),
)
torch.mps.synchronize()
torch.mps.empty_cache()
def restore_mps_cpu_non_layer_weights(self) -> None:
if not current_platform.is_mps():
return
with torch.inference_mode(False), torch.no_grad():
parameters = dict(self.named_parameters())
for name, tensor in self._mps_cpu_non_layer_parameters.items():
parameters[name].data = tensor
buffers = dict(self.named_buffers())
for name, tensor in self._mps_cpu_buffers.items():
buffers[name].data = tensor
def configure_layerwise_offload(self, server_args: ServerArgs):
self.layerwise_offload_managers = []
named_modules = dict(self.named_modules())
@@ -774,7 +987,9 @@ class LayerwiseOffloadableModuleMixin:
prefetch_value = (
server_args.dit_offload_prefetch_size if dit_tuning_enabled else 0.0
)
if prefetch_value < 1.0:
if current_platform.is_mps() and prefetch_value == 0.0:
prefetch_size = 0
elif prefetch_value < 1.0:
prefetch_size = 1 + int(round(prefetch_value * (num_layers - 1)))
else:
prefetch_size = int(prefetch_value)
@@ -797,6 +1012,7 @@ class LayerwiseOffloadableModuleMixin:
pin_cpu_memory=server_args.pin_cpu_memory,
prefetch_size=prefetch_size,
resident_layers=resident_layers,
initialize=not current_platform.is_mps(),
residency_policy=(
server_args.dit_layerwise_residency_policy
if dit_tuning_enabled
@@ -806,6 +1022,11 @@ class LayerwiseOffloadableModuleMixin:
self.layerwise_offload_managers.append(manager)
configured_layer_names.append(layer_name)
if current_platform.is_mps():
for manager in self.layerwise_offload_managers:
manager.initialize()
self._capture_mps_cpu_non_layer_weights()
if configured_layer_names:
logger.debug(
"Enabled layerwise offload for %s on modules: %s",
python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py
+200 -3
View File
@@ -77,6 +77,19 @@ logger = init_logger(__name__)
_ARCH_DEFAULTS = MiniMaxH3DiTArchConfig()
_BF16_DTYPE = torch.bfloat16
_FP32_DTYPE = torch.float32
_MPS_MLP_TOKEN_CHUNK_SIZE = 128
# keep MPS activation chunks below the allocator high-watermark; CUDA keeps
# its fused full-sequence projection
_MPS_QKV_PROJECTION_TOKEN_CHUNK_SIZE = 128
_MPS_ATTENTION_QUERY_TOKEN_CHUNK_SIZE = 128
_MPS_EMBED_WEIGHT_PREFIXES = (
"condition_proj",
"video_patch_proj",
"audio_patch_proj",
"time_embedder",
"token_refiner.final_norm",
)
_MINIMAX_H3_FP32_PARAM_NAMES_IN_MODEL_ORDER = (
"video_patch_proj.weight",
@@ -390,6 +403,18 @@ def _apply_rope_qk(
return q, k
def _apply_rope(
x: torch.Tensor,
cos_sin_cache: torch.Tensor,
) -> torch.Tensor:
"""Apply the eager (non-CUDA) H3 RoPE path to one Q or K tensor."""
half = cos_sin_cache.shape[-1] // 2
cos_half, sin_half = cos_sin_cache.split(half, dim=-1)
cos = torch.cat((cos_half, cos_half), dim=-1).unsqueeze(1)
sin = torch.cat((sin_half, sin_half), dim=-1).unsqueeze(1)
return _apply_rope_cos_sin(x, cos, sin)
class MiniMaxH3TimeEmbedder(nn.Module):
def __init__(
self,
@@ -597,6 +622,9 @@ class MiniMaxH3Attention(nn.Module):
def _install_qkv_weight_loader(self, arch: MiniMaxH3DiTArchConfig) -> None:
weight = self.qkv_proj.weight
# h3 checkpoints interleave each attention head's Q, K, and V rows
# this parameter needs reordering before the native QKV projection
weight.mps_zero_copy_unsafe = True
base_loader = weight.weight_loader
def _reorder_checkpoint_weight(loaded_weight: torch.Tensor) -> torch.Tensor:
@@ -630,6 +658,107 @@ class MiniMaxH3Attention(nn.Module):
# rank-local FSDP must reorder grouped QKV before selecting each shard
weight.rank_local_weight_transform = _reorder_checkpoint_weight
def _forward_mps_streamed_attention(
self,
x: torch.Tensor,
*,
rope_cache: tuple[torch.Tensor, torch.Tensor] | None,
cu_seqlens: torch.Tensor,
cu_seqlens_host: tuple[int, ...] | None,
max_seqlen: int,
) -> torch.Tensor:
"""Run MPS attention without materializing the full QKV activation.
H3's fused QKV output alone is roughly 1.45 GiB at 768px. MPS shares
unified memory with the host, so retaining it alongside the packed
residual and SDPA workspace can evict the OS. Build normalized K/V
once, then project Q a small chunk at a time and immediately consume it
through attention and the output projection. The formula, weights,
and complete K/V context are unchanged; this is intentionally limited
to single-device MPS where Ulysses collectives are not active.
"""
total = x.shape[0]
key = torch.empty(
(total, self.num_heads, self.head_dim), dtype=x.dtype, device=x.device
)
value = torch.empty_like(key)
cos_sin_cache = None if rope_cache is None else rope_cache[0]
# Do not retain Q while producing the K/V cache. Dropping the chunk
# before the next transfer keeps only two full-width attention tensors.
for start in range(0, total, _MPS_QKV_PROJECTION_TOKEN_CHUNK_SIZE):
stop = min(start + _MPS_QKV_PROJECTION_TOKEN_CHUNK_SIZE, total)
qkv, _ = self.qkv_proj(x[start:stop])
q_chunk, k_chunk, v_chunk = qkv.split(self.local_inner_dim, dim=-1)
del q_chunk
k_chunk = self.k_norm(k_chunk.view(-1, self.num_heads, self.head_dim))
if cos_sin_cache is not None:
k_chunk = _apply_rope(k_chunk, cos_sin_cache[start:stop])
key[start:stop].copy_(k_chunk)
value[start:stop].copy_(v_chunk.view(-1, self.num_heads, self.head_dim))
del qkv, k_chunk, v_chunk
torch.mps.synchronize()
torch.mps.empty_cache()
if self._attention_impl is None:
self._set_attention_backend(
get_attn_backend(
self.head_dim,
x.dtype,
attention_requirements=AttentionRequirements(packed_varlen=True),
)
)
bounds = (
cu_seqlens_host
if cu_seqlens_host is not None
else tuple(int(item) for item in cu_seqlens.tolist())
)
out = torch.empty_like(x)
for sequence_start, sequence_stop in zip(bounds[:-1], bounds[1:]):
if sequence_start == sequence_stop:
continue
keys = key[sequence_start:sequence_stop].unsqueeze(0)
values = value[sequence_start:sequence_stop].unsqueeze(0)
for start in range(
sequence_start,
sequence_stop,
_MPS_QKV_PROJECTION_TOKEN_CHUNK_SIZE,
):
stop = min(start + _MPS_QKV_PROJECTION_TOKEN_CHUNK_SIZE, sequence_stop)
qkv, _ = self.qkv_proj(x[start:stop])
q_chunk, k_chunk, v_chunk = qkv.split(self.local_inner_dim, dim=-1)
del k_chunk, v_chunk
for query_start in range(
start, stop, _MPS_ATTENTION_QUERY_TOKEN_CHUNK_SIZE
):
query_stop = min(
query_start + _MPS_ATTENTION_QUERY_TOKEN_CHUNK_SIZE, stop
)
q = self.q_norm(
q_chunk[query_start - start : query_stop - start].view(
-1, self.num_heads, self.head_dim
)
)
if cos_sin_cache is not None:
q = _apply_rope(q, cos_sin_cache[query_start:query_stop])
attention_out = self._attention_impl.forward(
q.unsqueeze(0), keys, values, None
)[0]
projected, _ = self.out_proj(
attention_out.reshape(
query_stop - query_start, self.local_inner_dim
)
)
out[query_start:query_stop].copy_(projected)
del q, attention_out, projected
del qkv, q_chunk
torch.mps.synchronize()
torch.mps.empty_cache()
del key, value
torch.mps.empty_cache()
return out
def forward(
self,
x: torch.Tensor,
@@ -652,6 +781,15 @@ class MiniMaxH3Attention(nn.Module):
so cu_seqlens retains global packed-document semantics. The inverse
all-to-all restores the row shard before the output projection.
"""
if x.device.type == "mps" and not ulysses_active:
return self._forward_mps_streamed_attention(
x,
rope_cache=rope_cache,
cu_seqlens=cu_seqlens,
cu_seqlens_host=cu_seqlens_host,
max_seqlen=max_seqlen,
)
total = x.shape[0]
qkv, _ = self.qkv_proj(x)
q, k, v = qkv.split(self.local_inner_dim, dim=-1)
@@ -747,6 +885,18 @@ class MiniMaxH3MLP(nn.Module):
self.reuse_fc1_activation = quant_config is None
def forward(self, x: torch.Tensor) -> torch.Tensor:
if x.device.type == "mps":
out = torch.empty_like(x)
for start in range(0, x.shape[0], _MPS_MLP_TOKEN_CHUNK_SIZE):
stop = min(start + _MPS_MLP_TOKEN_CHUNK_SIZE, x.shape[0])
hidden, _ = self.fc1(x[start:stop])
hidden = _silu_mul(hidden, reuse_input=self.reuse_fc1_activation)
chunk, _ = self.fc2(hidden)
out[start:stop].copy_(chunk)
del hidden, chunk
torch.mps.synchronize()
torch.mps.empty_cache()
return out
hidden, _ = self.fc1(x)
hidden = _silu_mul(hidden, reuse_input=self.reuse_fc1_activation)
out, _ = self.fc2(hidden)
@@ -1355,6 +1505,38 @@ class MiniMaxH3FinalLayer(nn.Module):
raise ValueError("MiniMax H3 AdaLN cache parameters are required")
adaln_params = self.adaln_proj(adaln_input)
shift, scale = adaln_params
if x.device.type == "mps":
video = audio = None
for start in range(0, x.shape[0], _MPS_MLP_TOKEN_CHUNK_SIZE):
stop = min(start + _MPS_MLP_TOKEN_CHUNK_SIZE, x.shape[0])
h = self.norm(x[start:stop])
h = _modulate_scale_shift(
h,
shift,
scale,
inverse_indices[start:stop],
dtype=_BF16_DTYPE,
).to(_FP32_DTYPE)
video_chunk, _ = self.video_out(h)
audio_chunk, _ = self.audio_out(h)
if video is None:
video = torch.empty(
(x.shape[0], video_chunk.shape[-1]),
dtype=video_chunk.dtype,
device=x.device,
)
audio = torch.empty(
(x.shape[0], audio_chunk.shape[-1]),
dtype=audio_chunk.dtype,
device=x.device,
)
video[start:stop].copy_(video_chunk)
audio[start:stop].copy_(audio_chunk)
del h, video_chunk, audio_chunk
torch.mps.synchronize()
torch.mps.empty_cache()
assert video is not None and audio is not None
return video, audio
h = self.norm(x)
h = _modulate_scale_shift(h, shift, scale, inverse_indices, dtype=_BF16_DTYPE)
# Preserve full precision through both final output projections.
@@ -1371,6 +1553,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
# parameters mix fp32 (patch projections, timestep embedder, and output
# heads) with bf16 blocks; FSDP must gather in each parameter's own dtype
_fsdp_mixed_dtype_params = True
mps_stream_non_layer_weights = True
_compile_conditions = [is_block]
param_names_mapping = _ARCH_DEFAULTS.param_names_mapping
reverse_param_names_mapping = _ARCH_DEFAULTS.reverse_param_names_mapping
@@ -1551,7 +1734,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
for index in range(arch.num_layers)
]
)
self.layer_names = ["blocks"]
self.layer_names = ["token_refiner.blocks", "blocks"]
self.final_layer = MiniMaxH3FinalLayer(
arch,
quant_config,
@@ -1655,6 +1838,9 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
device: torch.device,
) -> torch.Tensor:
"""Project and refine request-static text conditioning once."""
self.materialize_mps_non_layer_weights(
"condition_proj", "token_refiner.final_norm"
)
text_len = int(refiner_cu_seqlens[1].item())
if text_len <= 0 or text_len > int(prompt_embeds.shape[0]):
raise ValueError(
@@ -1670,12 +1856,14 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
)
)
text_embed, _ = self.condition_proj(text_rows)
return self.token_refiner(
refined = self.token_refiner(
text_embed,
cu_seqlens=true_refiner_cu,
cu_seqlens_host=(0, text_len, text_len),
max_seqlen=text_len,
)
self.release_mps_non_layer_weights("condition_proj", "token_refiner.final_norm")
return refined
def build_rope_cache(
self,
@@ -1690,6 +1878,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
chunk) -- see forward()'s row_start derivation for the identity
this must stay in sync with.
"""
self.materialize_mps_non_layer_weights("rope")
if img_position_ids.dim() != 3 or img_position_ids.shape[0] != 1:
raise ValueError(
"img_position_ids must be [1, S, 3], got "
@@ -1711,7 +1900,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
rope_freqs = self.rope(
img_position_ids[:, row_start : row_start + local_seq_len]
).to(device)
return (
result = (
_rope_cos_sin_cache(rope_freqs, dtype=_BF16_DTYPE),
torch.arange(
local_seq_len,
@@ -1719,6 +1908,8 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
dtype=torch.long,
),
)
self.release_mps_non_layer_weights("rope")
return result
@eager_on_graph(True)
def _embed(
@@ -1994,6 +2185,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
# request-static cache once; direct model callers use this fallback.
rope_cache = kwargs.get("rope_cache")
if rope_cache is None:
self.materialize_mps_non_layer_weights("rope")
rope_freqs = self.rope(img_position_ids[:, row_start:row_stop]).to(device)
rope_cache = (
_rope_cos_sin_cache(rope_freqs, dtype=_BF16_DTYPE),
@@ -2003,6 +2195,8 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
dtype=torch.long,
),
)
self.release_mps_non_layer_weights("rope")
self.materialize_mps_non_layer_weights(*_MPS_EMBED_WEIGHT_PREFIXES)
img_pos = img_pos.to(device)
audio_pos = audio_pos.to(device)
text_pos = text_pos.to(device)
@@ -2023,6 +2217,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
refined_prompt_embeds_length=kwargs.get("refined_prompt_embeds_length"),
local_embedding_layout=kwargs.get("local_embedding_layout"),
)
self.release_mps_non_layer_weights(*_MPS_EMBED_WEIGHT_PREFIXES)
# request-step AdaLN input shared by all blocks
adaln_input = nn.functional.silu(t_emb).to(_BF16_DTYPE)
inverse_indices = inverse_indices.to(device)
@@ -2091,6 +2286,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
None if block_adaln_params is None else block_adaln_params[index]
),
)
self.materialize_mps_non_layer_weights("final_layer")
video_logits, audio_logits = self.final_layer(
hidden,
adaln_input=adaln_input,
@@ -2104,6 +2300,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin):
)
),
)
self.release_mps_non_layer_weights("final_layer")
if sp_ws > 1:
from sglang.multimodal_gen.runtime.distributed.parallel_state import (
get_sp_group,
python/sglang/multimodal_gen/runtime/models/encoders/minimax_h3_qwen3vl.py
+15 -3
View File
@@ -41,9 +41,10 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder):
eight otherwise-idle ranks during encoding.
"""
layer_names = [*TextEncoder.layer_names, "model.visual.blocks"]
supports_dp_encode = True
# The inherited text-layer list covers Qwen's language stack; reference
# modes also execute the embedded visual tower.
layer_names = [*TextEncoder.layer_names, "model.visual.blocks"]
supported_checkpoint_quantization_methods = frozenset({"fp8"})
@staticmethod
@@ -191,7 +192,18 @@ class MiniMaxH3Qwen3VLEncoder(TextEncoder):
)
weight_loader = getattr(param, "weight_loader", default_weight_loader)
try:
weight_loader(param, loaded_weight.to(param.dtype))
can_keep_checkpoint_tensor = bool(
getattr(self, "_mps_zero_copy_weight_loading", False)
and weight_loader is default_weight_loader
and param.device.type == "cpu"
and loaded_weight.device.type == "cpu"
and loaded_weight.dtype == param.dtype
and tuple(loaded_weight.shape) == tuple(param.shape)
)
if can_keep_checkpoint_tensor:
param.data = loaded_weight
else:
weight_loader(param, loaded_weight.to(param.dtype))
except Exception as exc:
raise RuntimeError(
"Failed to load MiniMax H3 Qwen3-VL weight "
python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/base_module.py
+5 -1
View File
@@ -1,6 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# Transformer building blocks for the MiniMax H3 visual VAE ViT decoder.
import math
from contextlib import nullcontext
from typing import Optional
import torch
@@ -171,7 +172,10 @@ class RotaryEmbeddingND(nn.Module):
if D != self.n_dim:
raise ValueError(f"Expected {self.n_dim} dimensions, got {D}")
with torch.autocast("cuda", enabled=False):
autocast_context = (
torch.autocast("cuda", enabled=False) if img_ids.is_cuda else nullcontext()
)
with autocast_context:
angles = (
self.angle_scale
* img_ids[:, :, :, None]
python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/vae_vit.py
+8 -2
View File
@@ -1,5 +1,7 @@
# SPDX-License-Identifier: Apache-2.0
# ViT3D decoder for the MiniMax H3 visual VAE (inference-only bundle).
from contextlib import nullcontext
import torch
import torch.distributed as dist
import torch.nn as nn
@@ -22,6 +24,10 @@ def _linear_with_module_dtype(linear, tensor, out_dtype=None):
return output
def _cuda_autocast_disabled(tensor: torch.Tensor):
return torch.autocast("cuda", enabled=False) if tensor.is_cuda else nullcontext()
def _pack_tensors_3d(tensors, patch_size, patch_size_t):
batch_size, num_channels_tensors, temporal, height, width = tensors.shape
@@ -264,7 +270,7 @@ class ViT3DDecoder(ViTBase):
hidden_states = _pack_tensors_3d(x, 1, 1)
latent_size = (latent_T, latent_H, latent_W)
with torch.autocast("cuda", enabled=False):
with _cuda_autocast_disabled(hidden_states):
hidden_states = _linear_with_module_dtype(
self.x_embedder, hidden_states, hidden_states.dtype
)
@@ -339,7 +345,7 @@ class ViT3DDecoder(ViTBase):
hidden_states = self.apply_mask_postprocess(hidden_states, num_patches)
with torch.autocast("cuda", enabled=False):
with _cuda_autocast_disabled(hidden_states):
output = _linear_with_module_dtype(
self.proj_out, hidden_states, hidden_states.dtype
)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py
+5 -1
View File
@@ -1636,7 +1636,11 @@ class DenoisingStage(PipelineStage, RolloutDenoisingMixin):
# deallocate transformer if on mps
pipeline = self.pipeline() if self.pipeline else None
if torch.backends.mps.is_available() and not is_warmup:
if (
torch.backends.mps.is_available()
and not is_warmup
and not is_layerwise_offloaded_module(self.transformer)
):
logger.info(
"Memory before deallocating transformer: %s",
torch.mps.current_allocated_memory(),
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/reference_encoding.py
+4
View File
@@ -70,6 +70,8 @@ class _AudioVAEDeterminismContext:
_saved: tuple | None = None
def __enter__(self):
if not torch.cuda.is_available():
return self
if _AudioVAEDeterminismContext._depth == 0:
b = torch.backends
_AudioVAEDeterminismContext._saved = (
@@ -94,6 +96,8 @@ class _AudioVAEDeterminismContext:
return self
def __exit__(self, exc_type, exc, tb):
if not torch.cuda.is_available():
return
_AudioVAEDeterminismContext._depth -= 1
if _AudioVAEDeterminismContext._depth == 0:
b = torch.backends
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/release_metadata.py
+5
View File
@@ -147,6 +147,11 @@ class MiniMaxH3PartitionAdmissionStage(PipelineStage):
if not isinstance(task, str) or not task.strip():
raise ValueError("MiniMax H3 request task must be a non-empty string")
self.metadata.canonical_task(task)
if batch.num_inference_steps < 2:
raise ValueError(
"MiniMax H3 requires num_inference_steps >= 2 because its "
"video/audio sigma schedules include both interval endpoints"
)
quality = getattr(batch.sampling_params, "quality", "lossless")
if quality not in QUALITY_LEVELS:
raise ValueError(
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py
+21 -10
View File
@@ -3,6 +3,7 @@ from __future__ import annotations
import functools
from collections.abc import Mapping
from contextlib import nullcontext
import torch
@@ -296,11 +297,16 @@ class MiniMaxH3DecodingStage(DecodingStage):
audio_latent.device.type == "cuda"
and autocast_enabled(audio_vae_dtype, server_args.disable_autocast)
)
with torch.autocast(
device_type=audio_latent.device.type,
dtype=audio_vae_dtype,
enabled=audio_autocast_enabled,
):
autocast_context = (
torch.autocast(
device_type="cuda",
dtype=audio_vae_dtype,
enabled=audio_autocast_enabled,
)
if audio_latent.is_cuda
else nullcontext()
)
with autocast_context:
audio_decode = self._get_vae_decode_fn(
audio_vae,
server_args,
@@ -352,11 +358,16 @@ class MiniMaxH3DecodingStage(DecodingStage):
)
if visual_autocast_enabled:
selected_video_vae.prepare_decoder_autocast_weights(video_vae_dtype)
with torch.autocast(
device_type=visual_latent.device.type,
dtype=video_vae_dtype,
enabled=visual_autocast_enabled,
):
autocast_context = (
torch.autocast(
device_type="cuda",
dtype=video_vae_dtype,
enabled=visual_autocast_enabled,
)
if visual_latent.is_cuda
else nullcontext()
)
with autocast_context:
video_decode = self._get_vae_decode_fn(
selected_video_vae,
server_args,
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/denoising.py
+4 -3
View File
@@ -35,6 +35,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
VerificationResult,
)
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
from sglang.multimodal_gen.runtime.utils.nvtx_pytorch_hooks import maybe_nvtx_range
@@ -605,9 +606,9 @@ class MiniMaxH3DenoisingStage(DenoisingStage):
ctx = _resolve_full_loop_context(batch)
if not torch.cuda.is_available():
raise RuntimeError("MiniMax H3 full-loop denoise requires CUDA")
device = torch.device("cuda")
if not (current_platform.is_cuda() or current_platform.is_mps()):
raise RuntimeError("MiniMax H3 full-loop denoise requires CUDA or MPS")
device = current_platform.get_local_torch_device()
sigmas_video = [float(v) for v in ctx.sigmas["video"]]
self._maybe_enable_cache_dit_and_torch_compile(
len(sigmas_video) - 1,
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/text_encoding.py
+3
View File
@@ -18,6 +18,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.m
from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import (
TextEncodingStage,
)
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -55,6 +56,8 @@ class MiniMaxH3TextEncodingStage(TextEncodingStage):
try:
self._encode_from_plan(batch, plan)
self._publish_native_text_conditioning(batch)
if current_platform.is_mps():
self._finish_active_component_use()
except Exception:
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.material_io import (
minimax_h3_cleanup_temp_dirs,
python/sglang/multimodal_gen/runtime/platforms/mps.py
+3 -3
View File
@@ -46,15 +46,15 @@ class MpsPlatform(Platform):
@classmethod
def get_device_capability(cls, device_id: int = 0) -> DeviceCapability | None:
raise NotImplementedError
return None
@classmethod
def get_device_name(cls, device_id: int = 0) -> str:
raise NotImplementedError
return "Apple Silicon MPS"
@classmethod
def get_device_uuid(cls, device_id: int = 0) -> str:
raise NotImplementedError
return "mps"
@classmethod
@lru_cache(maxsize=1)
python/sglang/multimodal_gen/runtime/server_args/server_args.py
+6 -20
View File
@@ -1310,31 +1310,17 @@ class ServerArgs(DisaggServerArgsMixin):
def _adjust_platform_specific(self):
if current_platform.is_mps():
if self.num_gpus != 1:
raise ValueError("MPS currently supports only --num-gpus 1")
if self.component_residency is not None and any(
mode != RESIDENT for mode in self.component_residency.values()
mode not in (RESIDENT, LAYERWISE_OFFLOAD)
for mode in self.component_residency.values()
):
raise ValueError(
"--component-residency offload modes require CUDA; "
"MPS supports only resident components"
"MPS supports only resident or layerwise-offload component "
"residency"
)
self.use_fsdp_inference = False
self.dit_layerwise_offload = False
self.layerwise_offload_components = None
if (
self.dit_cpu_offload
or self.text_encoder_cpu_offload
or self.image_encoder_cpu_offload
or self.vae_cpu_offload
):
logger.warning(
"Disabling component CPU offload on MPS because the component "
"residency offload strategy is only validated on CUDA."
)
self.dit_cpu_offload = False
self.text_encoder_cpu_offload = False
self.image_encoder_cpu_offload = False
self.vae_cpu_offload = False
self.cpu_offload_components = None
def is_arg_explicitly_set(self, arg_name: str) -> bool:
return arg_name in self._explicit_arg_names
python/sglang/multimodal_gen/runtime/utils/hf_diffusers_utils.py
+24
View File
@@ -925,6 +925,7 @@ def maybe_download_model(
local_path = snapshot_download(
repo_id=model_name_or_path,
ignore_patterns=["*.onnx", "*.msgpack"],
allow_patterns=allow_patterns,
local_dir=local_dir,
local_files_only=True,
max_workers=8,
@@ -1034,6 +1035,7 @@ def maybe_download_model(
local_path = snapshot_download(
repo_id=model_name_or_path,
ignore_patterns=["*.onnx", "*.msgpack"],
allow_patterns=allow_patterns,
local_dir=local_dir,
max_workers=8,
force_download=True,
@@ -1084,6 +1086,28 @@ def maybe_download_model(
wait_time,
)
time.sleep(wait_time)
except RuntimeError as e:
if "client has been closed" not in str(e).lower():
raise ValueError(
f"Could not find model at {model_name_or_path} and failed to download from {_model_hub_name()}: {e}"
) from e
if attempt == MAX_RETRIES - 1:
raise ValueError(
f"Could not find model at {model_name_or_path} and failed to download from {_model_hub_name()} "
f"after {MAX_RETRIES} attempts due to network error: {e}"
) from e
from huggingface_hub.utils._http import close_session
close_session()
wait_time = 2**attempt
logger.warning(
"Download failed (attempt %d/%d) because the Hugging Face client was closed. "
"Retrying in %d seconds...",
attempt + 1,
MAX_RETRIES,
wait_time,
)
time.sleep(wait_time)
except Exception as e:
raise ValueError(
f"Could not find model at {model_name_or_path} and failed to download from {_model_hub_name()}: {e}"
python/sglang/multimodal_gen/runtime/utils/perf_logger.py
+10
View File
@@ -134,6 +134,16 @@ def capture_memory_snapshot() -> MemorySnapshot:
peak_reserved_mb=0.0,
)
if current_platform.is_mps():
allocated = torch.mps.current_allocated_memory()
reserved = torch.mps.driver_allocated_memory()
return MemorySnapshot(
allocated_mb=allocated / (1024**2),
reserved_mb=reserved / (1024**2),
peak_allocated_mb=allocated / (1024**2),
peak_reserved_mb=reserved / (1024**2),
)
allocated = torch.get_device_module().memory_allocated()
reserved = torch.get_device_module().memory_reserved()
peak_allocated = torch.get_device_module().max_memory_allocated()
python/sglang/multimodal_gen/test/unit/test_minimax_h3_admission.py
+31
View File
@@ -19,6 +19,10 @@ from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import (
from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import (
AttentionRequirements,
)
from sglang.multimodal_gen.runtime.managers.memory_managers.component_residency import (
LAYERWISE_OFFLOAD,
RESIDENT,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.release_metadata import (
MiniMaxH3PartitionAdmissionStage,
MiniMaxH3ReleaseMetadata,
@@ -369,3 +373,30 @@ def test_validate_server_args_requires_packed_varlen_backend():
):
with pytest.raises(ValueError, match="does not implement packed varlen"):
MiniMaxH3PipelineConfig.validate_server_args(config, server_args)
def test_mps_admission_requires_layerwise_residency_for_every_h3_component():
config = SimpleNamespace(
vae_config=SimpleNamespace(resolved_parallel_decode_mode=lambda: None),
dit_config=SimpleNamespace(arch_config=SimpleNamespace(attention_head_dim=128)),
_server_arg_value=MiniMaxH3PipelineConfig._server_arg_value,
)
modes = {
"transformer": LAYERWISE_OFFLOAD,
"text_encoder": LAYERWISE_OFFLOAD,
"video_vae": LAYERWISE_OFFLOAD,
"audio_vae": LAYERWISE_OFFLOAD,
}
server_args = SimpleNamespace(
component_attention_backends={},
attention_backend=None,
enable_torch_compile=False,
residency_mode=modes.get,
)
with patch.object(current_platform, "is_mps", return_value=True):
MiniMaxH3PipelineConfig.validate_server_args(config, server_args)
modes["audio_vae"] = RESIDENT
with pytest.raises(ValueError, match="audio_vae"):
MiniMaxH3PipelineConfig.validate_server_args(config, server_args)