You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Commit 08fdc0d
Browse filesBrowse the repository at this point in the historyBrowse files
parser.add_argument("--dmd2_fake_warmup_steps", type=int, default=100, help="Initial critic-only updates before training the student; 0 disables fake warmup.")
95
+
parser.add_argument("--dmd2_fake_learning_rate", type=float, default=None, help="Learning rate for fake and GAN classifier; defaults to learning_rate, which controls the student.")
96
+
parser.add_argument("--dmd2_generator_gan_loss_weight", type=float, default=0.005, help="Weight of the GAN loss for generator updates.")
97
+
parser.add_argument("--dmd2_discriminator_gan_loss_weight", type=float, default=0.01, help="Weight of real/fake GAN losses for fake model and classifier updates.")
98
+
parser.add_argument("--dmd2_student_sigmas", type=str, required=True, help="Sigmas of the student model.")
99
+
parser.add_argument("--dmd2_student_timesteps", type=str, required=True, help="Timesteps of the student model.")
100
+
parser.add_argument("--dmd2_teacher_cfg", type=float, default=3., help="CFG scale of the teacher model.")
101
+
parser.add_argument("--dmd2_noise_sampling", choices=["uniform", "beta"], default="uniform", help="DM and fake noise sampling: uniform or Beta(2.5, 1). GAN always uses uniform sampling.")
102
+
parser.add_argument("--dmd2_min_timestep", type=int, default=20, help="Lower randint bound for uniform DM sampling; beta DM/fake sigma is clamped below at (value + 1) / 1000.")
103
+
parser.add_argument("--dmd2_max_timestep", type=int, default=981, help="Exclusive upper randint bound for uniform DM sampling; beta DM/fake sigma is clamped above at value / 1000.")
Copy file name to clipboardExpand all lines: docs/en/Training/Distillation.md
+35-1Lines changed: 35 additions & 1 deletion
Display the source diff
Display the rich diff
Original file line number
Diff line number
Diff line change
@@ -27,7 +27,41 @@ Direct Distill is an end-to-end direct distillation training scheme that we desi
27
27
28
28
## DMD2
29
29
30
-
[DMD2](https://arxiv.org/abs/2405.14867) (Distribution Matching Distillation) is a product that combines the theories of Diffusion and GAN. Its training relies on a manually constructed discriminator model. We have only adapted DMD2 training on the FLUX.2 model. For details, see the [example code](https://github.com/modelscope/DiffSynth-Studio/tree/main/examples/flux2/model_training/special/dmd2).
30
+
[DMD2](https://arxiv.org/abs/2405.14867) (Distribution Matching Distillation) is a model distillation training scheme that combines distribution matching with GAN training. A frozen teacher predicts the target distribution, while a fake-score model estimates the student's generated distribution. Together, their predictions provide the student's distribution-matching gradient. The GAN loss supplies an additional training signal from real images to improve detail in images generated by the distilled model.
31
+
32
+
Training involves three model roles:
33
+
34
+
* Student: Generates samples along an ordinary Euler trajectory, retaining gradients only for a randomly selected sampling step, and optimizes distribution-matching and GAN losses.
35
+
* Teacher: Has frozen parameters, predicts the target distribution, and uses `--dmd2_teacher_cfg` to control CFG.
36
+
* Fake-score model: Learns to denoise student-generated samples and tracks the student's generated distribution. A GAN classifier head uses intermediate fake-score features to distinguish real images from student-generated images.
37
+
38
+
Real images are used only for GAN training to provide a signal for recovering image details; they do not contribute to the distribution-matching loss. With GAN disabled, dataset images serve only as placeholders for `height` and `width` to determine the generated image size. They are not encoded, and their pixel content is not used in training.
39
+
40
+
We provide the following training and validation examples for FLUX.2, Qwen-Image-2.1, MOSS-Image, and FLUX.3:
41
+
42
+
|Base Model|Trainable Parameters|Training Code|Post-Training Validation|Steps Before Distillation|Steps After Distillation|
*`--dmd2_student_sigmas` and `--dmd2_student_timesteps` specify the student schedule and must correspond to each other. Use `examples/dev_tools/compute_dmd2_sigmas.py` to calculate them. Validation should use the same schedule as training, with student sampling at CFG=1.
58
+
*`--dmd2_teacher_cfg` controls the CFG strength of teacher predictions and determines the guided target the student learns. A value of 1 applies no additional CFG guidance. This parameter affects only the teacher; student and fake-score predictions use CFG=1, and validation does not need to reuse the teacher's CFG.
59
+
*`--dmd2_fake_warmup_steps` controls the initial number of critic-only updates; set it to 0 to disable warmup. `--dmd2_fake_updates_per_generator` controls the update ratio after warmup. Steps here count optimization windows after gradient accumulation.
60
+
*`--learning_rate` controls the student learning rate. `--dmd2_fake_learning_rate` controls the fake-score and classifier learning rate and defaults to the student learning rate when omitted.
61
+
* GAN training is enabled by default and requires real images. `--dmd2_disable_gan` disables it, in which case real images are not encoded. For GAN training, `--dmd2_feature_module_of_dit_for_gan` selects the backbone layer for feature extraction, and `--dmd2_classifier_channels` controls the classifier head width. Each model's training module handles its feature layout.
62
+
*`--dmd2_noise_sampling` supports `uniform` and `beta`. The latter uses `Beta(2.5, 1)`, favoring larger sigma values in the high-noise region, with bounds controlled by `--dmd2_min_timestep` and `--dmd2_max_timestep`. This option affects DM and fake denoising; GAN noise sampling always remains uniform.
63
+
* Checkpoints export only student parameters. Use `--lora_checkpoint` to warm-start LoRA or `--resume_from_checkpoint` to load full student weights. These do not restore the fake-score model, classifier, optimizer, or random state and are not a full training resume. Replace the checkpoint paths in validation scripts with the actual training outputs.
64
+
* DMD2 training currently does not support DeepSpeed or `--enable_model_cpu_offload`.
0 commit comments