feat(mm-cpt): broaden multimodal CPT dataset paths

docs/multimodal.qmd

@@ -368,19 +368,27 @@ Here is an example of a multi-modal dataset:
 
 Raw image+text continued pretraining — no chat template, no conversational
 scaffolding. The model learns to emit raw text conditioned on visual patches.
-Intended for use cases like OCR/transcription corpora where every row is a
-tight `(image, target_text)` pair and any user/assistant framing would pollute
-the learned signal.
+This is intended for corpora where every row is a tight `(image, target_text)`
+pair and any user/assistant framing would pollute the learned signal.
 
 ::: {.callout-note}
-**Currently supported via the streaming `pretraining_dataset` route only.**
-A non-streaming `datasets:` route (`type: multimodal_pretrain` under
-`datasets:`) is intentionally not wired — the collator selection inside
-`build_collator` only routes MM CPT batches under the pretraining branch.
-Configure MM CPT under `pretraining_dataset:` with `streaming: true` as
-shown below.
+MM CPT supports two raw image+text routes:
+
+- `pretraining_dataset:` with `streaming: true` for iterable streaming.
+- `datasets:` with `type: multimodal_pretrain` for the standard non-streaming
+  prepared dataset pipeline.
+
+The non-streaming path stores token metadata plus image paths in the prepared
+dataset; decoded image tensors are loaded by the collator at batch time.
 :::
 
+Full config examples are available under `examples/qwen2_5-vl/`:
+
+- `mm-cpt-nonstreaming-qlora.yaml` for raw `datasets:`
+  preprocessing plus `num_epochs` based training.
+- `mm-cpt-streaming-qlora.yaml` for streaming `pretraining_dataset:`
+  training with prepared-cache resume.
+
 ### Dataset format (JSONL)
 
 Two keys per row: `text` (the raw string) and `images` (list of local paths).
@@ -413,7 +421,7 @@ Notes:
 Axolotl autodetects the placeholder from the loaded processor. If autodetection
 fails, supply `image_token: <your token>` on the dataset entry.
 
-### YAML example
+### Streaming YAML example
 
 ```yaml
 base_model: HuggingFaceTB/SmolVLM-500M-Instruct
@@ -438,6 +446,91 @@ micro_batch_size: 1
 gradient_accumulation_steps: 8
 ```
 
