MoECa: Aligning Feature Reuse with Expert Decomposition in Diffusion Transformers

🌟We have released the MoECa code, a fine-grained caching framework that reuses branch-level features across timesteps. Our code repository is based on EfficientMoE. Please refer to the original repository for model weights and environment configurations.

Repository Layout

MoECa/
├── config/          # DSMoE model configs
├── dataset/         # optional cached-latent dataset helper
├── diffusion/       # diffusion / rectified-flow implementation
├── evaluation/      # FID / sFID evaluation utilities
├── models/          # DSMoE model definition
├── scripts/         # training launch scripts
├── download.py      # checkpoint loading helper
├── sample_ddp_feature.py
├── train.py
├── train_fsdp.py
└── utils.py

Installation

pip install torch==2.6.0 torchvision==0.21.0
pip install  peft==0.17.1
pip install pexpect timm torchdiffeq tensorboard diffusers transformers
pip install tensorflow==2.15.0
pip install -e git+https://github.com/LTH14/torch-fidelity.git@master#egg=torch-fidelity

Please follow the installations of EfficentMoE, respectively.

Training

Single-node DDP example:

sh ./scripts/train_s_e16_dsmoe.sh

See details in EfficentMoE respectively.

Sampling

Example:

torchrun --nproc_per_node=4 sample_ddp_feature.py \
  --ckpt /path/to/checkpoint.pt \
  --config ./config/000_DSMoE_S_E16_Flow_half_misx2_s1a2_rope2d_l10.yaml \
  --vae stabilityai/sd-vae-ft-mse \
  --save-png \
  --save-activations

Evaluation

We follow the evaluation protocols provided by EfficentMoE.