fix(protrain): auto-offload memory-bound VL models; compose lora_mlp_kernel with offload

src/axolotl/integrations/protrain/api/model_wrapper.py

@@ -69,6 +69,20 @@
 # Reserve 2 GiB for CUDA context + allocator overhead.
 _DEFAULT_HEADROOM_BYTES = 2 * (1 << 30)
 
+# A vision tower the text-only profiling batch never exercises is invisible to
+# estimate_peak, so reserve extra capacity to force offload instead of OOM.
+_VL_VISION_HEADROOM_BYTES = 4 * (1 << 30)
+
+
+def _has_unprofiled_vision_tower(model: "nn.Module") -> bool:
+    cfg = getattr(model, "config", None)
+    if cfg is not None and getattr(cfg, "vision_config", None) is not None:
+        return True
+    for name, _ in model.named_modules():
+        if name.rsplit(".", 1)[-1] in {"visual", "vision_tower", "vision_model"}:
+            return True
+    return False
+
 # Slack for allocator frag, framework working set, dataloader workers.
 _DEFAULT_CPU_HEADROOM_BYTES = 2 * (1 << 30)
 
@@ -2204,8 +2218,17 @@ def protrain_model_wrapper(
     # Now derive capacity_bytes / cpu_capacity_bytes from the broadcast-synchronized
     # hardware_profile so every rank computes the identical searcher capacity-cutoff.
     if _capacity_bytes_caller is None:
+        _headroom = _DEFAULT_HEADROOM_BYTES
+        if _has_unprofiled_vision_tower(model):
+            _headroom += _VL_VISION_HEADROOM_BYTES
+            LOG.warning(
+                "ProTrain: vision tower not exercised by the text-only profiling "
+                "batch; reserving an extra %.1f GiB capacity headroom so the "
+                "searcher offloads instead of OOMing in the vision forward.",
+                _VL_VISION_HEADROOM_BYTES / (1 << 30),
+            )
         capacity_bytes = max(
-            0, int(hardware_profile.gpu_memory_bytes) - _DEFAULT_HEADROOM_BYTES
+            0, int(hardware_profile.gpu_memory_bytes) - _headroom
         )
     else:
         capacity_bytes = _capacity_bytes_caller

src/axolotl/integrations/protrain/args.py

@@ -186,26 +186,43 @@ class ProTrainArgs(BaseModel):
         ge=0,
         json_schema_extra={
             "description": (
-                "Debug override: force the number of persistent chunks. "
-                "Bypasses the exhaustive searcher when set alongside the other "
-                "three overrides."
+                "Debug override: force the number of persistent chunks. The "
+                "searcher is bypassed only when all four of "
+                "n_persist/n_buffer/n_swap/n_checkpoint are set; a partial set "
+                "is auto-completed with defaults (n_persist=0, n_buffer=2, "
+                "n_swap=0, n_checkpoint=0) and logs a warning."
             )
         },
     )
     protrain_n_buffer_override: int | None = Field(
         default=None,
         ge=0,
-        json_schema_extra={"description": "Debug override for n_buffer."},
+        json_schema_extra={
+            "description": (
+                "Debug override for n_buffer. Effective only as part of the full "
+                "4-knob override set (see protrain_n_persist_override)."
+            )
+        },
     )
     protrain_n_swap_override: int | None = Field(
         default=None,
         ge=0,
-        json_schema_extra={"description": "Debug override for n_swap."},
+        json_schema_extra={
+            "description": (
+                "Debug override for n_swap. Effective only as part of the full "
+                "4-knob override set (see protrain_n_persist_override)."
+            )
+        },
     )
     protrain_n_checkpoint_override: int | None = Field(
         default=None,
         ge=0,
-        json_schema_extra={"description": "Debug override for n_checkpoint."},
+        json_schema_extra={
+            "description": (
+                "Debug override for n_checkpoint. Effective only as part of the "
+                "full 4-knob override set (see protrain_n_persist_override)."
+            )
+        },
     )
     protrain_n_offload_override: int | None = Field(
         default=None,
@@ -506,6 +523,17 @@ class ProTrainArgs(BaseModel):
             )
         },
     )
