Publish Vortex Alpha research preview

README.md

@@ -0,0 +1,157 @@
+---
+language: en
+pipeline_tag: text-generation
+library_name: pytorch
+tags:
+- causal-lm
+- small-language-model
+- research
+- text-generation
+license: other
+thumbnail: Aurora-5.png
+---
+
+# Vortex Alpha
+
+Vortex is the working name for a compact, experimental language model. The
+final public name has not been decided. This release is intended for research,
+local experimentation, and further fine-tuning—not as a finished general
+assistant.
+
+## What is included
+
+- `model.safetensors`: the instruction/tool-format preview selected from the
+  best small internal behavior pilot.
+- `base_model.safetensors`: the corresponding pretrained text-completion
+  base.
+- `config.json`: the architecture configuration.
+- `tokenizer.model`: the 8,192-piece SentencePiece tokenizer used for both
+  checkpoints.
+- `vortex_model.py` and `inference.py`: a minimal dependency-light PyTorch
+  loader and sampler.
+- `Aurora-5.png`: the project thumbnail/banner.
+
+Optimizer state, private logs, local paths, credentials, and training-machine
+metadata are intentionally not included.
+
+## Architecture
+
+Vortex is a dense decoder-only Transformer with 174,942,720 trainable
+parameters:
+
+| Component | Parameters |
+| --- | ---: |
+| Shared token embedding and tied output head | 8,388,608 |
+| Attention Q projections | 12,582,912 |
+| Attention K projections | 3,145,728 |
+| Attention V projections | 3,145,728 |
+| Attention output projections | 12,582,912 |
+| Per-head QK RMSNorm parameters | 1,536 |
+| SwiGLU feed-forward networks | 135,069,696 |
+| Transformer-block RMSNorm parameters | 24,576 |
+| Final RMSNorm | 1,024 |
+| **Total** | **174,942,720** |
+
+Configuration: 12 layers, hidden size 1,024, 16 query heads, 4 key/value
+heads, 64-dimensional heads, SwiGLU with intermediate size 3,664, pre-layer
+RMSNorm, per-head QK-Norm, RoPE with base 100,000, bias-free projections,
+8,192-token vocabulary, and a 4,096-token training/inference limit.
+
+The published weights use the readable reference PyTorch layout. They do not
+require Transformer Engine to load. The input and output embeddings are tied.
+
+## Quick start
+
+The minimal reference runner uses PyTorch, SentencePiece, and
+`safetensors`:
+
+```bash
+python -m pip install torch sentencepiece safetensors
+python inference.py \
+  --weights model.safetensors \
+  --chat \
+  --prompt "Explain why the sky appears blue in two short paragraphs."
+```
+
+For ordinary next-token completion, use the base checkpoint:
+
+```bash
+python inference.py \
+  --weights base_model.safetensors \
+  --prompt "The sky appears blue because"
+```
+
+The reference runner recomputes the full prefix at each generated token and
+is deliberately simple. A production runner should add a KV cache and a
+fused attention implementation.
+
+The instruction preview was tuned with this compact serialization:
+
+```text
+[SYSTEM]
+You are a helpful assistant. Follow instructions, answer clearly, and say
+when information is missing.
+</s>
+[USER]
+Your question here
+</s>
+[ASSISTANT]
+```
+
+The instruction preview may emit a `CALL {json}` calculator/search request
+when prompted for tool use. No tool server is included in this repository;
+without a tool runner, treat such output as ordinary text. The base checkpoint
+is the better starting point for continued pretraining.
+
+## Training summary
+
+The base run used an approximate mixture of FineWeb, DCLM, educational/math
+material, and The Stack v3 code data. The recorded base checkpoint had seen
+about 8.48 billion pretraining tokens. Training used BF16 model computation,
+FP8-capable NVIDIA kernels where available, a WSD-style learning-rate tail,
+and token-budgeted batches designed for a 16 GB consumer GPU. The instruction
+preview is a lightweight supervised derivative of that base; its optimizer
+state and private training records are not part of this release.
+
+These data-mixture descriptions are a project-level summary, not a claim that
+every upstream document is suitable for every downstream use. Follow the
+licenses and terms of the upstream datasets.
+
+## Evaluation snapshot
+
+These are exploratory measurements, not official leaderboard submissions.
+The base results used greedy decoding, no tools, and the stated sample sizes:
+
+| Test | Result | Notes |
+| --- | ---: | --- |
+| MMLU cloze sample | 511/2,000 = 25.55% | Wilson 95% interval: 23.69–27.51% |
+| GSM8K strict numeric sample | 0/256 = 0.00% | Wilson 95% upper bound: 1.48% |
+| GSM8K fallback numeric sample | 2/256 = 0.78% | Wilson 95% interval: 0.21–2.80% |
+
+The instruction checkpoint reached 9/12 arithmetic, 4/4 grounding, 4/4
+abstention, 2/4 exact-format, and 1/2 JSON checks on a 26-prompt internal
+tool-format pilot when the calculator runner was available. That pilot is too
+small to support general capability claims, and the arithmetic result is not
+comparable to tool-free GSM8K.
+
+The results show why this is an alpha release: the model can produce useful
+local completions and structured tool calls, but it remains weak at reliable
+arithmetic, broad knowledge, long-form coherence, and hallucination control.
+
+## Limitations and intended use
+
+Vortex is a small research model. It can be repetitive, overconfident, or
+factually wrong; architecture and training scale do not guarantee reliable
+answers. Do not use it as the sole basis for medical, legal, financial,
+safety-critical, or other high-stakes decisions. It has not been evaluated for
+privacy, bias, cybersecurity, or comprehensive safety.
+
+The repository is public, but no open-source license is asserted yet. The
+final name, licensing terms, and a production release decision are still open.
+
+## Reproducibility note
+
+The conversion removed optimizer state and Transformer Engine-only auxiliary
+state, then wrote the model tensors in BF16 safetensors format. The exported
+reference tensors preserve the tied-embedding model weights and can be loaded
+with the included `vortex_model.py`.

