wehub-resource-sync
118 lines
3.4 KiB
Python
118 lines
3.4 KiB
Python
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License").
|
|
# You may not use this file except in compliance with the License.
|
|
# To view a copy of this license, visit http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# No warranties are given. The work is provided "AS IS", without warranty of any kind, express or implied.
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
|
|
def _is_main_process() -> bool:
|
|
return not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0
|
|
|
|
|
|
def _log_once(message: str) -> None:
|
|
if _is_main_process():
|
|
print(message)
|
|
|
|
|
|
class SafeCompiledCallable:
|
|
"""Lazy torch.compile wrapper that falls back to eager on compile/runtime errors."""
|
|
|
|
def __init__(
|
|
self,
|
|
fn,
|
|
*,
|
|
name: str,
|
|
backend: str = "inductor",
|
|
mode: str | None = "max-autotune-no-cudagraphs",
|
|
fullgraph: bool = False,
|
|
dynamic: bool | None = False,
|
|
options: dict | None = None,
|
|
suppress_errors: bool = True,
|
|
) -> None:
|
|
self.fn = fn
|
|
self.name = name
|
|
self.enabled = True
|
|
self.failed = False
|
|
self.failure_reason = None
|
|
|
|
if suppress_errors:
|
|
try:
|
|
import torch._dynamo as torch_dynamo
|
|
|
|
torch_dynamo.config.suppress_errors = True
|
|
except Exception as exc:
|
|
_log_once(f"[torch.compile] Could not enable suppress_errors: {exc}")
|
|
|
|
compile_kwargs = {
|
|
"backend": backend,
|
|
"fullgraph": fullgraph,
|
|
"dynamic": dynamic,
|
|
}
|
|
if mode:
|
|
compile_kwargs["mode"] = mode
|
|
if options:
|
|
compile_kwargs["options"] = options
|
|
|
|
_log_once(
|
|
"[torch.compile] Preparing "
|
|
f"{name}: backend={backend}, mode={mode}, "
|
|
f"fullgraph={fullgraph}, dynamic={dynamic}"
|
|
)
|
|
self.compiled_fn = torch.compile(fn, **compile_kwargs)
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
if not self.enabled:
|
|
return self.fn(*args, **kwargs)
|
|
|
|
try:
|
|
return self.compiled_fn(*args, **kwargs)
|
|
except Exception as exc:
|
|
self.enabled = False
|
|
self.failed = True
|
|
self.failure_reason = repr(exc)
|
|
_log_once(
|
|
f"[torch.compile][warn] {self.name} failed; "
|
|
f"falling back to eager. reason={exc}"
|
|
)
|
|
return self.fn(*args, **kwargs)
|
|
|
|
|
|
def configure_module_call_torch_compile(
|
|
module,
|
|
*,
|
|
name: str,
|
|
backend: str = "inductor",
|
|
mode: str | None = "max-autotune-no-cudagraphs",
|
|
fullgraph: bool = False,
|
|
dynamic: bool | None = False,
|
|
options: dict | None = None,
|
|
suppress_errors: bool = True,
|
|
):
|
|
if not torch.cuda.is_available():
|
|
_log_once(f"[torch.compile] Skipping {name}: CUDA is not available")
|
|
return None
|
|
|
|
try:
|
|
return SafeCompiledCallable(
|
|
module,
|
|
name=name,
|
|
backend=backend,
|
|
mode=mode,
|
|
fullgraph=fullgraph,
|
|
dynamic=dynamic,
|
|
options=options,
|
|
suppress_errors=suppress_errors,
|
|
)
|
|
except Exception as exc:
|
|
_log_once(
|
|
f"[torch.compile][warn] Could not prepare {name}; "
|
|
f"continuing in eager mode. reason={exc}"
|
|
)
|
|
return None
|