From a86edcdc0a69679b4927b4dcefbc64dd682bf0dd Mon Sep 17 00:00:00 2001 From: triple-mu Date: Fri, 14 Aug 2026 15:33:51 +0800 Subject: [PATCH] [diffusion] feat: rebuild minimax-h3 adaln outputs on demand (#34650) Co-authored-by: Mick --- .../cookbook/diffusion/MiniMax/MiniMax-H3.mdx | 46 ++ docs/docs/sglang-diffusion/api/cli.mdx | 1 + .../component_loaders/transformer_loader.py | 40 ++ .../runtime/loader/fsdp_load.py | 9 +- .../runtime/loader/weight_utils.py | 4 + .../runtime/models/dits/minimax_h3.py | 409 +++++++++++++++++- .../minimax_h3/denoise_loop.py | 5 + .../minimax_h3/stages/decoding.py | 14 +- .../runtime/server_args/server_args.py | 39 ++ .../test/unit/test_fsdp_load.py | 1 + .../test/unit/test_minimax_h3_adaln_cache.py | 59 +++ .../tools/build_minimax_h3_adaln_cache.py | 268 ++++++++++++ 12 files changed, 872 insertions(+), 23 deletions(-) create mode 100644 python/sglang/multimodal_gen/test/unit/test_minimax_h3_adaln_cache.py create mode 100644 python/sglang/multimodal_gen/tools/build_minimax_h3_adaln_cache.py diff --git a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx index 77711c9c5..69fd2aed7 100644 --- a/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx +++ b/docs/cookbook/diffusion/MiniMax/MiniMax-H3.mdx @@ -174,6 +174,52 @@ server environment. For MiniMax-H3, `--performance-mode speed` deliberately keeps the DiT eager. The current `torch.compile` path changes the model's numerical output, so it is not enabled implicitly by any recommended lossless preset. An explicit `--enable-torch-compile true` remains available for controlled experiments, but it should not be used to generate consistency ground truth. +### Advanced: precomputed AdaLN cache + +The [model card](https://huggingface.co/MiniMaxAI/MiniMax-H3) notes that about +13B H3 parameters are AdaLN branches whose outputs can be precomputed for +inference. The public base checkpoint contains the original branches, not a +ready-to-use cache. SGLang therefore keeps the standard path as the default. + + +This is an experimental deployment path. It is intentionally disabled unless +you provide an explicitly generated cache; end-to-end numerical and peak-memory +validation remains required before using it in production. + + +When an inference-only deployment has a fixed sampling schedule, build a cache +from the already materialized transformer directory on CUDA, then pass it to +the usual `sglang serve` command. This does not alter the denoising formula: +the cache stores the BF16 outputs of the original AdaLN linears. + +```bash Command +python -m sglang.multimodal_gen.tools.build_minimax_h3_adaln_cache \ + --transformer-path "$TRANSFORMER_PATH" \ + --model-variant fl2va \ + --mode t2va \ + --num-inference-steps 50 \ + --flow-shift 12 \ + --audio-flow-shift 3 \ + --output /models/minimax-h3-fl2va-adaln-50step.safetensors + +sglang serve \ + --model-path MiniMaxAI/MiniMax-H3 \ + --model-variant fl2va \ + --minimax-h3-adaln-cache-path /models/minimax-h3-fl2va-adaln-50step.safetensors \ + --num-gpus 4 \ + --tp-size 2 \ + --ulysses-degree 2 \ + --port 30010 +``` + +`$TRANSFORMER_PATH` is the `FL2VA/transformer` or `Ref2VA/transformer` +directory in the normal SGLang/Hugging Face snapshot; the builder never +downloads a second copy. A cache only covers the scheduler settings used to +create it, including its mode, step count, flow shifts, and condition noise +values. SGLang rejects a request outside that coverage instead of silently +changing conditioning. Cache mode supports the matching unquantized checkpoint +only. + ## 4. Generate video and audio MiniMax-H3 uses the asynchronous OpenAI-compatible video endpoint. Choose a diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 2e150b3f8..a383fc543 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -77,6 +77,7 @@ Use `sglang generate --help` and `sglang serve --help` for the full argument lis - `--model-path {MODEL}`: model path or Hugging Face model ID - `--served-model-name {NAME}`: stable model name exposed by serving APIs. Defaults to `--model-id` when set, otherwise `--model-path`. - `--model-variant {NAME}`: semantic checkpoint variant to load when one model repository contains multiple weight partitions. The pipeline maps this stable name to the repository layout before loading; for example, MiniMax-H3 accepts `fl2va` and `ref2va`. This is a server/load-time choice, unlike a request's `task`. +- `--minimax-h3-adaln-cache-path {FILE}`: advanced MiniMax-H3-only inference cache. It replaces the checkpoint's AdaLN projection weights with precomputed outputs and only accepts requests whose exact FP32 timestep plan is included in the cache. It requires unquantized weights and the matching model variant. - `--model-subfolder {PATH}`: advanced direct override for a component subfolder inside the model repository. Prefer `--model-variant` when the pipeline exposes semantic routing. If both are supplied, they must resolve to the same weight partition. - `--lora-path {PATH}` and `--lora-nickname {NAME}`: load a LoRA adapter - `--lora-weight-name {FILE}`: select one adapter file from a repository that contains multiple LoRA revisions. The Hub download is filtered to that file plus JSON metadata, so unused weights are not downloaded. diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py index ba2216eb2..caab1ac91 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/transformer_loader.py @@ -1,5 +1,6 @@ import copy import logging +from collections.abc import Callable from contextlib import nullcontext from typing import Any @@ -50,6 +51,10 @@ def _resolve_checkpoint_load_device( return runtime_device +def _minimax_h3_adaln_cache_key_filter(name: str) -> bool: + return ".adaln_proj.linear." not in name + + def _default_quantized_attention_backend( quant_spec: TransformerQuantLoadSpec, server_args: ServerArgs ) -> AttentionBackendEnum | None: @@ -203,6 +208,40 @@ class TransformerLoader(ComponentLoader): "hf_config": config, "quant_config": quant_spec.runtime_quant_config, } + checkpoint_key_filter: Callable[[str], bool] | None = None + adaln_cache_path = component_server_args.minimax_h3_adaln_cache_path + if adaln_cache_path is not None: + if cls_name != "MiniMaxH3DiTModel": + raise ValueError( + "--minimax-h3-adaln-cache-path is only supported by MiniMax H3" + ) + if component_server_args.model_variant not in ("fl2va", "ref2va"): + raise ValueError( + "MiniMax H3 AdaLN cache requires --model-variant fl2va or ref2va" + ) + init_params["adaln_cache_path"] = adaln_cache_path + init_params["adaln_cache_model_variant"] = ( + component_server_args.model_variant + ) + checkpoint_key_filter = _minimax_h3_adaln_cache_key_filter + if component_server_args.minimax_h3_adaln_online: + if cls_name != "MiniMaxH3DiTModel": + raise ValueError( + "--minimax-h3-adaln-online is only supported by MiniMax H3" + ) + if adaln_cache_path is not None: + raise ValueError( + "--minimax-h3-adaln-online and --minimax-h3-adaln-cache-path " + "are mutually exclusive" + ) + # Keep the weights off-device; the model rebuilds the AdaLN + # outputs from the checkpoint for each request's timestep plan. + init_params["adaln_weight_files"] = safetensors_list + init_params["adaln_plan_width"] = ( + component_server_args.minimax_h3_adaln_plan_width + ) + checkpoint_key_filter = _minimax_h3_adaln_cache_key_filter + if ( init_params["quant_config"] is None and component_server_args.transformer_weights_path is not None @@ -273,6 +312,7 @@ class TransformerLoader(ComponentLoader): output_dtype=None, strict=False, weight_load_plan=weight_load_plan, + checkpoint_key_filter=checkpoint_key_filter, ) # post-hooks (e.g., patch scales (nunchaku)) diff --git a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py index 695f30a80..e91052e28 100644 --- a/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py +++ b/python/sglang/multimodal_gen/runtime/loader/fsdp_load.py @@ -237,6 +237,7 @@ def maybe_load_fsdp_model( pin_cpu_memory: bool = True, strict: bool = True, weight_load_plan: WeightLoadPlan | None = None, + checkpoint_key_filter: Callable[[str], bool] | None = None, ) -> torch.nn.Module: """Load a model with optional FSDP (Fully Sharded Data Parallel) support. @@ -347,6 +348,7 @@ def maybe_load_fsdp_model( and use_fsdp and weight_dir_list and preprocess_loaded_state_dict is None + and checkpoint_key_filter is None and not is_bnb_quantized ): preconverted_state_dict = ( @@ -361,6 +363,7 @@ def maybe_load_fsdp_model( and not use_fsdp and weight_dir_list and preprocess_loaded_state_dict is None + and checkpoint_key_filter is None and not is_bnb_quantized ): preconverted_state_dict = ( @@ -375,10 +378,14 @@ def maybe_load_fsdp_model( if weight_load_plan.load_full_state_dict_on_device: weight_iterator = safetensors_weights_iterator( weight_dir_list, + key_filter=checkpoint_key_filter, weight_load_plan=weight_load_plan, ) else: - weight_iterator = safetensors_weights_iterator(weight_dir_list) + weight_iterator = safetensors_weights_iterator( + weight_dir_list, + key_filter=checkpoint_key_filter, + ) if preprocess_loaded_state_dict is not None: weight_iterator = preprocess_loaded_state_dict(weight_iterator) if is_bnb_quantized: diff --git a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py index ce0e6ea5e..18b27fae0 100644 --- a/python/sglang/multimodal_gen/runtime/loader/weight_utils.py +++ b/python/sglang/multimodal_gen/runtime/loader/weight_utils.py @@ -228,6 +228,10 @@ def safetensors_weights_iterator( use_runai_model_streamer = ( HAS_RUNAI_MODEL_STREAMER and envs.SGLANG_USE_RUNAI_MODEL_STREAMER ) + if key_filter is not None: + # streamer filters after materializing all tensors, so it cannot skip + # a checkpoint partition at load time + use_runai_model_streamer = False # Validate files before loading corrupted_files, duplicate_files_by_key = _scan_safetensors_files(hf_weights_files) diff --git a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py index 26d9f3dc5..1de3f0287 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py @@ -8,10 +8,14 @@ contract accepts packed inference keyword arguments and returns packed logits. from __future__ import annotations import math -from typing import Any +import os +import struct +from contextlib import ExitStack +from typing import Any, Callable import torch import torch.nn as nn +from safetensors.torch import safe_open from sglang.kernels.ops.activation.activation import ( silu_and_mul_with_activation_rounding_, @@ -40,6 +44,7 @@ from sglang.multimodal_gen.runtime.distributed import ( ) from sglang.multimodal_gen.runtime.distributed.parallel_state import ( get_ring_ctx, + get_tp_rank, get_ulysses_ctx, ) from sglang.multimodal_gen.runtime.layers.attention.backends.attention_backend import ( @@ -64,10 +69,13 @@ from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import ( eager_on_graph, ) +logger = init_logger(__name__) + _ARCH_DEFAULTS = MiniMaxH3DiTArchConfig() _BF16_DTYPE = torch.bfloat16 _FP32_DTYPE = torch.float32 @@ -802,6 +810,289 @@ class MiniMaxH3AdalnProj(nn.Module): return self.split_output(x) +# A ref2va request carrying both a visual and an audio reference reaches four +# distinct timesteps in one step: video, audio, the imgvid condition and the +# audio reference. That is the widest case, so it is the default; a deployment +# serving only narrower tasks (t2va reaches 2, fl2va 3) can shrink the slab +# proportionally via --minimax-h3-adaln-plan-width. +MINIMAX_H3_ADALN_MAX_PLAN_WIDTH = 4 + + +def _plan_key(timesteps: torch.Tensor) -> tuple[int, ...]: + """One denoise step's unique timesteps as their exact fp32 bit patterns.""" + return tuple( + struct.unpack(" None: + super().__init__() + if (path is None) == (weight_files is None): + raise ValueError( + "MiniMax H3 AdaLN cache takes exactly one of path (prebuilt " + "sidecar) or weight_files (rebuild from the checkpoint)" + ) + self.path = path + self.model_variant = model_variant + self.weight_files = weight_files + self.max_plans = max_plans + self.max_plan_width = max_plan_width + self.num_layers = arch.num_layers + self.hidden_size = arch.hidden_size + self.block_width = 6 * MINIMAX_H3_ADALN_MODALITY_NUM * arch.hidden_size + self.final_width = 2 * arch.hidden_size + # Rebuild path only: plan bit pattern -> slot, tracked on the host. + self._slots: dict[tuple[int, ...], int] = {} + self.rebuilds = 0 + + def load(self, device: torch.device) -> None: + if self.path is None: + self._allocate(device) + return + if not os.path.isfile(self.path): + raise ValueError(f"MiniMax H3 AdaLN cache does not exist: {self.path}") + + with safe_open(self.path, framework="pt", device="cpu") as cache_file: + metadata = cache_file.metadata() or {} + if metadata.get("format_version") != self._FORMAT_VERSION: + raise ValueError( + "MiniMax H3 AdaLN cache has an unsupported or missing format_version" + ) + cache_variant = metadata.get("model_variant") + if self.model_variant is not None and cache_variant != self.model_variant: + raise ValueError( + "MiniMax H3 AdaLN cache model_variant does not match the loaded " + f"variant ({cache_variant!r} != {self.model_variant!r})" + ) + plan_timesteps = cache_file.get_tensor("plan_timesteps") + plan_lengths = cache_file.get_tensor("plan_lengths") + block_params = cache_file.get_tensor("block_params") + final_params = cache_file.get_tensor("final_params") + + expected_block_width = 6 * MINIMAX_H3_ADALN_MODALITY_NUM * self.hidden_size + expected_final_width = 2 * self.hidden_size + if ( + plan_timesteps.dtype != _FP32_DTYPE + or plan_timesteps.ndim != 2 + or plan_lengths.dtype != torch.int64 + or plan_lengths.shape != (plan_timesteps.shape[0],) + or (plan_lengths < 1).any() + or (plan_lengths > plan_timesteps.shape[1]).any() + ): + raise ValueError("MiniMax H3 AdaLN cache has invalid timestep plans") + if block_params.dtype != _BF16_DTYPE or block_params.shape != ( + plan_timesteps.shape[0], + plan_timesteps.shape[1], + self.num_layers, + expected_block_width, + ): + raise ValueError("MiniMax H3 AdaLN cache has invalid block_params") + if final_params.dtype != _BF16_DTYPE or final_params.shape != ( + plan_timesteps.shape[0], + plan_timesteps.shape[1], + expected_final_width, + ): + raise ValueError("MiniMax H3 AdaLN cache has invalid final_params") + + self.register_buffer("plan_timesteps", plan_timesteps.to(device)) + self.register_buffer("plan_lengths", plan_lengths.to(device)) + self.register_buffer("block_params", block_params.to(device)) + self.register_buffer("final_params", final_params.to(device)) + + def _allocate(self, device: torch.device) -> None: + """Empty slab for the rebuild path; its pointers must never move. + + ``plan_lengths`` starts at zero and that is what keeps unused slots out + of ``lookup``: a real plan always has at least one timestep, so a zero + length can never match. Breakable CUDA graph keys its replay signature + on tensor pointers, so this is allocated once and only written in place. + """ + width = self.max_plan_width + self.register_buffer( + "plan_timesteps", + torch.zeros((self.max_plans, width), dtype=_FP32_DTYPE, device=device), + ) + self.register_buffer( + "plan_lengths", + torch.zeros((self.max_plans,), dtype=torch.int64, device=device), + ) + self.register_buffer( + "block_params", + torch.zeros( + (self.max_plans, width, self.num_layers, self.block_width), + dtype=_BF16_DTYPE, + device=device, + ), + ) + self.register_buffer( + "final_params", + torch.zeros( + (self.max_plans, width, self.final_width), + dtype=_BF16_DTYPE, + device=device, + ), + ) + logger.info( + "MiniMax H3 AdaLN rebuild slab: %d plans x %d timesteps = %.2f GiB", + self.max_plans, + width, + self.block_params.numel() * 2 / 2**30, + ) + + def build( + self, + step_timesteps: list[torch.Tensor], + *, + embed: Callable[[torch.Tensor], torch.Tensor], + ) -> None: + """Fill every plan this request will look up, in one streaming pass. + + Each plan keeps its own timestep count as the GEMM batch size, because + cuBLAS selects kernels by shape and the selection is not monotonic in M: + against the runtime's M == 2, results at M == 4/8/16/64/96 are + bit-identical while M == 32 differs in 11760 of 96768 elements and + M == 1 (the GEMV path the first denoise step takes) differs in 69. + Rebuilding a plan at any other batch size silently perturbs the output. + + The pass reads all 50 adaln_proj layers regardless of how many plans are + missing, so a request builds everything it needs before denoising rather + than filling in step by step. + """ + wanted: dict[tuple[int, ...], torch.Tensor] = {} + for timesteps in step_timesteps: + wanted.setdefault(_plan_key(timesteps), timesteps) + missing = {k: v for k, v in wanted.items() if k not in self._slots} + if not missing: + return + if len(self._slots) + len(missing) > self.max_plans: + # Every plan a request looks up has to stay resident for the whole + # denoise loop, so an overflow means the capacity is too small -- + # evicting part of it would only move the failure into lookup(). + self._slots.clear() + self.plan_lengths.zero_() + if len(missing) > self.max_plans: + raise ValueError( + f"MiniMax H3 AdaLN rebuild needs {len(missing)} plans but " + f"max_plans is {self.max_plans}" + ) + widest = max(timesteps.numel() for timesteps in missing.values()) + if widest > self.max_plan_width: + raise ValueError( + f"MiniMax H3 AdaLN rebuild hit a {widest}-timestep plan but the " + f"slab was allocated for {self.max_plan_width}; raise " + "--minimax-h3-adaln-plan-width (t2va needs 2, fl2va 3, ref2va 4)" + ) + + device = self.block_params.device + slots = [] + for key, timesteps in missing.items(): + slot = len(self._slots) + self._slots[key] = slot + slots.append((slot, timesteps.numel(), embed(timesteps.to(device)))) + self.plan_timesteps[slot, : timesteps.numel()] = timesteps.to(device) + + # adaln_proj is a ColumnParallelLinear: each rank owns a slice of the + # output features and all-gathers afterwards. The rebuild has to do the + # same rather than read the full width in one go -- a sharded GEMM has a + # different N, so cuBLAS picks a different kernel and the outputs stop + # matching. It also cuts per-rank checkpoint reads to 1/tp. + tp_size = get_tp_world_size() + tp_rank = get_tp_rank() if tp_size > 1 else 0 + + with ExitStack() as stack: + handles = [ + stack.enter_context(safe_open(f, framework="pt", device=str(device))) + for f in self.weight_files + ] + index = {name: h for h in handles for name in h.keys()} + + def read_shard(name: str, out_features: int) -> torch.Tensor: + if tp_size == 1: + return index[name].get_tensor(name) + shard = out_features // tp_size + start = tp_rank * shard + return index[name].get_slice(name)[start : start + shard] + + def project(adaln_input: torch.Tensor, weight, bias) -> torch.Tensor: + out = nn.functional.linear(adaln_input, weight, bias) + return tensor_model_parallel_all_gather(out) if tp_size > 1 else out + + for layer in range(self.num_layers): + prefix = f"blocks.{layer}.adaln_proj.linear" + weight = read_shard(f"{prefix}.weight", self.block_width) + bias = read_shard(f"{prefix}.bias", self.block_width) + for slot, length, adaln_input in slots: + self.block_params[slot, :length, layer] = project( + adaln_input, weight, bias + ) + del weight, bias + prefix = "final_layer.adaln_proj.linear" + weight = read_shard(f"{prefix}.weight", self.final_width) + bias = read_shard(f"{prefix}.bias", self.final_width) + for slot, length, adaln_input in slots: + self.final_params[slot, :length] = project(adaln_input, weight, bias) + del weight, bias + + for slot, length, _ in slots: + self.plan_lengths[slot] = length + self.rebuilds += 1 + logger.info( + "MiniMax H3 AdaLN: rebuilt %d plan(s), %d/%d resident, pass #%d", + len(missing), + len(self._slots), + self.max_plans, + self.rebuilds, + ) + + def lookup(self, unique_timesteps: torch.Tensor) -> torch.Tensor: + num_timesteps = unique_timesteps.shape[0] + matches = self.plan_lengths.eq(num_timesteps) & self.plan_timesteps[ + :, :num_timesteps + ].eq(unique_timesteps).all(dim=-1) + if not bool(matches.any()): + raise ValueError( + "MiniMax H3 AdaLN cache does not cover the request timestep plan" + ) + return matches.to(torch.int64).argmax() + + def block( + self, + index: int, + cache_plan_index: torch.Tensor, + num_timesteps: int, + ) -> tuple[torch.Tensor, ...]: + params = self.block_params[cache_plan_index, :num_timesteps, index] + params = params.reshape(-1, 6, self.hidden_size) + return tuple(params.unbind(dim=1)) + + def final( + self, + cache_plan_index: torch.Tensor, + num_timesteps: int, + ) -> tuple[torch.Tensor, ...]: + params = self.final_params[cache_plan_index, :num_timesteps] + return tuple(params.reshape(-1, 2, self.hidden_size).unbind(dim=1)) + + class MiniMaxH3TokenRefinerBlock(nn.Module): """Standard pre-norm transformer block without AdaLN or RoPE.""" @@ -890,6 +1181,7 @@ class MiniMaxH3DiTBlock(nn.Module): quant_config: QuantizationConfig | None, *, prefix: str, + use_adaln_cache: bool = False, ) -> None: super().__init__() self.norm1 = _norm(arch.hidden_size, eps=arch.norm_eps) @@ -900,13 +1192,17 @@ class MiniMaxH3DiTBlock(nn.Module): prefix=f"{prefix}.attn", ) self.mlp = MiniMaxH3MLP(arch, quant_config, prefix=f"{prefix}.mlp") - self.adaln_proj = MiniMaxH3AdalnProj( - arch, - arch.adaln_out_features, - quant_config, - prefix=f"{prefix}.adaln_proj", - expand_ratio=6, - modality_num=MINIMAX_H3_ADALN_MODALITY_NUM, + self.adaln_proj = ( + None + if use_adaln_cache + else MiniMaxH3AdalnProj( + arch, + arch.adaln_out_features, + quant_config, + prefix=f"{prefix}.adaln_proj", + expand_ratio=6, + modality_num=MINIMAX_H3_ADALN_MODALITY_NUM, + ) ) self.preserve_input_for_cache_dit = False @@ -932,6 +1228,8 @@ class MiniMaxH3DiTBlock(nn.Module): norm2 -> scale/shift -> MLP -> gated residual. """ if adaln_params is None: + if self.adaln_proj is None: + raise ValueError("MiniMax H3 AdaLN cache parameters are required") adaln_params = self.adaln_proj(adaln_input) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = adaln_params # Cache-DiT retains the inputs to its Fn and Mn block ranges. Only the @@ -984,6 +1282,7 @@ class MiniMaxH3FinalLayer(nn.Module): quant_config: QuantizationConfig | None, *, prefix: str, + use_adaln_cache: bool = False, ) -> None: super().__init__() video_patch_dim = ( @@ -993,13 +1292,17 @@ class MiniMaxH3FinalLayer(nn.Module): * arch.patch_size[2] ) self.norm = _norm(arch.hidden_size, eps=arch.final_norm_eps) - self.adaln_proj = MiniMaxH3AdalnProj( - arch, - arch.final_adaln_out_features, - quant_config, - prefix=f"{prefix}.adaln_proj", - expand_ratio=2, - modality_num=1, + self.adaln_proj = ( + None + if use_adaln_cache + else MiniMaxH3AdalnProj( + arch, + arch.final_adaln_out_features, + quant_config, + prefix=f"{prefix}.adaln_proj", + expand_ratio=2, + modality_num=1, + ) ) self.video_out = ColumnParallelLinear( arch.hidden_size, @@ -1026,6 +1329,7 @@ class MiniMaxH3FinalLayer(nn.Module): *, adaln_input: torch.Tensor, inverse_indices: torch.Tensor, + adaln_params: tuple[torch.Tensor, ...] | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """Project all rows into TP-local video/audio output shards. @@ -1034,7 +1338,11 @@ class MiniMaxH3FinalLayer(nn.Module): The model gathers output columns only after selecting live media rows, preserving the GEMM shape while reducing collective payload. """ - shift, scale = self.adaln_proj(adaln_input) + if adaln_params is None: + if self.adaln_proj is None: + raise ValueError("MiniMax H3 AdaLN cache parameters are required") + adaln_params = self.adaln_proj(adaln_input) + shift, scale = adaln_params h = self.norm(x) h = _modulate_scale_shift(h, shift, scale, inverse_indices, dtype=_BF16_DTYPE) # Preserve full precision through both final output projections. @@ -1056,9 +1364,25 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): reverse_param_names_mapping = _ARCH_DEFAULTS.reverse_param_names_mapping lora_param_names_mapping = _ARCH_DEFAULTS.lora_param_names_mapping + def prepare_adaln_plans(self, step_timesteps: list[torch.Tensor]) -> None: + """Fill the AdaLN cache for this request before denoising starts. + + No-op for a prebuilt sidecar; the rebuild path needs the model's own + timestep embedding so a filled plan is bit-identical to what resident + adaln_proj weights would have produced. + """ + if self.adaln_cache is None or self.adaln_cache.weight_files is None: + return + + def embed(timesteps: torch.Tensor) -> torch.Tensor: + return nn.functional.silu(self.time_embedder(timesteps)).to(_BF16_DTYPE) + + self.adaln_cache.build(step_timesteps, embed=embed) + def _can_batch_block_adaln(self) -> bool: return ( - get_tp_world_size() > 1 + self.adaln_cache is None + and get_tp_world_size() > 1 and not torch.compiler.is_compiling() and not envs.SGLANG_CACHE_DIT_ENABLED and not hasattr(self, "_sglang_cache_dit_adapter") @@ -1134,8 +1458,21 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): config: MiniMaxH3DiTConfig, hf_config: dict[str, Any], quant_config: QuantizationConfig | None = None, + adaln_cache_path: str | None = None, + adaln_cache_model_variant: str | None = None, + adaln_weight_files: list[str] | None = None, + adaln_plan_width: int = MINIMAX_H3_ADALN_MAX_PLAN_WIDTH, ) -> None: super().__init__(config=config, hf_config=hf_config) + if ( + adaln_cache_path is not None or adaln_weight_files is not None + ) and quant_config is not None: + raise ValueError( + "MiniMax H3 AdaLN cache is only compatible with unquantized weights" + ) + self._adaln_precomputed = ( + adaln_cache_path is not None or adaln_weight_files is not None + ) arch = self.config self.arch = arch self.hidden_size = arch.hidden_size @@ -1197,6 +1534,7 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): arch, quant_config, prefix=f"blocks.{index}", + use_adaln_cache=self._adaln_precomputed, ) for index in range(arch.num_layers) ] @@ -1206,6 +1544,18 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): arch, quant_config, prefix="final_layer", + use_adaln_cache=self._adaln_precomputed, + ) + self.adaln_cache = ( + MiniMaxH3AdalnCache( + arch, + path=adaln_cache_path, + model_variant=adaln_cache_model_variant, + weight_files=adaln_weight_files, + max_plan_width=adaln_plan_width, + ) + if self._adaln_precomputed + else None ) self._resolved_attention_backend: AttentionBackendEnum | None = None self._mark_missing_params_required() @@ -1256,6 +1606,8 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): raise ValueError( f"rope.inv_freq must stay fp32 after load, got {rope_inv_freq.dtype}." ) + if self.adaln_cache is not None: + self.adaln_cache.load(self.video_patch_proj.weight.device) @staticmethod def _pos_ids(pos_info: Any, key: str) -> torch.Tensor: @@ -1685,7 +2037,20 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): hidden = decoder_input cu_seqlens = cu_seqlens.to(device) block_adaln_params = None - if self._can_batch_block_adaln(): + adaln_cache_plan_index = None + if self.adaln_cache is not None: + adaln_cache_plan_index = self.adaln_cache.lookup( + unique_timesteps.view(-1).to(device) + ) + block_adaln_params = tuple( + self.adaln_cache.block( + index, + adaln_cache_plan_index, + adaln_input.shape[0], + ) + for index in range(len(self.blocks)) + ) + elif self._can_batch_block_adaln(): local_adaln = torch.stack( [block.adaln_proj.project_local(adaln_input) for block in self.blocks] ) @@ -1718,6 +2083,14 @@ class MiniMaxH3DiTModel(BaseDiT, LayerwiseOffloadableModuleMixin): hidden, adaln_input=adaln_input, inverse_indices=block_inverse, + adaln_params=( + None + if adaln_cache_plan_index is None + else self.adaln_cache.final( + adaln_cache_plan_index, + adaln_input.shape[0], + ) + ), ) if sp_ws > 1: from sglang.multimodal_gen.runtime.distributed.parallel_state import ( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py index 147340b61..8f914faf9 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/denoise_loop.py @@ -449,6 +449,11 @@ def minimax_h3_denoise_loop( imgvid_cond_noise_aug=float(imgvid_cond_noise_aug_for_inference), audio_ref_cond_noise_aug=float(audio_cond_noise_aug_for_inference), ) + # Every step's timesteps are settled by now. Rebuilding AdaLN reads all + # 24.2 GiB of adaln_proj whatever is missing, so fill the whole request in + # one pass here instead of topping up step by step inside the loop. + model.prepare_adaln_plans([entry[0] for entry in timestep_plan]) + # match the scheduler's device-fp32 math once, then reuse one denoised # scratch per modality instead of allocating intermediates every step video_sigmas = torch.tensor(sigmas_video, dtype=torch.float32, device=device) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py index 77d8e074a..424b6033c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/minimax_h3/stages/decoding.py @@ -61,7 +61,7 @@ def _cached_decode_mean_std( return mean, std -def _reverse_normalize_latents_( +def _reverse_normalize_latents( latents: torch.Tensor, *, mean_values, @@ -89,7 +89,13 @@ def _reverse_normalize_latents_( ) view_shape = [1] * latents.ndim view_shape[1] = int(mean.shape[0]) - return latents.mul_(std.view(*view_shape)).add_(mean.view(*view_shape)) + # Out of place on purpose. batch.latents / batch.audio_latents are + # inference tensors allocated inside the denoising stage's InferenceMode, + # while --vae-cpu-offload runs this stage under torch.inference_mode(False) + # (PipelineExecutor._stage_needs_version_counters), where writing to one in + # place raises. Keep the mul-then-add rounding order rather than addcmul so + # the result does not shift with FMA contraction. + return latents * std.view(*view_shape) + mean.view(*view_shape) def _crop_to_target_canvas(batch: Req, frames: torch.Tensor) -> torch.Tensor: @@ -276,7 +282,7 @@ class MiniMaxH3DecodingStage(DecodingStage): if audio_vae.training: audio_vae.eval() audio_arch_config = server_args.pipeline_config.audio_vae_config.arch_config - audio_decode_latent = _reverse_normalize_latents_( + audio_decode_latent = _reverse_normalize_latents( audio_latent, mean_values=audio_arch_config.latents_mean, std_values=audio_arch_config.latents_std, @@ -332,7 +338,7 @@ class MiniMaxH3DecodingStage(DecodingStage): if selected_video_vae.training: selected_video_vae.eval() visual_arch_config = server_args.pipeline_config.vae_config.arch_config - visual_decode_latent = _reverse_normalize_latents_( + visual_decode_latent = _reverse_normalize_latents( visual_latent, mean_values=visual_arch_config.latents_mean, std_values=visual_arch_config.latents_std, diff --git a/python/sglang/multimodal_gen/runtime/server_args/server_args.py b/python/sglang/multimodal_gen/runtime/server_args/server_args.py index 6c99daead..2675b805f 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args/server_args.py @@ -293,6 +293,13 @@ class ServerArgs(DisaggServerArgsMixin): # path to pre-quantized transformer weights (single .safetensors or directory). transformer_weights_path: str | None = None + # path to precomputed MiniMax H3 AdaLN outputs for inference-only serving. + minimax_h3_adaln_cache_path: str | None = None + # Rebuild AdaLN outputs per request from the checkpoint, no sidecar needed. + minimax_h3_adaln_online: bool = False + # Widest timestep plan the rebuild slab is sized for; see + # MINIMAX_H3_ADALN_MAX_PLAN_WIDTH. + minimax_h3_adaln_plan_width: int = 4 # Per-component transformer weight overrides (key = model_index.json component name). # Pipelines use this when a checkpoint ships separate quantized weights for # secondary DiT components; the generic loader consumes it without model-specific @@ -1495,6 +1502,38 @@ class ServerArgs(DisaggServerArgsMixin): "without exposing repository subfolder layout." ), ) + parser.add_argument( + "--minimax-h3-adaln-online", + action=StoreBoolean, + default=ServerArgs.minimax_h3_adaln_online, + help=( + "Rebuild MiniMax H3 AdaLN outputs from the checkpoint per " + "request instead of keeping the 24.2 GiB of adaln_proj weights " + "resident. Works with any step count or schedule and needs no " + "prebuilt artifact. Requires unquantized weights." + ), + ) + parser.add_argument( + "--minimax-h3-adaln-plan-width", + type=int, + default=ServerArgs.minimax_h3_adaln_plan_width, + help=( + "Widest timestep plan --minimax-h3-adaln-online sizes its slab " + "for. The default 4 covers every task; a deployment serving " + "only t2va (2) or fl2va (3) can shrink the slab proportionally. " + "A request exceeding it is rejected rather than truncated." + ), + ) + parser.add_argument( + "--minimax-h3-adaln-cache-path", + type=str, + default=ServerArgs.minimax_h3_adaln_cache_path, + help=( + "Path to a precomputed MiniMax H3 AdaLN cache. This only " + "supports the matching unquantized H3 checkpoint and rejects " + "requests whose timestep embeddings are not present in the cache." + ), + ) parser.add_argument( "--model-id", type=str, diff --git a/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py b/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py index 566552cb8..07ddc6e6a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py +++ b/python/sglang/multimodal_gen/test/unit/test_fsdp_load.py @@ -188,6 +188,7 @@ class TestOrdinaryWeightLoading(unittest.TestCase): rank_local_load.assert_not_called() weight_iterator.assert_called_once_with( ["model.safetensors"], + key_filter=None, weight_load_plan=load_plan, ) diff --git a/python/sglang/multimodal_gen/test/unit/test_minimax_h3_adaln_cache.py b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_adaln_cache.py new file mode 100644 index 000000000..c4dbc8296 --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_minimax_h3_adaln_cache.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 + +import torch +from safetensors.torch import save_file + +from sglang.multimodal_gen.configs.models.dits.minimax_h3 import ( + MiniMaxH3DiTArchConfig, +) +from sglang.multimodal_gen.runtime.models.dits.minimax_h3 import ( + MiniMaxH3AdalnCache, +) + + +def test_minimax_h3_adaln_cache_matches_bf16_embedding(tmp_path): + arch = MiniMaxH3DiTArchConfig( + num_layers=2, + hidden_size=4, + time_embed_dim=3, + ) + cache_path = tmp_path / "adaln.safetensors" + plan_timesteps = torch.tensor([[0.0, 0.0], [1.0, 2.0]]) + plan_lengths = torch.tensor([1, 2], dtype=torch.int64) + block_params = ( + torch.arange(2 * 2 * 2 * 72, dtype=torch.float32) + .reshape(2, 2, 2, 72) + .bfloat16() + ) + final_params = torch.arange(32, dtype=torch.float32).reshape(2, 2, 8).bfloat16() + save_file( + { + "plan_timesteps": plan_timesteps, + "plan_lengths": plan_lengths, + "block_params": block_params, + "final_params": final_params, + }, + cache_path, + metadata={"format_version": "2", "model_variant": "fl2va"}, + ) + + cache = MiniMaxH3AdalnCache( + arch, + path=str(cache_path), + model_variant="fl2va", + ) + cache.load(torch.device("cpu")) + + cache_plan_index = cache.lookup(plan_timesteps[1]) + block = cache.block(1, cache_plan_index, 2) + final = cache.final(cache_plan_index, 2) + + # block() hands the forward pass six [num_timesteps * modality, hidden] + # chunks, while the checkpoint stores a plan as one flat + # [num_timesteps, 6 * modality * hidden] row -- same elements, and the + # modality axis folds into the leading one rather than staying separate. + assert torch.equal( + torch.cat(block, dim=-1).reshape(block_params[1, :, 1].shape), + block_params[1, :, 1], + ) + assert torch.equal(torch.cat(final, dim=-1), final_params[1]) diff --git a/python/sglang/multimodal_gen/tools/build_minimax_h3_adaln_cache.py b/python/sglang/multimodal_gen/tools/build_minimax_h3_adaln_cache.py new file mode 100644 index 000000000..c4b58c923 --- /dev/null +++ b/python/sglang/multimodal_gen/tools/build_minimax_h3_adaln_cache.py @@ -0,0 +1,268 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Build an inference-only AdaLN cache from a local MiniMax H3 transformer.""" + +from __future__ import annotations + +import argparse +import json +import math +from contextlib import ExitStack +from pathlib import Path +from typing import Any + +import torch +import torch.nn.functional as F +from safetensors.torch import safe_open, save_file + +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.denoise_loop import ( + MINIMAX_H3_AUDIO_REF_COND_TIMESTEP, + MINIMAX_H3_IMGVID_COND_TIMESTEP, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.minimax_h3.time_request import ( + minimax_h3_time_shift_sigmas, +) + +_HIDDEN_SIZE = 5376 +_TIMESTEP_INPUT_DIM = 256 +_NUM_BLOCKS = 50 +_BLOCK_PARAM_WIDTH = 18 * _HIDDEN_SIZE +_FINAL_PARAM_WIDTH = 2 * _HIDDEN_SIZE +_CACHE_MODES = { + "t2va": ("video", "audio"), + "fl2va": ("video", "audio", "image"), + "ref2va-image": ("video", "audio", "image"), + "ref2va-audio": ("video", "audio", "audio_ref"), + "ref2va-mixed": ("video", "audio", "image", "audio_ref"), +} +_MODE_VARIANTS = { + "t2va": "fl2va", + "fl2va": "fl2va", + "ref2va-image": "ref2va", + "ref2va-audio": "ref2va", + "ref2va-mixed": "ref2va", +} + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Build a MiniMax H3 inference-only AdaLN cache from local weights." + ) + parser.add_argument("--transformer-path", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--model-variant", choices=("fl2va", "ref2va"), required=True) + parser.add_argument("--mode", choices=tuple(_CACHE_MODES), default="t2va") + parser.add_argument("--num-inference-steps", type=int, default=50) + parser.add_argument("--flow-shift", type=float, default=12.0) + parser.add_argument("--audio-flow-shift", type=float, default=3.0) + parser.add_argument( + "--imgvid-cond-noise-aug", + type=float, + default=MINIMAX_H3_IMGVID_COND_TIMESTEP, + ) + parser.add_argument( + "--audio-cond-noise-aug", + type=float, + default=MINIMAX_H3_AUDIO_REF_COND_TIMESTEP, + ) + parser.add_argument( + "--timesteps", + type=float, + nargs="+", + help="Override the scheduler-derived timestep plan.", + ) + parser.add_argument("--device", default="cuda") + return parser.parse_args() + + +def _cache_timestep_plans(args: argparse.Namespace) -> list[torch.Tensor]: + if args.timesteps is not None: + return [torch.tensor(args.timesteps, dtype=torch.float32).unique(sorted=True)] + + video_sigmas = minimax_h3_time_shift_sigmas( + num_steps=args.num_inference_steps, + shift_scale=args.flow_shift, + ) + audio_sigmas = minimax_h3_time_shift_sigmas( + num_steps=args.num_inference_steps, + shift_scale=args.audio_flow_shift, + ) + fields = _CACHE_MODES[args.mode] + plans = [] + for video_sigma, audio_sigma in zip(video_sigmas[:-1], audio_sigmas[:-1]): + video_timestep = 1.0 - video_sigma + audio_timestep = 1.0 - audio_sigma + candidates = { + "video": video_timestep, + "audio": audio_timestep, + "image": max(video_timestep, args.imgvid_cond_noise_aug), + "audio_ref": max(audio_timestep, args.audio_cond_noise_aug), + } + plans.append( + torch.tensor( + [candidates[field] for field in fields], dtype=torch.float32 + ).unique(sorted=True) + ) + + deduplicated = [] + seen = set() + for plan in plans: + key = tuple(plan.tolist()) + if key not in seen: + seen.add(key) + deduplicated.append(plan) + return deduplicated + + +def _time_embed( + timesteps: torch.Tensor, + *, + proj_in_weight: torch.Tensor, + proj_in_bias: torch.Tensor, + proj_out_weight: torch.Tensor, + proj_out_bias: torch.Tensor, +) -> torch.Tensor: + half = _TIMESTEP_INPUT_DIM // 2 + freqs = torch.exp( + -math.log(10000.0) + * torch.arange(half, dtype=torch.float32, device=timesteps.device) + / half + ) + args = timesteps[:, None] * freqs[None] + t_freq = torch.cat((torch.cos(args), torch.sin(args)), dim=-1) + return F.linear( + F.silu(F.linear(t_freq, proj_in_weight, proj_in_bias)), + proj_out_weight, + proj_out_bias, + ) + + +def _load_tensor( + name: str, + *, + weight_map: dict[str, str], + files: dict[str, Any], + device: torch.device, +) -> torch.Tensor: + tensor_file = files[weight_map[name]] + return tensor_file.get_tensor(name).to(device) + + +def main() -> None: + args = _parse_args() + if args.num_inference_steps < 2 and args.timesteps is None: + raise ValueError("--num-inference-steps must be at least 2") + mode_variant = _MODE_VARIANTS[args.mode] + if args.model_variant != mode_variant: + raise ValueError(f"--mode {args.mode} requires {mode_variant}") + device = torch.device(args.device) + if device.type != "cuda" or not torch.cuda.is_available(): + raise ValueError("MiniMax H3 AdaLN cache must be built on CUDA") + + index_path = args.transformer_path / "model.safetensors.index.json" + with index_path.open() as f: + weight_map = json.load(f)["weight_map"] + + plans = _cache_timestep_plans(args) + if not plans or any(plan.numel() == 0 for plan in plans): + raise ValueError("AdaLN cache must cover at least one timestep plan") + max_plan_length = max(plan.numel() for plan in plans) + plan_timesteps = torch.zeros((len(plans), max_plan_length), dtype=torch.float32) + plan_lengths = torch.tensor([plan.numel() for plan in plans], dtype=torch.int64) + block_params = torch.empty( + (len(plans), max_plan_length, _NUM_BLOCKS, _BLOCK_PARAM_WIDTH), + dtype=torch.bfloat16, + ) + final_params = torch.empty( + (len(plans), max_plan_length, _FINAL_PARAM_WIDTH), dtype=torch.bfloat16 + ) + + with ExitStack() as stack: + files = { + filename: stack.enter_context( + safe_open( + str(args.transformer_path / filename), + framework="pt", + device="cpu", + ) + ) + for filename in set(weight_map.values()) + } + time_kwargs = { + f"{module}_{name}": _load_tensor( + f"time_embedder.{module}.{name}", + weight_map=weight_map, + files=files, + device=device, + ) + for module, name in ( + ("proj_in", "weight"), + ("proj_in", "bias"), + ("proj_out", "weight"), + ("proj_out", "bias"), + ) + } + adaln_inputs = [] + for plan_index, plan in enumerate(plans): + plan_length = plan.numel() + plan_timesteps[plan_index, :plan_length].copy_(plan) + adaln_inputs.append( + F.silu(_time_embed(plan.to(device), **time_kwargs)).to(torch.bfloat16) + ) + + for index in range(_NUM_BLOCKS): + prefix = f"blocks.{index}.adaln_proj.linear" + weight = _load_tensor( + f"{prefix}.weight", + weight_map=weight_map, + files=files, + device=device, + ) + bias = _load_tensor( + f"{prefix}.bias", + weight_map=weight_map, + files=files, + device=device, + ) + for plan_index, adaln_input in enumerate(adaln_inputs): + plan_length = adaln_input.shape[0] + block_params[plan_index, :plan_length, index].copy_( + F.linear(adaln_input, weight, bias).cpu() + ) + + prefix = "final_layer.adaln_proj.linear" + weight = _load_tensor( + f"{prefix}.weight", + weight_map=weight_map, + files=files, + device=device, + ) + bias = _load_tensor( + f"{prefix}.bias", + weight_map=weight_map, + files=files, + device=device, + ) + for plan_index, adaln_input in enumerate(adaln_inputs): + plan_length = adaln_input.shape[0] + final_params[plan_index, :plan_length].copy_( + F.linear(adaln_input, weight, bias).cpu() + ) + + args.output.parent.mkdir(parents=True, exist_ok=True) + save_file( + { + "plan_timesteps": plan_timesteps, + "plan_lengths": plan_lengths, + "block_params": block_params, + "final_params": final_params, + }, + str(args.output), + metadata={ + "format_version": "2", + "model_variant": args.model_variant, + }, + ) + + +if __name__ == "__main__": + main()