Upload checkpoints/checkpoint-09B

checkpoints/checkpoint-09B/README.md

@@ -0,0 +1,148 @@
+---
+license: apache-2.0
+library_name: pytorch
+tags:
+- linear-attention
+- recurrent
+- gated-deltanet
+- gdn2
+- kaczmarz
+datasets:
+- HuggingFaceFW/fineweb-edu
+---
+
+# Gated_Linear_Attention2
+
+This repository stores milestone checkpoints for `gdn2_kla_1.3B` from the
+`gdn2_kla_1.3B_fineweb_edu_10bt` run.
+
+## What This Model Is
+
+`gdn2_kla_1.3B` is a recurrent-only linear attention experiment. It starts from
+Gated DeltaNet-2 and folds a Kaczmarz-style key-norm-normalized update step into
+the separate erase and write gates:
+
+```math
+\lambda_t = \frac{\eta_t}{\|k_t\|_2^2 + \epsilon}
+```
+
+```math
+S_t =
+\left(I - k_t(\lambda_t b_t \odot k_t)^\top\right)D_tS_{t-1}
++
+k_t(\lambda_t w_t \odot v_t)^\top
+```
+
+It is not a standard Transformers checkpoint and does not use softmax attention
+or SWA layers.
+
+## Code
+
+- GitHub: https://github.com/gyunggyung/Gated_Linear_Attention2
+
+## License
+
+The model weights in this Hugging Face repository are released under Apache-2.0.
+
+The standalone inference runtime linked above is also Apache-2.0. It does not
+import `lit_gpt`, `fla`, or the NVIDIA GatedDeltaNet-2 Triton kernels. The
+training code used during experimentation may contain NVIDIA GatedDeltaNet-2
+derived components under `Nvidia Source Code License-NC`, but this Hugging Face
+model repository is intended to be used with the standalone Apache-2.0 runtime.
+
+## Training Setup
+
+- Base architecture: recurrent-only GDN-2, 1.3B scale
+- Candidate: Kaczmarz-normalized GDN-2 gates
+- Training data source: FineWeb-Edu `sample/100BT` local parquet
+- Token budget for this run: 10B
+- Current milestone: 9,000,000,000 tokens
+- Sequence length: 4096 tokens
+- Global batch tokens: 1,048,576
+- Tokenizer: `TinyLlama/TinyLlama_v1.1`
+- Data shuffle seed: `3407`
+- Data shuffle buffer: `100000`
+
+## Checkpoint Format
+
+Each `checkpoints/checkpoint-XXB/` folder contains:
+
+- `model-ckpt.pth`: PyTorch model-only checkpoint
+- `training_metadata.json`: run metadata and model config
+- `README.md`: this model card snapshot
+
+This is not loadable with `transformers.AutoModelForCausalLM.from_pretrained`.
+
+## How To Use
+
+This is a causal language model: given a text prefix, it predicts the next token
+and can continue the text autoregressively. It was pretrained on FineWeb-Edu and
+is not instruction-tuned, RLHF-tuned, or chat-aligned.
+
+The checkpoint is a PyTorch `.pth` checkpoint, not a
+`transformers.AutoModelForCausalLM` checkpoint. Use the standalone runtime below
+to load it.
+
+Install and clone:
+
+```bash
+git clone https://github.com/gyunggyung/Gated_Linear_Attention2
+cd Gated_Linear_Attention2
+pip install -e .
+```
+
+Minimal text-generation example:
+
+```python
+import torch
+
+from gated_linear_attention2 import GatedLinearAttention2ForCausalLM, load_tokenizer
+from gated_linear_attention2.generation import generate
+
+repo_id = "gyung/Gated_Linear_Attention2"
+checkpoint_file = "checkpoints/checkpoint-01B/model-ckpt.pth"
+
+if not torch.cuda.is_available():
+    raise RuntimeError("CUDA is recommended for this 1.3B checkpoint; CPU will be very slow.")
+
+device = "cuda"
+dtype = torch.bfloat16
+
+model = GatedLinearAttention2ForCausalLM.from_hf(
+    repo_id=repo_id,
+    checkpoint=checkpoint_file,
+    device=device,
+    dtype=dtype,
+)
+tokenizer = load_tokenizer(repo_id, subfolder="tokenizer")
+
+prompt = "Artificial intelligence can help education by"
+print(generate(model, tokenizer, prompt, max_new_tokens=80, temperature=0.8, top_k=50))
+```
+
+For next-token scoring instead of generation, run one forward pass and inspect
+the final-position logits:
+
+```python
+prompt = "The capital of France is"
+input_ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+with torch.no_grad():
+    logits = model(input_ids)[:, -1, :]
+next_token_id = int(torch.argmax(logits, dim=-1)[0])
+print(tokenizer.decode([next_token_id]))
+```
+
+The standalone runtime uses a recurrent state cache during generation, so decode
+memory does not grow with generated token length like a Transformer KV cache.
+
+## Evaluation Plan
+
+Compare against the plain `gdn2_1.3B` baseline on the GDN-2 paper tasks:
+
+- WikiText and LAMBADA perplexity
+- LAMBADA and commonsense zero-shot accuracy
+- RULER S-NIAH and MK-NIAH
+- Real-world retrieval tasks: SWDE, SQuAD, FDA, TriviaQA, NQ, DROP
+
+The 10B run is an ablation, not a claim that it replaces the published 100B
+GDN-2 model.

