Spaces:
Running on Zero
Running on Zero
File size: 3,598 Bytes
311fd09 127a1e8 311fd09 127a1e8 311fd09 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """Small PyTorch helpers shared by offline autoencoder adapters."""
from __future__ import annotations
import fcntl
import os
from pathlib import Path
def default_device(torch) -> str:
if not torch.cuda.is_available():
return "cpu"
if torch.cuda.device_count() == 1:
return "cuda:0"
local_rank = int(os.environ.get("LOCAL_RANK", os.environ.get("SLURM_LOCALID", "0")))
return f"cuda:{local_rank}"
def materialize_weight_norm(torch, model) -> None:
"""Replace both legacy and parametrized weight normalization in eval models."""
for module in model.modules():
try:
if torch.nn.utils.parametrize.is_parametrized(module, "weight"):
torch.nn.utils.parametrize.remove_parametrizations(module, "weight", leave_parametrized=True)
continue
except (AttributeError, ValueError):
pass
try:
torch.nn.utils.remove_weight_norm(module)
except (AttributeError, ValueError):
pass
def compile_fixed(
torch,
function,
warmup,
*,
enabled: bool,
cache_name: str,
strategy: str = "inductor",
):
"""Optimize one fixed CUDA shape with Inductor or a CUDA Graph.
CUDA Graphs retain normal PyTorch kernels when discrete output stability is
important. Inductor is faster for continuous encoders, where small floating
point differences are expected and measured.
"""
warmup_args = warmup if isinstance(warmup, tuple) else (warmup,)
device = warmup_args[0].device
if strategy not in {"inductor", "cuda_graph"}:
raise ValueError(f"unknown fixed-shape strategy: {strategy}")
if device.type != "cuda":
with torch.inference_mode():
function(*warmup_args)
return None
with torch.inference_mode():
function(*warmup_args)
torch.cuda.synchronize(device)
if not enabled:
return None
if strategy == "inductor":
cache_root = Path(os.environ.get("TMPDIR", "/tmp"))
cache_dir = cache_root / f"bingua2-{cache_name}-inductor-u{os.getuid()}"
os.environ.setdefault("TORCHINDUCTOR_CACHE_DIR", str(cache_dir))
lock_path = cache_root / f"{cache_dir.name}.lock"
with lock_path.open("a+b") as compile_lock:
fcntl.flock(compile_lock, fcntl.LOCK_EX)
optimized = torch.compile(function, mode="reduce-overhead", fullgraph=True)
optimized(*warmup_args)
# reduce-overhead may defer its CUDA-graph recording until the
# next invocation; complete it before processing timers start.
optimized(*warmup_args)
torch.cuda.synchronize(device)
return optimized
for _ in range(2):
output = function(*warmup_args)
torch.cuda.synchronize(device)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
output = function(*warmup_args)
torch.cuda.synchronize(device)
def replay(*args):
if len(args) != len(warmup_args):
raise ValueError("fixed CUDA graph received the wrong argument count")
for static, current in zip(warmup_args, args):
if static.shape != current.shape or static.dtype != current.dtype:
raise ValueError("fixed CUDA graph received an incompatible tensor")
static.copy_(current)
graph.replay()
return output
return replay
__all__ = ["compile_fixed", "default_device", "materialize_weight_norm"]
|