Bexamask-v2-0.8B-Base

🚀 BexaMask-v2 (≈800M Parameters)

BexaMask-v2 is a pretrained base (foundation) decoder-only Transformer model trained on large-scale permissively licensed and uncopyrighted text data using the MaxText framework on TPU v4-16.

⚠️ This is a base model — it is not instruction-tuned and may not follow prompts like ChatGPT without further fine-tuning.

---

🧠 Model Overview

• Type: Pretrained Base Model (Foundation Model)
• Architecture: Decoder-only Transformer
• Parameters: ~800M
• Layers: 16
• Embedding Dimension: 2048
• MLP Dimension: 5120
• Attention Heads:
  • Query Heads: 16
  • KV Heads: 4 (Grouped Query Attention)
• Head Dimension: 128
• Activation: SiLU + Linear
• Max Context Length: 4096 tokens
• Vocabulary Size: 32,000 (SentencePiece)

---

⚙️ Training Details

• Framework: MaxText
• Hardware: TPU v4-16 (8 chips, 256GB HBM)

📦 Dataset

• Subset of The Pile (uncopyrighted / permissive sources only)
• Filtered to remove restricted or copyrighted data

🔧 Training Config

• Steps: 100,000
• Epochs: 2
• Batch Size: 16 per device
• Learning Rate: 3e-4
• Warmup Steps: 2,000
• Scheduler: Cosine decay

---

⚡ Optimization Techniques

• Flash Attention
• Full Rematerialization (Remat)
• Asynchronous Checkpointing
• Distributed GCS checkpointing
• IOTA embeddings

---

🧪 Inference

Run inference using MaxText:

python3 -m MaxText.decode \
  maxtext/configs/pretrain.yml \
  run_name=inference \
  load_parameters_path=/home/pynatic079/bexamask_v2_inference_local/items \
  tokenizer_path=/path/to/llama/tokenizer.model \
  max_target_length=512 \
  'prompt="<Your prompt>"' \
  decode_sampling_strategy="topk" \
  decode_sampling_top_k=4 \
  decode_sampling_temperature=1.9 \
  attention=dot_product