From 461f3d2f9e6ca0efd061a60234a53ec881b2c93b Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 1 Jun 2026 11:34:30 -0700 Subject: [PATCH 1/2] fix(protrain): auto-offload memory-bound VL models; compose lora_mlp_kernel with offload Three fixes so ProTrain auto-mode trains a memory-bound VL model (e.g. Qwen3.5-27B QLoRA on 2x24GB) with no manual tuning: - The VL vision tower is invisible to the text-only profiling batch, so estimate_peak under-counts ~3.7 GiB and the searcher picks an all-resident config that OOMs. Reserve extra capacity headroom when a vision_config / visual module is present so it offloads instead. - Lift the lora_mlp_kernel => no-offload pin (now opt-in via protrain_lora_mlp_forbid_offload). The v61 LoRA_MLPBackward shape mismatch it guarded is already fixed by the unconditional shape-preserving placeholders, so the fused MLP kernel composes with offload (~3% faster per step, validated). - Complete a partial residency override with safe defaults + warn instead of silently falling back to the auto search; document the all-four-knob requirement on the override fields. Validated on Qwen3.5-27B QLoRA seq-1024, 2x3090 Ti NVLink: auto-mode now offloads (n_persist=34) and trains at 5.05 s/it, vs OOM before. --- .../protrain/api/model_wrapper.py | 25 ++++++++++- src/axolotl/integrations/protrain/args.py | 40 ++++++++++++++--- src/axolotl/integrations/protrain/plugin.py | 43 ++++++++++++++++--- 3 files changed, 96 insertions(+), 12 deletions(-) diff --git a/src/axolotl/integrations/protrain/api/model_wrapper.py b/src/axolotl/integrations/protrain/api/model_wrapper.py index 0c99ca5338..256d10fddb 100644 --- a/src/axolotl/integrations/protrain/api/model_wrapper.py +++ b/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 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..2cfc9d9ffb 100644 --- a/src/axolotl/integrations/protrain/plugin.py +++ b/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. From 5cc43e545eb02c74cc445aa1841215fdaeed0c67 Mon Sep 17 00:00:00 2001 From: thad0ctor Date: Mon, 1 Jun 2026 11:47:24 -0700 Subject: [PATCH 2/2] style(protrain): ruff format + fix docstrings per review - ruff format model_wrapper.py and plugin.py (CI formatting drift) - update protrain_model_wrapper docstring: forbid_activation_offload now reads from protrain_lora_mlp_forbid_offload, not cfg.lora_mlp_kernel - add docstring to _has_unprofiled_vision_tower --- .../integrations/protrain/api/model_wrapper.py | 13 +++++++------ src/axolotl/integrations/protrain/plugin.py | 7 +++++-- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/src/axolotl/integrations/protrain/api/model_wrapper.py b/src/axolotl/integrations/protrain/api/model_wrapper.py index 256d10fddb..46915f420a 100644 --- a/src/axolotl/integrations/protrain/api/model_wrapper.py +++ b/src/axolotl/integrations/protrain/api/model_wrapper.py @@ -75,6 +75,7 @@ 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 @@ -83,6 +84,7 @@ def _has_unprofiled_vision_tower(model: "nn.Module") -> bool: return True return False + # Slack for allocator frag, framework working set, dataloader workers. _DEFAULT_CPU_HEADROOM_BYTES = 2 * (1 << 30) @@ -1916,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 @@ -2227,9 +2230,7 @@ def protrain_model_wrapper( "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 - ) + 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/plugin.py b/src/axolotl/integrations/protrain/plugin.py index 2cfc9d9ffb..460142d896 100644 --- a/src/axolotl/integrations/protrain/plugin.py +++ b/src/axolotl/integrations/protrain/plugin.py @@ -1668,8 +1668,11 @@ def post_model_load(self, cfg, model: "nn.Module") -> None: "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, + ", ".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.