[WIP] feat: add model architecture configuration

[aftersnow]

This commit introduces a new document detailing the model architecture configuration, including support for transformer architectures and their components. It outlines terminology, properties, and provides an example configuration for better understanding and implementation.

Description

Add an optional architecture configuration that describes the detailed structure and components of the model. Currently, only decoder-type transformer architectures are supported.

Motivation and Context

The current modelpack can serve as a standard for building, distributing, and managing AI models. Using this specification, we can:

• Build the model from a model checkpoint into a model image
• Push the model image to an OCI image registry
• Pull the model image, either as a standalone package or a read-only volume source in a container.

However, it lacks the architectural and runtime interfaces to run the model. For example, it lacks crucial architectural details such as:

• Architecture hyperparameters, such as number of layers, hidden size dimensions, number of attention heads, etc.
• Model components type, such as self-attention mechanism, feed-forward network, etc.
• Model runtime interfaces, such as quantization configuration, parallelism configuration, etc.
• And more

This creates a gap between the model image and the model runtime. Without this additional information, the model runtime cannot execute the model seamlessly (only if the model runtime has pre-built specific support). To address these limitations, we are defining the model architecture configuration as the first step.

This PR adds an optional architecture configuration that describes the detailed structure and components of the model.

Use Cases

• LLM Researchers: Avoid reimplementing similar code across multiple inference frameworks such as TensorRT, vLLM, sglang, TGI when releasing new models.
• Inference Engine Developers: Eliminate duplicate implementations when adding features across different model architectures such as Llama, DeepSeek, Qwen, etc.
• Infrastructure Providers: Access architectural details to optimize model deployment and orchestration.
• Hardware Vendors: Obtain architectural details for targeted model acceleration and optimization.

Details

The transformer is the most popular architecture for LLMs. It consists of a stack of structured layers, where each layer contains a self-attention block and a feed-forward network, with normalization layers and residual connections. The complete architecture includes a tokenizer, input embedding layer, position embedding layer, transformer layers, and output embedding layer. The transformer architecture has remained relatively stable since Attention is all you need. As shown in the table below, current open-weight model architectures are converging, making it feasible to define a common abstraction.

Model	Tokenizer	PE	Self-Attention	Norm	Feed-Forward	Residual
GPT2	BPE	Sinusoidal	MHA	Layer Norm	MLP	Yes
Llama2	BPE	RoPE	GQA	RMS Norm	MLP	Yes
Llama3	BPE	RoPE	GQA	RMS Norm	MLP	Yes
Qwen2	BPE	RoPE	GQA	RMS Norm	MoE	Yes
Qwen3	BPE	RoPE	GQA	RMS Norm	MoE	Yes
Gemma2	BPE	RoPE	GQA	RMS Norm	MLP	Yes
Gemma3	BPE	RoPE	GQA	RMS Norm	MLP	Yes
Mixtral	BPE	RoPE	SWA	RMS Norm	MoE	Yes
DeepseekV2	BPE	RoPE	MLA	RMS Norm	MoE	Yes
DeepseekV3	BPE	RoPE	MLA	RMS Norm	MoE	Yes
Kimi-K2	BPE	RoPE	MLA	RMS Norm	MoE	Yes

Note: Each model represents the largest variant within its respective series.

Next Steps

• Define the model architecture configuration for other architectures, such as:
  • Multi-modal language models
  • State Space Models
  • Diffusion Models
  • And more
• Define framework-agnostic runtime interfaces to enhance interoperability across different models and frameworks, providing a standard way to run LLMs in cloud-native environments.
• Integrate with popular LLM runtimes.

Acknowledgement

This PR is inspired by transformers, vllm, ollama, onnx, tgi, sglang, lmdeploy, curated-transformers, triton-inference-server, transformer-engine, ggml.

[gemini-code-assist]

Summary of Changes

Hello @aftersnow, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces a significant enhancement to the modelpack specification by adding an optional, detailed model architecture configuration. The primary goal is to provide explicit structural information about models, starting with decoder-type transformers, which is currently missing from the modelpack. This new configuration will enable model runtimes to seamlessly execute models by understanding their internal components and parameters, thereby improving interoperability and reducing the need for hardcoded details.

Highlights

• New Model Architecture Specification: I've added a comprehensive new document, docs/architecture.md, which defines an optional configuration for model architectures. This is a crucial step towards bridging the gap between model images and model runtimes by providing detailed structural information.
• Detailed Transformer Architecture Configuration: The new specification currently focuses on decoder-type transformer architectures, outlining detailed properties for components like tokenizers, embeddings (positional and token), transformer layers (attention, MLP, MoE), and normalization. It supports both uniform and mixed layer configurations, providing a granular way to describe model structure.
• Integration with Modelpack Configuration: I've updated docs/config.md to include an optional architecture_config field. This allows modelpack users to associate a detailed architectural description with their models, enhancing the existing specification and enabling better interoperability with model runtimes.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This PR introduces a new architecture_config field to the model configuration and a detailed specification for it in docs/architecture.md. This is a great addition that will improve model interoperability.

I've left a few comments on the documentation for clarity and consistency. The most important feedback is about a broken link in docs/config.md.

Additionally, it seems that the corresponding Go structs and JSON schema have not been updated to reflect these changes. The following files appear to be missing updates:

• specs-go/v1/config.go
• schema/config-schema.json

Keeping the specification documents, Go types, and JSON schema in sync is crucial for consumers of this spec. Please consider adding the necessary changes to these files in this PR.

[gemini-code-assist]

The interaction between mlp_layers and moe_frequency is unclear.

• Both are marked as REQUIRED, which seems contradictory if they are alternative ways of defining layer types.
• The description for mlp_layers says "If empty, moe_frequency determines sparsity", which suggests they might be mutually exclusive in some way.

Could you please clarify the logic? For example:

• Are both fields always required? If so, how do they work together?
• Should the user provide one or the other? If so, they should probably be OPTIONAL and a constraint like oneOf should be defined on the parent mixed_layers object.

The documentation should be updated to clearly explain how to configure mixed MLP and MoE layers. Also, note that mlp_layers is marked as REQUIRED but has a default value, which is inconsistent as mentioned in another comment.

[aftersnow]

Please take a look. More development, testing, and validation are needed. Recently, I have been busy with production environments and haven't had enough time. Everyone is welcome to contribute!

[github-actions]

This PR is stale because it has been open 40 days with no activity.