diff --git a/docs/multimodal.qmd b/docs/multimodal.qmd index aabff03f26..6fba75dc1b 100644 --- a/docs/multimodal.qmd +++ b/docs/multimodal.qmd @@ -360,6 +360,120 @@ Here is an example of a multi-modal dataset: ] ``` +## Continued Pre-training (CPT) with images {#sec-multimodal-cpt} + +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. + +### Dataset format (JSONL) + +Two keys per row: `text` (the raw string) and `images` (list of local paths). +The `text` must contain the model's placeholder token **once per image**, +placed immediately before the text it describes, followed by a newline: + +```json +{"text": "\nInvoice number: 10427. Total due: 148.32 USD.", "images": ["/dataset/crops/doc_14_p2.png"]} +{"text": "\nChapter 3 begins with a description of the storm over the harbor.", "images": ["/dataset/crops/doc_14_p3.png"]} +``` + +Notes: + +- Never wrap the row in `User:` / `Assistant:` / `Transcribe this:` scaffolding. That is the whole point of the CPT path. +- Do not manually append an EOS token. Axolotl appends one during tokenization. +- The newline between the placeholder and the real text preserves the BPE + boundary. Without it, some tokenizers merge the visual-token boundary with + the first real character. + +### The placeholder token varies by model + +| Model family | Placeholder | Notes | +|---|---|---| +| LLaVA-1.5 / 1.6 | `` | | +| SmolVLM / SmolVLM2 / Idefics3 | `` | Processor expands to 1088 tokens (17 tiles × 64) | +| Qwen2-VL / Qwen2.5-VL / Qwen3-VL | `<\|image_pad\|>` | Processor autowraps with `<\|vision_start\|>` / `<\|vision_end\|>` | +| Gemma-3 | `` | Processor expands to 256 `` | +| Gemma-4 | `<\|image\|>` | Processor expands to 256 `<\|image\|>` | + +Axolotl autodetects the placeholder from the loaded processor. If autodetection +fails, supply `image_token: ` on the dataset entry. + +### YAML example + +```yaml +base_model: HuggingFaceTB/SmolVLM-500M-Instruct +processor_type: AutoProcessor + +pretraining_dataset: + - path: /path/to/shards/*.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: "" # optional override; autodetect by default + +streaming: true +sequence_len: 2048 +sample_packing: false # REQUIRED — see below +remove_unused_columns: false # auto-set by validator + +max_steps: 10000 +micro_batch_size: 1 +gradient_accumulation_steps: 8 +``` + +### Eval datasets + +`test_datasets` accepts multimodal entries (`type: multimodal_pretrain` or +`multimodal: true`). Per-entry `text_column` and `image_column` are honored +independently. When more than one multimodal entry is provided, +`image_base_dir` and `image_token` must be either unset on every entry or +identical across them, because the eval collator resolves both once for the +merged eval stream. + +### Gates and rejections + +The following combinations are rejected at config-load time with a clear error: + +- `sample_packing: true` — cross-row packing would break the 1-to-1 alignment + between text placeholders and `pixel_values`. +- `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`. + +In addition, the following model families are **not supported** in v1 and will +be rejected when their processor is loaded: + +- **Llama-3.2-Vision (Mllama)** — uses cross-attention image injection, not + in-stream placeholders. Use chat-template SFT. +- **Pixtral** — requires `mistral_common` and a different API. +- **InternVL** — ships a custom processor that doesn't produce `pixel_values`. + +Per-row validation: at encode time the row's text is tokenized once and the +number of `image_token_id` occurrences in the resulting token-id list must +equal `len(images)`. Counting by token id (not by substring) avoids false +matches — e.g., `` would substring-match inside ``. +This is a critical guardrail — LLaVA and Qwen-VL processors silently +accept rows without placeholders and drop the image, which looks like +successful training but teaches nothing. If a row fails this check, +inspect the tokenized ids rather than the raw string. + +### Why masking image tokens in labels is automatic + +For this multimodal CPT path, the collator masks every image-family token id +(``, `<\|image_pad\|>`, `<\|vision_start\|>`, `<\|vision_end\|>`, +``, ``, ``, `<\|image\|>`, +etc.) to `-100` in the labels tensor. + +Supported processors expand image placeholders into vision-specific token ids. +If those ids contribute to loss, the model is trained to predict patch or +marker tokens rather than only the target text. On architectures like +Qwen-VL and SmolVLM, leaving them unmasked substantially increases loss and can +destabilize training. + ## FAQ 1. `PIL.UnidentifiedImageError: cannot identify image file ...` diff --git a/docs/multimodal_assistant_mask.md b/docs/multimodal_assistant_mask.md new file mode 100644 index 0000000000..339ab420f8 --- /dev/null +++ b/docs/multimodal_assistant_mask.md @@ -0,0 +1,84 @@ +# Multimodal assistant-only loss masking + +## Correct placement + +```yaml +# Top-level: only train_on_inputs lives here. +train_on_inputs: false + +datasets: + - path: data/train.jsonl + type: chat_template + roles_to_train: # per-dataset — this is what the MM scanner reads + - assistant + train_on_eos: turn # per-dataset — same + +test_datasets: + - path: data/val.jsonl + type: chat_template + split: train + roles_to_train: + - assistant + train_on_eos: turn +``` + +## How to verify at runtime + +`build_collator` logs the resolved knobs at INFO: + +```text +MM collator: train_on_inputs=False roles_to_train=['assistant'] train_on_eos=turn role_boundaries_override=none +``` + +If `roles_to_train` logs as `None`, the YAML knobs are not reaching the +scanner — check that they are under `datasets[0]`, not at the root. + +Each verified strategy additionally logs its resolved boundary token ids at +strategy init (e.g. `<|turn>model` → `[105, 4368]`, `` → `[106]` for +Gemma 4). If a strategy emits the "has no built-in role boundaries ... only +pad and media tokens are masked" one-shot warning instead, it is on the +fallback path — declare per-role markers in YAML via `cfg.role_boundaries` +(below) to activate masking. The strategies currently on this path are +listed in the audit table above under `fallback + warn`. + +## Config-based override: `cfg.role_boundaries` + +For the "unverified" strategies above, or for custom chat templates that +don't match a built-in strategy's markers, users can declare role boundaries +directly in YAML without subclassing: + +```yaml +role_boundaries: + - role: assistant + start: "<|turn>model" + end: "" + - role: user + start: "<|turn>user" + end: "" + # Optional keys: + # include_start: false # default False + # include_end: true # default True, respects cfg.train_on_eos + # end: eos_token # sentinel: resolves to tokenizer.eos_token_id + # end: null # span runs to end of sequence +``` + +Semantics: + +- `start` and `end` are literal strings; axolotl encodes them at strategy + init via `tokenizer.encode(..., add_special_tokens=False)` and logs the + resolved token-id sequences at INFO level. +- The special value `end: eos_token` is the portable way to express + "Pixtral-style assistant turns end at EOS" without hard-coding an id. +- `role_boundaries` is an **opt-in override**. A non-empty list **replaces** + the strategy's built-in declarations wholesale (partial overlays are + intentionally unsupported — they're hard to reason about at review time). + Leaving the field unset *or* setting it to an empty list (`[]`) both mean + "use the strategy's built-ins." Writing `role_boundaries: []` is almost + always a typo or leftover — honoring it literally would produce all-masked + labels and zero gradient, so it is treated the same as unset. +- `cfg.roles_to_train` still governs which declared roles contribute to + loss. You can declare `user` and `assistant` boundaries and set + `roles_to_train: ["assistant"]` to have the scanner correctly identify + user spans as masking boundaries without training on their content. +- Invalid specs fail loudly at strategy init (missing `role`/`start`, + unencodable markers), not silently at loss-compute time. diff --git a/src/axolotl/core/builders/causal.py b/src/axolotl/core/builders/causal.py index fe832dd452..168ce9d3d0 100644 --- a/src/axolotl/core/builders/causal.py +++ b/src/axolotl/core/builders/causal.py @@ -44,12 +44,34 @@ V2BatchSamplerDataCollatorForSeq2Seq, ) from axolotl.utils.collators.mm_chat import MultiModalChatDataCollator +from axolotl.utils.collators.mm_pretrain import MultiModalPretrainDataCollator from axolotl.utils.import_helper import get_cls_from_module_str from axolotl.utils.logging import get_logger LOG = get_logger(__name__) +def _is_multimodal_cpt(cfg) -> bool: + if not getattr(cfg, "pretraining_dataset", 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") + return (ds_type == "multimodal_pretrain") or bool(mm_flag) + + +def _mm_cpt_get(pt_cfg, key, default=None): + if isinstance(pt_cfg, dict): + return pt_cfg.get(key, default) + return getattr(pt_cfg, key, default) + + class HFCausalTrainerBuilder(TrainerBuilderBase): """ Build the HuggingFace training args/trainer for causal models and reward modeling @@ -451,6 +473,31 @@ def build(self, total_num_steps): return trainer + def _build_mm_pretrain_collator(self, pad_to_multiple_of=None, is_eval=False): + from axolotl.prompt_strategies.multimodal_pretrain import ( + 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 = {} + spec = build_image_token_spec( + self.processor, override=_mm_cpt_get(pt_cfg, "image_token") + ) + collator_kwargs = { + "tokenizer": self.tokenizer, + "processor": self.processor, + "image_token_spec": spec, + "image_base_dir": _mm_cpt_get(pt_cfg, "image_base_dir"), + "max_length": self.cfg.sequence_len, + } + if pad_to_multiple_of is not None: + collator_kwargs["pad_to_multiple_of"] = pad_to_multiple_of + return MultiModalPretrainDataCollator(**collator_kwargs) + def build_collator( self, training_args, # type: "AxolotlTrainingArguments" # type: ignore @@ -458,6 +505,16 @@ def build_collator( **kwargs, ): if training_args.pretraining: + # Intercept MM CPT before the text-only pretraining branches. + if ( + self.cfg.processor_type + and self.processor + and _is_multimodal_cpt(self.cfg) + ): + return self._build_mm_pretrain_collator( + pad_to_multiple_of=kwargs.get("pad_to_multiple_of"), + is_eval=is_eval, + ) if ( self.cfg.pretraining_sample_concatenation is False or self.cfg.micro_batch_size > 1 @@ -519,14 +576,64 @@ def build_collator( else: collator = BatchSamplerDataCollatorForSeq2Seq else: + if ( + self.cfg.processor_type + and self.processor + and _is_multimodal_cpt(self.cfg) + ): + return self._build_mm_pretrain_collator( + pad_to_multiple_of=kwargs.get("pad_to_multiple_of"), + is_eval=is_eval, + ) if self.cfg.processor_type and self.processor: collator = MultiModalChatDataCollator + # Mirror ChatTemplateStrategy: per-dataset masking knobs from first MM dataset, else global cfg. + # NOTE: Multi-dataset configs use the first dataset's masking knobs for all datasets; + # heterogeneous per-dataset overrides are not supported in the MM path today. + ds_entries = self.cfg.datasets or [] + ds_cfg = ds_entries[0] if ds_entries else None + + def _ds_get(cfg_obj, key): + # Handle DictDefault / dict / pydantic uniformly: + # dict-style .get first, then attribute access. + if cfg_obj is None: + return None + if hasattr(cfg_obj, "get"): + try: + return cfg_obj.get(key) + except (AttributeError, KeyError, TypeError): + pass + return getattr(cfg_obj, key, None) + + roles_to_train = _ds_get(ds_cfg, "roles_to_train") + train_on_eos = _ds_get(ds_cfg, "train_on_eos") + + # cfg.role_boundaries replaces the strategy's built-in markers. + role_boundaries_override = None + if self.cfg.role_boundaries: + role_boundaries_override = list(self.cfg.role_boundaries) + + # build() calls build_collator twice (eval + train); log once. + if not is_eval: + LOG.info( + "MM collator: train_on_inputs=%s roles_to_train=%s " + "train_on_eos=%s role_boundaries_override=%s", + bool(self.cfg.train_on_inputs), + roles_to_train, + train_on_eos, + "set" if role_boundaries_override else "none", + ) + kwargs["processing_strategy"] = get_processing_strategy( self.processor, training_args.chat_template, self.cfg.chat_template, image_size=training_args.image_size, image_resize_algorithm=training_args.image_resize_algorithm, + train_on_inputs=bool(self.cfg.train_on_inputs), + roles_to_train=roles_to_train, + train_on_eos=train_on_eos, + role_boundaries_override=role_boundaries_override, ) elif self.cfg.batch_flattening: collator = DataCollatorWithFlattening diff --git a/src/axolotl/processing_strategies.py b/src/axolotl/processing_strategies.py index cb1f9d984b..217bc765b5 100644 --- a/src/axolotl/processing_strategies.py +++ b/src/axolotl/processing_strategies.py @@ -1,6 +1,7 @@ """Module containing ProcessingStrategy classes and its derivative for different MultiModal Model types""" from copy import deepcopy +from dataclasses import dataclass, field from typing import Optional from PIL import Image, ImageOps @@ -17,9 +18,35 @@ LOG = get_logger(__name__) +# One-shot warning dedupe so opt-out subclasses don't spam per-batch. +_ROLE_MASK_WARNED: set[str] = set() + +# Supported values for ``train_on_eos`` — mirrors the text-only +# ChatTemplateStrategy (``turn`` = trainable turn ends only, ``all`` = every +# turn end, ``none`` = never, ``last`` = only the final trainable turn end). +_VALID_TRAIN_ON_EOS = ("turn", "all", "none", "last") + + +@dataclass(frozen=True) +class RoleBoundary: + """One role's token-level span markers for the masking scanner. + + Empty ``end_tokens`` means end-of-sequence terminates the span. + """ + + role: str + start_tokens: list[int] + end_tokens: list[int] = field(default_factory=list) + include_start: bool = False + include_end: bool = True + class ProcessingStrategy: - """Base Processing Strategy class""" + """Base Processing Strategy class. + + Subclasses opt in to role masking by overriding ``_build_role_boundaries``; + otherwise only pad + media tokens are masked (legacy behavior, one-shot warned). + """ def __init__( self, @@ -27,6 +54,10 @@ def __init__( chat_template: Optional[str] = None, image_size: int | tuple[int, int] | None = None, image_resize_algorithm: Resampling | None = None, + train_on_inputs: bool = False, + roles_to_train: Optional[list[str]] = None, + train_on_eos: Optional[str] = None, + role_boundaries_override: Optional[list[dict]] = None, ): self.processor = processor self.chat_template = chat_template @@ -38,54 +69,97 @@ def __init__( image_resize_algorithm or Image.Resampling.BILINEAR ) + # Defaults mirror the text-only ChatTemplateStrategy. An explicit + # empty list is honored as "no trainable roles" (masks everything); + # only ``None`` falls back to the default of assistant-only. + self.train_on_inputs = bool(train_on_inputs) + self.roles_to_train = ( + list(roles_to_train) if roles_to_train is not None else ["assistant"] + ) + self.train_on_eos = train_on_eos if train_on_eos is not None else "turn" + if self.train_on_eos not in _VALID_TRAIN_ON_EOS: + raise ValueError( + f"train_on_eos={self.train_on_eos!r} is not one of " + f"{_VALID_TRAIN_ON_EOS}." + ) + if hasattr(processor, "image_token"): self.image_token = processor.image_token self.image_token_id = processor.tokenizer.convert_tokens_to_ids( self.image_token ) + built_in = self._build_role_boundaries() + + # Truthiness check: empty list == unset (opt-in escape hatch), so + # `role_boundaries: []` in YAML falls through to built-ins instead of + # producing all-masked labels. + if role_boundaries_override: + overridden = _resolve_role_boundary_override( + role_boundaries_override, self.processor.tokenizer + ) + LOG.info( + "%s: overriding built-in role boundaries (%d decls) " + "with cfg.role_boundaries (%d decls).", + type(self).__name__, + len(built_in), + len(overridden), + ) + self.role_boundaries: list[RoleBoundary] = overridden + source = "override" + else: + self.role_boundaries = built_in + source = "built-in" + + # Single-line, grep-friendly summary of the resolved masking config so + # "why isn't masking firing?" is visible in training logs. For + # overrides we include the fully resolved (role, start_ids, end_ids) + # tuples; for built-ins we log a count (subclasses vary and logging + # every id sequence would be noisy on, e.g., Llama3 with five roles). + boundaries_repr: str | list[tuple[str, list[int], list[int]]] + if source == "override": + boundaries_repr = [ + (b.role, b.start_tokens, b.end_tokens) for b in self.role_boundaries + ] + else: + boundaries_repr = f"{len(self.role_boundaries)} built-in" + LOG.info( + "ProcessingStrategy init: class=%s train_on_inputs=%s " + "roles_to_train=%s train_on_eos=%s boundaries_source=%s " + "boundaries=%s", + type(self).__name__, + self.train_on_inputs, + self.roles_to_train, + self.train_on_eos, + source, + boundaries_repr, + ) + + def _build_role_boundaries(self) -> list[RoleBoundary]: + """Subclasses declare role boundaries here; [] opts out of role masking.""" + return [] + def __call__(self, examples: list[dict]) -> list[dict]: - """ - Preprocess conversation examples to ensure consistent format. - Converts different conversation formats to OpenAI format with 'messages'. - Supports two formats: - 1. OpenAI format with 'messages' - 2. Legacy format with 'conversations' - - Args: - examples: list of conversation dictionaries - - Returns: - list of dicts in OpenAI format with 'messages' key - - Raises: - ValueError: If the conversation format is not supported - """ + """Normalize examples to OpenAI ``messages`` format (accepts legacy ``conversations``).""" role_mapping = { "human": "user", "gpt": "assistant", } def normalize_role(role: str) -> str: - """Normalize role names to OpenAI format. Default to original role if not found.""" return role_mapping.get(role, role) def convert_legacy_format(example: dict) -> dict: - """Convert legacy 'conversations' format to OpenAI 'messages' format.""" messages = [ {"role": normalize_role(convo["from"]), "content": convo["value"]} for convo in example["conversations"] ] - - # Create new dict without 'conversations' key result = deepcopy(example) result.pop("conversations") result["messages"] = messages return result def convert_messages_to_multimedia_messages(messages: list[dict]) -> list[dict]: - """Convert regular messages format to Messages format with content type""" - new_messages = [] for message in messages: if isinstance(message["content"], str): @@ -119,21 +193,27 @@ def convert_messages_to_multimedia_messages(messages: list[dict]) -> list[dict]: "Only `messages` and `conversations` message keys are currently supported." ) - processed_example = None - if ( - "messages" in example and example["messages"] is not None - ): # OpenAI format - processed_example = example - else: # Legacy format + if "messages" in example and example["messages"] is not None: + # Deepcopy for symmetry with convert_legacy_format (which + # deepcopies internally) so downstream mutations of + # processed_example don't leak back to the caller's input. + processed_example = deepcopy(example) + elif "conversations" in example: processed_example = convert_legacy_format(example) + else: + # `messages` is present but None, and no `conversations` + # fallback exists — convert_legacy_format would KeyError on + # ["conversations"]. Surface a clear validation error instead. + raise ValueError( + "`messages` is present but None; provide non-null " + "`messages` or a `conversations` field." + ) - # convert regular messages format to Messages format with content type - # for compatibility with apply_chat_template + # Required for apply_chat_template compatibility. processed_example["messages"] = convert_messages_to_multimedia_messages( processed_example["messages"] ) - # find the image key if it exists possible_image_keys = ["images", "image"] image_key = None for key in possible_image_keys: @@ -141,11 +221,8 @@ def convert_messages_to_multimedia_messages(messages: list[dict]) -> list[dict]: image_key = key break - # if the image key exists, add the image to the first user message if image_key is not None and processed_example[image_key] is not None: - # TODO: check if it's normal to be single image only for common datasets - # From observation, it's usually a list of single image but some datasets may have several columns for images - # Temporary solution: take the first image and suggest people convert their datasets to use multi-content Messages + # TODO: support multi-image samples; for now we take the first. if len(processed_example[image_key]) > 1: LOG.warning( f"Found {len(processed_example[image_key])} images in a sample. Using the first one." @@ -155,7 +232,6 @@ def convert_messages_to_multimedia_messages(messages: list[dict]) -> list[dict]: image_value = processed_example[image_key][0] - # Handle image loading (Image, url, path, base64) image_value = load_image(image_value) if self.image_size is not None: @@ -168,11 +244,8 @@ def convert_messages_to_multimedia_messages(messages: list[dict]) -> list[dict]: self.image_size, self.image_resize_algorithm ) else: - # Set the padding value; here we use black (0, 0, 0) for RGB images + # Int image_size: preserve aspect ratio then pad to square (black) to avoid distortion. padding_color = (0, 0, 0) - - # When image_size is an int (square target), preserve aspect ratio then pad - # This is to prevent aspect ratio distortion when resizing to square image_value = ImageOps.pad( image_value, (self.image_size, self.image_size), @@ -180,8 +253,6 @@ def convert_messages_to_multimedia_messages(messages: list[dict]) -> list[dict]: color=padding_color, ) - # Look for any image type in the first message - # some dataset have an {type: "image"} in the first message msg_ind_to_add = None ind_to_add = None first_user_idx = None @@ -192,7 +263,7 @@ def convert_messages_to_multimedia_messages(messages: list[dict]) -> list[dict]: for i, content in enumerate( processed_example["messages"][msg_idx]["content"] ): - # Usually datasets created with image columns, don't have it in the messages itself + # Column-image datasets often leave a bare {type: "image"} placeholder. if content["type"] == "image" and all( k not in content for k in ["image", "url", "path", "base64"] ): @@ -200,13 +271,11 @@ def convert_messages_to_multimedia_messages(messages: list[dict]) -> list[dict]: ind_to_add = i break - # If an image type is found, add the image to that index if ind_to_add is not None and msg_ind_to_add is not None: processed_example["messages"][msg_ind_to_add]["content"][ ind_to_add ]["image"] = image_value else: - # if no image type is found, add it to end of the first user message if first_user_idx is None: first_user_idx = 0 processed_example["messages"][first_user_idx]["content"].append( @@ -221,28 +290,224 @@ def convert_messages_to_multimedia_messages(messages: list[dict]) -> list[dict]: return processed_examples def _mask_non_assistant(self, labels: Tensor) -> Tensor: - """ - Mask non assistant regions to -100. - To be implemented per subclass. - """ - return labels + """Mask non-trainable role regions to -100 using ``self.role_boundaries``.""" + if self.train_on_inputs: + return labels + + # Legacy no-op for boundary-less strategies; warn once so the miss shows up in logs. + if not self.role_boundaries: + key = type(self).__name__ + if key not in _ROLE_MASK_WARNED: + _ROLE_MASK_WARNED.add(key) + LOG.warning( + "%s has no built-in role boundaries; " + "cfg.train_on_inputs / cfg.roles_to_train / cfg.train_on_eos " + "will NOT restrict loss to assistant tokens for this " + "multimodal model — only pad and media tokens are masked, " + "every other token (system, user, assistant) contributes " + "to loss. To enable assistant-only masking, declare " + "per-role markers in YAML via cfg.role_boundaries — see " + "docs/multimodal_assistant_mask.md for the format and the " + "list of strategies on this fallback path.", + key, + ) + return labels + + return _apply_role_boundaries( + labels, + self.role_boundaries, + roles_to_train=set(self.roles_to_train), + train_on_eos=self.train_on_eos, + ) def process_labels(self, input_ids: Tensor) -> Tensor: labels = input_ids.clone() - labels = self._mask_non_assistant(labels) + pad_id = getattr(self.processor.tokenizer, "pad_token_id", None) + if pad_id is not None: + labels[labels == pad_id] = -100 + if self.image_token_id is not None: + labels[labels == self.image_token_id] = -100 + return labels - # The labels are the input_ids, and we mask the padding tokens in the loss computation - labels[labels == self.processor.tokenizer.pad_token_id] = -100 - # Ignore the image token index in the loss computation (model specific) - labels[labels == self.image_token_id] = -100 +def _apply_role_boundaries( + labels: Tensor, + role_boundaries: list[RoleBoundary], + roles_to_train: set[str], + train_on_eos: str, +) -> Tensor: + """Mask tokens outside trainable role spans to -100. + + Scan is greedy-left with longest-prefix-wins on start_tokens to disambiguate + nested markers (e.g. ``<|im_start|>assistant`` vs ``<|im_start|>``). + ``train_on_eos`` accepts ``"turn"`` (end marker in loss on trainable turns + only), ``"all"`` (always), ``"none"`` (never — overrides ``include_end``), + ``"last"`` (only on the last trainable turn in the sequence). + """ + mask = zeros_like(labels) + # For "last": remember each trainable turn's end-marker span so we can + # unmask only the final one after the scan finishes. + last_trainable_end_span: list[Optional[tuple[int, int]]] = [None] * labels.shape[0] + + # Work on a Python list per row — avoids O(n*boundaries) Tensor→list + # conversions in the hot prefix-match loop. + def _match_prefix(label: list[int], start_pos: int, tok_seq: list[int]) -> bool: + if not tok_seq or start_pos + len(tok_seq) > len(label): + return False + return label[start_pos : start_pos + len(tok_seq)] == tok_seq + + def _find_end( + label: list[int], start_pos: int, end_tok: list[int] + ) -> tuple[int, bool]: + # Empty end_tok means run to end-of-sequence. + if not end_tok: + return len(label), False + k = start_pos + while k < len(label): + if _match_prefix(label, k, end_tok): + return k + len(end_tok), True + k += 1 + return k, False + + for i in range(labels.shape[0]): + label = labels[i].tolist() + j = 0 + n = len(label) + while j < n: + best_match: Optional[RoleBoundary] = None + for b in role_boundaries: + if _match_prefix(label, j, b.start_tokens): + if best_match is None or len(b.start_tokens) > len( + best_match.start_tokens + ): + best_match = b + if best_match is None: + j += 1 + continue + + start_of_content = j + len(best_match.start_tokens) + end_after, found_end = _find_end( + label, start_of_content, best_match.end_tokens + ) - return labels + role_in_loss = best_match.role in roles_to_train + + if role_in_loss: + if best_match.include_start: + mask[i][j:start_of_content] = 1 + content_end = ( + end_after - len(best_match.end_tokens) if found_end else end_after + ) + mask[i][start_of_content:content_end] = 1 + # train_on_eos="none"/"last" override include_end during main + # loop; "last" is applied after the scan finishes. + if ( + found_end + and best_match.include_end + and train_on_eos not in ("none", "last") + ): + mask[i][content_end:end_after] = 1 + if found_end and best_match.include_end and train_on_eos == "last": + last_trainable_end_span[i] = (content_end, end_after) + else: + # Non-trainable role on train_on_eos="all": gate on include_end + # so Pixtral / Mistral V7 Tekken shared [/INST] doesn't leak. + if found_end and best_match.include_end and train_on_eos == "all": + content_end = end_after - len(best_match.end_tokens) + mask[i][content_end:end_after] = 1 + + # When include_end=False, do not consume the end marker: back up so + # the next iteration can re-match it as the next boundary's start + # marker (Pixtral / Mistral V7 Tekken share [/INST] between + # user-end and assistant-start). Requires end_tokens non-empty and + # actually found. + if found_end and not best_match.include_end and best_match.end_tokens: + j = end_after - len(best_match.end_tokens) + else: + j = end_after + + if train_on_eos == "last" and (span := last_trainable_end_span[i]) is not None: + s, e = span + mask[i][s:e] = 1 + + labels[i][mask[i] == 0] = -100 + + return labels + + +def _encode_markers(tokenizer, marker_strs: list[str]) -> list[list[int]]: + """Encode markers via ``encode(..., add_special_tokens=False)``; drops empty results.""" + result = [] + for s in marker_strs: + toks = tokenizer.encode(s, add_special_tokens=False) + if toks: + result.append(toks) + return result + + +def _resolve_role_boundary_override(specs: list[dict], tokenizer) -> list[RoleBoundary]: + """Resolve user ``cfg.role_boundaries`` specs into RoleBoundary objects. + + The sentinel ``end == "eos_token"`` resolves to ``eos_token_id`` (used by + Pixtral/Mistral v7 templates). ``end`` null/omitted runs to end-of-sequence. + """ + out: list[RoleBoundary] = [] + for i, spec in enumerate(specs): + if hasattr(spec, "model_dump"): + d = spec.model_dump() + else: + d = dict(spec) + + role = d.get("role") + start_str = d.get("start") + if not role or start_str is None: + raise ValueError( + f"cfg.role_boundaries[{i}] must have both 'role' and 'start' " + f"(got {d!r})." + ) + start_ids = tokenizer.encode(start_str, add_special_tokens=False) + if not start_ids: + raise ValueError( + f"cfg.role_boundaries[{i}]: start marker {start_str!r} " + f"tokenizes to an empty sequence; cannot match." + ) + + end_spec = d.get("end") + if end_spec is None: + end_ids: list[int] = [] + elif end_spec == "eos_token": + eos = getattr(tokenizer, "eos_token_id", None) + if eos is None: + raise ValueError( + f"cfg.role_boundaries[{i}] requested end='eos_token' but " + "the tokenizer has no eos_token_id." + ) + end_ids = [eos] + else: + end_ids = tokenizer.encode(end_spec, add_special_tokens=False) + if not end_ids: + raise ValueError( + f"cfg.role_boundaries[{i}]: end marker {end_spec!r} " + f"tokenizes to an empty sequence; cannot match. Use " + f"end=null to run to end-of-sequence or end='eos_token' " + f"to terminate at the tokenizer's EOS." + ) + + out.append( + RoleBoundary( + role=role, + start_tokens=start_ids, + end_tokens=end_ids, + include_start=bool(d.get("include_start", False)), + include_end=bool(d.get("include_end", True)), + ) + ) + return out class Qwen2VLProcessingStrategy(ProcessingStrategy): - """Processing Strategy class for Qwen2-VL""" + """Processing Strategy class for Qwen2-VL (ChatML ``<|im_start|>{role}\\n ... <|im_end|>``).""" def __init__( self, @@ -250,16 +515,44 @@ def __init__( chat_template: Optional[str] = None, image_size: int | tuple[int, int] | None = None, image_resize_algorithm: Resampling | None = None, + train_on_inputs: bool = False, + roles_to_train: Optional[list[str]] = None, + train_on_eos: Optional[str] = None, + role_boundaries_override: Optional[list[dict]] = None, ): - super().__init__(processor, chat_template, image_size, image_resize_algorithm) + super().__init__( + processor, + chat_template, + image_size, + image_resize_algorithm, + train_on_inputs=train_on_inputs, + roles_to_train=roles_to_train, + train_on_eos=train_on_eos, + role_boundaries_override=role_boundaries_override, + ) self.image_token = "<|image_pad|>" # nosec self.image_token_id = processor.tokenizer.convert_tokens_to_ids( self.image_token ) + def _build_role_boundaries(self) -> list[RoleBoundary]: + tok = self.processor.tokenizer + end = _encode_markers(tok, ["<|im_end|>"]) + if not end: + return [] + end_ids = end[0] + boundaries = [] + for role in ("system", "user", "assistant"): + start = _encode_markers(tok, [f"<|im_start|>{role}\n"]) + if start: + boundaries.append( + RoleBoundary(role=role, start_tokens=start[0], end_tokens=end_ids) + ) + return boundaries + -class Qwen3_5ProcessingStrategy(ProcessingStrategy): - """Processing Strategy class for Qwen3.5 (early-fusion VLM)""" +class Qwen3_5ProcessingStrategy(Qwen2VLProcessingStrategy): + """Processing Strategy class for Qwen3.5 (Qwen2-VL boundaries + ``<|video_pad|>`` mask).""" def __init__( self, @@ -267,11 +560,20 @@ def __init__( chat_template: Optional[str] = None, image_size: int | tuple[int, int] | None = None, image_resize_algorithm: Resampling | None = None, + train_on_inputs: bool = False, + roles_to_train: Optional[list[str]] = None, + train_on_eos: Optional[str] = None, + role_boundaries_override: Optional[list[dict]] = None, ): - super().__init__(processor, chat_template, image_size, image_resize_algorithm) - self.image_token = "<|image_pad|>" # nosec - self.image_token_id = processor.tokenizer.convert_tokens_to_ids( - self.image_token + super().__init__( + processor, + chat_template, + image_size, + image_resize_algorithm, + train_on_inputs=train_on_inputs, + roles_to_train=roles_to_train, + train_on_eos=train_on_eos, + role_boundaries_override=role_boundaries_override, ) self.video_token = "<|video_pad|>" # nosec self.video_token_id = processor.tokenizer.convert_tokens_to_ids( @@ -280,12 +582,44 @@ def __init__( def process_labels(self, input_ids): labels = super().process_labels(input_ids) - labels[labels == self.video_token_id] = -100 + if self.video_token_id is not None: + labels[labels == self.video_token_id] = -100 return labels -class Gemma3ProcessingStrategy(ProcessingStrategy): - """Processing Strategy class for Gemma3""" +class _GemmaTurnStrategy(ProcessingStrategy): + """Gemma3/3n ``{role} ... `` (Gemma 4 uses different markers).""" + + def _build_role_boundaries(self) -> list[RoleBoundary]: + tok = self.processor.tokenizer + end = _encode_markers(tok, [""]) + if not end: + return [] + end_ids = end[0] + boundaries = [] + # Template uses 'model'; external role knob stays 'assistant'. Gemma 3 + # and Gemma 3n jinja templates fold the system message into the first + # user's content prefix and never emit 'system', so we + # don't declare a system boundary here. + role_marker_pairs = [ + ("assistant", "model"), + ("user", "user"), + ] + for external_role, template_role in role_marker_pairs: + start = _encode_markers(tok, [f"{template_role}\n"]) + if start: + boundaries.append( + RoleBoundary( + role=external_role, + start_tokens=start[0], + end_tokens=end_ids, + ) + ) + return boundaries + + +class Gemma3ProcessingStrategy(_GemmaTurnStrategy): + """Processing Strategy class for Gemma3.""" def __init__( self, @@ -293,119 +627,247 @@ def __init__( chat_template: Optional[str] = None, image_size: int | tuple[int, int] | None = None, image_resize_algorithm: Resampling | None = None, + train_on_inputs: bool = False, + roles_to_train: Optional[list[str]] = None, + train_on_eos: Optional[str] = None, + role_boundaries_override: Optional[list[dict]] = None, ): - super().__init__(processor, chat_template, image_size, image_resize_algorithm) - self.image_token = processor.tokenizer.special_tokens_map["boi_token"] - self.image_token_id = processor.tokenizer.convert_tokens_to_ids( - self.image_token + super().__init__( + processor, + chat_template, + image_size, + image_resize_algorithm, + train_on_inputs=train_on_inputs, + roles_to_train=roles_to_train, + train_on_eos=train_on_eos, + role_boundaries_override=role_boundaries_override, ) + # Real Gemma3 tokenizers expose boi_token as a direct attribute, not + # via special_tokens_map (which only holds HF's standard slots). + boi = getattr(processor.tokenizer, "boi_token", None) + if boi is not None: + self.image_token = boi + self.image_token_id = processor.tokenizer.convert_tokens_to_ids(boi) def process_labels(self, input_ids): - labels = input_ids.clone() - - # Follows https://ai.google.dev/gemma/docs/core/huggingface_vision_finetune_qlora - labels[labels == self.processor.tokenizer.pad_token_id] = -100 - labels[labels == self.image_token_id] = -100 - labels[labels == 262144] = -100 # corresponds to - + labels = super().process_labels(input_ids) + # Resolve via tokenizer; fall back to default id + # if not in vocab. Matches Gemma4's pattern. + tok = self.processor.tokenizer + soft_id = tok.convert_tokens_to_ids("") + unk_id = getattr(tok, "unk_token_id", None) + if soft_id is not None and soft_id != unk_id: + labels[labels == soft_id] = -100 + else: + labels[labels == 262144] = -100 return labels -class Gemma3nProcessingStrategy(ProcessingStrategy): - """Processing Strategy class for Gemma3n""" - - def _mask_non_assistant(self, labels: Tensor) -> Tensor: - def _find_token_sequence(label, start_pos, token_sequence): - """Check if token_sequence appears at start_pos in label""" - if start_pos + len(token_sequence) > len(label): - return False - if label[start_pos] != token_sequence[0]: - return False - return ( - label[start_pos : start_pos + len(token_sequence)].tolist() - == token_sequence - ) - - def _find_assistant_end(label, start_pos, assistant_end_tok, mask, i): - """ - Find the end of assistant response and update mask accordingly - - Returns new position to continue from and whether the end seq is found - """ - k = start_pos - while k < len(label): - if not _find_token_sequence(label, k, assistant_end_tok): - mask[i][k] = 1 - k += 1 - continue - - return k + len(assistant_end_tok), True - - return k, False +class Gemma3nProcessingStrategy(_GemmaTurnStrategy): + """Gemma3n: same turn boundaries as Gemma3, additionally masks audio/delimiter tokens.""" - mask = zeros_like(labels) - - assistant_start_str = "model" - assistant_end_str = "" - include_assistant_start_tok = False - include_assistant_end_tok = True + def process_labels(self, input_ids): + labels = super().process_labels(input_ids) + tok = self.processor.tokenizer + # Follows huggingface-gemma-recipes fine_tune_gemma3n_on_t4 notebook. + for attr in ( + "image_token_id", + "audio_token_id", + "boi_token_id", + "eoi_token_id", + ): + tok_id = getattr(tok, attr, None) + if tok_id is not None: + labels[labels == tok_id] = -100 + return labels - # str to tokens - assistant_start_tok = self.processor.tokenizer.encode( - assistant_start_str, add_special_tokens=False - ) - assistant_end_tok = self.processor.tokenizer.encode( - assistant_end_str, add_special_tokens=False - ) - for i, label in enumerate(labels): - j = 0 - # while loop through each tok index in labels[i] - while j < len(label): - # Check until match start seq - if not _find_token_sequence(label, j, assistant_start_tok): - j += 1 - continue - - if include_assistant_start_tok: - mask[i][j : j + len(assistant_start_tok)] = 1 - - # Find where the assistant response ends - start_of_content = j + len(assistant_start_tok) - end_pos, found_end_seq = _find_assistant_end( - label, start_of_content, assistant_end_tok, mask, i +class Gemma4ProcessingStrategy(ProcessingStrategy): + """Processing Strategy class for Gemma 4. + + Boundary markers ``<|turn>model ... `` verified against + google/gemma-4-E2B-it. boi/eoi/boa/eoa ids are resolved via + ``convert_tokens_to_ids`` since only their string forms are on the processor. + """ + + def _build_role_boundaries(self) -> list[RoleBoundary]: + tok = self.processor.tokenizer + end = _encode_markers(tok, [""]) + if not end: + return [] + end_ids = end[0] + boundaries = [] + role_marker_pairs = [ + ("assistant", "model"), + ("user", "user"), + ("system", "system"), + ] + for external_role, template_role in role_marker_pairs: + # Include trailing ``\n`` for consistency with Qwen/Gemma3/Llama + # markers; the newline is part of the marker in the real + # google/gemma-4 tokenizer's chat template. + start = _encode_markers(tok, [f"<|turn>{template_role}\n"]) + if start: + boundaries.append( + RoleBoundary( + role=external_role, + start_tokens=start[0], + end_tokens=end_ids, + ) ) + return boundaries - # Include end token if requested - if include_assistant_end_tok and found_end_seq: - mask[i][end_pos - len(assistant_end_tok) : end_pos] = 1 - - j = end_pos + def process_labels(self, input_ids): + labels = super().process_labels(input_ids) - labels[i][mask[i] == 0] = -100 + tokenizer = self.processor.tokenizer + unk_id = getattr(tokenizer, "unk_token_id", None) + + if getattr(tokenizer, "image_token_id", None) is not None: + labels[labels == tokenizer.image_token_id] = -100 + if getattr(tokenizer, "audio_token_id", None) is not None: + labels[labels == tokenizer.audio_token_id] = -100 + + # boi/eoi/boa/eoa are only string attrs on the processor; resolve ids here. + for attr in ("boi_token", "eoi_token", "boa_token", "eoa_token"): + token_str = getattr(self.processor, attr, None) + if token_str is None: + continue + token_id = tokenizer.convert_tokens_to_ids(token_str) + if token_id is None or token_id == unk_id: + continue + labels[labels == token_id] = -100 + + # Video id lives on the processor, not the tokenizer. + video_token_id = getattr(self.processor, "video_token_id", None) + if video_token_id is not None and video_token_id != unk_id: + labels[labels == video_token_id] = -100 return labels - def process_labels(self, input_ids): - labels = input_ids.clone() - labels = self._mask_non_assistant(labels) - # Follows https://colab.research.google.com/github/huggingface/huggingface-gemma-recipes/blob/main/notebooks/fine_tune_gemma3n_on_t4.ipynb - labels[labels == self.processor.tokenizer.pad_token_id] = -100 - if hasattr(self.processor.tokenizer, "image_token_id"): - labels[labels == self.processor.tokenizer.image_token_id] = -100 - if hasattr(self.processor.tokenizer, "audio_token_id"): - labels[labels == self.processor.tokenizer.audio_token_id] = -100 - if hasattr(self.processor.tokenizer, "boi_token_id"): - labels[labels == self.processor.tokenizer.boi_token_id] = -100 - if hasattr(self.processor.tokenizer, "eoi_token_id"): - labels[labels == self.processor.tokenizer.eoi_token_id] = -100 +class Llama3_2VisionProcessingStrategy(ProcessingStrategy): + """Processing Strategy class for Llama-3.2 Vision (``<|start_header_id|>{role}<|end_header_id|>\\n\\n ... <|eot_id|>``).""" - return labels + def _build_role_boundaries(self) -> list[RoleBoundary]: + tok = self.processor.tokenizer + end = _encode_markers(tok, ["<|eot_id|>"]) + if not end: + return [] + end_ids = end[0] + boundaries = [] + for role in ("system", "user", "assistant", "ipython", "tool"): + start = _encode_markers( + tok, [f"<|start_header_id|>{role}<|end_header_id|>\n\n"] + ) + if start: + boundaries.append( + RoleBoundary(role=role, start_tokens=start[0], end_tokens=end_ids) + ) + return boundaries + + +class Llama4ProcessingStrategy(ProcessingStrategy): + """Processing Strategy class for Llama 4 (``<|header_start|>{role}<|header_end|>\\n\\n ... <|eot|>``).""" + + def _build_role_boundaries(self) -> list[RoleBoundary]: + tok = self.processor.tokenizer + end = _encode_markers(tok, ["<|eot|>"]) + if not end: + return [] + end_ids = end[0] + boundaries = [] + for role in ("system", "user", "assistant", "ipython", "tool"): + start = _encode_markers(tok, [f"<|header_start|>{role}<|header_end|>\n\n"]) + if start: + boundaries.append( + RoleBoundary(role=role, start_tokens=start[0], end_tokens=end_ids) + ) + return boundaries + + +class PixtralProcessingStrategy(ProcessingStrategy): + """Processing Strategy class for Pixtral (``[INST] ... [/INST]`` user, assistant terminates at ``eos_token``). + + ``[/INST]`` is shared between user-end and assistant-start. We declare user + with ``include_end=False`` so the scanner hands the ``[/INST]`` back to + assistant's start match on the next iteration. + """ + + def _build_role_boundaries(self) -> list[RoleBoundary]: + tok = self.processor.tokenizer + eos = getattr(tok, "eos_token_id", None) + if eos is None: + return [] + boundaries = [] + inst_start = _encode_markers(tok, ["[INST]"]) + inst_end = _encode_markers(tok, ["[/INST]"]) + if inst_start and inst_end: + boundaries.append( + RoleBoundary( + role="user", + start_tokens=inst_start[0], + end_tokens=inst_end[0], + include_end=False, + ) + ) + boundaries.append( + RoleBoundary( + role="assistant", + start_tokens=inst_end[0], + end_tokens=[eos], + ) + ) + return boundaries + + +class MistralV7TekkenProcessingStrategy(ProcessingStrategy): + """Processing Strategy class for Mistral v7 Tekken (Pixtral-style plus ``[SYSTEM_PROMPT]...[/SYSTEM_PROMPT]``). + + Same ``[/INST]``-shared-marker treatment as :class:`PixtralProcessingStrategy`. + """ + + def _build_role_boundaries(self) -> list[RoleBoundary]: + tok = self.processor.tokenizer + eos = getattr(tok, "eos_token_id", None) + if eos is None: + return [] + boundaries = [] + sys_start = _encode_markers(tok, ["[SYSTEM_PROMPT]"]) + sys_end = _encode_markers(tok, ["[/SYSTEM_PROMPT]"]) + if sys_start and sys_end: + boundaries.append( + RoleBoundary( + role="system", start_tokens=sys_start[0], end_tokens=sys_end[0] + ) + ) + inst_start = _encode_markers(tok, ["[INST]"]) + inst_end = _encode_markers(tok, ["[/INST]"]) + if inst_start and inst_end: + boundaries.append( + RoleBoundary( + role="user", + start_tokens=inst_start[0], + end_tokens=inst_end[0], + include_end=False, + ) + ) + boundaries.append( + RoleBoundary( + role="assistant", + start_tokens=inst_end[0], + end_tokens=[eos], + ) + ) + return boundaries class VoxtralProcessingStrategy(ProcessingStrategy): - """Processing Strategy class for Voxtral""" + """Processing Strategy class for Voxtral. + + Role boundaries NOT declared — mistral-common instruct tokenizer markers + unverified. Falls back to pad+audio masking with a one-shot warning. + """ def __init__( self, @@ -413,8 +875,21 @@ def __init__( chat_template: Optional[str] = None, image_size: int | tuple[int, int] | None = None, image_resize_algorithm: Resampling | None = None, + train_on_inputs: bool = False, + roles_to_train: Optional[list[str]] = None, + train_on_eos: Optional[str] = None, + role_boundaries_override: Optional[list[dict]] = None, ): - super().__init__(processor, chat_template, image_size, image_resize_algorithm) + super().__init__( + processor, + chat_template, + image_size, + image_resize_algorithm, + train_on_inputs=train_on_inputs, + roles_to_train=roles_to_train, + train_on_eos=train_on_eos, + role_boundaries_override=role_boundaries_override, + ) special_ids = ( processor.tokenizer.tokenizer.instruct_tokenizer.audio_encoder.special_ids ) @@ -424,16 +899,25 @@ def __init__( def process_labels(self, input_ids): labels = input_ids.clone() + labels = self._mask_non_assistant(labels) - labels[labels == self.processor.tokenizer.pad_token_id] = -100 - labels[labels == self.audio_token] = -100 - labels[labels == self.begin_audio_token] = -100 + pad_id = getattr(self.processor.tokenizer, "pad_token_id", None) + if pad_id is not None: + labels[labels == pad_id] = -100 + if self.audio_token is not None: + labels[labels == self.audio_token] = -100 + if self.begin_audio_token is not None: + labels[labels == self.begin_audio_token] = -100 return labels class SmolVLM2ProcessingStrategy(ProcessingStrategy): - """Processing Strategy class for SmolVLM2""" + """Processing Strategy class for SmolVLM2. + + Role boundaries NOT declared — SmolVLM2 chat_template varies per checkpoint + (HuggingFaceTB ships multiple variants), so we opt out rather than mis-mask. + """ def __init__( self, @@ -441,8 +925,21 @@ def __init__( chat_template: Optional[str] = None, image_size: int | tuple[int, int] | None = None, image_resize_algorithm: Resampling | None = None, + train_on_inputs: bool = False, + roles_to_train: Optional[list[str]] = None, + train_on_eos: Optional[str] = None, + role_boundaries_override: Optional[list[dict]] = None, ): - super().__init__(processor, chat_template, image_size, image_resize_algorithm) + super().__init__( + processor, + chat_template, + image_size, + image_resize_algorithm, + train_on_inputs=train_on_inputs, + roles_to_train=roles_to_train, + train_on_eos=train_on_eos, + role_boundaries_override=role_boundaries_override, + ) self.image_token = "" # nosec self.image_token_id = processor.tokenizer.additional_special_tokens_ids[ @@ -451,7 +948,11 @@ def __init__( class Mistral3ProcessingStrategy(ProcessingStrategy): - """Processing Strategy class for Mistral3""" + """Processing Strategy class for Mistral3. + + Role boundaries NOT declared (mistral-common instruct tokenizer unverified); + same fallback as VoxtralProcessingStrategy. + """ def __init__( self, @@ -459,8 +960,21 @@ def __init__( chat_template: Optional[str] = None, image_size: int | tuple[int, int] | None = None, image_resize_algorithm: Resampling | None = None, + train_on_inputs: bool = False, + roles_to_train: Optional[list[str]] = None, + train_on_eos: Optional[str] = None, + role_boundaries_override: Optional[list[dict]] = None, ): - super().__init__(processor, chat_template, image_size, image_resize_algorithm) + super().__init__( + processor, + chat_template, + image_size, + image_resize_algorithm, + train_on_inputs=train_on_inputs, + roles_to_train=roles_to_train, + train_on_eos=train_on_eos, + role_boundaries_override=role_boundaries_override, + ) special_ids = ( processor.tokenizer.tokenizer.instruct_tokenizer.image_encoder.special_ids ) @@ -471,17 +985,24 @@ def __init__( def process_labels(self, input_ids): labels = input_ids.clone() + labels = self._mask_non_assistant(labels) - labels[labels == self.processor.tokenizer.pad_token_id] = -100 - labels[labels == self.image_token] = -100 - labels[labels == self.image_break_token] = -100 - labels[labels == self.image_end_token] = -100 + pad_id = getattr(self.processor.tokenizer, "pad_token_id", None) + if pad_id is not None: + labels[labels == pad_id] = -100 + for tok_id in (self.image_token, self.image_break_token, self.image_end_token): + if tok_id is not None: + labels[labels == tok_id] = -100 return labels class InternVLProcessingStrategy(ProcessingStrategy): - """Processing Strategy class for InternVL""" + """Processing Strategy class for InternVL. + + Role boundaries NOT declared (InternLM-style template unverified); falls + back to pad + image-id masking with a one-shot warning. + """ def __init__( self, @@ -489,8 +1010,21 @@ def __init__( chat_template: Optional[str] = None, image_size: int | tuple[int, int] | None = None, image_resize_algorithm: Resampling | None = None, + train_on_inputs: bool = False, + roles_to_train: Optional[list[str]] = None, + train_on_eos: Optional[str] = None, + role_boundaries_override: Optional[list[dict]] = None, ): - super().__init__(processor, chat_template, image_size, image_resize_algorithm) + super().__init__( + processor, + chat_template, + image_size, + image_resize_algorithm, + train_on_inputs=train_on_inputs, + roles_to_train=roles_to_train, + train_on_eos=train_on_eos, + role_boundaries_override=role_boundaries_override, + ) if not hasattr(processor, "image_ids"): raise ValueError("'image_ids' missing from InternVL Processor.") @@ -499,20 +1033,26 @@ def __init__( def process_labels(self, input_ids): labels = input_ids.clone() + labels = self._mask_non_assistant(labels) - labels[labels == self.processor.tokenizer.pad_token_id] = -100 + pad_id = getattr(self.processor.tokenizer, "pad_token_id", None) + if pad_id is not None: + labels[labels == pad_id] = -100 for ids in self.image_token_ids: - labels[labels == ids] = -100 - - # Note: Check if need to mask 'video_token' as it gets converted to - # image patches during media processing + if ids is not None: + labels[labels == ids] = -100 + # Video tokens get converted to image patches during media processing; masking may be redundant. return labels class Glm4vProcessingStrategy(ProcessingStrategy): - """Processing Strategy class for GLM4V and GLM4V-MoE vision models.""" + """Shared strategy for Glm4vProcessor (GLM-4V / GLM-4.1V) and + Glm46VProcessor (GLM-4.6V / GLM-4.7V) — identical media-token markers. + + Role boundaries unverified; use cfg.role_boundaries to enable masking. + """ def __init__( self, @@ -520,8 +1060,21 @@ def __init__( chat_template: Optional[str] = None, image_size: int | tuple[int, int] | None = None, image_resize_algorithm: Resampling | None = None, + train_on_inputs: bool = False, + roles_to_train: Optional[list[str]] = None, + train_on_eos: Optional[str] = None, + role_boundaries_override: Optional[list[dict]] = None, ): - super().__init__(processor, chat_template, image_size, image_resize_algorithm) + super().__init__( + processor, + chat_template, + image_size, + image_resize_algorithm, + train_on_inputs=train_on_inputs, + roles_to_train=roles_to_train, + train_on_eos=train_on_eos, + role_boundaries_override=role_boundaries_override, + ) self.tokenizer = getattr(processor, "tokenizer", processor) @@ -549,16 +1102,22 @@ def __init__( def process_labels(self, input_ids): labels = input_ids.clone() + labels = self._mask_non_assistant(labels) - labels[labels == self.tokenizer.pad_token_id] = -100 - - labels[labels == self.image_token_id] = -100 - labels[labels == self.begin_image_token_id] = -100 - labels[labels == self.end_image_token_id] = -100 - - labels[labels == self.video_token_id] = -100 - labels[labels == self.begin_video_token_id] = -100 - labels[labels == self.end_video_token_id] = -100 + pad_id = getattr(self.tokenizer, "pad_token_id", None) + if pad_id is not None: + labels[labels == pad_id] = -100 + + for tok_id in ( + self.image_token_id, + self.begin_image_token_id, + self.end_image_token_id, + self.video_token_id, + self.begin_video_token_id, + self.end_video_token_id, + ): + if tok_id is not None: + labels[labels == tok_id] = -100 return labels @@ -569,14 +1128,20 @@ def get_processing_strategy( chat_template_type, image_size: int | tuple[int, int] | None = None, image_resize_algorithm: Resampling | None = None, + train_on_inputs: bool = False, + roles_to_train: Optional[list[str]] = None, + train_on_eos: Optional[str] = None, + role_boundaries_override: Optional[list[dict]] = None, ): - from axolotl.utils.mistral.mistral3_processor import Mistral3Processor - processing_kwargs = { "processor": processor, "chat_template": chat_template, "image_size": image_size, "image_resize_algorithm": image_resize_algorithm, + "train_on_inputs": train_on_inputs, + "roles_to_train": roles_to_train, + "train_on_eos": train_on_eos, + "role_boundaries_override": role_boundaries_override, } if chat_template_type in [None, "tokenizer_default"]: @@ -585,53 +1150,63 @@ def get_processing_strategy( processing_kwargs["chat_template"] = tokenizer.chat_template if chat_template_type == "qwen2_vl": - return Qwen2VLProcessingStrategy( - **processing_kwargs, - ) - if chat_template_type in ["qwen3_5", "qwen3_5_moe"]: - return Qwen3_5ProcessingStrategy( - **processing_kwargs, - ) + return Qwen2VLProcessingStrategy(**processing_kwargs) + if chat_template_type == "qwen3_5": + return Qwen3_5ProcessingStrategy(**processing_kwargs) if chat_template_type == "gemma3": - return Gemma3ProcessingStrategy( - **processing_kwargs, - ) + return Gemma3ProcessingStrategy(**processing_kwargs) if chat_template_type == "gemma3n": - return Gemma3nProcessingStrategy( - **processing_kwargs, - ) + return Gemma3nProcessingStrategy(**processing_kwargs) + if chat_template_type == "gemma4": + return Gemma4ProcessingStrategy(**processing_kwargs) + if chat_template_type == "llama3_2_vision": + return Llama3_2VisionProcessingStrategy(**processing_kwargs) + if chat_template_type == "llama4": + return Llama4ProcessingStrategy(**processing_kwargs) + if chat_template_type == "pixtral": + return PixtralProcessingStrategy(**processing_kwargs) + if chat_template_type == "mistral_v7_tekken": + return MistralV7TekkenProcessingStrategy(**processing_kwargs) if isinstance(processor, VoxtralProcessor): - return VoxtralProcessingStrategy( - **processing_kwargs, - ) + return VoxtralProcessingStrategy(**processing_kwargs) if isinstance(processor, SmolVLMProcessor): - return SmolVLM2ProcessingStrategy( - **processing_kwargs, - ) + return SmolVLM2ProcessingStrategy(**processing_kwargs) - if isinstance(processor, Mistral3Processor): - return Mistral3ProcessingStrategy( - **processing_kwargs, + # Lazy import: mistral_common is optional. Mirrors the Glm46V pattern below. + try: + from axolotl.utils.mistral.mistral3_processor import Mistral3Processor + + if isinstance(processor, Mistral3Processor): + return Mistral3ProcessingStrategy(**processing_kwargs) + except (ImportError, ModuleNotFoundError) as exc: + LOG.debug( + "Mistral3Processor import failed; Mistral3 strategy will be unavailable: %r", + exc, ) + + # Both Glm4vProcessor and Glm46VProcessor share markers; route to the same + # strategy. Independent try/except so either can be absent. + try: + from transformers.models.glm4v.processing_glm4v import Glm4vProcessor + + if isinstance(processor, Glm4vProcessor): + return Glm4vProcessingStrategy(**processing_kwargs) + except (ImportError, ModuleNotFoundError) as exc: + LOG.debug("Glm4vProcessor import failed: %r", exc) + try: from transformers.models.glm46v.processing_glm46v import Glm46VProcessor if isinstance(processor, Glm46VProcessor): - return Glm4vProcessingStrategy( - **processing_kwargs, - ) - except ImportError: - pass + return Glm4vProcessingStrategy(**processing_kwargs) + except (ImportError, ModuleNotFoundError) as exc: + LOG.debug("Glm46VProcessor import failed: %r", exc) if isinstance(processor, InternVLProcessor): - return InternVLProcessingStrategy( - **processing_kwargs, - ) + return InternVLProcessingStrategy(**processing_kwargs) - # llama3_2_vision, llama4, llava - # mistral_v7_tekken, pixtral, lfm2vl - return ProcessingStrategy( - **processing_kwargs, - ) + # Unregistered templates (llava, lfm2vl, mistral_v3_tekken, ...) use the + # base strategy; it warns once when train_on_inputs=False. + return ProcessingStrategy(**processing_kwargs) diff --git a/src/axolotl/prompt_strategies/multimodal_pretrain.py b/src/axolotl/prompt_strategies/multimodal_pretrain.py new file mode 100644 index 0000000000..c716c63aa3 --- /dev/null +++ b/src/axolotl/prompt_strategies/multimodal_pretrain.py @@ -0,0 +1,297 @@ +"""Multimodal CPT tokenization strategy.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +from transformers import BatchEncoding, PreTrainedTokenizerBase, ProcessorMixin + +from axolotl.prompt_strategies.pretrain import ( + PretrainTokenizationStrategy, + PretrainTokenizer, +) +from axolotl.utils.logging import get_logger + +LOG = get_logger(__name__) + + +def _get_incompatible_processor_classes() -> tuple[type, ...]: + classes: list[type] = [] + for mod_path, name in ( + ("transformers.models.mllama", "MllamaProcessor"), + ("transformers.models.pixtral", "PixtralProcessor"), + ("transformers.models.internvl", "InternVLProcessor"), + ): + try: + import importlib + + mod = importlib.import_module(mod_path) + cls = getattr(mod, name, None) + if cls is not None: + classes.append(cls) + except ImportError: + continue + return tuple(classes) + + +_KNOWN_IMAGE_TOKEN_CANDIDATES: tuple[str, ...] = ( + "", + "<|image|>", + "<|image_pad|>", + "", + "", + "[IMG]", + "", +) + +# Without masking these in labels, loss blows up ~10× on Qwen/SmolVLM. +_IMAGE_FAMILY_TOKEN_CANDIDATES: tuple[str, ...] = ( + "", + "<|image|>", + "<|image_pad|>", + "", + "", + "", + "<|vision_start|>", + "<|vision_end|>", + "[IMG]", + "[IMG_END]", + "", +) + +_INCOMPATIBLE_PROCESSOR_REASONS: dict[str, str] = { + "MllamaProcessor": ( + "Llama-3.2-Vision (Mllama) uses cross-attention image injection, not " + "in-stream placeholder tokens. Multimodal CPT is incompatible with " + "this architecture; use chat-template SFT instead." + ), + "PixtralProcessor": ( + "Pixtral's tokenizer goes through mistral_common with a different " + "API surface than AutoProcessor. Multimodal CPT not supported in v1; " + "use chat-template SFT or Mistral-Small-3.1." + ), + "InternVLProcessor": ( + "InternVL ships a custom processing pipeline (AutoProcessor returns " + "text-only); no pixel_values are produced. Multimodal CPT not " + "supported in v1." + ), +} +_INCOMPATIBLE_PROCESSOR_CLASSES = _get_incompatible_processor_classes() + + +@dataclass +class ImageTokenSpec: + image_token: str + image_token_id: int + image_family_token_ids: set[int] + + +def build_image_token_spec( + processor: ProcessorMixin, override: str | None = None +) -> ImageTokenSpec: + tokenizer = getattr(processor, "tokenizer", None) + if tokenizer is None: + raise ValueError( + "Processor has no `tokenizer` attribute — multimodal CPT " + "requires a processor with a text tokenizer (e.g. one produced " + "by AutoProcessor.from_pretrained for a VLM)." + ) + + def resolve_id(tok: str) -> int | None: + tid = tokenizer.convert_tokens_to_ids(tok) + unk = getattr(tokenizer, "unk_token_id", None) + if tid is None or tid == unk: + return None + return tid + + known_special_tokens: set[str] = set() + try: + known_special_tokens |= set(tokenizer.get_added_vocab().keys()) + except Exception: + pass + known_special_tokens |= set(getattr(tokenizer, "all_special_tokens", None) or []) + known_special_tokens |= set( + getattr(tokenizer, "additional_special_tokens", None) or [] + ) + + image_token: str | None = None + image_token_id: int | None = None + if override is not None: + # Reject plain words that BPE-tokenize cleanly but aren't placeholders. + if override not in known_special_tokens: + raise ValueError( + f"image_token override {override!r} is not a registered " + f"special token on this tokenizer. Pick one of the model's " + f"actual image tokens (e.g. '', '<|image_pad|>', " + f"''), or leave unset to autodetect." + ) + image_token_id = resolve_id(override) + if image_token_id is None: + raise ValueError( + f"image_token override {override!r} did not resolve to a " + f"token id (unk). Remove the override to autodetect." + ) + image_token = override + else: + proc_token = getattr(processor, "image_token", None) + if proc_token is not None: + image_token_id = resolve_id(proc_token) + if image_token_id is not None: + image_token = proc_token + if image_token is None: + for cand in _KNOWN_IMAGE_TOKEN_CANDIDATES: + tid = resolve_id(cand) + if tid is not None: + image_token = cand + image_token_id = tid + break + if image_token is None: + raise ValueError( + "Could not autodetect the image placeholder token for this " + "processor. Set `image_token: ` in the dataset config " + "(e.g. '' for LLaVA, '<|image_pad|>' for Qwen-VL, " + "'' for Gemma-3)." + ) + + # Filter to registered tokens so BPE-fallback ids don't get masked. + family: set[int] = {image_token_id} # type: ignore[arg-type] + for cand in _IMAGE_FAMILY_TOKEN_CANDIDATES: + if cand != image_token and cand not in known_special_tokens: + continue + tid = resolve_id(cand) + if tid is not None: + family.add(tid) + return ImageTokenSpec( + image_token=image_token, + image_token_id=image_token_id, # type: ignore[arg-type] + image_family_token_ids=family, + ) + + +def check_processor_compatibility(processor: ProcessorMixin) -> None: + if _INCOMPATIBLE_PROCESSOR_CLASSES and isinstance( + processor, _INCOMPATIBLE_PROCESSOR_CLASSES + ): + for cls in _INCOMPATIBLE_PROCESSOR_CLASSES: + if isinstance(processor, cls): + raise ValueError( + f"Multimodal CPT is not supported for {cls.__name__}: " + f"{_INCOMPATIBLE_PROCESSOR_REASONS.get(cls.__name__, '')}" + ) + # MRO-name fallback for test fakes and unimportable concrete classes. + for base_cls in type(processor).__mro__: + reason = _INCOMPATIBLE_PROCESSOR_REASONS.get(base_cls.__name__) + if reason is not None: + raise ValueError( + f"Multimodal CPT is not supported for {base_cls.__name__}: {reason}" + ) + + +class MultimodalPretrainTokenizationStrategy(PretrainTokenizationStrategy): + def __init__( + self, + *args: Any, + image_token: str, + image_token_id: int, + image_column: str = "images", + image_base_dir: str | None = None, + **kwargs: Any, + ) -> None: + super().__init__(*args, **kwargs) + self.image_token = image_token + self.image_token_id = image_token_id + self.image_column = image_column + self.image_base_dir = image_base_dir + + def _tokenize( + self, + prompt: str, + add_eos_token: bool = True, + strip_bos_token: bool = False, + ) -> BatchEncoding: + # No truncation: collator re-tokenizes the full text without truncation; + # truncating here decouples the stored ids from what the model receives. + res = self.tokenizer(prompt, add_special_tokens=True) + res["input_ids"] = [res["input_ids"] + [self.tokenizer.eos_token_id]] + res["attention_mask"] = [res["attention_mask"] + [1]] + return res + + def tokenize_prompt(self, prompt: dict[str, Any]) -> dict[str, list]: + text = prompt[self.text_column] + raw_images = prompt.get(self.image_column) + if raw_images is None: + images: list = [] + elif isinstance(raw_images, (list, tuple)): + images = list(raw_images) + else: + raise ValueError( + f"Row's `{self.image_column}` must be a list of image paths, " + f"got {type(raw_images).__name__}." + ) + + res = self._tokenize(text) + ids = res["input_ids"][0] + # Count by token id — `text.count` substring-matches `` in ``. + n_placeholders = sum(1 for t in ids if t == self.image_token_id) + if n_placeholders != len(images): + raise ValueError( + f"Multimodal CPT row has {n_placeholders} occurrence(s) of " + f"{self.image_token!r} in text but {len(images)} image path(s) " + f"in `{self.image_column}`. They must match — the text column " + f"must contain exactly one placeholder per image." + ) + if len(ids) > self.max_length: + raise ValueError( + f"Multimodal CPT row tokenizes to {len(ids)} tokens which " + f"exceeds sequence_len={self.max_length}. Pre-chunk your text " + f"or raise sequence_len." + ) + + n_chunks = len(res["input_ids"]) + res["images"] = [list(images)] * n_chunks + res["_mm_text"] = [text] * n_chunks + return res + + +def load( + tokenizer: PreTrainedTokenizerBase, + cfg: Any, + ds_cfg: dict | None = None, + processor: ProcessorMixin | None = None, +) -> MultimodalPretrainTokenizationStrategy: + if processor is None: + raise ValueError( + "multimodal_pretrain requires a processor. Set `processor_type: " + "AutoProcessor` (or the concrete processor class) in your config " + "so axolotl loads it at startup." + ) + check_processor_compatibility(processor) + + ds_cfg = dict(ds_cfg or {}) + text_column = ds_cfg.get("text_column") or ds_cfg.get("field") or "text" + image_column = ds_cfg.get("image_column") or "images" + image_base_dir = ds_cfg.get("image_base_dir") + image_token_override = ds_cfg.get("image_token") + + spec = build_image_token_spec(processor, override=image_token_override) + LOG.info( + f"multimodal_pretrain: placeholder={spec.image_token!r} " + f"(id={spec.image_token_id}), masking {len(spec.image_family_token_ids)} " + f"image-family token ids in labels" + ) + + strat = MultimodalPretrainTokenizationStrategy( + PretrainTokenizer(), + tokenizer, + cfg.train_on_inputs, + cfg.sequence_len, + text_column=text_column, + image_column=image_column, + image_base_dir=image_base_dir, + image_token=spec.image_token, + image_token_id=spec.image_token_id, + max_length=cfg.sequence_len, + ) + strat.image_token_spec = spec # type: ignore[attr-defined] + return strat diff --git a/src/axolotl/utils/collators/mm_pretrain.py b/src/axolotl/utils/collators/mm_pretrain.py new file mode 100644 index 0000000000..f6870434f4 --- /dev/null +++ b/src/axolotl/utils/collators/mm_pretrain.py @@ -0,0 +1,306 @@ +"""Collator for multimodal CPT.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass, field +from typing import Any, Literal, Optional, Union + +from PIL import Image +from torch import Tensor +from transformers import PreTrainedTokenizerBase, ProcessorMixin +from transformers.data.data_collator import DataCollatorMixin +from transformers.utils import PaddingStrategy + +from axolotl.prompt_strategies.multimodal_pretrain import ( + ImageTokenSpec, + check_processor_compatibility, +) +from axolotl.utils.logging import get_logger + +LOG = get_logger(__name__) + +# Decompression-bomb cap (~7070×7070). +_DEFAULT_MAX_IMAGE_PIXELS = 50_000_000 +_DEFAULT_MAX_IMAGES_PER_ROW = 32 + + +@dataclass +class MultiModalPretrainDataCollator(DataCollatorMixin): + tokenizer: PreTrainedTokenizerBase + processor: ProcessorMixin + image_token_spec: ImageTokenSpec + image_base_dir: Optional[str] = None + return_tensors: Literal["pt"] = "pt" + padding: Union[bool, str, PaddingStrategy] = True + pad_to_multiple_of: Optional[int] = None + max_length: Optional[int] = None + skip_bad_images: bool = False + max_image_pixels: int = _DEFAULT_MAX_IMAGE_PIXELS + max_images_per_row: int = _DEFAULT_MAX_IMAGES_PER_ROW + + _image_family_token_ids: set[int] = field(init=False, default_factory=set) + _base_dir_real: Optional[str] = field(init=False, default=None) + + def __post_init__(self) -> None: + if self.return_tensors != "pt": + raise ValueError( + "MultiModalPretrainDataCollator only supports " + "return_tensors='pt' (in-place torch ops are used downstream)." + ) + check_processor_compatibility(self.processor) + self._image_family_token_ids = set(self.image_token_spec.image_family_token_ids) + if self.image_base_dir is not None: + self._base_dir_real = os.path.realpath(self.image_base_dir) + + def _resolve_image_path(self, p: str) -> str: + if not isinstance(p, str): + raise ValueError(f"Image path must be str, got {type(p).__name__}.") + if "\x00" in p: + raise ValueError("Image path contains embedded NUL byte.") + p_lower = p.lower() + if p_lower.startswith( + ("http://", "https://", "ftp://", "ftps://", "file://", "data:") + ) or p.startswith(("\\\\", "//")): + raise ValueError( + f"Non-local image path scheme is not supported in v1 " + f"multimodal CPT (got {p!r})." + ) + if self._base_dir_real is not None: + if os.path.isabs(p): + raise ValueError( + f"Absolute image path {p!r} is rejected when " + f"`image_base_dir` is configured. All image paths must be " + f"relative to the configured base directory." + ) + resolved = os.path.realpath(os.path.join(self._base_dir_real, p)) + # commonpath (not startswith) so root-dir bases like "/" work. + try: + within_base = ( + os.path.commonpath([self._base_dir_real, resolved]) + == self._base_dir_real + ) + except ValueError: + within_base = False + if not within_base: + raise ValueError( + f"Image path {p!r} resolves outside `image_base_dir` " + f"after symlink resolution. Refusing to load." + ) + return resolved + return os.path.realpath(p) if os.path.isabs(p) else p + + def _open_image_hardened(self, resolved: str) -> Image.Image: + # O_NOFOLLOW closes the realpath→open TOCTOU window for the final component. + nofollow = getattr(os, "O_NOFOLLOW", 0) + try: + fd = os.open(resolved, os.O_RDONLY | nofollow) + except OSError as exc: + raise ValueError( + f"Cannot open image (os.open failed: {type(exc).__name__})." + ) from exc + file_obj = os.fdopen(fd, "rb") + try: + with Image.open(file_obj) as src: + w, h = src.size + if w * h > self.max_image_pixels: + raise ValueError( + f"Image pixels ({w}×{h}) exceed " + f"max_image_pixels ({self.max_image_pixels})." + ) + # Multi-frame bomb guard (GIF/TIFF/WebP). + n_frames = getattr(src, "n_frames", 1) + if n_frames > 1: + raise ValueError( + f"Multi-frame images are not supported (got {n_frames} frames)." + ) + img = src.convert("RGB") + img.load() + return img + finally: + if not file_obj.closed: + file_obj.close() + + def _load_images_for_row( + self, paths: list[str], row_index: int + ) -> list[Image.Image]: + if len(paths) > self.max_images_per_row: + raise ValueError( + f"Row {row_index}: {len(paths)} images exceeds " + f"`max_images_per_row={self.max_images_per_row}`. Split the " + f"row or raise the cap if this is expected." + ) + out: list[Image.Image] = [] + for raw in paths: + try: + resolved = self._resolve_image_path(raw) + img = self._open_image_hardened(resolved) + except Exception as exc: + # Top-level log gets basename only; full path stays on DEBUG. + basename = os.path.basename(str(raw)) + msg = ( + f"Row {row_index}: failed to load image {basename!r} " + f"({type(exc).__name__})" + ) + LOG.debug("failed image full path: %r; error: %s", raw, exc) + if self.skip_bad_images: + LOG.warning("%s — skipping", msg) + continue + raise RuntimeError(msg) from exc + out.append(img) + return out + + def torch_call(self, examples: list[dict]) -> dict[str, Any]: + if not examples: + raise ValueError("Empty batch passed to MultiModalPretrainDataCollator.") + + texts: list[str] = [] + images: list[list[Image.Image]] = [] + for i, ex in enumerate(examples): + if "_mm_text" not in ex or "images" not in ex: + raise KeyError( + f"MultiModalPretrainDataCollator: row {i} is missing " + f"'_mm_text' or 'images'. Did you wire the multimodal CPT " + f"encoder (encode_streaming_multimodal or " + f"MultimodalPretrainTokenizationStrategy)?" + ) + mm_text = ex["_mm_text"] + if not isinstance(mm_text, str): + raise TypeError( + f"Row {i}: `_mm_text` must be str, got " + f"{type(mm_text).__name__}. Check dataset encoding " + f"(Parquet BINARY columns may surface as bytes)." + ) + raw = ex["images"] + if raw is None: + raw_paths: list[str] = [] + elif isinstance(raw, (list, tuple)): + raw_paths = list(raw) + else: + raise TypeError( + f"Row {i}: `images` must be a list (or None), got " + f"{type(raw).__name__}." + ) + for j, rp in enumerate(raw_paths): + if not isinstance(rp, str): + raise TypeError( + f"Row {i}, image {j}: path must be str, got " + f"{type(rp).__name__}." + ) + texts.append(mm_text) + loaded = self._load_images_for_row(raw_paths, row_index=i) + if self.skip_bad_images and len(loaded) != len(raw_paths): + # Drop the row to avoid silent placeholder/image count mismatch. + LOG.warning( + "Row %d: %d/%d images failed to load; dropping row.", + i, + len(raw_paths) - len(loaded), + len(raw_paths), + ) + texts.pop() + continue + images.append(loaded) + + if not texts: + raise RuntimeError( + "All rows in the batch were dropped due to image load " + "failures. Check dataset integrity." + ) + + # All-text batch: bypass the processor and tokenize directly. + if all(len(im) == 0 for im in images): + LOG.debug( + "MultiModalPretrainDataCollator: all-text batch (%d rows); " + "using tokenizer-only fallback (no pixel_values).", + len(texts), + ) + tok_kwargs: dict[str, Any] = { + "text": texts, + "return_tensors": self.return_tensors, + "padding": self.padding, + } + if self.pad_to_multiple_of is not None: + tok_kwargs["pad_to_multiple_of"] = self.pad_to_multiple_of + batch = self.tokenizer(**tok_kwargs) + tok_input_ids: Tensor = batch["input_ids"] + tok_labels = tok_input_ids.clone() + pad_id = getattr(self.tokenizer, "pad_token_id", None) + if pad_id is not None: + tok_labels[tok_labels == pad_id] = -100 + for tid in self._image_family_token_ids: + tok_labels[tok_labels == tid] = -100 + batch["labels"] = tok_labels + return dict(batch) + + # No truncation: it chops input_ids mid-placeholder while pixel_values + # keep every image — silent text/pixel mismatch. We warn post-hoc instead. + proc_kwargs: dict[str, Any] = { + "text": texts, + "images": images, + "return_tensors": self.return_tensors, + "padding": self.padding, + } + if self.pad_to_multiple_of is not None: + proc_kwargs["pad_to_multiple_of"] = self.pad_to_multiple_of + try: + batch = self.processor(**proc_kwargs) + except Exception as exc: + # Pinpoint the bad row; bail to "inconclusive" if retry raises a different class. + offender_idx: Optional[int] = None + retry_ok = True + retry_kwargs: dict[str, Any] = { + "return_tensors": self.return_tensors, + "padding": self.padding, + } + if self.pad_to_multiple_of is not None: + retry_kwargs["pad_to_multiple_of"] = self.pad_to_multiple_of + for i, (t, imgs) in enumerate(zip(texts, images, strict=True)): + try: + self.processor(text=[t], images=[imgs], **retry_kwargs) + except Exception as retry_exc: + if isinstance(retry_exc, type(exc)) or isinstance( + exc, type(retry_exc) + ): + offender_idx = i + else: + retry_ok = False + break + if offender_idx is not None: + location = f"row {offender_idx}" + elif retry_ok: + location = ( + f"batch of {len(texts)} rows " + f"(individual rows all succeed; see __cause__ for details)" + ) + else: + location = f"batch of {len(texts)} rows (retry inconclusive)" + raise RuntimeError( + f"MultiModalPretrainDataCollator: processor call failed on " + f"{location} ({type(exc).__name__}: {exc}). Common causes: " + f"placeholder token absent from the row's text, image count " + f"mismatch, or an unsupported processor class." + ) from exc + + input_ids_len = batch["input_ids"].shape[-1] + if self.max_length is not None and input_ids_len > self.max_length: + LOG.warning( + "Batch input_ids length %d exceeds configured sequence_len %d " + "(image placeholder expansion). Reduce max_images_per_row or " + "raise sequence_len if this fires repeatedly.", + input_ids_len, + self.max_length, + ) + + input_ids: Tensor = batch["input_ids"] + labels = input_ids.clone() + + pad_id = getattr(self.tokenizer, "pad_token_id", None) + if pad_id is not None: + labels[labels == pad_id] = -100 + + # Without this, image-family ids dominate loss and blow it up ~10×. + for tid in self._image_family_token_ids: + labels[labels == tid] = -100 + + batch["labels"] = labels + return batch diff --git a/src/axolotl/utils/data/sft.py b/src/axolotl/utils/data/sft.py index 0b2ec2b5fb..809d08ac5f 100644 --- a/src/axolotl/utils/data/sft.py +++ b/src/axolotl/utils/data/sft.py @@ -10,6 +10,7 @@ DatasetDict, IterableDataset, IterableDatasetDict, + concatenate_datasets, load_dataset, ) from transformers import PreTrainedTokenizer, ProcessorMixin @@ -134,7 +135,9 @@ def _prepare_streaming_dataset( """ if cfg.pretraining_dataset: dataset_config = _extract_pretraining_config(cfg) - train_dataset = _load_streaming_dataset(dataset_config, cfg, tokenizer) + train_dataset = _load_streaming_dataset( + dataset_config, cfg, tokenizer, processor=processor + ) elif cfg.sample_packing: # TODO(djsaunde): Implement for multiple datasets dataset_config = DictDefault(cfg.datasets[0]) @@ -142,7 +145,9 @@ def _prepare_streaming_dataset( # Ensure we have a split set - default to 'train' if not specified if not hasattr(dataset_config, "split") or not dataset_config.split: dataset_config.split = "train" - train_dataset = _load_streaming_dataset(dataset_config, cfg, tokenizer) + train_dataset = _load_streaming_dataset( + dataset_config, cfg, tokenizer, processor=processor + ) else: # Use legacy loading function for non-packed streaming datasets train_dataset, eval_dataset, prompters = _load_and_prepare_datasets( @@ -160,35 +165,73 @@ def _prepare_streaming_dataset( # Load evaluation dataset if specified eval_dataset = None if cfg.test_datasets: - _, eval_dataset, _ = _load_and_prepare_datasets( - tokenizer, - cfg, - split="test", - processor=processor, - streaming=False, + test_dicts = [t if isinstance(t, dict) else dict(t) for t in cfg.test_datasets] + is_mm_cpt_eval = any( + t.get("type") == "multimodal_pretrain" or bool(t.get("multimodal")) + for t in test_dicts ) + if is_mm_cpt_eval: + eval_streams = [] + for entry in test_dicts: + if not ( + entry.get("type") == "multimodal_pretrain" + or bool(entry.get("multimodal")) + ): + raise ValueError( + "Mixing multimodal and non-multimodal entries in " + "`test_datasets` is not supported. All eval entries " + "must be MM (type: multimodal_pretrain or " + "multimodal: true) when training is MM CPT." + ) + eval_config = _pretraining_config_from_entry(entry) + eval_streams.append( + _load_streaming_dataset( + eval_config, cfg, tokenizer, processor=processor + ) + ) + eval_dataset = ( + eval_streams[0] + if len(eval_streams) == 1 + else concatenate_datasets(eval_streams) + ) + else: + _, eval_dataset, _ = _load_and_prepare_datasets( + tokenizer, + cfg, + split="test", + processor=processor, + streaming=False, + ) # For streaming, we return max_steps directly from config or -1 if not set total_num_steps = cfg.max_steps if cfg.max_steps else -1 return train_dataset, eval_dataset, total_num_steps, [] +def _pretraining_config_from_entry(entry: dict) -> DictDefault: + return DictDefault( + { + "path": entry["path"], + "name": entry.get("name"), + "skip": entry.get("skip"), + "split": entry.get("split", "train"), + "data_files": entry.get("data_files"), + "type": entry.get("type", "pretrain"), + "text_column": entry.get("text_column", "text"), + "multimodal": entry.get("multimodal"), + "image_column": entry.get("image_column", "images"), + "image_base_dir": entry.get("image_base_dir"), + "image_token": entry.get("image_token"), + } + ) + + def _extract_pretraining_config(cfg: DictDefault) -> DictDefault: """Extract pretraining configuration from the main config.""" if isinstance(cfg.pretraining_dataset, list) and isinstance( cfg.pretraining_dataset[0], dict ): - config = cfg.pretraining_dataset[0] - return DictDefault( - { - "path": config["path"], - "name": config["name"], - "skip": config["skip"], - "split": config.get("split", "train"), - "data_files": config.get("data_files"), - "type": config.get("type", "pretrain"), - } - ) + return _pretraining_config_from_entry(cfg.pretraining_dataset[0]) # Simple string path case return DictDefault( { @@ -198,12 +241,20 @@ def _extract_pretraining_config(cfg: DictDefault) -> DictDefault: "split": "train", "data_files": None, "type": "pretrain", + "text_column": "text", + "multimodal": None, + "image_column": "images", + "image_base_dir": None, + "image_token": None, # nosec } ) def _load_streaming_dataset( - pretraining_config: DictDefault, cfg: DictDefault, tokenizer: PreTrainedTokenizer + pretraining_config: DictDefault, + cfg: DictDefault, + tokenizer: PreTrainedTokenizer, + processor: ProcessorMixin | None = None, ) -> IterableDataset: """Load and prepare a streaming dataset for pretraining.""" # Create dataset wrapper partial function @@ -213,6 +264,7 @@ def _load_streaming_dataset( tokenizer=tokenizer, cfg=cfg, dataset_base_type=pretraining_config["type"], + processor=processor, ) # Load the actual dataset @@ -221,7 +273,7 @@ def _load_streaming_dataset( and cfg.accelerator_config.dispatch_batches and not is_local_main_process() ): - iter_dataset = _create_placeholder_dataset() + iter_dataset = _create_placeholder_dataset(pretraining_config) else: iter_dataset = load_dataset( pretraining_config["path"], @@ -242,19 +294,39 @@ def _load_streaming_dataset( tokenizer, cfg, dataset_wrapper_partial, + processor=processor, + pretraining_config=pretraining_config, ) # Format for PyTorch return train_dataset.with_format("torch") -def _create_placeholder_dataset() -> IterableDataset: +def _create_placeholder_dataset( + pretraining_config: DictDefault | None = None, +) -> IterableDataset: """Create a minimal placeholder dataset for non-main processes.""" - with tempfile.NamedTemporaryFile(mode="w+", delete=False) as f: - f.write("text\n") - f.write("lorem ipsum dolor sit amet\n") - f.seek(0) - return load_dataset("csv", data_files=f.name, split="train", streaming=True) + text_column = "text" + image_column: str | None = None + if pretraining_config is not None: + text_column = pretraining_config.get("text_column") or "text" + is_mm = pretraining_config.get("type") == "multimodal_pretrain" or bool( + pretraining_config.get("multimodal") + ) + if is_mm: + image_column = pretraining_config.get("image_column") or "images" + + if image_column is None: + with tempfile.NamedTemporaryFile(mode="w+", delete=False) as f: + f.write(f"{text_column}\n") + f.write("lorem ipsum dolor sit amet\n") + f.seek(0) + return load_dataset("csv", data_files=f.name, split="train", streaming=True) + + def _gen(): + yield {text_column: "lorem ipsum dolor sit amet", image_column: []} + + return IterableDataset.from_generator(_gen) def _load_tokenized_prepared_datasets( diff --git a/src/axolotl/utils/data/streaming.py b/src/axolotl/utils/data/streaming.py index 8b6b8a439b..966fa65719 100644 --- a/src/axolotl/utils/data/streaming.py +++ b/src/axolotl/utils/data/streaming.py @@ -7,7 +7,7 @@ import torch from datasets import Dataset from torch.utils.data import RandomSampler -from transformers import PreTrainedTokenizerBase +from transformers import PreTrainedTokenizerBase, ProcessorMixin from axolotl.utils.collators import PretrainingBatchSamplerDataCollatorForSeq2Seq from axolotl.utils.logging import get_logger @@ -176,11 +176,93 @@ def encode_streaming( return ret +def encode_streaming_multimodal( + examples: Dict[str, List], + tokenizer: PreTrainedTokenizerBase, + max_tokens: int, + image_token: str, + image_token_id: int, + text_column: str = "text", + image_column: str = "images", +) -> Dict[str, List]: + texts: List[str] = examples[text_column] + imgs_list: List[List[str]] = examples[image_column] + + if len(texts) != len(imgs_list): + raise ValueError( + f"encode_streaming_multimodal: text column has {len(texts)} rows " + f"but image column has {len(imgs_list)}" + ) + + input_ids: List[List[int]] = [] + labels: List[List[int]] = [] + attention_mask: List[List[int]] = [] + keep_images: List[List[str]] = [] + keep_text: List[str] = [] + + for text, imgs in zip(texts, imgs_list, strict=True): + if not isinstance(text, str): + raise TypeError( + f"encode_streaming_multimodal: `{text_column}` must be str, " + f"got {type(text).__name__}." + ) + if imgs is None: + imgs = [] + if not isinstance(imgs, (list, tuple)): + raise ValueError( + f"encode_streaming_multimodal: row's `{image_column}` must be " + f"a list; got {type(imgs).__name__}" + ) + for j, ip in enumerate(imgs): + if not isinstance(ip, str): + raise TypeError( + f"encode_streaming_multimodal: image {j} in row must be " + f"str, got {type(ip).__name__}." + ) + # No truncation: counting on truncated ids and storing untruncated text + # (which the collator re-tokenizes without truncation) silently produces + # oversize batches and confusing placeholder/image-count mismatches. + enc = tokenizer(text, add_special_tokens=True) + ids = list(enc["input_ids"]) + [tokenizer.eos_token_id] + mask = list(enc["attention_mask"]) + [1] + # Count by id — `text.count` substring-matches `` in ``. + n_placeholders = sum(1 for t in ids if t == image_token_id) + if n_placeholders != len(imgs): + raise ValueError( + f"Multimodal CPT row has {n_placeholders} occurrence(s) of " + f"{image_token!r} in text but {len(imgs)} image path(s). " + f"Text and image count must match (one placeholder per image)." + ) + if len(ids) > max_tokens: + raise ValueError( + f"Multimodal CPT row tokenizes to {len(ids)} tokens which " + f"exceeds sequence_len={max_tokens}. Pre-chunk your text or " + f"raise sequence_len (image patch expansion at the processor " + f"may push the final length even higher)." + ) + # Labels = ids; collator masks image-family ids after re-tokenization. + input_ids.append(ids) + labels.append(list(ids)) + attention_mask.append(mask) + keep_images.append(list(imgs)) + keep_text.append(text) + + return { + "input_ids": input_ids, + "labels": labels, + "attention_mask": attention_mask, + "images": keep_images, + "_mm_text": keep_text, + } + + def wrap_streaming_dataset( dataset, tokenizer, cfg, ds_wrapper_fn, + processor: Optional[ProcessorMixin] = None, + pretraining_config=None, ): if cfg.sample_packing: # For SFT (non-pretraining) datasets, always use multipack_attn=True to ensure @@ -213,17 +295,66 @@ def wrap_streaming_dataset( # NOTE: This is not reachable for SFT datasets since we use the pre-existing # loading function for non-packed streaming datasets. Refer to # _prepare_streaming_datasets in sft.py for that code path. - text_column = ( - getattr(cfg.pretraining_dataset[0], "text_column", "text") or "text" + # Prefer the resolved per-entry config so eval (test_datasets) doesn't + # silently inherit the training entry's columns/image_token. + if pretraining_config is not None: + ds_first = pretraining_config + elif cfg.pretraining_dataset: + ds_first = cfg.pretraining_dataset[0] + else: + ds_first = {} + # Plain dicts need `.get`; pydantic/DictDefault need `getattr`. + get_ds_value = ( + ds_first.get + if isinstance(ds_first, dict) + else lambda key, default=None: getattr(ds_first, key, default) ) - encode = functools.partial( - encode_streaming, - tokenizer=tokenizer, - max_tokens=cfg.sequence_len, - text_column=text_column, - concatenate=cfg.pretraining_sample_concatenation is True, + text_column = get_ds_value("text_column", "text") or "text" + ds_type = (get_ds_value("type", None) or "").strip() + is_mm_cpt = ds_type == "multimodal_pretrain" or bool( + get_ds_value("multimodal", False) ) + if is_mm_cpt: + if processor is None: + raise ValueError( + "Multimodal CPT (type: multimodal_pretrain) requires a " + "processor. Set `processor_type: AutoProcessor` (or the " + "concrete processor class) in your config." + ) + from axolotl.prompt_strategies.multimodal_pretrain import ( + build_image_token_spec, + check_processor_compatibility, + ) + + check_processor_compatibility(processor) + spec = build_image_token_spec( + processor, + override=get_ds_value("image_token", None), + ) + image_column = get_ds_value("image_column", None) or "images" + LOG.info( + f"multimodal streaming CPT: placeholder={spec.image_token!r} " + f"(id={spec.image_token_id})" + ) + encode = functools.partial( + encode_streaming_multimodal, + tokenizer=tokenizer, + max_tokens=cfg.sequence_len, + image_token=spec.image_token, + image_token_id=spec.image_token_id, + text_column=text_column, + image_column=image_column, + ) + else: + encode = functools.partial( + encode_streaming, + tokenizer=tokenizer, + max_tokens=cfg.sequence_len, + text_column=text_column, + concatenate=cfg.pretraining_sample_concatenation is True, + ) + if cfg.shuffle_merged_datasets: dataset = dataset.shuffle( seed=cfg.seed, buffer_size=cfg.streaming_multipack_buffer_size diff --git a/src/axolotl/utils/schemas/config.py b/src/axolotl/utils/schemas/config.py index c52ddce1a0..f3c15edf7b 100644 --- a/src/axolotl/utils/schemas/config.py +++ b/src/axolotl/utils/schemas/config.py @@ -20,6 +20,7 @@ DatasetConfig, DPODataset, KTODataset, + MultiModalEvalDataset, PretrainingDataset, SFTDataset, StepwiseSupervisedDataset, @@ -341,7 +342,8 @@ class AxolotlInputConfig( test_datasets: ( Annotated[ list[ - SFTDataset + MultiModalEvalDataset + | SFTDataset | DPODataset | KTODataset | StepwiseSupervisedDataset diff --git a/src/axolotl/utils/schemas/datasets.py b/src/axolotl/utils/schemas/datasets.py index 6114a63e0a..6c55a6a9d9 100644 --- a/src/axolotl/utils/schemas/datasets.py +++ b/src/axolotl/utils/schemas/datasets.py @@ -166,10 +166,10 @@ class SFTDataset(BaseModel): "description": "Roles to train on. The tokens from these roles will be considered for the loss." }, ) - train_on_eos: Literal["all", "turn", "last"] | None = Field( + train_on_eos: Literal["all", "turn", "last", "none"] | None = Field( default=None, json_schema_extra={ - "description": "Which EOS tokens to train on in the conversation. Possible values are: all: train on all EOS tokens, turn (default): train on the EOS token at the end of each trainable turn, last: train on the last EOS token in the conversation" + "description": "Which EOS tokens to train on in the conversation. Possible values are: all: train on all EOS tokens, turn (default): train on the EOS token at the end of each trainable turn, last: train on the last EOS token in the conversation, none: never train on EOS tokens" }, ) roles: dict[str, list[str]] | None = Field( @@ -238,6 +238,88 @@ class PretrainingDataset(BaseModel): data_files: str | None = None skip: int | None = None + # Multimodal CPT fields. Opt-in via `type: multimodal_pretrain` or `multimodal: true`. + multimodal: bool | None = Field( + default=None, + json_schema_extra={ + "description": "Opt in to multimodal CPT. Auto-enabled when type='multimodal_pretrain'." + }, + ) + image_column: str | None = Field( + default="images", + json_schema_extra={ + "description": "Column holding a list of image paths per row." + }, + ) + image_base_dir: str | None = Field( + default=None, + json_schema_extra={"description": "Base directory for relative image paths."}, + ) + image_token: str | None = Field( + default=None, + json_schema_extra={ + "description": "Override the image placeholder token (autodetected from processor if unset)." + }, + ) + + +class MultiModalEvalDataset(BaseModel): + """Multimodal CPT eval dataset configuration (test_datasets entry). + + Use type='multimodal_pretrain' (or multimodal=True). The dataset must + expose a text column and a list[str] image-paths column; their names + default to 'text' and 'images' and can be overridden per-entry. + """ + + path: str | None = None + name: str | None = None + split: str | None = "train" + data_files: str | list[str] | None = None + skip: int | None = None + type: str | None = None + trust_remote_code: bool | None = False + + multimodal: bool | None = Field( + default=None, + json_schema_extra={ + "description": "Opt in to multimodal eval. Auto-enabled when type='multimodal_pretrain'." + }, + ) + text_column: str | None = Field( + default="text", + json_schema_extra={"description": "Column holding the row's text."}, + ) + image_column: str | None = Field( + default="images", + json_schema_extra={ + "description": "Column holding a list of image paths per row." + }, + ) + image_base_dir: str | None = Field( + default=None, + json_schema_extra={"description": "Base directory for relative image paths."}, + ) + image_token: str | None = Field( + default=None, + json_schema_extra={ + "description": "Override the image placeholder token (autodetected from processor if unset)." + }, + ) + + @model_validator(mode="before") + @classmethod + def _require_mm_markers(cls, data): + if isinstance(data, BaseModel): + data = data.model_dump() + if not isinstance(data, dict): + return data + if data.get("type") != "multimodal_pretrain" and not data.get("multimodal"): + raise ValueError( + "MultiModalEvalDataset requires type='multimodal_pretrain' " + "or multimodal=True" + ) + return data + class UserDefinedDPOType(BaseModel): """User defined typing for DPO""" diff --git a/src/axolotl/utils/schemas/multimodal.py b/src/axolotl/utils/schemas/multimodal.py index a3449199f3..01ad5e5a3d 100644 --- a/src/axolotl/utils/schemas/multimodal.py +++ b/src/axolotl/utils/schemas/multimodal.py @@ -6,6 +6,57 @@ from pydantic import BaseModel, Field, field_validator +class RoleBoundarySpec(BaseModel): + """One ``cfg.role_boundaries`` row; see docs/multimodal_assistant_mask.md.""" + + role: str = Field( + json_schema_extra={ + "description": ( + "Role name as it appears in cfg.roles_to_train (e.g. " + "'assistant', 'user', 'system', 'tool', 'ipython')." + ) + }, + ) + start: str = Field( + json_schema_extra={ + "description": ( + "Literal string that marks the start of this role's span in " + "the rendered chat template. Tokenized via " + "``tokenizer.encode(..., add_special_tokens=False)`` at " + "strategy init." + ) + }, + ) + end: str | None = Field( + default=None, + json_schema_extra={ + "description": ( + "Literal string that marks the end of this role's span. " + "Set to ``eos_token`` to terminate at the tokenizer's EOS. " + "Leave unset / null to terminate at end-of-sequence." + ) + }, + ) + include_start: bool = Field( + default=False, + json_schema_extra={ + "description": ( + "Whether the start marker tokens contribute to loss on " + "trainable turns. Default False." + ) + }, + ) + include_end: bool = Field( + default=True, + json_schema_extra={ + "description": ( + "Whether the end marker tokens contribute to loss on " + "trainable turns (honoring cfg.train_on_eos). Default True." + ) + }, + ) + + class MultiModalConfig(BaseModel): """Multi-modal configuration subset""" @@ -26,6 +77,17 @@ class MultiModalConfig(BaseModel): "description": "The resampling algorithm to use for image resizing. Default is bilinear. Please refer to PIL.Image.Resampling for more details." }, ) + role_boundaries: list[RoleBoundarySpec] | None = Field( + default=None, + json_schema_extra={ + "description": ( + "Opt-in override for the MM mask scanner's per-role boundary " + "markers. Non-empty list replaces built-ins wholesale; unset " + "or empty falls back to built-ins. See " + "docs/multimodal_assistant_mask.md." + ) + }, + ) @field_validator("image_resize_algorithm", mode="before") @classmethod diff --git a/src/axolotl/utils/schemas/validation.py b/src/axolotl/utils/schemas/validation.py index fff69de260..76d36979b3 100644 --- a/src/axolotl/utils/schemas/validation.py +++ b/src/axolotl/utils/schemas/validation.py @@ -1340,6 +1340,93 @@ def check_streaming_w_multiple_datasets(cls, data): ) return data + @model_validator(mode="before") + @classmethod + def check_multimodal_cpt(cls, data): + pd = data.get("pretraining_dataset") + if not pd: + return data + + pd_list = pd if isinstance(pd, list) else [pd] + + def _entry_is_mm(entry) -> bool: + if isinstance(entry, dict): + ds_type_ = entry.get("type") + mm_flag_ = entry.get("multimodal") + else: + ds_type_ = getattr(entry, "type", None) + mm_flag_ = getattr(entry, "multimodal", None) + return ds_type_ == "multimodal_pretrain" or bool(mm_flag_) + + # MM config resolves from entry[0] only; multi-entry runs miscollate or silently demote. + if len(pd_list) > 1 and any(_entry_is_mm(e) for e in pd_list): + raise ValueError( + "Multimodal CPT supports exactly one `pretraining_dataset` " + f"entry (found {len(pd_list)}). Image settings " + "(`image_base_dir`, `image_token`) and MM-mode detection " + "both resolve from entry[0] only, so additional entries " + "would be silently miscollated or drop their MM config. " + "Split multimodal CPT into its own run." + ) + + first = pd_list[0] + if not isinstance(first, dict): + return data + + ds_type = first.get("type") + is_mm_cpt = ds_type == "multimodal_pretrain" or bool(first.get("multimodal")) + if not is_mm_cpt: + return data + + if not data.get("processor_type"): + raise ValueError( + "Multimodal CPT (type: multimodal_pretrain) requires " + "`processor_type` to be set — e.g. `processor_type: AutoProcessor`. " + "Without a processor, images in the dataset cannot be turned " + "into pixel tensors." + ) + if data.get("sample_packing"): + raise ValueError( + "Multimodal CPT is incompatible with `sample_packing: true`. " + "Each image's placeholder token expands to a variable number " + "of patch tokens at the processor, so cross-row packing would " + "break the 1-to-1 alignment between text placeholders and " + "pixel_values. Set `sample_packing: false`." + ) + if data.get("chat_template"): + raise ValueError( + "Multimodal CPT (raw image+text pretraining) is incompatible " + "with `chat_template`. The point of the CPT path is to avoid " + "conversational scaffolding entirely. Remove `chat_template` " + "or switch to chat-template SFT." + ) + # Keep `images` and `_mm_text` columns alive for the collator. + prev_remove_unused = data.get("remove_unused_columns") + if prev_remove_unused is not False: + LOG.info( + "Auto-set `remove_unused_columns: false` for multimodal CPT " + "to preserve `images` and `_mm_text` columns (previous value: %r)", + prev_remove_unused, + ) + data["remove_unused_columns"] = False + + test_datasets = data.get("test_datasets") or [] + mm_test = [t for t in test_datasets if isinstance(t, dict) and _entry_is_mm(t)] + if len(mm_test) > 1: + for key in ("image_base_dir", "image_token"): + values = {t.get(key) for t in mm_test} + if len(values) > 1: + raise ValueError( + f"Multimodal CPT eval requires `{key}` to be either " + f"unset on all `test_datasets` entries or identical " + f"across them. The eval collator resolves " + f"`image_base_dir` and `image_token` once from the " + f"first entry, so heterogeneous values would silently " + f"miscollate later entries. Got: {sorted(map(str, values))}." + ) + + return data + class ModelCompatibilityValidationMixin: """Validation methods for specific model compatibility.""" diff --git a/tests/conftest.py b/tests/conftest.py index 19e3dc3f05..96be276a96 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -112,6 +112,24 @@ def download_smollm2_135m_instruct_model(): snapshot_download_w_retry("HuggingFaceTB/SmolLM2-135M-Instruct", repo_type="model") +@pytest.fixture(scope="session", autouse=True) +def download_smolvlm_500m_instruct_model(): + # Processor/tokenizer only — skip ~1 GB of weight shards. + snapshot_download_w_retry( + "HuggingFaceTB/SmolVLM-500M-Instruct", + repo_type="model", + allow_patterns=[ + "*.json", + "*.txt", + "*.model", + "*.jinja", + "tokenizer*", + "vocab*", + "merges*", + ], + ) + + @pytest.fixture(scope="session", autouse=True) def download_smollm2_135m_gptq_model(): # download the model diff --git a/tests/prompt_strategies/test_multimodal_pretrain.py b/tests/prompt_strategies/test_multimodal_pretrain.py new file mode 100644 index 0000000000..e20dd9a4ae --- /dev/null +++ b/tests/prompt_strategies/test_multimodal_pretrain.py @@ -0,0 +1,242 @@ +"""Multimodal CPT prompt strategy + safety gate tests.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +from PIL import Image +from transformers import AutoProcessor + +from axolotl.prompt_strategies.multimodal_pretrain import ( + _INCOMPATIBLE_PROCESSOR_REASONS, + ImageTokenSpec, + MultimodalPretrainTokenizationStrategy, + build_image_token_spec, + check_processor_compatibility, + load, +) +from axolotl.prompt_strategies.pretrain import PretrainTokenizer + +from tests.hf_offline_utils import enable_hf_offline + +_SMOLVLM = "HuggingFaceTB/SmolVLM-500M-Instruct" + + +@pytest.fixture(scope="module", name="smolvlm_processor") +@enable_hf_offline +def fixture_smolvlm_processor( + download_smolvlm_500m_instruct_model, # pylint: disable=unused-argument +): + return AutoProcessor.from_pretrained(_SMOLVLM) + + +@pytest.fixture(scope="module", name="tiny_image_path") +def fixture_tiny_image_path(tmp_path_factory) -> Path: + d = tmp_path_factory.mktemp("mm_pretrain_imgs") + p = d / "dummy.png" + arr = np.random.default_rng(0).integers(0, 255, (64, 64, 3)).astype("uint8") + Image.fromarray(arr).save(p) + return p + + +# ---- build_image_token_spec ------------------------------------------------ + + +def test_build_image_token_spec_autodetects_smolvlm(smolvlm_processor): + spec = build_image_token_spec(smolvlm_processor) + assert isinstance(spec, ImageTokenSpec) + assert spec.image_token == "" + assert spec.image_token_id > 0 + assert spec.image_token_id in spec.image_family_token_ids + + +def test_build_image_token_spec_honors_override(smolvlm_processor): + spec = build_image_token_spec(smolvlm_processor, override="") + assert spec.image_token == "" + + +def test_build_image_token_spec_rejects_bad_override(smolvlm_processor): + with pytest.raises(ValueError, match="not a registered special token"): + build_image_token_spec(smolvlm_processor, override="") + + +def test_build_image_token_spec_rejects_plain_word_override(smolvlm_processor): + # Plain words BPE-tokenize but aren't placeholders. + with pytest.raises(ValueError, match="not a registered special token"): + build_image_token_spec(smolvlm_processor, override="image") + + +# ---- check_processor_compatibility (startup-time gate) --------------------- + + +@pytest.mark.parametrize("cls_name", list(_INCOMPATIBLE_PROCESSOR_REASONS.keys())) +def test_check_processor_compatibility_rejects_incompatible(cls_name): + fake = type(cls_name, (), {})() + with pytest.raises(ValueError) as exc: + check_processor_compatibility(fake) + assert cls_name in str(exc.value) + assert _INCOMPATIBLE_PROCESSOR_REASONS[cls_name] in str(exc.value) + + +def test_check_processor_compatibility_rejects_subclass(): + # MRO-name fallback must catch user-defined subclasses. + class BaseMllama: + pass + + BaseMllama.__name__ = "MllamaProcessor" + + class CustomUserProcessor(BaseMllama): + pass + + CustomUserProcessor.__name__ = "CustomUserProcessor" + + with pytest.raises(ValueError, match="MllamaProcessor"): + check_processor_compatibility(CustomUserProcessor()) + + +def test_check_processor_compatibility_accepts_supported(smolvlm_processor): + check_processor_compatibility(smolvlm_processor) + + +# ---- MultimodalPretrainTokenizationStrategy -------------------------------- + + +def _make_strategy( + smolvlm_processor: Any, + text_column: str = "text", + image_column: str = "images", +) -> MultimodalPretrainTokenizationStrategy: + spec = build_image_token_spec(smolvlm_processor) + return MultimodalPretrainTokenizationStrategy( + PretrainTokenizer(), + smolvlm_processor.tokenizer, + False, # train_on_inputs + 2048, # sequence_len + text_column=text_column, + image_column=image_column, + image_base_dir=None, + image_token=spec.image_token, + image_token_id=spec.image_token_id, + max_length=2048, + ) + + +def test_strategy_preserves_images_and_text(smolvlm_processor, tiny_image_path): + strat = _make_strategy(smolvlm_processor) + out = strat.tokenize_prompt( + { + "text": "\nsample transcription text", + "images": [str(tiny_image_path)], + } + ) + assert "input_ids" in out + assert "images" in out and "_mm_text" in out + assert len(out["input_ids"]) == 1 + assert len(out["images"]) == 1 + assert len(out["_mm_text"]) == 1 + assert out["images"][0] == [str(tiny_image_path)] + assert out["_mm_text"][0].startswith("") + + +def test_strategy_rejects_placeholder_count_mismatch( + smolvlm_processor, tiny_image_path +): + strat = _make_strategy(smolvlm_processor) + with pytest.raises(ValueError, match="occurrence"): + strat.tokenize_prompt( + { + "text": "\ntwo placeholders one image", + "images": [str(tiny_image_path)], + } + ) + + +def test_strategy_rejects_row_exceeding_max_length(smolvlm_processor, tiny_image_path): + spec = build_image_token_spec(smolvlm_processor) + strat = MultimodalPretrainTokenizationStrategy( + PretrainTokenizer(), + smolvlm_processor.tokenizer, + False, + 128, + text_column="text", + image_column="images", + image_base_dir=None, + image_token=spec.image_token, + image_token_id=spec.image_token_id, + max_length=128, + ) + huge = "word " * 5000 + with pytest.raises(ValueError, match="exceeds sequence_len"): + strat.tokenize_prompt( + { + "text": f"{spec.image_token} {huge}", + "images": [str(tiny_image_path)], + } + ) + + +def test_strategy_rejects_non_list_image_column(smolvlm_processor, tiny_image_path): + strat = _make_strategy(smolvlm_processor) + with pytest.raises(ValueError, match="list"): + strat.tokenize_prompt( + { + "text": "\nbad image field", + "images": str(tiny_image_path), # should be a list + } + ) + + +@pytest.mark.parametrize("bad_value", ["", 0, False]) +def test_strategy_rejects_falsy_non_none_image_column(smolvlm_processor, bad_value): + """Falsy non-None image cells (e.g. "") are rejected, not coerced to [].""" + strat = _make_strategy(smolvlm_processor) + with pytest.raises(ValueError, match="list"): + strat.tokenize_prompt( + { + "text": "no placeholder, but bad images cell", + "images": bad_value, + } + ) + + +def test_strategy_treats_none_image_column_as_empty(smolvlm_processor): + """images=None is the only falsy value treated as a text-only row.""" + strat = _make_strategy(smolvlm_processor) + out = strat.tokenize_prompt( + { + "text": "plain text-only row, no placeholder", + "images": None, + } + ) + assert out["images"][0] == [] + + +# ---- load() factory -------------------------------------------------------- + + +def test_load_requires_processor(smolvlm_processor): + class _Cfg: + train_on_inputs = False + sequence_len = 2048 + + with pytest.raises(ValueError, match="processor"): + load(smolvlm_processor.tokenizer, _Cfg(), ds_cfg={}, processor=None) + + +def test_load_returns_strategy_with_spec(smolvlm_processor): + class _Cfg: + train_on_inputs = False + sequence_len = 2048 + + strat = load( + smolvlm_processor.tokenizer, + _Cfg(), + ds_cfg={"text_column": "text", "image_column": "images"}, + processor=smolvlm_processor, + ) + assert isinstance(strat, MultimodalPretrainTokenizationStrategy) + assert hasattr(strat, "image_token_spec") + assert strat.image_token_spec.image_token == "" diff --git a/tests/test_multimodal_streaming.py b/tests/test_multimodal_streaming.py new file mode 100644 index 0000000000..9fc287acdf --- /dev/null +++ b/tests/test_multimodal_streaming.py @@ -0,0 +1,448 @@ +"""Multimodal CPT streaming encoder + collator tests.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +import torch +from PIL import Image +from transformers import AutoProcessor + +from axolotl.prompt_strategies.multimodal_pretrain import build_image_token_spec +from axolotl.utils.collators.mm_pretrain import MultiModalPretrainDataCollator +from axolotl.utils.data.streaming import ( + encode_streaming_multimodal, + wrap_streaming_dataset, +) +from axolotl.utils.dict import DictDefault + +from tests.hf_offline_utils import enable_hf_offline + +_SMOLVLM = "HuggingFaceTB/SmolVLM-500M-Instruct" + + +@pytest.fixture(scope="module", name="smolvlm_processor") +@enable_hf_offline +def fixture_smolvlm_processor( + download_smolvlm_500m_instruct_model, # pylint: disable=unused-argument +): + return AutoProcessor.from_pretrained(_SMOLVLM) + + +@pytest.fixture(scope="module", name="two_tiny_images") +def fixture_two_tiny_images(tmp_path_factory) -> list[Path]: + d = tmp_path_factory.mktemp("mm_stream_imgs") + out = [] + for i in range(2): + p = d / f"dummy_{i}.png" + arr = np.random.default_rng(i).integers(0, 255, (64, 64, 3)).astype("uint8") + Image.fromarray(arr).save(p) + out.append(p) + return out + + +# ---- encode_streaming_multimodal ------------------------------------------ + + +def test_encode_preserves_images_and_text(smolvlm_processor, two_tiny_images): + spec = build_image_token_spec(smolvlm_processor) + examples = { + "text": [ + f"{spec.image_token}\nrow one", + f"{spec.image_token}\nrow two slightly longer", + ], + "images": [[str(two_tiny_images[0])], [str(two_tiny_images[1])]], + } + out = encode_streaming_multimodal( + examples, + tokenizer=smolvlm_processor.tokenizer, + max_tokens=2048, + image_token=spec.image_token, + image_token_id=spec.image_token_id, + ) + assert set(out) >= {"input_ids", "labels", "attention_mask", "images", "_mm_text"} + assert len(out["input_ids"]) == 2 + assert out["images"] == [[str(two_tiny_images[0])], [str(two_tiny_images[1])]] + # EOS appended -> input_ids len equals attention_mask len and > text + for ids, mask in zip(out["input_ids"], out["attention_mask"], strict=True): + assert len(ids) == len(mask) and len(ids) > 0 + # CPT: labels == input_ids pre-masking. + for ids, lbls in zip(out["input_ids"], out["labels"], strict=True): + assert ids == lbls + + +def test_encode_rejects_mismatch(smolvlm_processor, two_tiny_images): + spec = build_image_token_spec(smolvlm_processor) + examples = { + "text": [f"{spec.image_token}{spec.image_token}\ntwo placeholders one image"], + "images": [[str(two_tiny_images[0])]], + } + with pytest.raises(ValueError, match="occurrence"): + encode_streaming_multimodal( + examples, + tokenizer=smolvlm_processor.tokenizer, + max_tokens=2048, + image_token=spec.image_token, + image_token_id=spec.image_token_id, + ) + + +def test_encode_rejects_row_without_list(smolvlm_processor, two_tiny_images): + spec = build_image_token_spec(smolvlm_processor) + with pytest.raises(ValueError, match="list"): + encode_streaming_multimodal( + { + "text": [f"{spec.image_token}\nrow one"], + "images": [str(two_tiny_images[0])], # scalar, not a list + }, + tokenizer=smolvlm_processor.tokenizer, + max_tokens=2048, + image_token=spec.image_token, + image_token_id=spec.image_token_id, + ) + + +def test_encode_counts_placeholders_on_full_text(smolvlm_processor, two_tiny_images): + # All 3 placeholders must be counted even when text would have been truncated. + spec = build_image_token_spec(smolvlm_processor) + long_filler = "lorem ipsum " * 20 + text = f"{spec.image_token} {long_filler} {spec.image_token} {long_filler} {spec.image_token}" + examples = { + "text": [text], + "images": [[str(two_tiny_images[0])] * 3], + } + out = encode_streaming_multimodal( + examples, + tokenizer=smolvlm_processor.tokenizer, + max_tokens=4096, + image_token=spec.image_token, + image_token_id=spec.image_token_id, + ) + assert sum(1 for t in out["input_ids"][0] if t == spec.image_token_id) == 3 + + +def test_encode_rejects_row_exceeding_max_tokens(smolvlm_processor, two_tiny_images): + spec = build_image_token_spec(smolvlm_processor) + huge = "word " * 5000 + examples = { + "text": [f"{spec.image_token} {huge}"], + "images": [[str(two_tiny_images[0])]], + } + with pytest.raises(ValueError, match="exceeds sequence_len"): + encode_streaming_multimodal( + examples, + tokenizer=smolvlm_processor.tokenizer, + max_tokens=512, + image_token=spec.image_token, + image_token_id=spec.image_token_id, + ) + + +# ---- wrap_streaming_dataset routing -------------------------------------- + + +def test_wrap_streaming_dataset_uses_pretraining_config_arg( + smolvlm_processor, monkeypatch +): + # Eval path passes a per-entry config that may differ from cfg.pretraining_dataset[0]. + # The MM-CPT branch must read from that arg, not re-resolve from cfg. + captured = {} + + def fake_partial(fn, **kwargs): + captured["encode_fn"] = fn + captured["kwargs"] = kwargs + return lambda batch: batch + + monkeypatch.setattr("axolotl.utils.data.streaming.functools.partial", fake_partial) + + class _Dataset: + features = {"text": None, "images": None} + + def shuffle(self, **_): + return self + + def map(self, *_args, **_kwargs): + return self + + cfg = DictDefault( + { + "sample_packing": False, + "pretraining_dataset": [ + { + "path": "train/ds", + "type": "multimodal_pretrain", + "text_column": "wrong_train_col", + "image_column": "wrong_train_imgs", + } + ], + "sequence_len": 256, + "shuffle_merged_datasets": False, + "streaming_multipack_buffer_size": 1000, + "seed": 42, + } + ) + eval_entry = DictDefault( + { + "path": "test/ds", + "type": "multimodal_pretrain", + "text_column": "eval_text", + "image_column": "eval_imgs", + } + ) + + wrap_streaming_dataset( + _Dataset(), + smolvlm_processor.tokenizer, + cfg, + ds_wrapper_fn=None, + processor=smolvlm_processor, + pretraining_config=eval_entry, + ) + + assert captured["encode_fn"] is encode_streaming_multimodal + assert captured["kwargs"]["text_column"] == "eval_text" + assert captured["kwargs"]["image_column"] == "eval_imgs" + + +# ---- MultiModalPretrainDataCollator --------------------------------------- + + +def test_collator_builds_batch_and_masks_labels(smolvlm_processor, two_tiny_images): + spec = build_image_token_spec(smolvlm_processor) + encoded = encode_streaming_multimodal( + { + "text": [ + f"{spec.image_token}\nrow one", + f"{spec.image_token}\nrow two slightly longer", + ], + "images": [[str(two_tiny_images[0])], [str(two_tiny_images[1])]], + }, + tokenizer=smolvlm_processor.tokenizer, + max_tokens=2048, + image_token=spec.image_token, + image_token_id=spec.image_token_id, + ) + rows = [ + { + k: encoded[k][i] + for k in ("input_ids", "labels", "attention_mask", "images", "_mm_text") + } + for i in range(2) + ] + collator = MultiModalPretrainDataCollator( + tokenizer=smolvlm_processor.tokenizer, + processor=smolvlm_processor, + image_token_spec=spec, + ) + batch = collator.torch_call(rows) + # Expected keys + for k in ("input_ids", "attention_mask", "pixel_values", "labels"): + assert k in batch, f"missing batch key {k}" + assert isinstance(batch["input_ids"], torch.Tensor) + # Label masking check: no image-family ids remaining as valid labels. + for tid in spec.image_family_token_ids: + assert int((batch["labels"] == tid).sum().item()) == 0, ( + f"label masking left id={tid} in labels" + ) + # Pad is also masked. + pad_id = smolvlm_processor.tokenizer.pad_token_id + if pad_id is not None: + assert int((batch["labels"] == pad_id).sum().item()) == 0 + + +def test_collator_raises_on_missing_columns(smolvlm_processor): + spec = build_image_token_spec(smolvlm_processor) + collator = MultiModalPretrainDataCollator( + tokenizer=smolvlm_processor.tokenizer, + processor=smolvlm_processor, + image_token_spec=spec, + ) + with pytest.raises(KeyError, match="encode_streaming_multimodal"): + collator.torch_call([{"input_ids": [1, 2, 3]}]) # no _mm_text / images + + +# ---- security gates ------------------------------------------------------- + + +def test_collator_rejects_path_traversal_with_base_dir( + smolvlm_processor, two_tiny_images, tmp_path +): + spec = build_image_token_spec(smolvlm_processor) + base = tmp_path / "images" + base.mkdir() + collator = MultiModalPretrainDataCollator( + tokenizer=smolvlm_processor.tokenizer, + processor=smolvlm_processor, + image_token_spec=spec, + image_base_dir=str(base), + ) + # Absolute path rejection + with pytest.raises(RuntimeError) as exc: + collator._load_images_for_row([str(two_tiny_images[0])], row_index=0) + assert isinstance(exc.value.__cause__, ValueError) + assert "Absolute image path" in str(exc.value.__cause__) + # Containment-escape rejection + with pytest.raises(RuntimeError) as exc: + collator._load_images_for_row(["../../../etc/passwd"], row_index=0) + assert isinstance(exc.value.__cause__, ValueError) + assert "outside" in str(exc.value.__cause__) + + +def test_collator_rejects_remote_urls(smolvlm_processor): + spec = build_image_token_spec(smolvlm_processor) + collator = MultiModalPretrainDataCollator( + tokenizer=smolvlm_processor.tokenizer, + processor=smolvlm_processor, + image_token_spec=spec, + ) + for url in ( + "http://example.com/a.png", + "https://x/y.jpg", + "file:///etc/passwd", + "ftp://x/y.png", + "data:image/png;base64,xxx", + # Case-variant bypass attempts. + "HTTP://evil.com/x.png", + "Https://x/y.jpg", + "FILE:///etc/passwd", + "DATA:image/png;base64,xxx", + ): + with pytest.raises(RuntimeError) as exc: + collator._load_images_for_row([url], row_index=0) + assert isinstance(exc.value.__cause__, ValueError) + assert "Non-local image path scheme" in str(exc.value.__cause__) + + +def test_collator_rejects_nul_byte_paths(smolvlm_processor): + spec = build_image_token_spec(smolvlm_processor) + collator = MultiModalPretrainDataCollator( + tokenizer=smolvlm_processor.tokenizer, + processor=smolvlm_processor, + image_token_spec=spec, + ) + with pytest.raises(RuntimeError) as exc: + collator._load_images_for_row(["bad\x00path.png"], row_index=0) + assert "NUL byte" in str(exc.value.__cause__) + + +def test_collator_rejects_non_string_image_entries(smolvlm_processor, two_tiny_images): + spec = build_image_token_spec(smolvlm_processor) + collator = MultiModalPretrainDataCollator( + tokenizer=smolvlm_processor.tokenizer, + processor=smolvlm_processor, + image_token_spec=spec, + ) + rows = [ + { + "_mm_text": f"{spec.image_token}\nrow", + "images": [None], # type: ignore[list-item] + } + ] + with pytest.raises(TypeError, match="path must be str"): + collator.torch_call(rows) + + +def test_collator_rejects_bytes_mm_text(smolvlm_processor, two_tiny_images): + spec = build_image_token_spec(smolvlm_processor) + collator = MultiModalPretrainDataCollator( + tokenizer=smolvlm_processor.tokenizer, + processor=smolvlm_processor, + image_token_spec=spec, + ) + rows = [ + { + "_mm_text": f"{spec.image_token}\nrow".encode(), + "images": [str(two_tiny_images[0])], + } + ] + with pytest.raises(TypeError, match="`_mm_text` must be str"): + collator.torch_call(rows) + + +def test_collator_sanitizes_error_message(smolvlm_processor, tmp_path): + spec = build_image_token_spec(smolvlm_processor) + collator = MultiModalPretrainDataCollator( + tokenizer=smolvlm_processor.tokenizer, + processor=smolvlm_processor, + image_token_spec=spec, + ) + missing = tmp_path / "subdir_with_secret_name" / "nope.png" + with pytest.raises(RuntimeError) as exc: + collator._load_images_for_row([str(missing)], row_index=3) + # basename appears, full directory path does NOT + assert "nope.png" in str(exc.value) + assert "subdir_with_secret_name" not in str(exc.value) + assert "Row 3" in str(exc.value) + + +def test_collator_rejects_too_many_images(smolvlm_processor, two_tiny_images): + spec = build_image_token_spec(smolvlm_processor) + collator = MultiModalPretrainDataCollator( + tokenizer=smolvlm_processor.tokenizer, + processor=smolvlm_processor, + image_token_spec=spec, + max_images_per_row=2, + ) + paths = [str(two_tiny_images[0])] * 3 + with pytest.raises(ValueError, match="max_images_per_row"): + collator._load_images_for_row(paths, row_index=0) + + +# ---- mixed / all-text batches -------------------------------------------- + + +def test_collator_all_text_batch_uses_tokenizer_fallback(smolvlm_processor): + """A batch where every row has images=[] tokenizes via the tokenizer; no pixel_values.""" + spec = build_image_token_spec(smolvlm_processor) + collator = MultiModalPretrainDataCollator( + tokenizer=smolvlm_processor.tokenizer, + processor=smolvlm_processor, + image_token_spec=spec, + ) + rows = [ + {"_mm_text": "first text-only row", "images": []}, + {"_mm_text": "second text-only row, slightly longer", "images": []}, + ] + batch = collator.torch_call(rows) + for k in ("input_ids", "attention_mask", "labels"): + assert k in batch, f"missing batch key {k}" + assert "pixel_values" not in batch + assert isinstance(batch["input_ids"], torch.Tensor) + pad_id = smolvlm_processor.tokenizer.pad_token_id + if pad_id is not None: + assert int((batch["labels"] == pad_id).sum().item()) == 0 + + +def test_collator_mixed_batch_still_succeeds(smolvlm_processor, two_tiny_images): + """A batch with one imaged row and one text-only row still produces pixel_values.""" + spec = build_image_token_spec(smolvlm_processor) + encoded = encode_streaming_multimodal( + { + "text": [ + f"{spec.image_token}\nimaged row", + "text-only row", + ], + "images": [[str(two_tiny_images[0])], []], + }, + tokenizer=smolvlm_processor.tokenizer, + max_tokens=2048, + image_token=spec.image_token, + image_token_id=spec.image_token_id, + ) + rows = [ + { + k: encoded[k][i] + for k in ("input_ids", "labels", "attention_mask", "images", "_mm_text") + } + for i in range(2) + ] + collator = MultiModalPretrainDataCollator( + tokenizer=smolvlm_processor.tokenizer, + processor=smolvlm_processor, + image_token_spec=spec, + ) + batch = collator.torch_call(rows) + for k in ("input_ids", "attention_mask", "pixel_values", "labels"): + assert k in batch, f"missing batch key {k}" diff --git a/tests/test_processing_strategies.py b/tests/test_processing_strategies.py new file mode 100644 index 0000000000..2d8f13fe57 --- /dev/null +++ b/tests/test_processing_strategies.py @@ -0,0 +1,1164 @@ +"""Tests for ``axolotl.processing_strategies`` using fake tokenizers (offline/CI-safe).""" + +import logging + +import pytest +import torch +from pydantic import ValidationError + +from axolotl.processing_strategies import ( + Gemma3nProcessingStrategy, + Gemma3ProcessingStrategy, + Gemma4ProcessingStrategy, + Llama3_2VisionProcessingStrategy, + Llama4ProcessingStrategy, + MistralV7TekkenProcessingStrategy, + PixtralProcessingStrategy, + ProcessingStrategy, + Qwen2VLProcessingStrategy, + Qwen3_5ProcessingStrategy, + RoleBoundary, + _apply_role_boundaries, + get_processing_strategy, +) + + +@pytest.fixture +def axolotl_caplog(caplog): + """caplog that also captures records from the ``axolotl`` logger. + + The axolotl logger sets ``propagate=False`` once ``configure_logging()`` is + called (which happens indirectly in many CI test paths), so the default + caplog handler installed on the root logger never sees these records. + Attaching ``caplog.handler`` to ``axolotl.processing_strategies`` directly + makes assertions reliable regardless of whether ``configure_logging`` has + already run on this worker. + """ + logger = logging.getLogger("axolotl.processing_strategies") + logger.addHandler(caplog.handler) + previous_level = logger.level + logger.setLevel(logging.DEBUG) + try: + yield caplog + finally: + logger.removeHandler(caplog.handler) + logger.setLevel(previous_level) + + +# --------------------------------------------------------------------------- # +# Generic fake tokenizer/processor scaffold +# --------------------------------------------------------------------------- # + + +class _Tokenizer: + """Minimal tokenizer stub; ``vocab`` maps marker strings to their id lists.""" + + def __init__( + self, + vocab: dict[str, list[int]], + pad_id: int = 0, + unk_id: int = 3, + eos_id: int | None = None, + ): + self.vocab = vocab + self._reverse = {} + for tok, ids in vocab.items(): + if len(ids) == 1: + self._reverse[ids[0]] = tok + self.pad_token_id = pad_id + self.unk_token_id = unk_id + if eos_id is not None: + self.eos_token_id = eos_id + + def encode(self, text, add_special_tokens=False): + # Unknown markers return [] so _encode_markers drops them silently. + return list(self.vocab.get(text, [])) + + def convert_tokens_to_ids(self, token): + v = self.vocab.get(token) + if v is None: + return self.unk_token_id + return v[0] if len(v) == 1 else self.unk_token_id + + +class _Processor: + def __init__(self, tokenizer: _Tokenizer): + self.tokenizer = tokenizer + + +# --------------------------------------------------------------------------- # +# Base scanner tests (train_on_inputs / roles_to_train / train_on_eos) +# --------------------------------------------------------------------------- # + + +def _scan(role_boundaries, seq, roles_to_train=("assistant",), train_on_eos="turn"): + labels = torch.tensor([seq]) + return _apply_role_boundaries( + labels, role_boundaries, set(roles_to_train), train_on_eos + ).tolist()[0] + + +def test_scanner_assistant_only_basic(): + boundaries = [ + RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]), + RoleBoundary(role="user", start_tokens=[1, 3], end_tokens=[9]), + ] + seq = [1, 3, 7, 9, 1, 2, 8, 8, 9, 5] + out = _scan(boundaries, seq) + assert out == [-100, -100, -100, -100, -100, -100, 8, 8, 9, -100] + + +def test_scanner_train_on_eos_none_excludes_end_marker(): + boundaries = [ + RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]), + ] + seq = [1, 2, 8, 8, 9] + out = _scan(boundaries, seq, train_on_eos="none") + assert out == [-100, -100, 8, 8, -100] + + +def test_scanner_train_on_eos_all_keeps_non_assistant_end_marker(): + boundaries = [ + RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]), + RoleBoundary(role="user", start_tokens=[1, 3], end_tokens=[9]), + ] + seq = [1, 3, 7, 9, 1, 2, 8, 9] + out = _scan(boundaries, seq, train_on_eos="all") + assert out == [-100, -100, -100, 9, -100, -100, 8, 9] + + +def test_scanner_train_on_eos_all_with_non_trainable_include_end_false(): + """Non-trainable + include_end=False must not leak end marker on 'all'.""" + boundaries = [ + RoleBoundary( + role="user", + start_tokens=[50], + end_tokens=[51], + include_end=False, # shared with assistant-start + ), + RoleBoundary( + role="assistant", + start_tokens=[51], + end_tokens=[99], # eos + ), + ] + seq = [50, 7, 51, 8, 8, 99] # [INST] 7 [/INST] 8 8 EOS + out = _scan(boundaries, seq, roles_to_train=("assistant",), train_on_eos="all") + # [/INST] at idx 2 must stay masked — user.include_end=False says so. + assert out == [-100, -100, -100, 8, 8, 99] + + +def test_scanner_roles_to_train_user_and_assistant(): + boundaries = [ + RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]), + RoleBoundary(role="user", start_tokens=[1, 3], end_tokens=[9]), + ] + seq = [1, 3, 7, 9, 1, 2, 8, 9] + out = _scan(boundaries, seq, roles_to_train=("user", "assistant")) + # include_start defaults to False so role-start markers stay masked. + assert out == [-100, -100, 7, 9, -100, -100, 8, 9] + + +def test_scanner_truncated_assistant(): + """Missing end marker: span runs to end-of-sequence, end marker not emitted.""" + boundaries = [ + RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]), + ] + seq = [1, 2, 8, 8, 8] + out = _scan(boundaries, seq) + assert out == [-100, -100, 8, 8, 8] + + +def test_scanner_longest_prefix_wins(): + boundaries = [ + RoleBoundary(role="assistant", start_tokens=[1, 2, 4], end_tokens=[9]), + RoleBoundary(role="user", start_tokens=[1, 2], end_tokens=[9]), + ] + seq = [1, 2, 4, 8, 9] + out = _scan(boundaries, seq) + assert out == [-100, -100, -100, 8, 9] + + +def test_scanner_no_boundaries_masks_everything(): + # Strategies short-circuit this in _mask_non_assistant; see test_base_strategy_warns_when_no_boundaries. + labels = torch.tensor([[1, 2, 3, 4]]) + out = _apply_role_boundaries(labels, [], {"assistant"}, "turn") + assert out.tolist() == [[-100, -100, -100, -100]] + + +def test_scanner_train_on_eos_last_only_final_trainable_turn(): + boundaries = [ + RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]), + ] + seq = [1, 2, 5, 9, 1, 2, 6, 9] + out = _scan(boundaries, seq, train_on_eos="last") + # Only the second assistant turn's end marker (index 7) is kept. + assert out == [-100, -100, 5, -100, -100, -100, 6, 9] + + +def test_scanner_train_on_eos_last_no_trainable_turn_is_noop(): + boundaries = [ + RoleBoundary(role="user", start_tokens=[1, 3], end_tokens=[9]), + ] + seq = [1, 3, 5, 9, 1, 3, 6, 9] + out = _scan(boundaries, seq, roles_to_train=("assistant",), train_on_eos="last") + assert out == [-100] * 8 + + +def test_strategy_rejects_unknown_train_on_eos(): + vocab = {"BOA": [50], "EOT": [60]} + with pytest.raises(ValueError, match="train_on_eos"): + ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + train_on_eos="bogus", + role_boundaries_override=[ + {"role": "assistant", "start": "BOA", "end": "EOT"} + ], + ) + + +def test_strategy_accepts_all_supported_train_on_eos_values(): + vocab = {"BOA": [50], "EOT": [60]} + for val in ("turn", "all", "none", "last"): + ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + train_on_eos=val, + role_boundaries_override=[ + {"role": "assistant", "start": "BOA", "end": "EOT"} + ], + ) + + +def test_empty_role_boundaries_override_falls_back_to_builtin(): + """Empty override must fall through to built-ins (opt-in semantics).""" + vocab = { + "<|im_start|>assistant\n": [101, 102, 103], + "<|im_start|>user\n": [101, 106, 103], + "<|im_end|>": [104], + } + strat_empty = Qwen2VLProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + role_boundaries_override=[], + ) + strat_default = Qwen2VLProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + ) + # Empty override === no override: both strategies keep the built-in boundaries. + assert strat_empty.role_boundaries == strat_default.role_boundaries + assert len(strat_empty.role_boundaries) > 0 # sanity: built-ins are non-empty + + +def test_sft_dataset_schema_accepts_all_supported_train_on_eos_values(): + """SFTDataset.train_on_eos accepts every value the scanner honors.""" + from axolotl.utils.schemas.datasets import SFTDataset + + for val in ("all", "turn", "last", "none"): + ds = SFTDataset(path="dummy", type="chat_template", train_on_eos=val) + assert ds.train_on_eos == val + + with pytest.raises(ValidationError): + SFTDataset(path="dummy", type="chat_template", train_on_eos="bogus") + + +def test_strategy_init_logs_resolved_masking_config_builtin(axolotl_caplog): + vocab = { + "<|im_start|>assistant\n": [101, 102, 103], + "<|im_start|>user\n": [101, 106, 103], + "<|im_end|>": [104], + } + with axolotl_caplog.at_level(logging.INFO, logger="axolotl.processing_strategies"): + Qwen2VLProcessingStrategy(_Processor(_Tokenizer(vocab, pad_id=0))) + msgs = [r.getMessage() for r in axolotl_caplog.records] + assert any( + "ProcessingStrategy init" in m + and "Qwen2VLProcessingStrategy" in m + and "boundaries_source=built-in" in m + for m in msgs + ) + + +def test_strategy_init_logs_resolved_masking_config_override(axolotl_caplog): + vocab = {"BOA": [50, 51], "EOT": [60]} + with axolotl_caplog.at_level(logging.INFO, logger="axolotl.processing_strategies"): + ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + role_boundaries_override=[ + {"role": "assistant", "start": "BOA", "end": "EOT"}, + ], + ) + msgs = [r.getMessage() for r in axolotl_caplog.records] + # Resolved start/end ids must appear in the log so users can verify what + # was actually matched. + assert any( + "ProcessingStrategy init" in m + and "boundaries_source=override" in m + and "[50, 51]" in m + and "[60]" in m + for m in msgs + ) + + +def test_process_labels_no_warning_when_image_token_id_none(): + """image_token_id=None must not trigger a UserWarning from ``labels == None``.""" + import warnings + + vocab = {"BOA": [50], "EOT": [60]} + strategy = ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + role_boundaries_override=[{"role": "assistant", "start": "BOA", "end": "EOT"}], + ) + assert strategy.image_token_id is None + with warnings.catch_warnings(): + warnings.simplefilter("error") + strategy.process_labels(torch.tensor([[1, 50, 2, 3, 60]])) + + +def test_roles_to_train_empty_list_masks_everything(): + """An explicit empty list is distinct from None and disables all roles.""" + vocab = {"BOA": [50], "EOT": [60]} + strategy = ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + roles_to_train=[], + role_boundaries_override=[{"role": "assistant", "start": "BOA", "end": "EOT"}], + ) + assert strategy.roles_to_train == [] + seq = [1, 50, 7, 8, 60, 9] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + assert out == [-100] * 6 + + +# --------------------------------------------------------------------------- # +# Qwen2VL / Qwen3.5 +# --------------------------------------------------------------------------- # + + +def _qwen_tokenizer(): + # ChatML-ish with image_pad=200, video_pad=201. + vocab = { + "<|im_start|>assistant\n": [101, 102, 103], + "<|im_start|>user\n": [101, 106, 103], + "<|im_start|>system\n": [101, 105, 103], + "<|im_end|>": [104], + "<|image_pad|>": [200], + "<|video_pad|>": [201], + } + return _Tokenizer(vocab, pad_id=0) + + +def _make_qwen2vl(): + tok = _qwen_tokenizer() + return Qwen2VLProcessingStrategy(_Processor(tok)) + + +def test_qwen2vl_masks_user_keeps_assistant_and_image_pad(): + strategy = _make_qwen2vl() + seq = [ + 101, + 105, + 103, + 77, + 104, + 101, + 106, + 103, + 7, + 104, + 101, + 102, + 103, + 200, + 8, + 104, + ] + labels = strategy.process_labels(torch.tensor([seq])) + out = labels.tolist()[0] + assert out[:10] == [-100] * 10 + assert out[10] == -100 and out[11] == -100 and out[12] == -100 + assert out[13] == -100 # image_pad masked post-scan + assert out[14] == 8 + assert out[15] == 104 + + +def test_qwen3_5_masks_video_pad_too(): + tok = _qwen_tokenizer() + strategy = Qwen3_5ProcessingStrategy(_Processor(tok)) + seq = [101, 102, 103, 201, 8, 104] + labels = strategy.process_labels(torch.tensor([seq])) + assert labels.tolist()[0] == [-100, -100, -100, -100, 8, 104] + + +def test_qwen2vl_train_on_inputs_true_keeps_everything(): + tok = _qwen_tokenizer() + strategy = Qwen2VLProcessingStrategy(_Processor(tok), train_on_inputs=True) + seq = [101, 106, 103, 7, 104, 101, 102, 103, 8, 104] + labels = strategy.process_labels(torch.tensor([seq])) + assert labels.tolist()[0] == seq + + +# --------------------------------------------------------------------------- # +# Gemma3 / Gemma3n +# --------------------------------------------------------------------------- # + + +def _gemma_tokenizer(): + vocab = { + "model\n": [1, 2, 3], + "user\n": [1, 10, 3], + "system\n": [1, 11, 3], + "": [4], + "": [50], # boi_token for Gemma3 + } + tok = _Tokenizer(vocab, pad_id=0) + # boi_token is a direct tokenizer attribute on real Gemma3. + tok.boi_token = "" + return tok + + +def test_gemma3_scanner_plus_soft_image_token(): + strategy = Gemma3ProcessingStrategy(_Processor(_gemma_tokenizer())) + seq = [1, 10, 3, 7, 4, 1, 2, 3, 50, 8, 262144, 4] + labels = strategy.process_labels(torch.tensor([seq])) + # boi(50) and soft-image-token(262144) masked post-scan. + assert labels.tolist()[0] == [ + -100, + -100, + -100, + -100, + -100, + -100, + -100, + -100, + -100, + 8, + -100, + 4, + ] + + +def test_gemma3n_masks_image_and_audio_attrs(): + tok = _gemma_tokenizer() + # Gemma3n exposes these as integer attrs on the tokenizer. + tok.image_token_id = 70 + tok.audio_token_id = 71 + tok.boi_token_id = 72 + tok.eoi_token_id = 73 + strategy = Gemma3nProcessingStrategy(_Processor(tok)) + seq = [1, 2, 3, 70, 71, 72, 73, 9, 4] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + assert out == [-100, -100, -100, -100, -100, -100, -100, 9, 4] + + +# --------------------------------------------------------------------------- # +# Gemma 4 +# --------------------------------------------------------------------------- # + + +class _FakeGemma4Tokenizer(_Tokenizer): + """Mirrors google/gemma-4-E2B-it token layout. Gemma4 role-start markers + include the trailing newline so the boundary matches the jinja template.""" + + VOCAB = { + "<|turn>model\n": [105, 4368, 108], + "<|turn>user\n": [105, 7777, 108], + "<|turn>system\n": [105, 8888, 108], + "": [106], + "<|image|>": [258880], + "<|video|>": [258884], + "<|audio|>": [258881], + "<|image>": [255999], + "": [258882], + "<|audio>": [256000], + "": [258883], + } + + def __init__(self): + # Pass a fresh dict so per-instance mutations (should any future + # code path introduce them) cannot leak across tests via the + # shared class-level VOCAB. + super().__init__( + {token: list(ids) for token, ids in self.VOCAB.items()}, + pad_id=0, + unk_id=3, + ) + + +class _FakeGemma4Processor: + def __init__(self): + self.tokenizer = _FakeGemma4Tokenizer() + self.tokenizer.image_token_id = self.tokenizer.vocab["<|image|>"][0] + self.tokenizer.audio_token_id = self.tokenizer.vocab["<|audio|>"][0] + self.image_token = "<|image|>" + self.image_token_id = self.tokenizer.vocab["<|image|>"][0] + self.boi_token = "<|image>" + self.eoi_token = "" + self.video_token = "<|video|>" + self.video_token_id = self.tokenizer.vocab["<|video|>"][0] + self.audio_token = "<|audio|>" + self.audio_token_id = self.tokenizer.vocab["<|audio|>"][0] + self.boa_token = "<|audio>" + self.eoa_token = "" + + +def test_gemma4_masks_everything_outside_assistant_span(): + strategy = Gemma4ProcessingStrategy(_FakeGemma4Processor()) + V = strategy.processor.tokenizer.vocab + user_start = V["<|turn>user\n"] + model_start = V["<|turn>model\n"] + turn_end = V[""][0] + seq = [ + 0, + *user_start, + 4444, + turn_end, + *model_start, + 5555, + turn_end, + 9999, + ] + labels = strategy.process_labels(torch.tensor([seq])) + expected = [-100] * (1 + len(user_start) + 1 + 1 + len(model_start)) + [ + 5555, + turn_end, + -100, + ] + assert labels.tolist()[0] == expected + + +def test_gemma4_masks_media_tokens_inside_assistant_span(): + strategy = Gemma4ProcessingStrategy(_FakeGemma4Processor()) + V = strategy.processor.tokenizer.vocab + model_start = V["<|turn>model\n"] + media = [ + V["<|image|>"][0], + V["<|video|>"][0], + V["<|audio|>"][0], + V["<|image>"][0], + V[""][0], + V["<|audio>"][0], + V[""][0], + ] + turn_end = V[""][0] + seq = [*model_start, *media, 9999, turn_end] + labels = strategy.process_labels(torch.tensor([seq])) + expected = [-100] * (len(model_start) + len(media)) + [9999, turn_end] + assert labels.tolist()[0] == expected + + +def test_gemma4_multiple_assistant_turns(): + strategy = Gemma4ProcessingStrategy(_FakeGemma4Processor()) + V = strategy.processor.tokenizer.vocab + turn_end = V[""][0] + + def user_turn(x): + return [*V["<|turn>user\n"], x, turn_end] + + def model_turn(x): + return [*V["<|turn>model\n"], x, turn_end] + + seq = user_turn(1111) + model_turn(2222) + user_turn(3333) + model_turn(4444) + labels = strategy.process_labels(torch.tensor([seq])) + kept = [t for t in labels.tolist()[0] if t != -100] + assert kept == [2222, turn_end, 4444, turn_end] + + +# --------------------------------------------------------------------------- # +# Llama 3.2 Vision / Llama 4 +# --------------------------------------------------------------------------- # + + +def test_llama3_2_vision_assistant_masking(): + vocab = { + "<|start_header_id|>assistant<|end_header_id|>\n\n": [1, 2, 3, 4, 5], + "<|start_header_id|>user<|end_header_id|>\n\n": [1, 2, 6, 4, 5], + "<|start_header_id|>system<|end_header_id|>\n\n": [1, 2, 7, 4, 5], + "<|start_header_id|>tool<|end_header_id|>\n\n": [1, 2, 8, 4, 5], + "<|start_header_id|>ipython<|end_header_id|>\n\n": [1, 2, 9, 4, 5], + "<|eot_id|>": [10], + } + strategy = Llama3_2VisionProcessingStrategy(_Processor(_Tokenizer(vocab, pad_id=0))) + seq = [1, 2, 6, 4, 5, 11, 10, 1, 2, 3, 4, 5, 12, 10] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + assert out == [-100] * 12 + [12, 10] + + +def test_llama4_assistant_masking(): + vocab = { + "<|header_start|>assistant<|header_end|>\n\n": [20, 21, 22, 23], + "<|header_start|>user<|header_end|>\n\n": [20, 21, 24, 23], + "<|header_start|>system<|header_end|>\n\n": [20, 21, 25, 23], + "<|header_start|>tool<|header_end|>\n\n": [20, 21, 26, 23], + "<|header_start|>ipython<|header_end|>\n\n": [20, 21, 27, 23], + "<|eot|>": [30], + } + strategy = Llama4ProcessingStrategy(_Processor(_Tokenizer(vocab, pad_id=0))) + seq = [20, 21, 24, 23, 100, 30, 20, 21, 22, 23, 200, 30] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + assert out == [-100] * 10 + [200, 30] + + +# --------------------------------------------------------------------------- # +# Pixtral / Mistral v7 Tekken (eos-terminated assistant) +# --------------------------------------------------------------------------- # + + +def test_pixtral_assistant_terminates_at_eos(): + # [/INST] is both user-end and assistant-start. Scanner backs up when + # user.include_end=False so the next iteration picks [/INST] up as + # assistant-start (Pixtral-specific handling in _build_role_boundaries). + vocab = { + "[INST]": [50], + "[/INST]": [51], + } + tok = _Tokenizer(vocab, pad_id=0, eos_id=99) + strategy = PixtralProcessingStrategy(_Processor(tok)) + seq = [50, 7, 51, 8, 8, 99] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + # Full-sequence expectation: user span masked; assistant content + eos kept. + assert out == [-100, -100, -100, 8, 8, 99] + + +def test_mistral_v7_tekken_system_user_assistant(): + vocab = { + "[SYSTEM_PROMPT]": [40], + "[/SYSTEM_PROMPT]": [41], + "[INST]": [50], + "[/INST]": [51], + } + tok = _Tokenizer(vocab, pad_id=0, eos_id=99) + strategy = MistralV7TekkenProcessingStrategy(_Processor(tok)) + seq = [40, 5, 41, 50, 7, 51, 8, 99] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + # Full-sequence expectation: system + user spans masked; assistant kept. + assert out == [-100, -100, -100, -100, -100, -100, 8, 99] + + +def test_pixtral_train_on_eos_all_respects_user_include_end_false(): + """Pixtral [/INST] (user-end include_end=False) stays masked on 'all'.""" + vocab = {"[INST]": [50], "[/INST]": [51]} + tok = _Tokenizer(vocab, pad_id=0, eos_id=99) + strategy = PixtralProcessingStrategy(_Processor(tok), train_on_eos="all") + seq = [50, 7, 51, 8, 8, 99] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + # [/INST] at idx 2 must stay masked — user.include_end=False says so. + # Assistant content (8, 8) + EOS (99) are unmasked as normal. + assert out == [-100, -100, -100, 8, 8, 99] + + +def test_mistral_v7_tekken_train_on_eos_all_respects_user_include_end_false(): + """System end (include_end=True) unmasked on 'all'; [/INST] stays masked.""" + vocab = { + "[SYSTEM_PROMPT]": [40], + "[/SYSTEM_PROMPT]": [41], + "[INST]": [50], + "[/INST]": [51], + } + tok = _Tokenizer(vocab, pad_id=0, eos_id=99) + strategy = MistralV7TekkenProcessingStrategy(_Processor(tok), train_on_eos="all") + seq = [40, 5, 41, 50, 7, 51, 8, 99] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + # system content masked, [/SYSTEM_PROMPT]=41 kept (include_end=True + all); + # user + [/INST]=51 masked (include_end=False); assistant 8 + eos 99 kept. + assert out == [-100, -100, 41, -100, -100, -100, 8, 99] + + +# --------------------------------------------------------------------------- # +# Dispatcher routing +# --------------------------------------------------------------------------- # + + +def _dispatch(processor, chat_template_type): + return get_processing_strategy( + processor=processor, + chat_template=None, + chat_template_type=chat_template_type, + ) + + +def test_dispatch_qwen2_vl(): + s = _dispatch(_Processor(_qwen_tokenizer()), "qwen2_vl") + assert isinstance(s, Qwen2VLProcessingStrategy) + + +def test_dispatch_qwen3_5(): + s = _dispatch(_Processor(_qwen_tokenizer()), "qwen3_5") + assert isinstance(s, Qwen3_5ProcessingStrategy) + + +def test_dispatch_gemma3(): + s = _dispatch(_Processor(_gemma_tokenizer()), "gemma3") + assert isinstance(s, Gemma3ProcessingStrategy) + + +def test_dispatch_gemma3n(): + s = _dispatch(_Processor(_gemma_tokenizer()), "gemma3n") + assert isinstance(s, Gemma3nProcessingStrategy) + + +def test_dispatch_gemma4(): + s = _dispatch(_FakeGemma4Processor(), "gemma4") + assert isinstance(s, Gemma4ProcessingStrategy) + + +def test_dispatch_llama3_2_vision(): + vocab = { + "<|start_header_id|>assistant<|end_header_id|>\n\n": [1, 2, 3, 4, 5], + "<|eot_id|>": [10], + } + s = _dispatch(_Processor(_Tokenizer(vocab, pad_id=0)), "llama3_2_vision") + assert isinstance(s, Llama3_2VisionProcessingStrategy) + + +def test_dispatch_llama4(): + vocab = { + "<|header_start|>assistant<|header_end|>\n\n": [20, 21, 22, 23], + "<|eot|>": [30], + } + s = _dispatch(_Processor(_Tokenizer(vocab, pad_id=0)), "llama4") + assert isinstance(s, Llama4ProcessingStrategy) + + +def test_dispatch_pixtral(): + vocab = {"[INST]": [50], "[/INST]": [51]} + s = _dispatch(_Processor(_Tokenizer(vocab, pad_id=0, eos_id=99)), "pixtral") + assert isinstance(s, PixtralProcessingStrategy) + + +def test_dispatch_mistral_v7_tekken(): + vocab = { + "[INST]": [50], + "[/INST]": [51], + "[SYSTEM_PROMPT]": [40], + "[/SYSTEM_PROMPT]": [41], + } + s = _dispatch( + _Processor(_Tokenizer(vocab, pad_id=0, eos_id=99)), "mistral_v7_tekken" + ) + assert isinstance(s, MistralV7TekkenProcessingStrategy) + + +def test_dispatch_unknown_falls_back_to_base(): + vocab = {"dummy": [1]} + s = _dispatch(_Processor(_Tokenizer(vocab, pad_id=0)), "llava") + assert type(s) is ProcessingStrategy + + +def _glm_vision_processor(cls_path): + """Spec'd MagicMock so isinstance(mock, cls) passes without real HF files.""" + from importlib import import_module + from unittest.mock import MagicMock + + mod_name, cls_name = cls_path.rsplit(".", 1) + cls = getattr(import_module(mod_name), cls_name) + + vocab = { + "<|image|>": [200], + "<|begin_of_image|>": [201], + "<|end_of_image|>": [202], + "<|video|>": [210], + "<|begin_of_video|>": [211], + "<|end_of_video|>": [212], + } + tok = _Tokenizer(vocab, pad_id=0) + proc = MagicMock(spec=cls) + proc.tokenizer = tok + # Drop processor.image_token so base class skips its probe. + del proc.image_token + return proc + + +def test_dispatch_glm4v_via_Glm4vProcessor(): + """Glm4vProcessor (GLM-4V) routes to Glm4vProcessingStrategy.""" + pytest.importorskip("transformers.models.glm4v.processing_glm4v") + from axolotl.processing_strategies import Glm4vProcessingStrategy + + proc = _glm_vision_processor( + "transformers.models.glm4v.processing_glm4v.Glm4vProcessor" + ) + s = _dispatch(proc, None) + assert isinstance(s, Glm4vProcessingStrategy) + + +def test_dispatch_glm4v_via_Glm46VProcessor(): + """Glm46VProcessor (GLM-4.6V) also routes to Glm4vProcessingStrategy.""" + pytest.importorskip("transformers.models.glm46v.processing_glm46v") + from axolotl.processing_strategies import Glm4vProcessingStrategy + + proc = _glm_vision_processor( + "transformers.models.glm46v.processing_glm46v.Glm46VProcessor" + ) + s = _dispatch(proc, None) + assert isinstance(s, Glm4vProcessingStrategy) + + +# --------------------------------------------------------------------------- # +# Config-based role-boundary override +# --------------------------------------------------------------------------- # + + +def test_role_boundaries_override_replaces_built_in(): + """Override swaps the built-in boundaries wholesale, not additively.""" + vocab = { + "<|im_start|>assistant\n": [101, 102, 103], + "<|im_start|>user\n": [101, 106, 103], + "<|im_end|>": [104], + ">>>A": [200, 201], + ">>>U": [200, 202], + "<<<": [210], + "<|image_pad|>": [250], + } + strategy = Qwen2VLProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + role_boundaries_override=[ + {"role": "assistant", "start": ">>>A", "end": "<<<"}, + {"role": "user", "start": ">>>U", "end": "<<<"}, + ], + ) + seq = [ + 101, + 106, + 103, + 7, + 104, + 200, + 201, + 9, + 9, + 210, + ] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + assert out == [-100, -100, -100, -100, -100, -100, -100, 9, 9, 210] + + +def test_role_boundaries_override_enables_unverified_strategy(): + """Override lets users opt in to role masking on strategies that default opt out.""" + vocab = { + "BOA": [50, 51], + "EOT": [60], + } + strategy = ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + role_boundaries_override=[ + {"role": "assistant", "start": "BOA", "end": "EOT"}, + ], + ) + seq = [1, 2, 3, 50, 51, 7, 8, 60, 9] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + assert out == [-100, -100, -100, -100, -100, 7, 8, 60, -100] + + +def test_role_boundaries_override_eos_token_sentinel(): + vocab = {"BOA": [50]} + tok = _Tokenizer(vocab, pad_id=0, eos_id=99) + strategy = ProcessingStrategy( + _Processor(tok), + role_boundaries_override=[ + {"role": "assistant", "start": "BOA", "end": "eos_token"}, + ], + ) + seq = [1, 50, 7, 7, 99, 2] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + assert out == [-100, -100, 7, 7, 99, -100] + + +def test_role_boundaries_override_end_null_runs_to_sequence_end(): + vocab = {"BOA": [50]} + strategy = ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + role_boundaries_override=[ + {"role": "assistant", "start": "BOA", "end": None}, + ], + ) + seq = [1, 2, 50, 7, 8, 9] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + assert out == [-100, -100, -100, 7, 8, 9] + + +def test_role_boundaries_override_rejects_bad_spec(): + vocab = {"BOA": [50]} + with pytest.raises(ValueError, match="must have both 'role' and 'start'"): + ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + role_boundaries_override=[{"role": "assistant"}], + ) + + +def test_role_boundaries_override_rejects_unencodable_start(): + vocab = {"BOA": [50]} + with pytest.raises(ValueError, match="tokenizes to an empty sequence"): + ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + role_boundaries_override=[ + {"role": "assistant", "start": "MISSING", "end": None} + ], + ) + + +def test_role_boundaries_override_rejects_unencodable_end(): + vocab = {"BOA": [50]} + with pytest.raises(ValueError, match="tokenizes to an empty sequence"): + ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + role_boundaries_override=[ + {"role": "assistant", "start": "BOA", "end": "MISSING"} + ], + ) + + +def test_role_boundaries_override_accepts_pydantic_models(): + # cfg.role_boundaries arrives as RoleBoundarySpec after pydantic parsing. + from axolotl.utils.schemas.multimodal import RoleBoundarySpec + + vocab = {"BOA": [50], "EOT": [60]} + strategy = ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + role_boundaries_override=[ + RoleBoundarySpec(role="assistant", start="BOA", end="EOT") + ], + ) + assert len(strategy.role_boundaries) == 1 + assert strategy.role_boundaries[0].role == "assistant" + assert strategy.role_boundaries[0].start_tokens == [50] + assert strategy.role_boundaries[0].end_tokens == [60] + + +def test_base_strategy_warns_when_no_boundaries(axolotl_caplog): + """No boundaries + train_on_inputs=False: one-shot warning, labels unchanged.""" + import axolotl.processing_strategies as mod + + mod._ROLE_MASK_WARNED.discard("ProcessingStrategy") + + vocab = {"dummy": [1]} + s = ProcessingStrategy(_Processor(_Tokenizer(vocab, pad_id=0))) + + with axolotl_caplog.at_level( + logging.WARNING, logger="axolotl.processing_strategies" + ): + labels = s.process_labels(torch.tensor([[1, 2, 3]])) + assert labels.tolist() == [[1, 2, 3]] + assert any("role boundaries" in rec.message for rec in axolotl_caplog.records) + + +# --------------------------------------------------------------------------- # +# Additional edge-case coverage +# --------------------------------------------------------------------------- # + + +def test_scanner_batch_size_greater_than_one(): + boundaries = [ + RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]), + RoleBoundary(role="user", start_tokens=[1, 3], end_tokens=[9]), + ] + labels = torch.tensor( + [ + [1, 3, 7, 9, 1, 2, 8, 9], + [1, 2, 5, 5, 9, 0, 0, 0], + ] + ) + out = _apply_role_boundaries(labels, boundaries, {"assistant"}, "turn").tolist() + assert out[0] == [-100, -100, -100, -100, -100, -100, 8, 9] + assert out[1] == [-100, -100, 5, 5, 9, -100, -100, -100] + + +def test_scanner_adjacent_trainable_turns(): + boundaries = [ + RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]), + ] + seq = [1, 2, 5, 9, 1, 2, 6, 9] + out = _scan(boundaries, seq) + assert out == [-100, -100, 5, 9, -100, -100, 6, 9] + + +def test_scanner_train_on_eos_none_multi_turn(): + boundaries = [ + RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]), + RoleBoundary(role="user", start_tokens=[1, 3], end_tokens=[9]), + ] + seq = [1, 3, 7, 9, 1, 2, 8, 9, 1, 3, 7, 9, 1, 2, 6, 9] + out = _scan(boundaries, seq, train_on_eos="none") + assert out == [ + -100, + -100, + -100, + -100, + -100, + -100, + 8, + -100, + -100, + -100, + -100, + -100, + -100, + -100, + 6, + -100, + ] + + +def test_scanner_train_on_eos_all_with_user_turn_no_end_marker(): + """Unclosed non-trainable span with train_on_eos='all': nothing included, no crash.""" + boundaries = [ + RoleBoundary(role="assistant", start_tokens=[1, 2], end_tokens=[9]), + RoleBoundary(role="user", start_tokens=[1, 3], end_tokens=[9]), + ] + seq = [1, 3, 7, 7, 7] + out = _scan(boundaries, seq, train_on_eos="all") + assert out == [-100, -100, -100, -100, -100] + + +def test_scanner_include_start_true_via_override(): + vocab = {"BOA": [50, 51], "EOT": [60]} + strategy = ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + role_boundaries_override=[ + { + "role": "assistant", + "start": "BOA", + "end": "EOT", + "include_start": True, + }, + ], + ) + seq = [1, 50, 51, 7, 8, 60, 9] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + assert out == [-100, 50, 51, 7, 8, 60, -100] + + +def test_scanner_include_end_false_via_override(): + """include_end=False drops end marker even with train_on_eos='turn'.""" + vocab = {"BOA": [50], "EOT": [60]} + strategy = ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + role_boundaries_override=[ + { + "role": "assistant", + "start": "BOA", + "end": "EOT", + "include_end": False, + }, + ], + ) + seq = [1, 50, 7, 8, 60, 9] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + assert out == [-100, -100, 7, 8, -100, -100] + + +def test_scanner_empty_start_tokens_is_defensive_noop(): + """Defensive: empty start_tokens matches nothing; everything masked.""" + boundaries = [ + RoleBoundary(role="assistant", start_tokens=[], end_tokens=[9]), + ] + seq = [1, 2, 3, 4, 9] + out = _scan(boundaries, seq) + assert out == [-100] * 5 + + +def test_process_labels_masks_pad_inside_assistant_span(): + """Pad inside a trainable span is still masked post-scan.""" + strategy = _make_qwen2vl() + seq = [101, 102, 103, 8, 0, 8, 104] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + assert out == [-100, -100, -100, 8, -100, 8, 104] + + +def test_process_labels_all_pad_sequence_does_not_crash(): + strategy = _make_qwen2vl() + seq = [0, 0, 0, 0] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + assert out == [-100, -100, -100, -100] + + +def test_qwen2vl_multiple_consecutive_assistant_turns(): + strategy = _make_qwen2vl() + seq = [101, 102, 103, 8, 104, 101, 102, 103, 9, 104] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + assert out == [ + -100, + -100, + -100, + 8, + 104, + -100, + -100, + -100, + 9, + 104, + ] + + +def test_qwen2vl_batch_of_two_rows(): + strategy = _make_qwen2vl() + row_a = [101, 106, 103, 7, 104, 101, 102, 103, 8, 104] + row_b = [101, 102, 103, 9, 104, 0, 0, 0, 0, 0] + out = strategy.process_labels(torch.tensor([row_a, row_b])).tolist() + assert out[0] == [-100, -100, -100, -100, -100, -100, -100, -100, 8, 104] + assert out[1] == [-100, -100, -100, 9, 104, -100, -100, -100, -100, -100] + + +def test_qwen3_5_train_on_inputs_true_still_masks_video_pad(): + """train_on_inputs=True skips role masking but media tokens are still masked.""" + tok = _qwen_tokenizer() + strategy = Qwen3_5ProcessingStrategy(_Processor(tok), train_on_inputs=True) + seq = [101, 106, 103, 201, 7, 104, 101, 102, 103, 201, 8, 104] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + expected = list(seq) + expected[3] = -100 + expected[9] = -100 + assert out == expected + + +def test_role_boundaries_override_role_not_in_roles_to_train(): + """Override covering only a non-trainable role masks everything.""" + vocab = {"BOU": [50], "EOT": [60]} + strategy = ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + role_boundaries_override=[ + {"role": "user", "start": "BOU", "end": "EOT"}, + ], + ) + seq = [1, 50, 7, 8, 60, 9] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + assert out == [-100] * 6 + + +def test_role_boundaries_override_include_start_flag_round_trips(): + from axolotl.utils.schemas.multimodal import RoleBoundarySpec + + vocab = {"BOA": [50], "EOT": [60]} + strategy = ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + role_boundaries_override=[ + RoleBoundarySpec( + role="assistant", start="BOA", end="EOT", include_start=True + ), + ], + ) + assert len(strategy.role_boundaries) == 1 + assert strategy.role_boundaries[0].include_start is True + assert strategy.role_boundaries[0].include_end is True + + +def test_multimodal_config_parses_dict_role_boundaries_to_specs(): + from axolotl.utils.schemas.multimodal import ( + MultiModalConfig, + RoleBoundarySpec, + ) + + cfg = MultiModalConfig( + role_boundaries=[ + {"role": "assistant", "start": "BOA", "end": "EOT"}, + {"role": "user", "start": "BOU", "end": "EOT"}, + ] + ) + assert cfg.role_boundaries is not None + assert len(cfg.role_boundaries) == 2 + assert all(isinstance(rb, RoleBoundarySpec) for rb in cfg.role_boundaries) + + vocab = {"BOA": [50], "BOU": [51], "EOT": [60]} + strategy = ProcessingStrategy( + _Processor(_Tokenizer(vocab, pad_id=0)), + role_boundaries_override=cfg.role_boundaries, + ) + seq = [51, 7, 60, 50, 8, 60] + out = strategy.process_labels(torch.tensor([seq])).tolist()[0] + assert out == [-100, -100, -100, -100, 8, 60] diff --git a/tests/utils/data/test_mm_cpt_eval.py b/tests/utils/data/test_mm_cpt_eval.py new file mode 100644 index 0000000000..57ff10897b --- /dev/null +++ b/tests/utils/data/test_mm_cpt_eval.py @@ -0,0 +1,186 @@ +"""Multimodal CPT eval-path tests.""" + +from __future__ import annotations + +import pytest + +from axolotl.utils.data.sft import ( + _create_placeholder_dataset, + _prepare_streaming_dataset, +) +from axolotl.utils.dict import DictDefault + +# ---- placeholder dataset for dispatch_batches ---------------------------- + + +def test_placeholder_text_only_keeps_existing_shape(): + """Without an MM config, the placeholder is a single-column text dataset.""" + ds = _create_placeholder_dataset() + row = next(iter(ds)) + assert "text" in row + assert "images" not in row + + +def test_placeholder_mm_emits_image_column(): + """MM placeholder rows carry the configured image column as an empty list.""" + pt_cfg = DictDefault( + { + "type": "multimodal_pretrain", + "text_column": "text", + "image_column": "images", + "multimodal": True, + } + ) + ds = _create_placeholder_dataset(pt_cfg) + row = next(iter(ds)) + assert "text" in row + assert "images" in row + assert row["images"] == [] + + +def test_placeholder_mm_honors_custom_columns(): + """Custom text_column / image_column on the MM config are reflected in the placeholder row.""" + pt_cfg = DictDefault( + { + "type": "multimodal_pretrain", + "text_column": "doc", + "image_column": "imgs", + } + ) + ds = _create_placeholder_dataset(pt_cfg) + row = next(iter(ds)) + assert "doc" in row + assert "imgs" in row + assert row["imgs"] == [] + + +# ---- multiple MM eval datasets are loaded -------------------------------- + + +def test_mm_eval_iterates_all_test_datasets(monkeypatch): + """All MM entries in test_datasets are loaded and concatenated into the eval stream.""" + cfg = DictDefault( + { + "streaming": True, + "pretraining_dataset": [ + {"path": "train/ds", "type": "multimodal_pretrain"} + ], + "test_datasets": [ + {"path": "eval/a", "type": "multimodal_pretrain"}, + {"path": "eval/b", "type": "multimodal_pretrain"}, + {"path": "eval/c", "type": "multimodal_pretrain"}, + ], + "max_steps": 10, + } + ) + + seen_eval_paths: list[str] = [] + + def fake_load_streaming(pretraining_config, *_a, **_kw): + path = pretraining_config["path"] + if path.startswith("eval/"): + seen_eval_paths.append(path) + return f"" + + def fake_concat(streams): + return tuple(streams) + + monkeypatch.setattr( + "axolotl.utils.data.sft._load_streaming_dataset", fake_load_streaming + ) + monkeypatch.setattr("axolotl.utils.data.sft.concatenate_datasets", fake_concat) + + train, eval_ds, _, _ = _prepare_streaming_dataset( + cfg, tokenizer=None, processor=None + ) + + assert seen_eval_paths == ["eval/a", "eval/b", "eval/c"] + assert eval_ds == ("", "", "") + + +def test_mm_eval_rejects_mixed_mm_and_non_mm_test_datasets(monkeypatch): + """MM CPT runs require every test_datasets entry to be MM; mixed lists raise.""" + cfg = DictDefault( + { + "streaming": True, + "pretraining_dataset": [ + {"path": "train/ds", "type": "multimodal_pretrain"} + ], + "test_datasets": [ + {"path": "eval/a", "type": "multimodal_pretrain"}, + # Plain text eval entry — not allowed alongside MM eval. + {"path": "eval/b", "type": "pretrain"}, + ], + "max_steps": 10, + } + ) + monkeypatch.setattr( + "axolotl.utils.data.sft._load_streaming_dataset", + lambda *_a, **_kw: "", + ) + with pytest.raises(ValueError, match="multimodal and non-multimodal"): + _prepare_streaming_dataset(cfg, tokenizer=None, processor=None) + + +# ---- eval collator pulls image settings from test_datasets --------------- + + +def test_eval_collator_uses_eval_image_settings(monkeypatch): + """Eval collator pulls image_base_dir / image_token from test_datasets[0]; train collator from pretraining_dataset[0].""" + from axolotl.core.builders.causal import HFCausalTrainerBuilder + + captured = {} + + class _FakeSpec: + image_token = "" + image_token_id = 7 + image_family_token_ids = (7,) + + def fake_build_image_token_spec(processor, override=None): + captured["override"] = override + return _FakeSpec() + + monkeypatch.setattr( + "axolotl.prompt_strategies.multimodal_pretrain.build_image_token_spec", + fake_build_image_token_spec, + ) + + class _FakeCollator: + def __init__(self, **kw): + captured["kwargs"] = kw + + monkeypatch.setattr( + "axolotl.core.builders.causal.MultiModalPretrainDataCollator", _FakeCollator + ) + + builder = HFCausalTrainerBuilder.__new__(HFCausalTrainerBuilder) + builder.tokenizer = object() + builder.processor = object() + builder.cfg = DictDefault( + { + "pretraining_dataset": [ + { + "type": "multimodal_pretrain", + "image_base_dir": "/train_images", + "image_token": "", + } + ], + "test_datasets": [ + { + "type": "multimodal_pretrain", + "image_base_dir": "/eval_images", + "image_token": "", + } + ], + "sequence_len": 2048, + } + ) + + builder._build_mm_pretrain_collator(is_eval=True) + assert captured["override"] == "" + assert captured["kwargs"]["image_base_dir"] == "/eval_images" + + captured.clear() + builder._build_mm_pretrain_collator(is_eval=False) + assert captured["override"] == "" + assert captured["kwargs"]["image_base_dir"] == "/train_images" diff --git a/tests/utils/schemas/validation/test_multimodal_cpt.py b/tests/utils/schemas/validation/test_multimodal_cpt.py new file mode 100644 index 0000000000..216b78f47a --- /dev/null +++ b/tests/utils/schemas/validation/test_multimodal_cpt.py @@ -0,0 +1,270 @@ +"""Multimodal CPT config validation gates.""" + +from __future__ import annotations + +import logging + +import pytest + +from axolotl.utils.config import validate_config +from axolotl.utils.dict import DictDefault + + +def _mm_cpt_cfg(min_base_cfg, **overrides) -> DictDefault: + base = DictDefault( + **( + min_base_cfg + | { + "datasets": None, + "pretraining_dataset": [ + { + "path": "some/ds", + "type": "multimodal_pretrain", + "image_column": "images", + } + ], + "streaming": True, + "max_steps": 10, + "processor_type": "AutoProcessor", + "sequence_len": 2048, + } + ) + ) + return base | DictDefault(overrides) + + +class TestMultimodalCPTGates: + def test_missing_processor_type_raises(self, min_base_cfg): + cfg = _mm_cpt_cfg(min_base_cfg) + cfg.pop("processor_type", None) + with pytest.raises(ValueError, match="processor_type"): + validate_config(cfg) + + def test_sample_packing_rejected(self, min_base_cfg): + cfg = _mm_cpt_cfg(min_base_cfg, sample_packing=True) + with pytest.raises(ValueError, match="sample_packing"): + validate_config(cfg) + + def test_chat_template_rejected(self, min_base_cfg): + cfg = _mm_cpt_cfg(min_base_cfg, chat_template="tokenizer_default") + with pytest.raises(ValueError, match="chat_template"): + validate_config(cfg) + + def test_multiple_pretraining_dataset_entries_rejected(self, min_base_cfg): + cfg = _mm_cpt_cfg(min_base_cfg) + cfg.pretraining_dataset.append({"path": "other/ds", "type": "pretrain"}) + with pytest.raises(ValueError, match="exactly one `pretraining_dataset`"): + validate_config(cfg) + + def test_multimodal_entry_in_non_first_slot_rejected(self, min_base_cfg): + cfg = DictDefault( + **( + min_base_cfg + | { + "datasets": None, + "pretraining_dataset": [ + {"path": "text/ds", "type": "pretrain"}, + { + "path": "mm/ds", + "type": "multimodal_pretrain", + "image_column": "images", + }, + ], + "streaming": True, + "max_steps": 10, + "processor_type": "AutoProcessor", + "sequence_len": 2048, + } + ) + ) + with pytest.raises(ValueError, match="exactly one `pretraining_dataset`"): + validate_config(cfg) + + def test_valid_cfg_passes_and_disables_remove_unused_columns(self, min_base_cfg): + cfg = _mm_cpt_cfg(min_base_cfg) + validated = validate_config(cfg) + assert validated.remove_unused_columns is False + pd = validated.pretraining_dataset[0] + assert pd.type == "multimodal_pretrain" + assert pd.image_column == "images" + + def test_multimodal_flag_triggers_gates(self, min_base_cfg): + cfg = _mm_cpt_cfg(min_base_cfg) + cfg.pretraining_dataset[0]["type"] = "pretrain" + cfg.pretraining_dataset[0]["multimodal"] = True + cfg.pop("processor_type", None) + with pytest.raises(ValueError, match="processor_type"): + validate_config(cfg) + + def test_non_mm_pretraining_dataset_unaffected(self, min_base_cfg): + cfg = DictDefault( + **( + min_base_cfg + | { + "datasets": None, + "pretraining_dataset": [{"path": "some/ds", "type": "pretrain"}], + "streaming": True, + "max_steps": 10, + "sequence_len": 2048, + } + ) + ) + validate_config(cfg) + + def test_mm_eval_dataset_keys_preserved_through_validation(self, min_base_cfg): + """MM-specific keys on a test_datasets entry survive validate_config.""" + cfg = _mm_cpt_cfg( + min_base_cfg, + test_datasets=[ + { + "path": "eval/ds", + "type": "multimodal_pretrain", + "text_column": "eval_text", + "image_column": "eval_imgs", + "image_base_dir": "/eval/images", + "image_token": "", + } + ], + ) + validated = validate_config(cfg) + td = validated.test_datasets[0] + assert td["text_column"] == "eval_text" + assert td["image_column"] == "eval_imgs" + assert td["image_base_dir"] == "/eval/images" + assert td["image_token"] == "" + + def test_mm_eval_dataset_via_multimodal_flag(self, min_base_cfg): + """`multimodal: true` (without type='multimodal_pretrain') opts an eval entry into MM.""" + cfg = _mm_cpt_cfg( + min_base_cfg, + test_datasets=[ + { + "path": "eval/ds", + "multimodal": True, + "image_column": "imgs2", + } + ], + ) + validated = validate_config(cfg) + td = validated.test_datasets[0] + assert td["image_column"] == "imgs2" + assert td["multimodal"] is True + + def test_non_mm_eval_entry_does_not_match_mm_model(self, min_base_cfg): + """SFT eval entries (no MM markers) still validate as SFTDataset.""" + cfg = DictDefault( + **( + min_base_cfg + | { + "test_datasets": [ + {"path": "eval/ds", "type": "alpaca", "split": "test"} + ], + "sequence_len": 2048, + } + ) + ) + validated = validate_config(cfg) + td = validated.test_datasets[0] + assert "message_property_mappings" in td + assert td["type"] == "alpaca" + + def test_mm_eval_rejects_mismatched_image_base_dir(self, min_base_cfg): + """Multiple MM eval entries with different image_base_dir are rejected.""" + cfg = _mm_cpt_cfg( + min_base_cfg, + test_datasets=[ + { + "path": "eval/a", + "type": "multimodal_pretrain", + "image_base_dir": "/images/a", + }, + { + "path": "eval/b", + "type": "multimodal_pretrain", + "image_base_dir": "/images/b", + }, + ], + ) + with pytest.raises(ValueError, match="image_base_dir"): + validate_config(cfg) + + def test_mm_eval_rejects_mismatched_image_token(self, min_base_cfg): + """Multiple MM eval entries with different image_token overrides are rejected.""" + cfg = _mm_cpt_cfg( + min_base_cfg, + test_datasets=[ + { + "path": "eval/a", + "type": "multimodal_pretrain", + "image_token": "", + }, + { + "path": "eval/b", + "type": "multimodal_pretrain", + "image_token": "", + }, + ], + ) + with pytest.raises(ValueError, match="image_token"): + validate_config(cfg) + + def test_mm_eval_accepts_matching_image_base_dir(self, min_base_cfg): + """Multiple MM eval entries sharing image_base_dir validate cleanly.""" + cfg = _mm_cpt_cfg( + min_base_cfg, + test_datasets=[ + { + "path": "eval/a", + "type": "multimodal_pretrain", + "image_base_dir": "/images/shared", + }, + { + "path": "eval/b", + "type": "multimodal_pretrain", + "image_base_dir": "/images/shared", + }, + ], + ) + validated = validate_config(cfg) + assert len(validated.test_datasets) == 2 + + def test_mm_eval_accepts_all_unset_image_settings(self, min_base_cfg): + """Multiple MM eval entries with image_base_dir / image_token unset everywhere validate.""" + cfg = _mm_cpt_cfg( + min_base_cfg, + test_datasets=[ + {"path": "eval/a", "type": "multimodal_pretrain"}, + {"path": "eval/b", "type": "multimodal_pretrain"}, + ], + ) + validated = validate_config(cfg) + assert len(validated.test_datasets) == 2 + + def test_remove_unused_columns_auto_set_emits_info_log(self, min_base_cfg, caplog): + """Auto-setting `remove_unused_columns: false` for MM CPT logs an INFO record naming the previous value.""" + cfg = _mm_cpt_cfg(min_base_cfg) + cfg.pop("remove_unused_columns", None) + with caplog.at_level(logging.INFO, logger="axolotl.utils.schemas.validation"): + validated = validate_config(cfg) + assert validated.remove_unused_columns is False + matches = [ + r + for r in caplog.records + if r.levelno == logging.INFO and "Auto-set" in r.getMessage() + ] + assert matches, "expected an INFO record about auto-setting remove_unused_columns" + msg = matches[0].getMessage() + assert "remove_unused_columns" in msg + assert "previous value: None" in msg + + def test_remove_unused_columns_already_false_does_not_log( + self, min_base_cfg, caplog + ): + """When the user already set `remove_unused_columns: false`, no auto-set log fires.""" + cfg = _mm_cpt_cfg(min_base_cfg, remove_unused_columns=False) + with caplog.at_level(logging.INFO, logger="axolotl.utils.schemas.validation"): + validate_config(cfg) + assert not any( + "Auto-set" in r.getMessage() and "remove_unused_columns" in r.getMessage() + for r in caplog.records + )