adjaxt 0.0.1__py3-none-any.whl

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/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/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/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/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)
+
+#=====================================================================================================

adjaxt/sharding.py

@@ -0,0 +1,251 @@
+import functools
+import glob
+import itertools
+import json
+import operator
+import os
+import re
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from safetensors.numpy import load_file, save_file
+
+# ==============================================================================
+# 1. Spec & Model Definitions
+# ==============================================================================
+
+@dataclass
+class WeightSpec:
+    jax_path: str
+    hf_path: str
+    transpose: bool = False
+
+@dataclass
+class ModelWeightMap:
+    specs: List[WeightSpec]
+
+# ==============================================================================
+# 2. Path Traversal & Pattern Parsing
+# ==============================================================================
+
+def get_nested_val(d: dict, path: List[Union[str, int]]) -> Any:
+    return functools.reduce(operator.getitem, path, d)
+
+def set_nested_val(d: dict, path: List[Union[str, int]], val: Any) -> None:
+    def traverse_or_create(curr: Any, pair: tuple) -> Any:
+        key, next_key = pair
+        default_next = [] if isinstance(next_key, int) else {}
+
+        if isinstance(curr, list):
+            while len(curr) <= key:
+                curr.append(default_next.copy())
+            return curr[key]
+        return curr.setdefault(key, default_next)
+
+    lookahead_pairs = list(zip(path[:-1], path[1:]))
+    target_container = functools.reduce(traverse_or_create, lookahead_pairs, d)
+    
+    last_key = path[-1]
+    if isinstance(target_container, list):
+        while len(target_container) <= last_key:
+            target_container.append(None)
+    target_container[last_key] = val
+
+def parse_path_pattern(pattern: str, **kwargs) -> List[Union[str, int]]:
+    formatted = pattern.format(**kwargs)
+    return [int(t) if t.isdigit() else t for t in formatted.split(".")]
+
+# ==============================================================================
+# 3. Automatic Stack/Dimension Expansion
+# ==============================================================================
+
+def _expand_spec(spec: WeightSpec, dim_sizes: Dict[str, int]):
+    """
+    Finds placeholders shared in jax_path (loop dimensions) vs placeholders 
+    only present in hf_path (dimensions that need stacking/unstacking).
+    """
+    jax_placeholders = set(re.findall(r"\{([a-zA-Z0-9_]+)\}", spec.jax_path))
+    hf_placeholders = set(re.findall(r"\{([a-zA-Z0-9_]+)\}", spec.hf_path))
+
+    shared_placeholders = list(jax_placeholders & hf_placeholders)
+    stack_placeholders = list(hf_placeholders - jax_placeholders)
+
+    # Validate that all required dimensions exist in dim_sizes
+    for p in shared_placeholders + stack_placeholders:
+        if p not in dim_sizes:
+            raise KeyError(f"Missing dimension limit for '{{{p}}}' in dim_sizes.")
+
+    # 1. Expand standard outer loop/layer combinations
+    if not shared_placeholders:
+        outer_combos = [{}]
+    else:
+        shared_ranges = [range(dim_sizes[p]) for p in shared_placeholders]
+        outer_combos = [dict(zip(shared_placeholders, idxs)) for idxs in itertools.product(*shared_ranges)]
+
+    # 2. Generate sub-combos for stacked dimensions (e.g., experts)
+    if not stack_placeholders:
+        stack_combos = [{}]
+    else:
+        stack_ranges = [range(dim_sizes[p]) for p in stack_placeholders]
+        stack_combos = [dict(zip(stack_placeholders, idxs)) for idxs in itertools.product(*stack_ranges)]
+
+    for outer in outer_combos:
+        yield spec, outer, stack_combos
+
+# ==============================================================================
+# 4. Checkpoint Loading & Saving
+# ==============================================================================
+
+def load_checkpoint(
+    checkpoint_path: str, 
+    weight_map: ModelWeightMap, 
+    dim_sizes: Optional[Dict[str, int]] = None
+) -> dict:
+    dim_sizes = dim_sizes or {}
+
+    if os.path.isdir(checkpoint_path):
+        index_file = os.path.join(checkpoint_path, "model.safetensors.index.json")
+        if os.path.exists(index_file):
+            with open(index_file, "r") as f:
+                index = json.load(f)
+            files = {os.path.join(checkpoint_path, v) for v in index["weight_map"].values()}
+        else:
+            files = glob.glob(os.path.join(checkpoint_path, "*.safetensors"))
+    else:
+        files = [checkpoint_path]
+
+    flat_tensors = {}
+    for f in files:
+        flat_tensors.update(load_file(f))
+
+    weights = {}
+    for spec in weight_map.specs:
+        for expanded_spec, outer_kwargs, stack_kwargs_list in _expand_spec(spec, dim_sizes):
+            jax_path_tokens = parse_path_pattern(expanded_spec.jax_path, **outer_kwargs)
+
+            # Case A: Multiple HuggingFace keys collapse/stack into one JAX tensor
+            if len(stack_kwargs_list) > 1:
+                stacked_tensors = []
+                for stack_kw in stack_kwargs_list:
+                    hf_key = expanded_spec.hf_path.format(**outer_kwargs, **stack_kw)
+                    t = flat_tensors[hf_key]
+                    if expanded_spec.transpose:
+                        t = t.T
+                    stacked_tensors.append(t)
+                set_nested_val(weights, jax_path_tokens, jnp.stack(stacked_tensors, axis=0))
+            
+            # Case B: Standard 1-to-1 weight mapping
+            else:
+                hf_key = expanded_spec.hf_path.format(**outer_kwargs)
+                t = flat_tensors[hf_key]
+                if expanded_spec.transpose:
+                    t = t.T
+                set_nested_val(weights, jax_path_tokens, jnp.array(t))
+
+    return weights
+
+
+def _parse_size_in_bytes(size: Union[int, str]) -> int:
+    """Converts human-readable size strings (e.g. '5GB', '500MB') to byte count."""
+    if isinstance(size, int):
+        return size
+    units = {"B": 1, "KB": 10**3, "MB": 10**6, "GB": 10**9, "TB": 10**12,
+             "KIB": 1 << 10, "MIB": 1 << 20, "GIB": 1 << 30, "TIB": 1 << 40}
+    match = re.match(r"^([0-9.]+)\s*([a-zA-Z]+)$", size.strip())
+    if not match:
+        raise ValueError(f"Invalid max_shard_size format: {size}")
+    val, unit = match.groups()
+    unit = unit.upper()
+    if unit not in units:
+        raise ValueError(f"Unknown size unit: {unit}")
+    return int(float(val) * units[unit])
+
+
+def save_checkpoint(
+    weights: dict, 
+    save_directory: str, 
+    weight_map: ModelWeightMap, 
+    dim_sizes: Optional[Dict[str, int]] = None,
+    max_shard_size: Union[str, int] = "5GB"
+) -> None:
+    """
+    Saves a nested JAX weight dictionary to single or multi-shard safetensors 
+    with a HuggingFace-compatible index file.
+    """
+    dim_sizes = dim_sizes or {}
+    max_shard_bytes = _parse_size_in_bytes(max_shard_size)
+    os.makedirs(save_directory, exist_ok=True)
+
+    # 1. Flatten and convert JAX weights into numpy dict
+    flat_tensors: Dict[str, np.ndarray] = {}
+    for spec in weight_map.specs:
+        for expanded_spec, outer_kwargs, stack_kwargs_list in _expand_spec(spec, dim_sizes):
+            jax_path_tokens = parse_path_pattern(expanded_spec.jax_path, **outer_kwargs)
+            val = get_nested_val(weights, jax_path_tokens)
+
+            # Case A: Tensor unstacks across missing dimensions (e.g., experts)
+            if len(stack_kwargs_list) > 1:
+                for idx, stack_kw in enumerate(stack_kwargs_list):
+                    hf_key = expanded_spec.hf_path.format(**outer_kwargs, **stack_kw)
+                    t = val[idx]
+                    if expanded_spec.transpose:
+                        t = t.T
+                    flat_tensors[hf_key] = np.asarray(jax.device_get(t))
+
+            # Case B: Standard 1-to-1 weight mapping
+            else:
+                hf_key = expanded_spec.hf_path.format(**outer_kwargs)
+                t = val
+                if expanded_spec.transpose:
+                    t = t.T
+                flat_tensors[hf_key] = np.asarray(jax.device_get(t))
+
+    # 2. Partition tensors into shards based on max byte size
+    total_size = sum(t.nbytes for t in flat_tensors.values())
+    
+    # Save as a single un-sharded file if within size limit
+    if total_size <= max_shard_bytes:
+        single_path = os.path.join(save_directory, "model.safetensors")
+        save_file(flat_tensors, single_path)
+        return
+
+    shards: List[Dict[str, np.ndarray]] = []
+    current_shard: Dict[str, np.ndarray] = {}
+    current_shard_size = 0
+
+    for name, tensor in flat_tensors.items():
+        tensor_bytes = tensor.nbytes
+        if current_shard and (current_shard_size + tensor_bytes > max_shard_bytes):
+            shards.append(current_shard)
+            current_shard = {}
+            current_shard_size = 0
+
+        current_shard[name] = tensor
+        current_shard_size += tensor_bytes
+
+    if current_shard:
+        shards.append(current_shard)
+
+    # 3. Write individual shard files and compile the index
+    num_shards = len(shards)
+    weight_to_file = {}
+
+    for idx, shard in enumerate(shards, start=1):
+        filename = f"model-{idx:05d}-of-{num_shards:05d}.safetensors"
+        file_path = os.path.join(save_directory, filename)
+        save_file(shard, file_path)
+
+        for tensor_name in shard.keys():
+            weight_to_file[tensor_name] = filename
+
+    # 4. Write model.safetensors.index.json
+    index_data = {
+        "metadata": {"total_size": total_size},
+        "weight_map": weight_to_file
+    }
+    index_path = os.path.join(save_directory, "model.safetensors.index.json")
+    with open(index_path, "w", encoding="utf-8") as f:
+        json.dump(index_data, f, indent=2)

