mixSGA

KO | EN

gitlite — search

#javascript #python #hacktoberfest #typescript #llm #rust #linux #macos #java #windows #awesome #kubernetes #ios #openai

(★ 11)

Official implementation of emnlp'25 paper Mixture of Weight-shared Heterogeneous Group Attention Experts for Dynamic Token-wise KV Optimization

mixSGA 최신버젼 다운로드

최종 버전 다운로드 (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

<!-- # Official PyTorch Implementation of **mixSGA**

## Overview
This repository provides the official implementation of **mixSGA**, a framework designed for dynamic token-wise key-value (KV) optimization in causal language modeling (CLM). Built upon the work from the [DistillM repository](https://github.com/jongwooko/distillm), this method introduces a mixture-of-experts (MoE) approach to efficiently allocate resources, dynamically adjusting the KV cache to token importance. -->

# mixSGA: Mixture of Weight-shared Heterogeneous Group Attention Experts for Dynamic Token-wise KV Optimization
[![Python](https://img.shields.io/badge/Python-3.10-green)](https://www.python.org/downloads/release/python-310/)
[![PyTorch](https://img.shields.io/badge/PyTorch-2.2.1-green)](https://pytorch.org/)
[![MIT](https://img.shields.io/badge/License-MIT-green)](https://github.com/jongwooko/mixSGA/blob/main/LICENSE)

**mixSGA** is a novel framework for dynamic token-wise key-value (KV) optimization in causal language modeling (CLM). This method, proposed in our paper at **EMNLP 2025**, introduces a mixture-of-experts (MoE) approach that dynamically adjusts the KV cache allocation based on the importance of individual tokens. By leveraging learned token-wise routing and weight-sharing techniques, **mixSGA** optimizes memory and computation resources, improving efficiency in large language models (LLMs).

<p align="center">
    <a href="https://arxiv.org/abs/2506.13541"> Paper (arXiv version) </a> •  <a href="https://github.com/GuanghuiSong/mixSGA"> GitHub Repository </a>
    <br/><br/>
    <img src="assets/overview.png" width="800px">
</p>

## Environment Setup
To get started with the **mixSGA** implementation, run the following command to install all dependencies:
```bash
bash install.sh

## Data Processing

Before training the model, you need to process the data. The script below tokenizes the data and stores it in binary files. Replace `${ckpt}` with the path to your checkpoint or model data.

```bash
bash scripts/gemma2/tools/process_data_dolly.sh ${ckpt}  # Process Dolly Train / Validation Data
```

## Base Pre-trained Models

To fine-tune the model, download the required pre-trained model checkpoints from \[Huggingface Model Hub] and place them in the `checkpoints/` directory.

Alternatively, you can change the `CKPT` variable in the script to automatically download the base model via Transformers.

## Training Instructions

The repository provides example scripts for training **mixSGA** on **Gemma2** models. You can find other model families' scripts in the `scripts/gemma2` folder.

### Fine-tune **mixSGA** models

```bash
bash scripts/gemma2/sft/sft_2B_lora_dygqa.sh ${PORT} ${LR} ${dynamic_gqa_sizes} ${dynamic_gqa_capacity}
```

## Evaluation

To evaluate the model after training, you can use the following command. It evaluates the model using **mixSGA** with different dynamic configurations.

```bash
bash scripts/gemma2/eval/eval_main_dolly_lora.sh ${PORT} ${ckpt_name} ${ckpt} ${LORA_PATH} ${dynamic_gqa_sizes} ${dynamic_gqa_capacity}
```

## How it Works

### Key Features:

* **Dynamic Token-Wise Routing:** **mixSGA** assigns tokens dynamically to experts based on learned importance scores, allowing efficient computation and memory usage.
* **Weight-sharing across Experts:** To minimize the overhead of additional parameters, attention projections are shared across all experts.
* **Auxiliary Loss for Consistency:** During training, a consistency loss ensures one-hot routing for both training and inference, ensuring routing decisions are consistent between the prefill and decoding phases.
* **Compatibility with KV Eviction:** The method can be integrated with existing KV cache eviction strategies such as **H2O** to further reduce memory usage.

### Built on `distillm`

The implementation of **mixSGA** is built upon the [DistillM repository](https://github.com/jongwooko/distillm), which streamlines the process of distilling knowledge in large language models. We extend the original framework to introduce dynamic token-wise KV optimization and heterogeneous expert capacities.

## Citation

If you use this code in your work, please cite the following paper:

```
@misc{song2025mixSGA,
      title={Mixture of Weight-shared Heterogeneous Group Attention Experts for Dynamic Token-wise KV Optimization}, 
      author={Guanghui Song and Dongping Liao and Yiren Zhao and Kejiang Ye and Cheng-zhong Xu and Xitong Gao},
      year={2025},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2506.13541}, 
}
```

mixSGA

mixSGA 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

mixSGA

mixSGA 최신버젼 다운로드

오늘의 추천 프로젝트

Was this content helpful?

Recent Feedback

Download README