DGXAI/driftcall

1"""Gemma 3n E2B boot via Unsloth FastModel (docs/modules/training.md §3.1).2 3Contract:4  - Base model: ``unsloth/gemma-3n-E2B-it`` (4-bit Dynamic5    NF4 quantization).6  - Precision: hardware-aware.7      V100 (sm_70) — explicit FP16 (``dtype=torch.float16``); Gemma 3n is8      BF16-native, so we force FP16 on V100 to avoid BF16 software-emulation9      slowdown / numerical instability.10      H100 (sm_90) — BF16 (``dtype=torch.bfloat16``); uses native tensor cores.11  - LoRA: r=16, α=32, dropout=0.05, vision towers frozen, language + attention12    + MLP trainable via Unsloth's multimodal API (``finetune_vision_layers=False,13    finetune_language_layers=True, finetune_attention_modules=True,14    finetune_mlp_modules=True``), Unsloth gradient checkpointing,15    ``random_state=3407``.16  - V100 halt: ``next(model.parameters()).dtype`` MUST be ``torch.float16``17    after FP16 load; any BF16 parameter triggers :class:`BF16SlippageError`18    before optimizer build.19  - H100 halt: ``next(model.parameters()).dtype`` MUST be ``torch.bfloat16``20    after BF16 load; any FP16 parameter triggers :class:`FP16SlippageError`21    before optimizer build.22 23Heavy imports (``unsloth``, ``torch``) are deferred inside functions so this24cell loads on CPU-only CI runners where Unsloth is not installed. Tests mock25``FastModel.from_pretrained`` and ``FastModel.get_peft_model``.26"""27 28from __future__ import annotations29 30from dataclasses import dataclass31from typing import Any, Literal32 33BASE_MODEL_ID: str = "unsloth/gemma-3n-E2B-it"34MAX_SEQ_LENGTH: int = 409635LORA_R: int = 1636LORA_ALPHA: int = 3237LORA_DROPOUT: float = 0.0538LORA_RANDOM_STATE: int = 340739 40# Gemma 3n multimodal LoRA flags — vision/audio towers stay frozen so GRPO41# trains only the language stack (Unsloth Gemma 3N notebook §fine-tune).42FINETUNE_VISION_LAYERS: bool = False43FINETUNE_LANGUAGE_LAYERS: bool = True44FINETUNE_ATTENTION_MODULES: bool = True45FINETUNE_MLP_MODULES: bool = True46 47HardwareT = Literal["v100", "h100"]48ALLOWED_HARDWARE: tuple[HardwareT, ...] = ("v100", "h100")49 50 51class BF16SlippageError(AssertionError):52    """Raised when the loaded model has any BF16 parameter on V100.53 54    V100 (sm_70) lacks BF16 tensor cores. Silent BF16 via software emulation55    causes ~10x slowdown plus numerical-instability patterns in56    ``docs/modules/training.md §7a``. Halt before the optimizer is built.57    """58 59 60class FP16SlippageError(AssertionError):61    """Raised when the loaded model has any FP16 parameter on H100.62 63    H100 (sm_90) has native BF16 tensor cores. Running FP16 on H100 means64    leaving native hardware capability unused and may cause gradient underflow65    at large batch sizes. Halt before the optimizer is built.66    """67 68 69@dataclass(frozen=True)70class BootConfig:71    """Arguments to :func:`boot_gemma`. Frozen per DriftCall immutability rule."""72 73    base_model_id: str = BASE_MODEL_ID74    max_seq_length: int = MAX_SEQ_LENGTH75    load_in_4bit: bool = True76    lora_r: int = LORA_R77    lora_alpha: int = LORA_ALPHA78    lora_dropout: float = LORA_DROPOUT79    lora_random_state: int = LORA_RANDOM_STATE80    finetune_vision_layers: bool = FINETUNE_VISION_LAYERS81    finetune_language_layers: bool = FINETUNE_LANGUAGE_LAYERS82    finetune_attention_modules: bool = FINETUNE_ATTENTION_MODULES83    finetune_mlp_modules: bool = FINETUNE_MLP_MODULES84    use_gradient_checkpointing: str = "unsloth"85    hardware: HardwareT = "v100"86 87 88def assert_dtype_for_hardware(model: Any, hardware: HardwareT) -> None:89    """Assert the first parameter dtype matches the expected precision for hardware.90 91    V100 must be ``torch.float16``; raises :class:`BF16SlippageError` otherwise.92    H100 must be ``torch.bfloat16``; raises :class:`FP16SlippageError` otherwise.93    Called once at ``boot_gemma`` entry, before any LoRA attach or optimizer build.94    """95    import torch96 97    params_iter = model.parameters()98    try:99        first_param = next(params_iter)100    except StopIteration as exc:  # pragma: no cover - defensive101        raise BF16SlippageError(102            "Model has no parameters; cannot verify dtype."103        ) from exc104 105    dtype = first_param.dtype106    if hardware == "v100":107        if dtype != torch.float16:108            raise BF16SlippageError(109                f"BF16 slipped through: V100 unsafe. "110                f"next(model.parameters()).dtype == {dtype}, expected torch.float16. "111                f"Root cause: Unsloth auto-picked BF16 despite dtype=torch.float16 kwarg. "112                f"Halt training; do NOT proceed on V100."113            )114    else:  # h100115        if dtype != torch.bfloat16:116            raise FP16SlippageError(117                f"FP16 slipped through: H100 should use BF16. "118                f"next(model.parameters()).dtype == {dtype}, expected torch.bfloat16. "119                f"Root cause: dtype kwarg may have forced FP16 on H100. "120                f"Halt training; do NOT proceed on H100 with FP16."121            )122 123 124def assert_fp16_dtype(model: Any) -> None:125    """Assert the first trainable parameter is torch.float16 (V100 safety).126 127    Thin wrapper around :func:`assert_dtype_for_hardware` for backwards128    compatibility with call sites that predate the hardware-aware API.129    Raises :class:`BF16SlippageError` with the halt message from130    ``docs/modules/training.md §3.1``.131    """132    assert_dtype_for_hardware(model, "v100")133 134 135def boot_gemma(config: BootConfig | None = None) -> tuple[Any, Any]:136    """Load Gemma 3n E2B in 4-bit + attach LoRA; return (model, tokenizer).137 138    Steps (training.md §3.1):139      1. ``FastModel.from_pretrained(base_model_id, max_seq_length=...,140         load_in_4bit=True, dtype=torch.float16)`` on V100141         or ``dtype=torch.bfloat16`` on H100.142      2. ``assert_dtype_for_hardware(model, hardware)`` — raises143         :class:`BF16SlippageError` or :class:`FP16SlippageError` if the dtype144         does not match the hardware.145      3. ``FastModel.get_peft_model(model, r=16, lora_alpha=32,146         finetune_vision_layers=False, finetune_language_layers=True,147         finetune_attention_modules=True, finetune_mlp_modules=True,148         use_gradient_checkpointing="unsloth", random_state=3407)``.149      4. Return ``(peft_model, tokenizer)``.150 151    All heavy imports are lazy so the module is importable on CPU-only CI.152    """153    cfg = config if config is not None else BootConfig()154 155    import torch156    from unsloth import FastModel157 158    dtype = torch.float16 if cfg.hardware == "v100" else torch.bfloat16159 160    model, tokenizer = FastModel.from_pretrained(161        cfg.base_model_id,162        max_seq_length=cfg.max_seq_length,163        load_in_4bit=cfg.load_in_4bit,164        dtype=dtype,165    )166 167    assert_dtype_for_hardware(model, cfg.hardware)168 169    peft_model = FastModel.get_peft_model(170        model,171        r=cfg.lora_r,172        lora_alpha=cfg.lora_alpha,173        lora_dropout=cfg.lora_dropout,174        finetune_vision_layers=cfg.finetune_vision_layers,175        finetune_language_layers=cfg.finetune_language_layers,176        finetune_attention_modules=cfg.finetune_attention_modules,177        finetune_mlp_modules=cfg.finetune_mlp_modules,178        use_gradient_checkpointing=cfg.use_gradient_checkpointing,179        random_state=cfg.lora_random_state,180    )181 182    return peft_model, tokenizer183 184 185__all__ = [186    "ALLOWED_HARDWARE",187    "BASE_MODEL_ID",188    "BF16SlippageError",189    "BootConfig",190    "FINETUNE_ATTENTION_MODULES",191    "FINETUNE_LANGUAGE_LAYERS",192    "FINETUNE_MLP_MODULES",193    "FINETUNE_VISION_LAYERS",194    "FP16SlippageError",195    "HardwareT",196    "LORA_ALPHA",197    "LORA_DROPOUT",198    "LORA_R",199    "LORA_RANDOM_STATE",200    "MAX_SEQ_LENGTH",201    "assert_dtype_for_hardware",202    "assert_fp16_dtype",203    "boot_gemma",204]205