# EDGC: Entropy-driven Dynamic Gradient Compression for Efficient LLM Training

EDGC是一个面向大规模语言模型分布式训练的**高效动态梯度压缩框架**。它通过“以训练阶段为纲，流水线时序为脉”的细粒度调度策略，在保障模型收敛的同时显著降低通信开销，提升训练吞吐。

在大模型训练中，梯度同步已成为扩展效率的瓶颈。虽然“**低秩分解 + 误差反馈**”能有效压缩梯度并防止误差累积，但静态压缩策略仍面临三大挑战：

1. **训练初期不稳定**：早期梯度方向变化剧烈，过早强压缩易导致优化偏移。
2. **静态策略不灵活**：固定压缩率难以兼顾训练前/后期对保真度与效率的不同需求。
3. **流水线并行存在时序差异**：非瓶颈阶段有“时间冗余”，但传统方法未加以利用。

> **核心理念**：高效的压缩不仅在于“如何压”，更在于 **“何时压”与“对谁压”**。

---

## 🌟 核心设计

### 动态压缩调度
- 训练初期禁用压缩（预热），避免干扰优化方向。
- 进入稳定期后，逐步启用并动态调整压缩率，实现“该保则保，该压则压”。

### 差异化压缩（基于流水线并行）
- 利用非瓶颈 Stage 的“时间冗余窗口”，提高其压缩强度；
- 在关键路径 Stage 保留低压缩比，确保同步效率与梯度质量。

### 误差反馈（Error Feedback）
- 持久化存储压缩残差，防止信息丢失；
- 保障长期训练中的收敛稳定性。

---

## 🚀 适用场景

EDGC 特别适用于以下分布式训练场景：

### 1. 大规模语言模型（LLM）训练
- **特征**：参数量巨大（>1B），梯度通信成为主要瓶颈。
- **收益**：显著减少梯度传输量；提升集群吞吐；支持分布式优化器与 `overlap-grad-sync` 等机制。
- ✅ **推荐用于**：BERT、LLaMA、GPT 等 Transformer 架构的预训练、后训练与微调。

### 2. 高延迟或低带宽网络环境
- **特征**：跨数据中心、云边协同、RoCE 受限。
- **收益**：缓解网络拥塞，减少同步等待时间，提升弱网下的训练稳定性。
- ✅ **推荐用于**：边缘设备协作训练、混合云部署、低成本 GPU 集群。

### 3. 计算-通信失衡系统
- **特征**：高端 GPU（如 A100/H100）搭配中等网络（如 25Gbps）。
- **问题**：GPU 算力过剩，通信拖累整体效率。
- **收益**：通过动态秩调节平衡通信与计算耗时，避免“算得快，传得慢”，提升 GPU 利用率。
- ✅ **推荐用于**：非对称硬件配置集群。

### 4. 长周期、高成本训练任务
- **特征**：训练持续数天甚至数周。
- **收益**：节省通信能耗与时间成本；自适应调节避免手动调参风险；支持容错恢复与断点续训。
- ✅ **推荐用于**：科研项目、商业级大模型训练。

---

## ⚠️ 不适用场景

尽管 EDGC 具有广泛适用性，但在以下情况需谨慎使用：

| 场景 | 原因 |
|------|------|
| 极短训练周期（<1 小时） | 压缩收益不足以覆盖初始化开销 |
| 模型极小（<10M 参数） | 通信开销本身很低，压缩反而引入额外计算负担 |
| 对收敛精度要求极高且无法容忍任何偏差 | 虽然 EF 机制已极大缓解失真，但仍属有损压缩 |
| 使用极稀疏梯度的算法（如 DDP + gradient accumulation + sparse update） | 可能与现有稀疏化策略冲突 |

---

## ⚙️ 使用说明

### 核心配置参数

