JetMoE: Reaching Llama2 Performance with 0.1M Dollars

• JetMoE-8B extends sparse activation to both attention and feed-forward layers, unlike prior approaches that sparsify only feed-forward layers.This design aims to reduce computational costs while maintaining performance.Source: §1 · p. 2
• JetMoE-8B uses a $100k budget, yet outperforms Llama2-7B, while JetMoE-8B-Chat outperforms Llama2-13B-Chat.The model uses 1.25T tokens and 30,000 H100 GPU hours.Source: §1 · p. 2
• The model has 8B parameters but activates only 2B per token, reducing inference computation by about 70% versus Llama2-7B.Source: §1 · p. 2
• JetMoE-8B is trained with public datasets and open-source code, making it accessible to many academic research settings.The report also highlights a comprehensive open-source data mixture.Source: §1 · p. 2
• The report details its training parameters and data mixture to encourage collaboration on accessible and efficient open foundation models.Source: §1 · p. 2

• Sparse routing selects the top k experts, and inactive experts need not be evaluated, reducing training and inference computation.The router assigns probabilities over experts and uses top-k selection; JetMoE therefore computes only selected expert outputs.Source: §2.1 · p. 3
• JetMoE replaces both self-attention and feed-forward layers with SMoE layers, unlike many open-source MoE models that replace only feed-forward layers.Source: §2.1 · p. 3
• JetMoE attention experts share key and value projections while retaining expert-specific query and output projections.The shared key and value matrices improve training and inference efficiency.Source: §2.3 · p. 3
• Mixture-of-Attention lets JetMoE add attention experts without increasing attention computation, preventing attention from becoming a bottleneck as MLP capacity grows.Source: §2.3 · p. 3
• Load-balancing and router z-losses regulate expert usage and improve training stability.The load-balancing loss discourages repeatedly selecting the same module, while the router z-loss is included to improve stability.Source: §2.4 · p. 4

3.1 Real-world Datasets

• 600 billion tokens from RefinedWeb provide a high-quality web-data component extracted from CommonCrawl using the MacroData Refinement pipeline.RefinedWeb contains 5 trillion tokens overall, while JetMoE uses its publicly available 600-billion-token extract.Source: §3.1 · p. 4
• The real-world mixture includes diverse English corpora such as StarCoder, Dolma, The Pile, mathematical datasets, and curated coding or instruction collections.These sources cover code, web text, academic and scientific material, mathematics, and natural-language-style prompts.Source: §3.1 · p. 4, §3.1.1 · pp. 4–5

3.2 Synthetic Datasets

• OpenHermes 2.5 supplies 1 million primarily synthetic instruction and chat samples compiled from numerous open-source and custom datasets.Its sources include coding, mathematics, medical tasks, instruction-following, and conversational datasets.Source: §3.2 · p. 5
• UltraTextbooks contributes synthetic and human-written textbooks, while UltraChat 200k provides a filtered subset of 1.4 million ChatGPT-generated dialogues.UltraChat filtering includes text truecasing and removal of dialogues with inappropriate assistant disclaimers.Source: §3.2 · p. 5
• Additional specialized datasets target mathematics and programming, including TemplateGSM, Magicoder, Evol-Code Alpaca, and Code-290k-ShareGPT.These collections provide grade-school math problems, automatically generated coding tasks, code instructions, and approximately 290,000 code conversations.Source: §3.2.1 · p. 5

4.1 Training Framework

• The training framework uses Megatron with Megablock MoE support and pipeline parallelism instead of expert parallelism for model parallel training.The framework was modified to support MoA and z-loss.Source: §4.1 · p. 6

4.2 Hyper-parameters

• JetMoE-8B replaces transformer self-attention and MLP layers with MoA and MoE, using 8 experts and top-k 2 in every layer.This configuration gives approximately twice the computation of a 1B model, with load-balancing and z-loss weights of 0.01 and 0.001.Source: §4.2 · p. 6
• The learning-rate schedule uses warmup, stable, and annealing stages, with 10B warmup tokens and 250B decay tokens.Initial and final learning rates are 10% of the maximum, with weight decay 0.1 and gradient clipping 1.0.Source: §4.2 · p. 6

