OPD做步数蒸馏的疑问

我在阅读 OPD 实现时注意到，align_sigmas 选取了学生当前区间内的教师节点，但 scheduler 会将最后一步直接更新到 sigma=0。

以Qwen-Image-2.1-Turbo的8步蒸馏到3步为例子，学生当前区间为 1 → 0.845148，教师实际执行的却是 1 → 0.978453 → 0.954180 → 0.926626 → 0.895080 → 0，然后用起始 sigma 归一化得到速度目标。

下面的表格写了两种方法，为什么diffsynth采用第一种方法，而不是“严格区间匹配”这种做法？第一种方法效果更好嘛？

方案	教师积分路径	最后一次教师前向的位置
diffsynth当前实现	1 → .978453 → .954180 → .926626 → .895080 → 0	.895080
严格区间匹配	1 → .978453 → .954180 → .926626 → .895080 → .845148	.895080

我试了第二种方法进行蒸馏，感觉效果还挺好的。第一种默认的方法没有尝试。

@651961

@651961 实际上两者实现在数学意义上完全等价，在教师模型推理结束后，我们需要把这多步推理合成一步。

将学生模型开始推理的起始点记为 sigma=u，教师模型推理的最后一个点记为 sigma=v，学生模型下一次推理的起始点记为 sigma=w

diffsynth 目前的实现方式是跳到 sigma=0 然后反推在开始时模型应该输出什么，此时速度目标为 $\frac{x_u-x_v+v\cdot e_v}{u}$
你提到的实现是跳到下一次学生模型推理的 sigma，此时速度目标为 $\frac{x_u-x_v+(v-w)\cdot e_v}{u-w}$

注意到 w 仅影响在哪里对齐，前者是后者设定为 w=0 的特例，清除了速度目标中的所有噪声部分。但后者在计算过程中需要除以区间 sigma 的差值，这个数值更小，容易出现较大的舍入误差，因此我们采用了前者。我们会考虑进行额外的实验，验证两者的训练效果，未来可能将其设定为一个超参数。

感谢回答。
不过，“数学上完全等价”应该是需要额外条件吧，他应该是不适用步数蒸馏的。
你给出的两个公式满足：

T_0 = ((u - w) / u) × T_w + (w / u) × e_v

因此，当 0 < w < u 时，两者相等的条件是 T_w = e_v。

比如一种情况，教师和学生的sigmas网格完全一致，你这个公式就成立。
但是步数蒸馏，教师和学生的sigmas网格不可能一致。

以qwen image 2.1 Turbo 8步蒸馏→3步蒸馏为例，如果学生是 1 → 0.845148 区间，

DiffSynth：
1 → .978453 → .954180 → .926626 → .895080 → 0

严格区间匹配：
1 → .978453 → .954180 → .926626 → .895080 → .845148

两者都做五次教师前向。考虑一个反例：起始状态 x_1 = 0，前四次教师预测速度均为 1，最后在 .895080 的预测速度为 2。

前四步完全相同，得到：

x_v = -(1 - .895080) = -.104920

DiffSynth 最后更新到 0：

x_0 = -.104920 - .895080 × 2
    = -1.895080

T_0 = (x_1 - x_0) / 1
    = 1.895080

严格区间匹配最后更新到 .845148：

x_w = -.104920 - (.895080 - .845148) × 2
    = -.204784

T_w = (x_1 - x_w) / (1 - .845148)
    = .204784 / .154852
    ≈ 1.322450

一个是1.895080，另一个是1.322450，两个速度目标不同。

我试了第二种方法进行蒸馏，感觉效果还挺好的。第一种默认的方法没有尝试。

蒸馏使用的数据，是什么规模呢

@Artiprocher

请教一个其他问题，我想在Qwen-Image-2.1-Turbo 基础上微调一个风格lora，直接sft lora 会破坏turbo本身的加速能力；能参考Z-Image-Turbo-DistillPatch进行Qwen-Image-2.1-Turbo的差分训练吗 @Artiprocher @wanglaofei

@qiu-pinggaizi

SFT配步数蒸馏的做法是在base模型上训练，也就是在Qwen-Image-2.1上面训练，得到Qwen-Image-2.1-SFT。然后做权重合并，把turbo合并进来。
@qiu-pinggaizi

