Question about the config for Flux.1-Dev

[yansenhan]

We are currently reimplementing this excellent work, including AWM, DiffusionNFT and FlowGRPO.

However, we cannot obtain the eval curve of the paper. Can you share your config for training Flux.1-Dev.
We find that KL beta = 0 in current config, but in AWM paper, they set it to be 0.01.
We have test both beta and still cannot obtain the result in the paper. I will appreciate very much if you can help us solve this problem.

Our config is in the following:

# >>>>>>>>>> 环境配置（按你的机器改）<<<<<<<<<<
launcher: "accelerate"
config_file: deepspeed/deepspeed_zero2.yaml   # 相对于 syf_exp/
num_processes: 8        # GPU 数量，按实际修改
main_process_port: 29500
mixed_precision: "bf16"

# >>>>>>>>>> 数据配置（必须改）<<<<<<<<<<
data:
  dataset_dir: "dataset/pickscore"   # 相对于 syf_exp/，放数据或建软链接
  preprocessing_batch_size: 8
  dataloader_num_workers: 16
  force_reprocess: true
  cache_dir: "cache/datasets"  # 相对于 syf_exp/
  max_dataset_size: 1024
  sampler_type: "auto"

# >>>>>>>>>> 模型配置（必须改）<<<<<<<<<<
model:
  finetune_type: 'lora'
  lora_rank: 64
  lora_alpha: 128
  target_modules: "default"
  model_name_or_path: "./models/FLUX.1-dev"  # HuggingFace ID 或本地绝对路径
  model_type: "flux1"
  resume_path: null
  resume_type: null

# >>>>>>>>>> 日志与保存 <<<<<<<<<<
log:
  run_name: null    # 为 null 时自动生成，建议改成有意义的名字
  project: "PERL-SYF-TEST"
  logging_backend: "swanlab"    # Options: wandb, swanlab, tensorboard, none
  save_dir: "saves/"          # 相对于 syf_exp/
  save_freq: 20
  save_model_only: true

# >>>>>>>>>> 训练配置 <<<<<<<<<<
train:
  # Trainer settings
  trainer_type: 'awm'
  advantage_aggregation: 'sum' # Options: 'sum', 'gdpo'
  off_policy: false
  awm_weighting: 'ghuber'
  ghuber_power: 0.25
  # Training Timestep distribution
  num_train_timesteps: 4 # Set null to all steps
  time_sampling_strategy: discrete # Options: uniform, logit_normal, discrete, discrete_with_init, discrete_wo_init
  time_shift: 3.0
  timestep_range: 0.7 # Select fraction of timesteps to train on
  # Clipping
  clip_range: 1  # PPO/GRPO clipping range
  adv_clip_range: 1.0  # Advantage clipping range
  # KL div
  kl_weight: 'Uniform'
  kl_type: 'v-based'
  kl_beta: 0.01 # KL divergence beta
  ref_param_device: 'cuda' # Options: cpu, cuda
  # EMA
  ema_kl_beta: 0.1 # Coefficient of KL-loss between current policy and EMA policy, used to stablize training
  ema_decay_schedule: "linear"  # Decay schedule for EMA. Options: ['constant', 'power', 'linear', 'piecewise_linear', 'cosine', 'warmup_cosine']
  ema_decay: 0.3  # EMA decay rate (0 to disable)
  ema_update_interval: 1  # EMA update interval (in epochs)
  warmup_steps: 300
  ema_device: "cuda"  # Device to store EMA model (options: cpu, cuda)

  # Sampling
  resolution: 512  # Can be int or [height, width]
  num_inference_steps: 10  # Number of timesteps
  guidance_scale: 3.5  # Guidance scale for sampling

  # Batch and sampling
  per_device_batch_size: 8  # Batch size per device
  group_size: 16  # Group size for GRPO sampling
  global_std: true  # Use global std for advantage normalization
  unique_sample_num_per_epoch: 48  # Unique samples per group
  gradient_step_per_epoch: 1  # Gradient steps per epoch. The first step is on-policy, the rest are off-policy.
    
  # Optimization
  learning_rate: 3.0e-4  # Initial learning rate
  adam_weight_decay: 1.0e-4  # AdamW weight decay
  adam_betas: [0.9, 0.999]  # AdamW betas
  adam_epsilon: 1.0e-8  # AdamW epsilon
  max_grad_norm: 1.0  # Max gradient norm for clipping

  # Gradient checkpointing
  enable_gradient_checkpointing: false  # Enable gradient checkpointing to save memory with extra compute

  # Seed
  seed: 42  # Random seed