RELEASE_MANIFEST.json

@@ -0,0 +1,35 @@
+{
+  "model": "Vortex Alpha",
+  "name_status": "working name; final public name undecided",
+  "repo_id": "North-ML1/vortex-alpha",
+  "parameter_count": 174942720,
+  "tensor_dtype": "bfloat16",
+  "architecture": {
+    "layers": 12,
+    "hidden_size": 1024,
+    "attention_heads": 16,
+    "kv_heads": 4,
+    "head_dim": 64,
+    "ffn_intermediate_size": 3664,
+    "vocab_size": 8192,
+    "context_length": 4096
+  },
+  "files": {
+    "model.safetensors": {
+      "bytes": 349898992,
+      "sha256": "330ae517cc9bd14d055f1bbde3b330b6d3206a8740bb4ef032a771af708a96b7"
+    },
+    "base_model.safetensors": {
+      "bytes": 349898984,
+      "sha256": "a00ad274d9e5b8608f69176937a47b4c55499bc810cb0311c4dc93cae1a890c3"
+    },
+    "tokenizer.model": {
+      "bytes": 123865,
+      "sha256": "8ebb450f2786d3153d64f874d97487ba5cb8e3450b799bbf6aeddd7d92f7ff73"
+    },
+    "Aurora-5.png": {
+      "bytes": 96932,
+      "sha256": "d444edf5ecab1a42e54719881182ba45627e268dff6b6b55fa1dff39f06d5ead"
+    }
+  }
+}

base_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a00ad274d9e5b8608f69176937a47b4c55499bc810cb0311c4dc93cae1a890c3
+size 349898984

config.json

@@ -0,0 +1,18 @@
+{
+  "model_type": "vortex",
+  "architectures": ["VortexForCausalLM"],
+  "vocab_size": 8192,
+  "hidden_size": 1024,
+  "num_hidden_layers": 12,
+  "num_attention_heads": 16,
+  "num_key_value_heads": 4,
+  "head_dim": 64,
+  "intermediate_size": 3664,
+  "hidden_act": "silu",
+  "max_position_embeddings": 4096,
+  "rope_theta": 100000.0,
+  "rms_norm_eps": 1e-5,
+  "tie_word_embeddings": true,
+  "attention_bias": false,
+  "torch_dtype": "bfloat16"
+}

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "eos_token_id": 1,
+  "do_sample": true,
+  "temperature": 0.8,
+  "top_p": 0.95,
+  "max_new_tokens": 256
+}

