Merge remote-tracking branch 'origin/main' into forward_kwargs
# Conflicts: # src/petals/__init__.py # src/petals/client/inference_session.pypull/467/head
Your Name
5 months ago
commit
3195579620
@ -0,0 +1,76 @@
|
|||||||
|
import torch
|
||||||
|
from torch.utils._pytree import tree_flatten as _tree_flatten, tree_unflatten as _tree_unflatten
|
||||||
|
|
||||||
|
|
||||||
|
def make_inference_graphed_callable(callable: callable, sample_args, num_warmup_iters=3):
|
||||||
|
"""Similar to torch.cuda.make_graphed_callables, but takes only one function and does not build a graph for the backward pass"""
|
||||||
|
assert not isinstance(callable, torch.nn.Module)
|
||||||
|
if torch.is_autocast_enabled() and torch.is_autocast_cache_enabled():
|
||||||
|
raise RuntimeError(
|
||||||
|
"make_graphed_callables does not support the autocast caching. Please set `cache_enabled=False`."
|
||||||
|
)
|
||||||
|
|
||||||
|
flatten_arg, _ = _tree_flatten(sample_args)
|
||||||
|
flatten_sample_args = tuple(flatten_arg)
|
||||||
|
assert all(
|
||||||
|
isinstance(arg, torch.Tensor) for arg in flatten_arg
|
||||||
|
), "In the beta API, sample_args for each callable must contain only Tensors. Other types are not allowed."
|
||||||
|
|
||||||
|
len_user_args = len(sample_args)
|
||||||
|
static_input_surface = flatten_sample_args
|
||||||
|
|
||||||
|
graph = torch.cuda.CUDAGraph()
|
||||||
|
|
||||||
|
# Warmup
|
||||||
|
# Hopefully prevents cudnn benchmarking and other lazy-initialization cuda work
|
||||||
|
# from ending up in any captures.
|
||||||
|
s = torch.cuda.Stream()
|
||||||
|
s.wait_stream(torch.cuda.current_stream())
|
||||||
|
with torch.cuda.stream(s):
|
||||||
|
for _ in range(num_warmup_iters):
|
||||||
|
outputs, _ = _tree_flatten(callable(*sample_args))
|
||||||
|
del outputs
|
||||||
|
torch.cuda.current_stream().wait_stream(s)
|
||||||
|
|
||||||
|
# Capture forward graph
|
||||||
|
with torch.cuda.graph(graph):
|
||||||
|
outputs = callable(*sample_args)
|
||||||
|
|
||||||
|
flatten_outputs, output_unflatten_spec = _tree_flatten(outputs)
|
||||||
|
static_outputs = tuple(flatten_outputs)
|
||||||
|
|
||||||
|
def make_graphed_function(
|
||||||
|
graph,
|
||||||
|
len_user_args,
|
||||||
|
output_unflatten_spec,
|
||||||
|
static_input_surface,
|
||||||
|
static_outputs,
|
||||||
|
):
|
||||||
|
def replay_graph(*inputs):
|
||||||
|
# At this stage, only the user args may (potentially) be new tensors.
|
||||||
|
for i in range(len_user_args):
|
||||||
|
if static_input_surface[i].data_ptr() != inputs[i].data_ptr():
|
||||||
|
static_input_surface[i].copy_(inputs[i])
|
||||||
|
graph.replay()
|
||||||
|
assert isinstance(static_outputs, tuple)
|
||||||
|
return tuple(o.detach() for o in static_outputs)
|
||||||
|
|
||||||
|
def functionalized(*user_args):
|
||||||
|
# Runs the autograd function with inputs == all inputs to the graph that might require grad
|
||||||
|
# (explicit user args + module parameters)
|
||||||
|
# Assumes module params didn't change since capture.
|
||||||
|
flatten_user_args, _ = _tree_flatten(user_args)
|
||||||
|
out = replay_graph(*flatten_user_args)
|
||||||
|
return _tree_unflatten(out, output_unflatten_spec)
|
||||||
|
|
||||||
|
return functionalized
|
||||||
|
|
||||||
|
# Put together the final graphed callable
|
||||||
|
graphed = make_graphed_function(
|
||||||
|
graph,
|
||||||
|
len_user_args,
|
||||||
|
output_unflatten_spec,
|
||||||
|
static_input_surface,
|
||||||
|
static_outputs,
|
||||||
|
)
|
||||||
|
return graphed
|
||||||
Loading…
Reference in New Issue