adjaxt-0.0.1.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,15 @@
+adjaxt/__init__.py,sha256=9ku486EWxqPGr1CgD7DGLwgnyAMM1bPAG00vonjLBsg,809
+adjaxt/approx.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+adjaxt/checkpoint.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+adjaxt/config.py,sha256=C-F3N866PiSxOAWQw5N1vta49-ReH7YEt6_svSK17xM,2375
+adjaxt/data.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+adjaxt/hub.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+adjaxt/init.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+adjaxt/layers.py,sha256=_B5nTfPgqDVYAAI5dJ_nZyD4R02sOQM9OrqGoU76QUI,10847
+adjaxt/model_maps.py,sha256=U4xLXywDya1vaDkycE7flKsVJsDNCWxi_cBU2_ewAoE,1761
+adjaxt/models.py,sha256=_q80rhx5sxNsKGLODl24NGkintjGIN4M27cucbb-9dI,5372
+adjaxt/sharding.py,sha256=9hpNiQNZOpgcD8teNV1WI9R8pIWIfhIyNCtVJKGUm0w,9840
+adjaxt/train.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+adjaxt-0.0.1.dist-info/METADATA,sha256=8A2_A-CD5Exbdf_iNa6EnfGYo_X30xihNiGL_S9jnRM,683
+adjaxt-0.0.1.dist-info/WHEEL,sha256=zOwg4jB6zX2kU910N-cMawjivD6tO8NEWvE12je1bVk,87
+adjaxt-0.0.1.dist-info/RECORD,,

adjaxt-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.0
+Root-Is-Purelib: true
+Tag: py3-none-any