scripts/offload_policy.py

from __future__ import annotations


def enable_h3_cpu_offload(
    components_manager, *, memory_reserve_margin: str = "64GB"
) -> dict[str, str]:
    device = "cuda"
    components_manager.enable_auto_cpu_offload(
        device=device,
        memory_reserve_margin=memory_reserve_margin,
    )
    return {
        "mode": "components_manager_auto_cpu_offload",
        "device": device,
        "memory_reserve_margin": memory_reserve_margin,
    }


def component_device(module) -> str:
    tensor = next(module.parameters(), None)
    if tensor is None:
        tensor = next(module.buffers(), None)
    return "none" if tensor is None else str(tensor.device)