checkpoints/checkpoint-09B/model-ckpt.pth

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

checkpoints/checkpoint-09B/training_metadata.json

@@ -0,0 +1,57 @@
+{
+  "checkpoint_format": "LitGPT/Fabric .pth, not a Transformers AutoModel checkpoint",
+  "data_shuffle_buffer": 100000,
+  "data_shuffle_seed": 3407,
+  "data_source": "HuggingFaceFW/fineweb-edu sample/100BT local parquet",
+  "experiment_name": "gdn2_kla_1.3B_fineweb_edu_10bt",
+  "global_batch_tokens": 1048576,
+  "learning_rate": 0.0004,
+  "milestone_index": 9,
+  "model_config": {
+    "_mlp_class": "LLaMAMLP",
+    "_norm_class": "FusedRMSNorm",
+    "activation_checkpointing": true,
+    "bias": false,
+    "block_size": 4096,
+    "condense_ratio": 1,
+    "gdn2_allow_neg_eigval": false,
+    "gdn2_conv_bias": false,
+    "gdn2_conv_size": 4,
+    "gdn2_expand_v": 1.0,
+    "gdn2_head_dim": 128,
+    "gdn2_kaczmarz_eps": 1e-06,
+    "gdn2_mode": "chunk",
+    "gdn2_num_heads": 16,
+    "gdn2_num_v_heads": null,
+    "gdn2_per_layer": 1,
+    "gdn2_use_kaczmarz_step": true,
+    "gdn2_use_qk_l2norm_in_kernel": false,
+    "gdn2_use_short_conv": true,
+    "intermediate_size": 6208,
+    "local_window": 2048,
+    "mamba_init": true,
+    "mlp": true,
+    "n_embd": 2304,
+    "n_head": 18,
+    "n_layer": 18,
+    "n_query_groups": 18,
+    "name": "gdn2_kla_1.3B",
+    "nope": true,
+    "norm_eps": 1e-05,
+    "org": "NVIDIA",
+    "padded_vocab_size": 32000,
+    "padding_multiple": 64,
+    "parallel_residual": false,
+    "rotary_percentage": 1.0,
+    "shared_attention_norm": false,
+    "vocab_size": 32000
+  },
+  "model_name": "gdn2_kla_1.3B",
+  "seed": 3407,
+  "sequence_length": 4096,
+  "tokenizer_name": "TinyLlama/TinyLlama_v1.1",
+  "tokenizer_path": "",
+  "train_config": "tsz128x4k_10B",
+  "trained_tokens": 9000000000,
+  "weight_decay": 0.1
+}