adjaxt 0.0.1__tar.gz

adjaxt-0.0.1/PKG-INFO

@@ -0,0 +1,22 @@
+Metadata-Version: 2.5
+Name: adjaxt
+Version: 0.0.1
+Summary: Distributed training with JAX on free compute, for humans
+License-Expression: MIT
+Requires-Python: >=3.11
+Requires-Dist: datasets>=2.19
+Requires-Dist: huggingface-hub>=0.23
+Requires-Dist: jax>=0.4.30
+Requires-Dist: numpy
+Requires-Dist: optax>=0.2.2
+Requires-Dist: orbax-checkpoint>=0.5
+Requires-Dist: safetensors>=0.4
+Provides-Extra: cuda
+Requires-Dist: jax[cuda12]>=0.4.30; extra == 'cuda'
+Provides-Extra: test
+Requires-Dist: pytest; extra == 'test'
+Requires-Dist: torch; extra == 'test'
+Requires-Dist: transformers; extra == 'test'
+Description-Content-Type: text/markdown
+
+Hi, this package is in development, not ready yet

adjaxt-0.0.1/README.md

@@ -0,0 +1 @@
+Hi, this package is in development, not ready yet

adjaxt-0.0.1/pyproject.toml

@@ -0,0 +1,33 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "adjaxt"
+version = "0.0.1"
+description = "Distributed training with JAX on free compute, for humans"
+readme = "README.md"
+requires-python = ">=3.11"
+license = "MIT"
+dependencies = [
+    "jax>=0.4.30",
+    "optax>=0.2.2",
+    "orbax-checkpoint>=0.5",
+    "datasets>=2.19",
+    "safetensors>=0.4",
+    "huggingface-hub>=0.23",
+    "numpy",
+]
+
+[project.optional-dependencies]
+cuda = [
+    "jax[cuda12]>=0.4.30",
+]
+test = [
+    "pytest",
+    "transformers",
+    "torch",
+]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/adjaxt"]

adjaxt-0.0.1/src/adjaxt/__init__.py

@@ -0,0 +1,36 @@
+"""adjaxt — distributed training with JAX on free compute, made not sloppy."""
+
+from adjaxt.config import (
+    RMSNormConfig,
+    SwiGLUConfig,
+    StandardAttnImplementation,
+    GQAAttnConfig,
+    Qwen3AttnConfig,
+    Qwen3MLPConfig,
+    Qwen3MoEBlockConfig,
+    Qwen3MoELayerConfig,
+    Qwen3MoEModelConfig,
+    compute_rope_freqs,
+)
+from adjaxt import layers, models, model_maps, sharding
+
+__version__ = "0.0.1"
+
+__all__ = [
+    # Configs
+    "RMSNormConfig",
+    "SwiGLUConfig",
+    "StandardAttnImplementation",
+    "GQAAttnConfig",
+    "Qwen3AttnConfig",
+    "Qwen3MLPConfig",
+    "Qwen3MoEBlockConfig",
+    "Qwen3MoELayerConfig",
+    "Qwen3MoEModelConfig",
+    "compute_rope_freqs",
+    # Submodules
+    "layers",
+    "models",
+    "model_maps",
+    "sharding",
+]

adjaxt-0.0.1/src/adjaxt/config.py

