[Feature] Setting number of steps for DMD distillation

[imagesite-account]

Hi,

I want to perform 4 step distillation on one of the example scripts and I wonder what I should modify to do so. Let's say we use this script: https://github.com/hao-ai-lab/FastVideo/blob/40ce2d72f5c74f199033643bb67f0a7c9c280a58/examples/distill/Wan2.2-TI2V-5B-Diffusers/crush_smol/distill_dmd_VSA_t2v_5B.sh

First of all, within the training scripts, I assume validation_sampling_steps should be set from 3 to 4. After that, dmd_denoising_steps should be changed to include a fourth timestep -- but what value should this be? Should we map a uniform noise schedule to a nonuniform one using of the schedulers (e.g. https://github.com/hao-ai-lab/FastVideo/blob/40ce2d72f5c74f199033643bb67f0a7c9c280a58/fastvideo/models/schedulers/scheduling_flow_unipc_multistep.py#L287)? 

Related resources

No response

[yuhezhang-ai]

Hi! I’m not a maintainer, just digging through the code to learn and hopefully contribute. But I looked into how the existing 3-step schedule relates to warping, and I think it’s not derived from a simple “uniform then warp” rule.

I tried reproducing what would happen if we started from uniform indices and relied on warp_denoising_step using the same scheduler Wan uses with the script:

from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler

# Create scheduler with Wan2.2's shift
scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0)
scheduler.set_timesteps(num_inference_steps=1000)

# 3-step uniform indices
indices = [0, 333, 666]
print([int(v) for v in scheduler.timesteps[indices]])

# 4-step uniform indices
indices = [0, 250, 500, 750]
print([int(v) for v in scheduler.timesteps[indices]])

This gives:

• 3-step warped values: [1000, 941, 803]
• 4-step warped values: [1000, 960, 890, 732]

These don’t match the existing config ([1000, 757, 522]) at all, which suggests the 3-step schedule was likely empirically chosen / hand-tuned, rather than derived from uniform spacing plus warping.

Because of that, I don’t think there’s a single “correct” automatic way to extend it to 4 steps unless the authors already validated one. If we want a reasonable heuristic that roughly preserves the spacing pattern of the existing 3-step schedule (instead of following uniform+warp), something like:

[1000, 800, 610, 430]

seems closer in spirit (monotonic, spread across noise levels, not overly concentrated at high noise).

That said, I’d defer to maintainers on whether there’s an official 4-step schedule they recommend.

---

BTW, if you do want to use uniform spacing and let the pipeline handle warping automatically, you can do that explicitly:

--dmd_denoising_steps '1000,750,500,250'
--warp_denoising_step

The warping logic lives in distillation_pipeline.py (around line 252):

if training_args.warp_denoising_step:
    timesteps = torch.cat(
        (self.noise_scheduler.timesteps.cpu(),
         torch.tensor([0], dtype=torch.float32))
    ).cuda()
    self.denoising_step_list = timesteps[1000 - self.denoising_step_list]
    logger.info("Warping denoising_step_list")