+For large streaming runs that need cheap checkpoint resume, set
+`dataset_prepared_path` so Axolotl can cache the tokenized MM CPT rows, and use
+`ignore_data_skip: true` when you prefer resuming optimizer/scheduler state
+without Trainer fast-forwarding through already-seen multimodal batches:
+
+```yaml
+pretraining_dataset:
+  - path: /path/to/shards/*.jsonl
+    ds_type: json
+    type: multimodal_pretrain
+    text_column: text
+    image_column: images
+
+streaming: true
+dataset_prepared_path: ./data/mm-cpt-stream-cache
+ignore_data_skip: true
+max_steps: 10000
+```
+
+### Non-streaming prepared YAML example
+
+```yaml
+base_model: HuggingFaceTB/SmolVLM-500M-Instruct
+processor_type: AutoProcessor
+
+datasets:
+  - path: /path/to/train.jsonl
+    ds_type: json
+    type: multimodal_pretrain
+    text_column: text
+    image_column: images
+    image_base_dir: /path/to/images   # optional, for relative paths
+    # image_token: "<image>"           # optional override; autodetect by default
+
+streaming: false
+dataset_prepared_path: ./data/prepared
+sequence_len: 2048
+sample_packing: false                  # REQUIRED — see below
+remove_unused_columns: false           # auto-set by validator
+
+num_epochs: 1                           # max_steps is optional in this path
+micro_batch_size: 1
+gradient_accumulation_steps: 8
+```
+
+Because this route creates a map-style dataset with a known length, Axolotl can
+infer scheduler/training steps from `num_epochs`; unlike streaming
+`pretraining_dataset`, you do not need to manually calculate `max_steps`.
+
+You can run `axolotl preprocess` on this mode. The prepared Arrow dataset keeps
+`images` as image references and `_mm_text` as raw text for the collator; the
+collator loads images and calls the configured processor during training. If
+your JSONL stores relative image paths, keep `image_base_dir` set in the train
+config that consumes the prepared dataset.
+
+### Already-tokenized MM CPT rows
+
+The non-streaming `datasets:` path also accepts rows that have already been
+tokenized, as long as they still include `_mm_text` and `images` for the
+collator. The collator re-tokenizes `_mm_text` with the processor at batch time
+so image placeholders expand consistently with the loaded model; the stored
+token columns are used by the prepared dataset pipeline and should match the raw
+text.
+
+```yaml
+base_model: HuggingFaceTB/SmolVLM-500M-Instruct
+processor_type: AutoProcessor
+
+datasets:
+  - path: /path/to/pretokenized_mm_cpt.jsonl
+    ds_type: json
+    type: multimodal_pretrain
+
+streaming: false
+skip_prepare_dataset: true
+sequence_len: 2048
+sample_packing: false
+remove_unused_columns: false
+num_epochs: 1
+```
+
+```json
+{"_mm_text": "<image>\nText target.", "images": ["/dataset/image.png"], "input_ids": [1, 2, 3], "attention_mask": [1, 1, 1], "labels": [1, 2, 3]}
+```
+
 ### Eval datasets
 
 `test_datasets` accepts multimodal entries (`type: multimodal_pretrain` or
@@ -473,6 +566,12 @@ The following combinations are rejected at config-load time with a clear error:
 - `chat_template` set to anything — defeats the purpose of the CPT path.
 - `processor_type` unset — no processor means no image tensors.
 - Multiple MM `test_datasets` entries with mismatched `image_base_dir` or `image_token`.
+- Multiple MM training entries under `datasets:` or `pretraining_dataset:`.
+- `datasets:` MM CPT with `streaming: true` — use `pretraining_dataset:` for
+  streaming and `datasets:` for the non-streaming prepared path.
+- `datasets:` MM CPT with `excess_length_strategy: truncate` — the collator
+  re-tokenizes `_mm_text` at batch time, so truncating only prepared token ids
+  would not truncate the actual model inputs.
 
 In addition, the following model families are **not supported** in v1 and will
 be rejected when their processor is loaded:

examples/qwen2_5-vl/mm-cpt-nonstreaming-qlora.yaml

@@ -0,0 +1,62 @@
+base_model: Qwen/Qwen2.5-VL-7B-Instruct
+processor_type: AutoProcessor
+trust_remote_code: true
+
+load_in_4bit: true
+
+# Raw JSONL row: {"text": "<|image_pad|>\nText target.", "images": ["image.png"]}
+datasets:
+  - path: /path/to/mm_cpt_train.jsonl
+    ds_type: json
+    type: multimodal_pretrain
+    split: train
+    text_column: text
+    image_column: images
+    image_base_dir: /path/to/images
+    # image_token: "<|image_pad|>"
+
+streaming: false
+dataset_prepared_path: last_run_prepared/qwen2_5_vl_mm_cpt
+val_set_size: 0.0
+output_dir: ./outputs/qwen2_5_vl_mm_cpt
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 8192
+pad_to_sequence_len: false
+sample_packing: false
+remove_unused_columns: false
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: true
+fp16:
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0

examples/qwen2_5-vl/mm-cpt-streaming-qlora.yaml

@@ -0,0 +1,63 @@
+base_model: Qwen/Qwen2.5-VL-7B-Instruct
+processor_type: AutoProcessor
+trust_remote_code: true
+
+load_in_4bit: true
+
+# Cache prepared shards to resume without replaying image preprocessing.
+pretraining_dataset:
+  - path: /path/to/mm_cpt_shards/*.jsonl
+    ds_type: json
+    type: multimodal_pretrain
+    split: train
+    text_column: text
+    image_column: images
+    image_base_dir: /path/to/images
+    # image_token: "<|image_pad|>"
+
+streaming: true
+dataset_prepared_path: last_run_prepared/qwen2_5_vl_mm_cpt_stream_cache
+ignore_data_skip: true
+max_steps: 10000
+val_set_size: 0.0
+output_dir: ./outputs/qwen2_5_vl_mm_cpt_stream
+
+adapter: qlora
+lora_model_dir:
+
+sequence_len: 8192
+pad_to_sequence_len: false
+sample_packing: false
+remove_unused_columns: false
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: true
+fp16:
+tf32: true
+
+gradient_checkpointing: true
+gradient_checkpointing_kwargs:
+  use_reentrant: false
+logging_steps: 1
+attn_implementation: flash_attention_2
+
+warmup_ratio: 0.1
+save_steps: 1000
+eval_steps: 1000
+weight_decay: 0.0

src/axolotl/core/builders/causal.py

@@ -52,20 +52,38 @@
 
 
 def _is_multimodal_cpt(cfg) -> bool:
-    if not getattr(cfg, "pretraining_dataset", None):
+    return _get_mm_cpt_config(cfg) is not None
+
+
+def _entry_is_multimodal_cpt(entry) -> bool:
+    if entry is None:
         return False
-    ds_first = cfg.pretraining_dataset[0]
     ds_type = None
     mm_flag = None
-    if hasattr(ds_first, "type"):
-        ds_type = getattr(ds_first, "type", None)
-        mm_flag = getattr(ds_first, "multimodal", None)
-    elif isinstance(ds_first, dict):
-        ds_type = ds_first.get("type")
-        mm_flag = ds_first.get("multimodal")
+    if hasattr(entry, "type"):
+        ds_type = getattr(entry, "type", None)
+        mm_flag = getattr(entry, "multimodal", None)
+    elif isinstance(entry, dict):
+        ds_type = entry.get("type")
+        mm_flag = entry.get("multimodal")
     return (ds_type == "multimodal_pretrain") or bool(mm_flag)
 
 
+def _get_mm_cpt_config(cfg, is_eval: bool = False):
+    if is_eval and getattr(cfg, "test_datasets", None):
+        for entry in cfg.test_datasets:
+            if _entry_is_multimodal_cpt(entry):
+                return entry
+    for collection_name in ("pretraining_dataset", "datasets"):
+        collection = getattr(cfg, collection_name, None)
+        if not collection:
+            continue
+        for entry in collection:
+            if _entry_is_multimodal_cpt(entry):
+                return entry
+    return None
+
+
 def _mm_cpt_get(pt_cfg, key, default=None):
     if isinstance(pt_cfg, dict):
         return pt_cfg.get(key, default)
@@ -476,12 +494,7 @@ def _build_mm_pretrain_collator(self, pad_to_multiple_of=None, is_eval=False):
             build_image_token_spec,
         )
 
-        if is_eval and self.cfg.test_datasets:
-            pt_cfg = self.cfg.test_datasets[0]
-        elif self.cfg.pretraining_dataset:
-            pt_cfg = self.cfg.pretraining_dataset[0]
-        else:
-            pt_cfg = {}
+        pt_cfg = _get_mm_cpt_config(self.cfg, is_eval=is_eval) or {}
         spec = build_image_token_spec(
             self.processor, override=_mm_cpt_get(pt_cfg, "image_token")
         )