@651961

6. SFT配步数蒸馏的做法是在base模型上训练，也就是在Qwen-Image-2.1上面训练，得到Qwen-Image-2.1-SFT。然后做权重合并，把turbo合并进来。
@651961
我已经得到了Qwen-Image-2.1-SFT，是针对具体某个风格做的lora微调；使用什么方法将turbo 合并到Qwen-Image-2.1-SFT中呢

@qiu-pinggaizi

额，最简单的权重合并就行吧，更好的方法也感知不到什么提升。我刚才让codex写了个脚本，你改一下路径什么的应该就能用。
@qiu-pinggaizi

DISTILLED_MODEL_PATH = "/models/Qwen-Image-2.1-Turbo"  # B：蒸馏模型
BASE_MODEL_PATH = "/models/Qwen-Image-2.1"  # W0：SFT 训练前的base基座
SFT_MODEL_PATH = "/datasets/codes_zsqiao/image_gen/DiffSynth-Studio/models/train/Qwen-Image-2.1_full_layered_multinode/epoch-39.safetensors"  # A：DiffSynth-Studio微调结果
MERGE_STRENGTH = 1.0
OUTPUT_PATH = (
    f"/models/ckpt_zsqiao/qwen-image-2.1-lambda-{MERGE_STRENGTH:.1f}.safetensors"
)
VALIDATE_ONLY = False  # True：只校验权重名称、形状和类型，不进行合并
CPU_THREADS = 8

import json
import math
import os
import struct
import tempfile
from contextlib import ExitStack
from pathlib import Path

import torch
from safetensors import safe_open

# 输入可以是单个 safetensors、包含分片的 transformer 目录，或完整模型目录。
# 输出仅包含 transformer 权重；推理时继续使用 B 的配置、scheduler 和 8 步设置。
# 不依赖 mergekit / diffusers。依赖：pip install torch safetensors numpy
DTYPES = {
    "BF16": (torch.bfloat16, 2),
    "F16": (torch.float16, 2),
    "F32": (torch.float32, 4),
}


def find_weight_files(path):
    path = Path(path).expanduser().resolve()
    if path.is_file():
        if path.suffix != ".safetensors":
            raise ValueError(f"只支持 safetensors 权重：{path}")
        return [path]
    if (path / "transformer").is_dir():
        path = path / "transformer"
    files = sorted(path.glob("*.safetensors"))
    if not files:
        raise FileNotFoundError(f"找不到 safetensors 权重：{path}")
    return files


def open_checkpoint(stack, path):
    """只读取元数据；实际张量在合并时逐个读取。"""
    tensors = {}
    for filename in find_weight_files(path):
        reader = stack.enter_context(safe_open(filename, framework="pt", device="cpu"))
        for key in reader.keys():
            if key in tensors:
                raise ValueError(f"多个分片中出现重复参数：{key} ({filename})")
            view = reader.get_slice(key)
            tensors[key] = (reader, view.get_shape(), view.get_dtype())
    if not tensors:
        raise ValueError(f"权重中没有张量：{path}")
    return tensors


def validate_checkpoints(distilled, base, sft):
    expected = set(distilled)
    for label, checkpoint in (("基座", base), ("SFT", sft)):
        missing = sorted(expected - set(checkpoint))
        extra = sorted(set(checkpoint) - expected)
        if missing or extra:
            raise ValueError(
                f"{label} 与 Turbo 参数名不一致："
                f"缺失 {len(missing)} 个 {missing[:5]}；"
                f"多出 {len(extra)} 个 {extra[:5]}"
            )
    for key, (_, shape, _) in distilled.items():
        for label, checkpoint in (("Turbo", distilled), ("基座", base), ("SFT", sft)):
            _, other_shape, dtype = checkpoint[key]
            if shape != other_shape:
                raise ValueError(
                    f"参数形状不一致：{key}，Turbo={shape}，{label}={other_shape}"
                )
            if dtype not in DTYPES:
                raise ValueError(f"不支持的参数类型：{label}/{key} 为 {dtype}")


