Skip to content

AnyFlow

O SimpleTuner implementa o NVIDIA AnyFlow como duas etapas explícitas de treinamento para modelos de flow matching. As duas etapas treinam um modelo que recebe o tempo de fluxo atual t e um endpoint de intervalo r.

  • stage=forward implementa o objetivo forward MeanFlow da NVIDIA.
  • stage=onpolicy implementa Flow Map Backward Simulation e DMD on-policy enquanto co-treina o objetivo forward.

Os modos removidos online_teacher e linear eram objetivos específicos do SimpleTuner e não são mais aceitos.

Para um exemplo de continuação Wan usando os checkpoints publicados pela NVIDIA, veja Quickstart de continuação AnyFlow.

Qwen Image 2.1 suporta condicionamento de intervalos FlowMap. O piloto Qwen de três etapas mistura captions longas do CC12M, e621 e poucos exemplos de Domokun, seguidos de um ajuste breve com regularização. A âncora na previsão da base testa preservação; não comprova que Qwen recebeu destilação de guidance.

O treinamento AnyFlow compilado permite temporariamente pelo menos 32 variantes Dynamo por frame para professor, gerador, discriminador, modos de gradiente, padding e batches diferentes. Limites maiores do usuário são mantidos, e o limite anterior é restaurado ao terminar. No Torch 2.11, esgotar as oito variantes padrão pode mudar a recomputação do checkpoint para execução eager e falhar na verificação dos metadados dos tensores salvos. Use dynamo_dynamic: true para captions de comprimento variável; o piloto Qwen já habilita essa opção.

Os checkpoints do AnyFlow incluem anyflow_rng_state_<rank>.pt para cada processo. Esse arquivo preserva os geradores aleatórios privados usados na amostragem de intervalos e nos rollouts on-policy. A retomada exige a mesma etapa, seed, configuração de dados e topologia distribuída. Checkpoints antigos sem esse estado não permitem continuação exata e são rejeitados; use init_lora com um novo otimizador para iniciar outra execução a partir do adapter.

init_lora usa por padrão o passo salvo no adapter. Ao iniciar uma nova etapa, defina init_lora_step: 0 para que o limite de passos e o agendador comecem do zero.

A exportação de adaptadores remove os wrappers regionais do compilador dos nomes dos tensores e preserva os tensores do embedding de intervalo. Modelos comuns e compilados podem recarregá-los. Com diffusion_ratio: 1.0, todos os intervalos têm r=t, portanto a preparação do alvo omite as duas previsões de diferenças finitas cuja contribuição é zero.

Etapa forward

{
  "model_type": "lora",
  "distillation_method": "anyflow",
  "distillation_config": {
    "anyflow": {
      "stage": "forward",
      "diffusion_ratio": 0.5,
      "consistency_ratio": 0.25,
      "central_difference_epsilon": 0.005,
      "meanflow_weight_type": "beta08",
      "meanflow_adaptive_weighting": true,
      "gate_value": 0.25,
      "deltatime_type": "r",
      "loss_weight": 1.0
    }
  }
}

Para cada batch global, a etapa forward:

  1. Amostra dois tempos de fluxo uniformes e os ordena como t >= r.
  2. Atribui 50% das amostras a intervalos de difusão (r=t), 25% a intervalos de endpoint (r=0) e o restante a intervalos arbitrários.
  3. Aplica o flow shift do scheduler do modelo aos dois endpoints.
  4. Avalia uma diferença central ao longo da trajetória latente reta.
  5. Constrói o target tangente MeanFlow e aplica o weighting normalizado beta08 da NVIDIA.
  6. Balanceia cada amostra não-diffusion contra a média global de loss da ramificação diffusion.

Etapa on-policy

Inicie esta etapa com um adapter AnyFlow forward usando init_lora e um novo estado do otimizador. Retome checkpoints de estado de treinamento apenas na mesma etapa, configuração de dataset e topologia distribuída:

{
  "model_type": "lora",
  "lora_type": "standard",
  "init_lora": "path-or-repo-to-forward-anyflow-adapter",
  "init_lora_step": 0,
  "learning_rate": 0.000002,
  "optimizer_beta1": 0.0,
  "optimizer_beta2": 0.999,
  "optimizer_config": "weight_decay=0.0",
  "distillation_method": "anyflow",
  "distillation_config": {
    "anyflow": {
      "stage": "onpolicy",
      "cotrain_forward": true,
      "rollout_step_counts": [2, 4, 8, 16, 50],
      "dmd_weight": 1.0,
      "dmd_batch_size": 1,
      "real_score_guidance_scale": 0.0,
      "discriminator_lr": 0.000002,
      "discriminator_betas": [0.0, 0.999],
      "discriminator_weight_decay": 0.0,
      "discriminator_grad_clip": 1.0
    }
  }
}

A etapa on-policy usa três papéis de score. O treinamento LoRA padrão compartilha um transformer base congelado entre eles:

  • O adapter AnyFlow carregado é o gerador.
  • O modelo base com adapters desativados é o score real congelado.
  • Um adapter anyflow_discriminator otimizado separadamente é o score fake.

Cada atualização do gerador escolhe um orçamento de rollout de rollout_step_counts, executa um rollout FlowMap diferenciável, adiciona ruído ao latent gerado em um tempo uniforme deslocado e aplica o gradiente DMD normalizado da NVIDIA. Cada atualização do discriminador executa um rollout do student sem gradientes, amostra um tempo deslocado logit-normal e treina o score fake no target flow normal. O adapter e o otimizador do discriminador são salvos ao lado de cada checkpoint do SimpleTuner como anyflow_discriminator.safetensors e anyflow_discriminator_optim.pt.

MiniMax-H3 já contém destilação CFG, então suas execuções on-policy normalmente devem manter real_score_guidance_scale=0. Modelos que exigem uma passada CFG externa para o score real precisam cachear embeddings de texto negativos e podem configurar a escala explicitamente.

Quando --seed está definido, o AnyFlow amostra intervalos MeanFlow, schedules de rollout, latentes de rollout, ruído DMD e sigmas DMD a partir de um gerador Torch isolado por dispositivo. Isso mantém as amostras AnyFlow estáveis quando código de treinamento não relacionado consome o RNG global do Torch. Isso não torna o backward de attention em CUDA bit-estável.

Configuração compartilhada

  • stage: forward ou onpolicy. Padrão: forward.
  • diffusion_ratio: fração do batch global usando r=t. Padrão: 0.5.
  • consistency_ratio: fração do batch global usando r=0. Padrão: 0.25.
  • central_difference_epsilon: offset normalizado no tempo deslocado. Padrão: 0.005, igual ao 5/1000 da NVIDIA.
  • meanflow_weight_type: beta08 ou uniform. Padrão: beta08.
  • meanflow_adaptive_weighting: balanceia amostras não-diffusion contra a ramificação diffusion. Padrão: true.
  • gate_value: mistura do embedding delta-timestep FlowMap. Padrão: 0.25.
  • deltatime_type: r ou t-r. Padrão: r.
  • loss_weight: multiplicador da loss forward MeanFlow. Padrão: 1.0.

Limites

  • AnyFlow requer um modelo flow-matching com conditioning de intervalo FlowMap específico do modelo.
  • O treinamento on-policy atualmente requer LoRA PEFT padrão. Compartilhar a base evita alocar cópias do gerador, score real e discriminador de um transformer grande em cada rank DDP.
  • Treinamento conjunto MiniMax-H3 audio-video é rejeitado. Video usa schedule shift 12 e audio usa shift 3; targets MeanFlow e rollouts nativos de duplo schedule precisam ser implementados antes que treinamento AV seja válido.
  • Treinamento do text encoder é desativado para todos os métodos de destilação do SimpleTuner.
  • A validação usa AnyFlowValidationScheduler, que fornece o próximo endpoint de intervalo aos componentes FlowMap registrados.

Logs

O treinamento forward adiciona anyflow_forward_loss, valores de timestep e intervalo, e frações globais de ramificação. O treinamento on-policy também adiciona anyflow_dmd_loss, anyflow_dmd_gradient_norm, anyflow_dmd_sigma e anyflow_rollout_steps.