yujiepan / glm-4-moe-tiny-random

This tiny model is for debugging. It is randomly initialized with the config adapted from zai-org/GLM-4.5 .

Note: The transformers implementation does not have multi-token prediction (MTP) support. So you might see some "weights not loaded" warnings. This is expected.

Example usage:

• vLLM

model_id=yujiepan/glm-4-moe-tiny-random
vllm serve $model_id \
    --tensor-parallel-size 1 \
    --tool-call-parser glm4_moe \
    --reasoning-parser glm4_moe \
    --enable-auto-tool-choice

• SGLang

# Multi-token prediction is supported
model_id=yujiepan/glm-4-moe-tiny-random
python3 -m sglang.launch_server \
    --model-path $model_id \
    --tp-size 1 \
    --cuda-graph-max-bs 4 \
    --tool-call-parser glm45  \
    --reasoning-parser glm45 \
    --speculative-algorithm EAGLE \
    --speculative-num-steps 3 \
    --speculative-eagle-topk 1 \
    --speculative-num-draft-tokens 4 \
    --mem-fraction-static 0.4

• Transformers

from transformers import pipeline
model_id = "yujiepan/glm-4-moe-tiny-random"
pipe = pipeline(
    "text-generation", model=model_id, device="cuda",
    trust_remote_code=True, max_new_tokens=20,
)
print(pipe("Hello World!"))

Codes to create this repo:

from copy import deepcopy

import torch
import torch.nn as nn
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    GenerationConfig,
    pipeline,
    set_seed,
)
from transformers.models.glm4_moe.modeling_glm4_moe import Glm4MoeDecoderLayer, Glm4MoeRMSNorm

source_model_id = "zai-org/GLM-4.5"
save_folder = "/tmp/yujiepan/glm-4-moe-tiny-random"

tokenizer = AutoTokenizer.from_pretrained(
    source_model_id, trust_remote_code=True,
)
tokenizer.save_pretrained(save_folder)

config = AutoConfig.from_pretrained(
    source_model_id, trust_remote_code=True,
)
config.hidden_size = 16
config.head_dim = 64
config.intermediate_size = 64
config.num_attention_heads = 4
config.num_hidden_layers = 2  # 1 dense, 1 moe
config.num_key_value_heads = 2
config.moe_intermediate_size = 64
config.n_routed_experts = 16
config.n_shared_experts = 1
config.first_k_dense_replace = 1
config.num_experts_per_tok = 8
config.num_nextn_predict_layers = 1  # after layer 0 and 1, there will be a another MTP layer
config.tie_word_embeddings = True

torch.set_default_dtype(torch.bfloat16)
model = AutoModelForCausalLM.from_config(
    config,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
)

class SharedHead(nn.Module):
    def __init__(self, config) -> None:
        super().__init__()
        self.norm = Glm4MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        # self.head = deepcopy(model.get_output_embeddings())

class Glm4MoeDecoderMTP(Glm4MoeDecoderLayer):
    def __init__(self, config, layer_idx):
        super().__init__(config, layer_idx=layer_idx)
        self.enorm = Glm4MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.hnorm = Glm4MoeRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
        self.shared_head = SharedHead(config=config)
        # self.embed_tokens = deepcopy(model.get_input_embeddings())

last_extra_layer = Glm4MoeDecoderMTP(config, layer_idx=config.num_hidden_layers)
model.model.layers.append(last_extra_layer)
model.generation_config = GenerationConfig.from_pretrained(
    source_model_id, trust_remote_code=True,
)
set_seed(42)
with torch.no_grad():
    for name, p in sorted(model.named_parameters()):
        torch.nn.init.normal_(p, 0, 0.2)
        print(name, p.shape)
model.save_pretrained(save_folder)