DBRX (language model) (CS)

DBRX — je otevřený velký jazykový model (LLM) vyvinutý výzkumným týmem Mosaic AI v rámci společnosti Databricks. Model byl oficiálně vydán 27. března 2024 a je prezentován jako vysoce výkonné řešení pro podnikové využití.

DBRX je postaven na jemnozrnné architektuře směsi expertů (MoE) a kombinuje vysoký výkon s efektivitou trénování a inference. V době vydání dosáhl DBRX nejlepších výsledků mezi všemi otevřenými modely na klíčových benchmarcích, překonal modely jako LLaMA 2, Mixtral a Grok-1 a prokázal konkurenceschopnost se zavřenými modely úrovně GPT-3.5 Turbo.

Historie vývoje

Vznik DBRX byl pokračováním strategie Databricks v oblasti rozvoje otevřených generativních modelů. V červnu 2023 společnost Databricks akvírovala startup MosaicML, specializující se na trénování velkých modelů, a na jeho základě bylo vytvořeno oddělení Mosaic AI.

Tým Mosaic AI, vedený předním architektem neuronových sítí Jonathanem Franklem, zahájil vývoj nového velkého LLM s cílem dosáhnout kvality srovnatelné s nejlepšími proprietárními systémy, avšak v otevřeném formátu. Projekt dostal název DBRX. Vývoj a předtrénování modelu trvaly přibližně 2,5 měsíce a podle odhadů stály přibližně 10 milionů dolarů.

Architektura

DBRX je transformer model typu pouze dekodér (decoder-only) a implementuje jemnozrnnou (fine-grained) architekturu směsi expertů (MoE).

Klíčové vlastnosti architektury:

• Celkový počet parametrů: 132 miliard.
• Experti: Model se skládá z 16 malých specializovaných podmodelů („expertů").
• Mechanismus aktivace: Pro každý vstupní token jsou aktivováni pouze 4 ze 16 expertů. To znamená, že při inferenci je aktivních pouze 36 miliard parametrů, což zajišťuje vysokou rychlost a efektivitu. Toto schéma poskytuje 65krát více možných kombinací expertů ve srovnání s modelem Mixtral (8 expertů s aktivací 2).
• Komponenty: Jsou použita moderní architektonická řešení, jako jsou rotační poziční embedding (RoPE), gated linear units (GLU) a grouped query attention (GQA).
• Délka kontextu: 32 768 tokenů.

Tato architektura umožňuje modelu kombinovat výhody obrovského počtu parametrů (pro ukládání znalostí) s efektivitou menších modelů (pro rychlost inference).

Trénování

Předtrénování DBRX bylo provedeno na pečlivě kurátorovaném datasetu o objemu 12 bilionů tokenů, složeném z textů a kódu. Kvalita dat byla klíčovou prioritou: vývojáři využili cloudovou platformu Databricks (Apache Spark, Databricks Notebooks, Unity Catalog) pro čištění, přípravu a audit dat.

Při trénování byla použita metoda kurikulárního učení (curriculum learning), při níž se v různých fázích měnil poměr typů dat. Například závěrečná část trénování byla věnována postupnému zavádění složitých úloh, což podle vývojářů přineslo znatelné zlepšení kvality. Trénování probíhalo na clusteru 3072 Nvidia H100 GPU.

Po předtrénování prošel základní model dodatečným doladěním (instruction tuning) pro vytvoření interaktivní verze DBRX Instruct, optimalizované pro plnění pokynů uživatele.

Výkon

V době vydání stanovil DBRX nový standard kvality pro otevřené LLM na širokém spektru benchmarků.

Srovnání s otevřenými modely

Výsledky DBRX Instruct na klíčových benchmarcích

Benchmark	Úloha	DBRX Instruct	Nejlepší následující (Mixtral/Grok-1)
Hugging Face Open LLM Leaderboard (AVG)	Obecné znalosti	74,5%	72,7% (Mixtral Instruct)
HumanEval	Programování	70,1%	63,2% (Grok-1)
GSM8K	Matematické uvažování	66,9%	62,9% (Grok-1)
MMLU	Obecné znalosti	73,7%	71,5% (Mixtral Instruct)

DBRX obsadil první místo jak v celkovém žebříčku Hugging Face Open LLM Leaderboard, tak v komplexním testu Databricks LLM Gauntlet, přičemž prokázal výrazný náskok před předchůdci.

Srovnání se zavřenými modely

DBRX Instruct překonává GPT-3.5 Turbo v řadě klíčových ukazatelů, včetně MMLU (73,7 % oproti 70,0 %) a HumanEval (70,1 % oproti 48,1 %). Z hlediska kvality odpovědí na některých benchmarcích (například MTBench) se model přibližuje úrovni Gemini 1.0 Pro a raných verzí GPT-4.

Efektivita trénování a inference

• Efektivita trénování: Použití architektury MoE umožnilo snížit náklady v FLOPS 2–4krát ve srovnání s hustými modely srovnatelné kvality.
• Efektivita inference: Díky aktivaci pouze 36 miliard parametrů poskytuje DBRX 2–3krát vyšší propustnost (rychlost inference) ve srovnání s hustými modely ekvivalentní velikosti (například LLaMA2-70B).

Licencování a dostupnost

DBRX je distribuován pod speciálně vytvořenou licencí Databricks Open Model License. Ta umožňuje volné použití a modifikaci, včetně komerčního využití, ale obsahuje určitá omezení. Zejména, podobně jako licence LLaMA 2, vyžaduje získání samostatného povolení od Databricks, pokud budou služby založené na DBRX využívány publikem přesahujícím 700 milionů aktivních uživatelů měsíčně.

Předtrénované váhy modelu (základní verze i verze Instruct) jsou dostupné ke stažení prostřednictvím repozitáře na Hugging Face.

Literatura

• Mosaic Research Team. (2024). Introducing DBRX: A New State‑of‑the‑Art Open LLM. Databricks Blog.
• Databricks. (2024). Databricks Open Model License (DBRX). Online specification.
• Fedus, W.; et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv:2101.03961.
• Lepikhin, D.; et al. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. arXiv:2006.16668.
• Ainslie, J.; et al. (2023). Grouped‑Query Attention: Efficient Training of Generalized Multi‑Query Transformers. arXiv:2305.13245.
• Su, J.; et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864.
• Dao, T. (2023). FlashAttention‑2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Cai, W.; et al. (2024). A Survey on Mixture of Experts in Large Language Models. arXiv:2407.06204.
• Chen, Y.; et al. (2024). Scaling Laws for Fine‑Grained Mixture of Experts. arXiv:2402.07871.
• Kundu, A.; et al. (2024). Strategic Data Ordering: Enhancing Large Language Model Training via Curriculum Learning. arXiv:2405.07490.