GTLM-1-2B-A350M-fp16

GTLM-1 is an experimental Sparse Mixture-of-Experts language model trained from scratch under extreme compute and budget constraints. The goal of this release is to document a viable MoE training recipe in the sub-2B parameter regime, focusing on data efficiency and engineering trade-offs rather than state-of-the-art performance.

🧠 Architecture & Engineering

GTLM-1 departs from standard implementations to prioritize efficiency in low-resource regimes:

• Architecture: Sparse Mixture of Experts (MoE) with 16 Experts per layer (Top-2 Routing).
• Vectorized Dispatch: Implemented a custom vectorized router in PyTorch for efficient token distribution.
• Optimizations:
  • Liger Kernels: Used for RMSNorm to reduce memory footprint.
  • Manual Implementations: SwiGLU and RoPE were implemented manually in PyTorch.
  • Engineering Note: We experimented with custom Fused Triton kernels for SwiGLU, Token Dispatch, and RoPE. However, benchmarks showed that in this specific model size regime, they offered no significant throughput gain over optimized manual PyTorch implementations, so we prioritized code simplicity and stability.

💻 Training Infrastructure & Cost

• Hardware: Trained on a single NVIDIA A100 (40GB/80GB partition).
• Duration: ~140 Hours.
• Throughput: Average of 50,000 tokens/second (TPS).
• Total Cost: ~$100 USD.
• Frameworks: PyTorch + DeepSpeed.

📚 Dataset (The Blend)

The model was trained on a highly curated mix of 15 Billion tokens. The data pipeline focused on density and reasoning quality rather than raw volume.

• Primary Language: English (>80%), with a specific subset of high-quality Portuguese (BR).
• Sources:
  • Reasoning/Math: Nemotron Math, FineWeb-Edu, FineMath.
  • General Knowledge: Dolma, RedPajama, SlimPajama, FineWeb.
  • Custom: A proprietary dataset of scraped and filtered Brazilian Web documents to ensure cultural alignment and syntax quality.

📊 Evaluation Benchmark

Comprehensive comparison of GTLM-1.2B-A350M (MoE) against Small Language Models (SLMs). Evaluations performed using lm-evaluation-harness (Zero-Shot).

🏆 Key Takeaways

• Science Specialist: GTLM achieves 87.6% on SciQ, outperforming the larger and heavily trained SmolLM-360M Instruct (85.9%).
• Data Efficiency: Beat Pythia-410M and GPT-2 Medium across the board with significantly fewer training tokens (15B vs 300B+).
• Architectural Win: Demonstrates that a sparse MoE architecture (350M active params) can compete with dense models trained on 40x more data.