Vedika-Vyom-4b-v1-mini

A Custom Architecture Implementation by Veda Labs

Website | Hugging Face | GitHub | X (Twitter)

Table of Contents

• Overview
• Model Information
• Custom Architecture
• Installation
• Quickstart
• Architecture Details
• Weight Mapping
• Citation
• Contact & Links

Overview

Vedika-Vyom-4b-v1-mini is a custom transformer architecture developed by Veda Labs. This project provides a clean, standalone PyTorch implementation featuring advanced components including RMSNorm normalization, Rotary Position Embeddings (RoPE), and grouped query attention mechanisms optimized for efficient text generation.

Key Features

• ✅ Custom Architecture: From-scratch implementation of all core transformer components
• ✅ Advanced Normalization: RMSNorm layers for stable training and inference
• ✅ Rotary Embeddings: RoPE for superior positional encoding
• ✅ Grouped Query Attention: Efficient GQA mechanism with 8 attention heads and 4 KV heads
• ✅ Hugging Face Integration: Seamless auto-registration with AutoConfig and AutoModelForCausalLM
• ✅ Standalone Portability: Runs with trust_remote_code=True without external dependencies
• ✅ Clean API: Intuitive module design following Hugging Face conventions

Model Information

Custom Architecture

This repository implements a complete custom architecture for Vedika-Vyom-4b-v1-mini with the following core components:

All components are natively designed as part of the Vedika-Vyom architecture, ensuring optimal performance and numerical precision.

Installation

Prerequisites

• Python 3.9+
• PyTorch 2.0+
• Transformers 4.50.0+

Install Dependencies

pip install torch>=2.0.0
pip install transformers>=4.50.0
pip install accelerate  # Optional, for multi-GPU support

Clone Repository (Optional)

git clone https://github.com/vedalabs-tech/vedika-vyom-4b-v1-mini.git
cd vedika-vyom-4b-v1-mini

Quickstart

Loading the Model

The model can be loaded directly from Hugging Face Hub using AutoTokenizer and AutoModelForCausalLM with trust_remote_code=True:

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Model identifier
model_id = "Veda-Labs/Vedika-Vyom-4b-v1-mini"

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)

# Load model
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    device_map="auto"  # Automatically move model to GPU if available
)

model.eval()

Running Inference

from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Initialize
model_id = "Veda-Labs/Vedika-Vyom-4b-v1-mini"
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    device_map="cuda"
)

# Prepare input
prompt = "Explain the concept of quantum entanglement in simple terms."
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

# Generate
with torch.inference_mode():
    outputs = model.generate(
        **inputs,
        max_new_tokens=256,
        do_sample=True,
        temperature=0.7,
        top_p=0.95,
        repetition_penalty=1.1
    )

# Decode output
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response)

Batch Inference Example

prompts = [
    "What is the capital of France?",
    "Write a haiku about artificial intelligence.",
    "Explain Newton's third law of motion."
]

inputs = tokenizer(prompts, return_tensors="pt", padding=True).to(model.device)

with torch.inference_mode():
    outputs = model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=False
    )

for i, prompt in enumerate(prompts):
    print(f"\nPrompt: {prompt}")
    print(f"Response: {tokenizer.decode(outputs[i], skip_special_tokens=True)}")

Architecture Details

Model Configuration

The VedikaVyomConfig class defines the following key hyperparameters for the Vedika-Vyom architecture:

VedikaVyomConfig(
    vocab_size=256000,        # Vocabulary size
    hidden_size=2048,         # Hidden dimension
    intermediate_size=8192,   # MLP intermediate dimension
    num_hidden_layers=26,     # Number of transformer layers
    num_attention_heads=8,    # Number of attention heads
    num_key_value_heads=4,    # Number of KV heads (GQA)
    head_dim=256,             # Head dimension
    rms_norm_eps=1e-6,        # RMSNorm epsilon
    max_position_embeddings=131072,  # Max context length (128K)
    rope_theta=1000000.0,     # RoPE theta
    hidden_act="gelu_pytorch_tanh"  # Activation function
)

Layer Structure

Each VedikaVyomDecoderLayer contains:

1. Input LayerNorm (input_layernorm): Pre-normalization before attention using RMSNorm
2. Self-Attention (self_attn): Multi-head attention with RoPE and Grouped Query Attention (GQA)
3. Post-Attention LayerNorm (post_attention_layernorm): Pre-normalization before MLP using RMSNorm
4. MLP (mlp): Gate-up-down projection network with GELU activation

The full model stack consists of:

• Token embeddings (embed_tokens)
• 26 decoder layers with native RMSNorm and RoPE
• Final normalization (norm)
• Language model head (lm_head)

Core Components

• RMSNorm: Root Mean Square normalization for improved training stability
• RoPE: Rotary Position Embeddings enabling excellent length extrapolation
• GQA: Grouped Query Attention with 8 query heads and 4 key-value heads for efficient inference
• SwiGLP-style MLP: Gate-projected feed-forward networks for enhanced expressivity

Weight Mapping

The custom architecture includes a built-in weight conversion utility that transparently maps checkpoint keys to the VedikaVyom parameter names.

Automatic Conversion

When loading from a compatible checkpoint, the following mapping is applied automatically:

Manual Weight Conversion (Optional)

If you need to convert weights manually:

from vedika_vyom import convert_checkpoint_to_vedikavyom

# Convert state dict
original_state_dict = torch.load("checkpoint.pt")
converted_state_dict = convert_checkpoint_to_vedikavyom(original_state_dict)

# Save converted weights
torch.save(converted_state_dict, "vedikavyom_weights.pt")

Verification

The implementation has been verified to produce consistent logits (within floating-point tolerance) across different runs with identical inputs and weights.

Citation

If you use Vedika-Vyom-4b-v1-mini in your research or applications, please cite:

@misc{vedika-vyom-4b-v1-mini,
  title={Vedika-Vyom-4b-v1-mini: A Custom Architecture for Efficient Text Generation},
  author={Veda Labs Team},
  year={2025},
  url={https://huggingface.co/Veda-Labs/Vedika-Vyom-4b-v1-mini}
}

Contact & Links

Official Resources

Support

• Email: vedalabs.veda@gmail.com
• Issues: Please report bugs or feature requests on our GitHub Issues

Built with ❤️ by Veda Labs

Licensed under Apache 2.0