+    protrain_lora_mlp_forbid_offload: bool = Field(
+        default=False,
+        json_schema_extra={
+            "description": (
+                "Re-enable the conservative pin that refuses n_offload>0 "
+                "candidates when lora_mlp_kernel is on. Off by default since "
+                "shape-preserving placeholders let the fused MLP kernel compose "
+                "with offload."
+            )
+        },
+    )
 
     # ------------------------------------------------------------------
     # Validators

src/axolotl/integrations/protrain/plugin.py

@@ -1642,6 +1642,36 @@ def post_model_load(self, cfg, model: "nn.Module") -> None:
         n_offload_override = getattr(cfg, "protrain_n_offload_override", None)
         zero3_shard = getattr(cfg, "protrain_zero3_shard", None)
 
+        # The wrapper bypasses the searcher only when all four of
+        # n_persist/n_buffer/n_swap/n_checkpoint are set; a partial set would
+        # otherwise silently fall back to the auto search. Complete it with safe
+        # defaults so a partial override does what the user intended.
+        _core_overrides = {
+            "protrain_n_persist_override": n_persist_override,
+            "protrain_n_buffer_override": n_buffer_override,
+            "protrain_n_swap_override": n_swap_override,
+            "protrain_n_checkpoint_override": n_checkpoint_override,
+        }
+        _set = [k for k, v in _core_overrides.items() if v is not None]
+        if 0 < len(_set) < 4:
+            if n_persist_override is None:
+                n_persist_override = 0
+            if n_buffer_override is None:
+                n_buffer_override = 2
+            if n_swap_override is None:
+                n_swap_override = 0
+            if n_checkpoint_override is None:
+                n_checkpoint_override = 0
+            LOG.warning(
+                "ProTrain: partial residency override (set: %s). The override "
+                "path needs all four of n_persist/n_buffer/n_swap/n_checkpoint "
+                "or it silently falls back to the auto search. Filled unset "
+                "knobs: n_persist=%d n_buffer=%d n_swap=%d n_checkpoint=%d. Set "
+                "all four explicitly to suppress this.",
+                ", ".join(_set), n_persist_override, n_buffer_override,
+                n_swap_override, n_checkpoint_override,
+            )
+
         # auto_mode default True; wrapper picks (force_persist, zero3) post-search.
         auto_mode = getattr(cfg, "protrain_auto_mode", True)
         if auto_mode is None:
@@ -1664,13 +1694,16 @@ def post_model_load(self, cfg, model: "nn.Module") -> None:
                 "(force_all_persistent=False, zero3_shard=False)."
             )
 
-        # Fused LoRA MLP backward kernel + offloaded-activation chunk placeholders
-        # crash with LoRA_MLPBackward shape-mismatch (v61); pin n_offload=0.
-        forbid_activation_offload = bool(getattr(cfg, "lora_mlp_kernel", False))
+        # The v61 LoRA_MLPBackward shape-mismatch is fixed by the unconditional
+        # shape-preserving placeholders, so the fused MLP kernel now composes
+        # with offload; opt back into the old conservative pin if needed.
+        forbid_activation_offload = bool(
+            getattr(cfg, "protrain_lora_mlp_forbid_offload", False)
+        )
         if forbid_activation_offload:
             LOG.info(
-                "ProTrain: cfg.lora_mlp_kernel=True; searcher will refuse "
-                "n_offload>0 candidates."
+                "ProTrain: protrain_lora_mlp_forbid_offload=True; searcher will "
+                "refuse n_offload>0 candidates."
             )
 
         # PR #17c: defensive searcher tie-break on non-NVLink multi-rank rigs.