4.3 Training Data Mixture

• 1.25T primarily English tokens cover web documents, mathematics, and code, with higher-quality data upweighted during learning-rate decay.Training is divided into warmup/stable and decay phases.Source: §4.3 · p. 6
• Phase 1 uses RefinedWeb, Starcoder, The Pile, peS2o from Dolma, and OpenWebMath, while Phase 2 adds additional high-quality data during decay.The detailed Phase 2 mixture is presented in Table 2.Source: §4.3 · pp. 6–7, Table 2 · p. 7
• The authors caution that the data mixture may not be ideal because of limited computing resources, although it provides a starting point for future optimization.This explicitly bounds the mixture’s intended status rather than presenting it as final.Source: §4.3 · p. 7

5.1 dSFT

• dSFT trains a student language model to respond to user prompts using data generated by a teacher model such as GPT-4 or Claude.The process includes data distillation from seed prompts and likelihood-based instruction tuning.Source: §5.1 · pp. 7–8

5.2 dDPO

• dDPO refines the dSFT model by incorporating preferences from an aligned teacher model through a reward objective.It uses KL-constrained optimization to maximize preferred outcomes while limiting divergence from a baseline policy.Source: §5.2 · p. 8
• Offline DPO offers stable learning and simpler tuning than RLHF, but static preference data can create distribution shifts as the policy evolves.Online and iterative DPO variants address this issue with greater computational complexity.Source: §5.2 · p. 8

5.3 Alignment Settings

• The alignment pipeline fine-tunes JetMoE-8B with dSFT on UltraChat 200k, Airoboros-3.2, Code-Feedback, Orca-math-word-problems-200k, SystemChat, and Capybara.The dSFT configuration uses a 2e-5 learning rate, batch size 128, and 3 epochs.Source: §5.3 · pp. 8–9
• A subsequent dDPO stage uses UltraFeedback binary preferences and produces JetMoE-8B-Chat after 60 H100 GPU hours.The dDPO configuration uses AdamW, learning rate 5e-7, batch size 128, and 1 epoch.Source: §5.3 · p. 9

• The evaluation covers OpenLLM tasks, physical and social reasoning, question answering, mathematics, commonsense reasoning, language modeling, reading comprehension, and code generation.Source: §6 · p. 9
• JetMoE-8B outperforms Gemma, LLaMA2, and DeepseekMoE on OpenLLM tasks except ARC-challenge and WinoGrande, while achieving the highest MBPP scores.Source: §6 · p. 9, Table 3 · p. 9
• JetMoE-8B-Chat achieves a higher MT-Bench score than Llama-2-13b-Chat after alignment.Source: §6 · p. 10, Table 4 · p. 10
• JetMoE-8B-Chat is relatively weak in coding and extraction compared with GPT-3.5-turbo, despite strong performance across other MT-Bench dimensions.Source: §6 · p. 10, Figure 3 · p. 10

• The $100k budget prevented architectural ablation studies, leaving the contributions of individual components to the final results for future investigation.Source: §7 · p. 10
• Hyperparameters and data mixtures were handpicked using empirical results from previous works, including Shen et al. (2023), Zoph et al. (2022), and Hu et al. (2024).Source: §7 · p. 10

• JetMoE-8B achieves state-of-the-art performance among open-source models while maintaining high efficiency.Source: §8 · p. 10
• Sparse activation in both attention and feed-forward layers reduces computational costs while maintaining strong performance across diverse tasks.Source: §8 · p. 11
• JetMoE-8B outperforms larger, more resource-intensive models on the OpenLLM Leaderboard, while JetMoE-8B-Chat remains competitive with other open-source chatbots.Source: §8 · p. 11
• Detailed training parameters and data mixtures are provided to encourage reproducibility and enable researchers to build upon the work.Source: §8 · p. 11