feat(model): add model support verification cards

docs/fern/versions/nightly/pages/models/qwen/qwen.mdx

@@ -218,9 +218,9 @@ uv run python examples/conversion/hf_to_megatron_generate_text.py \
 
 | Model | Mode | TP | PP | EP | Total GPUs | Use Case |
 |-------|------|----|----|----|-----------:|----------|
-| **Qwen3-30B-A3B** | Pretrain | 1 | 1 | 8 | 8 | Pre-training (single node) |
-| **Qwen3-30B-A3B** | Full SFT | 1 | 1 | 8 | 8 | Full supervised finetuning |
-| **Qwen3-30B-A3B** | LoRA/DoRA | 1 | 1 | 8 | 8 | PEFT finetuning (single node) |
+| **Qwen3-30B-A3B** | Pretrain | 1 | 1 | 16 | 16 | Pre-training (2 nodes) |
+| **Qwen3-30B-A3B** | Full SFT | 1 | 1 | 16 | 16 | Full supervised finetuning (2 nodes) |
+| **Qwen3-30B-A3B** | LoRA/DoRA | 4 | 1 | 4 | 4 | PEFT finetuning |
 | **Qwen3-235B-A22B** | Pretrain | 2 | 8 | 32 | 512 | Pre-training (64 nodes) |
 | **Qwen3-235B-A22B** | Full SFT | 2 | 8 | 32 | 512 | Full supervised finetuning (64 nodes) |
 | **Qwen3-235B-A22B** | LoRA/DoRA | 2 | 8 | 32 | 512 | PEFT finetuning (64 nodes) |
@@ -239,7 +239,7 @@ config = qwen3_30b_a3b_pretrain_config(
     train_iters=500_000,
     global_batch_size=2048,
     seq_length=4096,
-    # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically
+    # Uses TP=1, PP=1, EP=16 (16 GPUs) automatically
 )
 ```
 
@@ -272,7 +272,7 @@ config = qwen3_30b_a3b_sft_config(
     train_iters=1000,
     global_batch_size=64,
     finetune_lr=5e-6,
-    # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically
+    # Uses TP=1, PP=1, EP=16 (16 GPUs) automatically
 )
 ```
 
@@ -288,7 +288,7 @@ config = qwen3_30b_a3b_peft_config(
     train_iters=1000,
     global_batch_size=128,
     finetune_lr=1e-4,
-    # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically
+    # Uses TP=4, PP=1, EP=4 (4 GPUs) automatically
 )
 ```
 

docs/models/qwen/qwen.md

@@ -218,9 +218,9 @@ uv run python examples/conversion/hf_to_megatron_generate_text.py \
 
 | Model | Mode | TP | PP | EP | Total GPUs | Use Case |
 |-------|------|----|----|----|-----------:|----------|
-| **Qwen3-30B-A3B** | Pretrain | 1 | 1 | 8 | 8 | Pre-training (single node) |
-| **Qwen3-30B-A3B** | Full SFT | 1 | 1 | 8 | 8 | Full supervised finetuning |
-| **Qwen3-30B-A3B** | LoRA/DoRA | 1 | 1 | 8 | 8 | PEFT finetuning (single node) |
+| **Qwen3-30B-A3B** | Pretrain | 1 | 1 | 16 | 16 | Pre-training (2 nodes) |
+| **Qwen3-30B-A3B** | Full SFT | 1 | 1 | 16 | 16 | Full supervised finetuning (2 nodes) |
+| **Qwen3-30B-A3B** | LoRA/DoRA | 4 | 1 | 4 | 4 | PEFT finetuning |
 | **Qwen3-235B-A22B** | Pretrain | 2 | 8 | 32 | 512 | Pre-training (64 nodes) |
 | **Qwen3-235B-A22B** | Full SFT | 2 | 8 | 32 | 512 | Full supervised finetuning (64 nodes) |
 | **Qwen3-235B-A22B** | LoRA/DoRA | 2 | 8 | 32 | 512 | PEFT finetuning (64 nodes) |
@@ -239,7 +239,7 @@ config = qwen3_30b_a3b_pretrain_config(
     train_iters=500_000,
     global_batch_size=2048,
     seq_length=4096,
-    # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically
+    # Uses TP=1, PP=1, EP=16 (16 GPUs) automatically
 )
 ```
 
