Add VLM support to the dynamic-batching inference server

examples/multimodal/config.py

@@ -154,6 +154,9 @@ def get_language_model_config(config, enable_fusions=False, apply_rope_fusion=No
         )
         config.attention_softmax_in_fp32 = True
         config.ffn_hidden_size = 8192
+    elif config.language_model_type == "nemotron6-moe":
+        config.bias_activation_fusion = False
+        config.bias_dropout_fusion = False
     elif config.language_model_type.startswith("hf://"):
         # Loaded from HuggingFace config file.
         import transformers
@@ -392,6 +395,10 @@ def get_vision_projection_config(config, hidden_size, enable_fusions=False):
         config.ffn_hidden_size = 2048
         config.activation_func = torch.nn.functional.gelu
         config.normalization = "LayerNorm"
+    elif config.language_model_type == "nemotron6-moe":
+        config.ffn_hidden_size = 20480
+        config.bias_activation_fusion = False
+        config.bias_dropout_fusion = False
     elif config.language_model_type.startswith("hf://"):
         config.activation_func = torch.nn.functional.gelu
         config.ffn_hidden_size = 4096

examples/multimodal/layer_specs.py

@@ -3,7 +3,6 @@
 
 import torch
 
-from megatron.core.extensions.transformer_engine import HAVE_TE
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
 from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules
 from megatron.core.ssm.mamba_layer import MambaLayer, MambaLayerSubmodules
@@ -15,9 +14,11 @@
 from megatron.core.transformer.enums import AttnMaskType
 from megatron.core.transformer.identity_op import IdentityOp
 from megatron.core.transformer.mlp import MLP, MLPSubmodules
+from megatron.core.models.gpt.moe_module_specs import get_moe_module_spec
 from megatron.core.transformer.spec_utils import ModuleSpec
-from megatron.core.transformer.transformer_layer import TransformerLayer, TransformerLayerSubmodules
+from megatron.core.transformer.transformer_layer import MoETransformerLayer, TransformerLayer, TransformerLayerSubmodules
 from megatron.core.typed_torch import not_none
+from megatron.core.extensions.transformer_engine import HAVE_TE
 
 if HAVE_TE:
     from megatron.core.extensions.transformer_engine import (
@@ -127,12 +128,34 @@ def get_layer_spec_te(is_vit=False, padding=False) -> ModuleSpec:
     )
 
 
-def get_hybrid_layer_spec_te(padding=False) -> ModuleSpec:
+def get_hybrid_layer_spec_te(config=None, padding=False) -> ModuleSpec:
+    """Hybrid (Mamba + attention + MLP [+ MoE]) layer spec.
+
+    Args:
+        config: language-model ``TransformerConfig``. Required for MoE hybrids
+            (e.g. nemotron6-moe): the moe_layer branch reads
+            ``num_moe_experts`` / ``moe_grouped_gemm`` off it to match the
+            checkpoint's architecture. Non-MoE hybrids may pass ``None``;
+            they never traverse the moe_layer branch.
+        padding: use padding-causal attention mask (needed for context
+            parallel + sequence parallel).
+    """
     attn_mask_type = AttnMaskType.causal
     # Padding mask is needed for e.g. Context Parallel.
     if padding:
         attn_mask_type = AttnMaskType.padding_causal
 
+    if config is not None:
+        assert config.num_moe_experts is not None, (
+            "get_hybrid_layer_spec_te: config.num_moe_experts must be set to "
+            "build the MoE branch of the hybrid stack."
+        )
+        num_experts = config.num_moe_experts
+        moe_grouped_gemm = config.moe_grouped_gemm
+    else:
+        num_experts = None
+        moe_grouped_gemm = None
+
     return ModuleSpec(
         module=HybridStack,
         submodules=HybridStackSubmodules(
@@ -182,6 +205,19 @@ def get_hybrid_layer_spec_te(padding=False) -> ModuleSpec:
                     mlp_bda=get_bias_dropout_add,
                 ),
             ),
+            moe_layer=ModuleSpec(
+                module=MoETransformerLayer,
+                submodules=TransformerLayerSubmodules(
+                    pre_mlp_layernorm=TENorm,
+                    mlp=get_moe_module_spec(
+                        use_te=True,
+                        num_experts=num_experts,
+                        moe_grouped_gemm=moe_grouped_gemm,
+                        moe_use_legacy_grouped_gemm=False,
+                    ),
+                    mlp_bda=get_bias_dropout_add,
+                ),
+            ),
         ),
     )
 

examples/multimodal/model.py

@@ -8,16 +8,19 @@
 from layer_specs import (get_layer_spec, get_layer_spec_te, get_mlp_module_spec, get_norm_mlp_module_spec_te,
                          get_hybrid_layer_spec_te)
 
+from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
 from megatron.core.models.multimodal.llava_model import IMAGE_TOKEN, LLaVAModel
 from megatron.core.models.vision.clip_vit_model import get_num_image_embeddings
+from megatron.core.transformer.spec_utils import import_module
 from megatron.training import get_args, get_tokenizer, print_rank_0
 from megatron.training.arguments import core_transformer_config_from_args
 from megatron.core.utils import log_single_rank
 
 
+
 def model_provider(
     pre_process=True, post_process=True, add_encoder=True, add_decoder=True, parallel_output=True,
-    vp_stage=None, config=None, pg_collection=None
+    vp_stage=None, config=None, pg_collection=None,
 ) -> LLaVAModel:
     """Builds the model.
 
@@ -28,10 +31,10 @@ def model_provider(
             will live on only a subset of the pipeline stages (specifically, only the first stage).
         add_decoder (bool): Construct the decoder module (used with pipeline parallelism). Defaults to True. When we use pipelining, the decoder
             will live on only a subset of the pipeline stages (specifically, every stage after the first one).
-        parallel_output (bool): Enable parallel model output.
         vp_stage: Optional virtual pipeline stage. Used with virtual pipeline parallelism.
         config: Optional transformer config. If None, will be created from args.
         pg_collection: Optional process group collection. If None, will use default.
+        parallel_output (bool): Enable parallel model output.
 
     Returns:
         model: A multimodal model.
@@ -41,28 +44,38 @@ def model_provider(
 
     print_rank_0('building a multimodal model ...')
 
-    num_image_embeddings = get_num_image_embeddings(
-        args.img_h,
-        args.img_w,
-        args.patch_dim,
-        args.vision_model_type,
-        args.disable_vision_class_token,
-        1,
-        args.pixel_shuffle,
-        args.use_tile_tags,
-        args.max_num_tiles,
-        args.tokenizer_prompt_format
-    )
-    old_seq_length = args.seq_length
-    args.seq_length = args.encoder_seq_length = num_image_embeddings
-    if old_seq_length != args.seq_length:
-        log_single_rank(
-            logging.getLogger(__name__),
-            logging.WARNING,
-            f"Changed seq_length and encoder_seq_length (vision model sequence length) from {old_seq_length} to num_image_tokens ({num_image_embeddings})"
+    if getattr(args, 'dynamic_resolution', False):
+        max_num_image_embeddings = args.seq_length
+        num_image_embeddings = args.seq_length
+        if args.pixel_shuffle:
+            max_num_image_embeddings //= 4
+            num_image_embeddings //= 4
+        if getattr(args, 'conv_merging', False):
+            max_num_image_embeddings //= 4
+            num_image_embeddings //= 4
+    else:
+        num_image_embeddings = get_num_image_embeddings(
+            args.img_h,
+            args.img_w,
+            args.patch_dim,
+            args.vision_model_type,
+            args.disable_vision_class_token,
+            1,
+            args.pixel_shuffle,
+            args.use_tile_tags,
+            args.max_num_tiles,
+            args.tokenizer_prompt_format
         )
+        old_seq_length = args.seq_length
+        args.seq_length = args.encoder_seq_length = num_image_embeddings
+        if old_seq_length != args.seq_length:
+            log_single_rank(
+                logging.getLogger(__name__),
+                logging.WARNING,
+                f"Changed seq_length and encoder_seq_length (vision model sequence length) from {old_seq_length} to num_image_tokens ({num_image_embeddings})"
+            )
 
-    max_num_image_embeddings = max((args.max_num_tiles + int(args.use_thumbnail)), args.num_frames) * num_image_embeddings
+        max_num_image_embeddings = max((args.max_num_tiles + int(args.use_thumbnail)), args.num_frames) * num_image_embeddings
 
     assert (
         args.decoder_seq_length is not None
@@ -79,7 +92,7 @@ def model_provider(
     language_model_type = args.language_model_type
     vision_model_type = args.vision_model_type
 
-    base_config = core_transformer_config_from_args(get_args())
+    base_config = config or core_transformer_config_from_args(get_args())
     base_config.language_model_type = args.language_model_type
     base_config.vision_model_type = args.vision_model_type
     base_config.calculate_per_token_loss = True
@@ -98,8 +111,19 @@ def model_provider(
     elif use_te:
         # Padding mask needed for SP/CP.
         padding = args.context_parallel_size > 1 and args.sequence_parallel
-        if args.language_model_type.startswith('nemotron5-hybrid'):
-            language_transformer_layer_spec = get_hybrid_layer_spec_te(padding=padding)
+        if args.spec is not None:
+            language_transformer_layer_spec = import_module(args.spec)
+        elif args.language_model_type.startswith(('nemotron5-hybrid', 'nemotron6-moe')):
+            language_transformer_layer_spec = get_hybrid_layer_spec_te(
+                config=language_config, padding=padding
+            )
+        elif getattr(args, 'num_experts', None):
+            language_transformer_layer_spec = get_gpt_decoder_block_spec(
+                language_config,
+                use_transformer_engine=use_te,
+                normalization=args.normalization,
+                qk_l2_norm=getattr(args, 'qk_l2_norm', False),
+            )
         else:
             language_transformer_layer_spec = get_layer_spec_te(
                 is_vit=False, padding=padding
@@ -110,11 +134,12 @@ def model_provider(
         )
 
     vision_config = deepcopy(base_config)
-    vision_config = get_vision_model_config(
-        vision_config, apply_query_key_layer_scaling=args.apply_query_key_layer_scaling
-    )
+    vision_config = get_vision_model_config(vision_config)
+    # Most ViT checkpoints use bias in linear layers; override --disable-bias-linear.
+    # Pixtral (both sizes) uses no bias — config.py already sets add_bias_linear=False.
+    if vision_model_type not in ("pixtral-vit", "pixtral-vit-large"):
+        vision_config.add_bias_linear = True
     if vision_model_type.startswith("hf://"):
-        assert not args.sequence_parallel, "Huggingface models do not support --sequence-parallel"
         assert args.context_parallel_size < 2, "Huggingface models do not support --context-parallel-size > 1"
 
     if vision_model_type in ["clip", "siglip", "radio", "cradio-g"]:
@@ -141,6 +166,13 @@ def model_provider(
     elif vision_model_type == "internvit300M":
         from nvlm.internvit import get_internvit300M_layer_spec
         vision_transformer_layer_spec = get_internvit300M_layer_spec(use_te=use_te)
+    elif vision_model_type in ("pixtral-vit", "pixtral-vit-large", "qwen-vl", "kimi-vit"):
+        if use_te:
+            vision_transformer_layer_spec = get_layer_spec_te(is_vit=True)
+        else:
+            vision_transformer_layer_spec = get_layer_spec(
+                is_vit=True, normalization=vision_config.normalization
+            )
     elif vision_model_type.startswith("hf://"):
         vision_transformer_layer_spec = None
     else:
@@ -159,10 +191,13 @@ def model_provider(
     # Make sure the vision model does not inherit first and last pipeline num layers from the language model.
     vision_config.first_pipeline_num_layers = vision_config.last_pipeline_num_layers = None
 
+    # ``get_*_module_spec_te`` returns ``functools.partial(MLP.as_mlp_submodule,
+    # submodules=...)`` (see PR #3435). Pull the submodules out of the partial's
+    # bound kwargs so the vision projection sees an ``MLPSubmodules`` value.
     if vision_projection_config.normalization:
-        vision_projection_layer_spec = get_norm_mlp_module_spec_te().submodules
+        vision_projection_layer_spec = get_norm_mlp_module_spec_te().keywords["submodules"]
     else:
-        vision_projection_layer_spec = get_mlp_module_spec(use_te=use_te).submodules
+        vision_projection_layer_spec = get_mlp_module_spec(use_te=use_te).keywords["submodules"]
 
     # Toggle --recompute* for the vision and language model separately.
     if args.recompute_vision:
@@ -202,7 +237,7 @@ def model_provider(
         drop_vision_class_token=args.disable_vision_class_token,
         vision_projection_config=vision_projection_config,
         vision_projection_layer_spec=vision_projection_layer_spec,
-        vision_projection_type="mlp",
+        vision_projection_type=args.vision_projection_type,
         allow_missing_vision_projection_checkpoint=args.allow_missing_vision_projection_checkpoint,
         parallel_output=parallel_output,
         share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
@@ -221,9 +256,20 @@ def model_provider(
         fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
         image_token_index=image_token_index,
         pixel_shuffle=args.pixel_shuffle,
+        conv_merging=getattr(args, "conv_merging", False),
         tile_tags=tile_tags,
         max_num_tiles=args.max_num_tiles,
         tokenizer_type=args.tokenizer_prompt_format,
+        use_vision_backbone_fp8_arch=getattr(args, "use_vision_backbone_fp8_arch", False),
+        dynamic_resolution=getattr(args, "dynamic_resolution", False),
+        class_token_len=getattr(args, "class_token_len", None),
+        radio_force_eval_mode=getattr(args, "radio_force_eval_mode", False),
+        radio_force_cpe_eval_mode=getattr(args, "radio_force_cpe_eval_mode", False),
+        radio_interpolate_only_cpe=getattr(args, "radio_interpolate_only_cpe", False),
+        radio_cpe_aspect_ratio_select=getattr(args, "radio_cpe_aspect_ratio_select", False),
+        radio_disable_cpe=getattr(args, "radio_disable_cpe", False),
+        vp_stage=vp_stage,
+        pg_collection=pg_collection,
     )
 
     model.freeze(