[NewModel] Add Bagel pretrain

examples/bagel/conf/train.yaml

@@ -0,0 +1,37 @@
+defaults:
+  - train: 7b
+  - _self_
+
+experiment:
+  exp_name: bagel
+  exp_dir: outputs
+  task:
+    type: train
+    backend: megatron
+    entrypoint: ./flagscale/train/megatron/train_bagel.py
+  runner:
+    nnodes: 1
+    nproc_per_node: 8
+    rdzv_backend: static
+    hostfile: /share/project/zhaoyingli/hostfiles/hostfile.1
+  cmds:
+    before_start: "ulimit -n 1048576 && source /root/miniconda3/bin/activate /share/project/zhaoyingli/envs/bagel-fs-train"
+  envs:
+    CUDA_LAUNCH_BLOCKING: 1
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    CUDA_DEVICE_MAX_CONNECTIONS: 4
+    NVTE_APPLY_QK_LAYER_SCALING: 0
+    # NCCL_DEBUG: INFO
+    GLOO_SOCKET_IFNAME: eth0
+    NCCL_SOCKET_IFNAME: eth0
+    NCCL_IB_DISABLE: 0
+    NCCL_IB_GID_INDEX: 7
+    NCCL_IB_HCA: mlx5_100,mlx5_101,mlx5_102,mlx5_103,mlx5_104,mlx5_105,mlx5_106,mlx5_107
+    NCCL_IB_TIMEOUT: 22
+    NCCL_IB_RETRY_CNT: 13
+
+action: run
+
+hydra:
+  run:
+    dir: ${experiment.exp_dir}/hydra

examples/bagel/conf/train/7b.yaml

