diff --git a/olmocr/train/configs/v0.5.0/qwen3_vl_olmocrv5_1epoch_mix_1025_filtered.yaml b/olmocr/train/configs/v0.5.0/qwen3_vl_olmocrv5_1epoch_mix_1025_filtered.yaml new file mode 100644 index 0000000..a8ef965 --- /dev/null +++ b/olmocr/train/configs/v0.5.0/qwen3_vl_olmocrv5_1epoch_mix_1025_filtered.yaml @@ -0,0 +1,117 @@ +# Project metadata +project_name: olmocr-qwen-vl-training +run_name: qwen3_vl_olmocrv5_1epoch_mix_1025_filtered + +# Model configuration +model: + name: Qwen/Qwen3-VL-8B-Instruct + trust_remote_code: true + torch_dtype: bfloat16 + use_flash_attention: true + attn_implementation: flash_attention_2 + + # LoRA settings (disabled by default) + use_lora: false + # lora_rank: 8 + # lora_alpha: 32 + # lora_dropout: 0.1 + # lora_target_modules: + # - q_proj + # - v_proj + # - k_proj + # - o_proj + +# Dataset configuration +dataset: + + train: + - name: processed_01_books_train_iabooks + root_dir: /data/olmOCR-mix-1025/processed_01_books_train_iabooks/ + pipeline: &basic_pipeline + - name: FrontMatterParser + front_matter_class: PageResponse + - name: FilterOutRotatedDocuments + - name: ReformatLatexBoldItalic + - name: DatasetTextRuleFilter + - name: PDFRenderer + target_longest_image_dim: 1288 + - name: RotationAugmentation + probability: 0.02 + - name: NewYamlFinetuningPromptWithNoAnchoring + - name: FrontMatterOutputFormat + - name: InstructUserMessages + prompt_first: true + - name: Tokenizer + masking_index: -100 + end_of_message_token: "<|im_end|>" + - name: processed_00_documents_train_s2pdf + root_dir: /data/olmOCR-mix-1025/processed_00_documents_train_s2pdf/ + pipeline: *basic_pipeline + - name: processed_02_loc_transcripts + root_dir: /data/olmOCR-mix-1025/processed_02_loc_transcripts/ + pipeline: *basic_pipeline + - name: processed_03_national_archives + root_dir: /data/olmOCR-mix-1025/processed_03_national_archives/ + pipeline: *basic_pipeline + + eval: + - name: processed_00_documents_eval_s2pdf + root_dir: /data/olmOCR-mix-1025/processed_00_documents_eval_s2pdf/ + pipeline: *basic_pipeline + - name: processed_01_books_eval_iabooks + root_dir: /data/olmOCR-mix-1025/processed_01_books_eval_iabooks/ + pipeline: *basic_pipeline + - name: processed_02_loc_transcripts_eval + root_dir: /data/olmOCR-mix-1025/processed_02_loc_transcripts_eval/ + pipeline: *basic_pipeline + - name: processed_03_national_archives_eval + root_dir: /data/olmOCR-mix-1025/processed_03_national_archives_eval/ + pipeline: *basic_pipeline + + +# Training configuration +training: + output_dir: /weka/oe-training-default/jakep/olmocr-trainer/ + num_train_epochs: 1 + + # Batch size and accumulation + per_device_train_batch_size: 1 + per_device_eval_batch_size: 1 + gradient_accumulation_steps: 32 + + gradient_checkpointing: False + + collator_max_token_len: 8192 + + # Learning rate + learning_rate: 2e-5 + lr_scheduler_type: linear + warmup_ratio: 0.1 + + # Optimization + optim: adamw_torch + weight_decay: 0.01 + max_grad_norm: 1.0 + + # Torch compile settings + torch_compile: true + torch_compile_backend: inductor + torch_compile_mode: default + torch_compile_fullgraph: false + torch_compile_dynamic: false + + seed: 300 + data_seed: 301 + + # Evaluation and checkpointing + evaluation_strategy: steps + eval_steps: 500 + save_strategy: steps + save_steps: 500 + save_total_limit: 5 + load_best_model_at_end: false # Needs to be false because it has a problem restoring checkpoints for some reason + metric_for_best_model: eval_processed_00_documents_eval_s2pdf_loss + greater_is_better: false + + report_to: + - wandb \ No newline at end of file diff --git a/olmocr/train/train.py b/olmocr/train/train.py index 23f510f..6f579b7 100644 --- a/olmocr/train/train.py +++ b/olmocr/train/train.py @@ -18,6 +18,7 @@ from torch.utils.data import ConcatDataset, DataLoader from tqdm import tqdm from transformers import ( AutoProcessor, + Qwen3VLForConditionalGeneration, Qwen2_5_VLForConditionalGeneration, Qwen2VLForConditionalGeneration, get_scheduler, @@ -271,7 +272,10 @@ def main(): # Load model logger.info(f"Loading model: {config.model.name}") - if "qwen2.5-vl" in config.model.name.lower(): + if "qwen3-vl" in config.model.name.lower(): + model_class = Qwen3VLForConditionalGeneration + model = model_class.from_pretrained(config.model.name, **model_init_kwargs) + elif "qwen2.5-vl" in config.model.name.lower(): model_class = Qwen2_5_VLForConditionalGeneration model = model_class.from_pretrained(config.model.name, **model_init_kwargs) elif "qwen2-vl" in config.model.name.lower():