PowerFlow

Official implementation of PowerFlow: Unlocking the Dual Nature of LLMs via Principled Distribution Matching.

Figure: High-level illustration of PowerFlow and its key components.

Figure: Overview of the PowerFlow objective and training pipeline.

Released Models

We release official model weights on Hugging Face:

• PowerFlow collection 🤗🤗: crs25-tsinghua/powerflow

Usage

Install verl

This repo depends on verl (the RL training library).

• Option A (recommended): install the fork vendored in verl/.
• Option B: install the upstream verl: verl-project/verl

Note: using the latest upstream verl may lead to incompatibilities. We plan to upstream the required adaptations to the official verl repository.

Training (NuminaMath 20K, Qwen2.5-Math-1.5B)

After preparing your base model and environment, start training from the repo root:

cd PowerFlow
bash powerflow/run_model_nm_1.5b.sh

The training script is configurable via environment variables (e.g., MODEL_PATH, CKPTS_DIR, TRAIN_FILE, NNODES, RAY_ADDRESS). See powerflow/run_model_nm_1.5b.sh for details.

Merge checkpoints to a Hugging Face model

After training finishes and you have an FSDP checkpoint, merge it to a Hugging Face-format directory:

cd PowerFlow
bash merge.sh

merge.sh contains example paths (e.g., ckpt_path, hf_path, clean_path). Update them to your own checkpoint locations before running.

Evaluation (Qwen2.5-Math-1.5B)

Run evaluation scripts under eval/:

cd PowerFlow/eval
bash eval_qwen2.5_math_1.5b.sh

See the evaluation entrypoint: eval/eval_qwen2.5_math_1.5b.sh.

The evaluation script may activate a conda environment (default: lm_eval). Please edit the script if you use a different environment name.

lm_eval is a separate environment used to avoid interfering with the verl training environment. In practice, you can also run evaluation inside your verl environment by installing a few extra dependencies required by eval/ (see eval/requirements.txt and eval/README.md).