@@ -0,0 +1,152 @@
+system:
+  tensor_model_parallel_size: 1
+  pipeline_model_parallel_size: 1
+  use_distributed_optimizer: true
+  use_precision_aware_optimizer: true
+  overlap_grad_reduce: false
+  overlap_param_gather: false
+  use_megatron_fsdp: true
+  overlap_param_gather_with_optimizer_step: false
+  megatron_fsdp_main_params_dtype: fp32
+  megatron_fsdp_main_grads_dtype: bf16
+  megatron_fsdp_grad_comm_dtype: bf16
+  data_parallel_sharding_strategy: optim_grads_params # optim_grads, optim_grads_params
+
+  # full recompute
+  recompute_method: "uniform"
+  recompute_granularity: "full"
+  recompute_num_layers: 1
+
+  # # block recompute
+  # recompute_method: "block"
+  # recompute_granularity: "full"
+  # recompute_num_layers: 1
+
+  # # core recompute
+  # recompute_granularity: "selective"
+
+  no_gradient_accumulation_fusion: true
+  precision:
+    bf16: true
+    attention_softmax_in_fp32: true
+    accumulate_allreduce_grads_in_fp32: false
+    main_grads_dtype: bf16
+    main_params_dtype: fp32
+  logging:
+    log_interval: 1
+    tensorboard_log_interval: 1
+    # wandb_project: bagel_flagscale
+    # wandb_exp_name: bagel_7b_mot_8gpu
+  checkpoint:
+    save_interval: 5000000
+    load: null #/workspace/new_agent_v2/models/BAGEL-7B-MoT-megatron-tp2-ema
+    no_load_optim: true
+    no_load_rng: true
+    # ckpt_format: torch
+    ckpt_format: fsdp_dtensor
+
+model:
+  # LLM Architecture
+  language_model_type: qwen2.5_7B
+  transformer_impl: transformer_engine
+  num_layers: 28
+  hidden_size: 3584
+  ffn_hidden_size: 18944
+  num_attention_heads: 28
+  num_query_groups: 4
+  group_query_attention: true
+  seq_length: 30720
+  norm_epsilon: 1e-6
+  max_position_embeddings: 32768
+  position_embedding_type: rope
+  rotary_base: 1000000
+  swiglu: true
+  normalization: RMSNorm
+  init_method_std: 0.02
+  attention_dropout: 0.0
+  hidden_dropout: 0.0
+  untie_embeddings_and_output_weights: true
+  no_rope_fusion: true
+  qk_layernorm: true
+  disable_bias_linear: false
+  attention_backend: auto
+
+  # freeze model
+  freeze_LLM: false
+  freeze_ViT: true
+  freeze_VAE: true
+  freeze_text_embed: true
+
+  # freeze branch
+  freeze_und: false
+  freeze_gen: false
+
+  # vit
+  vision_model_type: siglip
+  patch_dim: 14
+  image_size: 980
+
+  # packing
+  max_num_tokens: 30720
+  max_num_tokens_per_sample: 30720
+  expected_num_tokens: 30720
+
+  # vit_hidden_size: 1152
+  # vit_ffn_hidden_size: 4304
+  # vit_num_layers: 27
+  # vit_num_attention_heads: 16
+  # vit_image_size: 980
+  # vit_patch_size: 14
+  # vit_num_channels: 3
+
+  # # VAE
+  # vae_z_channels: 16
+  # vae_downsample: 8
+
+  # Bagle extra
+  max_latent_size: 64
+  latent_patch_size: 2
+  vit_max_num_patch_per_side: 70
+  interpolate_pos: false
+  timestep_shift: 1.0
+  visual_und: true
+  visual_gen: true
+
+  # training
+  finetune: false
+  train_iters: 3000
+  eval_iters: 0
+  eval_interval: 999999
+  micro_batch_size: 1
+  global_batch_size: 8
+  ce_weight: 0.25
+  mse_weight: 1.0
+
+  optimizer:
+    lr: 1e-4
+    # min_lr: 1e-7
+    # weight_decay: 0.0
+    # lr_warmup_iters: 2500
+    # lr_decay_style: constant
+    adam_beta1: 0.9
+    adam_beta2: 0.95
+    adam_eps: 1e-15
+    clip_grad: 1.0
+
+data:
+  dataloader_type: external
+  # data_path: /share/project/zhaoyingli/dataset/bagel_example/vlm/wds/
+  data_path: /share/project/zhaoyingli/bagel-fs/FlagScale/examples/bagel/data/pretrain_dataset.yaml
+  split: "1,0,0"
+  num_workers: 2
+  packing_buffer_size: 24
+  tokenizer:
+    tokenizer_type: BagelTokenizerFS
+    tokenizer_path: /share/project/zhaoyingli/checkpoints/ByteDance-Seed/BAGEL-7B-MoT/
+    vocab_size: 152064
+    make_vocab_size_divisible_by: 64
+    tokenizer_prompt_format: qwen2p5
+
+  # bagel_model_path: /workspace/new_agent_v2/models/BAGEL-7B-MoT
+  # bagel_src_dir: /workspace/new_agent_v2/BAGEL
+  # dataset_config_file: /workspace/new_agent_v2/BAGEL/data/configs/example.yaml

examples/bagel/data/pretrain_dataset.yaml

@@ -0,0 +1,31 @@
+__module__: megatron.energon
+__class__: MetadatasetV2
+splits:
+  train:
+    blend:
+      - weight: 1.
+        path: /share/project/zhaoyingli/dataset/bagel_example/vlm/wds/
+        subflavors:
+          task: vlm_sft
+          image_transform_args:
+            image_stride: 14
+            max_image_size: 980
+            min_image_size: 379
+            max_pixels: 2_007_040
+            is_mandatory: true
+
+      - weight: 1.
+        path: /share/project/zhaoyingli/dataset/bagel_example/t2i/wds/
+        subflavors:
+          task: t2i
+          image_transform_args:
+            image_stride: 16
+            max_image_size: 1024
+            min_image_size: 512
+            is_mandatory: true
+  # val:
+  #   datasets:
+  #     - weight: 1.
+  #       path: <path_to_pretraining_dataset_in_energon_format>
+  #       subflavors:
+  #         augmentation: false