# Scheduler Configuration
scheduler:
  dynamics_type: "ODE"  # Options: Flow-SDE, Dance-SDE, CPS, ODE

# Evaluation settings
eval:
  resolution: 1024  # Evaluation resolution
  per_device_batch_size: 1  # Eval batch size
  guidance_scale: 3.5  # Guidance scale for sampling
  num_inference_steps: 20  # Number of eval timesteps
  eval_freq: 20  # Eval frequency in epochs (0 to disable)
  seed: 42  # Eval seed (defaults to training seed)

# Reward Model Configuration
rewards:
  - name: "pick_score"
    reward_model: "PickScore"
    batch_size: 16
    device: "cuda"
    dtype: bfloat16

# Optional Evaluation Reward Models
eval_rewards:
  - name: "pick_score"
    reward_model: "PickScore"
    batch_size: 32
    device: "cuda"
    dtype: bfloat16

The resulting curve is:

[Jayce-Ping]

I didn't fully and successfully reproduced AWM on Flux1-dev either. For SD3.5, it is close but not that fast as reported in the paper, same for the official implementation.

From my experience, if you want to make reward at least increase, you can set KL-beta as 0 to allow more reward increase, and then use ema_kl_beta=1 to stablize the training.

[yansenhan]

I didn't fully and successfully reproduced AWM on Flux1-dev either. For SD3.5, it is close but not that fast as reported in the paper, same for the official implementation.

From my experience, if you want to make reward at least increase, you can set KL-beta as 0 to allow more reward increase, and then use ema_kl_beta=1 to stablize the training.

Thanks for your quick and helpful response. I will follow your advice.
If we happens to reproduce AWM for Flux.1-Dev, I will make a pull request.
Besides AMW, I also plan to reimplement DiffusionNFT and FlowGRPO for Flux.1-Dev. Do these two methods also have the same problem?

[Jayce-Ping]

Yes. Any PRs for new configs or code bug fix are welcomed.

For DiffusionNFT and Flow-GRPO, the provided config should be able to reproduce the results. They are widely used in the community and regarded as real SOTA methods.

[Jayce-Ping]

BTW, I found a AWM config for FLUX.1-dev in my wandb corner, which provide the following curves:

The config is:

