Support NemotronH_Omni_Reasoning_V3 in SGLang (#35599)
Signed-off-by: Ryan Stewart <rystewart@nvidia.com> Signed-off-by: rystewart-nvidia <rystewart@nvidia.com> Co-authored-by: elvischenv <219235043+elvischenv@users.noreply.github.com> Co-authored-by: Po-Han Huang (NVIDIA) <53919306+nvpohanh@users.noreply.github.com>
This commit is contained in:
co-authored by
elvischenv
Po-Han Huang
parent
203d7e812c
commit
fae8cd84cb
@@ -1,6 +1,7 @@
|
||||
"""Config-time override declarations for nemotron_h.
|
||||
|
||||
Architectures: NemotronHForCausalLM, NemotronHPuzzleForCausalLM.
|
||||
Architectures: NemotronHForCausalLM, NemotronHPuzzleForCausalLM,
|
||||
NemotronH_Omni_Reasoning_V3.
|
||||
"""
|
||||
|
||||
import logging
|
||||
@@ -17,7 +18,11 @@ from sglang.srt.runtime_context import get_platform
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@_register_for("NemotronHForCausalLM", "NemotronHPuzzleForCausalLM")
|
||||
@_register_for(
|
||||
"NemotronHForCausalLM",
|
||||
"NemotronHPuzzleForCausalLM",
|
||||
"NemotronH_Omni_Reasoning_V3",
|
||||
)
|
||||
def _nemotron_h_overrides(server_args: Any, hf_config: Any) -> dict:
|
||||
"""NemotronH quantization / MoE runner / attention backend defaults
|
||||
(absorbed from the retired arg_groups/nemotron_h_hook.py; the mamba radix
|
||||
@@ -35,7 +40,12 @@ def _nemotron_h_overrides(server_args: Any, hf_config: Any) -> dict:
|
||||
]
|
||||
quantization = cfg.quantization
|
||||
if is_modelopt:
|
||||
assert model_config.hf_config.mlp_hidden_act == "relu2"
|
||||
language_config = (
|
||||
model_config.hf_text_config
|
||||
if model_arch == "NemotronH_Omni_Reasoning_V3"
|
||||
else hf_config
|
||||
)
|
||||
assert language_config.mlp_hidden_act == "relu2"
|
||||
if model_config.quantization == "modelopt":
|
||||
quant_algo = model_config.hf_config.quantization_config["quant_algo"]
|
||||
if quant_algo == "MIXED_PRECISION":
|
||||
|
||||
@@ -456,6 +456,7 @@ _MAMBA_RADIX_CACHE_ARCHS = frozenset(
|
||||
"MiniCPMV4_6ForConditionalGeneration",
|
||||
"NemotronHForCausalLM",
|
||||
"NemotronHPuzzleForCausalLM",
|
||||
"NemotronH_Omni_Reasoning_V3",
|
||||
"FalconH1ForCausalLM",
|
||||
"JetNemotronForCausalLM",
|
||||
"JetVLMForConditionalGeneration",
|
||||
@@ -488,6 +489,7 @@ _MAMBA_EXTRA_BUFFER_ARCHS = frozenset(
|
||||
"Glm5NextForConditionalGeneration",
|
||||
"NemotronHForCausalLM",
|
||||
"NemotronHPuzzleForCausalLM",
|
||||
"NemotronH_Omni_Reasoning_V3",
|
||||
# KDA-based: same MambaPool ping-pong machinery as GDN; requires the
|
||||
# KDA backend's track-snapshot writes (decode + extend) so donated
|
||||
# slots hold real states for prefix-cache restores.
|
||||
@@ -1001,6 +1003,7 @@ _FLASHINFER_ALLREDUCE_FUSION_ARCHS = frozenset(
|
||||
"Qwen3_5ForConditionalGeneration",
|
||||
"NemotronHForCausalLM",
|
||||
"NemotronHPuzzleForCausalLM",
|
||||
"NemotronH_Omni_Reasoning_V3",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@@ -57,6 +57,7 @@ from sglang.srt.configs.nanbeige import NanbeigeConfig
|
||||
from sglang.srt.configs.nano_nemotron_vl import (
|
||||
NemotronH_Nano_Omni_Reasoning_V3_Config,
|
||||
NemotronH_Nano_VL_V2_Config,
|
||||
NemotronH_Omni_Reasoning_V3_Config,
|
||||
)
|
||||
from sglang.srt.configs.nemotron_h import NemotronHConfig, NemotronHPuzzleConfig
|
||||
from sglang.srt.configs.olmo3 import Olmo3Config
|
||||
@@ -137,6 +138,7 @@ __all__ = [
|
||||
"NemotronHPuzzleConfig",
|
||||
"NemotronH_Nano_VL_V2_Config",
|
||||
"NemotronH_Nano_Omni_Reasoning_V3_Config",
|
||||
"NemotronH_Omni_Reasoning_V3_Config",
|
||||
"NanbeigeConfig",
|
||||
"JetNemotronConfig",
|
||||
"JetVLMConfig",
|
||||
|
||||
@@ -901,6 +901,14 @@ class ModelConfig:
|
||||
self.hf_config.architectures[0] = "ExaoneMoEForCausalLMMTP"
|
||||
self.hf_config.num_nextn_predict_layers = 1
|
||||
|
||||
if (
|
||||
is_draft_model
|
||||
and self.hf_config.architectures[0] == "NemotronH_Omni_Reasoning_V3"
|
||||
):
|
||||
self.hf_config = self.hf_text_config
|
||||
self.hf_config.architectures = ["NemotronHForCausalLMMTP"]
|
||||
self.hf_config.num_nextn_predict_layers = 1
|
||||
|
||||
if is_draft_model and self.hf_config.architectures[0] in [
|
||||
"NemotronHForCausalLM",
|
||||
"NemotronHPuzzleForCausalLM",
|
||||
@@ -2086,6 +2094,7 @@ multimodal_model_archs = [
|
||||
"MossVLForConditionalGeneration",
|
||||
"NemotronH_Nano_VL_V2",
|
||||
"NemotronH_Nano_Omni_Reasoning_V3",
|
||||
"NemotronH_Omni_Reasoning_V3",
|
||||
"MuseGlimmerForConditionalGeneration",
|
||||
"PixtralForConditionalGeneration",
|
||||
"Qwen2AudioForConditionalGeneration",
|
||||
|
||||
@@ -30,6 +30,15 @@ def float_triplet(seq: Any):
|
||||
return a, b, c
|
||||
|
||||
|
||||
def _nemotron_h_compatible_config(config: dict) -> dict:
|
||||
config = dict(config)
|
||||
aliases = {"linear_attention": "mamba", "full_attention": "attention"}
|
||||
for field in ("layers_block_type", "mtp_layers_block_type"):
|
||||
if config.get(field) is not None:
|
||||
config[field] = [aliases.get(value, value) for value in config[field]]
|
||||
return config
|
||||
|
||||
|
||||
class NemotronH_Nano_VL_V2_Config(PretrainedConfig):
|
||||
model_type = "NemotronH_Nano_VL_V2"
|
||||
is_composition = True
|
||||
@@ -166,3 +175,15 @@ class NemotronH_Nano_Omni_Reasoning_V3_Config(NemotronH_Nano_VL_V2_Config):
|
||||
# Explicit __init__ prevents PretrainedConfig.__init_subclass__ from
|
||||
# replacing the parent's custom __init__ with a dataclass-generated one.
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
|
||||
class NemotronH_Omni_Reasoning_V3_Config(NemotronH_Nano_Omni_Reasoning_V3_Config):
|
||||
model_type = "nemotron_h_omni"
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
args = list(args)
|
||||
if len(args) > 1 and args[1] is not None:
|
||||
args[1] = _nemotron_h_compatible_config(args[1])
|
||||
elif kwargs.get("llm_config") is not None:
|
||||
kwargs["llm_config"] = _nemotron_h_compatible_config(kwargs["llm_config"])
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
@@ -972,10 +972,11 @@ class Scheduler(
|
||||
initialize_mamba_selective_state_update_backend(self.server_args)
|
||||
|
||||
def init_moe_gemm_config(self):
|
||||
# For the MM models, check the text_config for MoE settings
|
||||
config_to_check = getattr(
|
||||
self.model_config.hf_config, "text_config", self.model_config.hf_config
|
||||
)
|
||||
config_to_check = self.model_config.hf_config
|
||||
if hasattr(self.model_config.hf_config, "text_config"):
|
||||
config_to_check = self.model_config.hf_config.text_config
|
||||
elif hasattr(self.model_config, "hf_text_config"):
|
||||
config_to_check = self.model_config.hf_text_config
|
||||
|
||||
# Different MoE architectures expose the per-token expert count under
|
||||
# different attribute names (e.g. Gemma4 uses ``top_k_experts``,
|
||||
|
||||
@@ -16,6 +16,7 @@
|
||||
# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/nano_nemotron_vl.py
|
||||
|
||||
import logging
|
||||
from collections import deque
|
||||
from typing import Iterable
|
||||
|
||||
import torch
|
||||
@@ -178,6 +179,18 @@ class NemotronH_Nano_VL_V2(EVS):
|
||||
x = x.permute(0, 2, 1, 3).contiguous()
|
||||
return x
|
||||
|
||||
def _normalize_vision_features(self, features: torch.Tensor) -> torch.Tensor:
|
||||
return features
|
||||
|
||||
def _load_extra_weight(self, name: str, weight: torch.Tensor) -> None:
|
||||
return
|
||||
|
||||
def get_embed_and_head(self):
|
||||
return self.language_model.get_embed_and_head()
|
||||
|
||||
def set_embed_and_head(self, embed, head):
|
||||
self.language_model.set_embed_and_head(embed, head)
|
||||
|
||||
def extract_feature_dynamic(self, pixel_values_list: list[torch.Tensor]):
|
||||
"""Extract features from variable-size images (dynamic resolution).
|
||||
|
||||
@@ -190,6 +203,7 @@ class NemotronH_Nano_VL_V2(EVS):
|
||||
offset = 0
|
||||
for i, num_patches in enumerate(num_patches_list):
|
||||
img_feats = features[0, offset : offset + num_patches]
|
||||
img_feats = self._normalize_vision_features(img_feats)
|
||||
h_patches = pixel_values_list[i].shape[-2] // patch_size
|
||||
w_patches = pixel_values_list[i].shape[-1] // patch_size
|
||||
img_feats = img_feats.reshape(1, h_patches, w_patches, -1)
|
||||
@@ -203,6 +217,7 @@ class NemotronH_Nano_VL_V2(EVS):
|
||||
def extract_video_feature_temporal(self, pixel_values, num_frames):
|
||||
"""Extract video features with temporal compression (tubelet grouping)."""
|
||||
vit_embeds = self.vision_model(pixel_values, num_frames=num_frames)
|
||||
vit_embeds = self._normalize_vision_features(vit_embeds)
|
||||
num_tubelets = vit_embeds.shape[0]
|
||||
patch_size = self.config.patch_size
|
||||
h_patches = pixel_values.shape[-2] // patch_size
|
||||
@@ -217,6 +232,9 @@ class NemotronH_Nano_VL_V2(EVS):
|
||||
def get_input_embeddings(self):
|
||||
return self.language_model.get_input_embeddings()
|
||||
|
||||
def set_dflash_layers_to_capture(self, layer_ids: list[int]) -> None:
|
||||
self.language_model.set_dflash_layers_to_capture(layer_ids)
|
||||
|
||||
def extract_feature(self, pixel_values):
|
||||
micro_batch_size = 128
|
||||
n = pixel_values.shape[0]
|
||||
@@ -229,6 +247,7 @@ class NemotronH_Nano_VL_V2(EVS):
|
||||
batch_size = chunk.shape[0]
|
||||
vit_embeds = self.vision_model(chunk)
|
||||
vit_embeds = vit_embeds.to(dtype=self.model_dtype)
|
||||
vit_embeds = self._normalize_vision_features(vit_embeds)
|
||||
vit_embeds = vit_embeds.reshape(batch_size, h_patches, w_patches, -1)
|
||||
vit_embeds = self.pixel_shuffle(
|
||||
vit_embeds, scale_factor=self.downsample_ratio
|
||||
@@ -358,36 +377,96 @@ class NemotronH_Nano_VL_V2(EVS):
|
||||
def is_sound_weights(name: str) -> bool:
|
||||
return name.startswith("sound")
|
||||
|
||||
# Separate weights by component
|
||||
llm_weights = []
|
||||
vision_weights = []
|
||||
sound_weights = []
|
||||
|
||||
for name, w in weights:
|
||||
if is_llm(name):
|
||||
# Strip 'language_model.' prefix for LLM weights
|
||||
llm_weights.append((".".join(name.split(".")[1:]), w))
|
||||
elif is_adapter_weights((name, w)):
|
||||
# Load vision-language adapter weights directly
|
||||
trimmed_name = ".".join(name.split(".")[1:])
|
||||
param = adapter_dict[trimmed_name]
|
||||
with torch.no_grad():
|
||||
default_weight_loader(param, w)
|
||||
elif is_vision_weights(name):
|
||||
# Convert: vision_model.radio_model.* → radio_model.*
|
||||
hf_key = name[len("vision_model.") :]
|
||||
vision_weights.append((hf_key, w))
|
||||
elif is_sound_weights(name):
|
||||
sound_weights.append((name, w))
|
||||
def iter_llm_weights():
|
||||
for name, w in weights:
|
||||
if is_llm(name):
|
||||
yield (".".join(name.split(".")[1:]), w)
|
||||
elif is_adapter_weights((name, w)):
|
||||
trimmed_name = ".".join(name.split(".")[1:])
|
||||
param = adapter_dict[trimmed_name]
|
||||
with torch.no_grad():
|
||||
default_weight_loader(param, w)
|
||||
elif is_vision_weights(name):
|
||||
hf_key = name[len("vision_model.") :]
|
||||
self.vision_model.load_weights([(hf_key, w)])
|
||||
elif is_sound_weights(name):
|
||||
if self.sound_encoder is not None:
|
||||
self.sound_encoder.load_weights([(name, w)])
|
||||
else:
|
||||
self._load_extra_weight(name, w)
|
||||
|
||||
llm_weights = iter_llm_weights()
|
||||
self.language_model.load_weights(llm_weights)
|
||||
self.vision_model.load_weights(vision_weights)
|
||||
if self.sound_encoder is not None and len(sound_weights) > 0:
|
||||
self.sound_encoder.load_weights(sound_weights)
|
||||
deque(llm_weights, maxlen=0)
|
||||
|
||||
|
||||
class NemotronH_Nano_Omni_Reasoning_V3(NemotronH_Nano_VL_V2):
|
||||
pass
|
||||
|
||||
|
||||
EntryClass = [NemotronH_Nano_VL_V2, NemotronH_Nano_Omni_Reasoning_V3]
|
||||
class NemotronH_Omni_Reasoning_V3(NemotronH_Nano_VL_V2):
|
||||
packed_modules_mapping = NemotronHForCausalLM.packed_modules_mapping
|
||||
_hf_projector_weight_names = {
|
||||
"vision_projector.mlp1.norm.": "mlp1.0.",
|
||||
"vision_projector.mlp1.linear1.": "mlp1.1.",
|
||||
"vision_projector.mlp1.linear2.": "mlp1.3.",
|
||||
}
|
||||
|
||||
def __init__(self, config, quant_config=None, prefix: str = ""):
|
||||
super().__init__(config, quant_config, prefix)
|
||||
self.vision_final_layernorm = (
|
||||
nn.LayerNorm(
|
||||
config.vit_hidden_size,
|
||||
eps=config.raw_vision_config.get("layer_norm_eps", 1e-6),
|
||||
).to(self.model_dtype)
|
||||
if (config.llm_config.num_nextn_predict_layers or 0) > 0
|
||||
else None
|
||||
)
|
||||
|
||||
@property
|
||||
def lm_head(self):
|
||||
return self.language_model.lm_head
|
||||
|
||||
def _normalize_vision_features(self, features: torch.Tensor) -> torch.Tensor:
|
||||
if self.vision_final_layernorm is None:
|
||||
return features
|
||||
return self.vision_final_layernorm(features)
|
||||
|
||||
def _load_extra_weight(self, name: str, weight: torch.Tensor) -> None:
|
||||
prefix = "vision_projector.vision_final_layernorm."
|
||||
if not name.startswith(prefix):
|
||||
raise ValueError(f"Unexpected Nemotron-H Omni weight: {name}")
|
||||
if self.vision_final_layernorm is None:
|
||||
raise ValueError(f"Unexpected vision projector weight: {name}")
|
||||
parameter_name = name.removeprefix(prefix)
|
||||
parameters = dict(self.vision_final_layernorm.named_parameters())
|
||||
if parameter_name not in parameters:
|
||||
raise ValueError(f"Unexpected vision projector weight: {name}")
|
||||
parameter = parameters[parameter_name]
|
||||
default_weight_loader(parameter, weight)
|
||||
|
||||
@classmethod
|
||||
def _remap_checkpoint_weight_name(cls, name: str) -> str:
|
||||
for source, target in cls._hf_projector_weight_names.items():
|
||||
if name.startswith(source):
|
||||
return name.replace(source, target, 1)
|
||||
if name.startswith("vision_model.") and not name.startswith(
|
||||
"vision_model.radio_model."
|
||||
):
|
||||
return name.replace(
|
||||
"vision_model.", "vision_model.radio_model.hf_model.", 1
|
||||
)
|
||||
return name
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
|
||||
super().load_weights(
|
||||
(self._remap_checkpoint_weight_name(name), weight)
|
||||
for name, weight in weights
|
||||
)
|
||||
|
||||
|
||||
EntryClass = [
|
||||
NemotronH_Nano_VL_V2,
|
||||
NemotronH_Nano_Omni_Reasoning_V3,
|
||||
NemotronH_Omni_Reasoning_V3,
|
||||
]
|
||||
|
||||
@@ -27,6 +27,7 @@ from sglang.srt.layers.layernorm import RMSNorm
|
||||
from sglang.srt.layers.linear import ColumnParallelLinear
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessor
|
||||
from sglang.srt.layers.quantization import QuantizationConfig
|
||||
from sglang.srt.layers.quantization.modelopt_quant import ModelOptNvFp4A16LinearMethod
|
||||
from sglang.srt.layers.vocab_parallel_embedding import (
|
||||
ParallelLMHead,
|
||||
VocabParallelEmbedding,
|
||||
@@ -38,6 +39,7 @@ from sglang.srt.models.nemotron_h import (
|
||||
NemotronHMoEDecoderLayer,
|
||||
)
|
||||
from sglang.srt.models.nemotron_h_utils import is_attn_layer
|
||||
from sglang.srt.models.utils import WeightsMapper
|
||||
from sglang.srt.runtime_context import get_parallel
|
||||
from sglang.srt.utils import add_prefix
|
||||
|
||||
@@ -293,7 +295,21 @@ class NemotronHMultiTokenPredictor(nn.Module):
|
||||
inputs_embeds: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds = self.get_input_embeddings(input_ids)
|
||||
inputs_embeds = forward_batch.mm_input_embeds
|
||||
if (
|
||||
forward_batch.forward_mode.is_extend()
|
||||
and forward_batch.contains_mm_inputs()
|
||||
and not forward_batch.forward_mode.is_draft_extend_v2()
|
||||
):
|
||||
assert inputs_embeds is not None
|
||||
last_indices = (
|
||||
forward_batch.extend_start_loc + forward_batch.extend_seq_lens - 1
|
||||
).long()
|
||||
inputs_embeds[last_indices] = self.get_input_embeddings(
|
||||
input_ids[last_indices]
|
||||
)
|
||||
if inputs_embeds is None:
|
||||
inputs_embeds = self.get_input_embeddings(input_ids)
|
||||
|
||||
hidden_states = forward_batch.spec_info.hidden_states
|
||||
residual = None
|
||||
@@ -309,6 +325,10 @@ class NemotronHMultiTokenPredictor(nn.Module):
|
||||
|
||||
|
||||
class NemotronHForCausalLMMTP(NemotronHForCausalLM):
|
||||
hf_to_sglang_mapper = NemotronHForCausalLM.hf_to_sglang_mapper | WeightsMapper(
|
||||
orig_to_new_prefix={"language_model.mtp.": "mtp."}
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config: NemotronHConfig,
|
||||
@@ -319,6 +339,7 @@ class NemotronHForCausalLMMTP(NemotronHForCausalLM):
|
||||
config = config.get_mtp_config()
|
||||
self.config = config
|
||||
self.quant_config = quant_config
|
||||
self._owns_lm_head = False
|
||||
# Required for parent's load_weights
|
||||
self.pp_group = get_pp_group()
|
||||
|
||||
@@ -366,10 +387,56 @@ class NemotronHForCausalLMMTP(NemotronHForCausalLM):
|
||||
def load_weights(
|
||||
self, weights: Iterable[tuple[str, torch.Tensor]], is_mtp: bool = False
|
||||
):
|
||||
super().load_weights(weights, is_mtp=True)
|
||||
has_mtp_layers = False
|
||||
has_target_layers = False
|
||||
head_weights = set()
|
||||
|
||||
def normalized_weights():
|
||||
nonlocal has_mtp_layers, has_target_layers
|
||||
for name, weight in weights:
|
||||
name = name.removeprefix("language_model.")
|
||||
has_mtp_layers |= name.startswith("mtp.layers.")
|
||||
has_target_layers |= name.startswith(
|
||||
("backbone.layers.", "model.layers.")
|
||||
)
|
||||
if name.startswith("lm_head."):
|
||||
head_weights.add(name)
|
||||
yield name, weight
|
||||
|
||||
# Inspect names while streaming: buffering a full target checkpoint here
|
||||
# would double its host-memory footprint during embedded MTP loading.
|
||||
super().load_weights(normalized_weights(), is_mtp=True)
|
||||
self._owns_lm_head = bool(
|
||||
has_mtp_layers and not has_target_layers and head_weights
|
||||
)
|
||||
if self._owns_lm_head:
|
||||
expected = {
|
||||
name
|
||||
for name, _ in self.named_parameters()
|
||||
if name.startswith("lm_head.")
|
||||
}
|
||||
if "lm_head.input_scale" in expected and isinstance(
|
||||
self.lm_head.quant_method, ModelOptNvFp4A16LinearMethod
|
||||
):
|
||||
# NVFP4A16 accepts this loader placeholder but never uses it.
|
||||
expected.remove("lm_head.input_scale")
|
||||
missing = (expected | {"lm_head.weight"}) - head_weights
|
||||
if missing:
|
||||
raise ValueError(
|
||||
f"Incomplete standalone MTP lm_head: missing {sorted(missing)}"
|
||||
)
|
||||
|
||||
def set_embed_and_head(self, embed, head):
|
||||
if not self._owns_lm_head:
|
||||
return super().set_embed_and_head(embed, head)
|
||||
# Standalone MTP checkpoints can supply a differently quantized head.
|
||||
# Share only the input embeddings; retain the entire loaded head module.
|
||||
self.model.embed_tokens.weight = embed
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
def set_lm_head_from_target(self, target_lm_head: nn.Module) -> None:
|
||||
if self.config.tie_word_embeddings:
|
||||
if self.config.tie_word_embeddings or self._owns_lm_head:
|
||||
return
|
||||
self.lm_head = target_lm_head
|
||||
|
||||
|
||||
@@ -42,6 +42,41 @@ input_dim_t: TypeAlias = int | tuple[int, int]
|
||||
norm_t: TypeAlias = tuple[float, float, float] | torch.Tensor
|
||||
|
||||
|
||||
def _map_hf_radio_weight_name(name: str) -> tuple[str, str | None] | None:
|
||||
prefix = "radio_model.hf_model."
|
||||
if not name.startswith(prefix):
|
||||
return None
|
||||
|
||||
name = name.removeprefix(prefix)
|
||||
if name == "summary_idxs":
|
||||
return None
|
||||
|
||||
embedding_names = {
|
||||
"embeddings.cls_register_token": "model.patch_generator.cls_token.token",
|
||||
"embeddings.patch_projection": "model.patch_generator.embedder",
|
||||
"embeddings.position_embedding": "model.patch_generator.pos_embed",
|
||||
"embeddings.video_patch_projection": "model.patch_generator.video_embedder",
|
||||
}
|
||||
for source, target in embedding_names.items():
|
||||
if name == source or name.startswith(f"{source}."):
|
||||
return name.replace(source, target, 1), None
|
||||
|
||||
name = name.replace("encoder.layer.", "model.encoder.layers.", 1)
|
||||
attention_names = {
|
||||
".attention.attention.query.": (".attn.attn.qkv_proj.", "q"),
|
||||
".attention.attention.key.": (".attn.attn.qkv_proj.", "k"),
|
||||
".attention.attention.value.": (".attn.attn.qkv_proj.", "v"),
|
||||
".attention.output.dense.": (".attn.attn.proj.", None),
|
||||
}
|
||||
for source, (target, shard_id) in attention_names.items():
|
||||
if source in name:
|
||||
return name.replace(source, target, 1), shard_id
|
||||
|
||||
name = name.replace(".layer_scale1.lambda1", ".ls1")
|
||||
name = name.replace(".layer_scale2.lambda1", ".ls2")
|
||||
return name, None
|
||||
|
||||
|
||||
def _ntuple(n):
|
||||
def parse(x):
|
||||
if isinstance(x, Iterable) and not isinstance(x, str):
|
||||
@@ -588,18 +623,34 @@ class RadioModel(nn.Module):
|
||||
weights_list = list(weights)
|
||||
|
||||
for name, weight in weights_list:
|
||||
if not name.startswith("radio_model."):
|
||||
# Skip non-radio weights
|
||||
continue
|
||||
name = replace_substrings(name, remap_substrings)
|
||||
name = replace_prefix(name, remap_prefixes)
|
||||
source_name = name
|
||||
is_hf_export = name.startswith("radio_model.hf_model.")
|
||||
loaded_shard_id = None
|
||||
if is_hf_export:
|
||||
mapped_weight = _map_hf_radio_weight_name(name)
|
||||
if mapped_weight is None:
|
||||
continue
|
||||
name, loaded_shard_id = mapped_weight
|
||||
else:
|
||||
if not name.startswith("radio_model."):
|
||||
# Skip non-radio weights
|
||||
continue
|
||||
name = replace_substrings(name, remap_substrings)
|
||||
name = replace_prefix(name, remap_prefixes)
|
||||
if name and name in params_dict:
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, weight)
|
||||
if loaded_shard_id is None:
|
||||
weight_loader(param, weight)
|
||||
else:
|
||||
weight_loader(param, weight, loaded_shard_id)
|
||||
loaded_params.add(name)
|
||||
if "video_embedder" in name:
|
||||
self.model.patch_generator._video_embedder_loaded = True
|
||||
elif is_hf_export:
|
||||
raise ValueError(
|
||||
f"Unexpected HF RADIO weight: {source_name} (mapped to {name})"
|
||||
)
|
||||
|
||||
return loaded_params
|
||||
|
||||
|
||||
@@ -22,6 +22,7 @@ from PIL import Image
|
||||
from sglang.srt.configs.nano_nemotron_vl import (
|
||||
NemotronH_Nano_Omni_Reasoning_V3_Config,
|
||||
NemotronH_Nano_VL_V2_Config,
|
||||
NemotronH_Omni_Reasoning_V3_Config,
|
||||
)
|
||||
from sglang.srt.managers.schedule_batch import (
|
||||
Modality,
|
||||
@@ -31,6 +32,7 @@ from sglang.srt.managers.schedule_batch import (
|
||||
from sglang.srt.models.nano_nemotron_vl import (
|
||||
NemotronH_Nano_Omni_Reasoning_V3,
|
||||
NemotronH_Nano_VL_V2,
|
||||
NemotronH_Omni_Reasoning_V3,
|
||||
)
|
||||
from sglang.srt.models.parakeet import ParakeetExtractor
|
||||
from sglang.srt.multimodal.audio_from_video import extract_audio_from_video_bytes
|
||||
@@ -58,7 +60,11 @@ MAX_FRAMES = 128
|
||||
|
||||
|
||||
class NanoNemotronVLImageProcessor(BaseMultimodalProcessor):
|
||||
models = [NemotronH_Nano_VL_V2, NemotronH_Nano_Omni_Reasoning_V3]
|
||||
models = [
|
||||
NemotronH_Nano_VL_V2,
|
||||
NemotronH_Nano_Omni_Reasoning_V3,
|
||||
NemotronH_Omni_Reasoning_V3,
|
||||
]
|
||||
gpu_image_decode = (
|
||||
False # NanoNemotronVL processes loaded image as PIL image explicitly
|
||||
)
|
||||
@@ -70,6 +76,7 @@ class NanoNemotronVLImageProcessor(BaseMultimodalProcessor):
|
||||
{
|
||||
NemotronH_Nano_VL_V2_Config: NemotronH_Nano_VL_V2,
|
||||
NemotronH_Nano_Omni_Reasoning_V3_Config: NemotronH_Nano_Omni_Reasoning_V3,
|
||||
NemotronH_Omni_Reasoning_V3_Config: NemotronH_Omni_Reasoning_V3,
|
||||
},
|
||||
)
|
||||
Image.MAX_IMAGE_PIXELS = None
|
||||
|
||||
@@ -68,6 +68,7 @@ from sglang.srt.configs import (
|
||||
NanbeigeConfig,
|
||||
NemotronH_Nano_Omni_Reasoning_V3_Config,
|
||||
NemotronH_Nano_VL_V2_Config,
|
||||
NemotronH_Omni_Reasoning_V3_Config,
|
||||
NemotronHConfig,
|
||||
NemotronHPuzzleConfig,
|
||||
Olmo3Config,
|
||||
@@ -138,6 +139,7 @@ _CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
|
||||
Dots3Config,
|
||||
NemotronH_Nano_VL_V2_Config,
|
||||
NemotronH_Nano_Omni_Reasoning_V3_Config,
|
||||
NemotronH_Omni_Reasoning_V3_Config,
|
||||
NemotronHConfig,
|
||||
NemotronHPuzzleConfig,
|
||||
NanbeigeConfig,
|
||||
|
||||
Reference in New Issue
Block a user