def make_header(distilled, metadata):
    """预先计算 safetensors 偏移，以逐张量写入单文件，避免缓存完整输出。"""
    header = {"__metadata__": metadata}
    offset = 0
    for key in sorted(distilled):
        _, shape, dtype = distilled[key]
        size = math.prod(shape) * DTYPES[dtype][1]
        header[key] = {
            "dtype": dtype,
            "shape": shape,
            "data_offsets": [offset, offset + size],
        }
        offset += size
    encoded = json.dumps(header, ensure_ascii=False, separators=(",", ":")).encode(
        "utf-8"
    )
    encoded += b" " * (-len(encoded) % 8)
    return encoded, offset


@torch.no_grad()
def merge_models(
    distilled_path, base_path, sft_path, output_path, strength=0.5, validate_only=False
):
    if not math.isfinite(strength):
        raise ValueError("合并强度必须为有限数值")
    output = Path(output_path).expanduser().absolute()
    if output.suffix != ".safetensors":
        raise ValueError("输出文件名必须以 .safetensors 结尾")
    if not validate_only and (output.exists() or output.is_symlink()):
        raise FileExistsError(f"输出已存在，请修改 OUTPUT_PATH（不会覆盖）：{output}")

    print(f"B（Distill）：{distilled_path}", flush=True)
    print(f"W0（Base）：{base_path}", flush=True)
    print(f"A（SFT）：{sft_path}", flush=True)
    print(f"合并公式：B + {strength:g} * (A - W0)", flush=True)

    with ExitStack() as stack:
        distilled = open_checkpoint(stack, distilled_path)
        base = open_checkpoint(stack, base_path)
        sft = open_checkpoint(stack, sft_path)
        validate_checkpoints(distilled, base, sft)
        metadata = {
            "format": "pt",
            "merge_formula": "B + strength * (A - W0)",
            "merge_strength": str(strength),
            "distilled_model": str(Path(distilled_path).expanduser().resolve()),
            "base_model": str(Path(base_path).expanduser().resolve()),
            "sft_model": str(Path(sft_path).expanduser().resolve()),
        }
        header, data_size = make_header(distilled, metadata)
        print(
            f"校验通过：{len(distilled)} 个参数张量，名称和形状完全一致。"
            f"输出权重约 {data_size / 1024**3:.2f} GiB，保存为 Turbo 的原始精度。",
            flush=True,
        )
        if validate_only:
            print("仅校验完成，未写入权重。", flush=True)
            return

        output.parent.mkdir(parents=True, exist_ok=True)
        # 临时文件与输出位于同一目录；成功后才发布最终文件，不覆盖已有文件。
        fd, temp_name = tempfile.mkstemp(
            prefix=f".{output.name}.", suffix=".tmp", dir=output.parent
        )
        temp_path = Path(temp_name)
        try:
            with os.fdopen(fd, "wb") as handle:
                handle.write(struct.pack("<Q", len(header)))
                handle.write(header)
                for index, key in enumerate(sorted(distilled), 1):
                    b_reader, _, dtype = distilled[key]
                    b_tensor = b_reader.get_tensor(key)
                    if strength == 0:
                        result = b_tensor
                    else:
                        # 先转 FP32 再相减，避免在 BF16 中计算增量时损失精度。
                        result = (
                            sft[key][0]
                            .get_tensor(key)
                            .to(dtype=torch.float32, copy=True)
                        )
                        result.sub_(base[key][0].get_tensor(key).float())
                        result.mul_(strength).add_(b_tensor.float())
                        result = result.to(DTYPES[dtype][0])
                    if not torch.isfinite(result).all().item():
                        raise ValueError(f"合并结果含 NaN/Inf：{key}")
                    # uint8 视图也支持 BF16；memoryview 避免复制成完整 bytes。
                    raw = result.contiguous().reshape(-1).view(torch.uint8).numpy()
                    handle.write(memoryview(raw))
                    del raw, result, b_tensor
                    print(f"[{index}/{len(distilled)}] {key}", flush=True)
                expected_size = 8 + len(header) + data_size
                if handle.tell() != expected_size:
                    raise RuntimeError("输出文件大小与 safetensors 元数据不一致")
                handle.flush()
                os.fsync(handle.fileno())

            # 用官方读取器验证文件头及全部张量的名称、形状和类型。
            with safe_open(temp_path, framework="pt", device="cpu") as reader:
                if set(reader.keys()) != set(distilled):
                    raise RuntimeError("输出参数名校验失败")
                for key, (_, shape, dtype) in distilled.items():
                    view = reader.get_slice(key)
                    if view.get_shape() != shape or view.get_dtype() != dtype:
                        raise RuntimeError(f"输出元数据校验失败：{key}")
            os.link(temp_path, output)  # 原子发布；若目标已存在，则失败而不覆盖。
        finally:
            temp_path.unlink(missing_ok=True)

    print(f"合并完成：{output}", flush=True)
    print(
        "输出仅包含 transformer 权重；推理请沿用 Turbo 的配置及 DMD 8 步设置。",
        flush=True,
    )


