Honor memory_mode over legacy cpu_offload and prefetch dense-quant transformer shards
plan_diffusion_memory only applies the legacy cpu_offload override when no memory_mode was supplied, matching the documented API contract that memory_mode overrides cpu_offload when set; an explicit fast request now stays resident even if the old flag is also enabled. The transformer-quant dense path fetches the base repo's transformer/ shards inside the locked finalize phase, where unload and cancellation cannot preempt the multi-GB download. The load worker now widens the preemptible prefetch to include those shards when that path can actually run: quant requested and supported for the device, scheme resolvable, and no pre-quantized checkpoint shortcutting the dense build.
This commit is contained in:
parent
dd792c6312
commit
a4277a01e4
4 changed files with 122 additions and 7 deletions
|
|
@ -270,6 +270,21 @@ def test_explicit_cpu_offload_overrides_resident_auto_choice():
|
|||
assert any("explicit cpu_offload" in r for r in plan.reasons)
|
||||
|
||||
|
||||
def test_explicit_memory_mode_wins_over_legacy_cpu_offload():
|
||||
# The API documents memory_mode as overriding cpu_offload when set: fast +
|
||||
# the legacy flag must stay resident, not silently downgrade to offload.
|
||||
plan = plan_diffusion_memory(
|
||||
target = _target(),
|
||||
device_memory = _discrete(80000),
|
||||
model_dense_mib = 4000,
|
||||
runtime_headroom_mib = 2000,
|
||||
requested_mode = MEMORY_MODE_FAST,
|
||||
explicit_offload = True,
|
||||
)
|
||||
assert plan.offload_policy == OFFLOAD_NONE
|
||||
assert not any("explicit cpu_offload" in r for r in plan.reasons)
|
||||
|
||||
|
||||
def test_explicit_cpu_offload_ignored_on_cpu_target():
|
||||
plan = plan_diffusion_memory(
|
||||
target = _target(device = "cpu", backend = "cpu", supports_offload = False),
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue