diff --git a/examples/auto_deploy/cookbooks/step_3.7_flash_trtllm_cookbook.ipynb b/examples/auto_deploy/cookbooks/step_3.7_flash_trtllm_cookbook.ipynb new file mode 100644 index 000000000000..c10eb8a6415f --- /dev/null +++ b/examples/auto_deploy/cookbooks/step_3.7_flash_trtllm_cookbook.ipynb @@ -0,0 +1,267 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Deploying Step-3.7-Flash with TensorRT-LLM\n", + "\n", + "This notebook walks you through deploying the `stepfun-ai/Step-3.7-Flash` model (text generation path) using TensorRT-LLM.\n", + "\n", + "[TensorRT-LLM](https://nvidia.github.io/TensorRT-LLM/) is NVIDIA's open-source library for accelerating and optimizing LLM inference on NVIDIA GPUs. Support for Step-3.7-Flash is enabled through the AutoDeploy workflow. More details about AutoDeploy can be found [here](https://nvidia.github.io/TensorRT-LLM/torch/auto_deploy/auto-deploy.html).\n", + "\n", + "**Model Resources:**\n", + "- [HuggingFace Model Card](https://huggingface.co/stepfun-ai/Step-3.7-Flash)\n", + "- [Technical Blog](https://static.stepfun.com/blog/step-3.7-flash/)\n", + "- [StepFun API Platform](https://platform.stepfun.ai)\n", + "- [Discord Community](https://discord.gg/RcMJhNVAQc)\n", + "\n", + "**Model Highlights:**\n", + "- 198B-parameter sparse Mixture-of-Experts (MoE) model, ~11B active parameters per token\n", + "- 288 routed experts (top-8) plus a shared expert; 45 decoder layers\n", + "- Mixed full / sliding-window grouped-query attention with a head-wise attention gate\n", + "- 256K token context length\n", + "- Three reasoning levels (low / medium / high) and tool-calling support\n", + "- Apache 2.0 License\n", + "\n", + "> **Note:** Step-3.7-Flash is a vision-language model. AutoDeploy onboards the **text generation path**; the vision encoder is not deployed through this workflow.\n", + "\n", + "**Prerequisites:**\n", + "- 8x NVIDIA GPUs with recent drivers (BF16 weights are ~400 GB total) and CUDA 12.x\n", + "- Python 3.10+\n", + "- TensorRT-LLM ([container](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tensorrt-llm/containers/release) or pip install)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Prerequisites & Environment\n", + "\n", + "Set up a containerized environment for TensorRT-LLM by running the following command in a terminal:\n", + "\n", + "```shell\n", + "docker run --rm -it --ipc=host --ulimit memlock=-1 --ulimit stack=67108864 --gpus=all -p 8000:8000 nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc1\n", + "```\n", + "\n", + "You now have TensorRT-LLM set up!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# If pip not found\n", + "!python -m ensurepip --default-pip" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "%pip install torch openai" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Verify GPU\n", + "\n", + "Check that CUDA is available and the GPUs are detected correctly. Step-3.7-Flash requires 8 GPUs for the BF16 weights." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Environment check\n", + "import sys\n", + "\n", + "import torch\n", + "\n", + "print(f\"Python: {sys.version}\")\n", + "print(f\"CUDA available: {torch.cuda.is_available()}\")\n", + "print(f\"Num GPUs: {torch.cuda.device_count()}\")\n", + "\n", + "if torch.cuda.is_available():\n", + " for i in range(torch.cuda.device_count()):\n", + " print(f\"GPU[{i}]: {torch.cuda.get_device_name(i)}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## OpenAI-Compatible Server\n", + "\n", + "Start a local OpenAI-compatible server with TensorRT-LLM via the terminal, within the running docker container.\n", + "\n", + "Ensure that the following commands are executed from the docker terminal.\n", + "\n", + "Start with the Step-3.7-Flash YAML here: `examples/auto_deploy/model_registry/configs/step-3.7-flash.yaml`" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Load the Model\n", + "\n", + "Launch the TensorRT-LLM server with Step-3.7-Flash:\n", + "\n", + "```shell\n", + "trtllm-serve \"stepfun-ai/Step-3.7-Flash\" \\\n", + " --host 0.0.0.0 \\\n", + " --port 8000 \\\n", + " --backend _autodeploy \\\n", + " --trust_remote_code \\\n", + " --extra_llm_api_options examples/auto_deploy/model_registry/configs/step-3.7-flash.yaml\n", + "```\n", + "\n", + "The same standalone config YAML can be validated locally with `build_and_run_ad.py`:\n", + "\n", + "```shell\n", + "python examples/auto_deploy/build_and_run_ad.py \\\n", + " --model stepfun-ai/Step-3.7-Flash \\\n", + " --args.yaml-extra examples/auto_deploy/model_registry/configs/step-3.7-flash.yaml\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Your server is now running!" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Use the API\n", + "\n", + "Use the OpenAI-compatible client to send requests to the TensorRT-LLM server." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from openai import OpenAI\n", + "\n", + "# Setup client\n", + "BASE_URL = \"http://0.0.0.0:8000/v1\"\n", + "API_KEY = \"null\"\n", + "client = OpenAI(base_url=BASE_URL, api_key=API_KEY)\n", + "\n", + "MODEL_ID = \"stepfun-ai/Step-3.7-Flash\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Basic chat completion\n", + "print(\"Chat Completion Example\")\n", + "print(\"=\" * 50)\n", + "\n", + "response = client.chat.completions.create(\n", + " model=MODEL_ID,\n", + " messages=[\n", + " {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n", + " {\"role\": \"user\", \"content\": \"What is 15% of 85? Show your reasoning.\"},\n", + " ],\n", + " temperature=1.0,\n", + " top_p=0.95,\n", + " max_tokens=512,\n", + ")\n", + "\n", + "print(\"Response:\")\n", + "print(response.choices[0].message.content)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Streaming chat completion\n", + "print(\"Streaming response:\")\n", + "print(\"=\" * 50)\n", + "\n", + "stream = client.chat.completions.create(\n", + " model=MODEL_ID,\n", + " messages=[\n", + " {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n", + " {\"role\": \"user\", \"content\": \"What are the first 5 prime numbers?\"},\n", + " ],\n", + " temperature=0.7,\n", + " max_tokens=1024,\n", + " stream=True,\n", + ")\n", + "\n", + "for chunk in stream:\n", + " if chunk.choices[0].delta.content:\n", + " print(chunk.choices[0].delta.content, end=\"\", flush=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Evaluation Parameters\n", + "\n", + "For optimal results, use the following parameters based on your task:\n", + "\n", + "**Default Settings (Reasoning Tasks)**\n", + "- `temperature`: 1.0\n", + "- `top_p`: 0.95\n", + "- `max_tokens`: up to 8192 (model supports a 256K context)\n", + "\n", + "**Deterministic Tasks**\n", + "- `temperature`: 0\n", + "- `max_tokens`: 16384\n", + "\n", + "Step-3.7-Flash exposes three reasoning levels (low / medium / high); see the model card for how to select them via the chat template." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Additional Resources\n", + "\n", + "- [TensorRT-LLM Documentation](https://nvidia.github.io/TensorRT-LLM/)\n", + "- [AutoDeploy Guide](https://nvidia.github.io/TensorRT-LLM/torch/auto_deploy/auto-deploy.html)\n", + "- [Step-3.7-Flash on HuggingFace](https://huggingface.co/stepfun-ai/Step-3.7-Flash)\n", + "- [Step-3.7-Flash Technical Blog](https://static.stepfun.com/blog/step-3.7-flash/)\n", + "- [StepFun Discord Community](https://discord.gg/RcMJhNVAQc)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/auto_deploy/model_registry/configs/step-3.7-flash.yaml b/examples/auto_deploy/model_registry/configs/step-3.7-flash.yaml new file mode 100644 index 000000000000..caaa6206eff4 --- /dev/null +++ b/examples/auto_deploy/model_registry/configs/step-3.7-flash.yaml @@ -0,0 +1,42 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +# Self-contained AutoDeploy config for StepFun Step-3.7-Flash (text decoder). +# +# Step-3.7-Flash is a vision-language model; AutoDeploy onboards the text generation +# path (model_type step3p7 / text config step3p5). The custom modeling code lives in +# tensorrt_llm/_torch/auto_deploy/models/custom/modeling_step3p7.py and is selected by +# registering Step3p7ForCausalLM under the Step3p7Config config class. +# +# 45 decoder layers (3 dense + 42 MoE, 288 routed experts top-8 + shared expert), +# GQA with mixed full/sliding attention and a head-wise attention gate. BF16 weights. +runtime: trtllm +compile_backend: torch-cudagraph +model_factory: AutoModelForCausalLM +attn_backend: flashinfer +max_seq_len: 4096 +max_num_tokens: 4096 +max_batch_size: 64 +world_size: 8 +enable_chunked_prefill: true +cuda_graph_config: + batch_sizes: [1, 2, 4, 8, 16, 32, 64] +kv_cache_config: + dtype: bfloat16 + enable_block_reuse: false + free_gpu_memory_fraction: 0.7 + tokens_per_block: 64 +model_kwargs: + torch_dtype: bfloat16 +# Hint-driven sharding IR: the modeling code carries explicit sharding hints +# (tp_mode / layer_type / tp_scaled_dim / all_reduce), so disable the legacy +# heuristic sharding and apply the hints instead. +transforms: + detect_sharding: + enabled: false + sharding_transform_executor: + enabled: false + apply_sharding_hints: + enabled: true + gather_logits_before_lm_head: + enabled: true diff --git a/examples/auto_deploy/model_registry/models.yaml b/examples/auto_deploy/model_registry/models.yaml index e71fef61adaa..5ecfd108f746 100644 --- a/examples/auto_deploy/model_registry/models.yaml +++ b/examples/auto_deploy/model_registry/models.yaml @@ -404,6 +404,18 @@ models: # - name: XiaomiMiMo/MiMo-V2-Flash # config_id: default_ws_8 # yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml'] +# --- StepFun Step-3.7-Flash (Feb 2026) --- +# Text decoder of the Step-3.7-Flash VLM (model_type step3p7 / step3p5). 45 layers, +# 288 routed experts top-8 + shared expert, mixed full/sliding GQA, head-wise attn gate. +- name: stepfun-ai/Step-3.7-Flash + config_id: step_3_7_flash + yaml_extra: ['step-3.7-flash.yaml'] +- name: stepfun-ai/Step-3.7-Flash-FP8 + config_id: step_3_7_flash + yaml_extra: ['step-3.7-flash.yaml'] +- name: stepfun-ai/Step-3.7-Flash-NVFP4 + config_id: step_3_7_flash + yaml_extra: ['step-3.7-flash.yaml'] # --- Kimi-K2.5 (Jan 2026) --- # TypeError: 'NoneType' object is not subscriptable # - name: moonshotai/Kimi-K2.5 diff --git a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py index 529556729c27..39042578f7e2 100644 --- a/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py +++ b/tensorrt_llm/_torch/auto_deploy/custom_ops/attention/flashinfer_attention.py @@ -416,6 +416,11 @@ def flashinfer_mha_with_cache( kv_layout=_GlobalFlashInferPlanner.kv_layout, ) + if num_prefill > 0 and not read_cache_only: + # FlashInfer planning depends on the preceding paged-KV append completing. + # Keep the same append-before-plan synchronization used by the PyTorch backend. + torch.cuda.current_stream().synchronize() + bs = b * s if out is not None: y = out.view(-1, n_heads, head_dim) diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py b/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py index e90767977c06..236da7375c8e 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py @@ -63,6 +63,7 @@ "modeling_skywork_r1v2": ["SkyworkR1V2ForConditionalGeneration"], "modeling_smollm3": ["SmolLM3ForCausalLM"], "modeling_starcoder2": ["Starcoder2ForCausalLM"], + "modeling_step3p7": ["Step3p7ForCausalLM"], } # AD_USE_IR_MODELS: opt-in flag for staging additional ``_ir.py`` modeling diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_step3p7.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_step3p7.py new file mode 100644 index 000000000000..33831a0dcc12 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_step3p7.py @@ -0,0 +1,841 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Slimmed-down PyTorch StepFun Step-3.7-Flash text model for AutoDeploy export (prefill only). + +Source: +https://huggingface.co/stepfun-ai/Step-3.7-Flash-FP8 (and bf16 sibling Step-3.7-Flash) + +Step-3.7-Flash is a vision-language model. This file ports ONLY the text decoder +(``Step3p5``-style ``model_type="step3p5"``); the vision tower is intentionally not +exported (AutoDeploy onboards the text generation path). + +Key text-architecture features: +* Per-layer attention type: ``full_attention`` and ``sliding_attention`` alternate + (1 full + 3 sliding per group). Full-attention layers use 64 Q heads; sliding-attention + layers use 96 Q heads. Both use 8 KV heads (GQA) and head_dim=128. +* Head-wise attention gate (``g_proj``): the attention output of each head is multiplied by + ``sigmoid(g_proj(hidden_states))`` before the output projection. +* Per-head QK RMSNorm over head_dim (Qwen3-style). +* Per-layer-type partial RoPE: full-attention layers rotate the first half of head_dim + (partial_rotary_factor=0.5, rope_theta=5e6, llama3 rope-scaling); sliding-attention layers + rotate the full head_dim (partial_rotary_factor=1.0, rope_theta=1e4, no scaling). +* Dense SwiGLU MLP on the first ``len(layers) - len(moe_layers)`` layers (layers 0-2); + the remaining layers are MoE (288 routed experts, top-8, sigmoid routing with a per-expert + bias used for *selection only*, fp32 gate, scaling 3.0) plus a dense shared expert. +* Gemma-style ``(1 + weight)`` RMSNorm convention for ALL norms (absorbed into the weight at + load time via a pre-hook so the graph uses plain ``torch_rmsnorm``). + +Differences from the HF reference (modeling_step3p7.py): +* Vision tower, multimodal merging, KV cache, training paths, dropout, and the MTP/spec + ``mtp_block`` layers (45-47) are all removed — prefill text decode only. +* Uses AD canonical ops: torch_rmsnorm, torch_attention, torch_rope_with_explicit_cos_sin, + torch_moe. No repeat_kv (torch_attention handles GQA natively). +* Stacked checkpoint MoE expert weights are split into per-expert Linear modules via a + load-state-dict pre-hook for torch_moe dispatch. +* The SwiGLU activation clamp (``swiglu_limits``) present on routed experts of the last two + MoE layers is NOT applied (the clamp limits are large numerical guards; see note in the + MoE block). It is still applied on the dense shared-expert path where it is a plain MLP. + +Tensor-parallel sharding (sharding-IR hints): +* Every shardable projection uses ``torch.ops.auto_deploy.torch_linear_simple`` with explicit + ``tp_mode`` / ``layer_type`` hints, head reshapes use ``torch.ops.auto_deploy.view`` with + ``tp_scaled_dim``, and rowwise outputs are followed by ``torch.ops.auto_deploy.all_reduce``. + The exported graph fully specifies sharding; ``apply_sharding_hints`` applies it. +* MHA: q/k/v/g colwise (k/v use ``tp_min_local_shape=head_dim`` for GQA; the head-wise gate + ``g_proj`` is a per-head column shard, ``tp_min_local_shape=1``), o_proj rowwise + all_reduce. +* MoE: routed experts via ``torch_moe(layer_type="moe")`` (EP/TP handled by the sharder); the + shared expert is a colwise/rowwise MLP with no internal all_reduce — a single all_reduce at + the ``routed + shared`` merge point covers both. Dense MLP layers reduce internally. +* The fp32 router gate is TP-replicated (kept as plain ``F.linear``). +""" + +import math +from dataclasses import dataclass +from typing import List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.activations import ACT2FN +from transformers.configuration_utils import PretrainedConfig +from transformers.generation import GenerationMixin +from transformers.modeling_utils import PreTrainedModel +from transformers.utils import ModelOutput + +from ... import custom_ops # noqa: F401 -- register all sharding-aware ops +from ..._compat import ActivationType +from ..hf import AutoModelForCausalLMFactory +from .rotary_utils import RotaryEmbeddingBase, build_rope_cos_sin_cache + +# --------------------------------------------------------------------------- +# Bundled config +# --------------------------------------------------------------------------- + + +class Step3p7Config(PretrainedConfig): + """Minimal flat text config for Step-3.7-Flash. + + Real deployments load the model's ``trust_remote_code`` config (the VLM wrapper + ``Step3p7Config`` with a nested ``text_config``); AutoDeploy passes that object straight to + ``_from_config`` and the model reads ``config.text_config`` via ``_get_text_config``. This + bundled class is the resolvable config the model registers under, used for standalone + construction and the offline sharding-IR equivalence harness (which builds a tiny instance and + overrides the universal dims). Its defaults are intentionally small and tensor-parallel + friendly so a 4-layer / 4-head tiny model shards cleanly; production values come from the + checkpoint config. + """ + + model_type = "step3p5" + + def __init__( + self, + vocab_size: int = 128896, + hidden_size: int = 64, + head_dim: int = 16, + num_attention_heads: int = 4, + num_attention_groups: int = 4, + attention_other_setting: Optional[dict] = None, + intermediate_size: int = 64, + num_hidden_layers: int = 4, + layer_types: Optional[list] = None, + moe_layers_enum: tuple = (2, 3), + moe_num_experts: int = 8, + moe_top_k: int = 2, + moe_intermediate_size: int = 16, + share_expert_dim: int = 16, + moe_router_scaling_factor: float = 3.0, + rms_norm_eps: float = 1e-5, + sliding_window: int = 4, + max_position_embeddings: int = 256, + rope_theta=(5e6, 1e4, 5e6, 1e4), + partial_rotary_factors=(0.5, 1.0, 0.5, 1.0), + rope_scaling: Optional[dict] = None, + yarn_only_types=("full_attention",), + swiglu_limits: Optional[list] = None, + swiglu_limits_shared: Optional[list] = None, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.head_dim = head_dim + self.num_attention_heads = num_attention_heads + self.num_attention_groups = num_attention_groups + self.attention_other_setting = attention_other_setting or { + "num_attention_heads": num_attention_heads, + "num_attention_groups": num_attention_groups, + "head_dim": head_dim, + } + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.layer_types = layer_types or [ + "full_attention" if i % 2 == 0 else "sliding_attention" + for i in range(num_hidden_layers) + ] + self.moe_layers_enum = moe_layers_enum + self.moe_num_experts = moe_num_experts + self.moe_top_k = moe_top_k + self.moe_intermediate_size = moe_intermediate_size + self.share_expert_dim = share_expert_dim + self.moe_router_scaling_factor = moe_router_scaling_factor + self.rms_norm_eps = rms_norm_eps + self.sliding_window = sliding_window + self.max_position_embeddings = max_position_embeddings + self.rope_theta = list(rope_theta) + self.partial_rotary_factors = list(partial_rotary_factors) + self.rope_scaling = rope_scaling or { + "rope_type": "llama3", + "factor": 2.0, + "original_max_position_embeddings": max_position_embeddings, + "low_freq_factor": 1.0, + "high_freq_factor": 32.0, + } + self.yarn_only_types = list(yarn_only_types) + self.swiglu_limits = swiglu_limits + self.swiglu_limits_shared = swiglu_limits_shared + super().__init__(**kwargs) + + +# --------------------------------------------------------------------------- +# Output dataclasses +# --------------------------------------------------------------------------- + + +@dataclass +class Step3p7ModelOutput(ModelOutput): + last_hidden_state: Optional[torch.FloatTensor] = None + + +@dataclass +class Step3p7CausalLMOutput(ModelOutput): + logits: Optional[torch.FloatTensor] = None + + +# --------------------------------------------------------------------------- +# Config access helper +# --------------------------------------------------------------------------- + + +def _get_text_config(config): + """Return the text sub-config (Step-3.7 wraps the LLM in a VLM ``Step3p7Config``).""" + return getattr(config, "text_config", config) + + +# --------------------------------------------------------------------------- +# Load-state-dict pre-hooks (run on the full ForCausalLM) +# --------------------------------------------------------------------------- + + +def _step3p7_norm_weight_load_hook(state_dict, prefix, *args, **kwargs): + """Absorb Step's ``(1 + weight)`` RMSNorm convention into the weight at load time. + + HF Step stores all norm weights as a bias around zero and applies ``(1 + weight)``. + Adding 1.0 here lets the forward use the standard ``torch_rmsnorm(x, weight, eps)`` + without an extra add node in the exported graph (matches the Gemma onboarding pattern). + """ + for key in list(state_dict.keys()): + if key.endswith("layernorm.weight") or key.endswith("norm.weight"): + state_dict[key] = state_dict[key] + 1.0 + + +def _step3p7_moe_split_load_hook(state_dict, prefix, *args, **kwargs): + """Split stacked routed-expert tensors into per-expert Linear tensors. + + The checkpoint stores routed experts as stacked tensors per projection: + * ``...moe.gate_proj.weight`` [E, moe_intermediate, hidden] + * ``...moe.up_proj.weight`` [E, moe_intermediate, hidden] + * ``...moe.down_proj.weight`` [E, hidden, moe_intermediate] + plus, for the FP8 checkpoint, block-wise dequant scales (one per projection): + * ``...moe.{proj}.weight_scale_inv`` [E, ceil(out/128), ceil(in/128)] + The custom model keeps per-expert ``nn.Linear`` modules for ``torch_moe`` dispatch, so split + every stacked ``[E, ...]`` tensor into per-expert tensors: + * ``...moe.experts.{e}.{gate,up,down}_proj.weight`` (and matching ``.weight_scale_inv``). + The FP8 ``quantize_finegrained_fp8_moe`` transform then consumes the per-expert FP8 weight + + ``weight_scale_inv`` exactly as it does for the (per-expert) DeepSeek-V3 checkpoint. + """ + for key in list(state_dict.keys()): + for proj in ("gate_proj", "up_proj", "down_proj"): + for leaf in ("weight", "weight_scale_inv"): + suffix = f".moe.{proj}.{leaf}" + if key.endswith(suffix) and state_dict[key].dim() == 3: + stacked = state_dict.pop(key) + base = key[: -len(suffix)] + for e in range(stacked.shape[0]): + state_dict[f"{base}.moe.experts.{e}.{proj}.{leaf}"] = stacked[e] + break + + +# --------------------------------------------------------------------------- +# RMSNorm (using AD canonical op; (1 + weight) absorbed at load time) +# --------------------------------------------------------------------------- + + +class Step3p7RMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.auto_deploy.torch_rmsnorm(x, self.weight, self.variance_epsilon) + + +# --------------------------------------------------------------------------- +# Rotary Embedding (per-layer-type: partial rotation + optional llama3 scaling) +# --------------------------------------------------------------------------- + + +def _compute_step3p7_inv_freq( + head_dim: int, + partial_rotary_factor: float, + base: float, + rope_scaling: Optional[dict], +) -> torch.Tensor: + """Inverse frequencies for Step RoPE (default or llama3-scaled).""" + dim = int(head_dim * partial_rotary_factor) + inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + + if not rope_scaling: + return inv_freq + + rope_type = rope_scaling.get("rope_type", rope_scaling.get("type")) + assert rope_type == "llama3", f"Step-3.7 only supports llama3 rope-scaling, got {rope_type!r}" + + # Faithful copy of transformers _compute_llama3_parameters scaling math. + factor = rope_scaling["factor"] + low_freq_factor = rope_scaling["low_freq_factor"] + high_freq_factor = rope_scaling["high_freq_factor"] + old_context_len = rope_scaling["original_max_position_embeddings"] + + low_freq_wavelen = old_context_len / low_freq_factor + high_freq_wavelen = old_context_len / high_freq_factor + + wavelen = 2 * math.pi / inv_freq + inv_freq_llama = torch.where(wavelen > low_freq_wavelen, inv_freq / factor, inv_freq) + smooth_factor = (old_context_len / wavelen - low_freq_factor) / ( + high_freq_factor - low_freq_factor + ) + smoothed_inv_freq = ( + 1 - smooth_factor + ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama + is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen) + inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama) + return inv_freq_llama + + +class Step3p7RotaryEmbedding(RotaryEmbeddingBase): + """RoPE table builder for one attention type (partial rotation, optional llama3 scaling). + + Keeps only the small ``inv_freq`` buffer; the cos/sin tables are graph-computed in forward + (so AD's ``optimize_rope`` can materialize a fused cache). For llama3 the attention scaling + factor is 1.0, so cos/sin are not rescaled. + """ + + def __init__( + self, + head_dim: int, + partial_rotary_factor: float, + base: float, + max_position_embeddings: int, + rope_scaling: Optional[dict] = None, + ): + super().__init__() + self.max_position_embeddings = max_position_embeddings + inv_freq = _compute_step3p7_inv_freq(head_dim, partial_rotary_factor, base, rope_scaling) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward( + self, x: torch.Tensor, position_ids: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + cos, sin = build_rope_cos_sin_cache(self.inv_freq, self.max_position_embeddings, x) + return cos[position_ids], sin[position_ids] + + +def _apply_partial_rope( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Apply RoPE to the first ``rotary_dim`` dims (bsnd layout), passing the rest through. + + ``rotary_dim = cos.shape[-1]`` may be < head_dim (full-attention layers) or == head_dim + (sliding-attention layers, in which case there is nothing to pass through). + """ + rotary_dim = cos.shape[-1] + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + q_rot, k_rot = torch.ops.auto_deploy.torch_rope_with_explicit_cos_sin( + q_rot, + k_rot, + cos, + sin, + 2, # unsqueeze_dim=2 for bsnd + ) + return torch.cat([q_rot, q_pass], dim=-1), torch.cat([k_rot, k_pass], dim=-1) + + +# --------------------------------------------------------------------------- +# Dense SwiGLU MLP (dense layers + shared expert) +# --------------------------------------------------------------------------- + + +class Step3p7MLP(nn.Module): + """SwiGLU MLP with an optional post-activation clamp (Step ``swiglu_limit``). + + Sharding: gate/up colwise, down rowwise. ``apply_all_reduce`` controls whether the rowwise + output is reduced here (True for a standalone dense MLP) or left partial for a downstream + merge-point all_reduce (False when used as a MoE shared expert). + """ + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + swiglu_limit: Optional[float] = None, + layer_type: str = "mlp", + apply_all_reduce: bool = True, + ): + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + self.act_fn = ACT2FN["silu"] + self.limit = swiglu_limit + self.layer_type = layer_type + self.apply_all_reduce = apply_all_reduce + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate = torch.ops.auto_deploy.torch_linear_simple( + x, self.gate_proj.weight, None, tp_mode="colwise", layer_type=self.layer_type + ) + up = torch.ops.auto_deploy.torch_linear_simple( + x, self.up_proj.weight, None, tp_mode="colwise", layer_type=self.layer_type + ) + gate = self.act_fn(gate) + if self.limit is not None: + gate = gate.clamp(max=self.limit) + up = up.clamp(min=-self.limit, max=self.limit) + down = torch.ops.auto_deploy.torch_linear_simple( + gate * up, self.down_proj.weight, None, tp_mode="rowwise", layer_type=self.layer_type + ) + if self.apply_all_reduce: + down = torch.ops.auto_deploy.all_reduce(down, layer_type=self.layer_type) + return down + + +# --------------------------------------------------------------------------- +# Sparse MoE block (routed experts + shared expert) +# --------------------------------------------------------------------------- + + +class Step3p7MoE(nn.Module): + """Routed MoE with sigmoid routing + per-expert bias (selection only). + + The dense shared expert is a sibling of this module on the decoder layer (matching the HF + hierarchy and the checkpoint layout ``model.layers.N.share_expert.*``), so it is NOT part of + this module. + + Routing (HF ``router_bias_func``): + 1. ``probs = sigmoid(fp32 router logits)`` + 2. select top-k experts by ``probs + router_bias`` + 3. gather the *un-biased* ``probs`` for the selected experts + 4. renormalize the gathered weights and scale by ``moe_router_scaling_factor`` + + NOTE on ``swiglu_limit``: the routed experts of the last two MoE layers carry a SwiGLU + activation clamp in the HF reference. ``torch_moe`` has no clamp parameter, so the routed + clamp is not applied here (the limits are large guards that rarely activate). The clamp IS + applied on the dense shared-expert path (a plain MLP, see ``Step3p7DecoderLayer``). + """ + + def __init__(self, config): + super().__init__() + self.num_experts = config.moe_num_experts + self.top_k = config.moe_top_k + self.hidden_size = config.hidden_size + self.routed_scaling_factor = getattr(config, "moe_router_scaling_factor", 1.0) + + self.gate = nn.Linear(self.hidden_size, self.num_experts, bias=False) + self.register_buffer( + "router_bias", torch.zeros(self.num_experts, dtype=torch.float32), persistent=True + ) + + self.experts = nn.ModuleList( + [ + Step3p7MLP(self.hidden_size, config.moe_intermediate_size) + for _ in range(self.num_experts) + ] + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + bsz, seq_len, hidden_dim = hidden_states.shape + hidden_flat = hidden_states.view(-1, hidden_dim) + + # fp32 router GEMM (config.need_fp32_gate) + router_logits = F.linear(hidden_flat.float(), self.gate.weight.float()) + probs = torch.sigmoid(router_logits) + + scores = probs + self.router_bias.unsqueeze(0) + _, selected_experts = torch.topk(scores, self.top_k, dim=-1) + routing_weights = torch.gather(probs, 1, selected_experts) + routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + 1e-20) + routing_weights = routing_weights * self.routed_scaling_factor + routing_weights = routing_weights.to(hidden_flat.dtype) + + routed = torch.ops.auto_deploy.torch_moe( + hidden_flat, + selected_experts, + routing_weights, + w1_weight=[e.gate_proj.weight for e in self.experts], + w2_weight=[e.down_proj.weight for e in self.experts], + w3_weight=[e.up_proj.weight for e in self.experts], + is_gated_mlp=True, + act_fn=int(ActivationType.Silu), + layer_type="moe", + ) + return routed.view(bsz, seq_len, hidden_dim) + + +# --------------------------------------------------------------------------- +# Attention (GQA + per-head QK norm + head-wise gate + partial RoPE) +# --------------------------------------------------------------------------- + + +class Step3p7Attention(nn.Module): + def __init__(self, config, layer_idx: int): + super().__init__() + self.layer_idx = layer_idx + self.head_dim = config.head_dim + self.attention_type = config.layer_types[layer_idx] + is_sliding = self.attention_type == "sliding_attention" + + if is_sliding: + other = config.attention_other_setting + self.num_heads = other["num_attention_heads"] + self.num_kv_heads = other["num_attention_groups"] + self.sliding_window = config.sliding_window + else: + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_attention_groups + self.sliding_window = None + + self.scaling = self.head_dim ** (-0.5) + + self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False) + # Head-wise attention gate: one sigmoid scalar per head, applied to the attention output + # before o_proj. Sharded as a per-head column shard (tp_min_local_shape=1) so it follows + # the same head partition as q/k/v under tensor parallelism. + self.g_proj = nn.Linear(config.hidden_size, self.num_heads, bias=False) + + # Per-head QK RMSNorm over head_dim. + self.q_norm = Step3p7RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = Step3p7RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + bsz, q_len, _ = hidden_states.size() + + q = torch.ops.auto_deploy.torch_linear_simple( + hidden_states, + self.q_proj.weight, + None, + tp_mode="colwise", + tp_min_local_shape=self.head_dim, + layer_type="mha", + ) + k = torch.ops.auto_deploy.torch_linear_simple( + hidden_states, + self.k_proj.weight, + None, + tp_mode="colwise", + tp_min_local_shape=self.head_dim, + layer_type="mha", + ) + v = torch.ops.auto_deploy.torch_linear_simple( + hidden_states, + self.v_proj.weight, + None, + tp_mode="colwise", + tp_min_local_shape=self.head_dim, + layer_type="mha", + ) + q = torch.ops.auto_deploy.view( + q, [bsz, q_len, self.num_heads, self.head_dim], tp_scaled_dim=2, layer_type="mha" + ) + k = torch.ops.auto_deploy.view( + k, [bsz, q_len, self.num_kv_heads, self.head_dim], tp_scaled_dim=2, layer_type="mha" + ) + v = torch.ops.auto_deploy.view( + v, [bsz, q_len, self.num_kv_heads, self.head_dim], tp_scaled_dim=2, layer_type="mha" + ) + + # Per-head QK norm over head_dim. + q = self.q_norm(q) + k = self.k_norm(k) + + cos, sin = position_embeddings + q, k = _apply_partial_rope(q, k, cos, sin) + + attn_output = torch.ops.auto_deploy.torch_attention( + q, + k, + v, + None, # attn_mask + 0.0, # dropout_p + True, # is_causal + self.scaling, # scale + None, # sinks + self.sliding_window, # sliding_window + None, # logit_cap + "bsnd", # layout + ) # [B, S, N, head_dim] + + # Head-wise gate: scale each head's output by sigmoid(per-head gate). g_proj is a per-head + # column shard (tp_min_local_shape=1), so its [B, S, N] output is sharded over the same + # head partition as the attention output. + gate = torch.ops.auto_deploy.torch_linear_simple( + hidden_states, + self.g_proj.weight, + None, + tp_mode="colwise", + tp_min_local_shape=1, + layer_type="mha", + ).sigmoid() # [B, S, N] + attn_output = attn_output * gate.unsqueeze(-1) + + attn_output = torch.ops.auto_deploy.view( + attn_output, + [bsz, q_len, self.num_heads * self.head_dim], + tp_scaled_dim=2, + layer_type="mha", + ) + attn_output = torch.ops.auto_deploy.torch_linear_simple( + attn_output, self.o_proj.weight, None, tp_mode="rowwise", layer_type="mha" + ) + attn_output = torch.ops.auto_deploy.all_reduce(attn_output, layer_type="mha") + return attn_output + + +# --------------------------------------------------------------------------- +# Decoder layer +# --------------------------------------------------------------------------- + + +class Step3p7DecoderLayer(nn.Module): + def __init__(self, config, layer_idx: int, is_moe_layer: bool): + super().__init__() + self.attention_type = config.layer_types[layer_idx] + self.self_attn = Step3p7Attention(config, layer_idx) + + _, shared_swiglu_limit = _layer_swiglu_limits(config, layer_idx) + self.is_moe_layer = is_moe_layer + if is_moe_layer: + self.moe = Step3p7MoE(config) + # Shared expert is a sibling of ``moe`` (checkpoint key model.layers.N.share_expert.*). + # No internal all_reduce: the single merge-point all_reduce (routed + shared) reduces it. + self.share_expert = Step3p7MLP( + config.hidden_size, + config.share_expert_dim, + swiglu_limit=shared_swiglu_limit, + layer_type="moe", + apply_all_reduce=False, + ) + else: + self.mlp = Step3p7MLP( + config.hidden_size, config.intermediate_size, swiglu_limit=shared_swiglu_limit + ) + + self.input_layernorm = Step3p7RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Step3p7RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + full_position_embeddings: Tuple[torch.Tensor, torch.Tensor], + sliding_position_embeddings: Tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + position_embeddings = ( + sliding_position_embeddings + if self.attention_type == "sliding_attention" + else full_position_embeddings + ) + + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, position_embeddings) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + if self.is_moe_layer: + # Single all_reduce at the routed + shared merge point (both are left partial above). + hidden_states = self.moe(hidden_states) + self.share_expert(hidden_states) + hidden_states = torch.ops.auto_deploy.all_reduce(hidden_states, layer_type="moe") + else: + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +def _moe_layer_indices(config) -> List[int]: + moe_layers_enum = getattr(config, "moe_layers_enum", None) + if moe_layers_enum is None: + return list(range(1, config.num_hidden_layers)) + if isinstance(moe_layers_enum, str): + return [int(i) for i in moe_layers_enum.split(",") if i.strip()] + return [int(i) for i in moe_layers_enum] + + +def _layer_swiglu_limits(config, layer_idx: int) -> Tuple[Optional[float], Optional[float]]: + """Return (routed-expert limit, shared/dense limit) for a layer, or None when disabled.""" + + def _val(values): + if not values or layer_idx >= len(values): + return None + v = values[layer_idx] + return float(v) if v else None + + return _val(getattr(config, "swiglu_limits", None)), _val( + getattr(config, "swiglu_limits_shared", None) + ) + + +# --------------------------------------------------------------------------- +# Model +# --------------------------------------------------------------------------- + + +class Step3p7PreTrainedModel(PreTrainedModel): + base_model_prefix = "model" + _no_split_modules = ["Step3p7DecoderLayer"] + supports_gradient_checkpointing = False + + +class Step3p7TextModel(Step3p7PreTrainedModel): + def __init__(self, config): + super().__init__(config) + text_config = _get_text_config(config) + self.config = config + + moe_layers = set(_moe_layer_indices(text_config)) + self.embed_tokens = nn.Embedding(text_config.vocab_size, text_config.hidden_size) + self.layers = nn.ModuleList( + [ + Step3p7DecoderLayer(text_config, idx, is_moe_layer=idx in moe_layers) + for idx in range(text_config.num_hidden_layers) + ] + ) + self.norm = Step3p7RMSNorm(text_config.hidden_size, eps=text_config.rms_norm_eps) + + self.full_rotary_emb, self.sliding_rotary_emb = _build_rotary_embeddings(text_config) + + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> Step3p7ModelOutput: + assert position_ids is not None, "position_ids is required for AD export" + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + inputs_embeds = inputs_embeds.to(self.norm.weight.dtype) + + full_pe = self.full_rotary_emb(inputs_embeds, position_ids) + sliding_pe = self.sliding_rotary_emb(inputs_embeds, position_ids) + + hidden_states = inputs_embeds + for decoder_layer in self.layers: + hidden_states = decoder_layer(hidden_states, full_pe, sliding_pe) + + hidden_states = self.norm(hidden_states) + return Step3p7ModelOutput(last_hidden_state=hidden_states) + + +def _build_rotary_embeddings(text_config): + """Build the two RoPE tables (full-attention and sliding-attention) from per-layer config. + + ``rope_theta`` and ``partial_rotary_factors`` are per-layer lists in the checkpoint config, + but they are constant within each attention type, so we read the value from a representative + layer of each type. llama3 rope-scaling applies only to ``yarn_only_types`` (full attention). + """ + layer_types = text_config.layer_types + rope_theta = text_config.rope_theta + partial_rotary_factors = getattr(text_config, "partial_rotary_factors", None) + rope_scaling = getattr(text_config, "rope_scaling", None) + yarn_only_types = getattr(text_config, "yarn_only_types", None) + head_dim = text_config.head_dim + max_pos = text_config.max_position_embeddings + + def _theta(idx): + return rope_theta[idx] if isinstance(rope_theta, (list, tuple)) else rope_theta + + def _partial(idx): + if partial_rotary_factors is not None: + return partial_rotary_factors[idx] + return getattr(text_config, "partial_rotary_factor", 1.0) + + def _scaling(layer_type): + if rope_scaling is None: + return None + if yarn_only_types is not None and layer_type not in yarn_only_types: + return None + return rope_scaling + + def _rep_index(layer_type): + return next(i for i, t in enumerate(layer_types) if t == layer_type) + + embeds = {} + for layer_type in ("full_attention", "sliding_attention"): + idx = _rep_index(layer_type) + embeds[layer_type] = Step3p7RotaryEmbedding( + head_dim=head_dim, + partial_rotary_factor=_partial(idx), + base=_theta(idx), + max_position_embeddings=max_pos, + rope_scaling=_scaling(layer_type), + ) + return embeds["full_attention"], embeds["sliding_attention"] + + +class Step3p7ForCausalLM(Step3p7PreTrainedModel, GenerationMixin): + def __init__(self, config, **kwargs): + super().__init__(config) + text_config = _get_text_config(config) + self.model = Step3p7TextModel(config) + self.vocab_size = text_config.vocab_size + self.lm_head = nn.Linear(text_config.hidden_size, text_config.vocab_size, bias=False) + + # Load-time checkpoint adapters: absorb (1 + weight) RMSNorm convention and split stacked + # MoE expert weights into per-expert Linear modules. + self._register_load_state_dict_pre_hook(_step3p7_norm_weight_load_hook) + self._register_load_state_dict_pre_hook(_step3p7_moe_split_load_hook) + + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def get_decoder(self): + return self.model + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> Step3p7CausalLMOutput: + assert position_ids is not None, "position_ids is required for AD export" + outputs = self.model( + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + logits = self.lm_head(outputs.last_hidden_state).float() + return Step3p7CausalLMOutput(logits=logits) + + +# --------------------------------------------------------------------------- +# Registration +# --------------------------------------------------------------------------- + +AutoModelForCausalLMFactory.register_custom_model_cls("Step3p7Config", Step3p7ForCausalLM) diff --git a/tests/unittest/auto_deploy/multigpu/transformations/library/test_step3p7_sharding_ir.py b/tests/unittest/auto_deploy/multigpu/transformations/library/test_step3p7_sharding_ir.py new file mode 100644 index 000000000000..30331894e290 --- /dev/null +++ b/tests/unittest/auto_deploy/multigpu/transformations/library/test_step3p7_sharding_ir.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Sharding-IR equivalence test for the StepFun Step-3.7-Flash AutoDeploy custom model. + +This is a model-specific, directly-runnable wrapper around the generic sharding-IR equivalence +harness in ``test_sharding_ir_equivalence.py``. Unlike the generic test (which is skipped unless +``--sharding-ir-modeling-file`` is supplied on the command line), this file pins the modeling file +to ``modeling_step3p7.py`` and parametrizes over every parallelism configuration, so it runs with a +plain:: + + pytest tests/unittest/auto_deploy/multigpu/transformations/library/test_step3p7_sharding_ir.py + +For each config it builds a tiny (4-layer, hidden_size=64) Step-3.7-Flash instance, exports it, +applies ``apply_sharding_hints`` to one copy, and asserts the sharded prefill matches the unsharded +reference within a relative-RMSE tolerance. It needs no PyExecutor / compile / checkpoint download +and skips automatically when fewer GPUs are available than a config requires. + +Tolerance: Step-3.7-Flash's correct-sharding rel_rmse is ~0.05 — higher than the generic default +(0.02, calibrated on plain dense/MoE models) because the head-wise attention gate and the MoE +``routed_scaling_factor=3.0`` amplify the bf16 ``all_reduce`` summation-order rounding. That this is +genuine finite-precision noise and not a sharding bug is confirmed by the harness's sabotage control +(removing the collectives drives rel_rmse to ~0.9, i.e. ~17x). We therefore set the per-model +tolerance below. +""" + +import sys +from functools import partial +from pathlib import Path + +import pytest + +# The generic harness and helpers live alongside this file / under _utils_test. +_THIS_DIR = Path(__file__).resolve().parent +if str(_THIS_DIR) not in sys.path: + sys.path.insert(0, str(_THIS_DIR)) + +from test_sharding_ir_equivalence import ( # noqa: E402 + _DIST_CONFIGS, + _gpu_check, + _run_equivalence_job, +) + +# Modeling file under test (bare path relative to repo root; the harness's +# ``spec_from_modeling_file`` also accepts a bare module short name). +_MODELING_FILE = "tensorrt_llm/_torch/auto_deploy/models/custom/modeling_step3p7.py" + +# Per-model relative-RMSE tolerance (see module docstring for the rationale). +STEP3P7_REL_RMSE_TOL = 0.08 + +pytestmark = pytest.mark.threadleak(enabled=False) + + +@pytest.mark.parametrize("dist_config", list(_DIST_CONFIGS)) +def test_step3p7_sharding_ir_equivalence(dist_config: str, monkeypatch) -> None: + """Sharded == unsharded prefill for Step-3.7-Flash under each parallelism config.""" + skip = _gpu_check(dist_config) + if skip: + pytest.skip(skip) + + # The per-rank worker reads the tolerance from this env var (spawned workers inherit it). + monkeypatch.setenv("SHARDING_IR_REL_RMSE_TOL", str(STEP3P7_REL_RMSE_TOL)) + + import tensorrt_llm._torch.auto_deploy.distributed.common as dist_common + + world_size = _DIST_CONFIGS[dist_config]["world_size"] + dist_common.spawn_multiprocess_job( + job=partial(_run_equivalence_job, _MODELING_FILE, dist_config), + size=world_size, + ) diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_step3p7_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_step3p7_modeling.py new file mode 100644 index 000000000000..8688282ca0e8 --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/models/test_step3p7_modeling.py @@ -0,0 +1,643 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Hierarchical equivalence tests for the StepFun Step-3.7-Flash AutoDeploy custom model. + +Step-3.7-Flash uses ``trust_remote_code`` (its ``step3p7`` / ``step3p5`` modeling code is not +in transformers natively), so the reference classes (``_Ref*``) below are minimal, faithful +standalone reimplementations of the HuggingFace text-decoder math (modeling_step3p7.py). + +Levels: MLP block -> Attention block (full + sliding) -> MoE block -> Decoder layer -> +Full model -> Export. +""" + +from typing import Tuple + +import pytest +import torch +import torch.nn.functional as F +from torch import nn +from torch.export import Dim +from transformers import PretrainedConfig +from transformers.activations import ACT2FN + +import tensorrt_llm._torch.auto_deploy.custom_ops # noqa: F401 — register ops +from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.models.custom.modeling_step3p7 import ( + Step3p7Attention, + Step3p7DecoderLayer, + Step3p7ForCausalLM, + Step3p7MLP, + Step3p7MoE, + Step3p7RotaryEmbedding, +) + + +@pytest.fixture(autouse=True) +def _deterministic_seed(): + """Seed RNG before every test so random weights/inputs are reproducible. + + The MoE routing top-k can otherwise be sensitive to inter-test RNG state on + near-tied expert scores, making the equivalence checks intermittently flaky. + """ + torch.manual_seed(0) + + +# --------------------------------------------------------------------------- +# Test-only config (minimal faithful copy of the HF Step3p7 text config) +# --------------------------------------------------------------------------- + + +class Step3p7TestConfig(PretrainedConfig): + """Flat text config for testing. + + The model's ``_get_text_config`` treats a config without a ``text_config`` attribute as the + text config itself, so no VLM wrapper is needed here. + """ + + model_type = "step3p5" + + def __init__( + self, + vocab_size=1000, + hidden_size=64, + head_dim=16, + num_attention_heads=4, + num_attention_groups=2, + attention_other_setting=None, + intermediate_size=128, + num_hidden_layers=4, + layer_types=None, + moe_layers_enum=(2, 3), + moe_num_experts=8, + moe_top_k=2, + moe_intermediate_size=32, + share_expert_dim=32, + moe_router_scaling_factor=3.0, + rms_norm_eps=1e-5, + sliding_window=4, + max_position_embeddings=64, + rope_theta=(5e6, 1e4, 5e6, 1e4), + partial_rotary_factors=(0.5, 1.0, 0.5, 1.0), + rope_scaling=None, + yarn_only_types=("full_attention",), + swiglu_limits=None, + swiglu_limits_shared=None, + **kwargs, + ): + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.head_dim = head_dim + self.num_attention_heads = num_attention_heads + self.num_attention_groups = num_attention_groups + self.attention_other_setting = attention_other_setting or { + "num_attention_heads": 6, + "num_attention_groups": 2, + "head_dim": 16, + } + self.intermediate_size = intermediate_size + self.num_hidden_layers = num_hidden_layers + self.layer_types = layer_types or [ + "full_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + ] + self.moe_layers_enum = moe_layers_enum + self.moe_num_experts = moe_num_experts + self.moe_top_k = moe_top_k + self.moe_intermediate_size = moe_intermediate_size + self.share_expert_dim = share_expert_dim + self.moe_router_scaling_factor = moe_router_scaling_factor + self.rms_norm_eps = rms_norm_eps + self.sliding_window = sliding_window + self.max_position_embeddings = max_position_embeddings + self.rope_theta = list(rope_theta) + self.partial_rotary_factors = list(partial_rotary_factors) + self.rope_scaling = rope_scaling or { + "rope_type": "llama3", + "factor": 2.0, + "original_max_position_embeddings": 64, + "low_freq_factor": 1.0, + "high_freq_factor": 32.0, + } + self.yarn_only_types = list(yarn_only_types) + self.swiglu_limits = swiglu_limits + self.swiglu_limits_shared = swiglu_limits_shared + super().__init__(**kwargs) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def assert_rmse_close(actual, expected, rmse_ratio_tol, msg=""): + diff = actual.float() - expected.float() + rmse_diff = torch.sqrt(torch.mean(diff**2)) + rmse_ref = torch.sqrt(torch.mean(expected.float() ** 2)) + ratio = (rmse_diff / rmse_ref).item() + assert ratio < rmse_ratio_tol, ( + f"{msg}RMSE ratio {ratio:.6f} exceeds tolerance {rmse_ratio_tol}. " + f"(rmse_diff={rmse_diff.item():.6f}, rmse_ref={rmse_ref.item():.6f})" + ) + + +def _device_and_dtype() -> Tuple[str, torch.dtype]: + if torch.cuda.is_available(): + return "cuda", torch.bfloat16 + return "cpu", torch.float32 + + +def _position_ids(batch: int, seq: int, device) -> torch.Tensor: + return torch.arange(seq, device=device).unsqueeze(0).expand(batch, -1) + + +def _small_config(**overrides) -> Step3p7TestConfig: + return Step3p7TestConfig(**overrides) + + +def _add_one_to_norms(state_dict: dict) -> dict: + """Mimic the AD load hook: absorb the (1 + weight) RMSNorm convention into norm weights.""" + out = {} + for k, v in state_dict.items(): + if k.endswith("layernorm.weight") or k.endswith("norm.weight"): + out[k] = v + 1.0 + else: + out[k] = v + return out + + +# --------------------------------------------------------------------------- +# Reference implementations (HF-faithful, plain PyTorch, bnsd layout) +# --------------------------------------------------------------------------- + + +class _RefRMSNorm(nn.Module): + """Step RMSNorm with the (1 + weight) convention.""" + + def __init__(self, hidden_size: int, eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter(torch.zeros(hidden_size)) + self.variance_epsilon = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + dtype = x.dtype + x = x.float() + variance = x.pow(2).mean(dim=-1, keepdim=True) + normed = x * torch.rsqrt(variance + self.variance_epsilon) + normed = normed * (self.weight.float() + 1) + return normed.to(dtype) + + +def _ref_rotate_half(x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def _ref_apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1): + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + rotary_dim = cos.shape[-1] + q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] + k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] + q_embed = (q_rot * cos) + (_ref_rotate_half(q_rot) * sin) + k_embed = (k_rot * cos) + (_ref_rotate_half(k_rot) * sin) + return torch.cat([q_embed, q_pass], dim=-1), torch.cat([k_embed, k_pass], dim=-1) + + +def _ref_repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + batch, num_kv_heads, slen, head_dim = hidden_states.shape + if n_rep == 1: + return hidden_states + hidden_states = hidden_states[:, :, None, :, :].expand( + batch, num_kv_heads, n_rep, slen, head_dim + ) + return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim) + + +class _RefMLP(nn.Module): + """Reference SwiGLU MLP with optional post-activation clamp (matches Step3p7MLP naming).""" + + def __init__(self, hidden_size, intermediate_size, swiglu_limit=None): + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + self.act_fn = ACT2FN["silu"] + self.limit = swiglu_limit + + def forward(self, x): + gate = self.act_fn(self.gate_proj(x)) + up = self.up_proj(x) + if self.limit is not None: + gate = gate.clamp(max=self.limit) + up = up.clamp(min=-self.limit, max=self.limit) + return self.down_proj(gate * up) + + +class _RefAttention(nn.Module): + """Reference Step attention: per-head QK norm, partial RoPE, head-wise gate, GQA, sliding.""" + + def __init__(self, config: Step3p7TestConfig, layer_idx: int): + super().__init__() + self.head_dim = config.head_dim + is_sliding = config.layer_types[layer_idx] == "sliding_attention" + if is_sliding: + other = config.attention_other_setting + self.num_heads = other["num_attention_heads"] + self.num_kv_heads = other["num_attention_groups"] + self.sliding_window = config.sliding_window + else: + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_attention_groups + self.sliding_window = None + self.num_kv_groups = self.num_heads // self.num_kv_heads + self.scaling = self.head_dim**-0.5 + + self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False) + self.g_proj = nn.Linear(config.hidden_size, self.num_heads, bias=False) + self.q_norm = _RefRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = _RefRMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward(self, hidden_states, position_embeddings): + bsz, q_len, _ = hidden_states.size() + hidden_shape = (bsz, q_len, -1, self.head_dim) + + q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + v = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + gate = self.g_proj(hidden_states) # [B, S, N] + + cos, sin = position_embeddings + q, k = _ref_apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1) + + k = _ref_repeat_kv(k, self.num_kv_groups) + v = _ref_repeat_kv(v, self.num_kv_groups) + + attn_weights = torch.matmul(q, k.transpose(2, 3)) * self.scaling + qpos = torch.arange(q_len, device=q.device) + pos_diff = qpos.unsqueeze(1) - qpos.unsqueeze(0) + if self.sliding_window is not None: + mask = (pos_diff < 0) | (pos_diff >= self.sliding_window) + else: + mask = pos_diff < 0 + attn_weights = attn_weights.masked_fill(mask.unsqueeze(0).unsqueeze(0), float("-inf")) + attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype) + attn_output = torch.matmul(attn_weights, v) # [B, N, S, D] + + attn_output = attn_output.transpose(1, 2) # [B, S, N, D] + attn_output = attn_output * gate.unsqueeze(-1).sigmoid() + attn_output = attn_output.reshape(bsz, q_len, -1) + return self.o_proj(attn_output) + + +class _RefMoE(nn.Module): + """Reference routed Step MoE: sigmoid routing + per-expert bias (selection only). + + The shared expert is a sibling of this module on the decoder layer (matching the checkpoint + hierarchy), so it lives on ``_RefDecoderLayer``, not here. + """ + + def __init__(self, config: Step3p7TestConfig): + super().__init__() + self.num_experts = config.moe_num_experts + self.top_k = config.moe_top_k + self.routed_scaling_factor = config.moe_router_scaling_factor + self.gate = nn.Linear(config.hidden_size, self.num_experts, bias=False) + self.register_buffer("router_bias", torch.zeros(self.num_experts, dtype=torch.float32)) + self.experts = nn.ModuleList( + [ + _RefMLP(config.hidden_size, config.moe_intermediate_size) + for _ in range(self.num_experts) + ] + ) + + def forward(self, hidden_states): + bsz, seq_len, hidden_dim = hidden_states.shape + hidden_flat = hidden_states.view(-1, hidden_dim) + + router_logits = F.linear(hidden_flat.float(), self.gate.weight.float()) + probs = torch.sigmoid(router_logits) + scores = probs + self.router_bias.unsqueeze(0) + _, idx = torch.topk(scores, self.top_k, dim=-1) + weights = torch.gather(probs, 1, idx) + weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-20) + weights = weights * self.routed_scaling_factor + weights = weights.to(hidden_flat.dtype) + + final = torch.zeros_like(hidden_flat) + expert_mask = F.one_hot(idx, num_classes=self.num_experts).permute(2, 1, 0) + for e in range(self.num_experts): + tok_idx, top_x = torch.where(expert_mask[e]) + if top_x.numel() == 0: + continue + out = self.experts[e](hidden_flat[top_x]) * weights[top_x, tok_idx, None] + final.index_add_(0, top_x, out.to(hidden_flat.dtype)) + + return final.view(bsz, seq_len, hidden_dim) + + +class _RefDecoderLayer(nn.Module): + def __init__(self, config: Step3p7TestConfig, layer_idx: int, is_moe: bool): + super().__init__() + self.self_attn = _RefAttention(config, layer_idx) + self.is_moe = is_moe + if is_moe: + self.moe = _RefMoE(config) + self.share_expert = _RefMLP(config.hidden_size, config.share_expert_dim) + else: + self.mlp = _RefMLP(config.hidden_size, config.intermediate_size) + self.input_layernorm = _RefRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = _RefRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward(self, hidden_states, position_embeddings): + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, position_embeddings) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + if self.is_moe: + hidden_states = self.moe(hidden_states) + self.share_expert(hidden_states) + else: + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +# --------------------------------------------------------------------------- +# RoPE helper: build cos/sin for a given attention type using the AD module +# --------------------------------------------------------------------------- + + +def _build_position_embeddings(config, layer_idx, B, S, device, dtype): + layer_type = config.layer_types[layer_idx] + scaling = config.rope_scaling if layer_type in config.yarn_only_types else None + rope = Step3p7RotaryEmbedding( + head_dim=config.head_dim, + partial_rotary_factor=config.partial_rotary_factors[layer_idx], + base=config.rope_theta[layer_idx], + max_position_embeddings=config.max_position_embeddings, + rope_scaling=scaling, + ).to(device) + dummy = torch.zeros(1, device=device, dtype=dtype) + return rope(dummy, _position_ids(B, S, device)) + + +def _transfer_block(ad_module, ref_module): + """Load reference weights into an AD block, absorbing the (1 + weight) norm convention.""" + ad_module.load_state_dict(_add_one_to_norms(ref_module.state_dict()), strict=True) + + +# --------------------------------------------------------------------------- +# Tests — Block equivalence +# --------------------------------------------------------------------------- + + +def test_mlp_equivalence(): + device, dtype = _device_and_dtype() + config = _small_config() + ref = _RefMLP(config.hidden_size, config.intermediate_size).to(device, dtype).eval() + ad = Step3p7MLP(config.hidden_size, config.intermediate_size).to(device, dtype).eval() + ad.load_state_dict(ref.state_dict()) + x = torch.randn(2, 8, config.hidden_size, device=device, dtype=dtype) + with torch.no_grad(): + torch.testing.assert_close(ad(x), ref(x), rtol=1e-3, atol=1e-3) + + +def test_mlp_clamped_equivalence(): + """SwiGLU clamp (swiglu_limit) on the dense/shared MLP path matches reference.""" + device, dtype = _device_and_dtype() + config = _small_config() + ref = _RefMLP(config.hidden_size, config.intermediate_size, swiglu_limit=0.5).to(device, dtype) + ref = ref.eval() + ad = Step3p7MLP(config.hidden_size, config.intermediate_size, swiglu_limit=0.5).to( + device, dtype + ) + ad = ad.eval() + ad.load_state_dict(ref.state_dict()) + x = torch.randn(2, 8, config.hidden_size, device=device, dtype=dtype) * 3.0 # trigger clamp + with torch.no_grad(): + torch.testing.assert_close(ad(x), ref(x), rtol=1e-3, atol=1e-3) + + +def test_attention_full_equivalence(): + """Full-attention layer (64-head config, partial RoPE w/ llama3 scaling) matches reference.""" + device, dtype = _device_and_dtype() + config = _small_config() + B, S = 2, 8 + layer_idx = 0 # full_attention + ref = _RefAttention(config, layer_idx).to(device, dtype).eval() + ad = Step3p7Attention(config, layer_idx).to(device, dtype).eval() + _transfer_block(ad, ref) + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + pe = _build_position_embeddings(config, layer_idx, B, S, device, dtype) + with torch.no_grad(): + ad_out = ad(x, position_embeddings=pe) + ref_out = ref(x, position_embeddings=pe) + assert_rmse_close(ad_out, ref_out, rmse_ratio_tol=0.10, msg="Full attention: ") + + +def test_attention_sliding_equivalence(): + """Sliding-attention layer (96-head config, full RoPE, sliding window) matches reference.""" + device, dtype = _device_and_dtype() + config = _small_config() + B, S = 2, 8 + layer_idx = 1 # sliding_attention + ref = _RefAttention(config, layer_idx).to(device, dtype).eval() + ad = Step3p7Attention(config, layer_idx).to(device, dtype).eval() + _transfer_block(ad, ref) + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + pe = _build_position_embeddings(config, layer_idx, B, S, device, dtype) + with torch.no_grad(): + ad_out = ad(x, position_embeddings=pe) + ref_out = ref(x, position_embeddings=pe) + assert_rmse_close(ad_out, ref_out, rmse_ratio_tol=0.10, msg="Sliding attention: ") + + +def test_moe_block_equivalence(): + device, dtype = _device_and_dtype() + config = _small_config() + ref = _RefMoE(config).to(device, dtype).eval() + ad = Step3p7MoE(config).to(device, dtype).eval() + ad.load_state_dict(ref.state_dict()) + x = torch.randn(2, 8, config.hidden_size, device=device, dtype=dtype) + with torch.no_grad(): + ad_out = ad(x) + ref_out = ref(x) + assert_rmse_close(ad_out, ref_out, rmse_ratio_tol=0.02, msg="MoE block: ") + + +# --------------------------------------------------------------------------- +# Tests — Layer equivalence +# --------------------------------------------------------------------------- + + +def test_decoder_layer_moe_equivalence(): + """MoE decoder layer (full-attention variant) matches reference.""" + device, dtype = _device_and_dtype() + config = _small_config() + B, S = 2, 8 + layer_idx = 2 # full_attention + moe + ref = _RefDecoderLayer(config, layer_idx, is_moe=True).to(device, dtype).eval() + ad = Step3p7DecoderLayer(config, layer_idx, is_moe_layer=True).to(device, dtype).eval() + _transfer_block(ad, ref) + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + full_pe = _build_position_embeddings(config, layer_idx, B, S, device, dtype) + with torch.no_grad(): + ad_out = ad(x, full_pe, full_pe) # full-attention layer ignores sliding pe + ref_out = ref(x, full_pe) + assert_rmse_close(ad_out, ref_out, rmse_ratio_tol=0.05, msg="MoE decoder layer: ") + + +def test_decoder_layer_dense_sliding_equivalence(): + """Dense decoder layer (sliding-attention variant) matches reference.""" + device, dtype = _device_and_dtype() + config = _small_config() + B, S = 2, 8 + layer_idx = 1 # sliding_attention + dense + ref = _RefDecoderLayer(config, layer_idx, is_moe=False).to(device, dtype).eval() + ad = Step3p7DecoderLayer(config, layer_idx, is_moe_layer=False).to(device, dtype).eval() + _transfer_block(ad, ref) + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + sliding_pe = _build_position_embeddings(config, layer_idx, B, S, device, dtype) + with torch.no_grad(): + ad_out = ad(x, sliding_pe, sliding_pe) + ref_out = ref(x, sliding_pe) + assert_rmse_close(ad_out, ref_out, rmse_ratio_tol=0.05, msg="Dense sliding decoder layer: ") + + +# --------------------------------------------------------------------------- +# Tests — Full model equivalence (also exercises the load-state-dict hooks) +# --------------------------------------------------------------------------- + + +class _RefForCausalLM(nn.Module): + def __init__(self, config: Step3p7TestConfig): + super().__init__() + self.config = config + moe_layers = set(config.moe_layers_enum) + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList( + [ + _RefDecoderLayer(config, i, is_moe=i in moe_layers) + for i in range(config.num_hidden_layers) + ] + ) + self.norm = _RefRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + def forward(self, input_ids, position_ids): + hidden = self.embed_tokens(input_ids) + B, S = input_ids.shape + device, dtype = input_ids.device, hidden.dtype + pe_cache = { + i: _build_position_embeddings(self.config, i, B, S, device, dtype) + for i in range(self.config.num_hidden_layers) + } + for i, layer in enumerate(self.layers): + hidden = layer(hidden, pe_cache[i]) + hidden = self.norm(hidden) + return self.lm_head(hidden).float() + + +def _ref_to_checkpoint_state_dict(ref: _RefForCausalLM) -> dict: + """Convert reference weights into the on-disk checkpoint form expected by the AD load hooks. + + * norm weights stay zero-centered (the hook adds 1.0) + * routed-expert weights are stacked into ``moe.{gate,up,down}_proj.weight`` + * everything else gets the ``model.`` prefix (lm_head stays top-level) + """ + ref_sd = ref.state_dict() + ckpt = {} + num_experts = ref.config.moe_num_experts + for k, v in ref_sd.items(): + # Stack per-expert routed weights -> moe.{proj}.weight + import re + + m = re.match( + r"layers\.(\d+)\.moe\.experts\.(\d+)\.(gate_proj|up_proj|down_proj)\.weight$", k + ) + if m: + continue # handled below + if k.startswith("lm_head."): + ckpt[k] = v + else: + ckpt[f"model.{k}"] = v + + for li in ref.config.moe_layers_enum: + for proj in ("gate_proj", "up_proj", "down_proj"): + ws = [ref_sd[f"layers.{li}.moe.experts.{e}.{proj}.weight"] for e in range(num_experts)] + ckpt[f"model.layers.{li}.moe.{proj}.weight"] = torch.stack(ws, 0) + return ckpt + + +def test_full_model_equivalence(): + device, dtype = _device_and_dtype() + config = _small_config() + ref = _RefForCausalLM(config).to(device, dtype).eval() + ad = Step3p7ForCausalLM(config).to(device, dtype).eval() + + missing, unexpected = ad.load_state_dict(_ref_to_checkpoint_state_dict(ref), strict=False) + assert not missing, f"Missing keys: {missing[:10]}" + assert not unexpected, f"Unexpected keys: {unexpected[:10]}" + + B, S = 2, 8 + input_ids = torch.randint(0, config.vocab_size, (B, S), device=device) + pos_ids = _position_ids(B, S, device) + with torch.no_grad(): + ref_logits = ref(input_ids, pos_ids) + ad_out = ad(input_ids=input_ids, position_ids=pos_ids) + + assert ad_out.logits.shape == (B, S, config.vocab_size) + assert torch.isfinite(ad_out.logits).all() + assert_rmse_close(ad_out.logits, ref_logits, rmse_ratio_tol=0.05, msg="Full model: ") + + +# --------------------------------------------------------------------------- +# Tests — Export +# --------------------------------------------------------------------------- + + +def test_export(): + device = "cpu" + dtype = torch.float32 + config = _small_config() + model = Step3p7ForCausalLM(config).to(device, dtype).eval() + + B, S = 2, 8 + input_ids = torch.randint(0, config.vocab_size, (B, S), device=device) + pos_ids = _position_ids(B, S, device) + dynamic_shapes = { + "input_ids": {0: Dim.DYNAMIC, 1: Dim.DYNAMIC}, + "position_ids": {0: Dim.DYNAMIC, 1: Dim.DYNAMIC}, + } + gm = torch_export_to_gm( + model, + args=(input_ids,), + kwargs={"position_ids": pos_ids}, + dynamic_shapes=dynamic_shapes, + ) + with torch.no_grad(): + pre = model(input_ids=input_ids, position_ids=pos_ids) + exported = gm(input_ids, position_ids=pos_ids) + logits = exported[0] if isinstance(exported, tuple) else getattr(exported, "logits", exported) + assert torch.isfinite(logits).all() + torch.testing.assert_close(logits, pre.logits, rtol=1e-3, atol=1e-3) + + # Second shape + B2, S2 = 1, 5 + ids2 = torch.randint(0, config.vocab_size, (B2, S2), device=device) + pos2 = _position_ids(B2, S2, device) + with torch.no_grad(): + out2 = gm(ids2, position_ids=pos2) + logits2 = out2[0] if isinstance(out2, tuple) else getattr(out2, "logits", out2) + assert logits2.shape == (B2, S2, config.vocab_size) + assert torch.isfinite(logits2).all()