{
  "log": {
    "value": {
      "verbose": true,
      "save_dir": "Flow-Factory",
      "save_freq": 0,
      "save_model_only": true
    }
  },
  "data": {
    "value": {
      "cache_dir": "~/.cache/flow_factory/datasets",
      "image_dir": null,
      "video_dir": null,
      "dataset_dir": "dataset/pickscore",
      "force_reprocess": false,
      "max_dataset_size": 1024,
      "enable_preprocess": true,
      "dataloader_num_workers": 16,
      "preprocessing_batch_size": 8
    }
  },
  "eval": {
    "value": {
      "seed": 42,
      "width": 512,
      "height": 512,
      "eval_freq": 20,
      "resolution": [
        512,
        512
      ],
      "guidance_scale": 3.5,
      "num_inference_steps": 20,
      "per_device_batch_size": 1
    }
  },
  "model": {
    "value": {
      "lora_rank": 64,
      "lora_alpha": 128,
      "model_type": "flux1",
      "resume_path": null,
      "resume_type": null,
      "attn_backend": "_flash_3_hub",
      "finetune_type": "lora",
      "target_modules": "default",
      "target_components": [
        "transformer"
      ],
      "model_name_or_path": "black-forest-labs/FLUX.1-dev",
      "master_weight_dtype": "bfloat16"
    }
  },
      },
      "m": [],
      "t": {
        "1": [
          1,
          11,
          41,
          49,
          51,
          71,
          83,
          98
        ],
        "2": [
          1,
          11,
          41,
          49,
          51,
          71,
          83,
          98
        ],
        "3": [
          13,
          16,
          61
        ],
        "4": "3.12.12",
        "5": "0.24.0",
        "6": "4.57.6",
        "12": "0.24.0",
        "13": "linux-x86_64"
      },
      "cli_version": "0.24.0",
      "python_version": "3.12.12"
    }
  },
  "reward": {
    "value": {
      "reward_0": {
        "name": "pick_score",
        "dtype": "bfloat16",
        "device": "cuda",
        "weight": 1,
        "batch_size": 16,
        "reward_model": "PickScore"
      }
    }
  },
  "project": {
    "value": "Workshop-Flow-Factory"
  },
  "launcher": {
    "value": "accelerate"
  },
  "run_name": {
    "value": "flux1_lora_awm_20260206_104150"
  },
  "training": {
    "value": {
      "seed": 42,
      "width": 512,
      "height": 512,
      "kl_beta": 0,
      "kl_type": "v-based",
      "nft_beta": 1,
      "ema_decay": 0.3,
      "kl_weight": "Uniform",
      "adam_betas": [
        0.9,
        0.999
      ],
      "clip_range": [
        -0.00001,
        0.00001
      ],
      "ema_device": "cuda",
      "global_std": true,
      "group_size": 16,
      "off_policy": false,
      "resolution": [
        512,
        512
      ],
      "time_shift": 3,
      "ema_kl_beta": 0.1,
      "adam_epsilon": 1e-8,
      "ghuber_power": 0.25,
      "trainer_type": "awm",
      "warmup_steps": 300,
      "awm_weighting": "ghuber",
      "learning_rate": 0.0003,
      "max_grad_norm": 1,
      "adv_clip_range": [
        -5,
        5
      ],
      "guidance_scale": 3.5,
      "timestep_range": [
        0,
        0.7
      ],
      "ref_param_device": "cuda",
      "adam_weight_decay": 0.0001,
      "ema_decay_schedule": "linear",
      "ema_update_interval": 1,
      "num_inference_steps": 10,
      "num_train_timesteps": 4,
      "advantage_aggregation": "sum",
      "num_batches_per_epoch": 48,
      "per_device_batch_size": 2,
      "time_sampling_strategy": "discrete",
      "gradient_step_per_epoch": 1,
      "gradient_accumulation_steps": 192,
      "unique_sample_num_per_epoch": 48,
      "enable_gradient_checkpointing": false
    }
  },
  "scheduler": {
    "value": {
      "seed": 42,
      "sde_steps": null,
      "noise_level": 0.7,
      "dynamics_type": "ODE",
      "num_sde_steps": 1
    }
  },
  "config_file": {
    "value": "config/deepspeed/deepspeed_zero2.yaml"
  },
  "eval_reward": {
    "value": {
      "reward_0": {
        "name": "pick_score",
        "dtype": "bfloat16",
        "device": "cuda",
        "weight": 1,
        "batch_size": 32,
        "reward_model": "PickScore"
      }
    }
  },
  "num_processes": {
    "value": 8
  },
  "logging_backend": {
    "value": "wandb"
  },
  "mixed_precision": {
    "value": "bf16"
  },
  "main_process_port": {
    "value": 29500
  }
}

where some personal info is deleted. The commit hash is 97eedd6f861836923e642b4277037b79f8f1a4bb. Hope it helps.

[yansenhan]

Thanks for your help. When we use git show 97eedd6f861836923e642b4277037b79f8f1a4bb, it cannot found the code corresponding to this hash.

[Jayce-Ping]

I see. Maybe it belongs to a branch that has been deleted. Have your tried to run the current example setting? I think it can provide correct results.

[yansenhan]

I have tried and still failed to obtain an expected result.

[Jayce-Ping]

That's weird. Do you have any curves for more information?

