Add Bagel Trainer and fix config, bagel data processor (#126)

* fix bagel trainer config processor

* fix bagel trainer config processor

* Add bagel fsdp trainer into register and in ema support trainer list

---------

Co-authored-by: kcz358 <kaichenzhang358@outlook.com>

Author: KemingWu

examples/bagel/edit_example_config.yaml

@@ -0,0 +1,215 @@
+trainer_type: bagel_fsdp2_trainer
+
+dataset_config:
+  extra_kwargs: {}
+  dataset_type: bagel_iterable
+  dataset_format: yaml
+  processor_config:
+    processor_name: lmms-lab/BAGEL-7B-MoT-ver.LE
+    processor_type: bagel
+    # Align with original Bagel edit dataset:
+    # - user/source image is encoded as BOTH VAE-condition (loss=0, timestep=-inf) and ViT-condition
+    # - optional CFG-style conditional dropouts (set to original defaults; set to 0.0 to disable)
+    extra_kwargs:
+      user_image_as_vae_condition: true
+      text_cond_dropout_prob: 0.0
+      vit_cond_dropout_prob: 0.0
+      vae_cond_dropout_prob: 0.0
+  # Update this to your dataset path
+  # Example dataset: https://huggingface.co/datasets/kcz358/bagel-example
+  datasets:
+    - path: /path/to/your/dataset/
+      data_folder: /path/to/your/data_folder/
+      data_type: arrow
+  shuffle: true
+  eval_dataset_path: null
+  object_storage: none
+  bucket_name: null
+  packing: true
+  packing_strategy: first_fit
+  packing_length: 10240
+  filter_overlong: true
+  filter_overlong_workers: 8
+  max_length: null
+  video_backend: qwen_vl_utils
+
+trainer_args:
+  output_dir: ./output/bagel_training
+  overwrite_output_dir: false
+  do_train: true
+  do_eval: false
+  do_predict: false
+  eval_strategy: 'no'
+  prediction_loss_only: false
+  per_device_train_batch_size: 2
+  per_device_eval_batch_size: 8
+  gradient_accumulation_steps: 1
+  eval_accumulation_steps: null
+  eval_delay: 0
+  torch_empty_cache_steps: null
+  learning_rate: 1.0e-06
+  weight_decay: 0.0
+  adam_beta1: 0.9
+  adam_beta2: 0.999
+  adam_epsilon: 1.0e-08
+  max_grad_norm: 1.0
+  num_train_epochs: 1
+  max_steps: 1000
+  lr_scheduler_type: cosine
+  lr_scheduler_kwargs: {}
+  warmup_ratio: 0.1
+  warmup_steps: 20
+  log_level: passive
+  log_level_replica: warning
+  log_on_each_node: true
+  logging_dir: ./output/bagel_training/runs
+  logging_strategy: steps
+  logging_first_step: false
+  logging_steps: 10
+  logging_nan_inf_filter: true
+  save_strategy: steps
+  save_steps: 500
+  save_total_limit: 2
+  save_safetensors: true
+  save_on_each_node: false
+  save_only_model: false
+  restore_callback_states_from_checkpoint: false
+  no_cuda: false
+  use_cpu: false
+  use_mps_device: false
+  seed: 42
+  data_seed: null
+  jit_mode_eval: false
+  bf16: true
+  fp16: false
+  fp16_opt_level: O1
+  half_precision_backend: auto
+  bf16_full_eval: false
+  fp16_full_eval: false
+  tf32: null
+  local_rank: 0
+  ddp_backend: null
+  tpu_num_cores: null
+  tpu_metrics_debug: false
+  debug: []
+  dataloader_drop_last: false
+  eval_steps: null
+  dataloader_num_workers: 8
+  dataloader_prefetch_factor: null
+  past_index: -1
+  run_name: bagel_training
+  disable_tqdm: false
+  remove_unused_columns: true
+  label_names: null
+  load_best_model_at_end: false
+  metric_for_best_model: null
+  greater_is_better: null
+  ignore_data_skip: false
+  fsdp: []
+  fsdp_min_num_params: 0
+  fsdp_config:
+    transformer_layer_cls_to_wrap:
+    - Qwen2MoTDecoderLayer
+    reshard_after_forward: false
+    min_num_params: 0
+    xla: false
+    xla_fsdp_v2: false
+    xla_fsdp_grad_ckpt: false
+  fsdp_transformer_layer_cls_to_wrap: null
+  accelerator_config:
+    split_batches: false
+    dispatch_batches: null
+    even_batches: true
+    use_seedable_sampler: true
+    non_blocking: false
+    gradient_accumulation_kwargs: null
+  parallelism_config: null
+  deepspeed: null
+  label_smoothing_factor: 0.0
+  optim: adamw_torch_fused
+  optim_args: null
+  adafactor: false
+  group_by_length: true
+  length_column_name: length
+  report_to:
+  - wandb
+  project: huggingface
+  trackio_space_id: trackio
+  ddp_find_unused_parameters: null
+  ddp_bucket_cap_mb: null
+  ddp_broadcast_buffers: null
+  dataloader_pin_memory: true
+  dataloader_persistent_workers: false
+  skip_memory_metrics: true
+  use_legacy_prediction_loop: false
+  push_to_hub: false
+  resume_from_checkpoint: null
+  hub_model_id: null
+  hub_strategy: every_save
+  hub_token: <HUB_TOKEN>
+  hub_private_repo: null
+  hub_always_push: false
+  hub_revision: null
+  gradient_checkpointing: true
+  gradient_checkpointing_kwargs: null
+  include_inputs_for_metrics: false
+  include_for_metrics: []
+  eval_do_concat_batches: true
+  fp16_backend: auto
+  push_to_hub_model_id: null
+  push_to_hub_organization: null
+  mp_parameters: ''
+  auto_find_batch_size: false
+  full_determinism: false
+  torchdynamo: null
+  ray_scope: last
+  ddp_timeout: 1800
+  torch_compile: false
+  torch_compile_backend: null
+  torch_compile_mode: null
+  include_tokens_per_second: false
+  include_num_input_tokens_seen: 'no'
+  neftune_noise_alpha: null
+  optim_target_modules: null
+  batch_eval_metrics: false
+  eval_on_start: false
+  use_liger_kernel: false
+  liger_kernel_config: null
+  eval_use_gather_object: false
+  average_tokens_across_devices: true
+  use_muon: false
+  freeze_modules:
+  - vae_model
+  use_rmpad: false
+  fsdp2: true
+  sp_ulysses_degree: 1
+  reduce_dtype: bfloat16
+  output_dtype: bfloat16
+  print_batch_input_steps: 5
+  enable_profiler: false
+  profiler_config:
+    start_step: 1
+    end_step: 3
+
+model_config:
+  extra_kwargs:
+    visual_und: false  # Enable/disable visual understanding
+  load_from_pretrained_path: lmms-lab/BAGEL-7B-MoT-ver.LE
+  load_from_config: null
+  attn_implementation: flash_attention_2
+  model_type: null
+  torch_dtype: bfloat16
+  overwrite_config: null
+  # Optional: Enable Native Sparse Attention
+  # monkey_patch_kwargs:
+  #   patch_type: ["nsa"]
+  #   block_size: 64
+  #   compress_type: "weightedpool"
+  #   kernel_size: 32
+  #   kernel_stride: 16
+  #   topk: 16
+  #   init_blocks: 1
+  #   local_blocks: 2
+  #   window_size: 512
+
+extra_kwargs: null