@@ -0,0 +1,88 @@
+from dataclasses import dataclass, field
+from typing import Dict, Any, Tuple
+from enum import StrEnum
+import jax
+import jax.numpy as jnp
+
+@dataclass(frozen=True)
+class RMSNormConfig:
+    dim: int
+    eps: float = 1e-6
+
+@dataclass(frozen=True)
+class SwiGLUConfig:
+    in_dim: int
+    hidden_dim: int
+
+class StandardAttnImplementation(StrEnum):
+    CUDNN = "cudnn"
+    XLA = "xla"
+
+@dataclass(frozen=True)
+class GQAAttnConfig:
+    implementation: StandardAttnImplementation
+    num_kv_groups: int
+    is_causal: bool
+
+def compute_rope_freqs(seq_len: int, head_dim: int, theta: float = 10000.0):
+    dim_indices = jnp.arange(0, head_dim, 2, dtype=jnp.float32)
+    inv_freq = 1.0 / (theta ** (dim_indices / head_dim))
+    t = jnp.arange(seq_len, dtype=jnp.float32)
+    freqs = jnp.outer(t, inv_freq)
+    
+    freqs = jnp.concatenate([freqs, freqs], axis=-1)
+    
+    cos = jnp.cos(freqs)[None, :, None, :]
+    sin = jnp.sin(freqs)[None, :, None, :]
+    return cos, sin
+
+@dataclass(frozen=True)
+class Qwen3AttnConfig:
+    gqa_conf: GQAAttnConfig
+    q_rms_conf: RMSNormConfig
+    k_rms_conf: RMSNormConfig
+    d: int
+    head_dim: int
+    rope_theta: float
+    cos_table: jax.Array = field(init=False, compare=False, hash=False)
+    sin_table: jax.Array = field(init=False, compare=False, hash=False)
+    max_position_embeddings: int
+    n_layers: int
+
+    def __post_init__(self):
+        cos, sin = compute_rope_freqs(
+            self.max_position_embeddings, 
+            self.head_dim, 
+            self.rope_theta
+        )
+        object.__setattr__(self, "cos_table", cos)
+        object.__setattr__(self, "sin_table", sin)
+
+@dataclass(frozen=True)
+class Qwen3MLPConfig:
+    act_fn: function
+    in_dim: int
+    hidden_dim: int
+
+@dataclass(frozen=True)
+class Qwen3MoEBlockConfig:
+    mlp_conf: Qwen3MLPConfig
+    top_k: int
+    num_experts: int
+    d_model: int
+
+@dataclass(frozen=True)
+class Qwen3MoELayerConfig:
+    input_rms_conf: RMSNormConfig
+    attn_conf: Qwen3AttnConfig
+    post_attn_rms_conf: RMSNormConfig
+    moe_block_conf: Qwen3MoEBlockConfig
+
+@dataclass(frozen=True)
+class Qwen3MoEModelConfig:
+    moe_layer_conf: Qwen3MoELayerConfig
+    final_rms_conf: RMSNormConfig
+    num_decoder_blocks: int
+    vocab_size: int
+    d_model: int
+    tie_word_embeddings: bool = False

adjaxt-0.0.1/src/adjaxt/layers.py