| 参数 | 类型 | 默认值 | 说明 |
|------|------|--------|------|
| `--enable-dynamic-grad-comp` | bool | `False` | **是否启用 EDGC 框架**。设为 `True` 后开启动态梯度压缩功能。只需添加此参数即可启用 EDGC，其余参数默认值已在多种典型场景中验证为最优，一般无需修改。 |
| `--grad-comp-warm-up` | float32 | `0.1` | **压缩预热比例**。表示在训练总步数的前 X% 阶段不启用压缩（即“预热期”）。例如，若总训练 100K 步且该值为 `0.1`，则前 10K 步不压缩，让模型稳定优化。避免训练初期梯度剧烈变化时压缩带来的方向偏移。 |
| `--rank-adjust-window-size` | int | `1000` | **动态秩调节的观测窗口大小**（单位：迭代次数）。系统每隔该窗口长度进行一次压缩策略评估，并决定是否调整压缩秩。较大的窗口更稳定，较小的窗口响应更快。默认值 `1000` 在精度与响应速度间取得良好平衡。 |
| `--iteration-sample-ratio` | float32 | `0.1` | **性能采样比例**。在每个调节窗口内，仅对其中 `X%` 的迭代进行实际性能探测（如通信时间、误差影响），以降低开销。例如 `0.1` 表示只采样 10% 的迭代来估算整体效果，兼顾效率与准确性。 |
| `--gradient-sample-ratio` | float32 | `0.1` | **梯度采样比例**。在性能探针阶段，仅对部分梯度（如 10%）进行低秩分解测试，用于快速估算当前最优压缩秩。减少探测过程中的计算负担，加快决策速度。 |
| `--collect-log-path` | str | `./logs` | **EDGC 系统日志输出路径**。用于保存压缩率、秩变化、通信时间等运行时监控信息，便于后续分析与调优。建议定期检查日志以确认压缩行为符合预期。 |
| `compression_dtype` | str | `bf16` | **压缩过程中间变量的数据类型**。例如低秩矩阵 `P` 和 `Q` 的计算精度。推荐使用 `bf16`，在精度损失极小的前提下显著降低显存占用和计算开销。可选 `fp32` 或 `fp16`，但通常不必要。 |
| `ef_store_dtype` | str | `bf16` | **误差反馈缓冲区的存储精度**。控制历史压缩残差的存储类型。使用 `bf16` 可节省显存（相比 `fp32` 减半），同时保留足够数值精度。对于显存紧张场景尤其重要。 |
| `use_error-feedback` | bool | `True` | **是否启用误差反馈机制**。强烈建议保持开启（`True`），否则压缩误差会累积，可能导致收敛失败或性能下降。该机制确保未传输的梯度信息被保留到下一轮，保障收敛稳定性。 |
| `rank_adjust_step` | int | `4` | **每次调整压缩秩的最大步长**。防止压缩强度突变导致训练不稳定。例如，当前秩为 64，最多调整 ±4，变为 60 或 68。小步长更安全，大步长响应更快。默认值 `4` 提供良好的动态适应性。 |
---
要启用 EDGC，仅需在训练命令中加入`--enable-dynamic-grad-comp`参数。所有其他参数的默认值均基于广泛实验调优，适用于大多数场景，无需手动配置。

## ❓ 常见问题与应对策略

### 1. 压缩后模型收敛变慢或不收敛
- **可能原因**：压缩过强（秩太低）。
- **解决方法**：
  - 提高 `--grad-comp-warm-up`（如从 `0.1` 改为 `0.2`），延长预热阶段；
  - 或暂时关闭动态压缩：`--enable-dynamic-grad-comp=False`。

### 2. 显存爆了（OOM）
- **可能原因**：EDGC 需存储误差缓冲区，带来额外显存占用（约 2%~6%）。
- **解决方法**：
  - 将 `compression_dtype` 和 `ef_store_dtype` 设为 `bf16`（比 `fp32` 节省 50% 显存）；
  - 若使用 ZeRO，确保开启分片模式（`distributed_sharding=True`）。

### 3. 通信时间未下降
- **可能原因**：
  - 模型太小或 batch size 太小，通信本就不成瓶颈；
  - 压缩逻辑未正确集成。
- **解决方法**：
  - EDGC 更适用于 >1B 参数模型或大 batch 训练；
  - 检查压缩是否在 `AllReduce` 前正确执行。

### 4. 初始训练阶段 `edgc_utils` 报错
- **可能原因**：通信未成为瓶颈，或压缩成本高于通信时间，导致无法建模。
- **解决方法**：关闭 EDGC 功能：`--enable-dynamic-grad-comp=False`。

---
