diff --git a/src/axolotl/integrations/protrain/api/model_wrapper.py b/src/axolotl/integrations/protrain/api/model_wrapper.py index 0c99ca5338..46915f420a 100644 --- a/src/axolotl/integrations/protrain/api/model_wrapper.py +++ b/src/axolotl/integrations/protrain/api/model_wrapper.py @@ -69,6 +69,22 @@ # 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: + """True if the model has a vision tower the text-only profiling batch skips.""" + 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) @@ -1902,9 +1918,10 @@ def protrain_model_wrapper( """Compose the ProTrain runtime around a standard ``nn.Module``. ``forbid_activation_offload``: when True, the searcher refuses any - CostConfig with ``n_offload > 0``. Set from ``cfg.lora_mlp_kernel`` — - the fused MLP backward kernel is incompatible with chunk-storage - placeholders. + CostConfig with ``n_offload > 0``. Set from + ``cfg.protrain_lora_mlp_forbid_offload`` (default False) — the fused MLP + backward kernel now composes with offload via shape-preserving + placeholders, so this is opt-in rather than implied by lora_mlp_kernel. ``prefer_no_offload_on_non_nvlink``: defensive searcher tie-break; auto-prefers ``n_offload=0`` configs within a 5% noise band when the @@ -2204,9 +2221,16 @@ 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: - capacity_bytes = max( - 0, int(hardware_profile.gpu_memory_bytes) - _DEFAULT_HEADROOM_BYTES - ) + _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) - _headroom) else: capacity_bytes = _capacity_bytes_caller if _cpu_capacity_bytes_caller is None: diff --git a/src/axolotl/integrations/protrain/args.py b/src/axolotl/integrations/protrain/args.py index ab750b4050..15dddb47b0 100644 --- a/src/axolotl/integrations/protrain/args.py +++ b/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 diff --git a/src/axolotl/integrations/protrain/plugin.py b/src/axolotl/integrations/protrain/plugin.py index 20057183ec..460142d896 100644 --- a/src/axolotl/integrations/protrain/plugin.py +++ b/src/axolotl/integrations/protrain/plugin.py @@ -1642,6 +1642,39 @@ 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 +1697,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.