# Example configuration for end-to-end multimodal supervised fine-tuning
# Intended training path: pretrain.yaml -> sft.yaml -> vlm.yaml -> vlm_sft.yaml

model:
  architecture_type: multimodal_wrapper
  llm_architecture_type: taonet
  vocab_size: 8192
  hidden_dim: 768
  num_layers: 10
  num_heads: 8
  max_seq_length: 1024
  d_latent_kv: 512
  d_rope: 64
  hidden_dim_ff: 2048
  dropout: 0.02
  gqa_groups: 1
  use_factorized_embedding: false
  d_embed_rank: 96
  init_std: 0.02

  vision_encoder_type: cnn
  vision_output_dim: 256
  image_size: 224
  vision_prefix_tokens: 10
  image_token: <image>
  cnn_channels: [32, 64, 128]
  cnn_kernel_size: 3

dataset:
  local: true
  jsonl_path: /home/student/Data/TaoData/vision_sft.jsonl
  image_path_column: image
  image_path_aliases: [images, image_path, image_file, file_name]
  caption_prompt: Describe the image.
  samples_per_chunk: 1000
  max_samples: 160000
  cache_dir: .cache/datasets

  tokenizer_type: sentencepiece
  tokenizer_path: tokenizer/tokenizer.model
  tokenizer_threads: 4

checkpoint_path: checkpoints/vlm/final_model.pt
user_token: <user>
assistant_token: <assistant>
response_loss_only: true

freeze_llm: false
unfreeze_last_n_layers: 0
vision_learning_rate: 5e-5
llm_learning_rate: 5e-5
vision_prefix_tokens: 10
image_token: <image>
image_size: 224

batch_size: 8
num_epochs: 1
gradient_accumulation_steps: 4
max_grad_norm: 1.0

optimizer:
  optimizer_type: adamw
  learning_rate: 5e-5
  weight_decay: 0.01

scheduler:
  scheduler_type: linearWarmup
  warmup_steps: 500

dtype: bfloat16
device: cuda

checkpoint_dir: checkpoints/vlm_sft
save_every_steps: 81920
save_best_model: true
keep_last_n_checkpoints: 2

eval_every_steps: 8192
eval_samples: 200

log_every_steps: 10
aim_repo: .aim

seed: 42
num_workers: 0
pin_memory: true