inference.py

@@ -0,0 +1,114 @@
+"""Minimal reference sampler for the public Vortex Alpha release."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from safetensors.torch import load_file
+
+from vortex_model import VortexConfig, VortexForCausalLM
+
+
+def load_config(path: Path) -> VortexConfig:
+    data = json.loads(path.read_text())
+    return VortexConfig(
+        vocab_size=int(data["vocab_size"]),
+        max_seq_len=int(data["max_position_embeddings"]),
+        n_layer=int(data["num_hidden_layers"]),
+        n_embd=int(data["hidden_size"]),
+        n_head=int(data["num_attention_heads"]),
+        n_kv_head=int(data["num_key_value_heads"]),
+        head_dim=int(data["head_dim"]),
+        intermediate_size=int(data["intermediate_size"]),
+        rope_theta=float(data["rope_theta"]),
+        norm_eps=float(data["rms_norm_eps"]),
+        logits_chunk_tokens=8192,
+        gradient_checkpointing=False,
+        use_transformer_engine=False,
+        attn_input_format="bshd",
+    )
+
+
+def choose_device(requested: str) -> torch.device:
+    if requested != "auto":
+        return torch.device(requested)
+    if torch.cuda.is_available():
+        return torch.device("cuda")
+    if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def sample_next(logits: torch.Tensor, temperature: float, top_p: float) -> torch.Tensor:
+    if temperature <= 0:
+        return logits.argmax(dim=-1, keepdim=True)
+    logits = logits / temperature
+    if 0 < top_p < 1:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+        sorted_probs = torch.softmax(sorted_logits, dim=-1)
+        cumulative = sorted_probs.cumsum(dim=-1)
+        remove = cumulative - sorted_probs > top_p
+        sorted_logits = sorted_logits.masked_fill(remove, float("-inf"))
+        logits = torch.full_like(logits, float("-inf"))
+        logits.scatter_(dim=-1, index=sorted_indices, src=sorted_logits)
+    return torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--weights", type=Path, default=Path("model.safetensors"))
+    parser.add_argument("--config", type=Path, default=Path("config.json"))
+    parser.add_argument("--tokenizer", type=Path, default=Path("tokenizer.model"))
+    parser.add_argument("--prompt", required=True)
+    parser.add_argument("--chat", action="store_true")
+    parser.add_argument(
+        "--system",
+        default="You are a helpful assistant. Answer clearly and say when information is missing.",
+    )
+    parser.add_argument("--max-new-tokens", type=int, default=128)
+    parser.add_argument("--temperature", type=float, default=0.0)
+    parser.add_argument("--top-p", type=float, default=0.95)
+    parser.add_argument("--device", default="auto")
+    args = parser.parse_args()
+
+    import sentencepiece as spm
+
+    device = choose_device(args.device)
+    dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
+    config = load_config(args.config)
+    model = VortexForCausalLM(config).to(device=device, dtype=dtype)
+    state = load_file(str(args.weights), device="cpu")
+    model.load_state_dict(state, strict=True)
+    model.eval()
+
+    prompt = args.prompt
+    if args.chat:
+        prompt = f"[SYSTEM]\n{args.system}\n</s>\n[USER]\n{prompt}\n</s>\n[ASSISTANT]\n"
+    tokenizer = spm.SentencePieceProcessor(model_file=str(args.tokenizer))
+    ids = tokenizer.encode(prompt, out_type=int)
+    if not ids:
+        raise ValueError("prompt encoded to zero tokens")
+    if len(ids) >= config.max_seq_len:
+        raise ValueError("prompt reaches the configured context limit")
+
+    input_ids = torch.tensor([ids], dtype=torch.long, device=device)
+    generated: list[int] = []
+    with torch.inference_mode():
+        for _ in range(max(0, args.max_new_tokens)):
+            logits, _ = model(input_ids)
+            next_id = sample_next(logits[:, -1, :].float(), args.temperature, args.top_p)
+            token_id = int(next_id.item())
+            if token_id == tokenizer.eos_id():
+                break
+            generated.append(token_id)
+            input_ids = torch.cat((input_ids, next_id), dim=1)
+            if input_ids.shape[1] >= config.max_seq_len:
+                break
+    print(tokenizer.decode(generated, out_type=str), end="")
+
+
+if __name__ == "__main__":
+    main()

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:330ae517cc9bd14d055f1bbde3b330b6d3206a8740bb4ef032a771af708a96b7
+size 349898992

special_tokens_map.json

@@ -0,0 +1,4 @@
+{
+  "unk_token": "<unk>",
+  "eos_token": "</s>"
+}

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8ebb450f2786d3153d64f874d97487ba5cb8e3450b799bbf6aeddd7d92f7ff73
+size 123865

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "model_type": "vortex",
+  "tokenizer_class": "SentencePieceProcessor",
+  "vocab_size": 8192,
+  "unk_token": "<unk>",
+  "eos_token": "</s>",
+  "add_bos_token": false,
+  "add_eos_token": false
+}

vortex_model.py

@@ -0,0 +1,404 @@
+"""Capability-first Vortex 175M decoder.
+
+This file intentionally uses only standard PyTorch CUDA primitives.  The
+model is deep-and-thin, uses GQA and SwiGLU, and ties the input/output
+embedding.  Those choices are much easier to train and export than a custom
+SSM kernel while retaining the main sub-billion-parameter wins reported by
+MobileLLM-style studies.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import asdict, dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+
+@dataclass
+class VortexConfig:
+    vocab_size: int = 8_192
+    max_seq_len: int = 4_096
+    n_layer: int = 12
+    n_embd: int = 1024
+    n_head: int = 16
+    n_kv_head: int = 4
+    head_dim: int = 64
+    intermediate_size: int = 3_664
+    rope_theta: float = 100_000.0
+    norm_eps: float = 1e-5
+    logits_chunk_tokens: int = 16_384
+    gradient_checkpointing: bool = False
+    use_transformer_engine: bool = False
+    attn_input_format: str = "bshd"
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float) -> None:
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return F.rms_norm(x, (x.shape[-1],), self.weight, self.eps)
+
+
+def _linear(config: VortexConfig, in_features: int, out_features: int) -> nn.Module:
+    """Create a bias-free projection, optionally backed by Transformer Engine."""
+    if not config.use_transformer_engine:
+        return nn.Linear(in_features, out_features, bias=False)
+    try:
+        import transformer_engine.pytorch as te
+    except ImportError as exc:  # pragma: no cover - exercised only on TE runs
+        raise RuntimeError(
+            "use_transformer_engine=True requires transformer-engine[pytorch]"
+        ) from exc
+    return te.Linear(
+        in_features,
+        out_features,
+        bias=False,
+        params_dtype=torch.bfloat16,
+        device="cuda",
+    )
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, max_seq_len: int, theta: float) -> None:
+        super().__init__()
+        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+        positions = torch.arange(max_seq_len, dtype=torch.float32)
+        frequencies = torch.outer(positions, inv_freq)
+        # NeoX-style rotate-half layout: the first and second halves share
+        # the same frequencies, so rotate_half() remains allocation-free in
+        # the hot attention path apart from its concatenation.
+        angles = torch.cat((frequencies, frequencies), dim=-1)
+        self.register_buffer("cos_cached", angles.cos()[None, None], persistent=False)
+        self.register_buffer("sin_cached", angles.sin()[None, None], persistent=False)
+
+    @staticmethod
+    def rotate_half(x: torch.Tensor) -> torch.Tensor:
+        half = x.shape[-1] // 2
+        return torch.cat((-x[..., half:], x[..., :half]), dim=-1)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        seq_len = q.shape[-2]
+        if seq_len > self.cos_cached.shape[-2]:
+            raise ValueError(f"sequence length {seq_len} exceeds configured maximum")
+        cos = self.cos_cached[:, :, :seq_len].to(dtype=q.dtype)
+        sin = self.sin_cached[:, :, :seq_len].to(dtype=q.dtype)
+        return (
+            q * cos + self.rotate_half(q) * sin,
+            k * cos + self.rotate_half(k) * sin,
+        )
+
+
+class GQAAttention(nn.Module):
+    def __init__(self, config: VortexConfig) -> None:
+        super().__init__()
+        if config.n_head % config.n_kv_head:
+            raise ValueError("n_head must be divisible by n_kv_head")
+        if config.n_head * config.head_dim != config.n_embd:
+            raise ValueError("n_head * head_dim must equal n_embd")
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.head_dim = config.head_dim
+        kv_dim = config.n_kv_head * config.head_dim
+        self.q_proj = _linear(config, config.n_embd, config.n_embd)
+        self.k_proj = _linear(config, config.n_embd, kv_dim)
+        self.v_proj = _linear(config, config.n_embd, kv_dim)
+        self.o_proj = _linear(config, config.n_embd, config.n_embd)
+        # QK-Norm keeps attention logits well-conditioned at the deliberately
+        # high pretraining learning rate.  These are per-head, parameter-light
+        # norms, not full hidden-size projections.
+        self.q_norm = RMSNorm(config.head_dim, config.norm_eps)
+        self.k_norm = RMSNorm(config.head_dim, config.norm_eps)
+        self.rope = RotaryEmbedding(config.head_dim, config.max_seq_len, config.rope_theta)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        batch, seq_len, _ = x.shape
+        q = self.q_proj(x).view(batch, seq_len, self.n_head, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(batch, seq_len, self.n_kv_head, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(batch, seq_len, self.n_kv_head, self.head_dim).transpose(1, 2)
+        q, k = self.rope(self.q_norm(q), self.k_norm(k))
+        # PyTorch dispatches this to the fused flash/efficient causal kernel
+        # when the local CUDA build supports it.  enable_gqa avoids material-
+        # izing repeated K/V heads.
+        y = F.scaled_dot_product_attention(
+            q, k, v, is_causal=True, enable_gqa=True
+        )
+        y = y.transpose(1, 2).contiguous().view(batch, seq_len, -1)
+        return self.o_proj(y)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, config: VortexConfig) -> None:
+        super().__init__()
+        self.gate_proj = _linear(config, config.n_embd, config.intermediate_size)
+        self.up_proj = _linear(config, config.n_embd, config.intermediate_size)
+        self.down_proj = _linear(config, config.intermediate_size, config.n_embd)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class VortexBlock(nn.Module):
+    def __init__(self, config: VortexConfig, layer_number: int | None = None) -> None:
+        super().__init__()
+        self.use_transformer_engine = config.use_transformer_engine
+        self.attn_input_format = config.attn_input_format
+        if self.use_transformer_engine:
+            try:
+                import transformer_engine.pytorch as te
+            except ImportError as exc:  # pragma: no cover - TE-only path
+                raise RuntimeError(
+                    "use_transformer_engine=True requires transformer-engine[pytorch]"
+                ) from exc
+            # This fused layer has the same parameter shapes as the explicit
+            # reference block: fused GQA QKV, RMSNorm, QK-Norm, SwiGLU, and
+            # the causal attention kernel.  It is used only for the TE backend;
+            # the reference PyTorch path remains readable and exportable.
+            self.te_layer = te.TransformerLayer(
+                hidden_size=config.n_embd,
+                ffn_hidden_size=config.intermediate_size,
+                num_attention_heads=config.n_head,
+                num_gqa_groups=config.n_kv_head,
+                layernorm_epsilon=config.norm_eps,
+                hidden_dropout=0.0,
+                attention_dropout=0.0,
+                kv_channels=config.head_dim,
+                layer_number=layer_number,
+                bias=False,
+                activation="swiglu",
+                normalization="RMSNorm",
+                qk_norm_type="RMSNorm",
+                qk_norm_before_rope=True,
+                fuse_qkv_params=True,
+                self_attn_mask_type="causal",
+                attn_input_format=config.attn_input_format,
+                params_dtype=torch.bfloat16,
+                device="cuda",
+            )
+            return
+        self.norm1 = RMSNorm(config.n_embd, config.norm_eps)
+        self.attn = GQAAttention(config)
+        self.norm2 = RMSNorm(config.n_embd, config.norm_eps)
+        self.ffn = SwiGLU(config)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        rotary_pos_emb: torch.Tensor | None = None,
+        is_first_microbatch: bool | None = None,
+        inference_params=None,
+        attention_mask: torch.Tensor | None = None,
+        inference_decode_bshd: bool = False,
+    ) -> torch.Tensor:
+        if self.use_transformer_engine:
+            if inference_params is not None:
+                # A packed THD prompt handles variable-length prefill
+                # efficiently.  Subsequent one-token decode is cheaper and
+                # more numerically stable through the regular BSHD cache
+                # path; switch the TE attention format explicitly between the
+                # two phases.
+                effective_format = "bshd" if inference_decode_bshd else self.attn_input_format
+                self.te_layer.self_attention.qkv_format = effective_format
+            te_kwargs = {
+                "attention_mask": attention_mask,
+                "self_attn_mask_type": (
+                    "padding_causal" if inference_params is not None else None
+                ),
+                "rotary_pos_emb": rotary_pos_emb,
+                "is_first_microbatch": is_first_microbatch,
+                "inference_params": inference_params,
+            }
+            if inference_params is not None and not inference_decode_bshd and self.attn_input_format == "thd":
+                batch_size = len(inference_params.sequences)
+                cu_seqlens = inference_params.cu_seqlens_q[: batch_size + 1]
+                sequence_lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+                te_kwargs.update(
+                    {
+                        "cu_seqlens_q": cu_seqlens,
+                        "cu_seqlens_q_padded": cu_seqlens,
+                        "max_seqlen_q": int(sequence_lengths.max().item()),
+                        "max_seqlen_kv": int(sequence_lengths.max().item()),
+                    }
+                )
+            return self.te_layer(
+                x,
+                **te_kwargs,
+            )
+        x = x + self.attn(self.norm1(x))
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+
+class VortexForCausalLM(nn.Module):
+    def __init__(self, config: VortexConfig | None = None) -> None:
+        super().__init__()
+        self.config = config or VortexConfig()
+        self.embed_tokens = nn.Embedding(self.config.vocab_size, self.config.n_embd)
+        self.layers = nn.ModuleList(
+            VortexBlock(self.config, layer_number=index + 1)
+            for index in range(self.config.n_layer)
+        )
+        self.norm = RMSNorm(self.config.n_embd, self.config.norm_eps)
+        if self.config.use_transformer_engine:
+            import transformer_engine.pytorch as te
+
+            self.rotary = te.RotaryPositionEmbedding(
+                self.config.head_dim,
+                rotary_base=self.config.rope_theta,
+                interleaved=False,
+            )
+        self._initialize_weights()
+
+    def _initialize_weights(self) -> None:
+        # Scale residual outputs down with depth; this gives a forgiving high-
+        # LR start without adding trainable parameters.
+        output_std = 0.02 / math.sqrt(2.0 * self.config.n_layer)
+        nn.init.normal_(self.embed_tokens.weight, mean=0.0, std=0.02)
+        for block in self.layers:
+            if self.config.use_transformer_engine:
+                for name, parameter in block.named_parameters():
+                    if parameter.ndim == 1:
+                        nn.init.ones_(parameter)
+                    else:
+                        is_output = (
+                            name.endswith("self_attention.proj.weight")
+                            or name.endswith("layernorm_mlp.fc2_weight")
+                        )
+                        nn.init.normal_(
+                            parameter,
+                            mean=0.0,
+                            std=output_std if is_output else 0.02,
+                        )
+            else:
+                for child in block.modules():
+                    if isinstance(child, nn.Linear):
+                        is_output = child is block.attn.o_proj or child is block.ffn.down_proj
+                        nn.init.normal_(
+                            child.weight,
+                            mean=0.0,
+                            std=output_std if is_output else 0.02,
+                        )
+
+    def parameter_count(self) -> int:
+        return sum(parameter.numel() for parameter in self.parameters())
+
+    def parameter_breakdown(self) -> dict[str, int]:
+        c = self.config
+        embedding = c.vocab_size * c.n_embd
+        q = c.n_layer * c.n_embd * c.n_embd
+        k = c.n_layer * c.n_embd * (c.n_kv_head * c.head_dim)
+        v = k
+        o = q
+        # One learned head-dimension scale is shared across all query heads,
+        # and another across all KV heads, matching the module definitions
+        # above (not one scale vector per physical head).
+        qk_norm = c.n_layer * 2 * c.head_dim
+        ffn = c.n_layer * 3 * c.n_embd * c.intermediate_size
+        block_norm = c.n_layer * 2 * c.n_embd
+        final_norm = c.n_embd
+        return {
+            "input_embedding_and_tied_output": embedding,
+            "attention_q_projection": q,
+            "attention_k_projection": k,
+            "attention_v_projection": v,
+            "attention_o_projection": o,
+            "attention_qk_norm": qk_norm,
+            "ffn_swiglu": ffn,
+            "block_rmsnorm": block_norm,
+            "final_rmsnorm": final_norm,
+            "total": self.parameter_count(),
+        }
+
+    def _chunked_tied_loss(self, hidden: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+        chunk = self.config.logits_chunk_tokens
+        total = hidden.new_zeros((), dtype=torch.float32)
+        for start in range(0, hidden.shape[0], chunk):
+            end = min(hidden.shape[0], start + chunk)
+            logits = F.linear(hidden[start:end], self.embed_tokens.weight)
+            total = total + F.cross_entropy(logits, targets[start:end]).float() * (end - start)
+        return total / hidden.shape[0]
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        labels: torch.Tensor | None = None,
+        is_first_microbatch: bool | None = None,
+        inference_params=None,
+        inference_attention_mask=None,
+        inference_decode_bshd: bool = False,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+        x = self.embed_tokens(input_ids)
+        rotary_pos_emb = None
+        attention_mask = None
+        if self.config.use_transformer_engine:
+            if inference_params is not None:
+                # TE's cached-attention path applies the correct absolute
+                # offset from inference_params.  Supplying the full table is
+                # necessary when the current query is only one token long but
+                # starts after a long cached prefix.
+                rotary_pos_emb = self.rotary(self.config.max_seq_len)
+                if self.config.attn_input_format == "thd" and not inference_decode_bshd:
+                    attention_mask = None
+                elif inference_attention_mask is None:
+                    query_padding_mask = torch.zeros(
+                        (x.shape[0], 1, 1, x.shape[1]),
+                        dtype=torch.bool,
+                        device=x.device,
+                    )
+                    key_padding_mask = torch.ones(
+                        (x.shape[0], 1, 1, inference_params.max_sequence_length),
+                        dtype=torch.bool,
+                        device=x.device,
+                    )
+                    for batch_index, sequence_length in enumerate(
+                        inference_params.sequences.values()
+                    ):
+                        key_padding_mask[batch_index, :, :, :sequence_length] = False
+                    # TE switches the cached self-attention implementation to
+                    # its cross-attention backend internally; that backend
+                    # expects the query and key padding masks as a pair.
+                    attention_mask = (query_padding_mask, key_padding_mask)
+                else:
+                    attention_mask = inference_attention_mask
+            else:
+                rotary_pos_emb = self.rotary(x.shape[1])
+        for block in self.layers:
+            if self.training and self.config.gradient_checkpointing:
+                x = checkpoint(
+                    block,
+                    x,
+                    rotary_pos_emb,
+                    is_first_microbatch,
+                    use_reentrant=False,
+                )
+            else:
+                x = block(
+                    x,
+                    rotary_pos_emb,
+                    is_first_microbatch,
+                    inference_params,
+                    attention_mask,
+                    inference_decode_bshd,
+                )
+        x = self.norm(x)
+        if labels is None:
+            return F.linear(x, self.embed_tokens.weight), None
+        hidden = x[:, :-1].reshape(-1, x.shape[-1])
+        targets = labels[:, 1:].reshape(-1)
+        return None, self._chunked_tied_loss(hidden, targets)
+
+
+if __name__ == "__main__":
+    config = VortexConfig()
+    model = VortexForCausalLM(config)
+    print(config.to_dict())
+    print(model.parameter_breakdown())