wehub-resource-sync
38 lines
1.3 KiB
Python
38 lines
1.3 KiB
Python
# Copyright (c) DeepSpeed Team.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# DeepSpeed Team
|
|
|
|
import pytest
|
|
import torch
|
|
from torch.fx import Graph
|
|
|
|
from deepspeed.compile.fx import add_end_backward, replace_reduce_outputs_with_none, get_output_node
|
|
from deepspeed.compile.util import get_deepcompile_handle, is_deepcompile_supported
|
|
|
|
|
|
@pytest.mark.skipif(not is_deepcompile_supported(), reason="DeepCompile requires CUDA and supported PyTorch")
|
|
@pytest.mark.sequential
|
|
def test_end_backward_depends_on_all_reduce_nodes():
|
|
get_deepcompile_handle()
|
|
|
|
graph = Graph()
|
|
grad = graph.placeholder("grad")
|
|
reduce_a = graph.create_node("call_function", torch.ops.dc.reduce_grad.default, (grad, 7, 11), name="reduce_a")
|
|
reduce_b = graph.create_node("call_function", torch.ops.dc.reduce_grad.default, (grad, 7, 12), name="reduce_b")
|
|
graph.output((grad, ))
|
|
|
|
add_end_backward(graph, 7)
|
|
replace_reduce_outputs_with_none(graph)
|
|
graph.lint()
|
|
|
|
end_backward = next(n for n in graph.nodes if n.target == torch.ops.dc.end_backward.default)
|
|
deps, graph_id, _ = end_backward.args
|
|
output_node = get_output_node(graph)
|
|
|
|
assert graph_id == 7
|
|
assert list(deps) == [reduce_a, reduce_b]
|
|
assert end_backward in reduce_a.users
|
|
assert end_backward in reduce_b.users
|
|
assert output_node.args == ((grad, ), )
|