Spaces:
Running on Zero
Running on Zero
| """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"] | |