@@ -272,7 +272,7 @@ config = qwen3_30b_a3b_sft_config(
     train_iters=1000,
     global_batch_size=64,
     finetune_lr=5e-6,
-    # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically
+    # Uses TP=1, PP=1, EP=16 (16 GPUs) automatically
 )
 ```
 
@@ -288,7 +288,7 @@ config = qwen3_30b_a3b_peft_config(
     train_iters=1000,
     global_batch_size=128,
     finetune_lr=1e-4,
-    # Uses TP=1, PP=1, EP=8 (8 GPUs) automatically
+    # Uses TP=4, PP=1, EP=4 (4 GPUs) automatically
 )
 ```
 

examples/conversion/compare_hf_and_megatron/compare.py

@@ -100,6 +100,8 @@
 import torch
 import torch.distributed as dist
 from megatron.core import parallel_state
+from megatron.core.inference.contexts import StaticInferenceContext
+from megatron.core.inference.utils import InferenceMode
 from megatron.core.pipeline_parallel.schedules import get_forward_backward_func
 from transformers import AutoConfig, AutoModelForCausalLM, AutoProcessor, AutoTokenizer
 
@@ -264,11 +266,20 @@ class SingleBatchIterator:
     then raises StopIteration. Used for single-step inference in the forward pass.
     """
 
-    def __init__(self, input_ids, position_ids, attention_mask, pixel_values=None, image_grid_thw=None):
+    def __init__(
+        self,
+        input_ids,
+        position_ids,
+        attention_mask,
+        pixel_values=None,
+        image_grid_thw=None,
+        inference_context=None,
+    ):
         self.batch = dict(
             tokens=input_ids,
             position_ids=position_ids,
             attention_mask=attention_mask,
+            inference_context=inference_context,
         )
 
         # Add vision inputs if provided
@@ -329,6 +340,47 @@ def loss_func(x, **kwargs):
     return output_tensor, loss_func
 
 
+def inference_forward_step(data_iterator, model, **kwargs) -> torch.Tensor:
+    """Run a text-model forward step with an explicit inference context."""
+    batch = next(data_iterator)
+
+    def loss_func(x, **kwargs):
+        return x
+
+    model_output = model(
+        input_ids=batch["tokens"],
+        position_ids=batch["position_ids"],
+        attention_mask=batch.get("attention_mask"),
+        inference_context=batch["inference_context"],
+        runtime_gather_output=True,
+    )
+    if isinstance(model_output, tuple):
+        model_output = model_output[0]
+    return model_output, loss_func
+
+
+def _run_megatron_forward(fwd_bwd_function, **kwargs):
+    """Run a Megatron forward pass with the inference execution paths active."""
+    with InferenceMode.active():
+        return fwd_bwd_function(**kwargs)
+
+
+def _maybe_gather_tensor_parallel_logits(megatron_output, hf_vocab_size: int, world_size: int, group):
+    """Gather sharded TP logits while preserving an already gathered full-vocabulary tensor."""
+    if megatron_output.size(-1) >= hf_vocab_size:
+        return megatron_output
+
+    gathered_tensors = [torch.zeros_like(megatron_output) for _ in range(world_size)]
+    dist.all_gather(gathered_tensors, megatron_output, group=group)
+    gathered_output = torch.cat(gathered_tensors, dim=2)
+    if gathered_output.size(-1) < hf_vocab_size:
+        raise ValueError(
+            f"Gathered Megatron vocabulary ({gathered_output.size(-1)}) is smaller than "
+            f"the Hugging Face vocabulary ({hf_vocab_size})."
+        )
+    return gathered_output
+
+
 def load_image(image_path: str) -> Image.Image:
     """Load an image from URL or file path.
 
@@ -671,6 +723,29 @@ def _setup_tokenizer_and_processor(args, is_vl_model: bool):
     return tokenizer, processor
 
 
+def _broadcast_hf_results(hf_logits, hf_next_token, device):
+    """Broadcast rank-0 HF results using the model's actual output vocabulary size."""
+    if hf_logits is not None:
+        hf_logits = hf_logits.float()
+
+    hf_logits_size = torch.tensor(
+        [hf_logits.numel() if hf_logits is not None else 0],
+        device=device,
+        dtype=torch.long,
+    )
+    torch.distributed.broadcast(hf_logits_size, 0)
+
+    if hf_next_token is None:
+        hf_next_token = torch.zeros(1, device=device, dtype=torch.long)
+    if hf_logits is None:
+        hf_logits = torch.zeros(hf_logits_size.item(), device=device, dtype=torch.float32)
+
+    torch.distributed.broadcast(hf_next_token, 0)
+    torch.distributed.broadcast(hf_logits, 0)
+    torch.distributed.barrier()
+    return hf_logits, hf_next_token
+
+
 def compare_models_one_step(args) -> None:
     """Compare 1-step generation between HF and Megatron models with debugging.
 
@@ -733,28 +808,10 @@ def compare_models_one_step(args) -> None:
 
     # Broadcast HF results to all ranks
     if torch.distributed.is_initialized():
-        # Ensure consistent dtype across ranks: rank 0 has bfloat16 logits from the HF model,
-        # so all ranks must use the same dtype for NCCL broadcast to work correctly.
-        if hf_logits is not None:
-            hf_logits = hf_logits.float()
-
-        # Create tensors for broadcasting if they don't exist on non-rank-0
-        if hf_next_token is None:
-            hf_next_token = torch.zeros(1, device=input_ids.device, dtype=torch.long)
-        if hf_logits is None:
-            # Get vocab size from tokenizer for proper tensor size
-            vocab_size = getattr(
-                tokenizer, "vocab_size", len(tokenizer.vocab) if hasattr(tokenizer, "vocab") else 32000
-            )
-            hf_logits = torch.zeros(vocab_size, device=input_ids.device, dtype=torch.float32)
-
-        # Ensure consistent dtype across ranks before broadcast
-        hf_logits = hf_logits.float()
-
-        # Broadcast from rank 0 to all ranks
-        torch.distributed.broadcast(hf_next_token, 0)
-        torch.distributed.broadcast(hf_logits, 0)
-        torch.distributed.barrier()
+        # The model's output vocabulary can be larger than the tokenizer vocabulary.
+        # Broadcast the actual logits length before allocating receive buffers so every
+        # rank participates in the logits broadcast with the same tensor shape.
+        hf_logits, hf_next_token = _broadcast_hf_results(hf_logits, hf_next_token, input_ids.device)
         print_rank_0("HF results broadcast complete.")
 
     # Run Megatron model forward pass
@@ -773,18 +830,44 @@ def compare_models_one_step(args) -> None:
         attention_mask = None
 
         fwd_bwd_function = get_forward_backward_func()
-        iterator = SingleBatchIterator(input_ids, position_ids, attention_mask, pixel_values, image_grid_thw)
-
-        megatron_output = fwd_bwd_function(
-            forward_step_func=vlm_forward_step,
-            data_iterator=iterator,
-            model=megatron_model,
-            num_microbatches=1,
-            forward_only=True,
-            seq_length=input_ids.size(1),
-            micro_batch_size=1,
-            collect_non_loss_data=True,
-        )
+        forward_kwargs = {
+            "model": megatron_model,
+            "num_microbatches": 1,
+            "forward_only": True,
+            "seq_length": input_ids.size(1),
+            "micro_batch_size": 1,
+            "collect_non_loss_data": True,
+        }
+        if is_vl_model:
+            iterator = SingleBatchIterator(
+                input_ids,
+                position_ids,
+                attention_mask,
+                pixel_values,
+                image_grid_thw,
+            )
+            megatron_output = fwd_bwd_function(
+                forward_step_func=vlm_forward_step,
+                data_iterator=iterator,
+                **forward_kwargs,
+            )
+        else:
+            inference_context = StaticInferenceContext(
+                max_batch_size=input_ids.size(0),
+                max_sequence_length=input_ids.size(1),
+            )
+            iterator = SingleBatchIterator(
+                input_ids,
+                position_ids,
+                attention_mask,
+                inference_context=inference_context,
+            )
+            megatron_output = _run_megatron_forward(
+                fwd_bwd_function,
+                forward_step_func=inference_forward_step,
+                data_iterator=iterator,
+                **forward_kwargs,
+            )
 
         if isinstance(megatron_output, list) and len(megatron_output) > 0:
             megatron_output = megatron_output[0]
@@ -796,11 +879,12 @@ def compare_models_one_step(args) -> None:
             # Gather tensor parallel results if using TP
             if torch.distributed.is_initialized() and parallel_state.get_tensor_model_parallel_world_size() > 1:
                 world_size = parallel_state.get_tensor_model_parallel_world_size()
-                gathered_tensors = [torch.zeros_like(megatron_output) for _ in range(world_size)]
-                dist.all_gather(
-                    gathered_tensors, megatron_output, group=parallel_state.get_tensor_model_parallel_group()
+                megatron_output = _maybe_gather_tensor_parallel_logits(
+                    megatron_output,
+                    hf_logits.size(0),
+                    world_size,
+                    parallel_state.get_tensor_model_parallel_group(),
                 )
-                megatron_output = torch.cat(gathered_tensors, dim=2)
 
             megatron_logits = megatron_output[0, -1, :]
             megatron_next_token = torch.argmax(megatron_logits, dim=-1)