You need to agree to share your contact information to access this model

This repository is publicly accessible, but you have to accept the conditions to access its files and content.

large-moe-1b

Mobile-optimized MoE model configuration from architecture search.

Metrics

Best Training Loss: 1.9496
Best Eval Loss: 4.7961
Total Parameters: 1003.0M
Active Parameters: 399.0M
Steps Completed: 2000

Usage

# Load the model
from safetensors.torch import load_file
from architecture.model import Qwen3Model  # Your custom model class

# Load config
import json
with open("config.json") as f:
    config = json.load(f)

# Initialize model
model = Qwen3Model(config)

# Load weights
state_dict = load_file("model.safetensors")
model.load_state_dict(state_dict)

# Load tokenizer
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("kshitijthakkar/moe-1002m-399m-8x2-16L-large-moe-1b")

Model Configuration

{
  "vocab_size": 151936,
  "emb_dim": 1024,
  "n_heads": 16,
  "n_layers": 16,
  "n_kv_groups": 4,
  "num_experts": 8,
  "num_experts_per_tok": 2,
  "moe_hidden_dim": 2048,
  "head_dim": 64,
  "max_position_embeddings": 4096,
  "rope_base": 1000000.0,
  "qk_norm": true
}

Training Configuration

{
  "model_config": {
    "vocab_size": 151936,
    "emb_dim": 1024,
    "n_heads": 16,
    "n_layers": 16,
    "n_kv_groups": 4,
    "num_experts": 8,
    "num_experts_per_tok": 2,
    "moe_hidden_dim": 2048,
    "head_dim": 64,
    "max_position_embeddings": 4096,
    "rope_base": 1000000.0,
    "qk_norm": true
  },
  "learning_rate": 0.0001,
  "batch_size": 4,
  "context_length": 1024,
  "warmup_ratio": 0.1,
  "warmup_steps": null,
  "weight_decay": 0.1,
  "gradient_clip": 1.0,
  "gradient_accumulation_steps": 1,
  "scheduler_type": "cosine",
  "wsd_decay_ratio": 0.1,
  "max_steps": 2000,
  "eval_steps": 500,
  "eval_batches": 20,
  "log_steps": 100,
  "early_stopping": true,
  "early_stopping_patience": 500,
  "early_stopping_min_delta": 0.01,
  "early_stopping_min_steps": 200,
  "track_expert_balance": true,
  "expert_balance_log_steps": 100,
  "use_wandb": true,
  "wandb_project": "moe-large-search",
  "wandb_entity": null,
  "wandb_tags": [
    "large-moe-1b",
    "architecture-search"
  ],
  "train_data_path": null,
  "val_data_path": null,
  "output_dir": null,
  "experiment_name": "large-moe-1b",
  "device": "cuda",
  "dtype": "bfloat16",
  "gradient_checkpointing": true,
  "architecture_name": "large-moe-1b",
  "mobile_estimate": {
    "tok_per_sec_fp16": 14.162501050987643,
    "tok_per_sec_q8": 23.60416841831274,
    "tok_per_sec_q4": 33.045835785637834,
    "ttft_ms_fp16": 364.72608581818184,
    "ttft_ms_q8": 243.1507238787879,
    "ttft_ms_q4": 194.52057910303031,
    "memory_mb_fp16": 1978.666015625,
    "memory_mb_q8": 1109.9862304687501,
    "memory_mb_q4": 627.8197265625,
    "total_params": 1002996736,
    "active_params": 399016960,
    "meets_ttft_target": false,
    "meets_throughput_target": false,
    "meets_memory_target": false
  }
}