diff --git a/examples/conversion/hf_megatron_roundtrip_multi_gpu.py b/examples/conversion/hf_megatron_roundtrip_multi_gpu.py index 80282daf53..5d2c495df1 100644 --- a/examples/conversion/hf_megatron_roundtrip_multi_gpu.py +++ b/examples/conversion/hf_megatron_roundtrip_multi_gpu.py @@ -70,8 +70,10 @@ # MiniMax-M2: QK norms stored as bf16 in HF, loaded as fp32 by Megatron config.params_dtype "q_norm.weight", "k_norm.weight", - # MiniMax-M2: router gate stored as fp32 in HF, loaded as bf16 via autocast_dtype + # MoE router gate stored as fp32 in Megatron, may be bf16 in HF "block_sparse_moe.gate.weight", + "mlp.gate.weight", + "moe.gate.weight", ] # FP8 dtypes whose dequantisation is inherently lossy — allclose is meaningless. diff --git a/examples/models/vlm/ernie_vl/README.md b/examples/models/vlm/ernie_vl/README.md new file mode 100644 index 0000000000..133322004f --- /dev/null +++ b/examples/models/vlm/ernie_vl/README.md @@ -0,0 +1,71 @@ +# ERNIE 4.5 VL Examples + +This directory contains example scripts for ERNIE 4.5 Vision-Language (VL) MoE models. + +## Supported Models + +| Model | Parameters | Active Parameters | Type | +|-------|-----------|-------------------|------| +| ERNIE-4.5-VL-28B-A3B-Instruct | 28B | 3B | VL MoE | +| ERNIE-4.5-VL-28B-A3B-Thinking | 28B | 3B | VL MoE | + +## Prerequisites + +- `--trust-remote-code` is required for the custom HuggingFace model class. +- All scripts use a `WORKSPACE` environment variable for checkpoints. Default: `/workspace`. + +```bash +export WORKSPACE=/your/custom/path +``` + +## Checkpoint Conversion + +### Import HF → Megatron + +```bash +uv run python examples/conversion/convert_checkpoints.py import \ + --hf-model baidu/ERNIE-4.5-VL-28B-A3B-Instruct \ + --megatron-path ${WORKSPACE}/ERNIE-4.5-VL-28B-A3B-Instruct \ + --torch-dtype bfloat16 \ + --trust-remote-code +``` + +### Export Megatron → HF + +```bash +uv run python examples/conversion/convert_checkpoints.py export \ + --hf-model baidu/ERNIE-4.5-VL-28B-A3B-Instruct \ + --megatron-path ${WORKSPACE}/ERNIE-4.5-VL-28B-A3B-Instruct/iter_0000000 \ + --hf-path ${WORKSPACE}/ERNIE-4.5-VL-28B-A3B-Instruct-hf-export \ + --trust-remote-code +``` + +See [conversion.sh](conversion.sh) for the full pipeline including multi-GPU round-trip validation. + +## Inference + +ERNIE 4.5 VL uses a processor API that differs from other VLMs (e.g., Qwen), so a +dedicated inference script is provided instead of the generic `hf_to_megatron_generate_vlm.py`. + +### Run Inference from HF Checkpoint + +```bash +uv run python -m torch.distributed.run --nproc_per_node=8 \ + examples/models/vlm/ernie_vl/hf_to_megatron_generate_ernie_vl.py \ + --hf_model_path baidu/ERNIE-4.5-VL-28B-A3B-Instruct \ + --image_path "https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16/resolve/main/images/table.png" \ + --prompt "Describe this image." \ + --max_new_tokens 100 \ + --tp 2 --pp 1 --ep 4 \ + --trust_remote_code +``` + +See [inference.sh](inference.sh) for a ready-to-use launch script. + +## Architecture Notes + +ERNIE 4.5 VL uses a **dual-pool MoE** architecture: +- Text and vision experts reside in separate pools within each MoE layer. +- Each pool has its own router and routes tokens independently. +- This design uses `SequentialMLP` (per-expert execution) rather than `GroupedMLP` + (batched GEMM), since the two pools cannot be merged into a single expert group. diff --git a/examples/models/vlm/ernie_vl/conversion.sh b/examples/models/vlm/ernie_vl/conversion.sh new file mode 100755 index 0000000000..3ea61d2517 --- /dev/null +++ b/examples/models/vlm/ernie_vl/conversion.sh @@ -0,0 +1,53 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +set -e + +# Workspace directory for checkpoints and results +WORKSPACE=${WORKSPACE:-/workspace} +# Supported model variants: +# ERNIE-4.5-VL-28B-A3B-Instruct, ERNIE-4.5-VL-28B-A3B-Thinking +MODEL_NAME=ERNIE-4.5-VL-28B-A3B-Instruct + +EP=4 +TP=2 +PP=1 + +# Import HF -> Megatron +uv run python examples/conversion/convert_checkpoints.py import \ + --hf-model baidu/${MODEL_NAME} \ + --megatron-path ${WORKSPACE}/${MODEL_NAME} \ + --torch-dtype bfloat16 \ + --trust-remote-code + +# HF and Megatron models logits comparison validation +uv run python -m torch.distributed.run --nproc_per_node=8 examples/conversion/compare_hf_and_megatron/compare.py \ + --hf_model_path baidu/${MODEL_NAME} \ + --megatron_model_path ${WORKSPACE}/${MODEL_NAME} \ + --model_class "Ernie4_5_VLMoeForConditionalGeneration" \ + --image_path "https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16/resolve/main/images/table.png" \ + --prompt "Describe this image." \ + --tp ${TP} --pp ${PP} --ep ${EP} \ + --trust-remote-code + +# Export Megatron -> HF +uv run python examples/conversion/convert_checkpoints.py export \ + --hf-model baidu/${MODEL_NAME} \ + --megatron-path ${WORKSPACE}/${MODEL_NAME}/iter_0000000 \ + --hf-path ${WORKSPACE}/${MODEL_NAME}-hf-export \ + --trust-remote-code + +# Round-trip validation +uv run python -m torch.distributed.run --nproc_per_node=8 examples/conversion/hf_megatron_roundtrip_multi_gpu.py \ + --hf-model-id baidu/${MODEL_NAME} --tp ${TP} --pp ${PP} --ep ${EP} --trust-remote-code diff --git a/examples/models/vlm/ernie_vl/ernie45_vl_fwd_bwd.py b/examples/models/vlm/ernie_vl/ernie45_vl_fwd_bwd.py new file mode 100644 index 0000000000..39d889110b --- /dev/null +++ b/examples/models/vlm/ernie_vl/ernie45_vl_fwd_bwd.py @@ -0,0 +1,520 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Standalone forward/backward test for ERNIE 4.5 VL MoE Megatron model. + +This script must be launched with torchrun (for multi-GPU tests): + torchrun --nproc_per_node=N ernie45_vl_fwd_bwd.py --hf-model-path --tp T --pp P --ep E + +It performs: +1. Load HF toy model -> convert to Megatron via AutoBridge +2. Forward pass with text-only or text+vision input +3. Backward pass: compute loss, check gradients exist +4. Print PASS/FAIL status + +With --with-vision, constructs a dummy image input that exercises the full +vision pipeline: ViT patch embedding -> vision transformer -> resampler -> +embedding injection -> language model forward. + +Exit code 0 = PASS, non-zero = FAIL. +""" + +import argparse +import os + + +# Disable torch.compile to avoid triton compatibility issues in some environments. +# Must be set before importing torch. +os.environ.setdefault("TORCH_COMPILE_DISABLE", "1") + +import torch +import torch.distributed as dist + +from megatron.bridge import AutoBridge +from megatron.bridge.models.decorators import torchrun_main + + +def _is_rank_0() -> bool: + if dist.is_initialized(): + return dist.get_rank() == 0 + return True + + +def print_rank_0(msg: str): + """Print a message only from rank 0.""" + if _is_rank_0(): + print(msg, flush=True) + + +def run_forward_backward( + hf_model_path: str, + tp: int = 1, + pp: int = 1, + ep: int = 1, + seq_len: int = 16, + backward: bool = True, + with_vision: bool = False, + prompt: str | None = None, +): + """Run forward (and optionally backward) pass on the ERNIE 4.5 VL MoE toy model. + + Args: + hf_model_path: Path to the HF toy model directory. + tp: Tensor parallelism size. + pp: Pipeline parallelism size. + ep: Expert parallelism size. + seq_len: Sequence length for the dummy input. + backward: Whether to also run backward pass. + with_vision: Whether to include a dummy image in the input to exercise + the vision tower and resampler forward path. + """ + print_rank_0("=== ERNIE 4.5 VL Forward/Backward Test ===") + print_rank_0(f" HF model: {hf_model_path}") + print_rank_0(f" TP={tp}, PP={pp}, EP={ep}, seq_len={seq_len}, backward={backward}, with_vision={with_vision}") + + # ------------------------------------------------------------------ # + # 1. Build Megatron model from HF checkpoint via AutoBridge + # ------------------------------------------------------------------ # + print_rank_0("Step 1: Loading HF model and converting to Megatron...") + + bridge = AutoBridge.from_hf_pretrained( + hf_model_path, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + ) + model_provider = bridge.to_megatron_provider(load_weights=True) + model_provider.tensor_model_parallel_size = tp + model_provider.pipeline_model_parallel_size = pp + model_provider.expert_model_parallel_size = ep + # Megatron requires sequence parallelism when MoE + TP > 1 during training + if tp > 1 and backward: + model_provider.sequence_parallel = True + model_provider.pipeline_dtype = torch.bfloat16 + model_provider.params_dtype = torch.bfloat16 + model_provider.finalize() + model_provider.initialize_model_parallel(seed=42) + + megatron_models = model_provider.provide_distributed_model(wrap_with_ddp=False) + + # Disable deallocate_pipeline_outputs: the no-pipelining schedule does not + # call deallocate_output_tensor(), so backward_step's custom_backward() + # would fail with "output should be pseudo-'freed' in schedule". + for m in megatron_models: + if hasattr(m, "config"): + m.config.deallocate_pipeline_outputs = False + + # Put models in train mode for backward pass + if backward: + for m in megatron_models: + m.train() + else: + for m in megatron_models: + m.eval() + + print_rank_0(f" Model built successfully. {len(megatron_models)} component(s).") + + # ------------------------------------------------------------------ # + # 2. Prepare input (text-only or text+vision) + # ------------------------------------------------------------------ # + print_rank_0(f"Step 2: Preparing {'text+vision' if with_vision else 'text-only'} input...") + + # Use vocab_size from model config to stay in valid range + vocab_size = getattr(model_provider, "padded_vocab_size", None) + if vocab_size is None: + vocab_size = getattr(model_provider, "vocab_size", None) + if vocab_size is None: + vocab_size = 2048 # toy model default + + pixel_values = None + image_grid_thw = None + + if with_vision: + # Read vision config and special token IDs from the model provider / HF config + import json + + with open(os.path.join(hf_model_path, "config.json")) as f: + hf_cfg = json.load(f) + + image_token_id = hf_cfg.get("image_token_id", 100295) + image_start_token_id = hf_cfg.get("image_start_token_id", 101304) + image_end_token_id = hf_cfg.get("image_end_token_id", 101305) + + # For toy models with small vocab_size, remap special token IDs to fit + # within the embedding table. This allows the vision pipeline to execute + # without index-out-of-range errors in the embedding layer. + text_cfg = hf_cfg.get("text_config", hf_cfg) + cfg_vocab_size = text_cfg.get("vocab_size", vocab_size) + if ( + image_token_id >= cfg_vocab_size + or image_start_token_id >= cfg_vocab_size + or image_end_token_id >= cfg_vocab_size + ): + # Use the last 3 tokens in the vocab as placeholders + image_token_id = cfg_vocab_size - 3 + image_start_token_id = cfg_vocab_size - 2 + image_end_token_id = cfg_vocab_size - 1 + print_rank_0( + f" Remapped special token IDs to fit vocab_size={cfg_vocab_size}: " + f"image_token={image_token_id}, start={image_start_token_id}, end={image_end_token_id}" + ) + + # Also update model config so get_placeholder_mask / get_rope_index use the remapped IDs + for m in megatron_models: + if hasattr(m, "config"): + m.config.image_token_id = image_token_id + m.config.image_start_token_id = image_start_token_id + m.config.image_end_token_id = image_end_token_id + + # Vision config for computing pixel_values shape + vis_cfg = hf_cfg.get("vision_config", {}) + patch_size = vis_cfg.get("patch_size", 14) + in_channels = vis_cfg.get("in_channels", 3) + spatial_merge_size = vis_cfg.get("spatial_merge_size", 2) + + # Use minimal valid image grid: 2x2 patches (must be divisible by spatial_merge_size) + grid_h, grid_w = 2, 2 + image_grid_thw = torch.tensor([[1, grid_h, grid_w]], dtype=torch.long, device="cuda") + + # pixel_values: [T*H*W, in_channels * patch_size * patch_size] + num_patches = grid_h * grid_w # 4 + patch_dim = in_channels * patch_size * patch_size # 3*14*14 = 588 + pixel_values = torch.randn(num_patches, patch_dim, dtype=torch.bfloat16, device="cuda") + + # Number of image placeholder tokens after resampler spatial merge + num_image_tokens = num_patches // (spatial_merge_size**2) # 4 // 4 = 1 + + # Build input_ids: [text..., image_start, , image_end, text...] + # Ensure seq_len is large enough + min_seq_len = num_image_tokens + 4 # at least: text + start + placeholders + end + text + actual_seq_len = max(seq_len, min_seq_len) + + # Number of text tokens before and after the image block + num_text_before = 2 + num_text_after = actual_seq_len - num_text_before - 1 - num_image_tokens - 1 + if num_text_after < 1: + num_text_after = 1 + actual_seq_len = num_text_before + 1 + num_image_tokens + 1 + num_text_after + + # Construct input_ids + # Use token IDs that don't collide with special tokens + max_text_token = min(vocab_size, 1024, image_token_id) + text_before = torch.randint(1, max(2, max_text_token), (num_text_before,), device="cuda") + img_start = torch.tensor([image_start_token_id], device="cuda") + img_placeholders = torch.full((num_image_tokens,), image_token_id, device="cuda") + img_end = torch.tensor([image_end_token_id], device="cuda") + text_after = torch.randint(1, max(2, max_text_token), (num_text_after,), device="cuda") + input_ids = torch.cat([text_before, img_start, img_placeholders, img_end, text_after]).unsqueeze(0) + actual_seq_len = input_ids.shape[1] + + # mm_token_type_ids: 0=text, 1=image placeholder + mm_token_type_ids = torch.zeros(1, actual_seq_len, dtype=torch.int32, device="cuda") + img_start_pos = num_text_before + 1 # position of first image placeholder + mm_token_type_ids[0, img_start_pos : img_start_pos + num_image_tokens] = 1 + + # Labels and loss mask + labels = torch.randint(0, min(vocab_size, 1024), (1, actual_seq_len), device="cuda") + loss_mask = torch.ones(1, actual_seq_len, dtype=torch.float32, device="cuda") + + print_rank_0(f" input_ids shape: {input_ids.shape}, vocab_size: {vocab_size}") + print_rank_0(f" pixel_values shape: {pixel_values.shape}") + print_rank_0(f" image_grid_thw: {image_grid_thw.tolist()}") + print_rank_0(f" num_image_tokens (placeholders): {num_image_tokens}") + seq_len = actual_seq_len + else: + # Text-only input + if prompt is not None: + # Use real text with tokenizer for meaningful loss measurement. + # Labels are the next-token targets (input shifted right by 1). + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(hf_model_path, trust_remote_code=True) + token_ids = tokenizer.encode(prompt, add_special_tokens=True) + # Need at least 2 tokens for next-token prediction + assert len(token_ids) >= 2, f"Prompt too short: {len(token_ids)} tokens" + input_ids = torch.tensor([token_ids], dtype=torch.long, device="cuda") + seq_len = input_ids.shape[1] + # Labels: shifted right by 1 (predict next token) + # For position i, label[i] = input_ids[i+1]; last position has no valid label + labels = input_ids.clone() + labels[:, :-1] = input_ids[:, 1:] + labels[:, -1] = -100 # ignore last position (no next token) + loss_mask = torch.ones(1, seq_len, dtype=torch.float32, device="cuda") + loss_mask[:, -1] = 0 # don't compute loss on last position + print_rank_0(f" Prompt: {prompt!r}") + print_rank_0(f" Tokenized: {len(token_ids)} tokens") + else: + input_ids = torch.randint(0, min(vocab_size, 1024), (1, seq_len), device="cuda") + labels = torch.randint(0, min(vocab_size, 1024), (1, seq_len), device="cuda") + loss_mask = torch.ones(1, seq_len, dtype=torch.float32, device="cuda") + mm_token_type_ids = torch.zeros(1, seq_len, dtype=torch.int32, device="cuda") + print_rank_0(f" input_ids shape: {input_ids.shape}, vocab_size: {vocab_size}") + + # ------------------------------------------------------------------ # + # 3. Forward pass + # ------------------------------------------------------------------ # + print_rank_0("Step 3: Running forward pass...") + + from megatron.core.pipeline_parallel.schedules import get_forward_backward_func + + # Build a data iterator compatible with get_forward_backward_func + class SingleBatchIterator: + def __init__(self, batch): + self.batch = batch + self._yielded = False + + def __iter__(self): + return self + + def __next__(self): + if self._yielded: + raise StopIteration + self._yielded = True + return self.batch + + def ernie_vl_forward_step(data_iterator, model, **kwargs): + """Forward step for ERNIE 4.5 VL model. + + Follows the Megatron forward_step / loss_func protocol: + - forward_step returns (output_tensor, loss_func) + - loss_func(output_tensor) returns a 3-tuple (scalar_loss, num_tokens, report) + or 2-tuple (scalar_loss, report) for backward compatibility. + - The scheduler calls loss_func on the last pipeline stage, normalises + the scalar loss, then calls .backward() on it. + + GPTModel.forward(labels=...) returns per-token cross-entropy loss of + shape [batch, seq_len]. loss_func must reduce it to a differentiable + scalar via loss_mask. + """ + batch = next(data_iterator) + + forward_args = { + "input_ids": batch["tokens"], + "mm_token_type_ids": batch["mm_token_type_ids"], + "attention_mask": None, # Let Megatron auto-generate causal mask + } + + # Pass vision inputs if present + if batch.get("pixel_values") is not None: + forward_args["pixel_values"] = batch["pixel_values"] + if batch.get("image_grid_thw") is not None: + forward_args["image_grid_thw"] = batch["image_grid_thw"] + + if backward: + forward_args["labels"] = batch["labels"] + forward_args["loss_mask"] = batch["loss_mask"] + + output = model(**forward_args) + + if backward: + # output is per-token loss [batch, seq_len] from GPTModel + per_token_loss = output + cur_loss_mask = batch["loss_mask"] + + def loss_func(output_tensor, **kwargs): + # Reduce per-token loss to a mean scalar using loss_mask + losses = output_tensor.view(-1).float() + mask = cur_loss_mask.view(-1).float() + num_tokens = mask.sum() + loss = torch.sum(losses * mask) / torch.clamp(num_tokens, min=1) + return loss, {"lm loss": loss.clone().detach()} + + return per_token_loss, loss_func + else: + if isinstance(output, tuple): + output = output[0] + + def loss_func(output_tensor, **kwargs): + # forward-only: return a dummy scalar loss and the logits as report + dummy_loss = output_tensor.sum() * 0 # zero-grad scalar on same device + return dummy_loss, {"logits": output_tensor} + + return output, loss_func + + batch = { + "tokens": input_ids, + "mm_token_type_ids": mm_token_type_ids, + "labels": labels, + "loss_mask": loss_mask, + "pixel_values": pixel_values, + "image_grid_thw": image_grid_thw, + } + + fwd_bwd_func = get_forward_backward_func() + iterator = SingleBatchIterator(batch) + + output = fwd_bwd_func( + forward_step_func=ernie_vl_forward_step, + data_iterator=iterator, + model=megatron_models, + num_microbatches=1, + forward_only=not backward, + seq_length=seq_len, + micro_batch_size=1, + ) + + print_rank_0(" Forward pass completed successfully.") + + # ------------------------------------------------------------------ # + # 4. Verify forward output + # ------------------------------------------------------------------ # + print_rank_0("Step 4: Verifying forward output...") + + from megatron.core import parallel_state + + is_last_stage = not dist.is_initialized() or parallel_state.is_pipeline_last_stage() + + if is_last_stage: + if isinstance(output, list) and len(output) > 0: + result = output[0] + else: + result = output + + if backward: + # In backward mode, output contains loss info + if isinstance(result, dict): + print_rank_0(f" Loss output: {result}") + elif isinstance(result, torch.Tensor): + print_rank_0(f" Loss value: {result.item():.6f}") + assert torch.isfinite(result), f"Loss is not finite: {result.item()}" + else: + print_rank_0(f" Output type: {type(result)}, value: {result}") + else: + # In forward-only mode, output is stored as {"logits": tensor} from loss_func + if isinstance(result, dict) and "logits" in result: + logits = result["logits"] + print_rank_0(f" Logits shape: {logits.shape}") + print_rank_0(f" Logits stats: mean={logits.float().mean():.4f}, std={logits.float().std():.4f}") + assert torch.isfinite(logits).all(), "Logits contain non-finite values" + elif isinstance(result, torch.Tensor): + print_rank_0(f" Output shape: {result.shape}") + print_rank_0(f" Output stats: mean={result.float().mean():.4f}, std={result.float().std():.4f}") + assert torch.isfinite(result).all(), "Output contains non-finite values" + else: + print_rank_0(f" Output type: {type(result)}") + + print_rank_0(" Forward output verification passed.") + + # ------------------------------------------------------------------ # + # 5. Verify gradients (backward pass) + # ------------------------------------------------------------------ # + if backward: + print_rank_0("Step 5: Verifying gradients from backward pass...") + + # Check that at least some parameters have gradients + total_params = 0 + params_with_grad = 0 + params_with_nonzero_grad = 0 + + for m in megatron_models: + for _name, param in m.named_parameters(): + if param.requires_grad: + total_params += 1 + if param.grad is not None: + params_with_grad += 1 + if param.grad.abs().sum() > 0: + params_with_nonzero_grad += 1 + + print_rank_0(f" Total trainable params: {total_params}") + print_rank_0(f" Params with gradient: {params_with_grad}") + print_rank_0(f" Params with non-zero gradient: {params_with_nonzero_grad}") + + # At least some params should have gradients + # (Not all will have gradients due to PP - only the local stage's params) + assert params_with_grad > 0, ( + f"No parameters have gradients! total_params={total_params}, params_with_grad={params_with_grad}" + ) + + # When vision is enabled, verify vision tower and resampler got gradients + if with_vision: + vision_params_with_grad = 0 + for m in megatron_models: + for name, param in m.named_parameters(): + if param.requires_grad and param.grad is not None: + if "vision_tower" in name or "resampler" in name: + if param.grad.abs().sum() > 0: + vision_params_with_grad += 1 + print_rank_0(f" Vision params with non-zero gradient: {vision_params_with_grad}") + assert vision_params_with_grad > 0, ( + "Vision tower/resampler parameters have no gradients! The vision forward path may not be exercised." + ) + + print_rank_0(" Gradient verification passed.") + + # ------------------------------------------------------------------ # + # Done + # ------------------------------------------------------------------ # + print_rank_0("=== ALL CHECKS PASSED ===") + + +@torchrun_main +def _run( + hf_model_path: str, + tp: int = 1, + pp: int = 1, + ep: int = 1, + seq_len: int = 16, + forward_only: bool = False, + with_vision: bool = False, + prompt: str | None = None, +): + """Entry point for torchrun-launched forward/backward test.""" + run_forward_backward( + hf_model_path=hf_model_path, + tp=tp, + pp=pp, + ep=ep, + seq_len=seq_len, + backward=not forward_only, + with_vision=with_vision, + prompt=prompt, + ) + + +def main(): + """Parse CLI arguments and launch the forward/backward test.""" + parser = argparse.ArgumentParser(description="ERNIE 4.5 VL MoE forward/backward test") + parser.add_argument("--hf-model-path", required=True, help="Path to HF toy model directory") + parser.add_argument("--tp", type=int, default=1, help="Tensor parallelism size") + parser.add_argument("--pp", type=int, default=1, help="Pipeline parallelism size") + parser.add_argument("--ep", type=int, default=1, help="Expert parallelism size") + parser.add_argument("--seq-len", type=int, default=16, help="Sequence length") + parser.add_argument("--forward-only", action="store_true", help="Skip backward pass") + parser.add_argument( + "--with-vision", + action="store_true", + help="Include a dummy image input to exercise the vision tower and resampler", + ) + parser.add_argument( + "--prompt", + type=str, + default=None, + help="Use real text prompt instead of random tokens for meaningful loss measurement", + ) + args = parser.parse_args() + + _run( + hf_model_path=args.hf_model_path, + tp=args.tp, + pp=args.pp, + ep=args.ep, + seq_len=args.seq_len, + forward_only=args.forward_only, + with_vision=args.with_vision, + prompt=args.prompt, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/models/vlm/ernie_vl/ernie45_vl_logit_compare.py b/examples/models/vlm/ernie_vl/ernie45_vl_logit_compare.py new file mode 100644 index 0000000000..e7dba5ac76 --- /dev/null +++ b/examples/models/vlm/ernie_vl/ernie45_vl_logit_compare.py @@ -0,0 +1,665 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Logit comparison between HuggingFace and Megatron ERNIE 4.5 VL MoE models. + +This script compares 1-step forward-pass logits between the HuggingFace ERNIE 4.5 +VL MoE model and its Megatron-Core conversion via AutoBridge. It is designed for +validating weight conversion correctness with the real 28B model. + +Strategy (sequential, same GPU): + 1. Load HF model on rank 0 GPU -> forward pass -> save logits to CPU + 2. Delete HF model, free GPU memory + 3. Load Megatron model (via AutoBridge) across GPUs -> forward pass -> gather logits + 4. Compare logits: cosine similarity, top-k token match, absolute diff + +Launch: + # Text-only comparison (TP=2, EP=2 -> 4 GPUs): + torchrun --nproc_per_node=4 ernie45_vl_logit_compare.py \ + --hf-model-path /path/to/ERNIE-4.5-VL-28B-A3B-Thinking \ + --prompt "Hello, how are you?" \ + --tp 2 --ep 2 + + # Single-GPU comparison: + torchrun --nproc_per_node=1 ernie45_vl_logit_compare.py \ + --hf-model-path /path/to/ERNIE-4.5-VL-28B-A3B-Thinking \ + --prompt "Hello" + + # With image (VL inference): + torchrun --nproc_per_node=1 ernie45_vl_logit_compare.py \ + --hf-model-path /path/to/ERNIE-4.5-VL-28B-A3B-Thinking \ + --prompt "Describe this image." \ + --image-path /path/to/image.jpg + +Exit code 0 = PASS (cosine similarity >= threshold), non-zero = FAIL. +""" + +import argparse +import gc +import os +import sys + + +os.environ.setdefault("TORCH_COMPILE_DISABLE", "1") + +import torch +import torch.distributed as dist +from megatron.core import parallel_state +from megatron.core.pipeline_parallel.schedules import get_forward_backward_func + +from megatron.bridge import AutoBridge +from megatron.bridge.utils.common_utils import disable_mtp_for_inference + + +SIMILARITY_THRESHOLD = 0.98 + + +def _is_rank_0() -> bool: + if dist.is_initialized(): + return dist.get_rank() == 0 + return True + + +def print_rank_0(msg: str): + """Print a message only from rank 0.""" + if _is_rank_0(): + print(msg, flush=True) + + +# ========================================================================== # +# Image+Text Preprocessing +# ========================================================================== # + + +def preprocess_image_text(hf_model_path: str, prompt: str, image_path: str): + """Use the HF processor to preprocess an image+text prompt. + + Builds the chat template with image placeholder, runs the processor to get + input_ids, position_ids, pixel patches, grid_thw, and token_type_ids. + + Returns a dict with all tensors needed for both HF and Megatron forward. + """ + from transformers import AutoProcessor + + processor = AutoProcessor.from_pretrained(hf_model_path, trust_remote_code=True) + + # Build chat messages with image + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": prompt}, + {"type": "image_url", "image_url": {"url": image_path}}, + ], + } + ] + + # Apply chat template + text = processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + print_rank_0(f" Chat template text (first 200 chars): {text[:200]}...") + + # Process vision info (loads and resizes images) + image_inputs, video_inputs = processor.process_vision_info(messages) + print_rank_0(f" Images loaded: {len(image_inputs) if image_inputs else 0}") + + # Run the processor to get all inputs + proc_out = processor( + text=[text], + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt", + ) + + print_rank_0(f" Processor output keys: {list(proc_out.keys())}") + print_rank_0(f" input_ids shape: {proc_out['input_ids'].shape}") + if "images" in proc_out: + print_rank_0(f" images (pixel patches) shape: {proc_out['images'].shape}") + if "grid_thw" in proc_out: + print_rank_0(f" grid_thw: {proc_out['grid_thw']}") + if "position_ids" in proc_out: + print_rank_0(f" position_ids shape: {proc_out['position_ids'].shape}") + if "token_type_ids" in proc_out: + print_rank_0(f" token_type_ids shape: {proc_out['token_type_ids'].shape}") + + return proc_out, processor + + +# ========================================================================== # +# Phase 1: HF Model Forward +# ========================================================================== # + + +def run_hf_forward( + hf_model_path: str, + input_ids: torch.Tensor, + tokenizer, + processor_output=None, + processor=None, +) -> torch.Tensor: + """Run HF model forward pass on rank 0 and return last-token logits (on CPU). + + If processor_output is provided, uses image+text inputs. + Otherwise, runs text-only with simple 3D M-RoPE position_ids. + + Returns None on non-rank-0 processes. + """ + if not _is_rank_0(): + return None + + print_rank_0("=== Phase 1: Loading HF model ===") + from transformers import AutoModelForCausalLM + + # Load HF model directly onto GPU 0. Using device_map={"": device} avoids + # accelerate's AlignDevicesHook / meta-device dispatch, which breaks for + # data-dependent ops (torch.nonzero) in the ERNIE VL MoE routing. + # The 28B model fits in ~56 GB bf16 on a single 80 GB GPU. + hf_local_rank = int(os.environ.get("LOCAL_RANK", 0)) + device = torch.device(f"cuda:{hf_local_rank}") + hf_model = AutoModelForCausalLM.from_pretrained( + hf_model_path, + torch_dtype=torch.bfloat16, + device_map={"": device}, + trust_remote_code=True, + ).eval() + + print_rank_0(f" HF model loaded: {type(hf_model).__name__} on {device}") + + # Remove accelerate dispatch hooks. Even with device_map={"": device}, + # accelerate may attach AlignDevicesHook to some modules. These hooks + # route tensors through meta-device shape inference, which breaks for + # data-dependent ops (torch.nonzero) used in ERNIE VL MoE routing. + from accelerate.hooks import remove_hook_from_module + + for _name, _module in hf_model.named_modules(): + remove_hook_from_module(_module) + print_rank_0(" Removed accelerate dispatch hooks from all modules.") + + # Disable use_correction_bias on all MoE layers. This flag controls + # accumulation of expert usage statistics (expert_num_local) for the + # auxiliary-loss correction bias -- a training-only feature. During + # inference the code path is unnecessary and triggers a torch.nonzero() + # error when experts_type_mask boolean tensors are used for fancy + # indexing on meta-dispatched tensors. + _fixed_moe = 0 + for _name, _module in hf_model.named_modules(): + if hasattr(_module, "use_correction_bias") and _module.use_correction_bias: + _module.use_correction_bias = False + _fixed_moe += 1 + if _fixed_moe: + print_rank_0(f" Disabled use_correction_bias on {_fixed_moe} MoE layers (inference-only).") + + # Safety: fix inv_freq if stuck on meta device (only happens with + # device_map="auto"). With device_map={"": device} this is a no-op. + for name, module in hf_model.named_modules(): + if hasattr(module, "inv_freq") and isinstance(module.inv_freq, torch.Tensor): + if module.inv_freq.device.type == "meta": + dim = module.inv_freq.shape[0] * 2 # inv_freq has shape [dim//2] + theta = 10000.0 + module.inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + print_rank_0(f" Fixed meta inv_freq in {name} -> CPU, shape={module.inv_freq.shape}") + + if processor_output is not None: + # Image+text VL forward + # Register image preprocessor for GPU-side pixel normalization + if processor is not None and hasattr(hf_model, "add_image_preprocess"): + hf_model.add_image_preprocess(processor) + print_rank_0(" Image preprocessor registered on HF model") + + print_rank_0(" Running HF forward pass (image+text)...") + forward_kwargs = { + "input_ids": processor_output["input_ids"].to(device), + } + + # Map processor output keys to HF model forward parameter names + if "images" in processor_output: + forward_kwargs["images"] = processor_output["images"].to(device) + if "grid_thw" in processor_output: + forward_kwargs["grid_thw"] = processor_output["grid_thw"].to(device) + if "position_ids" in processor_output: + forward_kwargs["position_ids"] = processor_output["position_ids"].to(device) + if "token_type_ids" in processor_output: + # HF forward() expects token_type_ids with shape [bsz, seq_len+1]. + # The extra trailing element is a "next-token type" lookahead + # (normally appended by prepare_inputs_for_generation, but we call + # forward() directly). Append a zero (TokenType.text) to match. + tti = processor_output["token_type_ids"] + if tti.shape[1] == processor_output["input_ids"].shape[1]: + tti = torch.cat([tti, torch.zeros(tti.shape[0], 1, dtype=tti.dtype)], dim=1) + forward_kwargs["token_type_ids"] = tti.to(device) + if "image_type_ids" in processor_output: + forward_kwargs["image_type_ids"] = processor_output["image_type_ids"].to(device) + + with torch.no_grad(): + hf_output = hf_model(**forward_kwargs) + else: + # Text-only forward + print_rank_0(" Running HF forward pass (text-only)...") + seq_len = input_ids.size(1) + # 3D M-RoPE position_ids: [batch, seq_len, 3] -- for text-only, all 3 dims identical + position_ids = ( + torch.arange(seq_len, dtype=torch.long, device=device) + .unsqueeze(0) # [1, seq_len] + .unsqueeze(-1) # [1, seq_len, 1] + .expand(1, seq_len, 3) # [1, seq_len, 3] + .clone() + ) + with torch.no_grad(): + hf_output = hf_model( + input_ids=input_ids.to(device), + attention_mask=torch.ones_like(input_ids, dtype=torch.bool, device=device), + position_ids=position_ids, + ) + + hf_logits = hf_output.logits[0, -1, :].float().cpu() # [vocab_size] + + hf_next_token = torch.argmax(hf_logits) + top5_vals, top5_ids = torch.topk(hf_logits, 5) + top5_tokens = [tokenizer.decode([idx]) for idx in top5_ids] + + print_rank_0(f" HF logits shape: {hf_output.logits.shape}") + print_rank_0(f" HF logits stats: mean={hf_logits.mean():.4f}, std={hf_logits.std():.4f}") + print_rank_0(f" HF next token: {hf_next_token.item()} ('{tokenizer.decode([hf_next_token.item()])}')") + print_rank_0(f" HF Top 5: {list(zip(top5_tokens, top5_vals.tolist()))}") + + # Free HF model + del hf_model, hf_output + gc.collect() + torch.cuda.empty_cache() + print_rank_0(" HF model freed.") + + return hf_logits + + +# ========================================================================== # +# Phase 2: Megatron Model Forward +# ========================================================================== # + + +class SingleBatchIterator: + """Iterator that yields a single batch for Megatron forward scheduling.""" + + def __init__(self, batch): + self.batch = batch + self._yielded = False + + def __iter__(self): + return self + + def __next__(self): + if self._yielded: + raise StopIteration + self._yielded = True + return self.batch + + +def ernie_vl_forward_step(data_iterator, model, **kwargs): + """Forward step for ERNIE 4.5 VL model (text or image+text, no loss).""" + batch = next(data_iterator) + forward_args = { + "input_ids": batch["tokens"], + "mm_token_type_ids": batch["mm_token_type_ids"], + "attention_mask": None, + } + if "pixel_values" in batch and batch["pixel_values"] is not None: + forward_args["pixel_values"] = batch["pixel_values"] + if "image_grid_thw" in batch and batch["image_grid_thw"] is not None: + forward_args["image_grid_thw"] = batch["image_grid_thw"] + # Pass moe_mm_token_type_ids for dual-pool MoE routing (text vs vision experts). + # Without this, all tokens are routed to text_moe_layer only, which causes + # significant logit divergence for image+text inputs. + if "moe_mm_token_type_ids" in batch and batch["moe_mm_token_type_ids"] is not None: + forward_args["moe_mm_token_type_ids"] = batch["moe_mm_token_type_ids"] + + output = model(**forward_args) + if isinstance(output, tuple): + output = output[0] + + def loss_func(x, **kwargs): + return x + + return output, loss_func + + +def run_megatron_forward( + hf_model_path: str, + input_ids: torch.Tensor, + tokenizer, + tp: int = 1, + pp: int = 1, + ep: int = 1, + pixel_values=None, + image_grid_thw=None, + mm_token_type_ids=None, +) -> torch.Tensor: + """Load Megatron model via AutoBridge, run forward, return last-token logits (on CPU). + + Returns logits only on last pipeline stage + TP rank 0 + EP rank 0. + Returns None on other ranks. + """ + print_rank_0("=== Phase 2: Loading Megatron model via AutoBridge ===") + + bridge = AutoBridge.from_hf_pretrained( + hf_model_path, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + ) + model_provider = bridge.to_megatron_provider(load_weights=True) + model_provider.tensor_model_parallel_size = tp + model_provider.pipeline_model_parallel_size = pp + model_provider.expert_model_parallel_size = ep + model_provider.pipeline_dtype = torch.bfloat16 + model_provider.params_dtype = torch.bfloat16 + model_provider.finalize() + model_provider.initialize_model_parallel(seed=42) + + megatron_models = model_provider.provide_distributed_model(wrap_with_ddp=False) + + for m in megatron_models: + disable_mtp_for_inference(m) + m.eval() + if hasattr(m, "config"): + m.config.deallocate_pipeline_outputs = False + + print_rank_0(f" Megatron model built. {len(megatron_models)} component(s).") + + # Prepare input + seq_len = input_ids.size(1) + input_ids_cuda = input_ids.cuda() + + if mm_token_type_ids is None: + mm_token_type_ids = torch.zeros(1, seq_len, dtype=torch.int32, device="cuda") + else: + mm_token_type_ids = mm_token_type_ids.to(dtype=torch.int32, device="cuda") + + batch = { + "tokens": input_ids_cuda, + "mm_token_type_ids": mm_token_type_ids, + # moe_mm_token_type_ids drives dual-pool MoE routing: text tokens (0) + # go to text_moe_layer, vision tokens (>=1) go to vision_moe_layer. + # For the logit comparison, it uses the same values as mm_token_type_ids. + "moe_mm_token_type_ids": mm_token_type_ids.clone(), + } + if pixel_values is not None: + batch["pixel_values"] = pixel_values.cuda() + print_rank_0(f" pixel_values shape: {batch['pixel_values'].shape}") + if image_grid_thw is not None: + batch["image_grid_thw"] = image_grid_thw.cuda() + print_rank_0(f" image_grid_thw: {batch['image_grid_thw']}") + + # Forward + print_rank_0(" Running Megatron forward pass...") + with torch.no_grad(): + fwd_bwd_func = get_forward_backward_func() + iterator = SingleBatchIterator(batch) + + output = fwd_bwd_func( + forward_step_func=ernie_vl_forward_step, + data_iterator=iterator, + model=megatron_models, + num_microbatches=1, + forward_only=True, + seq_length=seq_len, + micro_batch_size=1, + collect_non_loss_data=True, + ) + + # Process output on last pipeline stage + is_last_stage = not dist.is_initialized() or parallel_state.is_pipeline_last_stage() + + megatron_logits_cpu = None + + if is_last_stage: + if isinstance(output, list) and len(output) > 0: + output = output[0] + + # Gather TP shards + if dist.is_initialized() and parallel_state.get_tensor_model_parallel_world_size() > 1: + world_size = parallel_state.get_tensor_model_parallel_world_size() + gathered = [torch.zeros_like(output) for _ in range(world_size)] + dist.all_gather(gathered, output, group=parallel_state.get_tensor_model_parallel_group()) + output = torch.cat(gathered, dim=2) + + megatron_logits = output[0, -1, :].float() # [padded_vocab_size] + + is_primary = not dist.is_initialized() or ( + parallel_state.get_tensor_model_parallel_rank() == 0 + and parallel_state.get_expert_model_parallel_rank() == 0 + ) + + if is_primary: + megatron_next_token = torch.argmax(megatron_logits) + top5_vals, top5_ids = torch.topk(megatron_logits, 5) + top5_tokens = [tokenizer.decode([idx]) for idx in top5_ids] + + print_rank_0(f" Megatron output shape: {output.shape}") + print_rank_0( + f" Megatron logits stats: mean={megatron_logits.mean():.4f}, std={megatron_logits.std():.4f}" + ) + print_rank_0( + f" Megatron next token: {megatron_next_token.item()} ('{tokenizer.decode([megatron_next_token.item()])}')" + ) + print_rank_0(f" Megatron Top 5: {list(zip(top5_tokens, top5_vals.tolist()))}") + + megatron_logits_cpu = megatron_logits.cpu() + + return megatron_logits_cpu + + +# ========================================================================== # +# Phase 3: Comparison +# ========================================================================== # + + +def compare_logits( + hf_logits: torch.Tensor, megatron_logits: torch.Tensor, tokenizer, threshold: float = SIMILARITY_THRESHOLD +): + """Compare HF and Megatron logits. Returns True if pass.""" + print_rank_0("\n=== Phase 3: Comparing Logits ===") + + # Truncate Megatron logits to HF vocab size (Megatron may pad vocab) + hf_vocab_size = hf_logits.shape[0] + megatron_logits_cmp = megatron_logits[:hf_vocab_size] + + # Token match + hf_next = torch.argmax(hf_logits) + mg_next = torch.argmax(megatron_logits_cmp) + token_match = hf_next.item() == mg_next.item() + + hf_decoded = tokenizer.decode([hf_next.item()]) + mg_decoded = tokenizer.decode([mg_next.item()]) + + print_rank_0(f" HF next token: {hf_next.item()} ('{hf_decoded}')") + print_rank_0(f" Megatron next token: {mg_next.item()} ('{mg_decoded}')") + print_rank_0(f" Token match: {token_match}") + + # Cosine similarity + cosine_sim = torch.cosine_similarity( + hf_logits.unsqueeze(0).float(), + megatron_logits_cmp.unsqueeze(0).float(), + ).item() + print_rank_0(f" Cosine similarity: {cosine_sim:.6f} ({cosine_sim * 100:.2f}%)") + + # Absolute diff + diff = (hf_logits.float() - megatron_logits_cmp.float()).abs() + print_rank_0(f" Logits diff: max={diff.max():.6f}, mean={diff.mean():.6f}, median={diff.median():.6f}") + + # Top-5 overlap + hf_top5 = set(torch.topk(hf_logits, 5).indices.tolist()) + mg_top5 = set(torch.topk(megatron_logits_cmp, 5).indices.tolist()) + overlap = len(hf_top5 & mg_top5) + print_rank_0(f" Top-5 overlap: {overlap}/5") + + passed = cosine_sim >= threshold + status = "PASS" if passed else "FAIL" + within = "within" if passed else "outside" + print_rank_0(f"\n Result: {status} (cosine {cosine_sim:.4f} {within} threshold {threshold})") + + return passed + + +# ========================================================================== # +# Main +# ========================================================================== # + + +def main(): + """Run ERNIE 4.5 VL MoE logit comparison between HF and Megatron.""" + parser = argparse.ArgumentParser(description="ERNIE 4.5 VL MoE logit comparison (HF vs Megatron)") + parser.add_argument("--hf-model-path", required=True, help="Path to HF model directory") + parser.add_argument("--prompt", default="Hello, how are you?", help="Text prompt for comparison") + parser.add_argument("--tp", type=int, default=1, help="Tensor parallelism size") + parser.add_argument("--pp", type=int, default=1, help="Pipeline parallelism size") + parser.add_argument("--ep", type=int, default=1, help="Expert parallelism size") + parser.add_argument("--threshold", type=float, default=SIMILARITY_THRESHOLD, help="Cosine similarity threshold") + parser.add_argument("--image-path", type=str, default=None, help="Path to image file for VL inference") + args = parser.parse_args() + + print_rank_0("=== ERNIE 4.5 VL Logit Comparison ===") + print_rank_0(f" Model: {args.hf_model_path}") + print_rank_0(f" Prompt: '{args.prompt}'") + print_rank_0(f" Image: {args.image_path or '(none, text-only)'}") + print_rank_0(f" TP={args.tp}, PP={args.pp}, EP={args.ep}") + print_rank_0(f" Threshold: {args.threshold}") + + # Setup distributed (must happen before Phase 1 so _is_rank_0() works) + local_rank = int(os.environ.get("LOCAL_RANK", 0)) + torch.cuda.set_device(local_rank) + if not dist.is_initialized(): + dist.init_process_group(backend="nccl") + + # Load tokenizer + from transformers import AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained( + args.hf_model_path, + trust_remote_code=True, + ) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + # Prepare inputs depending on mode (text-only vs image+text) + processor_output = None + processor = None + pixel_values = None + image_grid_thw = None + mm_token_type_ids = None + + if args.image_path: + # ============================================================ + # Image+Text mode: use processor to prepare all inputs + # ============================================================ + print_rank_0("\n=== Preprocessing: Image+Text ===") + processor_output, processor = preprocess_image_text(args.hf_model_path, args.prompt, args.image_path) + input_ids = processor_output["input_ids"] + + # Extract vision tensors for Megatron side + if "images" in processor_output: + pixel_values = processor_output["images"] # [N_patches, 588] + if "grid_thw" in processor_output: + image_grid_thw = processor_output["grid_thw"] # [num_images, 3] + + # Build mm_token_type_ids for Megatron from processor's token_type_ids + # The HF processor outputs token_type_ids with shape [bsz, seq_len+1] + # (extra token for shifted labels). The Megatron model expects + # mm_token_type_ids with shape [bsz, seq_len]. + if "token_type_ids" in processor_output: + hf_token_type_ids = processor_output["token_type_ids"] + # Take the first seq_len values (drop the extra trailing token) + mm_token_type_ids = hf_token_type_ids[:, : input_ids.size(1)].to(torch.int32) + num_img_tokens = (mm_token_type_ids == 1).sum().item() + print_rank_0( + f" mm_token_type_ids: {mm_token_type_ids.shape}, image tokens: {num_img_tokens}/{input_ids.size(1)}" + ) + else: + # ============================================================ + # Text-only mode: simple tokenization + # ============================================================ + inputs = tokenizer(args.prompt, return_tensors="pt") + input_ids = inputs.input_ids + + # Pad sequence length to be divisible by TP size (needed for sequence parallel) + tp_size = args.tp + seq_len = input_ids.size(1) + remainder = seq_len % tp_size + if remainder != 0: + pad_len = tp_size - remainder + padding = torch.full( + (input_ids.shape[0], pad_len), + tokenizer.pad_token_id or 0, + dtype=input_ids.dtype, + ) + input_ids = torch.cat([input_ids, padding], dim=1) + # Also pad mm_token_type_ids if present + if mm_token_type_ids is not None: + mm_padding = torch.zeros( + mm_token_type_ids.shape[0], + pad_len, + dtype=mm_token_type_ids.dtype, + ) + mm_token_type_ids = torch.cat([mm_token_type_ids, mm_padding], dim=1) + + print_rank_0(f" Input IDs shape: {input_ids.shape} (original seq_len={seq_len})") + + # Phase 1: HF forward (rank 0 only) + hf_logits = run_hf_forward( + args.hf_model_path, + input_ids, + tokenizer, + processor_output=processor_output, + processor=processor, + ) + + # Synchronize all ranks before Phase 2 + dist.barrier() + + # Destroy the process group before Phase 2 since AutoBridge's + # initialize_model_parallel() will create its own process groups. + dist.destroy_process_group() + + # Phase 2: Megatron forward (all ranks) + megatron_logits = run_megatron_forward( + args.hf_model_path, + input_ids, + tokenizer, + tp=args.tp, + pp=args.pp, + ep=args.ep, + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, + mm_token_type_ids=mm_token_type_ids, + ) + + # Phase 3: Compare (rank 0 only, which has both logit tensors) + if _is_rank_0() and megatron_logits is not None and hf_logits is not None: + passed = compare_logits(hf_logits, megatron_logits, tokenizer, args.threshold) + else: + passed = True # Non-primary ranks don't compare + + # Broadcast pass/fail to all ranks + if dist.is_initialized(): + passed_tensor = torch.tensor([1 if passed else 0], device="cuda") + dist.broadcast(passed_tensor, 0) + passed = passed_tensor.item() == 1 + dist.barrier() + + if dist.is_initialized(): + dist.destroy_process_group() + + sys.exit(0 if passed else 1) + + +if __name__ == "__main__": + main() diff --git a/examples/models/vlm/ernie_vl/ernie45_vl_vit_compare.py b/examples/models/vlm/ernie_vl/ernie45_vl_vit_compare.py new file mode 100644 index 0000000000..859b880bcd --- /dev/null +++ b/examples/models/vlm/ernie_vl/ernie45_vl_vit_compare.py @@ -0,0 +1,474 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Vision encoder (ViT) alignment test: HF ViT vs MG-native ViT. + +Compares hidden state outputs from the HuggingFace +Ernie4_5_VLMoeVisionTransformerPretrainedModel and the Megatron-Core native +ErnieVLVisionModel using the same weights and identical input. + +Strategy (single GPU, no torchrun required): + 1. Load HF vision model from the ERNIE 4.5 VL checkpoint + 2. Construct MG ViT (ErnieVLVisionModel) with matching config + 3. Transfer weights from HF state_dict to MG state_dict + 4. Generate dummy pixel patches + grid_thw + 5. Forward both models + 6. Compare: cosine similarity, max absolute diff, relative diff + +This test validates that the MG-native ViT produces the same output as the +HF ViT, ensuring weight conversion correctness and architectural fidelity. + +Usage: + # With the real 28B model: + python ernie45_vl_vit_compare.py \ + --hf-model-path /path/to/ERNIE-4.5-VL-28B-A3B-Thinking + + # With a toy model (created by test_ernie45_vl_conversion.py): + python ernie45_vl_vit_compare.py \ + --hf-model-path /tmp/ernie45_vl_toy + +Exit code 0 = PASS, non-zero = FAIL. +""" + +import argparse +import logging +import os +import sys + + +os.environ.setdefault("TORCH_COMPILE_DISABLE", "1") + +import torch +import torch.nn.functional as F + + +def build_hf_vit(hf_model_path: str, device: torch.device): + """Load the HF vision encoder from an ERNIE 4.5 VL checkpoint. + + Returns (hf_vit_model, vision_config, hf_config). + """ + from transformers import AutoConfig + from transformers.models.ernie4_5_vl_moe.modeling_ernie4_5_vl_moe import ( + Ernie4_5_VLMoeVisionTransformerPretrainedModel, + ) + + hf_config = AutoConfig.from_pretrained(hf_model_path, trust_remote_code=True) + + # Extract vision config + vision_config = getattr(hf_config, "vision_config", None) + if vision_config is None: + raise ValueError("HF config does not have vision_config") + + # Normalize vision config (add missing attributes for compat) + from megatron.bridge.models.ernie_vl.modeling_ernie45_vl import ( + _normalize_vision_config, + ) + + _normalize_vision_config(vision_config, hf_config=hf_config) + + # Create HF ViT from config (random weights) + hf_vit = Ernie4_5_VLMoeVisionTransformerPretrainedModel._from_config(vision_config) + + # Load weights from the checkpoint + # HF ViT weights are prefixed with "vision_model." or "model.vision_model." on disk + from pathlib import Path + + from safetensors import safe_open + + model_dir = Path(hf_model_path) + safetensors_files = list(model_dir.glob("*.safetensors")) + + if not safetensors_files: + raise FileNotFoundError(f"No safetensors files found in {hf_model_path}") + + # Determine the on-disk prefix for vision weights + # Flat config (Thinking): "vision_model." + # Nested config (Instruct): "model.vision_model." + text_cfg_attr = getattr(hf_config, "text_config", None) + is_flat = (text_cfg_attr is None) or (text_cfg_attr is hf_config) + vision_prefix = "vision_model." if is_flat else "model.vision_model." + + vision_state_dict = {} + for st_file in safetensors_files: + with safe_open(str(st_file), framework="pt", device="cpu") as f: + for key in f.keys(): + if key.startswith(vision_prefix): + # Strip the prefix to get the HF in-memory key + local_key = key[len(vision_prefix) :] + vision_state_dict[local_key] = f.get_tensor(key) + + if not vision_state_dict: + raise ValueError(f"No vision weights found with prefix '{vision_prefix}' in {hf_model_path}") + + # Load state dict into HF ViT + missing, unexpected = hf_vit.load_state_dict(vision_state_dict, strict=False) + if missing: + logging.warning("Missing keys in HF ViT: %s...", missing[:5]) + if unexpected: + logging.warning("Unexpected keys in HF ViT: %s...", unexpected[:5]) + + hf_vit = hf_vit.to(device=device, dtype=torch.bfloat16).eval() + print(f"HF ViT loaded: {sum(p.numel() for p in hf_vit.parameters())} params") + return hf_vit, vision_config, hf_config + + +def build_mg_vit(vision_config, hf_config, device: torch.device): + """Construct the MG-native ErnieVLVisionModel. + + Returns the model on the specified device. + """ + from megatron.bridge.models.ernie_vl.vision_layer_spec import get_ernie_vit_layer_spec + from megatron.bridge.models.ernie_vl.vision_model import ErnieVLVisionModel + from megatron.bridge.models.ernie_vl.vision_transformer_config import ( + get_ernie_vision_config, + ) + + # Build transformer config from vision config + vit_config = get_ernie_vision_config(vision_config) + vit_layer_spec = get_ernie_vit_layer_spec() + + mg_vit = ErnieVLVisionModel( + transformer_config=vit_config, + transformer_layer_spec=vit_layer_spec, + ) + + mg_vit = mg_vit.to(device=device, dtype=torch.bfloat16).eval() + print(f"MG ViT created: {sum(p.numel() for p in mg_vit.parameters())} params") + return mg_vit + + +def _interleave_qkv(hf_qkv, num_heads): + """Convert HF contiguous QKV layout to Megatron interleaved layout. + + HF stores fused QKV as contiguous blocks: + [Q0, Q1, ..., Q_{H-1} | K0, K1, ..., K_{H-1} | V0, V1, ..., V_{H-1}] + + Megatron's get_query_key_value_tensors() expects GQA-interleaved layout: + [Q0, K0, V0 | Q1, K1, V1 | ... | Q_{H-1}, K_{H-1}, V_{H-1}] + + For MHA (num_query_groups == num_heads), each group has exactly one Q, K, V head. + + Args: + hf_qkv: Tensor of shape [3*hidden, ...] in HF contiguous layout. + For weight: [3*hidden, hidden], for bias: [3*hidden]. + num_heads: Number of attention heads. + + Returns: + Tensor of same shape but with Megatron interleaved layout. + """ + hidden = hf_qkv.shape[0] // 3 + head_dim = hidden // num_heads + + # Split into Q, K, V + q_all = hf_qkv[:hidden] + k_all = hf_qkv[hidden : 2 * hidden] + v_all = hf_qkv[2 * hidden :] + + # Reshape to per-head: [num_heads, head_dim, ...] + q_heads = q_all.reshape(num_heads, head_dim, *q_all.shape[1:]) + k_heads = k_all.reshape(num_heads, head_dim, *k_all.shape[1:]) + v_heads = v_all.reshape(num_heads, head_dim, *v_all.shape[1:]) + + # Interleave: [Q0, K0, V0, Q1, K1, V1, ...] + interleaved = torch.stack([q_heads, k_heads, v_heads], dim=1) + return interleaved.reshape(hf_qkv.shape) + + +def transfer_weights_hf_to_mg(hf_vit, mg_vit, vision_config): + """Transfer weights from HF ViT to MG ViT using direct state_dict mapping. + + The HF ViT state_dict keys map to MG ViT keys as follows: + HF: patch_embed.proj.weight -> MG: patch_embed.proj.weight + HF: blocks.{i}.attn.qkv.weight/bias -> MG: decoder.layers.{i}.self_attention.linear_qkv.weight/bias + HF: blocks.{i}.attn.proj.weight/bias -> MG: decoder.layers.{i}.self_attention.linear_proj.weight/bias + HF: blocks.{i}.norm1.weight/bias -> MG: decoder.layers.{i}.self_attention.linear_qkv.layer_norm_weight/bias + HF: blocks.{i}.norm2.weight/bias -> MG: decoder.layers.{i}.mlp.linear_fc1.layer_norm_weight/bias + HF: blocks.{i}.mlp.fc1.weight/bias -> MG: decoder.layers.{i}.mlp.linear_fc1.weight/bias + HF: blocks.{i}.mlp.fc2.weight/bias -> MG: decoder.layers.{i}.mlp.linear_fc2.weight/bias + HF: ln.weight/bias -> MG: decoder.final_layernorm.weight/bias + + Note: The fused QKV weight/bias must be interleaved from HF's contiguous + [Q|K|V] format to Megatron's per-head [Q0,K0,V0|Q1,K1,V1|...] format. + This interleaving is required because Megatron's get_query_key_value_tensors() + splits the fused QKV output by query groups (= per head for MHA). + """ + hf_sd = hf_vit.state_dict() + mg_sd = mg_vit.state_dict() + + num_heads = getattr(vision_config, "num_heads", getattr(vision_config, "num_attention_heads", 16)) + + # Build mapping: HF key -> MG key + key_mapping = {} + # Track which HF keys need QKV interleaving + qkv_keys = set() + + # Patch embed + key_mapping["patch_embed.proj.weight"] = "patch_embed.proj.weight" + + # Final layernorm + key_mapping["ln.weight"] = "decoder.final_layernorm.weight" + key_mapping["ln.bias"] = "decoder.final_layernorm.bias" + + # Per-block mappings + num_layers = getattr(vision_config, "depth", getattr(vision_config, "num_hidden_layers", 32)) + for i in range(num_layers): + # QKV (fused) - needs interleaving + qkv_w_key = f"blocks.{i}.attn.qkv.weight" + qkv_b_key = f"blocks.{i}.attn.qkv.bias" + key_mapping[qkv_w_key] = f"decoder.layers.{i}.self_attention.linear_qkv.weight" + key_mapping[qkv_b_key] = f"decoder.layers.{i}.self_attention.linear_qkv.bias" + qkv_keys.add(qkv_w_key) + qkv_keys.add(qkv_b_key) + # Proj + key_mapping[f"blocks.{i}.attn.proj.weight"] = f"decoder.layers.{i}.self_attention.linear_proj.weight" + key_mapping[f"blocks.{i}.attn.proj.bias"] = f"decoder.layers.{i}.self_attention.linear_proj.bias" + # Norm1 -> fused into linear_qkv + key_mapping[f"blocks.{i}.norm1.weight"] = f"decoder.layers.{i}.self_attention.linear_qkv.layer_norm_weight" + key_mapping[f"blocks.{i}.norm1.bias"] = f"decoder.layers.{i}.self_attention.linear_qkv.layer_norm_bias" + # Norm2 -> fused into linear_fc1 + key_mapping[f"blocks.{i}.norm2.weight"] = f"decoder.layers.{i}.mlp.linear_fc1.layer_norm_weight" + key_mapping[f"blocks.{i}.norm2.bias"] = f"decoder.layers.{i}.mlp.linear_fc1.layer_norm_bias" + # MLP fc1 + key_mapping[f"blocks.{i}.mlp.fc1.weight"] = f"decoder.layers.{i}.mlp.linear_fc1.weight" + key_mapping[f"blocks.{i}.mlp.fc1.bias"] = f"decoder.layers.{i}.mlp.linear_fc1.bias" + # MLP fc2 + key_mapping[f"blocks.{i}.mlp.fc2.weight"] = f"decoder.layers.{i}.mlp.linear_fc2.weight" + key_mapping[f"blocks.{i}.mlp.fc2.bias"] = f"decoder.layers.{i}.mlp.linear_fc2.bias" + + # Transfer weights + transferred = 0 + for hf_key, mg_key in key_mapping.items(): + if hf_key not in hf_sd: + print(f" WARNING: HF key not found: {hf_key}") + continue + if mg_key not in mg_sd: + print(f" WARNING: MG key not found: {mg_key}") + continue + + hf_tensor = hf_sd[hf_key] + mg_tensor = mg_sd[mg_key] + + if hf_tensor.shape != mg_tensor.shape: + print(f" WARNING: Shape mismatch for {hf_key} -> {mg_key}: {hf_tensor.shape} vs {mg_tensor.shape}") + continue + + # Apply QKV interleaving for fused QKV weight and bias + if hf_key in qkv_keys: + hf_tensor = _interleave_qkv(hf_tensor, num_heads) + + mg_sd[mg_key] = hf_tensor.to(dtype=mg_tensor.dtype) + transferred += 1 + + # Load the mapped state dict + mg_vit.load_state_dict(mg_sd, strict=True) + print(f"Transferred {transferred}/{len(key_mapping)} weight tensors") + + # Check for any MG keys that weren't covered + mapped_mg_keys = set(key_mapping.values()) + unmapped_mg_keys = [k for k in mg_sd if k not in mapped_mg_keys] + if unmapped_mg_keys: + print(f" WARNING: {len(unmapped_mg_keys)} MG keys not mapped: {unmapped_mg_keys[:5]}...") + + +def generate_dummy_input(vision_config, device: torch.device, num_images: int = 2): + """Generate dummy pixel patches and grid_thw for testing. + + Creates random pixel patches that mimic the output of the ERNIE 4.5 VL + processor (pre-flattened patches of shape [total_patches, C*P*P]). + + Args: + vision_config: HF vision config with patch_size, spatial_merge_size. + device: Target device. + num_images: Number of images to simulate. + + Returns: + (pixel_values, grid_thw): Dummy inputs. + """ + patch_size = getattr(vision_config, "patch_size", 14) + in_channels = getattr(vision_config, "in_channels", 3) + spatial_merge = getattr(vision_config, "spatial_merge_size", 2) + + # Create images of varying sizes (must be divisible by spatial_merge_size) + # Use small sizes for efficiency + grid_thw_list = [] + for i in range(num_images): + t = 1 # Single frame + h = spatial_merge * (2 + i) # e.g., 4, 6 for merge_size=2 + w = spatial_merge * (2 + i) + grid_thw_list.append([t, h, w]) + + grid_thw = torch.tensor(grid_thw_list, dtype=torch.long, device=device) + total_patches = int(torch.prod(grid_thw, dim=1).sum().item()) + + # Generate random pixel patches (simulating processor output) + # Shape: [total_patches, C * patch_size^2] + pixel_values = torch.randn( + total_patches, + in_channels * patch_size * patch_size, + dtype=torch.bfloat16, + device=device, + ) + + return pixel_values, grid_thw + + +def run_hf_vit_forward(hf_vit, pixel_values, grid_thw): + """Run HF ViT forward pass and return hidden states.""" + with torch.no_grad(): + output = hf_vit(pixel_values, grid_thw, return_dict=True) + return output.last_hidden_state + + +def run_mg_vit_forward(mg_vit, pixel_values, grid_thw): + """Run MG ViT forward pass and return hidden states.""" + with torch.no_grad(): + output = mg_vit(pixel_values, grid_thw) + return output + + +def compare_outputs(hf_out: torch.Tensor, mg_out: torch.Tensor, threshold: float = 0.99): + """Compare HF and MG ViT outputs and return (passed, stats_dict).""" + assert hf_out.shape == mg_out.shape, f"Shape mismatch: HF={hf_out.shape} vs MG={mg_out.shape}" + + hf_flat = hf_out.float().flatten() + mg_flat = mg_out.float().flatten() + + # Cosine similarity + cos_sim = F.cosine_similarity(hf_flat.unsqueeze(0), mg_flat.unsqueeze(0)).item() + + # Absolute difference + abs_diff = (hf_flat - mg_flat).abs() + max_abs_diff = abs_diff.max().item() + mean_abs_diff = abs_diff.mean().item() + + # Relative difference + denom = hf_flat.abs().clamp(min=1e-8) + rel_diff = abs_diff / denom + max_rel_diff = rel_diff.max().item() + mean_rel_diff = rel_diff.mean().item() + + stats = { + "cosine_similarity": cos_sim, + "max_abs_diff": max_abs_diff, + "mean_abs_diff": mean_abs_diff, + "max_rel_diff": max_rel_diff, + "mean_rel_diff": mean_rel_diff, + } + + passed = cos_sim >= threshold + return passed, stats + + +def main(): + """Run ERNIE 4.5 VL ViT alignment test.""" + parser = argparse.ArgumentParser(description="ERNIE 4.5 VL ViT alignment test") + parser.add_argument( + "--hf-model-path", + type=str, + required=True, + help="Path to the HF ERNIE 4.5 VL model directory", + ) + parser.add_argument( + "--num-images", + type=int, + default=2, + help="Number of dummy images to generate (default: 2)", + ) + parser.add_argument( + "--threshold", + type=float, + default=0.99, + help="Cosine similarity threshold for pass/fail (default: 0.99)", + ) + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"Device: {device}") + print(f"Model path: {args.hf_model_path}") + + # Initialize megatron parallel state for MG ViT (single GPU, TP=1) + # This is needed because TransformerBlock uses parallel_state internally. + if not torch.distributed.is_initialized(): + torch.distributed.init_process_group( + backend="nccl" if torch.cuda.is_available() else "gloo", + init_method="tcp://127.0.0.1:29500", + world_size=1, + rank=0, + ) + from megatron.core import parallel_state + + if not parallel_state.is_initialized(): + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + ) + + # Step 1: Load HF ViT + print("\n=== Step 1: Loading HF ViT ===") + hf_vit, vision_config, hf_config = build_hf_vit(args.hf_model_path, device) + + # Step 2: Build MG ViT + print("\n=== Step 2: Building MG ViT ===") + mg_vit = build_mg_vit(vision_config, hf_config, device) + + # Step 3: Transfer weights + print("\n=== Step 3: Transferring weights HF -> MG ===") + transfer_weights_hf_to_mg(hf_vit, mg_vit, vision_config) + + # Step 4: Generate dummy input + print("\n=== Step 4: Generating dummy input ===") + pixel_values, grid_thw = generate_dummy_input(vision_config, device, num_images=args.num_images) + print(f" pixel_values: {pixel_values.shape}, dtype={pixel_values.dtype}") + print(f" grid_thw: {grid_thw}") + print(f" total_patches: {pixel_values.shape[0]}") + + # Step 5: Forward pass + print("\n=== Step 5: Running forward passes ===") + hf_out = run_hf_vit_forward(hf_vit, pixel_values, grid_thw) + mg_out = run_mg_vit_forward(mg_vit, pixel_values, grid_thw) + print(f" HF output: {hf_out.shape}, dtype={hf_out.dtype}") + print(f" MG output: {mg_out.shape}, dtype={mg_out.dtype}") + + # Step 6: Compare + print("\n=== Step 6: Comparing outputs ===") + passed, stats = compare_outputs(hf_out, mg_out, threshold=args.threshold) + + print(f" Cosine similarity: {stats['cosine_similarity']:.8f}") + print(f" Max abs diff: {stats['max_abs_diff']:.6e}") + print(f" Mean abs diff: {stats['mean_abs_diff']:.6e}") + print(f" Max rel diff: {stats['max_rel_diff']:.6e}") + print(f" Mean rel diff: {stats['mean_rel_diff']:.6e}") + print(f" Threshold: {args.threshold}") + + if passed: + print("\nRESULT: PASS -- MG ViT output matches HF ViT") + else: + print("\nRESULT: FAIL -- MG ViT output does NOT match HF ViT") + + # Cleanup + del hf_vit, mg_vit + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + # Destroy parallel state + parallel_state.destroy_model_parallel() + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + sys.exit(0 if passed else 1) + + +if __name__ == "__main__": + main() diff --git a/examples/models/vlm/ernie_vl/ernie45_vl_vit_debug.py b/examples/models/vlm/ernie_vl/ernie45_vl_vit_debug.py new file mode 100644 index 0000000000..e6ec0d0bce --- /dev/null +++ b/examples/models/vlm/ernie_vl/ernie45_vl_vit_debug.py @@ -0,0 +1,366 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Layer-by-layer ViT alignment debug script. + +Compares HF ViT vs MG ViT at each stage: + 1. PatchEmbed output + 2. RoPE computation + 3. After each transformer block + 4. After final LayerNorm + +This helps isolate exactly WHERE the output divergence occurs. +""" + +import argparse +import os +import sys + + +os.environ.setdefault("TORCH_COMPILE_DISABLE", "1") + +import torch +import torch.nn.functional as F + + +def main(): + """Run layer-by-layer ViT alignment debug.""" + parser = argparse.ArgumentParser(description="ERNIE 4.5 VL ViT debug") + parser.add_argument("--hf-model-path", type=str, required=True) + parser.add_argument("--num-images", type=int, default=2) + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + print(f"Device: {device}") + + # Initialize megatron parallel state + if not torch.distributed.is_initialized(): + torch.distributed.init_process_group( + backend="nccl" if torch.cuda.is_available() else "gloo", + init_method="tcp://127.0.0.1:29501", + world_size=1, + rank=0, + ) + from megatron.core import parallel_state + + if not parallel_state.is_initialized(): + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + ) + + # ===================================================================== + # Load HF ViT + # ===================================================================== + from transformers import AutoConfig + from transformers.models.ernie4_5_vl_moe.modeling_ernie4_5_vl_moe import ( + Ernie4_5_VLMoeVisionTransformerPretrainedModel, + ) + + hf_config = AutoConfig.from_pretrained(args.hf_model_path, trust_remote_code=True) + vision_config = getattr(hf_config, "vision_config", None) + + from megatron.bridge.models.ernie_vl.modeling_ernie45_vl import _normalize_vision_config + + _normalize_vision_config(vision_config, hf_config=hf_config) + + hf_vit = Ernie4_5_VLMoeVisionTransformerPretrainedModel._from_config(vision_config) + + # Load weights from safetensors + from pathlib import Path + + from safetensors import safe_open + + model_dir = Path(args.hf_model_path) + text_cfg_attr = getattr(hf_config, "text_config", None) + is_flat = (text_cfg_attr is None) or (text_cfg_attr is hf_config) + vision_prefix = "vision_model." if is_flat else "model.vision_model." + + vision_state_dict = {} + for st_file in model_dir.glob("*.safetensors"): + with safe_open(str(st_file), framework="pt", device="cpu") as f: + for key in f.keys(): + if key.startswith(vision_prefix): + local_key = key[len(vision_prefix) :] + vision_state_dict[local_key] = f.get_tensor(key) + + hf_vit.load_state_dict(vision_state_dict, strict=False) + hf_vit = hf_vit.to(device=device, dtype=torch.bfloat16).eval() + print(f"HF ViT loaded: {sum(p.numel() for p in hf_vit.parameters())} params") + + # ===================================================================== + # Build MG ViT + # ===================================================================== + from megatron.bridge.models.ernie_vl.vision_layer_spec import get_ernie_vit_layer_spec + from megatron.bridge.models.ernie_vl.vision_model import ErnieVLVisionModel + from megatron.bridge.models.ernie_vl.vision_transformer_config import get_ernie_vision_config + + vit_config = get_ernie_vision_config(vision_config) + vit_layer_spec = get_ernie_vit_layer_spec() + mg_vit = ErnieVLVisionModel( + transformer_config=vit_config, + transformer_layer_spec=vit_layer_spec, + ) + mg_vit = mg_vit.to(device=device, dtype=torch.bfloat16).eval() + print(f"MG ViT created: {sum(p.numel() for p in mg_vit.parameters())} params") + + # Transfer weights + hf_sd = hf_vit.state_dict() + mg_sd = mg_vit.state_dict() + + # Build mapping + key_mapping = {} + key_mapping["patch_embed.proj.weight"] = "patch_embed.proj.weight" + key_mapping["ln.weight"] = "decoder.final_layernorm.weight" + key_mapping["ln.bias"] = "decoder.final_layernorm.bias" + + num_layers = getattr(vision_config, "depth", getattr(vision_config, "num_hidden_layers", 32)) + for i in range(num_layers): + key_mapping[f"blocks.{i}.attn.qkv.weight"] = f"decoder.layers.{i}.self_attention.linear_qkv.weight" + key_mapping[f"blocks.{i}.attn.qkv.bias"] = f"decoder.layers.{i}.self_attention.linear_qkv.bias" + key_mapping[f"blocks.{i}.attn.proj.weight"] = f"decoder.layers.{i}.self_attention.linear_proj.weight" + key_mapping[f"blocks.{i}.attn.proj.bias"] = f"decoder.layers.{i}.self_attention.linear_proj.bias" + key_mapping[f"blocks.{i}.norm1.weight"] = f"decoder.layers.{i}.self_attention.linear_qkv.layer_norm_weight" + key_mapping[f"blocks.{i}.norm1.bias"] = f"decoder.layers.{i}.self_attention.linear_qkv.layer_norm_bias" + key_mapping[f"blocks.{i}.norm2.weight"] = f"decoder.layers.{i}.mlp.linear_fc1.layer_norm_weight" + key_mapping[f"blocks.{i}.norm2.bias"] = f"decoder.layers.{i}.mlp.linear_fc1.layer_norm_bias" + key_mapping[f"blocks.{i}.mlp.fc1.weight"] = f"decoder.layers.{i}.mlp.linear_fc1.weight" + key_mapping[f"blocks.{i}.mlp.fc1.bias"] = f"decoder.layers.{i}.mlp.linear_fc1.bias" + key_mapping[f"blocks.{i}.mlp.fc2.weight"] = f"decoder.layers.{i}.mlp.linear_fc2.weight" + key_mapping[f"blocks.{i}.mlp.fc2.bias"] = f"decoder.layers.{i}.mlp.linear_fc2.bias" + + transferred = 0 + for hf_key, mg_key in key_mapping.items(): + if hf_key in hf_sd and mg_key in mg_sd and hf_sd[hf_key].shape == mg_sd[mg_key].shape: + mg_sd[mg_key] = hf_sd[hf_key].to(dtype=mg_sd[mg_key].dtype) + transferred += 1 + mg_vit.load_state_dict(mg_sd, strict=True) + print(f"Transferred {transferred}/{len(key_mapping)} weight tensors") + + # ===================================================================== + # Generate dummy input + # ===================================================================== + patch_size = getattr(vision_config, "patch_size", 14) + in_channels = getattr(vision_config, "in_channels", 3) + spatial_merge = getattr(vision_config, "spatial_merge_size", 2) + + grid_thw_list = [] + for i in range(args.num_images): + t, h, w = 1, spatial_merge * (2 + i), spatial_merge * (2 + i) + grid_thw_list.append([t, h, w]) + + grid_thw = torch.tensor(grid_thw_list, dtype=torch.long, device=device) + total_patches = int(torch.prod(grid_thw, dim=1).sum().item()) + + pixel_values = torch.randn( + total_patches, + in_channels * patch_size * patch_size, + dtype=torch.bfloat16, + device=device, + ) + print(f"\nInput: pixel_values {pixel_values.shape}, grid_thw {grid_thw}") + + # ===================================================================== + # Test 1: PatchEmbed + # ===================================================================== + print("\n" + "=" * 60) + print("TEST 1: PatchEmbed") + print("=" * 60) + with torch.no_grad(): + hf_patch_out = hf_vit.patch_embed(pixel_values) + mg_patch_out = mg_vit.patch_embed(pixel_values) + cos_sim = F.cosine_similarity( + hf_patch_out.float().flatten().unsqueeze(0), mg_patch_out.float().flatten().unsqueeze(0) + ).item() + max_diff = (hf_patch_out.float() - mg_patch_out.float()).abs().max().item() + print(f" PatchEmbed cosine_sim: {cos_sim:.8f}, max_diff: {max_diff:.6e}") + assert cos_sim > 0.999, f"PatchEmbed MISMATCH: cos_sim = {cos_sim}" + print(" PASS") + + # ===================================================================== + # Test 2: RoPE + # ===================================================================== + print("\n" + "=" * 60) + print("TEST 2: RoPE") + print("=" * 60) + with torch.no_grad(): + # HF RoPE + hf_rotary = hf_vit.rot_pos_emb(grid_thw) # [N, 40] + hf_emb = torch.cat((hf_rotary, hf_rotary), dim=-1) # [N, 80] + hf_cos = hf_emb.cos() + hf_sin = hf_emb.sin() + + # MG RoPE + mg_rotary = mg_vit.rot_pos_emb(grid_thw) # [N, 40] + mg_emb = torch.cat( + (mg_rotary.reshape(total_patches, 1, 1, -1), mg_rotary.reshape(total_patches, 1, 1, -1)), dim=-1 + ) + mg_cos = mg_emb.cos().flatten() + + cos_sim_rope = F.cosine_similarity( + hf_rotary.float().flatten().unsqueeze(0), mg_rotary.float().flatten().unsqueeze(0) + ).item() + max_diff_rope = (hf_rotary.float() - mg_rotary.float()).abs().max().item() + print(f" HF rot_pos_emb shape: {hf_rotary.shape}") + print(f" MG rot_pos_emb shape: {mg_rotary.shape}") + print(f" RoPE freqs cosine_sim: {cos_sim_rope:.8f}, max_diff: {max_diff_rope:.6e}") + + cos_sim_cos = F.cosine_similarity( + hf_cos.float().flatten().unsqueeze(0), mg_cos.float().flatten().unsqueeze(0) + ).item() + print(f" cos(emb) cosine_sim: {cos_sim_cos:.8f}") + + # ===================================================================== + # Test 3: RoPE application on dummy Q, K + # ===================================================================== + print("\n" + "=" * 60) + print("TEST 3: RoPE application") + print("=" * 60) + head_dim = 80 + num_heads = 16 + with torch.no_grad(): + # Create dummy Q, K + q_test = torch.randn(total_patches, num_heads, head_dim, dtype=torch.float32, device=device) + + # HF RoPE application + from transformers.models.ernie4_5_vl_moe.modeling_ernie4_5_vl_moe import ( + apply_rotary_pos_emb_vision, + ) + + hf_q_rot, _ = apply_rotary_pos_emb_vision(q_test, q_test, hf_cos, hf_sin) + + # MG RoPE application + from megatron.bridge.models.ernie_vl.vision_attention import apply_rotary_pos_emb_absolute + + mg_freqs = torch.cat((mg_rotary, mg_rotary), dim=-1) # [N, 80] + mg_freqs_4d = mg_freqs.reshape(total_patches, 1, 1, -1) # [N, 1, 1, 80] + # MG path: bshd format + mg_q_rot = apply_rotary_pos_emb_absolute( + q_test[:, None], # [N, 1, 16, 80] + mg_freqs_4d, + config=vit_config, + cu_seqlens=None, + ).squeeze(1) # [N, 16, 80] + + cos_sim_qrot = F.cosine_similarity( + hf_q_rot.float().flatten().unsqueeze(0), mg_q_rot.float().flatten().unsqueeze(0) + ).item() + max_diff_qrot = (hf_q_rot.float() - mg_q_rot.float()).abs().max().item() + print(f" RoPE-applied Q cosine_sim: {cos_sim_qrot:.8f}, max_diff: {max_diff_qrot:.6e}") + + # ===================================================================== + # Test 4: Single block comparison + # ===================================================================== + print("\n" + "=" * 60) + print("TEST 4: First transformer block") + print("=" * 60) + with torch.no_grad(): + # HF: first block + hf_hidden = hf_patch_out.clone() + hf_position_embeddings = (hf_cos, hf_sin) + # HF cu_seqlens + cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( + dim=0, dtype=torch.int32 + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + + hf_block0_out = hf_vit.blocks[0]( + hf_hidden, + cu_seqlens=cu_seqlens, + position_embeddings=hf_position_embeddings, + ) + + # MG: first block - need to trace through the TransformerBlock + mg_hidden = mg_patch_out.clone() + mg_hidden_3d = mg_hidden[:, None] # [N, 1, hidden] + + mg_packed_seq = mg_vit.build_packed_seq_params(grid_thw) + mg_rotary_for_block = torch.cat( + (mg_rotary.reshape(total_patches, 1, 1, -1), mg_rotary.reshape(total_patches, 1, 1, -1)), + dim=-1, + ) + + # Get first layer from MG decoder + mg_layer0 = mg_vit.decoder.layers[0] + + # Run MG first block + mg_block0_out, _ = mg_layer0( + hidden_states=mg_hidden_3d, + attention_mask=None, + rotary_pos_emb=mg_rotary_for_block, + packed_seq_params=mg_packed_seq, + ) + + cos_sim_block0 = F.cosine_similarity( + hf_block0_out.float().flatten().unsqueeze(0), mg_block0_out.squeeze(1).float().flatten().unsqueeze(0) + ).item() + max_diff_block0 = (hf_block0_out.float() - mg_block0_out.squeeze(1).float()).abs().max().item() + print(f" Block 0 cosine_sim: {cos_sim_block0:.8f}, max_diff: {max_diff_block0:.6e}") + + # ===================================================================== + # Test 5: Check the LayerNorm + QKV path specifically + # ===================================================================== + print("\n" + "=" * 60) + print("TEST 5: LayerNorm + QKV (pre-attention)") + print("=" * 60) + with torch.no_grad(): + # HF: norm1 then QKV + hf_normed = hf_vit.blocks[0].norm1(hf_patch_out) + hf_qkv = hf_vit.blocks[0].attn.qkv(hf_normed) + + # MG: TELayerNormColumnParallelLinear does fused LN + Linear + # But we can compare the QKV output by accessing the linear_qkv layer + mg_linear_qkv = mg_layer0.self_attention.linear_qkv + # This is TELayerNormColumnParallelLinear - it fuses LN + Linear + mg_qkv_out, _ = mg_linear_qkv(mg_patch_out[:, None]) # [N, 1, 3*hidden] + mg_qkv_flat = mg_qkv_out.squeeze(1) + + cos_sim_qkv = F.cosine_similarity( + hf_qkv.float().flatten().unsqueeze(0), mg_qkv_flat.float().flatten().unsqueeze(0) + ).item() + max_diff_qkv = (hf_qkv.float() - mg_qkv_flat.float()).abs().max().item() + print(f" QKV output cosine_sim: {cos_sim_qkv:.8f}, max_diff: {max_diff_qkv:.6e}") + + # ===================================================================== + # Test 6: Full model + # ===================================================================== + print("\n" + "=" * 60) + print("TEST 6: Full model comparison") + print("=" * 60) + with torch.no_grad(): + hf_out = hf_vit(pixel_values, grid_thw, return_dict=True).last_hidden_state + mg_out = mg_vit(pixel_values, grid_thw) + + cos_sim_full = F.cosine_similarity( + hf_out.float().flatten().unsqueeze(0), mg_out.float().flatten().unsqueeze(0) + ).item() + max_diff_full = (hf_out.float() - mg_out.float()).abs().max().item() + print(f" Full model cosine_sim: {cos_sim_full:.8f}, max_diff: {max_diff_full:.6e}") + + if cos_sim_full > 0.99: + print("\nOVERALL: PASS") + else: + print("\nOVERALL: FAIL") + + # Cleanup + parallel_state.destroy_model_parallel() + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() + + sys.exit(0 if cos_sim_full > 0.99 else 1) + + +if __name__ == "__main__": + main() diff --git a/examples/models/vlm/ernie_vl/hf_loss_check.py b/examples/models/vlm/ernie_vl/hf_loss_check.py new file mode 100644 index 0000000000..a86c4f73fa --- /dev/null +++ b/examples/models/vlm/ernie_vl/hf_loss_check.py @@ -0,0 +1,118 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Quick HF-side loss check for ERNIE 4.5 VL MoE. + +Computes next-token prediction loss using the HF model directly (with its native +softmax routing) to compare against the Megatron model (sigmoid routing). + +Usage: + python hf_loss_check.py --hf-model-path ./ERNIE-4.5-VL-28B-A3B-Thinking \ + --prompt "请介绍一下你自己。" +""" + +import argparse +import os + + +os.environ.setdefault("TORCH_COMPILE_DISABLE", "1") + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + + +def main(): + """Run HF model loss check.""" + parser = argparse.ArgumentParser() + parser.add_argument("--hf-model-path", required=True) + parser.add_argument("--prompt", required=True) + args = parser.parse_args() + + print(f"Loading tokenizer from {args.hf_model_path}...") + tokenizer = AutoTokenizer.from_pretrained(args.hf_model_path, trust_remote_code=True) + + print(f"Loading HF model from {args.hf_model_path}...") + model = AutoModelForCausalLM.from_pretrained( + args.hf_model_path, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + device_map="auto", + ) + model.eval() + + # Tokenize + token_ids = tokenizer.encode(args.prompt, add_special_tokens=True) + print(f"Prompt: {args.prompt!r}") + print(f"Tokenized: {len(token_ids)} tokens") + print(f"Token IDs: {token_ids}") + + device = model.device if hasattr(model, "device") else next(model.parameters()).device + input_ids = torch.tensor([token_ids], dtype=torch.long, device=device) + seq_len = input_ids.size(1) + + # 3D M-RoPE position_ids: [batch, seq_len, 3] — for text-only, all 3 dims identical + position_ids = ( + torch.arange(seq_len, dtype=torch.long, device=device) + .unsqueeze(0) # [1, seq_len] + .unsqueeze(-1) # [1, seq_len, 1] + .expand(1, seq_len, 3) # [1, seq_len, 3] + .clone() + ) + + # Next-token prediction labels + labels = input_ids.clone() + labels[:, :-1] = input_ids[:, 1:] + labels[:, -1] = -100 # ignore last position + + with torch.no_grad(): + outputs = model( + input_ids=input_ids, + attention_mask=torch.ones_like(input_ids, dtype=torch.bool, device=device), + position_ids=position_ids, + ) + + # The HF model may not compute loss even when labels are passed, + # so compute it manually from logits. + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + # logits: [1, seq_len, vocab_size] + # Shift: logits[:-1] predicts tokens[1:] + shift_logits = logits[:, :-1, :].contiguous().float() + shift_labels = input_ids[:, 1:].contiguous() + loss = torch.nn.functional.cross_entropy( + shift_logits.view(-1, shift_logits.size(-1)), + shift_labels.view(-1), + reduction="mean", + ) + + print(f"\nHF model loss (next-token prediction): {loss.item():.6f}") + print( + f"ln(vocab_size) = ln({tokenizer.vocab_size}) = {torch.log(torch.tensor(float(tokenizer.vocab_size))).item():.4f}" + ) + + print(f"\nLogits shape: {logits.shape}") + + # Show top-5 predictions for each position + for pos in range(logits.shape[1] - 1): + topk_vals, topk_ids = torch.topk(logits[0, pos], k=5) + actual_next = input_ids[0, pos + 1].item() + decoded_predictions = [tokenizer.decode([tid]) for tid in topk_ids.tolist()] + actual_decoded = tokenizer.decode([actual_next]) + print( + f" Position {pos}: actual_next='{actual_decoded}'({actual_next}), " + f"top5={list(zip(decoded_predictions, topk_ids.tolist(), topk_vals.tolist()))}" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/models/vlm/ernie_vl/hf_to_megatron_generate_ernie_vl.py b/examples/models/vlm/ernie_vl/hf_to_megatron_generate_ernie_vl.py new file mode 100755 index 0000000000..6c57e2f720 --- /dev/null +++ b/examples/models/vlm/ernie_vl/hf_to_megatron_generate_ernie_vl.py @@ -0,0 +1,352 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""ERNIE 4.5 VL: HF-to-Megatron generate (VLM). + +ERNIE 4.5 VL uses a custom processor API that differs from Qwen-style models: + - ``processor.tokenizer.apply_chat_template()`` instead of ``processor.apply_chat_template()`` + - ``processor.process_vision_info()`` for image pre-processing + - Output keys: "images" (pixel_values), "grid_thw" (image_grid_thw) + - mm_token_type_ids must be constructed from image_token_id positions + +This script mirrors ``hf_to_megatron_generate_vlm.py`` but handles ERNIE-specific +processor differences. + +Example: + # Single GPU: + torchrun --nproc_per_node=1 examples/models/vlm/ernie_vl/hf_to_megatron_generate_ernie_vl.py \ + --hf_model_path baidu/ERNIE-4.5-VL-28B-A3B-Instruct \ + --image_path /path/to/image.png \ + --prompt "Describe this image." + + # Multi-GPU (TP=2, EP=4): + torchrun --nproc_per_node=8 examples/models/vlm/ernie_vl/hf_to_megatron_generate_ernie_vl.py \ + --hf_model_path baidu/ERNIE-4.5-VL-28B-A3B-Instruct \ + --image_path /path/to/image.png \ + --prompt "Describe this image." \ + --tp 2 --ep 4 +""" + +import argparse +import json +import os +import sys +import types + + +os.environ.setdefault("TORCH_COMPILE_DISABLE", "1") +os.environ.setdefault("TORCHDYNAMO_DISABLE", "1") + +# Fake 'decord' module -- ERNIE VL processor imports it but only uses it for video. +if "decord" not in sys.modules: + + class _FakeVideoReader: + def __init__(self, *a, **kw): + raise RuntimeError("decord not installed; video processing unavailable") + + _decord_fake = types.ModuleType("decord") + _decord_fake.VideoReader = _FakeVideoReader + _decord_fake.cpu = lambda x=0: x + _bridge = types.ModuleType("decord.bridge") + _bridge.set_bridge = lambda *a, **kw: None + sys.modules["decord"] = _decord_fake + sys.modules["decord.bridge"] = _bridge + +import torch +import torch.distributed as dist +from megatron.core import parallel_state +from megatron.core.pipeline_parallel.schedules import get_forward_backward_func +from transformers import AutoProcessor + +from megatron.bridge import AutoBridge +from megatron.bridge.utils.common_utils import get_last_rank, print_rank_0, print_rank_last + + +# --------------------------------------------------------------------------- +# Forward step +# --------------------------------------------------------------------------- + + +class SingleBatchIterator: + """Iterator that yields a single batch then stops. Required by + ``get_forward_backward_func``.""" + + def __init__( + self, input_ids, position_ids, attention_mask, pixel_values=None, image_grid_thw=None, mm_token_type_ids=None + ): + self.batch = dict( + tokens=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + ) + if pixel_values is not None: + self.batch["pixel_values"] = pixel_values + if image_grid_thw is not None: + self.batch["image_grid_thw"] = image_grid_thw + if mm_token_type_ids is not None: + self.batch["mm_token_type_ids"] = mm_token_type_ids + self._yielded = False + + def __iter__(self): + return self + + def __next__(self): + if self._yielded: + raise StopIteration + self._yielded = True + return self.batch + + +def ernie_vl_forward_step(data_iterator, model, **kwargs) -> torch.Tensor: + """Forward step for ERNIE VL generation.""" + batch = next(data_iterator) + forward_args = { + "input_ids": batch["tokens"], + "mm_token_type_ids": batch.get("mm_token_type_ids"), + "moe_mm_token_type_ids": batch.get("mm_token_type_ids"), + "attention_mask": batch.get("attention_mask"), + } + if "pixel_values" in batch: + forward_args["pixel_values"] = batch["pixel_values"] + if "image_grid_thw" in batch: + forward_args["image_grid_thw"] = batch["image_grid_thw"] + + def loss_func(x, **kwargs): + return x + + model_output = model(**forward_args) + if isinstance(model_output, tuple): + output_tensor, _ = model_output + else: + output_tensor = model_output + return output_tensor, loss_func + + +# --------------------------------------------------------------------------- +# Input processing +# --------------------------------------------------------------------------- + + +def process_ernie_vl_inputs(processor, hf_model_path, image_path, prompt): + """Process inputs using ERNIE 4.5 VL processor API. + + Returns: + (input_ids, pixel_values, image_grid_thw, mm_token_type_ids) + All tensors are on CPU; caller moves to GPU. + """ + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": image_path}}, + {"type": "text", "text": prompt}, + ], + } + ] + text = processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + image_inputs, video_inputs = processor.process_vision_info(messages) + inputs = processor( + text=[text], + images=image_inputs, + videos=video_inputs, + padding=True, + return_tensors="pt", + ) + + input_ids = inputs["input_ids"] + pixel_values = inputs.get("images") # ERNIE uses "images" key + image_grid_thw = inputs.get("grid_thw") # ERNIE uses "grid_thw" key + + # Build mm_token_type_ids from image_token_id positions. + # The HF processor's token_type_ids marks IMAGE_START/END as type 1, + # but Megatron expects only actual image placeholder tokens to be type 1. + with open(os.path.join(hf_model_path, "config.json")) as f: + hf_cfg = json.load(f) + image_token_id = hf_cfg.get("image_token_id", 100295) + + mm_token_type_ids = torch.zeros(1, input_ids.shape[1], dtype=torch.int32) + image_placeholder_mask = input_ids == image_token_id + mm_token_type_ids[image_placeholder_mask] = 1 + + return input_ids, pixel_values, image_grid_thw, mm_token_type_ids + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def to_cuda(x): + """Move tensor to CUDA if not None.""" + if x is None: + return None + return x.cuda() + + +def main(args): + """Main generation function.""" + print_rank_0("=" * 60) + print_rank_0("ERNIE 4.5 VL -- HF to Megatron Generate") + print_rank_0("=" * 60) + + trust_remote = args.trust_remote_code + tp, pp, ep = args.tp, args.pp, args.ep + + # ------------------------------------------------------------------ + # Load model via AutoBridge + # ------------------------------------------------------------------ + print_rank_0(f"Loading model: {args.hf_model_path}") + bridge = AutoBridge.from_hf_pretrained( + args.hf_model_path, torch_dtype=torch.bfloat16, trust_remote_code=trust_remote + ) + model_provider = bridge.to_megatron_provider(load_weights=True) + model_provider.tensor_model_parallel_size = tp + model_provider.pipeline_model_parallel_size = pp + model_provider.expert_model_parallel_size = ep + model_provider.pipeline_dtype = torch.bfloat16 + model_provider.params_dtype = torch.bfloat16 + model_provider.finalize() + model_provider.initialize_model_parallel(seed=0) + model = model_provider.provide_distributed_model(wrap_with_ddp=False) + + model = [m.cuda() for m in model] + for m in model: + m.eval() + if hasattr(m, "config"): + m.config.grad_scale_func = None + m.config.deallocate_pipeline_outputs = False + + # ------------------------------------------------------------------ + # Processor + # ------------------------------------------------------------------ + processor = AutoProcessor.from_pretrained(args.hf_model_path, trust_remote_code=trust_remote) + eos_token_id = processor.tokenizer.eos_token_id + + # ------------------------------------------------------------------ + # Process inputs + # ------------------------------------------------------------------ + input_ids_raw, pixel_values, image_grid_thw, mm_token_type_ids = process_ernie_vl_inputs( + processor, args.hf_model_path, args.image_path, args.prompt + ) + + input_ids_raw = input_ids_raw.cuda() + pixel_values = to_cuda(pixel_values) + image_grid_thw = to_cuda(image_grid_thw) + mm_token_type_ids = to_cuda(mm_token_type_ids) + + print_rank_0(f"Input tokens: {input_ids_raw.shape[1]}, Image tokens: {(mm_token_type_ids == 1).sum().item()}") + + # ------------------------------------------------------------------ + # Greedy generation loop + # ------------------------------------------------------------------ + generated_ids = input_ids_raw.clone() + + for step in range(args.max_new_tokens): + with torch.no_grad(): + print_rank_0(f"Generation step {step}") + + real_seq_len = generated_ids.size(1) + input_ids = generated_ids + + position_ids = ( + torch.arange(input_ids.size(1), dtype=torch.long, device=input_ids.device) + .unsqueeze(0) + .expand_as(input_ids) + ) + + fwd_bwd_function = get_forward_backward_func() + iterator = SingleBatchIterator( + input_ids, + position_ids, + None, + pixel_values, + image_grid_thw, + mm_token_type_ids, + ) + + output = fwd_bwd_function( + forward_step_func=ernie_vl_forward_step, + data_iterator=iterator, + model=model, + num_microbatches=1, + forward_only=True, + seq_length=input_ids.size(1), + micro_batch_size=1, + collect_non_loss_data=True, + ) + if isinstance(output, list) and len(output) > 0: + output = output[0] + + if parallel_state.is_pipeline_last_stage(): + world_size = parallel_state.get_tensor_model_parallel_world_size() + gathered_tensors = [torch.zeros_like(output) for _ in range(world_size)] + dist.all_gather( + gathered_tensors, + output, + group=parallel_state.get_tensor_model_parallel_group(), + ) + output = torch.cat(gathered_tensors, dim=2) + + last_pos = real_seq_len - 1 + next_token_ids = torch.argmax(output[:, last_pos], dim=-1, keepdim=True) + + if step < 5: + logits = output[0, last_pos, :] + top5_vals, top5_ids = torch.topk(logits, 5) + top5_tokens = [processor.tokenizer.decode([idx]) for idx in top5_ids] + print_rank_last(f"Top 5: {list(zip(top5_tokens, top5_vals.tolist()))}") + print_rank_last( + f"Selected: '{processor.tokenizer.decode([next_token_ids.item()])}' " + f"(id={next_token_ids.item()})" + ) + else: + next_token_ids = torch.ones((1, 1), device=generated_ids.device, dtype=generated_ids.dtype) + + torch.distributed.broadcast(next_token_ids, get_last_rank()) + generated_ids = torch.cat([generated_ids, next_token_ids], dim=-1) + + if mm_token_type_ids is not None: + mm_token_type_ids = torch.cat( + [mm_token_type_ids, torch.zeros_like(next_token_ids, dtype=mm_token_type_ids.dtype)], + dim=-1, + ) + + if next_token_ids.item() == eos_token_id: + break + + generated_text = processor.tokenizer.decode(list(generated_ids[0, input_ids_raw.shape[1] :])) + print_rank_0("======== GENERATED TEXT OUTPUT ========") + if args.image_path: + print_rank_0(f"Image: {args.image_path}") + print_rank_0(f"Prompt: {args.prompt}") + print_rank_0(f"Generated: {generated_text}") + print_rank_0("=======================================") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description="ERNIE 4.5 VL: HF to Megatron Generate") + parser.add_argument("--hf_model_path", type=str, required=True, help="Path to the HuggingFace ERNIE 4.5 VL model.") + parser.add_argument("--prompt", type=str, default="Describe this image.", help="Input prompt.") + parser.add_argument("--max_new_tokens", type=int, default=50, help="Maximum number of new tokens to generate.") + parser.add_argument("--tp", type=int, default=1, help="Tensor parallelism size") + parser.add_argument("--pp", type=int, default=1, help="Pipeline parallelism size") + parser.add_argument("--ep", type=int, default=1, help="Expert parallelism size") + parser.add_argument("--image_path", type=str, default=None, help="Path or URL to image (optional).") + parser.add_argument("--trust_remote_code", action="store_true", help="Trust remote code for HF model loading") + args = parser.parse_args() + + main(args) + + if torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() diff --git a/examples/models/vlm/ernie_vl/inference.sh b/examples/models/vlm/ernie_vl/inference.sh new file mode 100755 index 0000000000..fc3f65b1bd --- /dev/null +++ b/examples/models/vlm/ernie_vl/inference.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -e + +# Supported model variants: +# ERNIE-4.5-VL-28B-A3B-Instruct, ERNIE-4.5-VL-28B-A3B-Thinking +MODEL_NAME=ERNIE-4.5-VL-28B-A3B-Instruct + +EP=4 +TP=2 +PP=1 + +# ERNIE 4.5 VL uses a custom processor API that differs from Qwen-style models, +# so we use a dedicated generate script instead of the generic hf_to_megatron_generate_vlm.py. + +# Inference with Hugging Face checkpoints +uv run python -m torch.distributed.run --nproc_per_node=8 \ + examples/models/vlm/ernie_vl/hf_to_megatron_generate_ernie_vl.py \ + --hf_model_path baidu/${MODEL_NAME} \ + --image_path "https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16/resolve/main/images/table.png" \ + --prompt "Describe this image." \ + --max_new_tokens 50 \ + --tp ${TP} --pp ${PP} --ep ${EP} \ + --trust_remote_code diff --git a/src/megatron/bridge/models/__init__.py b/src/megatron/bridge/models/__init__.py index 9791e1b225..4e68e06c26 100644 --- a/src/megatron/bridge/models/__init__.py +++ b/src/megatron/bridge/models/__init__.py @@ -32,6 +32,14 @@ DeepSeekV2Bridge, DeepSeekV3Bridge, ) +from megatron.bridge.models.ernie import ( + Ernie45Bridge, +) +from megatron.bridge.models.ernie_vl import ( + Ernie45VLBridge, + Ernie45VLModel, + Ernie45VLModelProvider, +) from megatron.bridge.models.falcon_h1 import ( FalconH1Bridge, FalconH1ModelProvider, @@ -170,6 +178,12 @@ # DeepSeek Models "DeepSeekV2Bridge", "DeepSeekV3Bridge", + # ERNIE Text-Only Models + "Ernie45Bridge", + # ERNIE VL Models + "Ernie45VLBridge", + "Ernie45VLModel", + "Ernie45VLModelProvider", "FalconH1Bridge", "FalconH1ModelProvider", "Gemma3ModelProvider", diff --git a/src/megatron/bridge/models/conversion/model_bridge.py b/src/megatron/bridge/models/conversion/model_bridge.py index 1e4256b330..54a4965cda 100644 --- a/src/megatron/bridge/models/conversion/model_bridge.py +++ b/src/megatron/bridge/models/conversion/model_bridge.py @@ -315,6 +315,28 @@ def _update_grouped_expert_number(param_name: str, param_type: str) -> str: param_name = _update_grouped_expert_number(param_name, "weight") elif re.search(r"\.bias\d+(?=$|\.)", param_name): param_name = _update_grouped_expert_number(param_name, "bias") + + # EP for SequentialMLP: expert index is in the module path as local_experts.N. + # This covers both standard SequentialMLP (e.g., quantization) and dual-pool MoE + # (e.g., text_moe_layer.experts.local_experts.N or vision_moe_layer.experts.local_experts.N). + elif ( + ".experts.local_experts." in param_name + and ep_group is not None + and get_pg_size(ep_group) > 1 + and ".adapter." not in param_name + ): + num_experts = config.num_moe_experts + num_experts_per_rank = num_experts // ep_group.size() + + match = re.search(r"\.local_experts\.(\d+)\.", param_name) + if match: + local_expert_number = int(match.group(1)) + global_expert_number = num_experts_per_rank * ep_group.rank() + local_expert_number + param_name = param_name.replace( + f".local_experts.{local_expert_number}.", + f".local_experts.{global_expert_number}.", + ) + return param_name diff --git a/src/megatron/bridge/models/conversion/param_mapping.py b/src/megatron/bridge/models/conversion/param_mapping.py index 0aca8bd5ee..4dcb5c56aa 100644 --- a/src/megatron/bridge/models/conversion/param_mapping.py +++ b/src/megatron/bridge/models/conversion/param_mapping.py @@ -201,10 +201,12 @@ def etp_size(self) -> int: def is_expert(self) -> bool: """Check if this mapping is for an expert parameter. - Matches both TEGroupedMLP (.mlp.experts.linear_fc) and - SequentialMLP (.mlp.experts.local_experts.*.linear_fc) patterns. + Matches both TEGroupedMLP (.experts.linear_fc) and + SequentialMLP (.experts.local_experts.*.linear_fc) patterns. + Uses ``.experts.`` rather than ``.mlp.experts.`` so models with an + intermediate sub-module (e.g. ``.mlp..experts.``) are matched too. """ - return ".mlp.experts.linear_fc" in self.megatron_param or ".mlp.experts.local_experts." in self.megatron_param + return ".experts.linear_fc" in self.megatron_param or ".experts.local_experts." in self.megatron_param @property def is_adapter(self) -> bool: @@ -386,7 +388,11 @@ def broadcast_from_pp_rank( if target_tensor_spec is None: # No rank had the tensor – this is an error in the caller. - raise ValueError("Object must exist on at least one PP rank") + raise ValueError( + f"Object must exist on at least one PP rank. " + f"megatron_param={self.megatron_param}, hf_param={self.hf_param}, " + f"cache_key={cache_key}" + ) # ------------------------------------------------------------------ # 3. Ensure every rank has an allocated tensor with the right shape @@ -752,6 +758,10 @@ def gather_from_ep_ranks( Dict[str, torch.Tensor]: Mapping from HF parameter names (one per EP rank) to the corresponding expert tensors gathered from each EP rank. """ + # Fast path for EP=1: no gathering needed, just return with the given name. + if self.ep_size == 1: + return {str(hf_param_name): megatron_weights} + if megatron_module is None: num_experts_per_rank = self.broadcast_obj_from_pp_rank(None, "num_experts_per_rank") else: @@ -763,8 +773,8 @@ def gather_from_ep_ranks( global_expert_number = extract_expert_number_from_param(self.megatron_param) local_expert_number = global_expert_number % num_experts_per_rank - # Compute global expert numbers for all EP ranks. HF MoE params use - # experts.N naming here; local_experts.N would require a separate pattern. + # Compute global expert numbers for all EP ranks + # use regex to replace the local expert number with the global expert number gathered_expert_param_names = [ re.sub( r"experts\.(\d+)", f"experts.{int(local_expert_number) + num_experts_per_rank * i}", str(hf_param_name) @@ -1344,8 +1354,11 @@ def hf_to_megatron( megatron_module: nn.Module, ) -> torch.Tensor: """Delegate to appropriate mapping based on module type.""" - # Apply permutation if specified (before distribution) - if self.permute_dims is not None and self.tp_rank == 0: + # Apply permutation if specified (before distribution). + # Must be applied on ALL ranks (not just tp_rank 0) because some delegate + # mappings (e.g. ReplicatedMapping) expect hf_weights to have the correct + # shape on every rank. + if self.permute_dims is not None: hf_weights = torch.permute(hf_weights, self.permute_dims).contiguous() # Detect type and create delegate on first use @@ -1376,6 +1389,12 @@ def megatron_to_hf( # PP group likely has 1 member - skipping. return {} + # If no PP rank detected a type (e.g. Megatron parameter without an + # HF counterpart, such as MoE modules on dense layers created by + # moe_layer_freq), skip export gracefully. + if self._detected_type is None: + return {} + self._mapping = self._get_or_create_mapping(self._detected_type) result = self._mapping.megatron_to_hf(megatron_weights, megatron_module) diff --git a/src/megatron/bridge/models/ernie/__init__.py b/src/megatron/bridge/models/ernie/__init__.py new file mode 100644 index 0000000000..70ac0a7635 --- /dev/null +++ b/src/megatron/bridge/models/ernie/__init__.py @@ -0,0 +1,20 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from megatron.bridge.models.ernie.ernie_45_bridge import Ernie45Bridge + + +__all__ = [ + "Ernie45Bridge", +] diff --git a/src/megatron/bridge/models/ernie/ernie_45_bridge.py b/src/megatron/bridge/models/ernie/ernie_45_bridge.py new file mode 100644 index 0000000000..2ca9f1f64f --- /dev/null +++ b/src/megatron/bridge/models/ernie/ernie_45_bridge.py @@ -0,0 +1,368 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Megatron Bridge for ERNIE 4.5 text-only MoE model. + +Maps HuggingFace Ernie4_5_MoeForCausalLM weights and config to +Megatron-Core GPTModel with single-pool MoE (64 experts, top-6 routing, +shared experts, expert bias for aux-free load balancing). +""" + +import torch.nn.functional as F +from megatron.core.models.gpt.gpt_model import GPTModel + +from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.conversion.param_mapping import ( + AutoMapping, + GatedMLPMapping, + QKVMapping, + ReplicatedMapping, +) +from megatron.bridge.models.gpt_provider import GPTModelProvider + + +def _ernie45_decoder_block_spec(config: "GPTModelProvider", vp_stage: int | None = None): + """Create a decoder block spec that respects ``moe_layer_freq``. + + The default ``GPTModelProvider.transformer_layer_spec`` calls + ``get_gpt_layer_with_transformer_engine_spec`` which returns a single + MoE layer spec applied uniformly to ALL layers, ignoring + ``moe_layer_freq``. + + ERNIE 4.5 has mixed dense/MoE layers (layer 0 is dense, layers 1-N + are MoE). This function uses ``get_gpt_decoder_block_spec`` which + calls ``get_gpt_decoder_layer_specs`` — the code path that parses + ``config.moe_layer_freq`` and creates per-layer specs (dense for + pattern=0, MoE for pattern=1). + """ + from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec + + return get_gpt_decoder_block_spec( + config=config, + use_transformer_engine=True, + vp_stage=vp_stage, + ) + + +# HF class name string; avoids requiring the HF modeling module at import time. +_ERNIE45_MOE_HF_CLASS_NAME = "Ernie4_5_MoeForCausalLM" + + +class _PPSafeMixin: + """Mixin that makes ``megatron_to_hf`` safe for PP export of MoE-only params. + + When ``moe_layer_freq`` makes some layers dense and others MoE, + MoE-only parameters (router weight, expert bias, shared/routed expert + weights) do not exist on dense layers. With PP > 1, + ``broadcast_from_pp_rank`` raises ``ValueError`` because no PP rank + owns the tensor. + + This mixin catches that error and returns ``{}`` so the conversion + loop simply omits the parameter from the output. + + **Must be listed before the base mapping class in the MRO** so that + ``super().megatron_to_hf`` resolves to the concrete mapping's method. + """ + + def megatron_to_hf(self, megatron_weights, megatron_module): + try: + return super().megatron_to_hf(megatron_weights, megatron_module) + except ValueError: + # Parameter doesn't exist on any PP rank (dense layer). + return {} + + +class _PPSafeAutoMapping(_PPSafeMixin, AutoMapping): + """AutoMapping that skips export for missing parameters.""" + + pass + + +class _PPSafeReplicatedMapping(_PPSafeMixin, ReplicatedMapping): + """ReplicatedMapping that skips export for missing parameters.""" + + pass + + +class _PPSafeGatedMLPMapping(_PPSafeMixin, GatedMLPMapping): + """GatedMLPMapping that skips export for missing parameters.""" + + pass + + +class _SqueezeBiasMapping(_PPSafeReplicatedMapping): + """Mapping for the single-pool expert bias tensor. + + The HF text-only model stores ``moe_statics.e_score_correction_bias`` + with shape ``[1, num_experts]`` (1 expert group for text-only). + Megatron stores ``router.expert_bias`` as a flat ``[num_experts]`` tensor. + + This mapping squeezes dim-0 on import and unsqueezes on export. + + Inherits from ``_PPSafeReplicatedMapping`` to gracefully skip dense + layers during PP export. + """ + + def hf_to_megatron(self, hf_weights, megatron_module): + # [1, N] -> [N] + if hf_weights.ndim == 2 and hf_weights.shape[0] == 1: + hf_weights = hf_weights.squeeze(0) + return super().hf_to_megatron(hf_weights, megatron_module) + + def megatron_to_hf(self, megatron_weights, megatron_module): + result = super().megatron_to_hf(megatron_weights, megatron_module) + if result: + # [N] -> [1, N] + out = {} + for k, v in result.items(): + out[k] = v.unsqueeze(0) if v.ndim == 1 else v + return out + return result + + +@MegatronModelBridge.register_bridge( + source=_ERNIE45_MOE_HF_CLASS_NAME, + target=GPTModel, + provider=GPTModelProvider, + model_type="ernie4_5_moe", +) +class Ernie45Bridge(MegatronModelBridge): + """ + Megatron Bridge for ERNIE 4.5 text-only MoE Causal LM. + + This bridge handles the conversion between HuggingFace Ernie4_5_MoeForCausalLM + and Megatron-Core GPTModel formats with single-pool MoE architecture. + + Key architectural features: + - Single-pool MoE: 64 experts, top-6 routing, shared experts + - Softmax routing with expert bias for aux-free load balancing + - Interleaved RoPE (base=500000) + - GQA with 20 query heads, 4 KV heads, kv_channels=128 + - RMSNorm, SiLU-gated MLP + - Router gate weight stored as [H, E] in HF (transposed for Megatron [E, H]) + + Example: + >>> from megatron.bridge import AutoBridge + >>> bridge = AutoBridge.from_hf_pretrained("baidu/ERNIE-4.5-0.3B-PT") + >>> provider = bridge.to_megatron_provider() + """ + + @staticmethod + def _get_num_experts(hf_config) -> int: + """Extract num_experts as an int. + + The config may store moe_num_experts as a plain int or as a list + ``[N]`` (single pool) or ``[N, N]`` (dual pool -- take first). + """ + raw = getattr(hf_config, "moe_num_experts", 64) + if isinstance(raw, (list, tuple)): + return raw[0] + return int(raw) + + def provider_bridge(self, hf_pretrained): + """Convert HuggingFace ERNIE 4.5 MoE config to GPTModelProvider. + + Uses super().provider_bridge() for standard CONFIG_MAPPING fields + (hidden_size, num_layers, rope_theta, tie_word_embeddings, etc.) + and then overrides ERNIE-specific settings. + """ + provider = super().provider_bridge(hf_pretrained) + hf_config = hf_pretrained.config + + # --- Architecture overrides --- + provider.normalization = "RMSNorm" + provider.activation_func = F.silu + provider.gated_linear_unit = True + provider.add_bias_linear = False + provider.add_qkv_bias = False + provider.hidden_dropout = 0.0 + provider.position_embedding_type = "rope" + provider.rotary_base = 500000.0 + provider.rotary_interleaved = True + provider.moe_router_load_balancing_type = "aux_loss" + # Mixed dense/MoE layers (layer 0 dense, rest MoE): use decoder + # block spec that parses moe_layer_freq per-layer instead of the + # default spec which applies MoE uniformly to all layers. + provider.transformer_layer_spec = _ernie45_decoder_block_spec + + # --- MoE settings (ERNIE uses non-standard HF config field names) --- + num_experts = self._get_num_experts(hf_config) + provider.num_moe_experts = num_experts + provider.moe_router_topk = getattr(hf_config, "moe_k", 6) + + # Expert intermediate size: may be int or list (text-only uses first/only). + moe_intermediate = getattr(hf_config, "moe_intermediate_size", None) + if isinstance(moe_intermediate, (list, tuple)): + provider.moe_ffn_hidden_size = moe_intermediate[0] + elif moe_intermediate is not None: + provider.moe_ffn_hidden_size = moe_intermediate + else: + provider.moe_ffn_hidden_size = getattr(hf_config, "intermediate_size", 5120) + + # Shared experts + moe_num_shared = getattr(hf_config, "moe_num_shared_experts", 2) + provider.moe_shared_expert_intermediate_size = provider.moe_ffn_hidden_size * moe_num_shared + + # Router settings + provider.moe_aux_loss_coeff = getattr(hf_config, "router_aux_loss_coef", 0.001) + + # MoE runtime settings — same as DeepSeek V3 (sigmoid routing + expert bias) + provider.moe_grouped_gemm = True + provider.moe_router_pre_softmax = False + provider.moe_router_score_function = "sigmoid" + provider.moe_router_enable_expert_bias = True + provider.moe_router_dtype = "fp32" + provider.moe_token_dispatcher_type = "alltoall" + provider.moe_permute_fusion = True + # gradient_accumulation_fusion: use the auto-detected default from + # GPTModelProvider (checks for APEX or TE availability) rather than + # overriding it here. For conversion jobs (no backward pass) the + # flag is irrelevant; for training it will be enabled whenever + # the required extensions are present. + + # Disable MTP (Multi-Token Prediction) for inference -- the ERNIE HF + # model stores num_nextn_predict_layers in config but does not ship + # MTP weights, so we must not create MTP layers in Megatron. + provider.mtp_num_layers = None + + # Determine which layers are dense vs MoE. + mlp_layer_types = getattr(hf_config, "mlp_layer_types", None) + if mlp_layer_types is not None: + provider.moe_layer_freq = [0 if t == "dense" else 1 for t in mlp_layer_types] + else: + num_layers = hf_config.num_hidden_layers + moe_start = getattr(hf_config, "moe_layer_start_index", None) + if moe_start is not None: + start = moe_start[0] if isinstance(moe_start, (list, tuple)) else moe_start + provider.moe_layer_freq = [0] * start + [1] * (num_layers - start) + else: + # Default: layer 0 dense, rest MoE + provider.moe_layer_freq = [0] + [1] * (num_layers - 1) + + return provider + + def mapping_registry(self) -> MegatronMappingRegistry: + """Return MegatronMappingRegistry with parameter mappings for ERNIE 4.5 MoE.""" + # Simple 1:1 parameter mappings + param_mappings = { + # Embeddings & output + "embedding.word_embeddings.weight": "model.embed_tokens.weight", + "output_layer.weight": "lm_head.weight", + "decoder.final_layernorm.weight": "model.norm.weight", + # Attention + "decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": ("model.layers.*.input_layernorm.weight"), + "decoder.layers.*.self_attention.linear_proj.weight": ("model.layers.*.self_attn.o_proj.weight"), + # Dense MLP layernorm (layer 0 -- fused into linear_fc1) + "decoder.layers.*.mlp.linear_fc1.layer_norm_weight": ("model.layers.*.post_attention_layernorm.weight"), + # Dense MLP down_proj (layer 0) + "decoder.layers.*.mlp.linear_fc2.weight": ("model.layers.*.mlp.down_proj.weight"), + # MoE layers: pre_mlp_layernorm + "decoder.layers.*.pre_mlp_layernorm.weight": ("model.layers.*.post_attention_layernorm.weight"), + } + + mapping_list = [] + for megatron_param, hf_param in param_mappings.items(): + mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) + + # MoE-only parameters use PP-safe variants that gracefully return {} + # when the parameter doesn't exist on any PP rank (dense layers + # created by moe_layer_freq). + + # Shared experts: down_proj (MoE-only) + mapping_list.append( + _PPSafeAutoMapping( + megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc2.weight", + hf_param="model.layers.*.mlp.shared_experts.down_proj.weight", + ) + ) + + mapping_list.extend( + [ + # ============================================================= + # QKV: Combine separate Q, K, V into fused QKV + # ============================================================= + QKVMapping( + megatron_param="decoder.layers.*.self_attention.linear_qkv.weight", + q="model.layers.*.self_attn.q_proj.weight", + k="model.layers.*.self_attn.k_proj.weight", + v="model.layers.*.self_attn.v_proj.weight", + ), + # ============================================================= + # Dense MLP (layer 0): gate_proj + up_proj -> fused linear_fc1 + # ============================================================= + GatedMLPMapping( + megatron_param="decoder.layers.*.mlp.linear_fc1.weight", + gate="model.layers.*.mlp.gate_proj.weight", + up="model.layers.*.mlp.up_proj.weight", + ), + # ============================================================= + # Router weight: HF text-only model stores [E, H] (same as + # Megatron), so no transpose needed. Note: the VL model + # stores the gate weight transposed as [H, E] and needs + # permute_dims=(1, 0); the text-only model does not. + # + # Uses ``_PPSafeReplicatedMapping`` because TopKRouter.weight + # is replicated across TP ranks, and dense layers (created + # by ``moe_layer_freq``) have no router — the PP-safe + # variant gracefully returns ``{}`` for those layers. + # ============================================================= + _PPSafeReplicatedMapping( + megatron_param="decoder.layers.*.mlp.router.weight", + hf_param="model.layers.*.mlp.gate.weight", + ), + # ============================================================= + # MoE expert mappings for TEGroupedMLP (fused 3D tensors) + # ============================================================= + _PPSafeGatedMLPMapping( + megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*", + gate="model.layers.*.mlp.experts.*.gate_proj.weight", + up="model.layers.*.mlp.experts.*.up_proj.weight", + ), + _PPSafeAutoMapping( + megatron_param="decoder.layers.*.mlp.experts.linear_fc2.weight*", + hf_param="model.layers.*.mlp.experts.*.down_proj.weight", + ), + # ============================================================= + # MoE expert mappings for SequentialMLP (per-expert, for quantization) + # ============================================================= + _PPSafeGatedMLPMapping( + megatron_param="decoder.layers.*.mlp.experts.local_experts.*.linear_fc1.weight", + gate="model.layers.*.mlp.experts.*.gate_proj.weight", + up="model.layers.*.mlp.experts.*.up_proj.weight", + ), + _PPSafeAutoMapping( + megatron_param="decoder.layers.*.mlp.experts.local_experts.*.linear_fc2.weight", + hf_param="model.layers.*.mlp.experts.*.down_proj.weight", + ), + # ============================================================= + # Shared experts: gate+up -> fused linear_fc1 + # ============================================================= + _PPSafeGatedMLPMapping( + megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc1.weight", + gate="model.layers.*.mlp.shared_experts.gate_proj.weight", + up="model.layers.*.mlp.shared_experts.up_proj.weight", + ), + # ============================================================= + # Expert bias: [1, N] on disk -> [N] in Megatron + # ============================================================= + _SqueezeBiasMapping( + megatron_param="decoder.layers.*.mlp.router.expert_bias", + hf_param="model.layers.*.mlp.moe_statics.e_score_correction_bias", + ), + ] + ) + + return MegatronMappingRegistry(*mapping_list) diff --git a/src/megatron/bridge/models/ernie_vl/__init__.py b/src/megatron/bridge/models/ernie_vl/__init__.py new file mode 100644 index 0000000000..9bdbd49170 --- /dev/null +++ b/src/megatron/bridge/models/ernie_vl/__init__.py @@ -0,0 +1,34 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from megatron.bridge.models.ernie_vl.ernie45_vl_bridge import Ernie45VLBridge +from megatron.bridge.models.ernie_vl.ernie45_vl_provider import Ernie45VLModelProvider +from megatron.bridge.models.ernie_vl.modeling_ernie45_vl import Ernie45VLModel +from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.ernie_decoder_layer_spec import ( + get_ernie45_vl_decoder_block_spec, +) +from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.ernie_moe_layer import ( + ErnieMultiTypeMoE, + MultiTypeMoeSubmodules, +) + + +__all__ = [ + "Ernie45VLBridge", + "Ernie45VLModel", + "Ernie45VLModelProvider", + "ErnieMultiTypeMoE", + "MultiTypeMoeSubmodules", + "get_ernie45_vl_decoder_block_spec", +] diff --git a/src/megatron/bridge/models/ernie_vl/ernie45_vl_bridge.py b/src/megatron/bridge/models/ernie_vl/ernie45_vl_bridge.py new file mode 100644 index 0000000000..cf670355f2 --- /dev/null +++ b/src/megatron/bridge/models/ernie_vl/ernie45_vl_bridge.py @@ -0,0 +1,908 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Megatron Bridge for ERNIE 4.5 VL MoE (Vision-Language with Mixture of Experts). + +This bridge handles conversion between HuggingFace Ernie4_5_VLMoeForConditionalGeneration +and Megatron-Core Ernie45VLModel formats, including: + +- Language model weights with heterogeneous dual-pool MoE: + * text_moe: 64 experts with intermediate_size=1536 (text tokens) + -> mapped to ErnieMultiTypeMoE.text_moe_layer (standard MoELayer) + * vision_moe: 64 experts with intermediate_size=512 (vision tokens) + -> mapped to ErnieMultiTypeMoE.vision_moe_layer (standard MoELayer) + * shared_experts: 2 shared experts with intermediate_size=3072 + -> mapped to ErnieMultiTypeMoE.shared_experts +- Vision encoder weights: + * HF ViT (use_mg_vit=False): replicated across TP ranks via ReplicatedMapping + * MG ViT (use_mg_vit=True): TP-sharded with ConcatenatedQKVMapping for fused QKV +- Resampler / projector weights (replicated across TP ranks) +- 3D M-RoPE position embedding configuration + +The ErnieMultiTypeMoE module contains two separate MoELayer instances (one per +expert pool), each with its own router and SequentialMLP experts. This gives both +pools full TP support through standard Megatron-Core infrastructure. + +HF on-disk (safetensors) keys -- after ``_checkpoint_conversion_mapping`` reversal: + model.layers.{i}.mlp.gate.weight (text router) + model.layers.{i}.mlp.gate.weight_1 (vision router) + model.layers.{i}.mlp.moe_statics.e_score_correction_bias (concat text+vision) + model.layers.{i}.mlp.experts.{j}.gate_proj.weight (j=0..N-1 text, j=N..2N-1 vision) + model.layers.{i}.mlp.experts.{j}.up_proj.weight + model.layers.{i}.mlp.experts.{j}.down_proj.weight + model.layers.{i}.mlp.shared_experts.{gate,up,down}_proj.weight + model.vision_model.** + +Megatron Weight Naming (per-expert SequentialMLP within ErnieMultiTypeMoE): + language_model.decoder.layers.{i}.mlp.text_moe_layer.router.weight + language_model.decoder.layers.{i}.mlp.text_moe_layer.router.expert_bias + language_model.decoder.layers.{i}.mlp.text_moe_layer.experts.local_experts.{j}.linear_fc1.weight + language_model.decoder.layers.{i}.mlp.text_moe_layer.experts.local_experts.{j}.linear_fc2.weight + language_model.decoder.layers.{i}.mlp.vision_moe_layer.router.weight + language_model.decoder.layers.{i}.mlp.vision_moe_layer.router.expert_bias + language_model.decoder.layers.{i}.mlp.vision_moe_layer.experts.local_experts.{j}.linear_fc1.weight + language_model.decoder.layers.{i}.mlp.vision_moe_layer.experts.local_experts.{j}.linear_fc2.weight + language_model.decoder.layers.{i}.mlp.shared_experts.linear_fc1.weight + language_model.decoder.layers.{i}.mlp.shared_experts.linear_fc2.weight + +MG-native ViT Weight Naming (use_mg_vit=True, TP-sharded): + vision_model.decoder.layers.{i}.self_attention.linear_qkv.weight (fused QKV, ConcatenatedQKVMapping) + vision_model.decoder.layers.{i}.self_attention.linear_qkv.bias + vision_model.decoder.layers.{i}.self_attention.linear_qkv.layer_norm_weight (fused norm1) + vision_model.decoder.layers.{i}.self_attention.linear_qkv.layer_norm_bias + vision_model.decoder.layers.{i}.self_attention.linear_proj.weight + vision_model.decoder.layers.{i}.self_attention.linear_proj.bias + vision_model.decoder.layers.{i}.mlp.linear_fc1.weight + vision_model.decoder.layers.{i}.mlp.linear_fc1.bias + vision_model.decoder.layers.{i}.mlp.linear_fc1.layer_norm_weight (fused norm2) + vision_model.decoder.layers.{i}.mlp.linear_fc1.layer_norm_bias + vision_model.decoder.layers.{i}.mlp.linear_fc2.weight + vision_model.decoder.layers.{i}.mlp.linear_fc2.bias + vision_model.patch_embed.proj.weight (replicated) + vision_model.decoder.final_layernorm.weight + vision_model.decoder.final_layernorm.bias + +Note on Expert Parallelism: + EP>1 is supported for dual-pool MoE. The bridge handles the expert offset + between text and vision pools correctly: text experts use indices 0..N-1 and + vision experts use N..2N-1 in HF on-disk format. The framework's + `_megatron_local_name_to_global` function handles SequentialMLP-style expert + numbering, and `gather_from_ep_ranks` preserves pool offsets when + reconstructing HF parameter names during export. +""" + +import logging +import re +from typing import Dict, Optional, Tuple + +import torch + +from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.conversion.param_mapping import ( + AutoMapping, + ConcatenatedQKVMapping, + GatedMLPMapping, + QKVMapping, + ReplicatedMapping, + RowParallelMapping, +) +from megatron.bridge.models.ernie_vl.ernie45_vl_provider import Ernie45VLModelProvider +from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.model import Ernie45VLModel +from megatron.bridge.models.hf_pretrained.vlm import PreTrainedVLM +from megatron.bridge.utils.common_utils import extract_expert_number_from_param + + +logger = logging.getLogger(__name__) + +# Use string-based registration since the HF model class may not be importable +# if transformers is an older version or the model isn't registered yet. +_ERNIE45_VL_MOE_HF_CLASS_NAME = "Ernie4_5_VLMoeForConditionalGeneration" + + +# --------------------------------------------------------------------------- +# Vision pool expert offset mappings +# --------------------------------------------------------------------------- +# In ERNIE VL's dual-pool MoE, vision expert j maps to HF flat expert +# (j + num_text_experts). Two offset-aware mapping classes handle both +# directions: +# - resolve(): shifts the expert index wildcard for the HF side only +# - gather_from_ep_ranks(): reconstructs offset HF indices during EP export +# --------------------------------------------------------------------------- + + +def _offset_gather_from_ep_ranks( + mapping, + megatron_weights: Optional[torch.Tensor], + megatron_module, + hf_param_name: Optional[str] = None, +) -> Dict[str, torch.Tensor]: + """EP all-gather with pool offset for dual-pool MoE vision experts. + + Per EP rank *i* the HF expert index is: + expert_offset + local_expert_number + num_experts_per_rank * i + """ + if mapping.ep_size == 1: + return {str(hf_param_name): megatron_weights} + + if megatron_module is None: + num_experts_per_rank = mapping.broadcast_obj_from_pp_rank(None, "num_experts_per_rank") + else: + model_config = mapping._get_config(megatron_module) + num_experts = model_config.num_moe_experts + num_experts_per_rank = num_experts // mapping.ep_size + num_experts_per_rank = mapping.broadcast_obj_from_pp_rank(num_experts_per_rank, "num_experts_per_rank") + + global_expert_number = extract_expert_number_from_param(mapping.megatron_param) + local_expert_number = global_expert_number % num_experts_per_rank + + gathered_expert_param_names = [ + re.sub( + r"experts\.(\d+)", + f"experts.{mapping._expert_offset + local_expert_number + num_experts_per_rank * i}", + str(hf_param_name), + ) + for i in range(mapping.ep_size) + ] + assert str(hf_param_name) in gathered_expert_param_names, ( + f"hf_param_name {hf_param_name} not in gathered_expert_param_names {gathered_expert_param_names}" + ) + + gathered_weights = [torch.empty_like(megatron_weights) for _ in range(mapping.ep_size)] + torch.distributed.all_gather(gathered_weights, megatron_weights, group=mapping.ep_group) + + weights_dict: Dict[str, torch.Tensor] = {} + for i, param_name in enumerate(gathered_expert_param_names): + if param_name in weights_dict: + weights_dict[param_name] = torch.cat([weights_dict[param_name], gathered_weights[i].unsqueeze(0)], dim=0) + else: + weights_dict[param_name] = gathered_weights[i].unsqueeze(0) + for param_name in weights_dict: + weights_dict[param_name] = weights_dict[param_name].squeeze() + return weights_dict + + +def _resolve_with_offset( + megatron_pattern: str, + hf_pattern, + captures: Tuple[str, ...], + expert_offset: int, +) -> Tuple[str, ...]: + """Resolve wildcard captures, shifting the 2nd capture (expert index) for HF side.""" + if expert_offset and len(captures) >= 2: + shifted_expert = str(int(captures[1]) + expert_offset) + hf_captures = (captures[0], shifted_expert) + captures[2:] + else: + hf_captures = captures + + resolved_megatron = megatron_pattern + idx = 0 + while "**" in resolved_megatron and idx < len(captures): + resolved_megatron = resolved_megatron.replace("**", captures[idx], 1) + idx += 1 + while "*" in resolved_megatron and idx < len(captures): + resolved_megatron = resolved_megatron.replace("*", captures[idx], 1) + idx += 1 + + if isinstance(hf_pattern, dict): + resolved_hf: dict | str = {} + for k, v in hf_pattern.items(): + resolved_v = v + idx = 0 + while "**" in resolved_v and idx < len(hf_captures): + resolved_v = resolved_v.replace("**", hf_captures[idx], 1) + idx += 1 + while "*" in resolved_v and idx < len(hf_captures): + resolved_v = resolved_v.replace("*", hf_captures[idx], 1) + idx += 1 + resolved_hf[k] = resolved_v + else: + resolved_hf = hf_pattern + idx = 0 + while "**" in resolved_hf and idx < len(hf_captures): + resolved_hf = resolved_hf.replace("**", hf_captures[idx], 1) + idx += 1 + while "*" in resolved_hf and idx < len(hf_captures): + resolved_hf = resolved_hf.replace("*", hf_captures[idx], 1) + idx += 1 + + return resolved_megatron, resolved_hf + + +class _OffsetGatedMLPMapping(GatedMLPMapping): + """GatedMLPMapping with expert index offset for vision pool. + + Handles both directions: + - resolve(): shifts expert index for HF side only + - gather_from_ep_ranks(): reconstructs offset HF indices during EP export + """ + + def __init__(self, megatron_param: str, gate: str, up: str, expert_offset: int = 0): + super().__init__(megatron_param=megatron_param, gate=gate, up=up) + self._expert_offset = expert_offset + + def resolve(self, captures: Tuple[str, ...]): + resolved_megatron, resolved_hf = _resolve_with_offset( + self.megatron_param, + self.hf_param, + captures, + self._expert_offset, + ) + return _OffsetGatedMLPMapping( + megatron_param=resolved_megatron, + gate=resolved_hf["gate"], + up=resolved_hf["up"], + expert_offset=self._expert_offset, + ) + + def gather_from_ep_ranks(self, megatron_weights, megatron_module, hf_param_name=None): + return _offset_gather_from_ep_ranks(self, megatron_weights, megatron_module, hf_param_name) + + +class _OffsetRowParallelMapping(RowParallelMapping): + """RowParallelMapping with expert index offset for vision pool. + + Used for vision expert down_proj (linear_fc2), which is always + row-parallel in SequentialMLP. Using explicit RowParallelMapping + avoids the AutoMapping delegation issue where the delegate's + gather_from_ep_ranks bypasses offset logic. + """ + + def __init__(self, megatron_param: str, hf_param: str, expert_offset: int = 0): + super().__init__(megatron_param=megatron_param, hf_param=hf_param) + self._expert_offset = expert_offset + + def resolve(self, captures: Tuple[str, ...]): + resolved_megatron, resolved_hf = _resolve_with_offset( + self.megatron_param, + self.hf_param, + captures, + self._expert_offset, + ) + return _OffsetRowParallelMapping( + megatron_param=resolved_megatron, + hf_param=resolved_hf, + expert_offset=self._expert_offset, + ) + + def gather_from_ep_ranks(self, megatron_weights, megatron_module, hf_param_name=None): + return _offset_gather_from_ep_ranks(self, megatron_weights, megatron_module, hf_param_name) + + +class _ConcatBiasMapping(AutoMapping): + """Mapping for the concatenated text+vision expert bias tensor. + + The on-disk HF format stores a single ``moe_statics.e_score_correction_bias`` + tensor of shape ``[2, num_experts]`` where row 0 is the text pool bias and + row 1 is the vision pool bias. This mapping extracts the appropriate row + based on ``slice_name``. + + For export (megatron_to_hf), the text mapping buffers its bias in a + class-level dict keyed by resolved HF param name. The vision mapping + retrieves the buffered text bias, stacks them into ``[2, N]``, and + returns the merged tensor. This ensures only one entry per HF key. + """ + + # Class-level buffer: {resolved_hf_param: text_bias_tensor} + _export_buffer: Dict[str, torch.Tensor] = {} + + @classmethod + def clear_export_buffer(cls): + """Remove any stale entries from the class-level export buffer.""" + cls._export_buffer.clear() + + def __init__(self, megatron_param: str, hf_param: str, slice_name: str, num_experts: int): + super().__init__(megatron_param=megatron_param, hf_param=hf_param) + self._slice_name = slice_name # "text" or "vision" + self._num_experts = num_experts + self.allow_hf_name_mismatch = True + + def resolve(self, captures: Tuple[str, ...]): + resolved_megatron_param, resolved_hf_param = self._resolve_names(captures) + result = _ConcatBiasMapping( + megatron_param=resolved_megatron_param, + hf_param=resolved_hf_param, + slice_name=self._slice_name, + num_experts=self._num_experts, + ) + return result + + def hf_to_megatron(self, hf_weights, megatron_module): + """Extract the text or vision slice from the concatenated bias. + + On-disk shape is [2, num_experts]: row 0 = text, row 1 = vision. + """ + if self._slice_name == "text": + sliced = hf_weights[0] + else: + sliced = hf_weights[1] + return super().hf_to_megatron(sliced, megatron_module) + + def megatron_to_hf(self, megatron_weights, megatron_module): + """Export text+vision expert bias as concatenated [2, N] tensor. + + The text mapping buffers its bias; the vision mapping retrieves it + and stacks into [2, N]. If the text bias is not yet buffered + (shouldn't happen in practice), falls back to exporting as-is. + """ + result = super().megatron_to_hf(megatron_weights, megatron_module) + if result is None: + return result + + hf_key = str(self.hf_param) + + if self._slice_name == "text": + # Buffer the text bias, return empty dict (don't emit yet) + for _, tensor in result.items(): + _ConcatBiasMapping._export_buffer[hf_key] = tensor + return {} + else: + # Vision: retrieve buffered text bias and merge + text_bias = _ConcatBiasMapping._export_buffer.pop(hf_key, None) + if text_bias is not None: + for _, vision_bias in result.items(): + merged = torch.stack([text_bias, vision_bias], dim=0) + return {hf_key: merged} + # Fallback: no buffered text bias, just return vision as-is + return result + + +@MegatronModelBridge.register_bridge( + source=_ERNIE45_VL_MOE_HF_CLASS_NAME, + target=Ernie45VLModel, + provider=Ernie45VLModelProvider, + model_type="ernie4_5_vl_moe", +) +class Ernie45VLBridge(MegatronModelBridge): + """ + Megatron Bridge for ERNIE 4.5 VL MoE Conditional Generation. + + This bridge handles the conversion between HuggingFace Ernie4_5_VLMoeForConditionalGeneration + and Megatron-Core Ernie45VLModel formats, including weight mappings and + configuration translation for this vision-language MoE model. + + Key architectural features handled: + - Heterogeneous dual-pool MoE via ErnieMultiTypeMoE: + * text_moe_layer: standard Megatron MoELayer (TP support) + * vision_moe_layer: standard Megatron MoELayer (TP support) + - Shared experts across modalities + - 3D Multimodal RoPE (M-RoPE) + - Variable-resolution vision resampler (spatial + temporal merging) + - GQA with configurable query/KV heads + - HF on-disk per-expert weights <-> Megatron per-expert SequentialMLP weights + + Example: + >>> from megatron.bridge import AutoBridge + >>> bridge = AutoBridge.from_hf_pretrained("baidu/ERNIE-4.5-VL-28B-A3B-Instruct") + >>> provider = bridge.to_megatron_provider() + """ + + @staticmethod + def _get_text_config(hf_config): + """Extract the text/language config from either nested or flat HF config. + + The transformers-builtin ``Ernie4_5_VLMoeConfig`` (model_type=ernie4_5_vl_moe) + uses a nested ``text_config`` sub-object, while the custom auto_map config + ``Ernie4_5_VLMoEConfig`` (model_type=ernie4_5_moe_vl, e.g. the Thinking model) + uses a flat layout where all LLM fields live directly on the top-level config. + + Returns the appropriate config object (nested text_config or the config itself). + """ + text_config = getattr(hf_config, "text_config", None) + if text_config is not None: + return text_config + # Flat config: LLM fields are on hf_config itself + return hf_config + + @staticmethod + def _get_num_experts(text_config) -> int: + """Extract the per-pool number of experts as an int. + + The nested config stores ``moe_num_experts`` as a plain int (e.g. 4), + while the flat/Thinking config stores it as a list ``[64, 64]`` + (text pool, vision pool -- both values are always equal). + """ + raw = getattr(text_config, "moe_num_experts", 4) + if isinstance(raw, (list, tuple)): + return raw[0] + return raw + + def provider_bridge(self, hf_pretrained: PreTrainedVLM) -> Ernie45VLModelProvider: + """ + Create an Ernie45VLModelProvider from a HuggingFace pretrained model. + + Maps HuggingFace Ernie4_5_VLMoeConfig fields to Megatron provider parameters, + including vision config, MoE settings, M-RoPE sections, and token IDs. + + Supports both nested config (transformers builtin, model_type=ernie4_5_vl_moe) + and flat config (auto_map custom, model_type=ernie4_5_moe_vl). + + Args: + hf_pretrained: HuggingFace pretrained VLM model. + + Returns: + Ernie45VLModelProvider configured with the HF model's parameters. + """ + hf_config = hf_pretrained.config + text_config = self._get_text_config(hf_config) + + # Extract common config fields via base class utility + provider_kwargs = self.hf_config_to_provider_kwargs(text_config) + + # ERNIE 4.5 VL has moe_intermediate_size=[1536, 512] (list of 2 values + # for text/vision expert pools). CONFIG_MAPPING would auto-map this to + # moe_ffn_hidden_size, but that field expects a single int. Pop it here + # and set it explicitly below with the text expert size. + provider_kwargs.pop("moe_ffn_hidden_size", None) + + # Similarly, the attribute_map on the HF config aliases num_experts -> + # moe_num_experts, so CONFIG_MAPPING might double-set num_moe_experts. + # Pop MoE fields that we will set explicitly. + provider_kwargs.pop("num_moe_experts", None) + provider_kwargs.pop("moe_router_topk", None) + + provider = Ernie45VLModelProvider(**provider_kwargs) + + # --- Common LLM settings --- + provider.normalization = "RMSNorm" + provider.gated_linear_unit = True + provider.add_qkv_bias = False + provider.add_bias_linear = False + provider.hidden_dropout = 0.0 + # ERNIE 4.5 VL language model uses interleaved RoPE (pairs even/odd dims), + # unlike the LLaMA-style first-half/second-half split. + provider.rotary_interleaved = True + + # Extract rope_theta: nested config uses rope_parameters dict, flat config + # may use rope_scaling.mrope_section or a top-level rope_theta attribute. + rope_params = getattr(text_config, "rope_parameters", None) or {} + if isinstance(rope_params, dict): + provider.rotary_base = rope_params.get("rope_theta", getattr(text_config, "rope_theta", 500000.0)) + else: + provider.rotary_base = getattr(text_config, "rope_theta", 500000.0) + + # For VLMs, tie_word_embeddings lives on the top-level config, not text_config + provider.share_embeddings_and_output_weights = getattr(hf_config, "tie_word_embeddings", True) + + # --- MoE settings --- + num_experts = self._get_num_experts(text_config) + provider.moe_ffn_hidden_size = text_config.moe_intermediate_size[0] # 1536 (text experts) + provider.num_moe_experts = num_experts + provider.moe_router_topk = text_config.moe_k # 6 + provider.moe_router_pre_softmax = False + provider.moe_token_dispatcher_type = "alltoall" + provider.moe_router_load_balancing_type = "aux_loss" + provider.moe_aux_loss_coeff = getattr(text_config, "router_aux_loss_coef", 0.001) + # ERNIE 4.5 MoE uses sigmoid gating with expert bias for + # aux-free load balancing: + # 1. scores = sigmoid(logits) -- per-expert independent scores + # 2. scores_ = scores + e_score_correction_bias -- biased for top-k selection + # 3. weights = gather(scores, topk_indices) -- unbiased sigmoid scores + # 4. weights = weights / sum(weights) -- normalize to sum=1 + provider.moe_router_score_function = "sigmoid" + provider.moe_router_enable_expert_bias = True + provider.moe_router_dtype = "fp32" + provider.gradient_accumulation_fusion = False + + # Dual-pool MoE intermediate sizes + provider.moe_intermediate_size = tuple(text_config.moe_intermediate_size) # (1536, 512) + + # Shared experts: intermediate_size = moe_intermediate_size[0] * moe_num_shared_experts + # e.g. 1536 * 2 = 3072 + moe_num_shared_experts = getattr(text_config, "moe_num_shared_experts", 2) + provider.moe_shared_expert_intermediate_size = text_config.moe_intermediate_size[0] * moe_num_shared_experts + + # Determine which layers are dense vs MoE. + # Nested config (Instruct): mlp_layer_types = ["dense", "sparse", ...] + # Flat config (Thinking): moe_layer_start_index = [1, 1], + # moe_layer_end_index = [29, 28] + mlp_layer_types = getattr(text_config, "mlp_layer_types", None) + if mlp_layer_types is not None: + provider.moe_layer_freq = [0 if t == "dense" else 1 for t in mlp_layer_types] + else: + num_layers = text_config.num_hidden_layers + moe_start = getattr(text_config, "moe_layer_start_index", None) + if moe_start is not None: + # moe_layer_start_index can be a list (per-pool) or int. + # Take the first value: this is the first MoE layer index. + start = moe_start[0] if isinstance(moe_start, (list, tuple)) else moe_start + provider.moe_layer_freq = [0] * start + [1] * (num_layers - start) + else: + # Default: layer 0 dense, rest MoE + provider.moe_layer_freq = [0] + [1] * (num_layers - 1) + + # --- VL-specific overrides --- + provider.position_embedding_type = "mrope" + provider.vision_config = hf_config.vision_config + provider.hf_config = hf_config + + # M-RoPE section: [height, width, temporal] frequency allocation + # Nested config: rope_parameters.mrope_section + # Flat config: rope_scaling.mrope_section + mrope_section = None + if isinstance(rope_params, dict): + mrope_section = rope_params.get("mrope_section") + if mrope_section is None: + rope_scaling = getattr(text_config, "rope_scaling", None) or {} + if isinstance(rope_scaling, dict): + mrope_section = rope_scaling.get("mrope_section") + provider.mrope_section = mrope_section or [22, 22, 20] + + # Token IDs -- these live on the top-level config in both formats + provider.image_start_token_id = getattr(hf_config, "image_start_token_id", 101304) + provider.image_end_token_id = getattr(hf_config, "image_end_token_id", 101305) + provider.image_token_id = getattr(hf_config, "image_token_id", getattr(hf_config, "im_patch_id", 100295)) + provider.video_start_token_id = getattr(hf_config, "video_start_token_id", 101306) + provider.video_end_token_id = getattr(hf_config, "video_end_token_id", 101307) + provider.video_token_id = getattr(hf_config, "video_token_id", 103367) + + return provider + + def stream_weights_megatron_to_hf(self, *args, **kwargs): + """Override to clear the _ConcatBiasMapping export buffer before each run.""" + _ConcatBiasMapping.clear_export_buffer() + return super().stream_weights_megatron_to_hf(*args, **kwargs) + + def mapping_registry(self) -> MegatronMappingRegistry: + """ + Return MegatronMappingRegistry with parameter mappings for ERNIE 4.5 VL MoE. + + Uses the HF **on-disk (safetensors)** key format, which differs from the + in-memory ``state_dict()`` format due to HuggingFace's ``_checkpoint_conversion_mapping``. + + On-disk format: + - No ``language_model.`` prefix: ``model.layers.*`` not ``model.language_model.layers.*`` + - Per-expert flat-indexed weights: ``experts.{j}.gate_proj.weight`` + - Text experts indices 0..N-1, vision experts indices N..2N-1 + - Single ``gate.weight`` (text router) and ``gate.weight_1`` (vision router) + - Concatenated ``moe_statics.e_score_correction_bias`` for text+vision + - ``model.vision_model.**`` not ``model.vision_tower.**`` + - Resampler: ``spatial_linear.0/2/3`` not ``spatial_linear.fc1/fc2/ln`` + (same for ``temporal_linear``) + + Returns: + MegatronMappingRegistry with all parameter mappings. + """ + # Get num_experts from the HF config (injected by the bridge dispatch). + # Falls back to 4 for toy model / direct instantiation. + num_experts = 4 + is_flat_config = False + use_mg_vit = False + if hasattr(self, "hf_config"): + text_config = self._get_text_config(self.hf_config) + num_experts = self._get_num_experts(text_config) + # Detect flat config (Thinking / auto_map) vs nested config (Instruct). + # + # Simple ``not hasattr(hf_config, 'text_config')`` is unreliable because + # ``_normalize_hf_config()`` in modeling_ernie45_vl.py mutates the config + # object to add ``text_config = hf_config`` (self-reference) so that the + # HF resampler can access ``config.text_config.hidden_size``. After this + # mutation ``hasattr`` returns True even for flat configs. + # + # Instead, detect flat config by checking whether ``text_config`` is absent + # OR is a self-reference (points back to hf_config itself). A genuinely + # nested config has ``text_config`` as a distinct sub-object. + text_cfg_attr = getattr(self.hf_config, "text_config", None) + is_flat_config = (text_cfg_attr is None) or (text_cfg_attr is self.hf_config) + + # Check use_mg_vit: set externally on the bridge instance (like hf_config). + # When True, the Megatron model uses MG-native ViT (TP-sharded weights). + # When False (default), it uses HF-wrapped ViT (replicated weights). + use_mg_vit = getattr(self, "use_mg_vit", False) + + # Determine on-disk vision key prefix based on config format. + # Flat config (Thinking/auto_map): "vision_model.**" + # Nested config (Instruct/transformers-builtin): "model.vision_model.**" + vision_hf_prefix = "vision_model.**" if is_flat_config else "model.vision_model.**" + # Per-param prefix (without glob suffix) for MG ViT block-level mappings. + vision_hf_block_prefix = "vision_model" if is_flat_config else "model.vision_model" + + # ===================================================================== + # Simple 1:1 parameter mappings (AutoMapping detects parallelism) + # ===================================================================== + param_mappings = { + # ================================================================= + # Language Model: Embeddings and output + # ================================================================= + "language_model.embedding.word_embeddings.weight": "model.embed_tokens.weight", + "language_model.decoder.final_layernorm.weight": "model.norm.weight", + # ================================================================= + # Language Model: Self-attention (all layers) + # input_layernorm is fused into TELayerNormColumnParallelLinear + # ================================================================= + "language_model.decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": ( + "model.layers.*.input_layernorm.weight" + ), + "language_model.decoder.layers.*.self_attention.linear_proj.weight": ( + "model.layers.*.self_attn.o_proj.weight" + ), + # ================================================================= + # Dense MLP (layer 0): post_attention_layernorm fused into + # TELayerNormColumnParallelLinear (linear_fc1) + # ================================================================= + "language_model.decoder.layers.*.mlp.linear_fc1.layer_norm_weight": ( + "model.layers.*.post_attention_layernorm.weight" + ), + "language_model.decoder.layers.*.mlp.linear_fc2.weight": ("model.layers.*.mlp.down_proj.weight"), + # ================================================================= + # MoE layers: pre_mlp_layernorm (separate, not fused) + # ================================================================= + "language_model.decoder.layers.*.pre_mlp_layernorm.weight": ( + "model.layers.*.post_attention_layernorm.weight" + ), + # ================================================================= + # Shared experts down projection + # ================================================================= + "language_model.decoder.layers.*.mlp.shared_experts.linear_fc2.weight": ( + "model.layers.*.mlp.shared_experts.down_proj.weight" + ), + } + + mapping_list = [] + for megatron_param, hf_param in param_mappings.items(): + mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param)) + + # ===================================================================== + # Vision encoder mappings + # + # Two modes depending on use_mg_vit: + # - HF ViT (use_mg_vit=False): single ReplicatedMapping for all + # vision_tower weights (replicated across TP ranks). + # - MG ViT (use_mg_vit=True): per-parameter TP-sharded mappings + # using ConcatenatedQKVMapping for fused QKV and AutoMapping + # for other TP-auto-detected parameters. + # + # On-disk key prefix differs by config format: + # Flat (Thinking/auto_map): "vision_model.*" + # Nested (Instruct/transformers): "model.vision_model.*" + # ===================================================================== + if use_mg_vit: + # MG-native ViT: TP-sharded weight mappings + # Megatron key prefix: "vision_model.decoder.layers.*..." + # HF on-disk prefix: "{vision_hf_block_prefix}.blocks.*..." + vit_param_mappings = { + # Attention: proj weight/bias (TP-sharded via AutoMapping) + "vision_model.decoder.layers.*.self_attention.linear_proj.weight": ( + f"{vision_hf_block_prefix}.blocks.*.attn.proj.weight" + ), + "vision_model.decoder.layers.*.self_attention.linear_proj.bias": ( + f"{vision_hf_block_prefix}.blocks.*.attn.proj.bias" + ), + # MLP: fc1 weight/bias (TP column-parallel via AutoMapping) + "vision_model.decoder.layers.*.mlp.linear_fc1.weight": ( + f"{vision_hf_block_prefix}.blocks.*.mlp.fc1.weight" + ), + "vision_model.decoder.layers.*.mlp.linear_fc1.bias": ( + f"{vision_hf_block_prefix}.blocks.*.mlp.fc1.bias" + ), + # MLP: fc2 weight/bias (TP row-parallel via AutoMapping) + "vision_model.decoder.layers.*.mlp.linear_fc2.weight": ( + f"{vision_hf_block_prefix}.blocks.*.mlp.fc2.weight" + ), + "vision_model.decoder.layers.*.mlp.linear_fc2.bias": ( + f"{vision_hf_block_prefix}.blocks.*.mlp.fc2.bias" + ), + # LayerNorm: norm1 fused into linear_qkv (TELayerNormColumnParallelLinear) + "vision_model.decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": ( + f"{vision_hf_block_prefix}.blocks.*.norm1.weight" + ), + "vision_model.decoder.layers.*.self_attention.linear_qkv.layer_norm_bias": ( + f"{vision_hf_block_prefix}.blocks.*.norm1.bias" + ), + # LayerNorm: norm2 fused into linear_fc1 (TELayerNormColumnParallelLinear) + "vision_model.decoder.layers.*.mlp.linear_fc1.layer_norm_weight": ( + f"{vision_hf_block_prefix}.blocks.*.norm2.weight" + ), + "vision_model.decoder.layers.*.mlp.linear_fc1.layer_norm_bias": ( + f"{vision_hf_block_prefix}.blocks.*.norm2.bias" + ), + # Final LayerNorm (post_layer_norm in TransformerBlock) + "vision_model.decoder.final_layernorm.weight": (f"{vision_hf_block_prefix}.ln.weight"), + "vision_model.decoder.final_layernorm.bias": (f"{vision_hf_block_prefix}.ln.bias"), + } + for mg_param, hf_param in vit_param_mappings.items(): + mapping_list.append(AutoMapping(megatron_param=mg_param, hf_param=hf_param)) + + # Fused QKV: ConcatenatedQKVMapping handles the [Q|K|V] -> interleaved + # GQA layout conversion, with TP-aware splitting. + # ERNIE ViT uses fused attn.qkv.weight/bias on disk. + mapping_list.extend( + [ + ConcatenatedQKVMapping( + megatron_param="vision_model.decoder.layers.*.self_attention.linear_qkv.weight", + hf_param=f"{vision_hf_block_prefix}.blocks.*.attn.qkv.weight", + ), + ConcatenatedQKVMapping( + megatron_param="vision_model.decoder.layers.*.self_attention.linear_qkv.bias", + hf_param=f"{vision_hf_block_prefix}.blocks.*.attn.qkv.bias", + ), + ] + ) + + # Patch embedding: replicated across TP ranks (not TP-sharded). + # ERNIE ViT PatchEmbed is nn.Linear with weight only (no bias). + mapping_list.append( + ReplicatedMapping( + megatron_param="vision_model.patch_embed.proj.**", + hf_param=f"{vision_hf_block_prefix}.patch_embed.proj.**", + ), + ) + else: + # HF-wrapped ViT: all weights replicated across TP ranks. + mapping_list.append( + ReplicatedMapping( + megatron_param="vision_tower.**", + hf_param=vision_hf_prefix, + ), + ) + + # ===================================================================== + # Special mappings requiring parameter transformation + # ===================================================================== + mapping_list.extend( + [ + # ============================================================= + # Resampler / projector: replicated across TP ranks + # + # On-disk keys use sequential indices (0, 2, 3) for + # spatial_linear and temporal_linear sub-modules, while + # Megatron/HF in-memory uses named attributes (fc1, fc2, ln). + # HF's _checkpoint_conversion_mapping reverses this: + # spatial_linear.0 <-> spatial_linear.fc1 + # spatial_linear.2 <-> spatial_linear.fc2 + # spatial_linear.3 <-> spatial_linear.ln + # (same for temporal_linear) + # + # We must use on-disk key format since SafeTensorsStateSource + # returns raw on-disk keys (no HF renaming applied). + # ============================================================= + # spatial_linear: fc1 -> 0, fc2 -> 2, ln -> 3 + ReplicatedMapping( + megatron_param="resampler_model.spatial_linear.fc1.**", + hf_param="model.resampler_model.spatial_linear.0.**", + ), + ReplicatedMapping( + megatron_param="resampler_model.spatial_linear.fc2.**", + hf_param="model.resampler_model.spatial_linear.2.**", + ), + ReplicatedMapping( + megatron_param="resampler_model.spatial_linear.ln.**", + hf_param="model.resampler_model.spatial_linear.3.**", + ), + # temporal_linear: fc1 -> 0, fc2 -> 2, ln -> 3 + ReplicatedMapping( + megatron_param="resampler_model.temporal_linear.fc1.**", + hf_param="model.resampler_model.temporal_linear.0.**", + ), + ReplicatedMapping( + megatron_param="resampler_model.temporal_linear.fc2.**", + hf_param="model.resampler_model.temporal_linear.2.**", + ), + ReplicatedMapping( + megatron_param="resampler_model.temporal_linear.ln.**", + hf_param="model.resampler_model.temporal_linear.3.**", + ), + # Remaining resampler params (no on-disk renaming) + ReplicatedMapping( + megatron_param="resampler_model.mlp.**", + hf_param="model.resampler_model.mlp.**", + ), + ReplicatedMapping( + megatron_param="resampler_model.after_norm.**", + hf_param="model.resampler_model.after_norm.**", + ), + # ============================================================= + # Text MoE router weight (transposed on disk: + # on-disk [hidden_size, num_experts] -> Megatron [num_experts, hidden_size]) + # ============================================================= + AutoMapping( + megatron_param="language_model.decoder.layers.*.mlp.text_moe_layer.router.weight", + hf_param="model.layers.*.mlp.gate.weight", + permute_dims=(1, 0), + ), + # ============================================================= + # Vision MoE router weight (saved as gate.weight_1 on disk, also transposed) + # ============================================================= + AutoMapping( + megatron_param="language_model.decoder.layers.*.mlp.vision_moe_layer.router.weight", + hf_param="model.layers.*.mlp.gate.weight_1", + permute_dims=(1, 0), + ), + # ============================================================= + # QKV (fused Q, K, V into single QKV matrix) + # ============================================================= + QKVMapping( + megatron_param="language_model.decoder.layers.*.self_attention.linear_qkv.weight", + q="model.layers.*.self_attn.q_proj.weight", + k="model.layers.*.self_attn.k_proj.weight", + v="model.layers.*.self_attn.v_proj.weight", + ), + # ============================================================= + # Dense MLP (layer 0): gate_proj + up_proj -> fused linear_fc1 + # ============================================================= + GatedMLPMapping( + megatron_param="language_model.decoder.layers.*.mlp.linear_fc1.weight", + gate="model.layers.*.mlp.gate_proj.weight", + up="model.layers.*.mlp.up_proj.weight", + ), + # ============================================================= + # Text MoE experts (per-expert, using DeepSeek-style wildcards) + # Megatron: local_experts.{j}.linear_fc1.weight + # HF on-disk: experts.{j}.gate_proj.weight + experts.{j}.up_proj.weight + # Direct index mapping (text expert j -> HF expert j) + # ============================================================= + GatedMLPMapping( + megatron_param=( + "language_model.decoder.layers.*.mlp.text_moe_layer.experts.local_experts.*.linear_fc1.weight" + ), + gate="model.layers.*.mlp.experts.*.gate_proj.weight", + up="model.layers.*.mlp.experts.*.up_proj.weight", + ), + AutoMapping( + megatron_param=( + "language_model.decoder.layers.*.mlp.text_moe_layer.experts.local_experts.*.linear_fc2.weight" + ), + hf_param="model.layers.*.mlp.experts.*.down_proj.weight", + ), + # ============================================================= + # Vision MoE experts (per-expert, with expert index offset) + # Megatron vision expert j -> HF expert (j + num_experts) + # ============================================================= + _OffsetGatedMLPMapping( + megatron_param=( + "language_model.decoder.layers.*.mlp.vision_moe_layer" + ".experts.local_experts.*.linear_fc1.weight" + ), + gate="model.layers.*.mlp.experts.*.gate_proj.weight", + up="model.layers.*.mlp.experts.*.up_proj.weight", + expert_offset=num_experts, + ), + _OffsetRowParallelMapping( + megatron_param=( + "language_model.decoder.layers.*.mlp.vision_moe_layer" + ".experts.local_experts.*.linear_fc2.weight" + ), + hf_param="model.layers.*.mlp.experts.*.down_proj.weight", + expert_offset=num_experts, + ), + # ============================================================= + # Shared experts: gate+up -> fused linear_fc1 + # ============================================================= + GatedMLPMapping( + megatron_param="language_model.decoder.layers.*.mlp.shared_experts.linear_fc1.weight", + gate="model.layers.*.mlp.shared_experts.gate_proj.weight", + up="model.layers.*.mlp.shared_experts.up_proj.weight", + ), + # ============================================================= + # Expert bias: concatenated [text; vision] on disk + # Text router expert_bias -> first N entries + # Vision router expert_bias -> last N entries + # ============================================================= + _ConcatBiasMapping( + megatron_param=("language_model.decoder.layers.*.mlp.text_moe_layer.router.expert_bias"), + hf_param="model.layers.*.mlp.moe_statics.e_score_correction_bias", + slice_name="text", + num_experts=num_experts, + ), + _ConcatBiasMapping( + megatron_param=("language_model.decoder.layers.*.mlp.vision_moe_layer.router.expert_bias"), + hf_param="model.layers.*.mlp.moe_statics.e_score_correction_bias", + slice_name="vision", + num_experts=num_experts, + ), + ] + ) + + return MegatronMappingRegistry(*mapping_list) diff --git a/src/megatron/bridge/models/ernie_vl/ernie45_vl_provider.py b/src/megatron/bridge/models/ernie_vl/ernie45_vl_provider.py new file mode 100644 index 0000000000..ce71eacf76 --- /dev/null +++ b/src/megatron/bridge/models/ernie_vl/ernie45_vl_provider.py @@ -0,0 +1,152 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Provider for ERNIE 4.5 VL MoE model. + +Maps HuggingFace Ernie4_5_VLMoeConfig to Megatron-Core TransformerConfig +and provides model instantiation logic for the dual-pool MoE architecture. + +The language model uses a custom ErnieMultiTypeMoE layer containing both +text_moe_layer and vision_moe_layer as separate MoELayer instances, each +with their own router, experts, and EP support. +""" + +from dataclasses import dataclass, field +from typing import Any, Callable, List, Tuple, Union + +from megatron.core.models.gpt import GPTModel as MCoreGPTModel +from megatron.core.transformer.spec_utils import ModuleSpec + + +try: + from transformers.models.ernie4_5_vl_moe.configuration_ernie4_5_vl_moe import ( + Ernie4_5_VLMoeConfig, + Ernie4_5_VLMoeVisionConfig, + ) +except ImportError: + # Fallback for environments where the builtin transformers class is not available + # (e.g. auto_map models only). Use generic types for type hints. + Ernie4_5_VLMoeConfig = None + Ernie4_5_VLMoeVisionConfig = None + +from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.ernie_decoder_layer_spec import ( + get_ernie45_vl_decoder_block_spec, +) +from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.model import Ernie45VLModel +from megatron.bridge.models.gpt_provider import GPTModelProvider + + +@dataclass +class Ernie45VLModelProvider(GPTModelProvider): + """ + Model provider for ERNIE 4.5 VL MoE. + + This provider extends GPTModelProvider with ERNIE 4.5 VL-specific fields: + - Vision configuration for the ViT encoder and resampler + - 3D M-RoPE parameters (mrope_section) + - Dual-pool MoE configuration (moe_intermediate_size as tuple) + - Custom decoder layer spec with ErnieMultiTypeMoE + - Token IDs for image/video placeholder tokens + - Freeze options for vision/language components + """ + + # VL models shouldn't scatter embeddings across sequence parallel regions + # because the vision embeddings are going to be inserted into the language embeddings. + scatter_embedding_sequence_parallel: bool = False + + # Position embedding type: M-RoPE for multimodal 3D positions + position_embedding_type: str = "mrope" + mrope_section: List[int] = field(default_factory=lambda: [22, 22, 20]) + + # Vision configuration -- accepts either Ernie4_5_VLMoeVisionConfig (nested) or + # DFNRopeVisionTransformerConfig (flat/auto_map) or any config-like object. + vision_config: Any = field( + default_factory=lambda: Ernie4_5_VLMoeVisionConfig() if Ernie4_5_VLMoeVisionConfig else None + ) + hf_config: Any = None + + # Dual-pool MoE intermediate sizes: (text_ffn_size, vision_ffn_size) + # This is passed to ErnieMultiTypeMoE which creates separate configs per pool. + moe_intermediate_size: Tuple[int, int] = (1536, 512) + + # Token IDs (from Ernie4_5_VLMoeConfig defaults) + image_start_token_id: int = 101304 + image_end_token_id: int = 101305 + image_token_id: int = 100295 + video_start_token_id: int = 101306 + video_end_token_id: int = 101307 + video_token_id: int = 103367 + + # Freeze options + freeze_language_model: bool = False + freeze_vision_model: bool = False + freeze_vision_projection: bool = False + + # Use MG-native ViT instead of HF-wrapped ViT for better TP performance. + # When False (default), the vision encoder uses the HuggingFace implementation + # replicated across TP ranks. When True, uses Megatron-Core TransformerBlock + # with TE modules for TP-native attention and MLP layers. + use_mg_vit: bool = False + + # Use custom decoder block spec for heterogeneous layers (dense + dual-pool MoE) + transformer_layer_spec: Union[ModuleSpec, Callable[["GPTModelProvider"], ModuleSpec]] = ( + get_ernie45_vl_decoder_block_spec + ) + + def provide(self, pre_process=None, post_process=None, vp_stage=None) -> Ernie45VLModel: + """Build the composite VLM model (vision + resampler + language model). + + Args: + pre_process: Whether to include pre-processing (embedding + vision). Defaults to first PP stage. + post_process: Whether to include post-processing (output layer). Defaults to last PP stage. + vp_stage: Virtual pipeline stage index. + + Returns: + Ernie45VLModel: Configured ERNIE 4.5 VL MoE model instance. + """ + model = Ernie45VLModel( + self, + pre_process=pre_process, + post_process=post_process, + vp_stage=vp_stage, + ) + + # Apply freeze options if any are enabled + if self.freeze_language_model or self.freeze_vision_model or self.freeze_vision_projection: + model.freeze( + freeze_language_model=self.freeze_language_model, + freeze_vision_model=self.freeze_vision_model, + freeze_vision_projection=self.freeze_vision_projection, + ) + + return model + + def provide_language_model(self, pre_process=None, post_process=None, vp_stage=None) -> MCoreGPTModel: + """Build only the language model (MCoreGPTModel) for weight conversion. + + This uses GPTModelProvider.provide() which builds a standard MCoreGPTModel + but with the custom ErnieMultiTypeMoE layer spec set via transformer_layer_spec. + The resulting model has both text_moe_layer and vision_moe_layer as proper + submodules of each MoE transformer layer. + + Args: + pre_process: Whether to include pre-processing. + post_process: Whether to include post-processing. + vp_stage: Virtual pipeline stage index. + + Returns: + MCoreGPTModel: Configured Megatron-Core GPT model instance with dual-pool MoE. + """ + return GPTModelProvider.provide(self, pre_process=pre_process, post_process=post_process, vp_stage=vp_stage) diff --git a/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/__init__.py b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/__init__.py new file mode 100644 index 0000000000..bfbbb35bcd --- /dev/null +++ b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/__init__.py @@ -0,0 +1,18 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.model import Ernie45VLModel + + +__all__ = ["Ernie45VLModel"] diff --git a/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/ernie_decoder_layer_spec.py b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/ernie_decoder_layer_spec.py new file mode 100644 index 0000000000..7b8c6761d7 --- /dev/null +++ b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/ernie_decoder_layer_spec.py @@ -0,0 +1,276 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Decoder layer spec for ERNIE 4.5 VL MoE. + +Creates heterogeneous transformer block specs where: +- Layer 0: dense MLP +- Layers 1+: ErnieMultiTypeMoE (dual-pool MoE with text + vision expert pools) + +The text and vision MoE pools each use standard Megatron MoELayer with +SequentialMLP experts, enabling full TP/EP compatibility through standard +Megatron-Core infrastructure. +""" + +from typing import Optional + +from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add +from megatron.core.transformer.attention import SelfAttention, SelfAttentionSubmodules +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.mlp import MLP, MLPSubmodules +from megatron.core.transformer.moe.experts import SequentialMLP +from megatron.core.transformer.moe.moe_layer import MoELayer, MoESubmodules +from megatron.core.transformer.moe.shared_experts import SharedExpertMLP +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_block import ( + TransformerBlockSubmodules, + get_num_layers_to_build, +) +from megatron.core.transformer.transformer_layer import ( + TransformerLayer, + TransformerLayerSubmodules, + get_transformer_layer_offset, +) + +from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.ernie_moe_layer import ( + ErnieMultiTypeMoE, + MultiTypeMoeSubmodules, +) + + +try: + import transformer_engine # noqa: F401 + + HAVE_TE = True +except (ImportError, ModuleNotFoundError): + HAVE_TE = False + + +def _get_linear_modules(): + """Get appropriate linear module classes based on TE availability.""" + if HAVE_TE: + from megatron.core.extensions.transformer_engine import ( + TEColumnParallelLinear, + TEDotProductAttention, + TELayerNormColumnParallelLinear, + TENorm, + TERowParallelLinear, + ) + + return ( + TEColumnParallelLinear, + TERowParallelLinear, + TEDotProductAttention, + TELayerNormColumnParallelLinear, + TENorm, + ) + else: + from megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear + from megatron.core.transformer.dot_product_attention import DotProductAttention + from megatron.core.transformer.torch_norm import WrappedTorchNorm + + return ( + ColumnParallelLinear, + RowParallelLinear, + DotProductAttention, + ColumnParallelLinear, # No fused LN+Linear without TE + WrappedTorchNorm, + ) + + +def _get_mlp_module_spec( + num_experts: Optional[int] = None, + moe_grouped_gemm: bool = False, +) -> ModuleSpec: + """Get MLP module spec for dense or dual-pool MoE layers. + + Args: + num_experts: Number of experts per pool. None for dense MLP. + moe_grouped_gemm: Whether to use grouped GEMM for experts. + + Returns: + ModuleSpec for either dense MLP or ErnieMultiTypeMoE. + """ + ColumnParallel, RowParallel, _, LayerNormColumnParallel, _ = _get_linear_modules() + + if num_experts is None: + # Dense MLP (layer 0): uses TELayerNormColumnParallelLinear to fuse + # post_attention_layernorm with linear_fc1 + return ModuleSpec( + module=MLP, + submodules=MLPSubmodules( + linear_fc1=LayerNormColumnParallel, + linear_fc2=RowParallel, + ), + ) + + # Expert MLP spec (used by both text and vision pools) + if moe_grouped_gemm: + try: + from megatron.core.transformer.moe.experts import TEGroupedMLP + + experts_spec = ModuleSpec(module=TEGroupedMLP) + except ImportError: + experts_spec = ModuleSpec( + module=SequentialMLP, + submodules=MLPSubmodules( + linear_fc1=ColumnParallel, + linear_fc2=RowParallel, + ), + ) + else: + experts_spec = ModuleSpec( + module=SequentialMLP, + submodules=MLPSubmodules( + linear_fc1=ColumnParallel, + linear_fc2=RowParallel, + ), + ) + + # Each pool is a standard MoELayer + base_moe_spec = ModuleSpec( + module=MoELayer, + submodules=MoESubmodules( + experts=experts_spec, + ), + ) + + # Shared experts MLP + shared_experts_spec = ModuleSpec( + module=SharedExpertMLP, + submodules=MLPSubmodules( + linear_fc1=ColumnParallel, + linear_fc2=RowParallel, + ), + ) + + # Dual-pool MoE + return ModuleSpec( + module=ErnieMultiTypeMoE, + submodules=MultiTypeMoeSubmodules( + text_moe_layer=base_moe_spec, + vision_moe_layer=base_moe_spec, + shared_experts=shared_experts_spec, + ), + ) + + +def _get_ernie_decoder_layer_spec( + num_experts: Optional[int] = None, + moe_grouped_gemm: bool = False, +) -> ModuleSpec: + """Get a single transformer layer spec. + + Args: + num_experts: Number of experts per pool. None for dense layer. + moe_grouped_gemm: Whether to use grouped GEMM. + + Returns: + ModuleSpec for a TransformerLayer. + """ + _, RowParallel, DotProductAttention, LayerNormColumnParallel, Norm = _get_linear_modules() + + mlp_spec = _get_mlp_module_spec( + num_experts=num_experts, + moe_grouped_gemm=moe_grouped_gemm, + ) + + # For dense layers, the post_attention_layernorm is fused into + # TELayerNormColumnParallelLinear (linear_fc1), so pre_mlp_layernorm = IdentityOp (default). + # For MoE layers, a separate pre_mlp_layernorm is needed since the MoE + # does not have a fused layernorm path. + layer_submodules = TransformerLayerSubmodules( + self_attention=ModuleSpec( + module=SelfAttention, + params={"attn_mask_type": AttnMaskType.causal}, + submodules=SelfAttentionSubmodules( + linear_qkv=LayerNormColumnParallel, + core_attention=DotProductAttention, + linear_proj=RowParallel, + ), + ), + self_attn_bda=get_bias_dropout_add, + mlp=mlp_spec, + mlp_bda=get_bias_dropout_add, + ) + + # MoE layers need separate pre_mlp_layernorm (not fused into MLP) + if num_experts is not None: + layer_submodules.pre_mlp_layernorm = Norm + + return ModuleSpec( + module=TransformerLayer, + submodules=layer_submodules, + ) + + +def get_ernie45_vl_decoder_block_spec( + config, + use_transformer_engine: bool = True, +) -> TransformerBlockSubmodules: + """Get the full decoder block spec for ERNIE 4.5 VL MoE. + + Creates a heterogeneous block where layer types are determined by + config.moe_layer_freq (list of 0/1 per layer): + - 0: dense MLP layer + - 1: ErnieMultiTypeMoE layer (dual-pool MoE) + + Args: + config: TransformerConfig with moe_layer_freq, num_moe_experts, etc. + use_transformer_engine: Whether to use TE modules. + + Returns: + TransformerBlockSubmodules with heterogeneous layer specs. + """ + num_experts = getattr(config, "num_moe_experts", None) + moe_grouped_gemm = getattr(config, "moe_grouped_gemm", False) + + # Dense layer spec (no MoE) + dense_layer_spec = _get_ernie_decoder_layer_spec( + num_experts=None, + moe_grouped_gemm=False, + ) + + # MoE layer spec (dual-pool) + moe_layer_spec = _get_ernie_decoder_layer_spec( + num_experts=num_experts, + moe_grouped_gemm=moe_grouped_gemm, + ) + + # Build per-layer specs based on moe_layer_freq + moe_layer_freq = getattr(config, "moe_layer_freq", None) + if moe_layer_freq is None: + # Default: all MoE + moe_layer_freq = [1] * config.num_layers + + layer_specs = [] + for i in range(config.num_layers): + if isinstance(moe_layer_freq, list): + is_moe = moe_layer_freq[i] + else: + is_moe = moe_layer_freq + layer_specs.append(moe_layer_spec if is_moe else dense_layer_spec) + + # Slice for pipeline parallelism + offset = get_transformer_layer_offset(config) + num_layers_to_build = get_num_layers_to_build(config) + layer_specs = layer_specs[offset : offset + num_layers_to_build] + + # Get the Norm class for final_layernorm (TENorm or WrappedTorchNorm). + # Without this, TransformerBlock.final_layernorm would be None because + # TransformerBlockSubmodules.layer_norm defaults to None. + _, _, _, _, Norm = _get_linear_modules() + + return TransformerBlockSubmodules(layer_specs=layer_specs, layer_norm=Norm) diff --git a/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/ernie_moe_layer.py b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/ernie_moe_layer.py new file mode 100644 index 0000000000..09087c1f53 --- /dev/null +++ b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/ernie_moe_layer.py @@ -0,0 +1,288 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Custom MoE modules for ERNIE 4.5 VL MoE dual-pool architecture. + +ERNIE 4.5 VL uses a heterogeneous dual-pool MoE where each transformer layer +(except layer 0 which is dense) contains: + +- text_moe_layer: 64 experts with intermediate_size=1536 for text tokens +- vision_moe_layer: 64 experts with intermediate_size=512 for vision tokens +- shared_experts: 2 shared experts with intermediate_size=3072 for all tokens + +Both pools use separate routers and expert sets. Tokens are dispatched to their +respective pool based on modality (token_type_ids: 0=text, 1=vision). + +Module hierarchy (MoE layers): + decoder.layers.{i}.mlp = ErnieMultiTypeMoE + .text_moe_layer = MoELayer (standard Megatron) + .router = TopKRouter + .experts = SequentialMLP + .local_experts.{j} = MLP (with linear_fc1, linear_fc2) + .vision_moe_layer = MoELayer (standard Megatron) + .router = TopKRouter + .experts = SequentialMLP + .local_experts.{j} = MLP (with linear_fc1, linear_fc2) + .shared_experts = SharedExpertMLP + .linear_fc1, .linear_fc2 + +Communication pattern for moe_mm_token_type_ids: + Megatron-Core's TransformerBlock / TransformerLayer do not propagate extra + kwargs to MLP layers. To pass ``moe_mm_token_type_ids`` from + ``Ernie45VLModel.forward()`` to ``ErnieMultiTypeMoE.forward()`` we use a + module-level context variable ``_current_moe_mm_token_type_ids`` that is set + before the language model forward and cleared afterwards. +""" + +from copy import deepcopy +from dataclasses import dataclass +from typing import Optional, Union + +import torch +from megatron.core import parallel_state +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.moe.moe_utils import get_default_pg_collection +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_config import TransformerConfig + + +# Module-level context variable for passing moe_mm_token_type_ids from +# Ernie45VLModel.forward() to ErnieMultiTypeMoE.forward() without modifying +# Megatron-Core's TransformerBlock / TransformerLayer forward signatures. +# Set by Ernie45VLModel.forward() before calling language_model.forward(), +# cleared after. Read by ErnieMultiTypeMoE.forward() when token_type_ids=None. +_current_moe_mm_token_type_ids: "torch.Tensor | None" = None + + +def set_moe_mm_token_type_ids(token_type_ids): + """Set the current moe_mm_token_type_ids for MoE routing. + + Called by ``Ernie45VLModel.forward()`` before ``language_model.forward()``. + """ + global _current_moe_mm_token_type_ids + _current_moe_mm_token_type_ids = token_type_ids + + +def clear_moe_mm_token_type_ids(): + """Clear the current moe_mm_token_type_ids after forward pass. + + Called by ``Ernie45VLModel.forward()`` after ``language_model.forward()``. + """ + global _current_moe_mm_token_type_ids + _current_moe_mm_token_type_ids = None + + +@dataclass +class MultiTypeMoeSubmodules: + """Submodule specs for the dual-pool MoE layer. + + Attributes: + text_moe_layer: Spec for the text MoE pool (larger FFN). + vision_moe_layer: Spec for the vision MoE pool (smaller FFN). + shared_experts: Spec for the shared expert MLP. + """ + + text_moe_layer: Union[ModuleSpec, type] = None + vision_moe_layer: Union[ModuleSpec, type] = None + shared_experts: Union[ModuleSpec, type] = None + + +class ErnieMultiTypeMoE(MegatronModule): + """Dual-pool Mixture of Experts layer for ERNIE 4.5 VL. + + Routes text tokens to text_moe_layer and vision tokens to vision_moe_layer, + then combines outputs with shared expert output. + + Each pool is a standard Megatron MoELayer with its own router and experts, + supporting TP and EP parallelism natively. + + Args: + config: TransformerConfig with moe_intermediate_size as a tuple/list + of [text_ffn_size, vision_ffn_size]. + submodules: MultiTypeMoeSubmodules containing specs for both pools. + layer_number: Layer index in the transformer stack. + pg_collection: Process group collection for parallelism. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: Optional[MultiTypeMoeSubmodules] = None, + layer_number: Optional[int] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + ): + super().__init__(config=config) + self.layer_number = layer_number + + # TransformerLayer only passes pg_collection to known MLP types (MoELayer, + # TEGroupedMLP, SequentialMLP). ErnieMultiTypeMoE is not in that list, so + # pg_collection may be None. Fall back to default MoE process groups. + if pg_collection is None: + pg_collection = get_default_pg_collection() + + # Create separate configs for each pool with different FFN sizes + self.text_config = deepcopy(config) + self.vision_config = deepcopy(config) + self.text_config.moe_ffn_hidden_size = config.moe_intermediate_size[0] + self.vision_config.moe_ffn_hidden_size = config.moe_intermediate_size[1] + + # Disable shared experts within each MoELayer since ErnieMultiTypeMoE + # manages its own shared_experts externally (not inside the per-pool MoELayer). + self.text_config.moe_shared_expert_intermediate_size = None + self.vision_config.moe_shared_expert_intermediate_size = None + + # Build the two MoE pools and shared experts + self.text_moe_layer = build_module(submodules.text_moe_layer, self.text_config) + self.vision_moe_layer = build_module(submodules.vision_moe_layer, self.vision_config) + self.shared_experts = build_module( + submodules.shared_experts, + config=config, + pg_collection=pg_collection, + gate=False, + ) + + def forward( + self, + hidden_states: torch.Tensor, + token_type_ids: torch.Tensor = None, + padding_mask: Optional[torch.Tensor] = None, + ): + """Forward pass for dual-pool MoE. + + Args: + hidden_states: Input tensor [seq_len, batch, hidden_size]. + When Sequence Parallel (SP) is enabled, seq_len is the local + partition size (full_seq_len / tp_size). + token_type_ids: Modality indicator [batch, seq_len]. + 0 = text token -> text_moe_layer + 1 or 2 = vision token -> vision_moe_layer + When SP is enabled, this must already be sliced to match the + local sequence partition (done by Ernie45VLModel.forward()). + padding_mask: Optional padding mask [batch, seq_len] passed by + Megatron's TransformerLayer. Forwarded to each MoE pool's + router for filtering out padding tokens during routing. + + Returns: + Tuple of (output, bias). bias is always None. + """ + if token_type_ids is None: + # Read from the module-level context variable set by + # Ernie45VLModel.forward() before language_model.forward(). + token_type_ids = _current_moe_mm_token_type_ids + + if token_type_ids is None: + # Ultimate fallback: all text tokens (no vision routing) + token_type_ids = torch.zeros( + hidden_states.shape[1], + hidden_states.shape[0], + dtype=torch.long, + device=hidden_states.device, + ) + + # hidden_states: [seq_len, batch, hidden_size] + # token_type_ids: [batch, seq_len] (0=text, >=1 = vision) + seq_len, batch_size, hidden_size = hidden_states.shape + + # Flatten batch and seq dims to create a flat token stream. + # This supports batch_size > 1 correctly, whereas the old code + # assumed batch=1 via squeeze(0). + # flat_hidden: [seq_len * batch, 1, hidden_size] + # We reshape to [N, 1, H] so that each MoELayer still sees + # a 3D tensor with batch_dim=1 (required by MoELayer internals). + flat_hidden = hidden_states.permute(1, 0, 2).reshape( + batch_size * seq_len, 1, hidden_size + ) # [batch * seq_len, 1, hidden] + + # Flatten token_type_ids: [batch, seq_len] -> [batch * seq_len] + flat_type_ids = token_type_ids.reshape(-1) # [batch * seq_len] + + # Build per-token modality mask. True = vision token. + vision_mask = flat_type_ids.bool() # >=1 means vision + + # ---------- Filter tokens by modality BEFORE routing ---------- + # This matches HF behaviour: text_moe only sees text tokens, + # vision_moe only sees vision tokens. The routers and expert + # computation never touch wrong-modality tokens. + text_indices = (~vision_mask).nonzero(as_tuple=True)[0] + vision_indices = vision_mask.nonzero(as_tuple=True)[0] + + text_hidden = flat_hidden[text_indices] # [N_text, 1, hidden] + vision_hidden = flat_hidden[vision_indices] # [N_vision, 1, hidden] + + # Route filtered tokens through their respective MoE pools. + # + # When EP > 1, both pools must ALWAYS be called on all EP ranks, + # even if N_text=0 or N_vision=0 on this rank, because each + # MoELayer's token_dispatcher performs alltoall/allgather + # collectives that require all EP ranks to participate. + # + # MCore's RouterGatingLinearFunction cannot reshape [0]-element + # tensors (shape [0, 1, -1]). When a pool has 0 tokens but EP > 1 + # requires collective participation, we inject a single dummy token, + # run the MoE forward (so collectives execute), then discard the + # dummy output. + # + # When EP == 1 (no expert parallelism), we can safely skip the + # empty pool since there are no inter-rank collectives. + ep_size = parallel_state.get_expert_model_parallel_world_size() + + if text_indices.numel() > 0: + text_output, _ = self.text_moe_layer(text_hidden, padding_mask=None) + elif ep_size > 1: + # Inject a dummy token to participate in EP collectives + dummy = torch.zeros(1, 1, hidden_size, dtype=flat_hidden.dtype, device=flat_hidden.device) + _, _ = self.text_moe_layer(dummy, padding_mask=None) + text_output = None + else: + text_output = None + + if vision_indices.numel() > 0: + vision_output, _ = self.vision_moe_layer(vision_hidden, padding_mask=None) + elif ep_size > 1: + # Inject a dummy token to participate in EP collectives + dummy = torch.zeros(1, 1, hidden_size, dtype=flat_hidden.dtype, device=flat_hidden.device) + _, _ = self.vision_moe_layer(dummy, padding_mask=None) + vision_output = None + else: + vision_output = None + + # Scatter back to full flat sequence + flat_moe_output = torch.zeros_like(flat_hidden) + if text_output is not None and text_indices.numel() > 0: + flat_moe_output[text_indices] = text_output + if vision_output is not None and vision_indices.numel() > 0: + flat_moe_output[vision_indices] = vision_output + + # Reshape back to [seq_len, batch, hidden] from [batch * seq_len, 1, hidden] + moe_output = flat_moe_output.reshape(batch_size, seq_len, hidden_size).permute(1, 0, 2).contiguous() + + # Shared experts see ALL tokens (same as HF) + # SharedExpertMLP.forward() returns a single tensor (not a tuple), + # unlike MLP.forward() which returns (output, output_bias). + shared_result = self.shared_experts(hidden_states) + shared_output = shared_result[0] if isinstance(shared_result, tuple) else shared_result + + moe_output = moe_output + shared_output + + return moe_output, None + + def set_layer_number(self, layer_number: int): + """Set the layer number for both MoE pools.""" + self.layer_number = layer_number + if hasattr(self.text_moe_layer, "router"): + self.text_moe_layer.router.set_layer_number(layer_number) + if hasattr(self.vision_moe_layer, "router"): + self.vision_moe_layer.router.set_layer_number(layer_number) diff --git a/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/model.py b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/model.py new file mode 100644 index 0000000000..38046154b9 --- /dev/null +++ b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/model.py @@ -0,0 +1,608 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Megatron-Core compatible ERNIE 4.5 VL MoE model. + +This module wraps the HuggingFace ERNIE 4.5 VL MoE vision encoder and resampler +with a Megatron-Core GPT language model to create a distributable VLM. + +Architecture: + - Vision Tower: Ernie4_5_VLMoeVisionTransformerPretrainedModel (HF, replicated across TP) + - Resampler: Ernie4_5_VLMoeVariableResolutionResamplerModel (HF, replicated across TP) + - Language Model: MCoreGPTModel (Megatron-Core, distributed across TP/PP/EP) + with custom ErnieMultiTypeMoE layers supporting dual-pool MoE: + * text_moe_layer: 64 experts (FFN=1536) for text tokens + * vision_moe_layer: 64 experts (FFN=512) for vision tokens + * shared_experts: shared MLP for all tokens +""" + +import types +from typing import Optional + +import torch +from megatron.core import parallel_state +from megatron.core.inference.contexts import BaseInferenceContext +from megatron.core.models.common.embeddings.rotary_pos_embedding import ( + MultimodalRotaryEmbedding, + get_pos_emb_on_this_cp_rank, +) +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.tensor_parallel import scatter_to_sequence_parallel_region +from megatron.core.transformer.module import MegatronModule +from torch import Tensor +from transformers.models.ernie4_5_vl_moe.modeling_ernie4_5_vl_moe import ( + Ernie4_5_VLMoeModel, + Ernie4_5_VLMoeVariableResolutionResamplerModel, + Ernie4_5_VLMoeVisionTransformerPretrainedModel, +) + +from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.ernie_moe_layer import ( + clear_moe_mm_token_type_ids, + set_moe_mm_token_type_ids, +) +from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.vision_layer_spec import get_ernie_vit_layer_spec +from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.vision_model import ErnieVLVisionModel +from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.vision_transformer_config import get_ernie_vision_config +from megatron.bridge.models.gpt_provider import GPTModelProvider +from megatron.bridge.utils.common_utils import hook_hf_module_setattr_for_tp_grad_sync + + +def _normalize_hf_config(hf_config): + """Ensure the HF config has a ``text_config`` attribute. + + The transformers-builtin ``Ernie4_5_VLMoeVariableResolutionResamplerModel`` accesses + ``config.text_config.hidden_size`` and ``config.text_config.rms_norm_eps``. The + nested config (Instruct model) has ``text_config`` as a sub-object, but the flat + config (Thinking model) stores all LLM fields directly on the top-level config. + + For flat configs, we set ``text_config`` to point to the config itself so that + ``config.text_config.hidden_size`` resolves to ``config.hidden_size``. + """ + if hf_config is not None and not hasattr(hf_config, "text_config"): + hf_config.text_config = hf_config + return hf_config + + +def _normalize_vision_config(vision_config, hf_config=None): + """Ensure the vision config has all attributes required by the transformers-builtin + vision model classes (Ernie4_5_VLMoeVisionBlock, Ernie4_5_VLMoeVisionTransformerPretrainedModel, + Ernie4_5_VLMoeVariableResolutionResamplerModel). + + The Thinking model's custom ``DFNRopeVisionTransformerConfig`` (auto_map) uses + ``mlp_ratio`` + ``embed_dim`` instead of ``intermediate_size``, omits ``rms_norm_eps``, + and omits ``temporal_merge_size``. This function adds the missing attributes so + the same config object works with the transformers-builtin vision model code. + """ + # rms_norm_eps: used by LayerNorm in vision blocks and final LN + if not hasattr(vision_config, "rms_norm_eps"): + # Prefer the top-level config's rms_norm_eps (LLM side), default 1e-6 + rms_norm_eps = 1e-6 + if hf_config is not None: + rms_norm_eps = getattr(hf_config, "rms_norm_eps", 1e-6) + vision_config.rms_norm_eps = rms_norm_eps + + # intermediate_size: used by vision MLP (= mlp_ratio * embed_dim) + if not hasattr(vision_config, "intermediate_size"): + mlp_ratio = getattr(vision_config, "mlp_ratio", 4) + embed_dim = getattr(vision_config, "embed_dim", getattr(vision_config, "hidden_size", 1280)) + vision_config.intermediate_size = int(mlp_ratio * embed_dim) + + # temporal_merge_size: used by resampler for temporal pooling + if not hasattr(vision_config, "temporal_merge_size"): + vision_config.temporal_merge_size = getattr(vision_config, "spatial_merge_size", 2) + + return vision_config + + +class _MgVitTowerAdapter(torch.nn.Module): + """Thin adapter that makes the MG-native ErnieVLVisionModel compatible + with the HF-bound ``get_image_features`` / ``get_video_features`` methods. + + The HF methods call ``self.vision_tower(pixel_values, grid_thw, return_dict=True)`` + and expect a ``BaseModelOutputWithPooling`` with ``.last_hidden_state``. They also + access ``self.vision_tower.spatial_merge_size``. + + This adapter wraps ``ErnieVLVisionModel`` to match that interface exactly. + """ + + def __init__(self, mg_vision_model: "ErnieVLVisionModel"): + super().__init__() + self.mg_vision_model = mg_vision_model + self.spatial_merge_size = mg_vision_model.spatial_merge_size + + def forward(self, pixel_values, grid_thw, return_dict=True, **kwargs): + from transformers.modeling_outputs import BaseModelOutputWithPooling + + hidden_states = self.mg_vision_model(pixel_values, grid_thw) + return BaseModelOutputWithPooling(last_hidden_state=hidden_states) + + +class ErnieMultimodalRotaryEmbedding(MultimodalRotaryEmbedding): + """ERNIE-specific 3D M-RoPE with interleaved H/W frequency allocation. + + ERNIE 4.5 VL uses a custom RoPE layout that differs from the standard + Qwen2VL-style contiguous block layout used by ``MultimodalRotaryEmbedding``. + + Standard (Qwen2VL) layout with mrope_section=[22, 22, 20]: + head dims [0:44] -> T (temporal) axis, freq bands 0-21 + head dims [44:88] -> H (height) axis, freq bands 0-21 + head dims [88:128] -> W (width) axis, freq bands 0-19 + + ERNIE layout with freq_allocation=20: + head dims [0:44] -> H,W interleaved: even freq bands -> H, odd -> W + head dims [44:88] -> (same interleaving continues) + head dims [88:128] -> T (temporal) axis, freq bands 44-63 + + More precisely, for freq band index f (0..63): + f in {0,2,4,...,42} (even, f<44) -> H position + f in {1,3,5,...,43} (odd, f<44) -> W position + f in {44,45,...,63} (last 20) -> T position + + For text tokens (T=H=W=p), both layouts produce identical results since + all axes have the same position value. The difference only manifests for + image/video tokens where T, H, W have distinct values. + + This subclass overrides ``forward()`` to implement ERNIE's interleaved + layout while reusing the parent's ``inv_freq`` and infrastructure. + """ + + def __init__(self, freq_allocation: int = 20, **kwargs): + super().__init__(**kwargs) + self.freq_allocation = freq_allocation + + def forward( + self, + position_ids: torch.Tensor, + mrope_section, + cp_group=None, + ) -> Tensor: + """Compute ERNIE-style interleaved M-RoPE embeddings. + + Args: + position_ids: [3, batch, seq_len] where axis 0=T, 1=H, 2=W + mrope_section: Ignored (kept for API compatibility). ERNIE uses + freq_allocation instead. + cp_group: Context parallel group. + + Returns: + Tensor: RoPE embedding of shape [seq_len, batch, 1, head_dim]. + """ + device = self.inv_freq.device + seq = position_ids.to(device=device, dtype=self.inv_freq.dtype) + # seq: [3, bs, seq_len] + + if self.seq_len_interpolation_factor is not None: + seq = seq * (1.0 / self.seq_len_interpolation_factor) + + bs = seq.shape[1] + seq_len = seq.shape[2] + num_freqs = self.inv_freq.shape[0] # head_dim // 2 = 64 + + # Compute freqs for each axis: freqs[axis, bs, seq_len, num_freqs] + # Each axis independently: theta_f = inv_freq[f] * pos[axis] + inv_freq_exp = self.inv_freq[None, None, :, None].expand(3, bs, -1, 1) + seq_exp = seq[:, :, None, :].float() + freqs = (inv_freq_exp @ seq_exp).transpose(2, 3) + # freqs: [3, bs, seq_len, num_freqs=64] + + # ERNIE interleaved layout: + # For freq band f (0-indexed into inv_freq): + # f < (num_freqs - freq_allocation) AND f is even -> H (axis 1) + # f < (num_freqs - freq_allocation) AND f is odd -> W (axis 2) + # f >= (num_freqs - freq_allocation) -> T (axis 0) + # + # Build the combined freq tensor by selecting the right axis per band. + hw_bands = num_freqs - self.freq_allocation # 44 + # H bands: even indices 0,2,4,...,42 -> 22 bands + h_freq_indices = torch.arange(0, hw_bands, 2, device=device) + # W bands: odd indices 1,3,5,...,43 -> 22 bands + w_freq_indices = torch.arange(1, hw_bands, 2, device=device) + # T bands: last freq_allocation indices 44,...,63 -> 20 bands + t_freq_indices = torch.arange(hw_bands, num_freqs, device=device) + + # Gather per-axis freqs for their assigned bands + # freqs[axis]: [bs, seq_len, 64] + h_freqs = freqs[1, :, :, h_freq_indices] # [bs, seq_len, 22] + w_freqs = freqs[2, :, :, w_freq_indices] # [bs, seq_len, 22] + t_freqs = freqs[0, :, :, t_freq_indices] # [bs, seq_len, 20] + + # Interleave H and W: [H0, W0, H1, W1, ..., H21, W21] -> 44 values + hw_interleaved = torch.stack([h_freqs, w_freqs], dim=-1).reshape(bs, seq_len, hw_bands) # [bs, seq_len, 44] + + # Concatenate HW + T + combined_freqs = torch.cat([hw_interleaved, t_freqs], dim=-1) + # combined_freqs: [bs, seq_len, 64] + + # Apply interleaved doubling (matching rotary_interleaved=True): + # Each freq band f expands to two consecutive head dims: [f, f] + combined_flat = combined_freqs.reshape(bs, -1, 1) + emb = torch.stack((combined_flat, combined_flat), dim=-1).reshape(bs, seq_len, -1) + # emb: [bs, seq_len, 128] + + # Reshape to match MCore expected output: [seq_len, bs, 1, head_dim] + emb = emb[..., None, :].transpose(0, 1).contiguous() + # emb: [seq_len, bs, 1, 128] + + if cp_group is None: + cp_group = self.cp_group + if cp_group is not None and cp_group.size() > 1: + emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group) + return emb + + +class Ernie45VLModel(MegatronModule): + """ + ERNIE 4.5 VL MoE Model (Vision-Language with Mixture of Experts). + + This model combines: + - A HuggingFace ERNIE 4.5 vision encoder (32-layer ViT with 2D RoPE) + - A variable-resolution resampler (spatial + temporal merging) + - A Megatron-Core GPT language model with heterogeneous dual-pool MoE + + The vision tower and resampler are borrowed directly from HuggingFace + and replicated across TP ranks. The language model uses standard + Megatron-Core distributed infrastructure. + + Args: + config (GPTModelProvider): Language model provider configuration. + pre_process (bool): Include embedding layer (used with pipeline parallelism). + post_process (bool): Include output layer (used with pipeline parallelism). + vp_stage (int, optional): Virtual pipeline stage index. + """ + + def __init__( + self, + config: GPTModelProvider, + pre_process: bool = True, + post_process: bool = True, + vp_stage: Optional[int] = None, + ) -> None: + super().__init__(config=config) + + self.pre_process = pre_process + self.post_process = post_process + self.vp_stage = vp_stage + + # HF bound methods (get_image_features, get_video_features, etc.) access + # self.config.return_dict via the @can_return_tuple decorator. + # Ensure the attribute exists on the provider config. + if not hasattr(config, "return_dict"): + config.return_dict = True + + self.use_mg_vit = getattr(config, "use_mg_vit", False) + + if pre_process: + # Normalize configs for compatibility with transformers-builtin + # vision model and resampler classes. + # 1. Vision config: DFNRopeVisionTransformerConfig may lack rms_norm_eps, + # intermediate_size, temporal_merge_size. + _normalize_vision_config(config.vision_config, hf_config=config.hf_config) + # 2. HF config: flat config (Thinking) lacks text_config sub-object + # that the resampler accesses as config.text_config.hidden_size, etc. + _normalize_hf_config(config.hf_config) + + if self.use_mg_vit: + # Megatron-Core native ViT: TP-native attention and MLP via + # TransformerBlock with TE modules. + vision_transformer_config = get_ernie_vision_config( + config.vision_config, + megatron_config=config, + ) + vision_layer_spec = get_ernie_vit_layer_spec() + self.vision_model = ErnieVLVisionModel( + transformer_config=vision_transformer_config, + transformer_layer_spec=vision_layer_spec, + ) + # Wrap MG ViT with an adapter that matches the HF + # vision_tower interface (forward signature, spatial_merge_size + # attribute) so the bound HF get_image/video_features methods + # work transparently. + self.vision_tower = _MgVitTowerAdapter(self.vision_model) + else: + # HF-wrapped ViT: replicated across TP ranks. + self.vision_tower = Ernie4_5_VLMoeVisionTransformerPretrainedModel._from_config(config.vision_config) + # Ensure HF vision tower params are tracked for TP gradient sync + hook_hf_module_setattr_for_tp_grad_sync(self.vision_tower) + + # Instantiate the HF resampler (spatial + temporal merging + projection). + # The resampler is kept as an HF module regardless of use_mg_vit because + # it is small, replicated, and already has bridge weight mappings. + self.resampler_model = Ernie4_5_VLMoeVariableResolutionResamplerModel(config.hf_config) + hook_hf_module_setattr_for_tp_grad_sync(self.resampler_model) + + # Build the Megatron-Core GPT language model + self.language_model = self.config.provide_language_model( + pre_process=pre_process, post_process=post_process, vp_stage=vp_stage + ) + + # Replace the default MultimodalRotaryEmbedding with ERNIE's custom + # interleaved variant. GPTModel.__init__ creates rotary_pos_emb as a + # standard MultimodalRotaryEmbedding, but ERNIE 4.5 VL uses a non-standard + # interleaved H/W frequency layout (see ErnieMultimodalRotaryEmbedding). + if hasattr(self.language_model, "rotary_pos_emb") and isinstance( + self.language_model.rotary_pos_emb, MultimodalRotaryEmbedding + ): + old_rope = self.language_model.rotary_pos_emb + freq_allocation = getattr(config.hf_config, "freq_allocation", None) + if freq_allocation is None: + # Derive from mrope_section: last element is T (temporal) allocation. + # Real model: mrope_section=[22,22,20] -> freq_allocation=20 + # Toy model: mrope_section=[2,2,2] -> freq_allocation=2 + mrope_section = getattr(config, "mrope_section", None) + freq_allocation = mrope_section[-1] if mrope_section else 20 + self.language_model.rotary_pos_emb = ErnieMultimodalRotaryEmbedding( + freq_allocation=freq_allocation, + kv_channels=config.kv_channels, + rotary_percent=1.0, + rotary_interleaved=config.rotary_interleaved, + seq_len_interpolation_factor=old_rope.seq_len_interpolation_factor, + rotary_base=config.rotary_base, + ) + + # Required for finalize_model_grads and tied weights + self.share_embeddings_and_output_weights = config.share_embeddings_and_output_weights + self.shared_embedding_or_output_weight = self.language_model.shared_embedding_or_output_weight + + # Bind utility methods from HF Ernie4_5_VLMoeModel to this instance + self.get_placeholder_mask = types.MethodType(Ernie4_5_VLMoeModel.get_placeholder_mask, self) + self.get_image_features = types.MethodType(Ernie4_5_VLMoeModel.get_image_features, self) + self.get_video_features = types.MethodType(Ernie4_5_VLMoeModel.get_video_features, self) + self.get_rope_index = types.MethodType(Ernie4_5_VLMoeModel.get_rope_index, self) + self.get_vision_position_ids = types.MethodType(Ernie4_5_VLMoeModel.get_vision_position_ids, self) + + if pre_process: + # Register pixel normalization buffers for the vision encoder. + # The ERNIE 4.5 VL processor outputs raw (unnormalized) pixel patches + # with do_rescale=False, do_normalize=False. Normalization is expected + # to happen on-device before the ViT, matching the HF custom model's + # vision_forward() + add_image_preprocess() logic. + # + # OPENAI_CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073] + # OPENAI_CLIP_STD = [0.26862954, 0.26130258, 0.27577711] + patch_size = getattr(config.vision_config, "patch_size", 14) + pixels_per_patch = patch_size * patch_size # 196 for patch_size=14 + + clip_mean = torch.tensor([0.48145466, 0.4578275, 0.40821073], dtype=torch.float32) + clip_std = torch.tensor([0.26862954, 0.26130258, 0.27577711], dtype=torch.float32) + + # Expand to match flattened patch layout: [C * patch_size^2] = [588] + # Each channel's mean/std is repeated patch_size^2 times + pixel_mean = clip_mean.repeat_interleave(pixels_per_patch) # [588] + pixel_std = clip_std.repeat_interleave(pixels_per_patch) # [588] + + self.register_buffer("pixel_mean", pixel_mean, persistent=False) + self.register_buffer("pixel_std", pixel_std, persistent=False) + + @property + def decoder(self): + """Expose language model decoder for mcore inference compatibility.""" + return getattr(self.language_model, "decoder", None) + + def set_input_tensor(self, input_tensor) -> None: + """Set model chunk input tensor.""" + self.language_model.set_input_tensor(input_tensor) + + def _normalize_pixel_values(self, pixel_values: torch.Tensor) -> torch.Tensor: + """Normalize raw pixel patches for the vision encoder. + + The ERNIE 4.5 VL processor outputs raw pixel patches (0-255 range, + ``do_rescale=False, do_normalize=False``). This method applies CLIP + normalization on-device, matching the HF custom model's + ``vision_forward()`` + ``add_image_preprocess()`` logic: + + pixel_values = pixel_values / 255.0 + pixel_values = (pixel_values - CLIP_MEAN) / CLIP_STD + + Args: + pixel_values: Raw pixel patches [total_patches, C*patch_size^2]. + Values in 0-255 range (any dtype). + + Returns: + Normalized pixel patches in bfloat16, values in ~(-2, 2.5) range. + """ + # Rescale: divide by 255 (in float32 for precision) + pixel_values = pixel_values.to(torch.float32) * (1.0 / 255.0) + # Normalize: (x - mean) / std using CLIP mean/std + pixel_values = (pixel_values - self.pixel_mean.to(pixel_values.device)) / self.pixel_std.to( + pixel_values.device + ) + # Cast to bfloat16 for the ViT + return pixel_values.to(torch.bfloat16) + + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + pixel_values: Optional[torch.Tensor] = None, + pixel_values_videos: Optional[torch.FloatTensor] = None, + image_grid_thw: Optional[torch.LongTensor] = None, + video_grid_thw: Optional[torch.LongTensor] = None, + mm_token_type_ids: Optional[torch.IntTensor] = None, + moe_mm_token_type_ids: Optional[torch.IntTensor] = None, + labels: Tensor = None, + inference_context: BaseInferenceContext = None, + packed_seq_params: PackedSeqParams = None, + extra_block_kwargs: dict = None, + runtime_gather_output: Optional[bool] = None, + *, + inference_params: Optional[BaseInferenceContext] = None, + loss_mask: Optional[Tensor] = None, + ) -> Tensor: + r""" + Forward pass for ERNIE 4.5 VL MoE. + + Args: + input_ids: Token IDs [batch_size, seq_len]. + pixel_values: Image pixel values for the vision encoder. + pixel_values_videos: Video pixel values for the vision encoder. + image_grid_thw: Grid dimensions (T, H, W) per image [num_images, 3]. + video_grid_thw: Grid dimensions (T, H, W) per video [num_videos, 3]. + mm_token_type_ids: Token type IDs for M-RoPE computation (0=text, 1=image, 2=video). + moe_mm_token_type_ids: Token type IDs for MoE routing (0=text, 1/2=vision). + labels: Labels for language modeling loss. + loss_mask: Mask for loss computation. + """ + + if self.pre_process: + if inputs_embeds is None: + # Get text embeddings from language model embedding layer + inputs_embeds = self.language_model.embedding( + input_ids=input_ids, position_ids=None + ) # [seq_len, batch, hidden] + + inputs_embeds = inputs_embeds.transpose(1, 0).contiguous() # [batch, seq_len, hidden] + + # Process images through vision tower + resampler + if pixel_values is not None: + # Normalize raw pixel patches from the processor (0-255 uint8-valued). + # The custom ERNIE processor outputs do_rescale=False, do_normalize=False, + # so normalization must happen on-device before the ViT. + pixel_values = self._normalize_pixel_values(pixel_values) + image_embeds = self.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output + image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) + image_mask, _ = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, image_features=image_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + + # Process videos through vision tower + resampler + if pixel_values_videos is not None: + pixel_values_videos = self._normalize_pixel_values(pixel_values_videos) + video_embeds = self.get_video_features( + pixel_values_videos, video_grid_thw, return_dict=True + ).pooler_output + video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype) + _, video_mask = self.get_placeholder_mask( + input_ids, inputs_embeds=inputs_embeds, video_features=video_embeds + ) + inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + + # Transpose back to [seq_len, batch, hidden] for Megatron-Core + inputs_embeds = inputs_embeds.transpose(1, 0) + + if self.config.sequence_parallel: + inputs_embeds = scatter_to_sequence_parallel_region(inputs_embeds) + + # Compute 3D MRoPE position IDs on ALL pipeline stages + # Each stage has input_ids and visual grid info from the data iterator + # + # The custom ERNIE processor marks IMAGE_START/IMAGE_END/VIDEO_START/VIDEO_END + # boundary tokens as image/video type in token_type_ids (for MoE routing). + # But get_rope_index expects these boundary tokens to be text type (0), + # because it generates exactly T*H*W/merge^2 vision positions per image + # group from grid_thw — boundary tokens should get sequential text positions. + rope_mm_token_type_ids = mm_token_type_ids + if mm_token_type_ids is not None and input_ids is not None: + boundary_token_ids = [ + getattr(self.config, "image_start_token_id", 101304), + getattr(self.config, "image_end_token_id", 101305), + getattr(self.config, "video_start_token_id", 101306), + getattr(self.config, "video_end_token_id", 101307), + ] + boundary_mask = torch.zeros_like(input_ids, dtype=torch.bool) + for tid in boundary_token_ids: + boundary_mask |= input_ids == tid + if boundary_mask.any(): + rope_mm_token_type_ids = mm_token_type_ids.clone() + rope_mm_token_type_ids[boundary_mask] = 0 + + position_ids, _rope_deltas = self.get_rope_index( + input_ids, + rope_mm_token_type_ids, + image_grid_thw, + video_grid_thw, + attention_mask=None, + ) + + # Set moe_mm_token_type_ids in the module-level context so that + # ErnieMultiTypeMoE layers can read it during forward. This avoids + # modifying Megatron-Core's TransformerBlock/TransformerLayer signatures. + # + # When Sequence Parallel (SP) is enabled, hidden_states entering each + # TransformerLayer's MLP are scattered across TP ranks: + # hidden_states shape = [seq_len / tp_size, batch, hidden] + # The moe_mm_token_type_ids must be sliced to match the local sequence + # partition so that ErnieMultiTypeMoE sees the correct modality labels + # for the tokens on this TP rank. + sp_moe_mm_token_type_ids = moe_mm_token_type_ids + if ( + moe_mm_token_type_ids is not None + and self.config.sequence_parallel + and parallel_state.get_tensor_model_parallel_world_size() > 1 + ): + tp_size = parallel_state.get_tensor_model_parallel_world_size() + tp_rank = parallel_state.get_tensor_model_parallel_rank() + full_seq_len = moe_mm_token_type_ids.shape[-1] + local_seq_len = full_seq_len // tp_size + start = tp_rank * local_seq_len + end = start + local_seq_len + sp_moe_mm_token_type_ids = moe_mm_token_type_ids[..., start:end].contiguous() + + set_moe_mm_token_type_ids(sp_moe_mm_token_type_ids) + + try: + outputs = self.language_model.forward( + input_ids=None, + position_ids=position_ids, + attention_mask=attention_mask, + decoder_input=inputs_embeds, + labels=labels, + loss_mask=loss_mask, + runtime_gather_output=runtime_gather_output, + packed_seq_params=packed_seq_params, + ) + finally: + # Always clear after forward to avoid holding stale references + clear_moe_mm_token_type_ids() + + return outputs + + def freeze( + self, + freeze_language_model: bool, + freeze_vision_model: bool, + freeze_vision_projection: bool, + ): + """Freeze model modules. + + Args: + freeze_language_model: Freeze the language model module. + freeze_vision_model: Freeze the vision encoder (patch_embed + blocks). + freeze_vision_projection: Freeze the resampler / projector. + """ + modules = [] + + if freeze_language_model and hasattr(self, "language_model") and self.language_model is not None: + modules.append(self.language_model) + + if freeze_vision_model: + if hasattr(self, "vision_model") and self.vision_model is not None: + # MG-native ViT: freeze the entire vision model + modules.append(self.vision_model) + elif hasattr(self, "vision_tower") and self.vision_tower is not None: + # HF ViT: freeze patch_embed and blocks + if hasattr(self.vision_tower, "patch_embed"): + modules.append(self.vision_tower.patch_embed) + if hasattr(self.vision_tower, "blocks"): + modules.append(self.vision_tower.blocks) + + if freeze_vision_projection and hasattr(self, "resampler_model") and self.resampler_model is not None: + modules.append(self.resampler_model) + + for module in modules: + for param in module.parameters(): + param.requires_grad = False diff --git a/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_attention.py b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_attention.py new file mode 100644 index 0000000000..6e3c819952 --- /dev/null +++ b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_attention.py @@ -0,0 +1,248 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Custom SelfAttention for ERNIE 4.5 VL vision encoder. + +Overrides the standard MCore SelfAttention to apply absolute 2D RoPE +positional embeddings instead of the standard relative RoPE. + +ERNIE ViT uses non-interleaved RoPE (rotate_half style, splitting at +the midpoint: [-x2, x1]), corresponding to ``rotary_interleaved=False`` +in MCore. The RoPE frequencies are pre-computed as absolute position +embeddings based on 2D (height, width) grid coordinates. + +This approach mirrors Qwen3VLSelfAttention but with ERNIE-specific +non-interleaved rotation. +""" + +from typing import Optional, Tuple, Union + +from megatron.core.models.common.embeddings.rope_utils import _apply_rotary_pos_emb_bshd +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.transformer.attention import ( + BaseInferenceContext, + SelfAttention, + deprecate_inference_params, + nvtx_range_pop, + nvtx_range_push, +) +from torch import Tensor + + +def _apply_rotary_pos_emb_thd_absolute( + t: Tensor, cu_seqlens: Tensor, freqs: Tensor, rotary_interleaved: bool = False +) -> Tensor: + """Apply RoPE to ``thd`` (packed) format tensors using absolute position embeddings. + + Args: + t: Input tensor of shape [total_tokens, num_heads, head_dim]. + cu_seqlens: Cumulative sequence lengths (currently unused, kept for API consistency). + freqs: Rotary embedding frequencies of shape [total_tokens, 1, 1, head_dim]. + rotary_interleaved: Whether to use interleaved rotation. + + Returns: + Tensor of shape [total_tokens, num_heads, head_dim] with RoPE applied. + """ + # Unsqueeze to [total_tokens, 1, num_heads, head_dim] for bshd RoPE, then squeeze back + return _apply_rotary_pos_emb_bshd(t[:, None], freqs, rotary_interleaved=rotary_interleaved).squeeze(1) + + +def apply_rotary_pos_emb_absolute( + t: Tensor, + freqs: Tensor, + config, + cu_seqlens: Optional[Tensor] = None, +) -> Tensor: + """Apply absolute RoPE, routing to bshd or thd format as appropriate. + + For ERNIE ViT, the freqs tensor has shape [total_tokens, 1, 1, head_dim] + (absolute position embeddings, where the raw frequencies of shape + [head_dim//2] are tiled 2x to cover the full head_dim), unlike standard + relative RoPE where freqs is [max_seqlen, 1, 1, rotary_dim]. + + Args: + t: Input tensor (Q or K). + freqs: Pre-computed RoPE frequencies. + config: TransformerConfig (used for rotary_interleaved flag). + cu_seqlens: If provided, indicates packed sequence (thd) format. + + Returns: + Tensor with RoPE applied, same shape as input. + """ + orig_dtype = t.dtype + # Compute RoPE in fp32 for numerical stability + t = t.float() + + if cu_seqlens is None: + result = _apply_rotary_pos_emb_bshd(t, freqs, rotary_interleaved=config.rotary_interleaved) + else: + result = _apply_rotary_pos_emb_thd_absolute(t, cu_seqlens, freqs, rotary_interleaved=config.rotary_interleaved) + + return result.to(orig_dtype) + + +class ErnieVLSelfAttention(SelfAttention): + """SelfAttention with absolute 2D RoPE for ERNIE ViT. + + Overrides the standard MCore SelfAttention.forward() to apply + ``apply_rotary_pos_emb_absolute`` instead of the standard + ``apply_rotary_pos_emb`` which expects relative position embeddings. + + This is necessary because ERNIE ViT pre-computes absolute 2D (H, W) + position embeddings and passes them as rotary_pos_emb through the + TransformerBlock, rather than using the standard MCore RoPE infrastructure + that computes frequencies from sequential position IDs. + """ + + def forward( + self, + hidden_states: Tensor, + attention_mask: Tensor, + key_value_states: Optional[Tensor] = None, + inference_context: Optional[BaseInferenceContext] = None, + rotary_pos_emb: Optional[Union[Tensor, Tuple[Tensor, Tensor]]] = None, + rotary_pos_cos: Optional[Tensor] = None, + rotary_pos_sin: Optional[Tensor] = None, + attention_bias: Optional[Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + sequence_len_offset: Optional[int] = None, + *, + inference_params: Optional[BaseInferenceContext] = None, + rotary_pos_cos_sin: Optional[Tensor] = None, + ) -> Tuple[Tensor, Tensor]: + """Forward pass with absolute 2D RoPE for vision encoder. + + The main difference from the parent class is in the RoPE application + section: we use ``apply_rotary_pos_emb_absolute`` which handles + absolute position embeddings properly for both bshd and thd formats. + + Args: + hidden_states: Input tensor [seq_len, batch, hidden_size]. + attention_mask: Attention mask (typically None for ViT). + rotary_pos_emb: Pre-computed absolute 2D RoPE frequencies. + packed_seq_params: Parameters for per-image packed sequence attention. + (other args): See parent class SelfAttention. + + Returns: + Tuple of (output, bias) where output is [seq_len, batch, hidden_size]. + """ + inference_context = deprecate_inference_params(inference_context, inference_params) + + # For self attention, duplicate the rotary_pos_emb if it isn't already a tuple + if rotary_pos_emb is not None and not isinstance(rotary_pos_emb, tuple): + rotary_pos_emb = (rotary_pos_emb,) * 2 + + # ===================== + # Query, Key, and Value + # ===================== + nvtx_range_push(suffix="qkv") + query, key, value = self.get_query_key_value_tensors(hidden_states, key_value_states) + nvtx_range_pop(suffix="qkv") + + # =================================================== + # Adjust key, value, and rotary_pos_emb for inference + # =================================================== + nvtx_range_push(suffix="adjust_key_value") + query, key, value, rotary_pos_emb, attn_mask_type, _block_table = self._adjust_key_value_for_inference( + inference_context, + query, + key, + value, + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + sequence_len_offset, + ) + + if packed_seq_params is not None: + query = query.squeeze(1) + key = key.squeeze(1) + value = value.squeeze(1) + nvtx_range_pop(suffix="adjust_key_value") + + # ================================================ + # Apply absolute 2D RoPE (the key difference) + # ================================================ + nvtx_range_push(suffix="rotary_pos_emb") + if rotary_pos_emb is not None: + q_pos_emb, k_pos_emb = rotary_pos_emb + + if packed_seq_params is not None: + if packed_seq_params.cu_seqlens_q_padded is not None: + cu_seqlens_q = packed_seq_params.cu_seqlens_q_padded + else: + cu_seqlens_q = packed_seq_params.cu_seqlens_q + if packed_seq_params.cu_seqlens_kv_padded is not None: + cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded + else: + cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + else: + cu_seqlens_q = cu_seqlens_kv = None + + if q_pos_emb is not None: + query = apply_rotary_pos_emb_absolute( + query, + q_pos_emb, + config=self.config, + cu_seqlens=cu_seqlens_q, + ) + if k_pos_emb is not None: + key = apply_rotary_pos_emb_absolute( + key, + k_pos_emb, + config=self.config, + cu_seqlens=cu_seqlens_kv, + ) + nvtx_range_pop(suffix="rotary_pos_emb") + + # ================================== + # Core attention computation + # ================================== + nvtx_range_push(suffix="core_attention") + if self.checkpoint_core_attention and self.training: + core_attn_out = self._checkpointed_attention_forward( + query, + key, + value, + attention_mask, + attn_mask_type=attn_mask_type, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + ) + else: + core_attn_out = self.core_attention( + query, + key, + value, + attention_mask, + attn_mask_type=attn_mask_type, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + ) + + if packed_seq_params is not None and packed_seq_params.qkv_format == "thd": + # Reshape to same output shape as unpacked case + # (t, np, hn) -> (t, b=1, h=np*hn) + core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1) + nvtx_range_pop(suffix="core_attention") + + # ================= + # Output. [sq, b, h] + # ================= + nvtx_range_push(suffix="linear_proj") + output, bias = self.linear_proj(core_attn_out) + nvtx_range_pop(suffix="linear_proj") + + return output, bias diff --git a/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_layer_spec.py b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_layer_spec.py new file mode 100644 index 0000000000..f88962816c --- /dev/null +++ b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_layer_spec.py @@ -0,0 +1,54 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Layer spec for the ERNIE 4.5 VL Megatron-native Vision Transformer (ViT). + +Provides ``get_ernie_vit_layer_spec()`` which returns a ``ModuleSpec`` for a +single ViT transformer layer using Transformer Engine modules. + +The spec is identical to the standard MCore ViT spec from +``megatron.core.models.vision.vit_layer_specs.get_vit_layer_with_transformer_engine_spec`` +except that ``self_attention.module`` is overridden with ``ErnieVLSelfAttention`` +to handle absolute 2D RoPE (non-interleaved rotate_half style). + +Architecture details: + - Attention: TELayerNormColumnParallelLinear (fused QKV + LN) + + TEDotProductAttention + + TERowParallelLinear + - MLP: TELayerNormColumnParallelLinear (fused fc1 + LN) + + TERowParallelLinear + - Mask type: AttnMaskType.no_mask (bidirectional attention for ViT) + - pre_mlp_layernorm: IdentityOp (LN is fused into TE linear layers) +""" + +from megatron.core.models.vision.vit_layer_specs import get_vit_layer_with_transformer_engine_spec + +from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.vision_attention import ErnieVLSelfAttention + + +def get_ernie_vit_layer_spec(): + """Return a TransformerLayer ModuleSpec for ERNIE ViT. + + This reuses the standard MCore ViT TE spec and only overrides the + self-attention module with ``ErnieVLSelfAttention`` to apply absolute + 2D RoPE embeddings instead of the standard relative RoPE. + + Returns: + ModuleSpec: Spec for one ERNIE ViT transformer layer. + """ + spec = get_vit_layer_with_transformer_engine_spec() + # Override self-attention module with ERNIE's absolute 2D RoPE variant + spec.submodules.self_attention.module = ErnieVLSelfAttention + return spec diff --git a/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_model.py b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_model.py new file mode 100644 index 0000000000..dce0f58bb6 --- /dev/null +++ b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_model.py @@ -0,0 +1,343 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Megatron-Core native Vision Transformer for ERNIE 4.5 VL. + +This module implements the ERNIE 4.5 VL DFN-style ViT using Megatron-Core +TransformerBlock infrastructure instead of the HuggingFace implementation. +This enables TP-native attention and MLP layers for better distributed +training performance. + +Architecture (matching HF DFNRopeVisionTransformerPretrainedModel): + - PatchEmbed: nn.Linear(C * P * P, embed_dim, bias=False) + - 2D RoPE: Non-interleaved rotate_half with spatial_merge_size reordering + - 32x TransformerLayer (TE-backed): + - LayerNorm(1280, eps=1e-6) -> QKV(1280, 3*1280, bias=True) -> Attention -> Proj + - LayerNorm(1280, eps=1e-6) -> FC1(1280, 5120) -> quick_gelu -> FC2(5120, 1280) + - Final LayerNorm(1280, eps=1e-6) + - Per-image packed sequence attention via PackedSeqParams (thd format) + +Key differences from Qwen3VL MG ViT: + - PatchEmbed uses nn.Linear (not Conv3d) on pre-flattened patches + - No positional embedding interpolation (ERNIE uses pure 2D RoPE) + - No deepstack feature extraction + - Non-interleaved RoPE (rotate_half style, rotary_interleaved=False) + - No PatchMerger (merging is done by the resampler) +""" + +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from megatron.core import InferenceParams +from megatron.core.models.common.vision_module.vision_module import VisionModule +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.transformer.enums import ModelType +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_block import TransformerBlock + +from megatron.bridge.models.ernie_vl.modeling_ernie45_vl.vision_transformer_config import ErnieVisionTransformerConfig + + +class ErnieVisionPatchEmbed(nn.Module): + """Patch embedding for ERNIE 4.5 VL ViT. + + Unlike Qwen3VL which uses Conv3d on raw image tensors, ERNIE's processor + pre-flattens each patch into a vector of size [C * patch_size^2] = [588], + so patch embedding is a simple linear projection. + + Args: + in_channels: Number of input channels (default 3). + patch_size: Patch size in pixels (default 14). + embed_dim: Embedding dimension (default 1280). + """ + + def __init__(self, in_channels: int = 3, patch_size: int = 14, embed_dim: int = 1280): + super().__init__() + self.patch_size = patch_size + self.in_channels = in_channels + self.embed_dim = embed_dim + # HF PatchEmbed: nn.Linear(in_channels * patch_size^2, embed_dim, bias=False) + self.proj = nn.Linear(in_channels * patch_size * patch_size, embed_dim, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Project pre-flattened patches to embedding space. + + Args: + hidden_states: [total_patches, C * patch_size^2] (e.g., [N, 588]) + + Returns: + [total_patches, embed_dim] (e.g., [N, 1280]) + """ + return self.proj(hidden_states.to(dtype=self.proj.weight.dtype)) + + +class ErnieVisionRotaryEmbedding(nn.Module): + """1D rotary embedding frequency table for ERNIE ViT 2D RoPE. + + Computes a frequency table of shape [max_seqlen, dim//2] which is then + indexed by 2D (H, W) position IDs to produce per-token RoPE embeddings. + + This matches HF's ``VisionRotaryEmbedding`` in the ERNIE 4.5 VL model. + + Args: + dim: Half of the per-head dimension (head_dim // 2). + For ERNIE ViT: head_dim = 1280 / 16 = 80, so dim = 40. + theta: RoPE base frequency (default 10000.0). + """ + + def __init__(self, dim: int, theta: float = 10000.0): + super().__init__() + self.dim = dim + self.theta = theta + + def forward(self, seqlen: int) -> torch.Tensor: + """Compute frequency table for positions 0..seqlen-1. + + Args: + seqlen: Maximum sequence length to compute frequencies for. + + Returns: + Tensor of shape [seqlen, dim] containing outer product of + position indices and inverse frequencies. + """ + if not hasattr(self, "inv_freq"): + inv_freq = 1.0 / (self.theta ** (torch.arange(0, self.dim, 2, dtype=torch.float) / self.dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype) + freqs = torch.outer(seq, self.inv_freq) + return freqs + + +class ErnieVLVisionModel(VisionModule): + """Megatron-Core native ERNIE 4.5 VL Vision Transformer. + + Implements the DFN-style ViT with 2D RoPE using MCore TransformerBlock + for TP-native distributed training. + + Architecture: + 1. PatchEmbed (nn.Linear, replicated) + 2. 2D RoPE computation (per-image H/W position lookup with spatial merge reordering) + 3. TransformerBlock (32 ViT layers with TE modules) + 4. Final LayerNorm + + Unlike the HF-wrapped version, this implementation: + - Uses TE-backed attention and MLP layers for TP support + - Leverages MCore's PackedSeqParams for per-image variable-length attention + - Enables activation recomputation through TransformerBlock + + Args: + transformer_config: ErnieVisionTransformerConfig with ViT hyperparameters. + transformer_layer_spec: ModuleSpec for each ViT transformer layer. + """ + + def __init__( + self, + transformer_config: ErnieVisionTransformerConfig, + transformer_layer_spec: ModuleSpec, + ) -> None: + super().__init__(config=transformer_config) + + self.spatial_merge_size = transformer_config.spatial_merge_size + self.patch_size = transformer_config.patch_size + + # Patch embedding: nn.Linear (replicated across TP ranks) + self.patch_embed = ErnieVisionPatchEmbed( + in_channels=transformer_config.in_channels, + patch_size=transformer_config.patch_size, + embed_dim=transformer_config.hidden_size, + ) + + # 1D frequency table for 2D RoPE lookup + head_dim = transformer_config.hidden_size // transformer_config.num_attention_heads + self.rotary_pos_emb = ErnieVisionRotaryEmbedding(head_dim // 2) + + self.model_type = ModelType.encoder_or_decoder + + # Transformer layers (32 ViT blocks with TE modules) + self.decoder = TransformerBlock( + config=transformer_config, + spec=transformer_layer_spec, + pre_process=True, + post_process=True, + post_layer_norm=True, # Apply final LN after last block + ) + + self.input_tensor = None + + def set_input_tensor(self, input_tensor: torch.Tensor) -> None: + """Set input tensor (for pipeline parallelism, currently not used for ViT).""" + self.input_tensor = input_tensor + + def rot_pos_emb(self, grid_thw: torch.Tensor) -> torch.Tensor: + """Compute 2D RoPE positional embeddings for all tokens. + + For each image/video frame, computes (H, W) position IDs with + spatial_merge_size reordering (grouping merge_size x merge_size + patches together), then looks up the frequency table. + + The spatial merge reordering ensures that patches within each + spatial merge unit (2x2 by default) are consecutive in the + sequence, matching the resampler's spatial pooling pattern. + + Args: + grid_thw: [num_images, 3] tensor of (T, H, W) grid dimensions + for each image/video. + + Returns: + Tensor of shape [total_tokens, head_dim] containing the + concatenated cos/sin frequencies for 2D RoPE. + """ + merge_size = self.spatial_merge_size + + # Compute frequency table up to the maximum spatial dimension + max_hw = int(grid_thw[:, 1:].max().item()) + freq_table = self.rotary_pos_emb(max_hw) # [max_hw, dim//2] + device = freq_table.device + + total_tokens = int(torch.prod(grid_thw, dim=1).sum().item()) + pos_ids = torch.empty((total_tokens, 2), dtype=torch.long, device=device) + + offset = 0 + for num_frames, height, width in grid_thw: + merged_h = height // merge_size + merged_w = width // merge_size + + # Build position IDs with spatial merge reordering: + # Within each merged_h x merged_w block, iterate over + # merge_size x merge_size sub-positions + block_rows = torch.arange(merged_h, device=device) + block_cols = torch.arange(merged_w, device=device) + intra_row = torch.arange(merge_size, device=device) + intra_col = torch.arange(merge_size, device=device) + + # Full-resolution (H, W) positions + row_idx = block_rows[:, None, None, None] * merge_size + intra_row[None, None, :, None] + col_idx = block_cols[None, :, None, None] * merge_size + intra_col[None, None, None, :] + + row_idx = row_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1) + col_idx = col_idx.expand(merged_h, merged_w, merge_size, merge_size).reshape(-1) + + coords = torch.stack((row_idx, col_idx), dim=-1) # [H*W, 2] + + if num_frames > 1: + coords = coords.repeat(num_frames, 1) + + num_tokens = coords.shape[0] + pos_ids[offset : offset + num_tokens] = coords + offset += num_tokens + + # Look up frequency table by position IDs: [total_tokens, 2, dim//2] + embeddings = freq_table[pos_ids] + # Flatten to [total_tokens, head_dim//2 * 2 = head_dim] + embeddings = embeddings.flatten(1) + return embeddings + + def build_packed_seq_params( + self, + grid_thw: torch.Tensor, + ) -> PackedSeqParams: + """Build PackedSeqParams for per-image variable-length attention. + + Each frame in each image/video is treated as a separate sequence + for attention computation. This enables per-image attention without + cross-image contamination. + + Args: + grid_thw: [num_images, 3] tensor of (T, H, W) grid dimensions. + + Returns: + PackedSeqParams with cu_seqlens for thd-format attention. + """ + # Each frame is a separate sequence: seqlen = H * W + seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]) + cu_seqlens = seqlens.cumsum(dim=0) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0).int() + + max_seqlen_q = seqlens.max() + return PackedSeqParams( + cu_seqlens_q=cu_seqlens, + cu_seqlens_kv=cu_seqlens, + qkv_format="thd", + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_q, + ) + + def forward( + self, + hidden_states: torch.Tensor, + grid_thw: torch.Tensor, + inference_params: Optional[InferenceParams] = None, + extra_block_kwargs: Optional[dict] = None, + ) -> torch.Tensor: + """Forward pass of the ERNIE ViT. + + Args: + hidden_states: Pre-flattened pixel patches [total_patches, C*P*P]. + grid_thw: [num_images, 3] tensor of (T, H, W) grid dimensions. + inference_params: Inference parameters (currently unused for ViT). + extra_block_kwargs: Extra kwargs to pass to TransformerBlock. + + Returns: + Vision features of shape [total_patches, hidden_size]. + """ + assert grid_thw is not None + assert self.input_tensor is None + assert inference_params is None + + # 1. Patch embedding: [total_patches, C*P*P] -> [total_patches, embed_dim] + hidden_states = self.patch_embed(hidden_states) + + seq_len = hidden_states.size(0) + + # 2. Compute 2D RoPE frequencies: [total_patches, head_dim] + # rot_pos_emb() returns raw frequency values (theta * position) of + # shape [total_patches, head_dim//2] = [N, 40]. + # + # HF's apply_rotary_pos_emb_vision tiles cos/sin along the last dim: + # cos = freqs.cos().unsqueeze(1).tile(1, 1, 2) # (N,40) -> (N,1,80) + # so the effective frequency pattern is [f0..f39, f0..f39] (doubled). + # This means ALL head_dim dimensions get rotated (rot_dim == head_dim). + # + # MCore's _apply_rotary_pos_emb_bshd uses freqs.shape[-1] as rot_dim. + # To match HF's tiling, we must duplicate the frequencies: + # [f0..f39] -> [f0..f39, f0..f39] (shape 80) + # so that rot_dim == head_dim and the rotation covers all dimensions. + # + # With rotary_interleaved=False, MCore's _rotate_half splits at the + # midpoint [-x2, x1] which is identical to HF's rotate_half. + rotary_pos_emb = self.rot_pos_emb(grid_thw) + rotary_pos_emb = rotary_pos_emb.reshape(seq_len, 1, 1, -1) + rotary_pos_emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1) + + # 3. Reshape for TransformerBlock: [total_tokens, 1, embed_dim] + # (seq_len, batch=1, hidden_size) + hidden_states = hidden_states[:, None] + + # 4. Forward through transformer layers + hidden_states = self.decoder( + hidden_states=hidden_states, + attention_mask=None, + inference_params=inference_params, + rotary_pos_emb=rotary_pos_emb, + packed_seq_params=self.build_packed_seq_params(grid_thw), + **(extra_block_kwargs or {}), + ) + + # 5. Remove batch dimension: [total_tokens, 1, hidden_size] -> [total_tokens, hidden_size] + hidden_states = hidden_states.squeeze(1) + + return hidden_states diff --git a/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_transformer_config.py b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_transformer_config.py new file mode 100644 index 0000000000..4b3636ffdb --- /dev/null +++ b/src/megatron/bridge/models/ernie_vl/modeling_ernie45_vl/vision_transformer_config.py @@ -0,0 +1,158 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +TransformerConfig for the ERNIE 4.5 VL vision encoder (DFN-style ViT with 2D RoPE). + +This config inherits from Megatron-Core's TransformerConfig and adds +vision-specific fields (patch_size, spatial_merge_size, etc.). It is +constructed from the HF vision config via ``get_ernie_vision_config()``. +""" + +from dataclasses import dataclass + +import torch +from megatron.core.transformer.transformer_config import TransformerConfig + + +@dataclass +class ErnieVisionTransformerConfig(TransformerConfig): + """TransformerConfig for ERNIE 4.5 VL vision encoder. + + Extends Megatron-Core TransformerConfig with ERNIE vision-specific fields. + + Architecture constants from HF DFNRopeVisionTransformerConfig: + embed_dim=1280, depth=32, num_heads=16, mlp_ratio=4, + patch_size=14, in_channels=3, spatial_merge_size=2, + hidden_act="quick_gelu" + """ + + patch_size: int = 14 + """Vision patch size (pixels per side).""" + + in_channels: int = 3 + """Number of input image channels.""" + + spatial_merge_size: int = 2 + """Spatial merge factor for the resampler (2x2 pooling).""" + + +def _quick_gelu(x): + """Quick GELU activation: x * sigmoid(1.702 * x). + + This is the activation function used by ERNIE 4.5 VL ViT (and OpenAI CLIP). + It is a fast approximation of GELU but is NOT equivalent to + ``F.gelu(x, approximate="tanh")``, which uses a different formula: + 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3))) + The two differ by up to ~2% per element. + """ + return x * torch.sigmoid(1.702 * x) + + +def get_ernie_vision_config( + hf_vision_config, + megatron_config=None, +) -> ErnieVisionTransformerConfig: + """Construct an ErnieVisionTransformerConfig from a HF vision config. + + Args: + hf_vision_config: HF DFNRopeVisionTransformerConfig or equivalent + with fields: embed_dim, depth, num_heads, mlp_ratio, patch_size, + in_channels, spatial_merge_size, hidden_act. + megatron_config: Optional language model TransformerConfig to copy + recompute / CUDA-graph / TP settings from. + + Returns: + ErnieVisionTransformerConfig ready for ErnieVLVisionModel. + """ + embed_dim = getattr(hf_vision_config, "embed_dim", getattr(hf_vision_config, "hidden_size", 1280)) + num_heads = getattr(hf_vision_config, "num_heads", getattr(hf_vision_config, "num_attention_heads", 16)) + mlp_ratio = getattr(hf_vision_config, "mlp_ratio", 4) + depth = getattr(hf_vision_config, "depth", getattr(hf_vision_config, "num_hidden_layers", 32)) + + config = ErnieVisionTransformerConfig( + num_layers=depth, + hidden_size=embed_dim, + num_attention_heads=num_heads, + ffn_hidden_size=int(embed_dim * mlp_ratio), + add_bias_linear=True, # ERNIE ViT: all linear layers have bias=True + add_qkv_bias=True, # ERNIE ViT: QKV projection has bias=True + ) + + # Copy parallelism / recompute settings from language model config + if megatron_config is not None: + config.recompute_granularity = megatron_config.recompute_granularity + config.recompute_method = megatron_config.recompute_method + config.recompute_num_layers = megatron_config.recompute_num_layers + config.tensor_model_parallel_size = megatron_config.tensor_model_parallel_size + config.enable_cuda_graph = megatron_config.enable_cuda_graph + config.cuda_graph_use_single_mempool = megatron_config.cuda_graph_use_single_mempool + config.cuda_graph_retain_backward_graph = megatron_config.cuda_graph_retain_backward_graph + config.cuda_graph_warmup_steps = megatron_config.cuda_graph_warmup_steps + config.external_cuda_graph = megatron_config.external_cuda_graph + config.cuda_graph_impl = megatron_config.cuda_graph_impl + config.cuda_graph_scope = megatron_config.cuda_graph_scope + + # Vision encoder specific: no MoE, no EP + config.num_moe_experts = None + config.expert_model_parallel_size = 1 + config.moe_ffn_hidden_size = None + + # No dropout in vision encoder + config.hidden_dropout = 0.0 + config.attention_dropout = 0.0 + + # LayerNorm with eps=1e-6 (matching HF DFNRopeVisionBlock) + config.layernorm_epsilon = 1e-6 + config.normalization = "LayerNorm" + + # ERNIE ViT uses quick_gelu: x * sigmoid(1.702 * x) + # Note: quick_gelu is NOT the same as F.gelu(x, approximate="tanh"). + # F.gelu(tanh) uses: 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715*x^3))) + # The two differ by up to ~2% per element, which compounds over 32 layers. + config.activation_func = _quick_gelu + config.gated_linear_unit = False # No gated MLP in ViT + + # Derived head dimension + config.kv_channels = embed_dim // num_heads + config.num_query_groups = num_heads # No GQA in ViT + + # Disable various fusions/features not needed for ViT + config.layernorm_zero_centered_gamma = False + config.apply_query_key_layer_scaling = False + config.bias_activation_fusion = False + config.bias_dropout_fusion = False + config.attention_softmax_in_fp32 = True + config.apply_rope_fusion = False + + # No TP comm overlap or SP for vision encoder + config.tp_comm_overlap = False + config.sequence_parallel = False + + # No pipeline parallelism for vision encoder + config.context_parallel_size = 1 + config.pipeline_model_parallel_size = 1 + config.num_layers_in_first_pipeline_stage = None + config.num_layers_in_last_pipeline_stage = None + config.virtual_pipeline_model_parallel_size = 1 + config.pipeline_model_parallel_layout = None + config.account_for_embedding_in_pipeline_split = None + config.account_for_loss_in_pipeline_split = None + + # Vision-specific fields + config.patch_size = getattr(hf_vision_config, "patch_size", 14) + config.in_channels = getattr(hf_vision_config, "in_channels", 3) + config.spatial_merge_size = getattr(hf_vision_config, "spatial_merge_size", 2) + + return config diff --git a/tests/functional_tests/launch_scripts/gb200/active/L1_Launch_models_ernie.sh b/tests/functional_tests/launch_scripts/gb200/active/L1_Launch_models_ernie.sh new file mode 100755 index 0000000000..ac613921eb --- /dev/null +++ b/tests/functional_tests/launch_scripts/gb200/active/L1_Launch_models_ernie.sh @@ -0,0 +1,18 @@ +#!/bin/bash +set -xeuo pipefail + +export CUDA_VISIBLE_DEVICES="0,1" + +uv run coverage run \ + --data-file=/opt/Megatron-Bridge/.coverage \ + --source=/opt/Megatron-Bridge/ \ + --parallel-mode \ + -m pytest \ + -o log_cli=true \ + -o log_cli_level=INFO \ + -v -s -x \ + -m "not pleasefixme" \ + --tb=short -rA \ + tests/functional_tests/test_groups/models/ernie/test_ernie45_moe_conversion.py + +coverage combine -q diff --git a/tests/functional_tests/launch_scripts/gb200/active/L1_Launch_models_ernie_vl.sh b/tests/functional_tests/launch_scripts/gb200/active/L1_Launch_models_ernie_vl.sh new file mode 100755 index 0000000000..92ad0cc98e --- /dev/null +++ b/tests/functional_tests/launch_scripts/gb200/active/L1_Launch_models_ernie_vl.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +set -xeuo pipefail # Exit immediately if a command exits with a non-zero status + +export CUDA_VISIBLE_DEVICES="0,1" + +uv run coverage run --data-file=/opt/Megatron-Bridge/.coverage --source=/opt/Megatron-Bridge/ --parallel-mode -m pytest \ + -o log_cli=true -o log_cli_level=INFO -v -s -x -m "not pleasefixme" --tb=short -rA \ + tests/functional_tests/test_groups/models/ernie_vl/test_ernie45_vl_conversion.py +coverage combine -q diff --git a/tests/functional_tests/test_groups/models/ernie/__init__.py b/tests/functional_tests/test_groups/models/ernie/__init__.py new file mode 100644 index 0000000000..341a77c5bc --- /dev/null +++ b/tests/functional_tests/test_groups/models/ernie/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/functional_tests/test_groups/models/ernie/test_ernie45_moe_conversion.py b/tests/functional_tests/test_groups/models/ernie/test_ernie45_moe_conversion.py new file mode 100644 index 0000000000..52dc86aa60 --- /dev/null +++ b/tests/functional_tests/test_groups/models/ernie/test_ernie45_moe_conversion.py @@ -0,0 +1,350 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import importlib +import json +import shutil +import subprocess +import sys +from pathlib import Path + +import pytest +import torch +from transformers import AutoTokenizer, Ernie4_5_MoeConfig, Ernie4_5_MoeForCausalLM + + +# Path to the local ERNIE VL model that contains tokenizer files. +# Used as a tokenizer source when running offline (CI/air-gapped environments). +_ERNIE_VL_MODEL_PATH = Path(__file__).parent.parent.parent.parent.parent.parent / ("ERNIE-4.5-VL-28B-A3B-Thinking") +# Tokenizer files to copy from the reference model directory. +_TOKENIZER_FILES = [ + "tokenizer_config.json", + "tokenizer.model", + "added_tokens.json", + "special_tokens_map.json", +] + + +# Toy config: 4 layers (layer 0 dense, layers 1-3 MoE), 4 experts, top-2 routing, +# 1 shared expert, small hidden/intermediate sizes for fast CI testing. +HF_ERNIE45_MOE_TOY_MODEL_CONFIG = { + "architectures": ["Ernie4_5_MoeForCausalLM"], + "model_type": "ernie4_5_moe", + "hidden_size": 256, + "intermediate_size": 512, + "num_hidden_layers": 4, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "hidden_act": "silu", + "max_position_embeddings": 1024, + "initializer_range": 0.02, + "rms_norm_eps": 1e-05, + "use_cache": True, + "use_bias": False, + "vocab_size": 1024, + "rope_theta": 500000.0, + "tie_word_embeddings": False, + # MoE settings + "moe_num_experts": 4, + "moe_k": 2, + "moe_intermediate_size": 128, + "moe_num_shared_experts": 1, + "moe_layer_start_index": 1, + "moe_layer_end_index": 3, + "moe_layer_interval": 1, + "router_aux_loss_coef": 0.001, + "output_router_logits": False, + # Token IDs + "bos_token_id": 1, + "eos_token_id": 2, + "pad_token_id": 0, + # Dtype + "torch_dtype": "bfloat16", +} + + +class TestErnie45MoEConversion: + """ + Test ERNIE 4.5 MoE model conversion from local HuggingFace model + with different parallelism configurations. + """ + + @pytest.fixture(scope="class") + def ernie45_moe_toy_model_path(self, tmp_path_factory): + """ + Create and save a HuggingFace ERNIE 4.5 MoE toy model from config + to a temporary directory. + + Args: + tmp_path_factory: Pytest temporary path factory for class-scoped fixtures + + Returns: + str: Path to the saved HuggingFace model directory + """ + # Create a temporary directory for this test class + temp_dir = tmp_path_factory.mktemp("ernie45_moe_toy_model") + model_dir = temp_dir / "ernie45_moe_toy" + + # Create ERNIE 4.5 MoE config from the toy model config + config = Ernie4_5_MoeConfig(**HF_ERNIE45_MOE_TOY_MODEL_CONFIG) + config.torch_dtype = torch.bfloat16 + + # Create model with random weights and convert to bfloat16 + model = Ernie4_5_MoeForCausalLM(config) + model = model.bfloat16() + + # Copy tokenizer files from the local ERNIE VL model (works offline). + # Falls back to downloading from HuggingFace Hub if the local model is absent. + if _ERNIE_VL_MODEL_PATH.exists(): + model_dir.mkdir(parents=True, exist_ok=True) + for fname in _TOKENIZER_FILES: + src = _ERNIE_VL_MODEL_PATH / fname + if src.exists(): + shutil.copy2(src, model_dir / fname) + # Sanitize tokenizer_config.json: the VL model's config references + # a custom tokenizer class (Ernie4_5_VLTokenizer) via auto_map, + # which triggers trust_remote_code checks. Strip these fields so + # the toy model uses the standard LlamaTokenizer instead. + tok_cfg_path = model_dir / "tokenizer_config.json" + if tok_cfg_path.exists(): + with open(tok_cfg_path) as f: + tok_cfg = json.load(f) + tok_cfg.pop("auto_map", None) + tok_cfg.pop("tokenizer_class", None) + with open(tok_cfg_path, "w") as f: + json.dump(tok_cfg, f, indent=2) + else: + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") + tokenizer.save_pretrained(model_dir) + + # Save model and config to directory + model.save_pretrained(model_dir, safe_serialization=True) + + # Also save config.json explicitly to ensure compatibility + config_to_save = HF_ERNIE45_MOE_TOY_MODEL_CONFIG.copy() + config_path = model_dir / "config.json" + with open(config_path, "w") as f: + json.dump(config_to_save, f, indent=2) + + return str(model_dir) + + def test_toy_model_creation(self, ernie45_moe_toy_model_path): + """ + Test that the toy ERNIE 4.5 MoE model is created correctly and can be loaded. + + Args: + ernie45_moe_toy_model_path: Path to the toy ERNIE 4.5 MoE model (from fixture) + """ + # Verify the model directory exists + model_path = Path(ernie45_moe_toy_model_path) + assert model_path.exists(), f"Model directory not found at {model_path}" + + # Check essential files exist + config_file = model_path / "config.json" + assert config_file.exists(), f"config.json not found at {config_file}" + + # Check for model weights (safetensors preferred) + weights_file = model_path / "model.safetensors" + if not weights_file.exists(): + weights_file = model_path / "pytorch_model.bin" + + # If neither single file exists, check for sharded files + if not weights_file.exists(): + sharded_files = list(model_path.glob("model-*-of-*.safetensors")) + if sharded_files: + weights_file = sharded_files[0] + else: + sharded_files = list(model_path.glob("pytorch_model-*-of-*.bin")) + if sharded_files: + weights_file = sharded_files[0] + + assert weights_file.exists(), f"Model weights file not found in {model_path}" + + # Check for tokenizer files + tokenizer_config_file = model_path / "tokenizer_config.json" + assert tokenizer_config_file.exists(), f"tokenizer_config.json not found at {tokenizer_config_file}" + + # Load and verify config + with open(config_file) as f: + config_data = json.load(f) + + assert config_data["model_type"] == "ernie4_5_moe" + assert config_data["hidden_size"] == 256 + assert config_data["num_hidden_layers"] == 4 + assert config_data["num_attention_heads"] == 4 + assert config_data["vocab_size"] == 1024 + # Verify MoE specific parameters + assert config_data["moe_num_experts"] == 4 + assert config_data["moe_k"] == 2 + assert config_data["moe_intermediate_size"] == 128 + assert config_data["moe_num_shared_experts"] == 1 + assert config_data["moe_layer_start_index"] == 1 + + # Try loading the model to verify it's valid + try: + model = Ernie4_5_MoeForCausalLM.from_pretrained( + ernie45_moe_toy_model_path, + torch_dtype=torch.bfloat16, + low_cpu_mem_usage=False, + ) + + # Try loading the tokenizer as well + try: + tokenizer = AutoTokenizer.from_pretrained(ernie45_moe_toy_model_path) + print(f"Tokenizer loaded successfully with vocab_size: {tokenizer.vocab_size}") + except Exception as e: + print(f"Warning: Could not load tokenizer (this might be OK for conversion testing): {e}") + + # Verify model structure + assert hasattr(model, "model") + assert hasattr(model.model, "layers") + assert len(model.model.layers) == 4 # num_hidden_layers + + # Verify MoE structure: layer 0 should be dense, layers 1-3 should be MoE + layer0 = model.model.layers[0] + layer1 = model.model.layers[1] + assert hasattr(layer0, "mlp") + assert hasattr(layer1, "mlp") + # Layer 0 is dense (Ernie4_5_MoeMLP), layer 1 is MoE (Ernie4_5_MoeSparseMoeBlock) + assert "MLP" in type(layer0.mlp).__name__ + assert "SparseMoeBlock" in type(layer1.mlp).__name__ or "Moe" in type(layer1.mlp).__name__ + + print(f"SUCCESS: ERNIE 4.5 MoE toy model created and validated at {ernie45_moe_toy_model_path}") + print("Model weights are correctly in bfloat16 format") + print( + f"MoE structure validated: {config_data['moe_num_experts']} experts, " + f"top-{config_data['moe_k']} routing" + ) + + except Exception as e: + pytest.fail(f"Failed to load created toy MoE model: {e}") + + @pytest.mark.run_only_on("GPU") + @pytest.mark.parametrize( + "tp,pp,ep,test_name", + [ + (2, 1, 1, "TP"), + (1, 2, 1, "PP"), + pytest.param(1, 1, 2, "EP", marks=pytest.mark.pleasefixme), + ], + ) + def test_ernie45_moe_conversion_parallelism(self, ernie45_moe_toy_model_path, tmp_path, tp, pp, ep, test_name): + """ + Test ERNIE 4.5 MoE model conversion with different parallelism configurations. + + Args: + ernie45_moe_toy_model_path: Path to the toy ERNIE 4.5 MoE model (from fixture) + tmp_path: Pytest temporary path fixture + tp: Tensor parallelism size + pp: Pipeline parallelism size + ep: Expert parallelism size + test_name: Name of the test for identification + """ + + # Create temporary output directory for conversion results + test_output_dir = tmp_path / f"ernie45_moe_{test_name}" + test_output_dir.mkdir(exist_ok=True) + + # Run hf_megatron_roundtrip_multi_gpu.py with specified parallelism. + # Use coverage wrapper when available (CI), skip it otherwise (local dev). + _has_coverage = importlib.util.find_spec("coverage") is not None + cmd = [ + sys.executable, + "-m", + "torch.distributed.run", + "--nproc_per_node=2", + "--nnodes=1", + ] + if _has_coverage: + cmd += [ + "-m", + "coverage", + "run", + "--data-file=/opt/Megatron-Bridge/.coverage", + "--source=/opt/Megatron-Bridge/", + "--parallel-mode", + ] + cmd += [ + "examples/conversion/hf_megatron_roundtrip_multi_gpu.py", + "--hf-model-id", + ernie45_moe_toy_model_path, + "--output-dir", + str(test_output_dir), + "--tp", + str(tp), + "--pp", + str(pp), + "--ep", + str(ep), + "--trust-remote-code", + ] + + try: + result = subprocess.run( + cmd, + capture_output=True, + text=True, + cwd=Path(__file__).parent.parent.parent.parent.parent.parent, + ) + + # Check that the conversion completed successfully + if result.returncode != 0: + print(f"STDOUT: {result.stdout}") + print(f"STDERR: {result.stderr}") + pytest.fail(f"ERNIE 4.5 MoE {test_name} conversion failed with return code {result.returncode}") + + # Verify that the converted model was saved + model_name = Path(ernie45_moe_toy_model_path).name # "ernie45_moe_toy" + converted_model_dir = test_output_dir / model_name + assert converted_model_dir.exists(), f"Converted model directory not found at {converted_model_dir}" + + # Check that essential model files exist + config_file = converted_model_dir / "config.json" + assert config_file.exists(), f"config.json not found in converted model at {config_file}" + + # Check for model weights file + weights_file_safetensors = converted_model_dir / "model.safetensors" + weights_file_pytorch = converted_model_dir / "pytorch_model.bin" + + weights_found = weights_file_safetensors.exists() or weights_file_pytorch.exists() + + if not weights_found: + sharded_safetensors = list(converted_model_dir.glob("model-*-of-*.safetensors")) + sharded_pytorch = list(converted_model_dir.glob("pytorch_model-*-of-*.bin")) + weights_found = len(sharded_safetensors) > 0 or len(sharded_pytorch) > 0 + + assert weights_found, f"Model weights file not found in converted model at {converted_model_dir}" + + # Verify the config contains ERNIE 4.5 MoE-specific parameters + with open(config_file) as f: + saved_config = json.load(f) + + assert saved_config["model_type"] == "ernie4_5_moe", "Model type should be ernie4_5_moe" + assert saved_config["hidden_size"] == 256, "Hidden size should match toy config" + assert saved_config["num_attention_heads"] == 4, "Number of attention heads should match toy config" + # Verify MoE specific parameters are preserved + assert saved_config["moe_num_experts"] == 4, "Number of experts should match toy config" + assert saved_config["moe_k"] == 2, "moe_k (top-k routing) should match toy config" + assert saved_config["moe_intermediate_size"] == 128, "MoE intermediate size should match toy config" + + print(f"SUCCESS: ERNIE 4.5 MoE {test_name} conversion test completed successfully") + print(f"Converted model saved at: {converted_model_dir}") + print( + f"MoE parameters preserved: {saved_config['moe_num_experts']} experts, " + f"top-{saved_config['moe_k']} routing" + ) + + except Exception as e: + print(f"Error during ERNIE 4.5 MoE {test_name} conversion test: {e}") + raise diff --git a/tests/functional_tests/test_groups/models/ernie_vl/__init__.py b/tests/functional_tests/test_groups/models/ernie_vl/__init__.py new file mode 100644 index 0000000000..341a77c5bc --- /dev/null +++ b/tests/functional_tests/test_groups/models/ernie_vl/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/functional_tests/test_groups/models/ernie_vl/test_ernie45_vl_conversion.py b/tests/functional_tests/test_groups/models/ernie_vl/test_ernie45_vl_conversion.py new file mode 100644 index 0000000000..fa64316eba --- /dev/null +++ b/tests/functional_tests/test_groups/models/ernie_vl/test_ernie45_vl_conversion.py @@ -0,0 +1,440 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Tests for ERNIE 4.5 VL MoE model conversion between HuggingFace and Megatron-Core formats. + +Usage: + uv run python -m torch.distributed.run --nproc_per_node=1 -m pytest \ + tests/functional_tests/test_groups/models/ernie_vl/test_ernie45_vl_conversion.py::TestErnie45VLConversion::test_toy_model_creation +""" + +import json +import subprocess +from pathlib import Path + +import pytest +import torch + + +try: + from transformers import AutoTokenizer + from transformers.models.ernie4_5_vl_moe.configuration_ernie4_5_vl_moe import ( + Ernie4_5_VLMoeConfig, + ) + from transformers.models.ernie4_5_vl_moe.modeling_ernie4_5_vl_moe import ( + Ernie4_5_VLMoeForConditionalGeneration, + ) + + _HAS_ERNIE45_VL = True +except ImportError: + _HAS_ERNIE45_VL = False + + +# Tiny model config optimized for fast testing. +# ERNIE 4.5 VL MoE architecture: dual-pool MoE (text_moe + vision_moe) + shared_experts. +# Reduced from 28B to ~10M parameters for fast CI testing. +# +# Key architectural constraints preserved: +# - Layer 0 is dense MLP, layers 1+ are sparse MoE +# - Dual-pool MoE: text_moe (intermediate=1536->64) + vision_moe (intermediate=512->32) +# - Shared experts: moe_num_shared_experts=2 +# - GQA: num_attention_heads must be divisible by num_key_value_heads +# - 3D M-RoPE: mrope_section sums to head_dim//2 (here [2, 2, 2] sums to 6 for head_dim=12) +# - num_key_value_heads >= 2 for TP=2 compatibility +HF_ERNIE45_VL_MOE_TOY_MODEL_CONFIG = { + "architectures": ["Ernie4_5_VLMoeForConditionalGeneration"], + "model_type": "ernie4_5_vl_moe", + "tie_word_embeddings": True, + "image_start_token_id": 101304, + "image_end_token_id": 101305, + "image_token_id": 100295, + "video_start_token_id": 101306, + "video_end_token_id": 101307, + "video_token_id": 103367, + "text_config": { + "model_type": "ernie4_5_vl_moe_text", + "vocab_size": 2048, + "hidden_size": 48, + "intermediate_size": 128, + "num_hidden_layers": 4, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "hidden_act": "silu", + "max_position_embeddings": 4096, + "initializer_range": 0.02, + "rms_norm_eps": 1e-5, + "use_cache": True, + "use_bias": False, + "rope_parameters": { + "rope_type": "default", + "rope_theta": 500000.0, + "mrope_section": [2, 2, 2], + }, + # MoE settings: scaled down from production (64 experts -> 4) + "moe_intermediate_size": [64, 32], + "moe_k": 2, + "moe_num_experts": 4, + "moe_num_shared_experts": 2, + "moe_norm_min": 1e-12, + "output_router_logits": False, + "router_aux_loss_coef": 0.001, + # Layer 0 = dense, layers 1-3 = sparse + "mlp_layer_types": ["dense", "sparse", "sparse", "sparse"], + }, + "vision_config": { + "model_type": "ernie4_5_vl_moe_vision", + "depth": 1, + "hidden_size": 48, + "hidden_act": "quick_gelu", + "num_heads": 4, + "in_channels": 3, + "patch_size": 14, + "spatial_merge_size": 2, + "intermediate_size": 128, + "temporal_merge_size": 2, + "rms_norm_eps": 1e-6, + }, +} + + +@pytest.mark.skipif(not _HAS_ERNIE45_VL, reason="ERNIE 4.5 VL MoE model not available in transformers") +class TestErnie45VLConversion: + """ + Test ERNIE 4.5 VL MoE model conversion from local HuggingFace model + with different parallelism configurations. + """ + + @pytest.fixture(scope="class") + def ernie45_vl_toy_model_path(self, tmp_path_factory): + """ + Create and save a HuggingFace ERNIE 4.5 VL MoE toy model from config + to a temporary directory. + + Args: + tmp_path_factory: Pytest temporary path factory for class-scoped fixtures. + + Returns: + str: Path to the saved HuggingFace model directory. + """ + temp_dir = tmp_path_factory.mktemp("ernie45_vl_toy_model") + model_dir = temp_dir / "ernie45_vl_toy" + + # Create ERNIE 4.5 VL MoE config from the toy model config dict + config = Ernie4_5_VLMoeConfig(**HF_ERNIE45_VL_MOE_TOY_MODEL_CONFIG) + config.torch_dtype = torch.bfloat16 + + # Ensure rope_parameters is set on text_config + if hasattr(config, "text_config") and config.text_config is not None: + config.text_config.rope_parameters = { + "rope_type": "default", + "rope_theta": 500000.0, + "mrope_section": [2, 2, 2], + } + + # Create model with random weights and convert to bfloat16 + model = Ernie4_5_VLMoeForConditionalGeneration(config) + model = model.to(dtype=torch.bfloat16) + + # Load tokenizer from local model or a reference model, or create minimal fallback + _local_tokenizer_path = ( + Path(__file__).parent.parent.parent.parent.parent.parent / "ERNIE-4.5-VL-28B-A3B-Thinking" + ) + try: + if _local_tokenizer_path.exists(): + tokenizer = AutoTokenizer.from_pretrained(str(_local_tokenizer_path)) + else: + tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B") + tokenizer.save_pretrained(model_dir) + except (OSError, ValueError): + # Create a functional dummy tokenizer from a readily available model + # so that save_hf_pretrained can re-save it without errors. + try: + from tokenizers import Tokenizer as TkTokenizer + from tokenizers import models as tk_models + from transformers import PreTrainedTokenizerFast + + # Build a minimal BPE tokenizer with a few tokens + tk = TkTokenizer(tk_models.BPE()) + tk.add_special_tokens(["", "", "", ""]) + # Pad the vocab to the expected size with dummy tokens + dummy_tokens = [f"" for i in range(103420)] + tk.add_tokens(dummy_tokens) + fast_tokenizer = PreTrainedTokenizerFast( + tokenizer_object=tk, + bos_token="", + eos_token="", + pad_token="", + unk_token="", + ) + model_dir.mkdir(parents=True, exist_ok=True) + fast_tokenizer.save_pretrained(model_dir) + except (OSError, ValueError): + # Last resort: just create the directory so save_pretrained can write weights + model_dir.mkdir(parents=True, exist_ok=True) + + # Save model weights (safetensors format) + model.save_pretrained(model_dir, safe_serialization=True) + + # Overwrite config.json with the toy config to ensure exact key structure + config_to_save = HF_ERNIE45_VL_MOE_TOY_MODEL_CONFIG.copy() + config_path = model_dir / "config.json" + with open(config_path, "w") as f: + json.dump(config_to_save, f, indent=2) + + return str(model_dir) + + def test_toy_model_creation(self, ernie45_vl_toy_model_path): + """ + Test that the ERNIE 4.5 VL MoE toy model is created correctly and can be loaded. + + Args: + ernie45_vl_toy_model_path: Path to the toy model (from fixture). + """ + model_path = Path(ernie45_vl_toy_model_path) + assert model_path.exists(), f"Model directory not found at {model_path}" + + # Check essential files exist + config_file = model_path / "config.json" + assert config_file.exists(), f"config.json not found at {config_file}" + + # Check for model weights + weights_file = model_path / "model.safetensors" + if not weights_file.exists(): + weights_file = model_path / "model.safetensors.index.json" + if not weights_file.exists(): + weights_file = model_path / "pytorch_model.bin" + assert weights_file.exists(), f"Model weights file not found in {model_path}" + + # Check for tokenizer files + tokenizer_config_file = model_path / "tokenizer_config.json" + assert tokenizer_config_file.exists(), f"tokenizer_config.json not found at {tokenizer_config_file}" + + # Load and verify config + with open(config_file) as f: + config_data = json.load(f) + + assert config_data["model_type"] == "ernie4_5_vl_moe" + assert "text_config" in config_data + assert "vision_config" in config_data + assert config_data["text_config"]["hidden_size"] == 48 + assert config_data["text_config"]["num_hidden_layers"] == 4 + assert config_data["text_config"]["num_attention_heads"] == 4 + assert config_data["text_config"]["moe_num_experts"] == 4 + assert config_data["text_config"]["moe_intermediate_size"] == [64, 32] + assert config_data["text_config"]["mlp_layer_types"] == ["dense", "sparse", "sparse", "sparse"] + + # Verify model can be loaded from pretrained + _ = Ernie4_5_VLMoeForConditionalGeneration.from_pretrained( + ernie45_vl_toy_model_path, + torch_dtype=torch.bfloat16, + low_cpu_mem_usage=False, + ) + + # Try loading the tokenizer + try: + tokenizer = AutoTokenizer.from_pretrained(ernie45_vl_toy_model_path) + print(f"Tokenizer loaded successfully with vocab_size: {tokenizer.vocab_size}") + except Exception as e: + print(f"Warning: Could not load tokenizer (OK for conversion testing): {e}") + + print(f"SUCCESS: ERNIE 4.5 VL MoE toy model created and validated at {ernie45_vl_toy_model_path}") + + @pytest.mark.run_only_on("GPU") + @pytest.mark.parametrize( + "tp,pp,ep,test_name", + [ + (2, 1, 1, "TP"), + (1, 2, 1, "PP"), + (1, 1, 2, "EP"), + ], + ) + def test_ernie45_vl_conversion_parallelism(self, ernie45_vl_toy_model_path, tmp_path, tp, pp, ep, test_name): + """ + Test ERNIE 4.5 VL MoE model conversion with different parallelism configurations. + + Covers: + - TP (Tensor Parallelism): splits attention heads and MLP across GPUs + - PP (Pipeline Parallelism): splits transformer layers across GPUs + - EP (Expert Parallelism): splits MoE experts across GPUs + + The EP test validates that dual-pool MoE (text_moe_layer + vision_moe_layer) + correctly handles per-pool expert offset when sharding across EP ranks. + + Args: + ernie45_vl_toy_model_path: Path to the toy model (from fixture). + tmp_path: Pytest temporary path fixture. + tp: Tensor parallelism size. + pp: Pipeline parallelism size. + ep: Expert parallelism size. + test_name: Name of the test for identification. + """ + test_output_dir = tmp_path / f"ernie45_vl_{test_name}" + test_output_dir.mkdir(exist_ok=True) + + # Run HF-to-Megatron roundtrip conversion as a subprocess + # Use sys.executable to ensure the same Python interpreter (venv-aware) + import sys + + cmd = [ + sys.executable, + "-m", + "torch.distributed.run", + "--nproc_per_node=2", + "--nnodes=1", + "examples/conversion/hf_megatron_roundtrip_multi_gpu.py", + "--hf-model-id", + ernie45_vl_toy_model_path, + "--output-dir", + str(test_output_dir), + "--tp", + str(tp), + "--pp", + str(pp), + "--ep", + str(ep), + ] + + try: + result = subprocess.run( + cmd, + capture_output=True, + text=True, + cwd=Path(__file__).parent.parent.parent.parent.parent.parent, + ) + + if result.returncode != 0: + print(f"STDOUT: {result.stdout}") + print(f"STDERR: {result.stderr}") + pytest.fail(f"ERNIE 4.5 VL {test_name} conversion failed with return code {result.returncode}") + + # Verify converted model directory exists + model_name = Path(ernie45_vl_toy_model_path).name # "ernie45_vl_toy" + converted_model_dir = test_output_dir / model_name + assert converted_model_dir.exists(), f"Converted model directory not found at {converted_model_dir}" + + # Check that essential model files exist + config_file = converted_model_dir / "config.json" + assert config_file.exists(), f"config.json not found in converted model at {config_file}" + + # Verify the config contains ERNIE 4.5 VL-specific parameters + with open(config_file) as f: + saved_config = json.load(f) + + assert saved_config["model_type"] == "ernie4_5_vl_moe", "Model type should be ernie4_5_vl_moe" + assert "text_config" in saved_config, "VL model should have text_config" + assert "vision_config" in saved_config, "VL model should have vision_config" + assert saved_config["text_config"]["hidden_size"] == 48, "Hidden size should match toy config" + assert saved_config["text_config"]["num_attention_heads"] == 4, ( + "Number of attention heads should match toy config" + ) + + print(f"SUCCESS: ERNIE 4.5 VL {test_name} conversion test completed successfully") + print(f"Converted model saved at: {converted_model_dir}") + + except Exception as e: + print(f"Error during ERNIE 4.5 VL {test_name} conversion test: {e}") + raise + + @pytest.mark.run_only_on("GPU") + @pytest.mark.parametrize( + "tp,pp,ep,nproc,with_vision,test_name", + [ + (1, 1, 1, 1, False, "single_gpu"), + (1, 1, 2, 2, False, "EP2"), + (1, 1, 1, 1, True, "single_gpu_vision"), + (1, 1, 2, 2, True, "EP2_vision"), + ], + ) + def test_ernie45_vl_forward_backward( + self, ernie45_vl_toy_model_path, tmp_path, tp, pp, ep, nproc, with_vision, test_name + ): + """ + Test ERNIE 4.5 VL MoE model forward and backward pass. + + Builds the Megatron model from the toy HF checkpoint via AutoBridge, + runs a forward pass and backward pass, and verifies: + - Forward produces finite output + - Backward produces gradients on trainable parameters + + When with_vision=True, a dummy image is injected to exercise the full + vision pipeline: ViT patch embedding -> vision transformer -> resampler + -> embedding injection -> language model forward. Vision tower and + resampler gradients are also verified. + + Args: + ernie45_vl_toy_model_path: Path to the toy model (from fixture). + tmp_path: Pytest temporary path fixture. + tp: Tensor parallelism size. + pp: Pipeline parallelism size. + ep: Expert parallelism size. + nproc: Number of processes for torchrun. + with_vision: Whether to include a dummy image input. + test_name: Name of the test for identification. + """ + import sys + + fwd_bwd_script = str(Path(__file__).parent / "ernie45_vl_fwd_bwd.py") + + cmd = [ + sys.executable, + "-m", + "torch.distributed.run", + f"--nproc_per_node={nproc}", + "--nnodes=1", + fwd_bwd_script, + "--hf-model-path", + ernie45_vl_toy_model_path, + "--tp", + str(tp), + "--pp", + str(pp), + "--ep", + str(ep), + "--seq-len", + "16", + ] + + if with_vision: + cmd.append("--with-vision") + + try: + result = subprocess.run( + cmd, + capture_output=True, + text=True, + cwd=Path(__file__).parent.parent.parent.parent.parent.parent, + timeout=300, + ) + + if result.returncode != 0: + print(f"STDOUT: {result.stdout[-3000:]}") + print(f"STDERR: {result.stderr[-3000:]}") + pytest.fail( + f"ERNIE 4.5 VL {test_name} forward/backward test failed with return code {result.returncode}" + ) + + # Verify the output contains the success marker + assert "ALL CHECKS PASSED" in result.stdout, ( + f"Forward/backward test did not complete successfully. STDOUT tail: {result.stdout[-1000:]}" + ) + + print(f"SUCCESS: ERNIE 4.5 VL {test_name} forward/backward test passed") + + except subprocess.TimeoutExpired: + print(f"TIMEOUT: ERNIE 4.5 VL {test_name} forward/backward test timed out after 300s") + raise + except Exception as e: + print(f"Error during ERNIE 4.5 VL {test_name} forward/backward test: {e}") + raise diff --git a/tests/unit_tests/models/ernie/__init__.py b/tests/unit_tests/models/ernie/__init__.py new file mode 100644 index 0000000000..341a77c5bc --- /dev/null +++ b/tests/unit_tests/models/ernie/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unit_tests/models/ernie/test_ernie45_bridge.py b/tests/unit_tests/models/ernie/test_ernie45_bridge.py new file mode 100644 index 0000000000..58e9aa0f79 --- /dev/null +++ b/tests/unit_tests/models/ernie/test_ernie45_bridge.py @@ -0,0 +1,373 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for ERNIE 4.5 text-only MoE bridge.""" + +from unittest.mock import Mock, patch + +import pytest +import torch + +from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.ernie.ernie_45_bridge import ( + Ernie45Bridge, + _PPSafeAutoMapping, + _PPSafeGatedMLPMapping, + _PPSafeReplicatedMapping, + _SqueezeBiasMapping, +) +from megatron.bridge.models.gpt_provider import GPTModelProvider + + +@pytest.fixture +def ernie45_config(): + """Minimal ERNIE 4.5 MoE config mock.""" + config = Mock(spec=[]) + config.hidden_size = 2560 + config.intermediate_size = 6912 + config.num_attention_heads = 20 + config.num_key_value_heads = 4 + config.num_hidden_layers = 28 + config.max_position_embeddings = 32768 + config.rms_norm_eps = 1e-05 + config.rope_theta = 500000.0 + config.tie_word_embeddings = False + config.torch_dtype = "bfloat16" + config.vocab_size = 189440 + config.initializer_range = 0.02 + config.model_type = "ernie4_5_moe" + config.architectures = ["Ernie4_5_MoeForCausalLM"] + # MoE fields + config.moe_num_experts = 64 + config.moe_k = 6 + config.moe_intermediate_size = 1408 + config.moe_num_shared_experts = 2 + config.router_aux_loss_coef = 0.001 + config.mlp_layer_types = ["dense"] + ["moe"] * 27 + return config + + +@pytest.fixture +def ernie45_config_list_format(): + """Config with list-format MoE fields.""" + config = Mock(spec=[]) + config.hidden_size = 2560 + config.intermediate_size = 6912 + config.num_attention_heads = 20 + config.num_key_value_heads = 4 + config.num_hidden_layers = 28 + config.max_position_embeddings = 32768 + config.rms_norm_eps = 1e-05 + config.rope_theta = 500000.0 + config.tie_word_embeddings = False + config.torch_dtype = "bfloat16" + config.vocab_size = 189440 + config.initializer_range = 0.02 + config.model_type = "ernie4_5_moe" + config.architectures = ["Ernie4_5_MoeForCausalLM"] + # MoE fields as lists + config.moe_num_experts = [64] + config.moe_k = 6 + config.moe_intermediate_size = [1408] + config.moe_num_shared_experts = 2 + config.router_aux_loss_coef = 0.001 + config.moe_layer_start_index = [1] + # No mlp_layer_types -- should derive from moe_layer_start_index + del config.mlp_layer_types + return config + + +@pytest.fixture +def mock_pretrained(ernie45_config): + """Mock PreTrainedCausalLM for ERNIE 4.5.""" + pretrained = Mock() + pretrained.config = ernie45_config + return pretrained + + +class TestErnie45BridgeRegistration: + """Test bridge class and registration.""" + + def test_is_subclass_of_megatron_model_bridge(self): + assert issubclass(Ernie45Bridge, MegatronModelBridge) + + def test_bridge_instantiation(self): + bridge = Ernie45Bridge() + assert bridge is not None + + +class TestErnie45BridgeGetNumExperts: + """Test _get_num_experts static method.""" + + def test_int_input(self): + config = Mock(spec=[]) + config.moe_num_experts = 64 + assert Ernie45Bridge._get_num_experts(config) == 64 + + def test_single_element_list(self): + config = Mock(spec=[]) + config.moe_num_experts = [64] + assert Ernie45Bridge._get_num_experts(config) == 64 + + def test_dual_pool_list_returns_first(self): + config = Mock(spec=[]) + config.moe_num_experts = [64, 32] + assert Ernie45Bridge._get_num_experts(config) == 64 + + def test_default_when_missing(self): + # Mock with spec=[] means arbitrary attrs return Mock; use a simple object + class EmptyConfig: + pass + + cfg = EmptyConfig() + assert Ernie45Bridge._get_num_experts(cfg) == 64 + + +class TestErnie45BridgeProviderBridge: + """Test provider_bridge method.""" + + def test_returns_gpt_model_provider(self, mock_pretrained): + bridge = Ernie45Bridge() + provider = bridge.provider_bridge(mock_pretrained) + assert isinstance(provider, GPTModelProvider) + + def test_basic_dimensions(self, mock_pretrained): + bridge = Ernie45Bridge() + provider = bridge.provider_bridge(mock_pretrained) + assert provider.num_layers == 28 + assert provider.hidden_size == 2560 + assert provider.num_attention_heads == 20 + + def test_vocabulary(self, mock_pretrained): + bridge = Ernie45Bridge() + provider = bridge.provider_bridge(mock_pretrained) + assert provider.vocab_size == 189440 + assert provider.share_embeddings_and_output_weights is False + + def test_normalization(self, mock_pretrained): + bridge = Ernie45Bridge() + provider = bridge.provider_bridge(mock_pretrained) + assert provider.normalization == "RMSNorm" + assert provider.layernorm_epsilon == 1e-05 + + def test_rope_config(self, mock_pretrained): + bridge = Ernie45Bridge() + provider = bridge.provider_bridge(mock_pretrained) + assert provider.position_embedding_type == "rope" + assert provider.rotary_base == 500000.0 + assert provider.rotary_interleaved is True + + def test_moe_basic(self, mock_pretrained): + bridge = Ernie45Bridge() + provider = bridge.provider_bridge(mock_pretrained) + assert provider.num_moe_experts == 64 + assert provider.moe_router_topk == 6 + assert provider.moe_ffn_hidden_size == 1408 + + def test_moe_router_settings(self, mock_pretrained): + bridge = Ernie45Bridge() + provider = bridge.provider_bridge(mock_pretrained) + assert provider.moe_grouped_gemm is True + assert provider.moe_router_score_function == "sigmoid" + assert provider.moe_router_enable_expert_bias is True + assert provider.moe_router_dtype == "fp32" + + def test_shared_experts(self, mock_pretrained): + bridge = Ernie45Bridge() + provider = bridge.provider_bridge(mock_pretrained) + # shared_expert_intermediate_size = moe_ffn_hidden_size * moe_num_shared_experts + assert provider.moe_shared_expert_intermediate_size == 1408 * 2 + + def test_moe_layer_freq_from_mlp_layer_types(self, mock_pretrained): + bridge = Ernie45Bridge() + provider = bridge.provider_bridge(mock_pretrained) + expected = [0] + [1] * 27 + assert provider.moe_layer_freq == expected + + def test_moe_layer_freq_from_moe_layer_start_index(self, ernie45_config_list_format): + pretrained = Mock() + pretrained.config = ernie45_config_list_format + bridge = Ernie45Bridge() + provider = bridge.provider_bridge(pretrained) + expected = [0] + [1] * 27 + assert provider.moe_layer_freq == expected + + def test_mlp_config(self, mock_pretrained): + bridge = Ernie45Bridge() + provider = bridge.provider_bridge(mock_pretrained) + assert provider.gated_linear_unit is True + assert provider.add_bias_linear is False + + def test_moe_intermediate_size_from_list(self, ernie45_config_list_format): + pretrained = Mock() + pretrained.config = ernie45_config_list_format + bridge = Ernie45Bridge() + provider = bridge.provider_bridge(pretrained) + assert provider.moe_ffn_hidden_size == 1408 + + def test_moe_aux_loss_coeff(self, mock_pretrained): + bridge = Ernie45Bridge() + provider = bridge.provider_bridge(mock_pretrained) + assert provider.moe_aux_loss_coeff == 0.001 + + +class TestErnie45BridgeMappingRegistry: + """Test mapping_registry method.""" + + def test_returns_mapping_registry(self): + bridge = Ernie45Bridge() + registry = bridge.mapping_registry() + assert isinstance(registry, MegatronMappingRegistry) + + def test_has_mappings(self): + bridge = Ernie45Bridge() + registry = bridge.mapping_registry() + assert len(registry.mappings) > 0 + + def test_has_embedding_mappings(self): + bridge = Ernie45Bridge() + registry = bridge.mapping_registry() + hf_params = [m.hf_param for m in registry.mappings if hasattr(m, "hf_param") and isinstance(m.hf_param, str)] + assert "model.embed_tokens.weight" in hf_params + assert "lm_head.weight" in hf_params + + def test_has_qkv_mapping(self): + from megatron.bridge.models.conversion.param_mapping import QKVMapping + + bridge = Ernie45Bridge() + registry = bridge.mapping_registry() + qkv_mappings = [m for m in registry.mappings if isinstance(m, QKVMapping)] + assert len(qkv_mappings) > 0 + + qkv = qkv_mappings[0] + assert "q" in qkv.hf_param + assert "k" in qkv.hf_param + assert "v" in qkv.hf_param + + def test_has_router_mapping(self): + bridge = Ernie45Bridge() + registry = bridge.mapping_registry() + megatron_params = [ + m.megatron_param + for m in registry.mappings + if hasattr(m, "megatron_param") and isinstance(m.megatron_param, str) + ] + assert "decoder.layers.*.mlp.router.weight" in megatron_params + + def test_has_expert_bias_mapping(self): + bridge = Ernie45Bridge() + registry = bridge.mapping_registry() + megatron_params = [ + m.megatron_param + for m in registry.mappings + if hasattr(m, "megatron_param") and isinstance(m.megatron_param, str) + ] + assert "decoder.layers.*.mlp.router.expert_bias" in megatron_params + + def test_has_shared_expert_mappings(self): + bridge = Ernie45Bridge() + registry = bridge.mapping_registry() + megatron_params = [ + m.megatron_param + for m in registry.mappings + if hasattr(m, "megatron_param") and isinstance(m.megatron_param, str) + ] + shared = [p for p in megatron_params if "shared_experts" in p] + assert len(shared) > 0 + + def test_has_expert_fc1_mapping(self): + bridge = Ernie45Bridge() + registry = bridge.mapping_registry() + megatron_params = [ + m.megatron_param + for m in registry.mappings + if hasattr(m, "megatron_param") and isinstance(m.megatron_param, str) + ] + expert_fc1 = [p for p in megatron_params if "experts" in p and "linear_fc1" in p] + assert len(expert_fc1) > 0 + + +class TestSqueezeBiasMapping: + """Test _SqueezeBiasMapping hf_to_megatron / megatron_to_hf.""" + + def _make_mock_module(self): + """Create a mock megatron_module with valid device attributes.""" + mock_module = Mock() + mock_module.weight = torch.nn.Parameter(torch.zeros(1)) + return mock_module + + def test_squeeze_2d_to_1d(self): + mapping = _SqueezeBiasMapping( + megatron_param="decoder.layers.*.mlp.router.expert_bias", + hf_param="model.layers.*.mlp.moe_statics.e_score_correction_bias", + ) + hf_weights = torch.randn(1, 64) + result = mapping.hf_to_megatron(hf_weights, self._make_mock_module()) + assert result.shape == (64,) + + def test_already_1d_passthrough(self): + mapping = _SqueezeBiasMapping( + megatron_param="decoder.layers.*.mlp.router.expert_bias", + hf_param="model.layers.*.mlp.moe_statics.e_score_correction_bias", + ) + hf_weights = torch.randn(64) + result = mapping.hf_to_megatron(hf_weights, self._make_mock_module()) + assert result.shape == (64,) + + def test_unsqueeze_on_export(self): + mapping = _SqueezeBiasMapping( + megatron_param="decoder.layers.*.mlp.router.expert_bias", + hf_param="model.layers.*.mlp.moe_statics.e_score_correction_bias", + ) + megatron_weights = torch.randn(64) + # megatron_to_hf returns a dict; patch super() to return the weights + with patch.object( + _PPSafeReplicatedMapping, + "megatron_to_hf", + return_value={"model.layers.0.mlp.moe_statics.e_score_correction_bias": megatron_weights}, + ): + result = mapping.megatron_to_hf(megatron_weights, Mock()) + # Result values should be [1, 64] + for v in result.values(): + assert v.shape == (1, 64) + + +class TestPPSafeMappings: + """Test PP-safe mapping variants.""" + + def test_pp_safe_auto_mapping_exists(self): + from megatron.bridge.models.conversion.param_mapping import AutoMapping + + assert issubclass(_PPSafeAutoMapping, AutoMapping) + mapping = _PPSafeAutoMapping( + megatron_param="decoder.layers.*.mlp.shared_experts.linear_fc2.weight", + hf_param="model.layers.*.mlp.shared_experts.down_proj.weight", + ) + assert mapping is not None + + def test_pp_safe_gated_mlp_mapping_exists(self): + mapping = _PPSafeGatedMLPMapping( + megatron_param="decoder.layers.*.mlp.experts.linear_fc1.weight*", + gate="model.layers.*.mlp.experts.*.gate_proj.weight", + up="model.layers.*.mlp.experts.*.up_proj.weight", + ) + assert mapping is not None + + def test_pp_safe_replicated_mapping_exists(self): + mapping = _PPSafeReplicatedMapping( + megatron_param="decoder.layers.*.mlp.router.weight", + hf_param="model.layers.*.mlp.gate.weight", + ) + assert mapping is not None diff --git a/tests/unit_tests/models/ernie_vl/__init__.py b/tests/unit_tests/models/ernie_vl/__init__.py new file mode 100644 index 0000000000..341a77c5bc --- /dev/null +++ b/tests/unit_tests/models/ernie_vl/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/tests/unit_tests/models/ernie_vl/test_ernie45_vl_bridge.py b/tests/unit_tests/models/ernie_vl/test_ernie45_vl_bridge.py new file mode 100644 index 0000000000..46daae368f --- /dev/null +++ b/tests/unit_tests/models/ernie_vl/test_ernie45_vl_bridge.py @@ -0,0 +1,479 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for ERNIE 4.5 VL (Vision-Language) MoE bridge.""" + +from unittest.mock import Mock + +import pytest +import torch + +from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry +from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge +from megatron.bridge.models.ernie_vl.ernie45_vl_bridge import ( + Ernie45VLBridge, + _ConcatBiasMapping, + _OffsetGatedMLPMapping, + _OffsetRowParallelMapping, +) +from megatron.bridge.models.ernie_vl.ernie45_vl_provider import Ernie45VLModelProvider + + +def _make_vision_config(): + """Create a mock vision config.""" + vision_config = Mock(spec=[]) + vision_config.hidden_size = 1280 + vision_config.num_attention_heads = 16 + vision_config.num_hidden_layers = 32 + vision_config.patch_size = 14 + vision_config.image_size = 384 + vision_config.intermediate_size = 5120 + return vision_config + + +def _make_flat_vl_config(): + """Create a flat (auto_map) VL config -- used by Thinking model.""" + config = Mock(spec=[]) + config.hidden_size = 2560 + config.intermediate_size = 6912 + config.num_attention_heads = 20 + config.num_key_value_heads = 4 + config.num_hidden_layers = 28 + config.max_position_embeddings = 32768 + config.rms_norm_eps = 1e-05 + config.rope_theta = 500000.0 + config.tie_word_embeddings = False + config.torch_dtype = "bfloat16" + config.vocab_size = 189440 + config.initializer_range = 0.02 + config.model_type = "ernie4_5_moe_vl" + config.architectures = ["Ernie4_5_VLMoeForConditionalGeneration"] + # Dual-pool MoE + config.moe_num_experts = [64, 64] + config.moe_k = 6 + config.moe_intermediate_size = [1536, 512] + config.moe_num_shared_experts = 2 + config.router_aux_loss_coef = 0.001 + config.mlp_layer_types = ["dense"] + ["moe"] * 27 + # Vision + config.vision_config = _make_vision_config() + # VL tokens + config.image_token_id = 151859 + config.video_token_id = 151860 + config.image_start_token_id = 101304 + config.image_end_token_id = 101305 + config.video_start_token_id = 101306 + config.video_end_token_id = 101307 + # Flat config: no text_config attribute + config.text_config = config # self-reference (as done by _normalize_hf_config) + # RoPE + config.rope_parameters = {"rope_theta": 500000.0, "mrope_section": [22, 22, 20]} + config.rope_scaling = None + # auto_map + config.auto_map = { + "AutoModelForCausalLM": "modeling_ernie4_5_vl.Ernie4_5_VLMoeForConditionalGeneration", + } + return config + + +def _make_nested_vl_config(): + """Create a nested (transformers-builtin) VL config.""" + text_config = Mock(spec=[]) + text_config.hidden_size = 2560 + text_config.intermediate_size = 6912 + text_config.num_attention_heads = 20 + text_config.num_key_value_heads = 4 + text_config.num_hidden_layers = 28 + text_config.max_position_embeddings = 32768 + text_config.rms_norm_eps = 1e-05 + text_config.rope_theta = 500000.0 + text_config.torch_dtype = "bfloat16" + text_config.vocab_size = 189440 + text_config.initializer_range = 0.02 + text_config.model_type = "ernie4_5_moe" + # MoE + text_config.moe_num_experts = 4 # nested config uses smaller num for toy + text_config.moe_k = 6 + text_config.moe_intermediate_size = [1536, 512] + text_config.moe_num_shared_experts = 2 + text_config.router_aux_loss_coef = 0.001 + text_config.mlp_layer_types = ["dense"] + ["moe"] * 27 + text_config.rope_parameters = {"rope_theta": 500000.0, "mrope_section": [22, 22, 20]} + text_config.rope_scaling = None + + config = Mock(spec=[]) + config.model_type = "ernie4_5_vl_moe" + config.architectures = ["Ernie4_5_VLForConditionalGeneration"] + config.text_config = text_config + config.vision_config = _make_vision_config() + config.tie_word_embeddings = True + config.image_token_id = 151859 + config.video_token_id = 151860 + config.image_start_token_id = 101304 + config.image_end_token_id = 101305 + config.video_start_token_id = 101306 + config.video_end_token_id = 101307 + return config + + +@pytest.fixture +def flat_config(): + return _make_flat_vl_config() + + +@pytest.fixture +def nested_config(): + return _make_nested_vl_config() + + +@pytest.fixture +def mock_pretrained_flat(flat_config): + pretrained = Mock() + pretrained.config = flat_config + return pretrained + + +@pytest.fixture +def mock_pretrained_nested(nested_config): + pretrained = Mock() + pretrained.config = nested_config + return pretrained + + +class TestErnie45VLBridgeRegistration: + """Test bridge class and registration.""" + + def test_is_subclass(self): + assert issubclass(Ernie45VLBridge, MegatronModelBridge) + + def test_instantiation(self): + bridge = Ernie45VLBridge() + assert bridge is not None + + +class TestErnie45VLBridgeGetTextConfig: + """Test _get_text_config static method.""" + + def test_nested_config_returns_text_config(self, nested_config): + result = Ernie45VLBridge._get_text_config(nested_config) + # nested_config.text_config is a distinct object + assert result is nested_config.text_config + assert result is not nested_config + + def test_flat_config_returns_self(self, flat_config): + result = Ernie45VLBridge._get_text_config(flat_config) + # flat config: text_config is self-reference + assert result is flat_config + + +class TestErnie45VLBridgeGetNumExperts: + """Test _get_num_experts static method.""" + + def test_list_input(self): + config = Mock(spec=[]) + config.moe_num_experts = [64, 64] + assert Ernie45VLBridge._get_num_experts(config) == 64 + + def test_int_input(self): + config = Mock(spec=[]) + config.moe_num_experts = 4 + assert Ernie45VLBridge._get_num_experts(config) == 4 + + +class TestErnie45VLBridgeProviderBridge: + """Test provider_bridge method.""" + + def test_returns_vl_provider_flat(self, mock_pretrained_flat): + bridge = Ernie45VLBridge() + provider = bridge.provider_bridge(mock_pretrained_flat) + assert isinstance(provider, Ernie45VLModelProvider) + + def test_returns_vl_provider_nested(self, mock_pretrained_nested): + bridge = Ernie45VLBridge() + provider = bridge.provider_bridge(mock_pretrained_nested) + assert isinstance(provider, Ernie45VLModelProvider) + + def test_basic_dimensions_flat(self, mock_pretrained_flat): + bridge = Ernie45VLBridge() + provider = bridge.provider_bridge(mock_pretrained_flat) + assert provider.num_layers == 28 + assert provider.hidden_size == 2560 + assert provider.num_attention_heads == 20 + + def test_moe_config_flat(self, mock_pretrained_flat): + bridge = Ernie45VLBridge() + provider = bridge.provider_bridge(mock_pretrained_flat) + assert provider.num_moe_experts == 64 + assert provider.moe_router_topk == 6 + assert provider.moe_ffn_hidden_size == 1536 # text expert intermediate + assert provider.moe_router_score_function == "sigmoid" + assert provider.moe_router_enable_expert_bias is True + + def test_dual_pool_intermediate_sizes(self, mock_pretrained_flat): + bridge = Ernie45VLBridge() + provider = bridge.provider_bridge(mock_pretrained_flat) + assert provider.moe_intermediate_size == (1536, 512) + + def test_shared_experts_flat(self, mock_pretrained_flat): + bridge = Ernie45VLBridge() + provider = bridge.provider_bridge(mock_pretrained_flat) + # 1536 * 2 = 3072 + assert provider.moe_shared_expert_intermediate_size == 1536 * 2 + + def test_rope_mrope(self, mock_pretrained_flat): + bridge = Ernie45VLBridge() + provider = bridge.provider_bridge(mock_pretrained_flat) + assert provider.position_embedding_type == "mrope" + assert provider.rotary_base == 500000.0 + assert provider.rotary_interleaved is True + assert provider.mrope_section == [22, 22, 20] + + def test_moe_layer_freq(self, mock_pretrained_flat): + bridge = Ernie45VLBridge() + provider = bridge.provider_bridge(mock_pretrained_flat) + expected = [0] + [1] * 27 + assert provider.moe_layer_freq == expected + + def test_token_ids_flat(self, mock_pretrained_flat): + bridge = Ernie45VLBridge() + provider = bridge.provider_bridge(mock_pretrained_flat) + assert provider.image_token_id == 151859 + assert provider.video_token_id == 151860 + assert provider.image_start_token_id == 101304 + assert provider.image_end_token_id == 101305 + + def test_tie_word_embeddings_nested(self, mock_pretrained_nested): + bridge = Ernie45VLBridge() + provider = bridge.provider_bridge(mock_pretrained_nested) + # nested config: tie_word_embeddings comes from top-level config + assert provider.share_embeddings_and_output_weights is True + + def test_tie_word_embeddings_flat(self, mock_pretrained_flat): + bridge = Ernie45VLBridge() + provider = bridge.provider_bridge(mock_pretrained_flat) + assert provider.share_embeddings_and_output_weights is False + + def test_normalization(self, mock_pretrained_flat): + bridge = Ernie45VLBridge() + provider = bridge.provider_bridge(mock_pretrained_flat) + assert provider.normalization == "RMSNorm" + + def test_vision_config_propagated(self, mock_pretrained_flat, flat_config): + bridge = Ernie45VLBridge() + provider = bridge.provider_bridge(mock_pretrained_flat) + assert provider.vision_config is flat_config.vision_config + + +class TestErnie45VLBridgeMappingRegistry: + """Test mapping_registry method.""" + + def _get_bridge_with_config(self, config, num_experts=64, is_flat=True): + """Create bridge with hf_config injected.""" + bridge = Ernie45VLBridge() + bridge.hf_config = config + return bridge + + def test_returns_mapping_registry_flat(self, flat_config): + bridge = self._get_bridge_with_config(flat_config) + registry = bridge.mapping_registry() + assert isinstance(registry, MegatronMappingRegistry) + assert len(registry.mappings) > 0 + + def test_returns_mapping_registry_nested(self, nested_config): + bridge = self._get_bridge_with_config(nested_config, num_experts=4, is_flat=False) + registry = bridge.mapping_registry() + assert isinstance(registry, MegatronMappingRegistry) + assert len(registry.mappings) > 0 + + def test_has_language_model_embeddings(self, flat_config): + bridge = self._get_bridge_with_config(flat_config) + registry = bridge.mapping_registry() + megatron_params = [ + m.megatron_param + for m in registry.mappings + if hasattr(m, "megatron_param") and isinstance(m.megatron_param, str) + ] + assert "language_model.embedding.word_embeddings.weight" in megatron_params + + def test_has_qkv_mapping(self, flat_config): + from megatron.bridge.models.conversion.param_mapping import QKVMapping + + bridge = self._get_bridge_with_config(flat_config) + registry = bridge.mapping_registry() + qkv_mappings = [m for m in registry.mappings if isinstance(m, QKVMapping)] + assert len(qkv_mappings) > 0 + + def test_has_vision_mappings(self, flat_config): + bridge = self._get_bridge_with_config(flat_config) + registry = bridge.mapping_registry() + hf_params = [] + for m in registry.mappings: + if hasattr(m, "hf_param") and isinstance(m.hf_param, str): + hf_params.append(m.hf_param) + vision_params = [p for p in hf_params if "vision" in p] + assert len(vision_params) > 0, "Should have vision encoder mappings" + + def test_has_offset_expert_mappings(self, flat_config): + bridge = self._get_bridge_with_config(flat_config) + registry = bridge.mapping_registry() + offset_mappings = [ + m for m in registry.mappings if isinstance(m, (_OffsetGatedMLPMapping, _OffsetRowParallelMapping)) + ] + assert len(offset_mappings) > 0, "Should have offset expert mappings for vision pool" + + def test_has_concat_bias_mapping(self, flat_config): + bridge = self._get_bridge_with_config(flat_config) + registry = bridge.mapping_registry() + concat_bias_mappings = [m for m in registry.mappings if isinstance(m, _ConcatBiasMapping)] + assert len(concat_bias_mappings) > 0, "Should have ConcatBiasMapping for expert bias" + + def test_has_shared_expert_mappings(self, flat_config): + bridge = self._get_bridge_with_config(flat_config) + registry = bridge.mapping_registry() + megatron_params = [ + m.megatron_param + for m in registry.mappings + if hasattr(m, "megatron_param") and isinstance(m.megatron_param, str) + ] + shared = [p for p in megatron_params if "shared_experts" in p] + assert len(shared) > 0 + + +class TestConcatBiasMapping: + """Test _ConcatBiasMapping class.""" + + def test_clear_export_buffer(self): + _ConcatBiasMapping.clear_export_buffer() + # Should not raise + + def test_hf_to_megatron_text_slice(self): + """Test that text slice extracts row 0 from [2, N] bias.""" + from unittest.mock import patch + + mapping = _ConcatBiasMapping( + megatron_param="language_model.decoder.layers.*.mlp.text_moe_layer.router.expert_bias", + hf_param="model.layers.*.mlp.moe_statics.e_score_correction_bias", + slice_name="text", + num_experts=4, + ) + # [2, 4] - row 0 is text, row 1 is vision + concat_bias = torch.randn(2, 4) + # Patch AutoMapping.hf_to_megatron to return the sliced input as-is + with patch( + "megatron.bridge.models.conversion.param_mapping.AutoMapping.hf_to_megatron", + side_effect=lambda w, m: w, + ): + result = mapping.hf_to_megatron(concat_bias, Mock()) + assert result.shape == (4,) + assert torch.allclose(result, concat_bias[0]) + + def test_hf_to_megatron_vision_slice(self): + """Test that vision slice extracts row 1 from [2, N] bias.""" + from unittest.mock import patch + + mapping = _ConcatBiasMapping( + megatron_param="language_model.decoder.layers.*.mlp.vision_moe_layer.router.expert_bias", + hf_param="model.layers.*.mlp.moe_statics.e_score_correction_bias", + slice_name="vision", + num_experts=4, + ) + concat_bias = torch.randn(2, 4) + with patch( + "megatron.bridge.models.conversion.param_mapping.AutoMapping.hf_to_megatron", + side_effect=lambda w, m: w, + ): + result = mapping.hf_to_megatron(concat_bias, Mock()) + assert result.shape == (4,) + assert torch.allclose(result, concat_bias[1]) + + def test_hf_to_megatron_squeeze_2d(self): + """Test slicing logic with text (row 0).""" + from unittest.mock import patch + + mapping = _ConcatBiasMapping( + megatron_param="language_model.decoder.layers.*.mlp.text_moe_layer.router.expert_bias", + hf_param="model.layers.*.mlp.moe_statics.e_score_correction_bias", + slice_name="text", + num_experts=4, + ) + concat_bias = torch.randn(2, 4) + with patch( + "megatron.bridge.models.conversion.param_mapping.AutoMapping.hf_to_megatron", + side_effect=lambda w, m: w, + ): + result = mapping.hf_to_megatron(concat_bias, Mock()) + assert result.shape == (4,) + + +class TestOffsetMappings: + """Test _OffsetGatedMLPMapping and _OffsetRowParallelMapping.""" + + def test_offset_gated_mlp_mapping_instantiation(self): + mapping = _OffsetGatedMLPMapping( + megatron_param="language_model.decoder.layers.*.mlp.vision_moe_layer.experts.local_experts.*.linear_fc1.weight", + gate="model.layers.*.mlp.experts.*.gate_proj.weight", + up="model.layers.*.mlp.experts.*.up_proj.weight", + expert_offset=64, + ) + assert mapping is not None + assert mapping._expert_offset == 64 + + def test_offset_row_parallel_mapping_instantiation(self): + mapping = _OffsetRowParallelMapping( + megatron_param="language_model.decoder.layers.*.mlp.vision_moe_layer.experts.local_experts.*.linear_fc2.weight", + hf_param="model.layers.*.mlp.experts.*.down_proj.weight", + expert_offset=64, + ) + assert mapping is not None + assert mapping._expert_offset == 64 + + def test_offset_gated_mlp_mapping_resolve(self): + """Test that resolve offsets the expert index.""" + mapping = _OffsetGatedMLPMapping( + megatron_param="language_model.decoder.layers.*.mlp.vision_moe_layer.experts.local_experts.*.linear_fc1.weight", + gate="model.layers.*.mlp.experts.*.gate_proj.weight", + up="model.layers.*.mlp.experts.*.up_proj.weight", + expert_offset=64, + ) + # Resolve with captures: (layer_idx, local_expert_idx) + resolved = mapping.resolve(("1", "0")) + # The HF expert index should be offset: local_expert 0 -> HF expert 64 + assert "64" in resolved.hf_param["gate"] + assert "64" in resolved.hf_param["up"] + + def test_offset_row_parallel_mapping_resolve(self): + """Test that resolve offsets the expert index for row parallel.""" + mapping = _OffsetRowParallelMapping( + megatron_param="language_model.decoder.layers.*.mlp.vision_moe_layer.experts.local_experts.*.linear_fc2.weight", + hf_param="model.layers.*.mlp.experts.*.down_proj.weight", + expert_offset=64, + ) + resolved = mapping.resolve(("1", "0")) + # The resolved HF param should reference expert 64 + assert "64" in resolved.hf_param + + +class TestErnie45VLModelProvider: + """Test Ernie45VLModelProvider class.""" + + def test_instantiation(self): + provider = Ernie45VLModelProvider() + assert provider is not None + + def test_has_expected_fields(self): + provider = Ernie45VLModelProvider() + # Should have standard GPT fields + assert hasattr(provider, "num_layers") + assert hasattr(provider, "hidden_size") + assert hasattr(provider, "num_attention_heads") diff --git a/tests/unit_tests/models/test_param_mapping.py b/tests/unit_tests/models/test_param_mapping.py index 6050ea28ff..f7208c139b 100644 --- a/tests/unit_tests/models/test_param_mapping.py +++ b/tests/unit_tests/models/test_param_mapping.py @@ -890,17 +890,18 @@ def test_transpose_non_rank_zero_hf_to_megatron(self, mock_distributed_env, tran mapping = AutoMapping("transpose.weight", "hf.weight", permute_dims=(1, 0)) hf_weight = torch.randn(4, 8) - megatron_module = MockModule(transformer_config, weight_shape=(4, 4)) + megatron_module = MockModule(transformer_config, weight_shape=(8, 4)) with patch.object(mapping, "_mapping") as mock_delegate: - mock_delegate.hf_to_megatron.return_value = torch.randn(4, 4) + mock_delegate.hf_to_megatron.return_value = torch.randn(8, 4) with patch.object(mapping, "_detect_parallelism_type", return_value="column"): mapping.hf_to_megatron(hf_weight, megatron_module) - # On non-rank-0, permutation is skipped, original tensor passed to delegate + # Permutation is applied on ALL ranks so delegate mappings + # (e.g. ReplicatedMapping) always receive the correct shape. mock_delegate.hf_to_megatron.assert_called_once() passed_tensor = mock_delegate.hf_to_megatron.call_args[0][0] - assert torch.equal(passed_tensor, hf_weight) + assert torch.equal(passed_tensor, hf_weight.permute(1, 0).contiguous()) def test_transpose_identity_permutation(self, mock_distributed_env, transformer_config): """Test AutoMapping with identity permutation."""