# RS-MoE: Collaborative Compression for Mixture-of-Experts LLMs based on Low-Rank and Sparse Approximation

简要说明  
本项目实现了一套面向 MoE（Mixture-of-Experts）架构的压缩与运行时替换方案，包含激活感知的重要性评估、按维度/列的差异化压缩（高/中/低等级别）、以及用于推理的轻量化 OAGE runtime 模块。代码以可复现的流水线组织，支持模型蒸馏式/低秩+稀疏混合压缩并保留推理行为。

主要特点
- ACI（Anomalous Contribution Integration）：基于中间激活与下游投影能量的维度重要性评分。
- 差异化压缩：按重要性分配高/中/低组，高组保留原列/行，中组低秩分解，低组置零。
- down-proj 基准拟合 + 桥接矩阵，减少重构误差。
- 推理时替换为 OAGEDeepseekMLP / OAGEDeepseekMoE（将权重以 buffer 形式保存以减少训练参数计数）。
- 吞吐量（throughput）与 PPL 评估脚本集成，支持按批次计时与 warmup。
- 支持多种 MoE 实现风格（项目中包括对 DeepSeek-like 与 Phi-like 路径的兼容逻辑）。

目录结构（简要）
- compressor/deepseek_compressor.py — 压缩主实现 + OAGE runtime 模块
- compress_deepseek.py — 高层压缩流水线入口（集成分组、MiNE、分配、压缩与评估）
- expert_group.py — 专家分组与行为画像提取
- mine.py — 互信息估计（MINE）与层间重要性评估
- test.py — PPL / 吞吐量测试与示例生成
- utils/ — 若干工具（数据处理、评估、保存/加载等）
- results/ — 运行/评估输出目录

环境与依赖
- Linux (开发与评估在 Linux 上测试)
- Python 3.8+
- PyTorch（与 CUDA）  
- transformers（trust_remote_code 支持自定义模型实现）
- datasets, accelerate, scikit-learn, tqdm, numpy, scipy
- 建议 GPU：支持大显存卡（多卡部署可显著提高吞吐量）

快速开始（示例）
1. 安装依赖（建议使用虚拟环境）
   ```
   pip install -r requirements.txt
   ```
2. 运行专家分组与压缩流水线（示例）
   ```
   # 可通过 --compression_ratio, --num_samples 等控制行为
   python compress_deepseek.py --compression_ratio 0.2 --num_samples 256
   ```
3. 评估与吞吐量测量（在 test.py 中已有开关）
   ```
   # 在 test.py 中可启用 measure_throughput 参数
   python test.py
   ```
常用运行提示
- 若需固定使用哪些 GPU，可先设置环境变量：
  ```
  export CUDA_VISIBLE_DEVICES="0,1,2"
  ```
  或在脚本顶部设置 os.environ["CUDA_VISIBLE_DEVICES"]。
- 生成时报错与 DynamicCache 相关：可在生成时禁用 kv-cache
  ```
  model.config.use_cache = False
  model.generate(..., use_cache=False, ...)
  ```
- 若想查看每批耗时用于吞吐量诊断，请在 timed loop 中记录并输出 per-batch times（test.py 已集成示例）。

可配置超参（在各入口脚本与类中）
- compression_ratio、layer_group_params、num_calibration_samples、max_seq_length、gamma（ACI 下游权重）、w_mean/w_var/w_peak（ACI 三权重）、sample caps、SVD backend 等。多数以 argparse 或 DeepSeekCompressor 构造参数暴露。

复现与调试建议
- 首先用小样本（num_calibration_samples=8..64）进行快速 smoke test，再扩大到完整设置。
- 为避免显存峰值，采用按专家按需移动数据到 GPU（代码已实现按专家调度）。
- 调试中关注：importance 分组结果（high/medium/low），中组的重构误差（可打印相对重建误差），以及 runtime 替换后的非零参数统计（用于验证目标稀疏率）。

评估输出
- PPL、吞吐量（tokens/sec、samples/sec）、参数/非零参数统计、压缩比/保存比（compress ratio / save ratio）会保存到 results/ 下的日志文件中。

贡献与扩展
- 欢迎按需添加对其它 MoE 实现的兼容路径（修改 `expert_group.get_model_config` 与钩子注册逻辑）。
- 可将运行时推理模块导出为 ONNX / TensorRT 以进一步提升吞吐量。

许可证与联系
- 请在使用/发布时遵守本仓库所依赖模型与数据集的许可证与使用条款。
