Finetuning Model Onboarding

Finetuning model onboarding

All finetuning runs on Modal with Unsloth (overbae/services/sft_assets/engine_unsloth.py). Every model and family in this pipeline goes through the same engine — the rules below exist to keep it that way.

1. Never branch on model/family in the shared script

engine_unsloth.py and overbae/services/finetuning_runner.py/finetuning_policy.py are shared by every model. If a new model needs different behavior:

• Do not add if model_id == "..." or if family == "..." branches to the engine, runner, or policy.
• Do add/extend a family hooks module under overbae/services/sft_assets/families/<family>.py, subclassing DefaultHooks (families/__init__.py) and overriding only the hook methods you need (env_overrides, load_kwargs, post_load, device_map, sft_config_overrides, peft_lora_dropout, peft_gradient_checkpointing, fast_model_cls). Export a single hooks = XxxHooks().
• Register the family in modal_shared/modelfam/families.py as a FamilySpec (patterns, images, hooks_module="families.<family>"), then add it to the ordered _FAMILIES tuple in modal_shared/modelfam/registry.py — specific patterns before general ones (e.g. qwen35/qwen_coder before qwen_mm before qwen).
• The engine calls hooks uniformly (load_hooks(_spec.hooks_module) then _hooks.<method>(...)) — it never knows which family it's running. That's the contract to preserve.

Reference existing hooks files for scale: qwen35.py (two-line SDPA override) up to gpt_oss.py (monkeypatches + adapter base-model swap) — match the size of the override to the size of the actual quirk, don't build more than the model needs.

2. Chat template must be training-compatible (pretok Path A)

Unsloth labels come from overbae/services/sft_assets/pretok.py. Path A is TRL get_training_chat_template + return_assistant_tokens_mask. That only works if the live tokenizer template already has {% generation %} / {% endgeneration %} markers, or TRL can exact-match it to a bundled original, or we swap in a hand-patched twin.

If the log shows:

diag: pretok_trl unsupported (ValueError('The chat template is not training-compatible (missing prefix-preservation or `{% generation %}` markers) and patching is not supported for this template. ...')); using multi_header fallback

the job may still train, but assistant-only masks are the weaker Path B. Onboarding is not done until Path A succeeds (pretok paths: contains trl_training_template or native_generation_markers, not only multi_header_fallback).

How to check (tokenizer issue — LoRA is enough; skip a second Full run):

1. Probe the catalog hf_model_id, not the catalog id. Unsloth repos often ship a different chat_template.jinja than the upstream id; TRL and our patches match the live string after from_pretrained.
2. Tiny Modal LoRA (MAX_STEPS=2, short MAX_LENGTH) is enough. Read runs/<run_id>/train_stdout.log for pretok paths: / pretok_trl unsupported.
3. One size is not enough when Unsloth templates differ inside a family (Gemma 4 E2B vs 12B, Qwen3-0.6B vs 8B, Qwen3.5-0.8B vs 4B). Hash the Hub template for each hf_model_id you enable.

How to fix — do not branch in engine_unsloth.py / pretok.py:

1. Save the live Hub template as a base jinja under overbae/services/sft_assets/ (existing dirs: llama_templates/, qwen_templates/, gemma_templates/, …).
2. Copy it to *_training.jinja and wrap only assistant-generated spans in {% generation %} / {% endgeneration %}. {% generation %} is a real Jinja block: it cannot open inside {% if %} and close after {% endif %}, and {% if %}/{% else %}/{% endif %} must sit wholly inside or wholly outside it. The training file must render byte-identical text to the base; markers are invisible in the rendered string. Qwen3's twins omit the empty <think> block the base inserts on a final assistant turn: that block is not prefix-preserving, and serving keeps thinking off. Compile every twin (see test_every_training_template_compiles).
3. Register (base, training) in KNOWN_TEMPLATE_PATCHES in overbae/services/sft_assets/training_chat_template.py. pretok.py and engine_unsloth.py both call patch_known_training_template — that is the only dispatch table.
4. Add a pair assertion in tests/test_sft_training_chat_template.py (or rely on test_every_patch_pair_exists_and_training_has_markers).
5. sft_assets is baked into the train image (add_local_dir). A pretok/jinja change does nothing until modal deploy of overbae/modal/modal_sft_worker.py. Re-run the LoRA probe after deploy.

3. If a model fails on its assigned environment, re-derive the GPU, don't patch around it

GPU selection for training lives in finetuning_runner.py's _GPU_TABLE / _GPU_TABLE_FULL (params_b → gpu_type/count) plus the _LONG_CONTEXT_H200 bump and any family-specific clamp (e.g. clamp_gemma4_training_gpu — Gemma4 forced onto 1×H200 because Unsloth's device_map="balanced" multi-GPU split is broken for it).

If a new model OOMs or errors on its assigned GPU:

1. Check whether the failure is architectural (needs a real device_map/env override → family hook, §1) or capacity (needs a different GPU/context ceiling → §4).
2. If it's a one-off model/family quirk, add a clamp or override scoped to that family only (see the Gemma4 precedent), not a new general rule in the shared table.
3. Re-run the context-length probe (§4) rather than guessing a new ceiling.

4. Deriving real_max_context_length and validated_context_length

For every new model, before it's usable for finetuning:

1. Find the real max context length. This is the HF max_position_embeddings (no RoPE/YaRN scaling) — it goes in finetuning.context_length in models.json, distinct from the top-level context_length (published inference window, which may be YaRN-extended).
2. Probe each enabled training type (full, lora) independently. Use scripts/calibrate_activation_budget.py (--experiment e3 context sweep, spawns real Modal sft_unsloth jobs and reads peak VRAM) to find the largest context that actually trains without OOM.
3. Try H100 first, then H200. TRAINING_GPU_VRAM_GB = {"H100": 80.0, "H200": 141.0} (finetuning_policy.py) — H100 has less VRAM, so it's the cheaper GPU and must be tried first. Only fall back to H200 if the model can't reach its real max context length on H100.
4. Record the result per training type in models.json under finetuning.training_type.<full|lora>:
   • context_length: the largest context that trained successfully (equal to finetuning.context_length if the full ceiling was reached, lower otherwise).
   • validated_context_length: true once probed — this flag means "this number came from an actual training run on H100/H200," not an assumption.
5. Set finetuning.real_max_context_length to the max across validated training types, and keep every training_type.*.context_length <= real_max_context_length.

tests/test_modelfam.py::test_baseten_real_max_context_length enforces this shape for every backend: "baseten" entry — run it before considering onboarding done.