[yansenhan]

 {
  "launcher": "accelerate",
  "config_file": "deepspeed/deepspeed_zero2.yaml",
  "num_processes": 8,
  "main_process_port": 29500,
  "mixed_precision": "bf16",
  "data": {
    "cache_dir": "cache/datasets",
    "image_dir": null,
    "video_dir": null,
    "dataset_dir": "dataset/pickscore",
    "sampler_type": "auto",
    "force_reprocess": false,
    "max_dataset_size": 1024,
    "enable_preprocess": true,
    "dataloader_num_workers": 16,
    "preprocess_parallelism": "local",
    "preprocessing_batch_size": 8
  },
  "model": {
    "lora_rank": 64,
    "lora_alpha": 128,
    "model_type": "flux1",
    "resume_path": null,
    "resume_type": null,
    "attn_backend": null,
    "finetune_type": "lora",
    "target_modules": "default",
    "target_components": [
      "transformer"
    ],
    "model_name_or_path": "../models/FLUX.1-dev",
    "master_weight_dtype": "bfloat16"
  },
  "scheduler": {
    "seed": 42,
    "sde_steps": [0, 1, 2, 3, 4, 5, 6, 7, 8],
    "noise_level": 0.7,
    "dynamics_type": "ODE",
    "num_sde_steps": 1
  },
  "training": {
    "seed": 42,
    "width": 512,
    "height": 512,
    "kl_beta": 0,
    "kl_type": "v-based",
    "ema_decay": 0.3,
    "kl_weight": "Uniform",
    "adam_betas": [0.9, 0.999],
    "clip_range": [-0.00001, 0.00001],
    "ema_device": "cuda",
    "global_std": true,
    "group_size": 16,
    "max_epochs": null,
    "off_policy": false,
    "resolution": [512, 512],
    "time_shift": 3,
    "ema_kl_beta": 0.1,
    "adam_epsilon": 1e-8,
    "ghuber_power": 0.25,
    "trainer_type": "awm",
    "warmup_steps": 300,
    "awm_weighting": "ghuber",
    "learning_rate": 0.0003,
    "max_grad_norm": 1,
    "adv_clip_range": [-5, 5],
    "guidance_scale": 3.5,
    "timestep_range": [0, 0.7],
    "num_inner_epochs": 1,
    "ref_param_device": "cuda",
    "adam_weight_decay": 0.0001,
    "ema_decay_schedule": "linear",
    "ema_update_interval": 1,
    "num_inference_steps": 10,
    "num_train_timesteps": 4,
    "latent_storage_dtype": "fp16",
    "advantage_aggregation": "sum",
    "num_batches_per_epoch": 12,
    "per_device_batch_size": 8,
    "time_sampling_strategy": "discrete",
    "gradient_step_per_epoch": 1,
    "gradient_accumulation_steps": 48,
    "unique_sample_num_per_epoch": 48,
    "enable_gradient_checkpointing": false
  },
  "eval": {
    "seed": 42,
    "width": 512,
    "height": 512,
    "eval_freq": 20,
    "resolution": [512, 512],
    "guidance_scale": 3.5,
    "num_inference_steps": 20,
    "per_device_batch_size": 1
  },
  "log": {
    "verbose": true,
    "run_name": "flux1_lora_awm_20260415_161245",
    "save_dir": "saves/",
    "save_freq": 20,
    "logging_backend": "swanlab",
    "save_model_only": true,
    "eval_samples_local_dir": null,
    "save_eval_samples_local": false,
    "eval_samples_local_limit": 16,
    "log_eval_samples_to_backend": true
  },
  "reward": {
    "reward_0": {
      "name": "pick_score",
      "dtype": "bfloat16",
      "device": "cuda",
      "weight": 1,
      "batch_size": 16,
      "num_workers": 1,
      "async_reward": false,
      "reward_model": "PickScore",
      "model_name_or_path": "../PickScore_v1",
      "processor_name_or_path": "../models/CLIP-ViT-H-14-laion2B-s32B-b79K"
    }
  },
  "eval_reward": {
    "reward_0": {
      "name": "pick_score",
      "dtype": "bfloat16",
      "device": "cuda",
      "weight": 1,
      "batch_size": 32,
      "num_workers": 1,
      "async_reward": false,
      "reward_model": "PickScore",
      "model_name_or_path": "../models/PickScore_v1",
    }
  }
}

[Jayce-Ping]

clip_range is set 1 in official setting. And more importantly, adv_clip_range should be set 1. The AWM algorithm is very sensitive to it and I believe the spike in the training curve is caused by this.

[yansenhan]

clip_range is set 1 in official setting. And more importantly, adv_clip_range should be set 1. The AWM algorithm is very sensitive to it and I believe the spike in the training curve is caused by this.

Thanks very much for your help. I will try it again. But I find in your config, these two hyperparameters are not 1.