if __name__ == "__main__":
    torch.set_num_threads(CPU_THREADS)
    merge_models(
        DISTILLED_MODEL_PATH,
        BASE_MODEL_PATH,
        SFT_MODEL_PATH,
        OUTPUT_PATH,
        strength=MERGE_STRENGTH,
        validate_only=VALIDATE_ONLY,
    )

@651961 实际上两者实现在数学意义上完全等价，在教师模型推理结束后，我们需要把这多步推理合成一步。

将学生模型开始推理的起始点记为 sigma=u，教师模型推理的最后一个点记为 sigma=v，学生模型下一次推理的起始点记为 sigma=w
diffsynth 目前的实现方式是跳到 sigma=0 然后反推在开始时模型应该输出什么，此时速度目标为
    x
    u
  
  −
  
    x
    v
  
  +
  v
  ⋅
  
    e
    v
  

u
你提到的实现是跳到下一次学生模型推理的 sigma，此时速度目标为
    x
    u
  
  −
  
    x
    v
  
  +
  (
  v
  −
  w
  )
  ⋅
  
    e
    v
  


  u
  −
  w
注意到 w 仅影响在哪里对齐，前者是后者设定为 w=0 的特例，清除了速度目标中的所有噪声部分。但后者在计算过程中需要除以区间 sigma 的差值，这个数值更小，容易出现较大的舍入误差，因此我们采用了前者。我们会考虑进行额外的实验，验证两者的训练效果，未来可能将其设定为一个超参数。

看danceopd的做法，

            dt = (t_next - t_cur) / 1000.0
            if stochastic:
                sigma = (t_cur.float() / 1000.0).reshape(1)
                sigma_next = (t_next.float() / 1000.0).reshape(1)
                next_latents, old_log_prob, _, _ = sde_step(
                    v, sigma, sigma_next, latents,
                    noise_level=float(self.cfg.training.get("flowopd_noise_level", 0.7)),
                )
                trajectory.append(FlowState(
                    latents=latents.detach().clone(), timestep=t_cur.detach().clone(),
                    extra={
                        "next_latents": next_latents.detach().clone(),
                        "timestep_next": t_next.detach().clone(),
                        "sigma": sigma.detach().clone(),
                        "sigma_next": sigma_next.detach().clone(),
                        "old_log_prob": old_log_prob.detach().clone(),
                    },
                ))
                latents = next_latents
            else:
                if selected is None:
                    trajectory.append(FlowState(latents=latents.detach().clone(), timestep=t_cur.detach().clone()))
                latents = latents + dt * v
            latents = latents.to(self.torch_dtype)
        return trajectory, latents.detach()

看起来不太一样。是出于什么考虑？

Navigation Menu

Repository navigation

OPD做步数蒸馏的疑问 #1786

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

OPD做步数蒸馏的疑问 #1786

Description

Activity

wanglaofei commented on Oct 27, 2026

Artiprocher commented on Oct 27, 2026

651961 commented on Oct 27, 2026

qiu-pinggaizi commented on Oct 30, 2026

qiu-pinggaizi commented on Oct 30, 2026

651961 commented on Oct 30, 2026

qiu-pinggaizi commented on Oct 30, 2026

651961 commented on Oct 30, 2026

dongdk commented on Nov 2, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions