WAM-Diff2/configs/training/block32.yaml at master · fudan-generative-vision/WAM-Diff2 · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
step_scheduler:
  global_batch_size: 128
  local_batch_size: 4
  ckpt_every_steps: 100
  num_epochs: 5

dist_env:
  backend: nccl
  timeout_minutes: 20

rng:
  _target_: wam_diff.components.training.rng.StatefulRNG
  seed: 42
  ranked: true

model:
  _target_: wam_diff.NeMoAutoModelForImageTextToText.from_pretrained
  pretrained_model_name_or_path: checkpoints/WAM-Diff2-B16-2B
  attn_implementation: sdpa
  torch_dtype: torch.bfloat16

processor:
  _target_: transformers.AutoProcessor.from_pretrained
  pretrained_model_name_or_path: checkpoints/WAM-Diff2-B16-2B
  min_pixels: 12544
  max_pixels: 2073600

checkpoint:
  enabled: true
  checkpoint_dir: outputs/WAM-Diff2-B32-2B
  model_save_format: torch_save

distributed:
  _target_: wam_diff.components.distributed.fsdp2.FSDP2Manager
  dp_size: null
  tp_size: 1
  cp_size: 1
  sequence_parallel: false
  activation_checkpointing: true

loss_fn:
  _target_: wam_diff.losses.weighted_ce.WeightedCrossEntropy

dataset:
  _target_: wam_diff.data.datasets.qwen_vl_nav_dataset
  path_or_dataset: data/train.json
  max_len: 8192
  drop_last: true
  root_dir: data
  prior_dist: Mask
  min_pixels: 12544
  max_pixels: 2073600

dataloader:
  _target_: torchdata.stateful_dataloader.StatefulDataLoader
  num_workers: 8
  pin_memory: true
  prefetch_factor: 4
  collate_fn:
    _target_: wam_diff.data.collate_fns.wam_diff2_block_collate_fn
    noise_scheduler:
      _target_: wam_diff.utils.scheduler.CondOTScheduler
    model_type: qwen3-vl
    mask_token_id: 151671
    vocab_size: 151646
    pad_token: <|im_end|>
    im_end_token: <|im_end|>
    block_size: 32

optimizer:
  _target_: torch.optim.AdamW
  lr: 4.0e-5
  merger_lr: 4.0e-5
  visual_lr: 2.0e-6
  weight_decay: 0.05
  betas: [0.9, 0.95]

lr_scheduler:
  lr_decay_style: cosine

freeze_config:
  freeze_embeddings: false
  freeze_vision_tower: false
  freeze_language_model: false