@@ -0,0 +1,322 @@
+import jax, jax.numpy as jnp
+import math
+from adjaxt.config import *
+
+#def init_normal_layer(key, shape, std: float) -> jax.Array:
+#    return jax.random.normal(key, shape, jnp.bfloat16) * std
+
+
+#===================================================================================
+#RMS Norm
+
+def rms_norm(x, w, cfg) -> jax.Array:
+    """
+    Args:
+        x:   [..., D]
+        w:   [D]
+    Returns:
+        out: [..., D]
+    """
+    x32 = x.astype(jnp.float32)
+    ms = jnp.mean(jnp.square(x32), axis=-1, keepdims=True)
+    return w * (x32 * jax.lax.rsqrt(ms + cfg.eps)).astype(x.dtype)
+
+def rms_norm_init(key, cfg: RMSNormConfig) -> jax.Array:
+    """
+    Returns:
+        w: [dim]
+    """
+    return jnp.ones((cfg.dim,), dtype=jnp.float32)
+
+#===================================================================================
+#RoPE
+
+def rotate_half(x: jax.Array) -> jax.Array:
+    """
+    Args:
+        x:   [B, S, H, D]
+    Returns:
+        out: [B, S, H, D]
+    """
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return jnp.concatenate((-x2, x1), axis=-1)
+
+def apply_rope(x: jax.Array, cos: jax.Array, sin: jax.Array) -> jax.Array:
+    return x * cos + rotate_half(x) * sin
+
+def apply_rot_pos_emb(q: jax.Array, k: jax.Array, cos: jax.Array, sin: jax.Array):
+    """
+    Args:
+        q:   [B, S, H_q, D]
+        k:   [B, S, H_kv, D]
+        cos: [1, S, 1, D]
+        sin: [1, S, 1, D]
+    Returns:
+        q_rot, k_rot: ([B, S, H_q, D], [B, S, H_kv, D])
+    """
+    # q: (batch, seq_len, num_heads, head_dim)
+    # cos: (1, seq_len, 1, head_dim)
+    q_rot = (q * cos) + (rotate_half(q) * sin)
+    k_rot = (k * cos) + (rotate_half(k) * sin)
+    return q_rot, k_rot
+
+#===================================================================================
+#SwiGLU
+
+def swiglu(x: jax.Array, w: dict, cfg: SwiGLUConfig) -> jax.Array:
+    return (jax.nn.silu(x @ w["w_gate"]) * (x @ w["w_up"])) @ w["w_down"]
+
+def swiglu_init(key, cfg: SwiGLUConfig) -> dict:
+    k1, k2, k3 = jax.random.split(key, 3)
+    def kaiming(k, shape):
+        return jax.random.normal(k, shape, dtype=jnp.float32) * jnp.sqrt(2.0 / shape[0])
+    return  {
+                "w_gate": kaiming(k1, (cfg.in_dim, cfg.hidden_dim)),
+                "w_up": kaiming(k2, (cfg.in_dim, cfg.hidden_dim)),
+                "w_down": kaiming(k3, (cfg.hidden_dim, cfg.in_dim))
+            }
+
+#===================================================================================
+#RepeatKV
+
+def repeat_kv(hidden_states: jax.Array, n_rep: int) -> jax.Array:
+    """
+    Expands KV heads to match query heads.
+    Args:
+        x:   [B, S, H_kv, D]
+    Returns:
+        out: [B, S, H_kv * n_rep, D]
+    """
+    if n_rep == 1:
+        return hidden_states
+    return jnp.repeat(hidden_states, n_rep, axis=2)
+
+#===================================================================================
+#GQAAttention
+
+def gqa_attn(
+        q: jax.Array,
+        k: jax.Array,
+        v: jax.Array,
+        cfg: GQAAttnConfig,
+        attn_mask: jax.Array = None
+    ) -> jax.Array:
+    """
+    Args:
+        q:         [B, S_q, H_q, D]
+        k:         [B, S_kv, H_kv, D]
+        v:         [B, S_kv, H_kv, D]
+        attn_mask: [B, H_q, S_q, S_kv] or None
+    Returns:
+        out:       [B, S_q, H_q, D]
+    """
+
+    batch, q_len, num_q_heads, head_dim = q.shape
+    _, kv_len, num_kv_heads, _ = k.shape
+
+    num_groups = num_q_heads // num_kv_heads
+
+    if num_groups > 1:
+        k = jnp.repeat(k, repeats=num_groups, axis=2)
+        v = jnp.repeat(v, repeats=num_groups, axis=2)
+
+    causal_flag = cfg.is_causal if attn_mask is None else False
+
+    out = jax.nn.dot_product_attention(
+        query=q,
+        key=k,
+        value=v,
+        mask=attn_mask,
+        is_causal=causal_flag,
+        implementation=cfg.implementation,
+    )
+
+    return out
+
+#===================================================================================
+#Qwen3Attn
+
+def qwen3_attn(
+    x: jax.Array,
+    w: dict,
+    cfg: Qwen3AttnConfig
+):
+    """
+    Args:
+        x:   [B, S, D_model]
+        w:   Dict containing projection weights and head norms
+    Returns:
+        out: [B, S, D_model]
+    """
+    seq_len = x.shape[1]
+    input_shape = x.shape[:-1]
+    hidden_shape = (*input_shape, -1, cfg.head_dim)
+    
+    # Shapes: (batch, seq_len, num_heads, head_dim)
+    q = rms_norm((x @ w["q_proj"]).reshape(hidden_shape), w["q_norm"], cfg.q_rms_conf)
+    k = rms_norm((x @ w["k_proj"]).reshape(hidden_shape), w["k_norm"], cfg.k_rms_conf)
+    v = (x @ w["v_proj"]).reshape(hidden_shape)  # keep (B, S, KV_H, D)
+    
+    # cos/sin shape: (1, seq_len, 1, head_dim)
+    cos = cfg.cos_table[:, :seq_len, :, :]
+    sin = cfg.sin_table[:, :seq_len, :, :]
+    q, k = apply_rot_pos_emb(q, k, cos, sin)
+    
+    attn_out = gqa_attn(
+        q,
+        k,
+        v,
+        cfg.gqa_conf
+    )
+    attn_out = attn_out.reshape(*input_shape, -1)
+    attn_out = attn_out @ w["o_proj"]
+    return attn_out
+
+
+def qwen3_attn_init(
+    key,
+    cfg: Qwen3AttnConfig
+) -> dict:
+    k1, k2, k3, k4, k5, k6 = jax.random.split(key, 6)
+    res = 1.0 / math.sqrt(2 * cfg.n_layers)
+    
+    # Calculate dimensions based on GQA heads
+    q_dim = cfg.d  # num_heads * head_dim
+    kv_dim = (cfg.d // cfg.gqa_conf.num_kv_groups)  # num_kv_heads * head_dim
+    
+    w = {
+        "q_proj": jax.random.normal(k1, (cfg.d, q_dim), jnp.bfloat16) / math.sqrt(cfg.d),
+        "k_proj": jax.random.normal(k2, (cfg.d, kv_dim), jnp.bfloat16) / math.sqrt(cfg.d),
+        "v_proj": jax.random.normal(k3, (cfg.d, kv_dim), jnp.bfloat16) / math.sqrt(cfg.d),
+        "o_proj": jax.random.normal(k4, (q_dim, cfg.d), jnp.bfloat16) * res / math.sqrt(q_dim),
+        "q_norm": jnp.ones((cfg.head_dim,), dtype=jnp.bfloat16),
+        "k_norm": jnp.ones((cfg.head_dim,), dtype=jnp.bfloat16),
+    }
+    return w
+
+#===================================================================================
+#Qwen3MLP
+
+def qwen3_mlp(x: jax.Array, w: dict, cfg: Qwen3MLPConfig) -> jax.Array:
+    """
+    Args:
+        x:   [..., D_in]
+        w:   {"w_gate": [D_in, D_hidden], "w_up": [D_in, D_hidden], "w_down": [D_hidden, D_in]}
+    Returns:
+        out: [..., D_in]
+    """
+    return (cfg.act_fn(x @ w["w_gate"]) * (x @ w["w_up"])) @ w["w_down"]
+
+def qwen3_mlp_init(key, cfg: Qwen3MLPConfig) -> dict:
+    k1, k2, k3 = jax.random.split(key, 3)
+    def kaiming(k, shape):
+        return jax.random.normal(k, shape, dtype=jnp.float32) * jnp.sqrt(2.0 / shape[0])
+    return  {
+                "w_gate": kaiming(k1, (cfg.in_dim, cfg.hidden_dim)),
+                "w_up": kaiming(k2, (cfg.in_dim, cfg.hidden_dim)),
+                "w_down": kaiming(k3, (cfg.hidden_dim, cfg.in_dim))
+            }
+
+#===================================================================================
+#Qwen3MoEBlock
+
+def qwen3_moe_block(x: jax.Array, w: dict, cfg: Qwen3MoEBlockConfig) -> jax.Array:
+    """
+    Args:
+        x:   [B, S, D_model]
+    Returns:
+        out: [B, S, D_model]
+    """
+    tokens = x.reshape(-1, x.shape[-1])                     # (T, d)
+    router_logits = tokens @ w["router"]                    # (T, E)
+    probs = jax.nn.softmax(router_logits.astype(jnp.float32), -1)
+    topk_probs, topk_idx = jax.lax.top_k(probs, cfg.top_k)      # (T, k)
+    topk_probs = topk_probs / topk_probs.sum(-1, keepdims=True)  # norm_topk_prob
+    # scatter the renormalized probs back to (T, E), zeros elsewhere
+    weights = jnp.zeros_like(probs).at[
+        jnp.arange(tokens.shape[0])[:, None], topk_idx
+    ].set(topk_probs)
+    expert_out = jax.vmap(
+        lambda we: qwen3_mlp(tokens, we, cfg.mlp_conf)
+    )(w["experts"])                                          # (E, T, d)
+    out = jnp.einsum("te,etd->td", weights.astype(x.dtype), expert_out)
+    return out.reshape(x.shape)
+
+def qwen3_moe_block_init(key, cfg: Qwen3MoEBlockConfig) -> dict:
+    k_router, k_experts = jax.random.split(key, 2)
+    expert_keys = jax.random.split(k_experts, cfg.num_experts)
+    return {
+        "router": jax.random.normal(k_router, (cfg.d_model, cfg.num_experts), dtype=jnp.bfloat16) * (cfg.d_model ** -0.5),
+        "experts": jax.vmap(lambda k: qwen3_mlp_init(k, cfg.mlp_conf))(expert_keys),
+    }
+
+#===================================================================================
+#Qwen3MoELayer
+
+def qwen3_moe_layer(x: jax.Array, w: dict, cfg: Qwen3MoELayerConfig) -> jax.Array:
+    """
+    Args:
+        x:   [B, S, D_model]
+    Returns:
+        out: [B, S, D_model]
+    """
+    residual = x
+    x = rms_norm(x, w["input_layernorm"], cfg.input_rms_conf)
+    x = qwen3_attn(x, w["attn"], cfg.attn_conf)
+    x = x + residual
+    residual = x
+    x = rms_norm(x, w["post_attn_layernorm"], cfg.post_attn_rms_conf)
+    x = qwen3_moe_block(x, w["mlp"], cfg.moe_block_conf)
+    x = residual + x
+    return x
+
+def qwen3_moe_layer_init(key, cfg: Qwen3MoELayerConfig) -> dict:
+    attn_key, mlp_key, k3, k4 = jax.random.split(key, 4)
+    return {
+        "attn": qwen3_attn_init(attn_key, cfg.attn_conf),
+        "mlp": qwen3_moe_block_init(mlp_key, cfg.moe_block_conf),
+        "input_layernorm": rms_norm_init(k3, cfg.input_rms_conf),
+        "post_attn_layernorm": rms_norm_init(k4, cfg.post_attn_rms_conf)
+    }
+
+#===================================================================================
+#Qwen3MoeModel
+
+def qwen3_moe_model(x: jax.Array, w: dict, cfg: Qwen3MoEModelConfig) -> jax.Array:
+    x = w["embeds"][x]
+    for l in w["decoder_blocks"]:
+        x = qwen3_moe_layer(
+            x,
+            l,
+            cfg.moe_layer_conf
+        )
+    x = rms_norm(x, w["norm"], cfg.final_rms_conf)
+    logits = x @ w["lm_head"]
+    return logits
+
+def qwen3_moe_model_init(key, cfg: Qwen3MoEModelConfig) -> dict:
+    k_embed, k_layers, k_norm, k_head = jax.random.split(key, 4)
+    embeds = jax.random.normal(k_embed, shape=(cfg.vocab_size, cfg.d_model)) * (
+        cfg.d_model**-0.5
+    )
+    layer_keys = jax.random.split(k_layers, cfg.num_decoder_blocks)
+    decoder_blocks = [
+        qwen3_moe_layer_init(k, cfg.moe_layer_conf) for k in layer_keys
+    ]
+    norm = rms_norm_init(k_norm, cfg.final_rms_conf)
+    if cfg.tie_word_embeddings:
+        lm_head = embeds.T
+    else:
+        lm_head = jax.random.normal(k_head, shape=(cfg.d_model, cfg.vocab_size)) * (
+            cfg.d_model**-0.5
+        )
+
+    return {
+        "embeds": embeds,
+        "decoder_blocks": decoder_blocks,
+        "norm": norm,
+        "lm_head": lm_head,
+    }
+
+#===================================================================================

adjaxt-0.0.1/src/adjaxt/model_maps.py

@@ -0,0 +1,27 @@
+from adjaxt.sharding import *
+
+QWEN3_MOE_WEIGHT_MAP = ModelWeightMap(
+    specs=[
+        WeightSpec("embeds", "model.embed_tokens.weight"),
+        WeightSpec("norm", "model.norm.weight"),
+        WeightSpec("lm_head", "lm_head.weight", transpose=True),
+        
+        # Layer Norms
+        WeightSpec("decoder_blocks.{i}.input_layernorm", "model.layers.{i}.input_layernorm.weight"),
+        WeightSpec("decoder_blocks.{i}.post_attn_layernorm", "model.layers.{i}.post_attention_layernorm.weight"),
+
+        # Attention Projections & QK Norms
+        WeightSpec("decoder_blocks.{i}.attn.q_proj", "model.layers.{i}.self_attn.q_proj.weight", transpose=True),
+        WeightSpec("decoder_blocks.{i}.attn.k_proj", "model.layers.{i}.self_attn.k_proj.weight", transpose=True),
+        WeightSpec("decoder_blocks.{i}.attn.v_proj", "model.layers.{i}.self_attn.v_proj.weight", transpose=True),
+        WeightSpec("decoder_blocks.{i}.attn.o_proj", "model.layers.{i}.self_attn.o_proj.weight", transpose=True),
+        WeightSpec("decoder_blocks.{i}.attn.q_norm", "model.layers.{i}.self_attn.q_norm.weight"),
+        WeightSpec("decoder_blocks.{i}.attn.k_norm", "model.layers.{i}.self_attn.k_norm.weight"),
+
+        # MoE Router & Stacked Experts
+        WeightSpec("decoder_blocks.{i}.mlp.router", "model.layers.{i}.mlp.gate.weight", transpose=True),
+        WeightSpec("decoder_blocks.{i}.mlp.experts.w_gate", "model.layers.{i}.mlp.experts.{e}.gate_proj.weight", transpose=True),
+        WeightSpec("decoder_blocks.{i}.mlp.experts.w_up", "model.layers.{i}.mlp.experts.{e}.up_proj.weight", transpose=True),
+        WeightSpec("decoder_blocks.{i}.mlp.experts.w_down", "model.layers.{i}.mlp.experts.{e}.down_proj.weight", transpose=True),
+    ]
+)

adjaxt-0.0.1/src/adjaxt/models.py

@@ -0,0 +1,148 @@
+from adjaxt.config import *
+from adjaxt.layers import *
+from adjaxt.sharding import *
+from adjaxt.model_maps import *
+from huggingface_hub import snapshot_download, hf_hub_download
+from typing import Callable
+
+#=====================================================================================================
+#Qwen3 MoE
+
+def qwen3_moe_model_load(cfg: Qwen3MoEModelConfig, ckpt_path: str, token: Optional[str] = None):
+    """
+    Loads weights into a Qwen3 MoE JAX model from either a local path 
+    or directly from a Hugging Face repository repo_id.
+    """
+    # 1. Resolve checkpoint path (download from HF hub if not a local folder/file)
+    if not os.path.exists(ckpt_path):
+        resolved_path = snapshot_download(
+            repo_id=ckpt_path,
+            allow_patterns=["*.safetensors", "*.json"],
+            token=token,
+        )
+    else:
+        resolved_path = ckpt_path
+
+    # 2. Extract loop/stack dimension limits dynamically from cfg
+    dim_sizes = {
+        "i": cfg.num_decoder_blocks,
+        "e": cfg.moe_layer_conf.moe_block_conf.num_experts,
+    }
+
+    # 3. Load flat safetensors and reconstruct the nested JAX tree
+    weights = load_checkpoint(
+        checkpoint_path=resolved_path,
+        weight_map=QWEN3_MOE_WEIGHT_MAP,
+        dim_sizes=dim_sizes,
+    )
+
+    # 4. Handle tied word embeddings if configured
+    if cfg.tie_word_embeddings:
+        weights["lm_head"] = weights["embeds"].T
+
+    return weights
+
+def _get_act_fn(act_name: str) -> Callable[[jax.Array], jax.Array]:
+    act_map = {
+        "silu": jax.nn.silu,
+        "swish": jax.nn.swish,
+        "gelu": jax.nn.gelu,
+        "relu": jax.nn.relu,
+    }
+    if act_name not in act_map:
+        raise ValueError(f"Unsupported activation function: {act_name}")
+    return act_map[act_name]
+
+def create_qwen3_moe_config(
+    hf_config: dict,
+    attn_implementation: StandardAttnImplementation = StandardAttnImplementation.XLA,
+) -> Qwen3MoEModelConfig:
+    """Builds an MoE Qwen3 configuration dataclass from a parsed HF config.json dictionary."""
+    hidden_size = hf_config["hidden_size"]
+    num_heads = hf_config["num_attention_heads"]
+    num_kv_heads = hf_config.get("num_key_value_heads", num_heads)
+    head_dim = hf_config.get("head_dim", hidden_size // num_heads)
+    rms_eps = hf_config.get("rms_norm_eps", 1e-6)
+    moe_intermediate_size = hf_config.get("moe_intermediate_size", hf_config.get("intermediate_size"))
+    num_layers = hf_config["num_hidden_layers"]
+    vocab_size = hf_config["vocab_size"]
+    rope_theta = float(hf_config.get("rope_theta", 1000000.0))
+    max_pos = hf_config.get("max_position_embeddings", 32768)
+    num_experts = hf_config.get("num_experts", hf_config.get("num_local_experts", 64))
+    top_k = hf_config.get("num_experts_per_tok", 8)
+    act_fn = _get_act_fn(hf_config.get("hidden_act", "silu"))
+
+    gqa_conf = GQAAttnConfig(
+        implementation=attn_implementation,
+        num_kv_groups=num_heads // num_kv_heads,
+        is_causal=True,
+    )
+    q_rms_conf = RMSNormConfig(dim=head_dim, eps=rms_eps)
+    k_rms_conf = RMSNormConfig(dim=head_dim, eps=rms_eps)
+
+    attn_conf = Qwen3AttnConfig(
+        gqa_conf=gqa_conf,
+        q_rms_conf=q_rms_conf,
+        k_rms_conf=k_rms_conf,
+        d=hidden_size,
+        head_dim=head_dim,
+        rope_theta=rope_theta,
+        max_position_embeddings=max_pos,
+        n_layers=num_layers,
+    )
+
+    mlp_conf = Qwen3MLPConfig(
+        act_fn=act_fn,
+        in_dim=hidden_size,
+        hidden_dim=moe_intermediate_size,
+    )
+
+    moe_block_conf = Qwen3MoEBlockConfig(
+        mlp_conf=mlp_conf,
+        top_k=top_k,
+        num_experts=num_experts,
+        d_model=hidden_size,
+    )
+
+    moe_layer_conf = Qwen3MoELayerConfig(
+        input_rms_conf=RMSNormConfig(dim=hidden_size, eps=rms_eps),
+        attn_conf=attn_conf,
+        post_attn_rms_conf=RMSNormConfig(dim=hidden_size, eps=rms_eps),
+        moe_block_conf=moe_block_conf,
+    )
+
+    return Qwen3MoEModelConfig(
+        moe_layer_conf=moe_layer_conf,
+        final_rms_conf=RMSNormConfig(dim=hidden_size, eps=rms_eps),
+        num_decoder_blocks=num_layers,
+        vocab_size=vocab_size,
+        d_model=hidden_size,
+        tie_word_embeddings=hf_config.get("tie_word_embeddings", False),
+    )
+
+
+def qwen3_moe_config_from_pretrained(
+    model_name_or_path: str,
+    attn_implementation: StandardAttnImplementation = StandardAttnImplementation.XLA,
+    token: Optional[str] = None,
+) -> Qwen3MoEModelConfig:
+    """Loads config.json from local disk or HuggingFace Hub and constructs the corresponding dataclass."""
+    if os.path.exists(model_name_or_path):
+        config_file = (
+            model_name_or_path
+            if os.path.isfile(model_name_or_path)
+            else os.path.join(model_name_or_path, "config.json")
+        )
+    else:
+        config_file = hf_hub_download(
+            repo_id=model_name_or_path,
+            filename="config.json",
+            token=token,
+        )
+
+    with open(config_file, "r", encoding="utf-8") as f:
+        hf_config = json.load(f)
+
+    return create_qwen3_moe_config(hf_config, attn_implementation)
+
+#=====================================================================================================